Add floquet
#681
Add floquet
#681
Conversation
…nstantly using conditional logic, and it made more sense to just split them into separate functions since they are doing different things. Code is running but have not updated the tests yet
…ifferentiation (still don't like the name tho)
added more tests, all passing
…o avoid repeated calls to e.g. _flat_vectorize which was yielding weird/incorrect results
also save T as part of FloquetResult added to mkdocs.yml
…he end use my definition of sigmap
|
Quick comment raised by @BenjaminDAnjou: Is there any other way to compute the Floquet basis @dkweiss31? Otherwise, maybe we could perform the workaround suggested in the linked issue, that is callback to the CPU implementation of |
|
Thanks for this, I hadn't appreciated that |
There was a problem hiding this comment.
Thanks a lot @dkweiss31 for the PR, and sorry again for the late review.
In the most recent commits, I propose several changes to your PR:
- Unified the implementation of
floquet_t0andfloquet_tin a single integrator. This is mainly such that we don't need to recompute the propagator multiple times. Here, it is computed only once both for the t=t0 modes and for the later propagation of the t=t modes. - Removed batching over
T. This is kinda annoying, but related to not having support for batching overtsave. I think it really simplifies the implementation. Again, I'm quite down to rethink this in a later PR. In the meantime, I would propose to simply provide an example in the main docstring to show how to batch over multiple drive periods using an externaljax.vmap. This is rather straightforward to achieve actually. - Enforced that
tsave[0]andtsave[-1]live within one period of each other. The other solution would indeed be to use ajnp.mod(tsave, T)to simplify the life of users, but for now, it simplifies our life haha. - Proposed to rename
quasienstoquasienergiesandfloquet_modestomodesin an effort to have short but meaningful names. Happy to reconsider this one if you think it was better previously. - Removed the
safeargument in favor of equinox runtime error checking, which is quite neat.
Let me know if all these changes work with you!
To finish the PR, I think there's mainly two things remaining:
- Adding the docstring example for batching over
T/tsave(drive period) - Fixing tests with this new version
… issue with jax versions >0.4.31 jax-ml/jax#24255 . It seems however that there is a GPU implementation of eig in the works jax-ml/jax#24663, making this pure callback stuff unnecesary
|
Hey @gautierronan , thank you very much for this review and for the significant simplifications you introduced. I am very happy with all of your proposed changes and will write up an example for T batching shortly. I had been working with jax 0.4.29, and upgraded based on your comment that the |
There was a problem hiding this comment.
Thanks a lot @dkweiss31! All good on my end! Let's wait for the review of @pierreguilmin and aim for merging next week :)
|
Sounds great, thanks very much @gautierronan ! |
- Nit fixes to documentation. - Add str representation for modes and quasienergies. - Add example of cartesian batching.
There was a problem hiding this comment.
Beautiful work, congratulations @dkweiss31, that's an amazing addition to the library! 😍 And @gautierronan I loved how your thorough review helped it converge into this clean and simple API. I'm adding a few nits on top. Merging now!
Closes DYN-231
Add Floquet solvers and associated tests. The API as it stands now supplies two functions:
floquetandfloquet_tfor obtaining the floquet modes (and associated quasienergies) at t=0, and obtaining the floquet modes at other times t, respectively. I'm not thrilled about the namefloquet_t, nor the name of e.g. the respective integratorFloquetIntegrator_t, so any suggestions for better names are welcome.