Arraylias integration - Rotating frame -

[to24toro]

Summary

I started to migrate Arraylias to rotating_frame.py.
The first difficulty is dealing with csr_matrix and BCOO.

Details and comments

[TODO]

• how to deal with Operator. I think the way is whether converting to op.data or registering Operator to ArrayLike.
• how to use matmul in sparse matrix or BCOO.
• how to set spmatrix. Is it necessary to add some concrete classes of a sparse matrix.

[DanPuzzuoli]

The to_array, to_csr, to_BCOO, and to_numeric_matrix_type are things that we'll need to think carefully about when converting to arraylias. Dynamics has been very loose with array types, and that has been partly enabled by these functions, which are used throughout the package to funnel a variety of types (including lists, lists of arrays, etc) into array and sparse types. It would be nice to not need these methods, but I'm not sure to what extent we can avoid using them.

One option would be to just start being very strict about types. E.g. the RotatingFrame methods could all require the inputs be ArrayLike (assuming csr_matrix and BCOO become registered Arraylias types). The reason I hesitate with fully going this direction is that it's natural to specify operators in models in a variety of ways, either as raw arrays/sparse matrices, qiskit Operator/QuantumState instances, qutip Qobj, or lists of these. We could just force users to convert these to their desired array types in most places, and make the to_* functions into publically available helper functions. I wouldn't want to do this everywhere, e.g. in Solver.solve, there is some high level handling of QuantumState types that I wouldn't want to get rid of.

Aside from these concerns, I don't think we want to register csr_matrix or BCOO as straight 'numpy'/'jax' types. Both require special syntax, so I think it will be necessary to register them as their own libs. E.g., in arraylias_state.py, I think we can do something like (with suitable try blocks in case JAX isn't installed):

DYNAMICS_ALIAS.register_type(csr_matrix, lib="scipy_sparse")

"""
register required custom versions of functions for csr type here
"""

DYNAMICS_ALIAS.register_type(BCOO, lib="jax_sparse")
"""
register required custom versions of functions for BCOO type here
"""

I think for scipy, there are also multiple different sparse types that we'll need to register (they get transformed into each other). The scipy.sparse function issparse is used to check for all of these types. Maybe we can look there for all of the types.

Some considerations with functions that need to be registered:

• For some reason, the sparse types in scipy use * for matrix multiplication instead of @. There are places in RotatingFrame where the issparse checks for a sparse type, and a line is implemented using *. I think this type checking can be eliminated with Arraylias if we convert matrix multiplications in RotatingFrame to be written with unp.matmul, and we register lambda x, y: x * y as the definition of matmul for lib='scipy_sparse'.
• The BCOO type has been set up interestingly in JAX to use the sparsify transformation. This takes a JAX function meant to act on dense arrays, and converts it into one that also works with sparse arrays. (It doesn't work with absolutely everything in JAX, but it does for everything we'll care about.) Checkout the JAX documentation on this, and also the operator_collections.py file for examples in Dynamics. I think the things being registered in operator_collections.py can probably be registered in the main arraylias_state.py file to be used throughout the rest of Dynamics.
• The asarray method will need to be registered in both cases, and I think this is what your 'csr_matrix' registered function should be.

[CLAassistant]

All committers have signed the CLA.

[to24toro]

close and move to #282.