sandialabs / pyGSTi

A python implementation of Gate Set Tomography
http://www.pygsti.info
Apache License 2.0
137 stars 55 forks source link

Slimmed-down Levenberg-Marquardt for nonlinear least squares #500

Open rileyjmurray opened 1 month ago

rileyjmurray commented 1 month ago

pyGSTi's nonlinear least squares solver is a custom implementation with Levenberg-Marquardt with a dizzying array of options. The core algorithm is over 800 lines long and has 10 levels of indentation at its deepest point. The complexity of this existing implementation makes it basically impossible to extend.

This PR introduces slimmed-down infrastructure for nonlinear least squares. I started by copying customlm.py into a new file called simplerlm.py, then I removed one feature at a time in a series of commits. "Removal" of a feature meant making the default behavior for a given option the only behavior.

Notes

I tried to use ChatGPT to help split the monolithic custom_leastsq function in customlm.py into simpler functions. This turned out to not work for reasons that I'll discuss off-GitHub if people are interested.

For the curious, I'm doing this because I want to extend our nonlinear least squares solver in two ways. First, I want it to support pytorch Tensors; right now we only support numpy arrays and whatever home-cooked thing we have for distributed-memory computations. Second, I want to experiment with optimization algorithms that rarely (if ever!) require evaluating the full Jacobian of circuit outcome probabilities.

I also edited a scaling test for MPI because it raised errors when I ran pytest.

Finally, I made slight extensions to the ArraysInterface classes. The new functions are there so that the simpler LM implementation never needs to directly index into JTJ.

rileyjmurray commented 1 month ago

Ping @enielse for review, when you get a chance. Now is a good time for us to document why options are worth keeping in the SimplerLMOptimizer class and the simplish_leastsq function. I'd be all for removing even more options if you think some are unnecessary.