[WIP] Arraylias integration -Model-

Summary

Updating Model folders

~generator_model~
~hamiltonian_model~
~lindblad_model~
~operator_collections~
~rotating_frame~
~rotating_wave_approximation~

Details and comments

[TODO]

check if we can deprecate to_numeric_matrix_type

[Points to discuss]

1. The necessity of `to_numeric_matrix_type`

Can its role be effectively covered by asarray?

2. Registering `list` type to alias

The list type will be used for List[sparse_matrix], and it may be replaced if a new class like CsrMatrixList is defined. However, it's unclear where this type is used. Conversion to CsrMatrixList or BCOOList is necessary in certain places. In dynamic, dealing with List[Sparse matrix], and it's uncertain whether registering CsrMatrixList to scipy_sparse and custom functions in scipy_sparse can adequately cover the role of List[Sparse matrix]. (Ref: Link)

3. Replacing `Array.default_backend() == jax`

Array will be deprecated, and an alternative method needs consideration. The current code implementation involves checking if operators is related to jax using is_operators_jax = any(lib in ("jax", "jax_sparse") for lib in DYNAMICS_NUMPY_ALIAS.infer_libs(operators)). Functionalizing this process is desired.

4. `isinstance(, ArrayLike)`

Encounter TypeError: Subscripted generics cannot be used with class and instance checks. The aim is to use this if possible.

5. `unp.to_dense`

Assuming that Jax-related functions or classes can be used only when using jnp.array(). For instance, the difference between to_BCOO and unp.to_sparse. to_BCOO covers BCOO even with any array, while unp.to_sparse only changes to BCOO if jnp.array() is used.

6. `Operator.from_label()`

Similar to the points above. Operator.from_label().data supports np.array(), and if using the Jax model, Operator.from_label() is utilized as jnp.array(Operator.from_label()).

1. The necessity of to_numeric_matrix_type

Can its role be effectively covered by asarray?

I think yes - as far as I can tell this is all to_numeric_matrix_type is doing (converting to one of the accepted types). It is worth trying to remove it and rebuild the module without it.

2. Registering list type to alias

The list type will be used for List[sparse_matrix], and it may be replaced if a new class like CsrMatrixList is defined. However, it's unclear where this type is used. Conversion to CsrMatrixList or BCOOList is necessary in certain places. In dynamic, dealing with List[Sparse matrix], and it's uncertain whether registering CsrMatrixList to scipy_sparse and custom functions in scipy_sparse can adequately cover the role of List[Sparse matrix]. (Ref: Link)

I don't remember all the places where List[sparse_matrix] is used. The only place I can remember, which I think is the only important place, is within operator_collections.py. I searched for uses of to_csr, and the only place it is used in the entire package (aside from tests) is within SparseOperatorCollection.

I think, unfortunately, handling of scipy CSR matrices is always going to be a weird edge case. I had wanted to handle these gracefully like all other array types, and to use arraylias to refactor operator_collections.py so that, e.g. DenseOperatorCollection, SparseOperatorCollection, and JAXSparseOperatorCollection could all be merged into a single class, but there may not be a completely clean way of doing this. It seems like a minor technical issue, but the fact that a "list of CSR matrices" is not nicely captured by an identifiable type is makes it hard to naturally fit it into the arraylias framework. I think the options are:

Create the type CsrMatrixList, which will simply wrap a numpy array of CSR matrices. This class could be completely contained within operator_collections.py, and would only require registering the few operations required in SparseOperatorCollection. If we do this, we could merge all of the OperatorCollection classes together.
Keep CSR matrices as a special case in operator_collections.py. In this case, we will be able to merge DenseOperatorCollection and JAXSparseOperatorCollection, but we will need to keep SparseOperatorCollection for performing operations with CSR matrices.

I'm not sure which of the two options I prefer. I originally wanted to implement the first to simplify the operator_collections.py file, but now I am willing to consider the second, as it will keep the arraylias infrastructure simpler.

3. Replacing Array.default_backend() == jax

Array will be deprecated, and an alternative method needs consideration. The current code implementation involves checking if operators is related to jax using is_operators_jax = any(lib in ("jax", "jax_sparse") for lib in DYNAMICS_NUMPY_ALIAS.infer_libs(operators)). Functionalizing this process is desired.

Yeah maybe we can just add an internal function like _preferred_lib. What is Array.default_backend() == jax this typically used for? Before checking if JAX tracing is happening?

4. isinstance(, ArrayLike)

Encounter TypeError: Subscripted generics cannot be used with class and instance checks. The aim is to use this if possible.

Is what you're trying to do here equivalent to:
DYNAMICS_NUMPY_ALIAS.infer_libs(x) != tuple()
We can't check isinstance(x, ArrayLike), but I guess what we're really checking is if DYNAMICS_NUMPY_ALIAS.infer_libs can identify the type of x?

5. unp.to_dense

Assuming that Jax-related functions or classes can be used only when using jnp.array(). For instance, the difference between to_BCOO and unp.to_sparse. to_BCOO covers BCOO even with any array, while unp.to_sparse only changes to BCOO if jnp.array() is used.

I'm actually starting to wonder if we even need to_BCOO, to_sparse, or to_dense at all. These are all primarily used in operator_collections.py for setting up OperatorCollection instances, which are controlled by the evaluation_mode argument when constructing model classes (GeneratorModel and related classes). We are going to need to change the way evaluation_mode works anyway, as currently the only options are "dense" or "sparse", and these work in conjunction with Array.default_backend(). As Array.default_backend() will no longer exist, I think we could change the options here to "numpy", "jax", "jax_sparse", or "scipy_sparse", and otherwise the default behaviour should just be to call asarray. Similarly with to_numeric_matrix_type it may make sense to try to not include to_sparse or to_dense initially and see if we can make things work without them.

As an aside: I see to_BCOO is used in rotating_frame.py in lines like

elif type(operator).__name__ == "BCOO":
    op_to_add_in_fb = to_BCOO(op_to_add_in_fb)

but I think lines like this could just be changed to

op_to_add_in_fb = alias(like=operator).asarray(op_to_add_in_fb)

so to_BCOO isn't needed here either.

6. Operator.from_label()

Similar to the points above. Operator.from_label().data supports np.array(), and if using the Jax model, Operator.from_label() is utilized as jnp.array(Operator.from_label()).

I think this is okay - if a user is creating an array using Operator.from_label() then that won't be JAX traceable anyway, so it's okay if unp.asarray(Operator.from_label("X")) returns a numpy array.

qiskit-community / qiskit-dynamics