Added second-order Magnus solver in the fixed-step exponential solvers

[malekakhlagh]

Summary

Second order Magnus solver is added to the following fixed-step exponential solvers:
scipy_expm_solver
jax_expm_solver
jax_expm_parallel_solver

Details and comments

[CLAassistant]

All committers have signed the CLA.

[DanPuzzuoli]

Awesome, thanks for submitting this! There are a few things to address to proceed with this PR:

• Documentation
  • This optional argument needs to be described in the solve_lmde documentation (in solvers/solver_functions.py) as one of the optional parameters of the exponentiation-based solvers - something along the lines of "Additionally, the optional argument magnus_order provides control over...". We should also add a reference to whatever textbook or literature you used for this. To set this up properly this will require setting up some extra files associated with how the docs build (I've set this up in the perturbation PR, but it will likely take longer than this one).
• Testing
  • The tests for solvers are setup to make it "easy" to add tests for new solvers. Tests for this solver can be added by modifying the files test/dynamics/solvers/test_solver_classes.py and test/dynamics/solvers/test_fixed_step_solvers.py. I put "easy" in quotes as it may not be super obvious how to do this, but it can be done by adding new test classes that inherit from existing ones. All this new test class would need to do is modify the name of the solver being tested along with required options, then all of the existing test cases will automatically run on it.
• Formatting
  • This will be more important later in the PR process, but there are tools built into the repo for making sure that the code is written with a consistent style. To use these you'll need to install tox (pip install tox) and run commands tox -e black (which will auto format some things) and tox -e lint (which checks things that can't be auto-formatted). For a PR to be mergable the command tox -e lint needs to run without raising any warnings.

We can discuss more about each individual point - I'm happy to help with any part of this, especially the parts involving more detailed knowledge of the repo itself (e.g. setting up the proper citation method or adding tests). Aside from this, I also have some direct comments on the code itself:

• Adding the magnus_order == 2 case introduces many new lines of code, and repeating this in every expm-based solver is redundant (especially if even higher order methods are added). I think the setup of the take_step function for exponentiation-based solver should be moved to a separate function, that each solver can call. I'm imagining something of the form:

def get_exponential_take_step(magnus_order, expm_func):
    
    # put the if-statements here for constructing the function 
    
    # return the function
    return take_step

Here, expm_func is included as an argument so that each solver can pass whatever matrix exponential function it wants, e.g. either jexpm or expm. E.g. for the jax solver we will now have:

take_step = get_exponential_take_step(magnus_order, expm_func=jexpm)

This will also be useful if the suggestion in #83 is implemented (so that a diagonalization-based method can optionally be used.)

• Some additional minor comments on the code itself:
  • Numerical constants used in take_step should be defined outside of the function itself, so that they aren't computed every time the function is called. E.g. c1=1/2+np.sqrt(3)/6 should be set before defining take_step.
  • I think term1=(1/2)*h*(g1+g2) should either be 0.5*h*(g1 + g2) or h * (g1 + g2) / 2. I don't know if there is any preference of one over the other, but in either case one operation is eliminated (i.e. (1/2)*... is a division and a multiplication, v.s. only one of the two).

[DanPuzzuoli]

New commit is has nice improvements. Some comments:

• Move the get_exponential_take_step method to later in the file, to right before fixed_step_solver_template.
  • Also, add some documentation for this function.
• Add the magnus_order argument to the documentation for scipy_expm_solver and jax_expm_solver. State that only order 1 and 2 are supported. Note that this file is totally "internal" (i.e. users technically shouldn't see it), so the documentation here is for developers.
• Add the magnus_order argument to the documentation for the exponential solvers in the solve_lmde function in qiskit_dynamics/solvers/solver_functions.py. You can see how max_dt is described there and add a similar sentence. You should also add a reference here to the geometric methods book. You can do this by adding an entry to the bibtex file docs/refs.bib, then cite it in the documentation using the footcite command. (See qiskit_dynamics/perturbation/__init__.py for an example usage of footcite.)
• Clean up the if/else for get_exponential_take_step, and raise an error if magnus_order > 2. E.g. do:

if magnus_order == 1:
    ...
elif magnus_order == 2:
    ...
else:
    raise QiskitError("Only magnus_order 1 or 2 are supported.")

For this last point you'll need to add from qiskit import QiskitError to the imports at the top of the file.

[DanPuzzuoli]

Added some standalone comments, including direct code suggestions. Don't forget to add the magnus_order kwarg to jax_expm_parallel_solver as well.

For the solve_lmde doc string, the description of the magnus_order argument should be moved into the first bullet point describing the 'scipy_expm' function. Currently it says:

    - ``'scipy_expm'``: A fixed-step matrix-exponential solver using ``scipy.linalg.expm``.
      Requires additional kwarg ``max_dt`` indicating the maximum step
      size to take. This solver will break integration periods into even
      sub-intervals no larger than ``max_dt``, and solve over each sub-interval via
      matrix exponentiation of the generator sampled at the midpoint.

and it should now say something like (please review/edit this):

    - ``'scipy_expm'``: A fixed-step matrix-exponential solver using ``scipy.linalg.expm``.
      Requires additional kwarg ``max_dt`` indicating the maximum step
      size to take. This solver will break integration periods into even
      sub-intervals no larger than ``max_dt`` and solve over each sub-interval.
      The optional kwarg ``magnus_order`` controls the integration rule: if 
      ``magnus_order==1``, the generator is sampled at the interval midpoint
      and exponentiated, and if ``magnus_order==2``, a second-order exponentiation
      rule is adopted from :footcite:`hairer_geometric_2006`.

Also it seems like something strange has happened with your refs.bib file - it says you're adding the whole thing, which may be because your branch was out of date. We can resolve this later.

After these changes, the next thing is to start adding tests. One thing that will make this easier is that you're adding functionality to already-existing functions, so you'll just need to add some additional tests to handle the case magnus_order=2. There are a few places where tests need to be added. Each file tests different aspects/levels of the interfaces in Dynamics, and each one is mostly set up to make it easy to subject new solver methods to a suite of tests. All of the relevant testing files are in the test/dynamics/solvers folder.

• The file test_fixed_step_solvers.py contains tests for directly testing the low level solver functions like scipy_expm_solver. To add a test for the scipy_expm_solver function with magnus_order==2, you can copy the class TestScipyExpmSolver (giving it a new name, like TestScipyExpmSolver_Magnus2), and update the take_step method so that it implements the second order rule, and the solve method so that it calls scipy_expm_solver with magnus_order=2. This same process can be repeated in this file for the tests for jax_expm_solver and jax_expm_parallel_solver.
• A similar process needs to be done in the test_solver_functions.py file, which tests things at the level of the solve_ode and solve_lmde functions, which users will actually interact with.
• Lastly, you'll need to add a test in test_solver_functions_interface.py, to validate that an error is raised if a user attempts to pass something other than the integers 1 or 2 to the magnus_order argument.

The last thing I want to ask at this stage, is if there are orders higher than 2 that we can support. Are there higher order rules described in the geometric integration methods book? If we could also add order 3 and/or 4 that'd be pretty cool, and now would be the right time to do it.

[DanPuzzuoli]

This may seem pedantic, but for comments, try to follow the general style in the repo. E.g. Capitalize/punctuate. Indent new lines.

Suggested change

	magnus_order: the expansion order in the Magnus method. Only order 1 and 2
	are supported
	magnus_order: The expansion order in the Magnus method. Only order 1 and 2
	are supported.

[DanPuzzuoli]

unnecessary comment

[DanPuzzuoli]

no need for this comment

[DanPuzzuoli]

Suggested change

	magnus_order: the expansion order in the Magnus method
	expm_func: method of matrix exponentian; can e.g. be expm, jexp.
	magnus_order: The expansion order in the Magnus method.
	expm_func: Method of matrix exponentian.

[DanPuzzuoli]

Undo these changes - see main comment.

[DanPuzzuoli]

Suggested change

	magnus_order: the expansion order in the Magnus method. Only order 1 and 2
	are supported
	magnus_order: The expansion order in the Magnus method. Only order 1 and 2
	are supported.

[DanPuzzuoli]

I've double checked and verified that I believe the constants/integration rules are implemented correctly. Also had some other very minor comments. There are a couple of other stylistic things we'll need to address to make the automated style-checker pass before merging, but we can do that at the end after the tests are added.

[DanPuzzuoli]

This is obviously very minor but change all 1/2 to 0.5. Obviously a trivial operation but it still saves the construction of the ints and then the division.

[DanPuzzuoli]

I don't think any clarity is lost by making this a single line:

terms = a1 + a3/12 + matrix_commutator(-20*a1-a3+com1,a2+com2)/240

[DanPuzzuoli]

Suggested change

	raise QiskitError("Only magnus_order 1, 2 and 3 are supported.")
	raise QiskitError("Only magnus_order 1, 2, and 3 are supported.")

[ihincks]

Looks like a great addition! Apologies for all the trivial comments, but it's the best way to maintain coding standards.

Note that I did not look at blanes_magnus_2009 and so cannot vouch for correctness.

[DanPuzzuoli]

Looks like a great addition! Apologies for all the trivial comments, but it's the best way to maintain coding standards.

Note that I did not look at blanes_magnus_2009 and so cannot vouch for correctness.

I think I've changed everything you said. No need to apologize, I appreciate the thoroughness and commitment to consistent formatting.