Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Operator collections arraylias integration #291

Conversation

DanPuzzuoli
Copy link
Collaborator

@DanPuzzuoli DanPuzzuoli commented Nov 24, 2023

Summary

Updating operator_collections.py to work with arraylias.

General list of changes:

  • All operator collections corresponding to numpy, jax, and jax_sparse are merged into a single class using arraylias dispatching (one each for OperatorCollection, LindbladCollection, and VectorizedLindbladCollection. However, I've decided to keep scipy_sparse as a separate special case in each instance, as the interface is just too different from the others. It would be nice to have it all handled in a single class, but I'm worried that this would require building too much scipy sparse into our numpy alias that wouldn't be used anywhere else.
  • I've added registration of some more functions for "jax_sparse": conjugate and transpose.
  • I've added a linear_combo function to the registration for "numpy", "jax", and "jax_sparse". This is just tensordot for the first two but for jax_sparse a somewhat-hacky implementation is necessary.
  • Updated _preferred_lib to prefer "jax_sparse" if present.
  • The model classes (GeneratorModel, HamiltonianModel, and LindbladModel) are currently broken. I've commented out the lines using the old versions of the operator collection classes to avoid errors. I'll address these in the next PR updating these classes.
  • I've rearranged the tests for operator collections to reflect the new class structure. I tried to keep track and make sure all of the tests got transferred over.

One more general change:

  • This is somewhat orthogonal to the main point of this PR, but as we are making wholesale changes it make sense to address it now: I'm taking this opportunity to make the operator collections AND the model classes immutable. This should simplify various things, even in the Solver class, in which signals need to be added to, and removed from, models. This PR won't be the full change, but I'm going to mark this PR as closing Parameter dependent simulation for optimal bayesian experimentation - JAX @jit compatibility #243, in which the immutability of model classes would have prevented the weird behaviour from occurring in the first place. For this PR, that means getting rid of all of the setter methods for the different operator types. I will do this in the model classes in a subsequent PR as well.

Details and comments

Current list of test files we need to make sure are passing before merge to feature branch:

  • test_rotating_frame.py
  • test_operator_collections.py
  • test_alias.py

We should keep track of this list in subsequent PRs.

One technical issue:

  • Unfortunately, for LindbladCollection, it isn't possible to directly pass JAX sparse arrays at construction. This was always the case, but I had hoped to change that here. The handling of the n_batch argument at BCOO construction needs to be set to a particular value to get it to work here, and this isn't something I want the user to need to do. For now, they will need to pass dense arrays and select array_library="jax_sparse".

@DanPuzzuoli DanPuzzuoli marked this pull request as ready for review December 2, 2023 00:32
Copy link
Contributor

@to24toro to24toro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is definitely better code than before.
I think the addition of array_library and the use of arraylias has improved the maintainability of the code considerably.

qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
qiskit_dynamics/models/operator_collections.py Outdated Show resolved Hide resolved
test/dynamics/arraylias/test_alias.py Outdated Show resolved Hide resolved
@DanPuzzuoli DanPuzzuoli requested a review from to24toro December 4, 2023 22:49
Copy link
Contributor

@to24toro to24toro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you, Dan. I've approved.

@DanPuzzuoli DanPuzzuoli merged commit c0d6b95 into qiskit-community:models_arraylias_integration Dec 5, 2023
2 checks passed
DanPuzzuoli added a commit that referenced this pull request Jan 19, 2024
DanPuzzuoli added a commit that referenced this pull request Jan 24, 2024
DanPuzzuoli added a commit that referenced this pull request Feb 5, 2024
DanPuzzuoli added a commit that referenced this pull request Feb 20, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants