Open DanPuzzuoli opened 1 month ago
One thing to note here is that the "vectorized" transformations can be implemented by the standard RotatingFrame
methods but with the frame operator being the vectorized version of the map X -> [F, X]
. So, an elegant solution to implementing these methods may be to just internally construct a RotatingFrame
instance using the vectorized frame, and then call the ordinary methods.
The problem
This method is not JAX compatible (simple fix).
It is also "ambiguous" in the sense that "
vectorized_map
" does not indicate the nature of the transformation - a vectorized map in this context could either be the state of a differential equation or an operator that acts on states, and both are subject to different transformation rules. It currently implements the operator transformation, and there is no way to do a vectorized state transformation.Suggestion
JAX fix is easy - I think there is a single usage of
np
that can be changed tounp
. Along with this though, tests need to be added that would have caught this. I don't believe this method is being tested anywhere.For the ambiguity, this method should be renamed
vectorized_map
->vectorized_operator
, and the equivalentvectorized_state
method should also be added. There also isn't anout_of_frame
version of this function, so one should be added for each of theoperator
andstate
methods.