qutip / qutip-jax

JAX backend for QuTiP
BSD 3-Clause "New" or "Revised" License
15 stars 7 forks source link

Add a jax backend based on dia sparse matrix, `JaxDia`. #20

Closed Ericgig closed 3 months ago

Ericgig commented 1 year ago

Add a dia sparse data layer for Jax. JaxDia's data are jax.array, (differentiable), but the structure is meta data. This allow most functions to support jit. tidyup does not jit, so I did not include automated tidyup. The dia format is good for operator with few diagonals, passable for ket and horrible for bra. I made the function so it can be supported as the operators in mesolve.

Included specialisations:

Won't do:

Maybe:

19 is a branch off this one, and should be merged before this one.

coveralls commented 1 year ago

Pull Request Test Coverage Report for Build 8504122727

Details


Changes Missing Coverage Covered Lines Changed/Added Lines %
src/qutip_jax/ode.py 8 9 88.89%
src/qutip_jax/qutip_trees.py 3 4 75.0%
src/qutip_jax/binops.py 152 156 97.44%
src/qutip_jax/properties.py 40 44 90.91%
src/qutip_jax/create.py 49 60 81.67%
src/qutip_jax/jaxdia.py 89 102 87.25%
src/qutip_jax/qobjevo.py 83 107 77.57%
<!-- Total: 527 585 90.09% -->
Files with Coverage Reduction New Missed Lines %
src/qutip_jax/qobjevo.py 1 81.36%
src/qutip_jax/jaxarray.py 2 85.71%
<!-- Total: 3 -->
Totals Coverage Status
Change from base Build 7964056840: 0.06%
Covered Lines: 1195
Relevant Lines: 1321

💛 - Coveralls
Ericgig commented 3 months ago

We can change to PIDController before making a release, but not in this PR.