qiskit-community / qiskit-dynamics

Tools for building and solving models of quantum systems in Qiskit
https://qiskit-community.github.io/qiskit-dynamics/
Apache License 2.0
105 stars 61 forks source link

RotatingFrame.vectorized_map_into_frame is not JAX compatible and has ambiguous name #364

Open DanPuzzuoli opened 1 month ago

DanPuzzuoli commented 1 month ago

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 to unp. 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 equivalent vectorized_state method should also be added. There also isn't an out_of_frame version of this function, so one should be added for each of the operator and state methods.

DanPuzzuoli commented 6 days 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.