DanPuzzuoli
commented
1 year ago 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 inRotatingFramewhere theissparsechecks 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 inRotatingFrameto be written withunp.matmul, and we registerlambda x, y: x * yas the definition ofmatmulforlib='scipy_sparse'. - The
BCOOtype has been set up interestingly in JAX to use thesparsifytransformation. 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 inoperator_collections.pycan probably be registered in the mainarraylias_state.pyfile to be used throughout the rest of Dynamics. - The
asarraymethod will need to be registered in both cases, and I think this is what your'csr_matrix'registered function should be.
CLAassistant
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]
Operator. I think the way is whether converting toop.dataor registeringOperatorto ArrayLike.sparse matrixorBCOO.spmatrix. Is it necessary to add some concrete classes of a sparse matrix.