PyTorch-backed forward simulation

rileyjmurray commented 6 months ago

This PR introduces TorchForwardSimulator, a forward simulator (for computing circuit outcome probabilities) based on PyTorch. It uses automatic differentiation to compute the Jacobian of the map from model parameters to circuit outcome probabilities. In the future we could extend it to do computations on a system's GPU, or to use PyTorch-based optimization algorithms instead of pyGSTi's custom algorithms for MLE.

Approach

My approach required creating a new ModelMember subclass called Torchable. This subclass adds two required functions, called stateless_data and torch_base. Their meanings are given below: https://github.com/sandialabs/pyGSTi/blob/1ec6909a7f9149c4f6758f9cd46256dbbff984c1/pygsti/modelmembers/torchable.py#L18-L43 In principle, TorchForwardSimulator can handle all models for which constituent parameterized ModelMembers are Torchable. So far I've only extended TPState, FullTPOp, and TPPOVM to be Torchable; these are the classes used in "full TP" GST.

The Python file that contains TorchForwardSimulator also defines two helper classes: StatelessCircuit and StatelessModel. I think it's fine to keep these classes as purely internal implementation-specific constructs for now. Depending on future performance optimizations of TorchForwardSimulator we might want to put them elsewhere in pyGSTi.

What should come after this PR

We should compare performance of TorchForwardSimulator to MapForwardSimulator on problems of interest. There's a chance that the former isn't faster than the latter with the current implementation. If that's the case then I should look at possible performance optimizations specifically inside TorchForwardSimulator.

We should add implementations of stateless_data and torch_base to GST models beyond "Full TP" (in particular I'd like to try CPTP).

Incidental changes

My implementation originally interacted with the following evotype classes

    <class 'pygsti.evotypes.densitymx[_slow].statereps.StateRepDense'>
    <class 'pygsti.evotypes.densitymx[_slow].opreps.OpRepDenseSuperop'>
    <class 'pygsti.evotypes.densitymx[_slow].effectreps.EffectRepConjugatedState'>

When I write [_slow] in the class names above you can put the empty string or just _slow, depending on the default evotype specified in evotypes.py.

To my surprise, I found that interacting with evotypes was neither necessary nor sufficient for what I wanted to accomplish. So while I did make changes in evotypes/densitymx_slow/ to remove unnecessary class inheritances and to add documentation, those changes were only to make life a little easier for future pyGSTi contributors.

sserita commented 5 months ago

I'm giving a few quick comments here because I think this PR will actually take me some time to get through.

The main thing I am concerned and thinking about is the choice to extend the ModelMember class as opposed to adjusting the evotypes. The general purpose of the split between modelmembers and evotypes is so that we don't have to go through and implement these abstract methods in all modelmembers - we can make the change in the evotype and then any modelmember works. It is totally possible that extending the API is the best way to do this. The TermSimulator follows a similar pattern where the API extensions are the cleanest way to do it. But I'll probably spend some time thinking about whether this is way we want to implement this. Probably a point of discussion for us in our dev meetings in the upcoming weeks.

rileyjmurray commented 5 months ago

@sserita, regarding

The main thing I am concerned and thinking about is the choice to extend the ModelMember class as opposed to adjusting the evotypes. The general purpose of the split between modelmembers and evotypes is so that we don't have to go through and implement these abstract methods in all modelmembers - we can make the change in the evotype and then any modelmember works.

Unfortunately there's no way to do this just through evotypes. Using pytorch's AD capabilities requires knowing the free parameters in a modelmember and how those free parameters map to the common parameterization-agnostic representation (i.e., representations in evotypes).

rileyjmurray commented 5 months ago

Notes from today's meeting:

I removed empty base classes for OpRep, EffectRep, StateRep. I should make sure we aren't doing type checking elsewhere in pyGSTi that would be broken by this. (There's a chance we are and it simply isn't showing up in current tests.)

rileyjmurray commented 2 months ago

@coreyostrove, @sserita, @enielse: this is ready for review.

rileyjmurray commented 1 month ago

@coreyostrove and @sserita this is ready for review.

rileyjmurray commented 1 month ago

@coreyostrove I still plan on adding proper unit tests (in contrast to the existing tests, which are really integration tests).

coreyostrove commented 1 month ago

I'm withholding explicit approval until 0.9.13 goes out just as a reminder that we don't want this to merge in before then, and it sounds like you are still working on getting in a few tests as well.

Minor clarification, did you mean 0.9.12.3?

sserita commented 1 month ago

Yes I absolutely meant 0.9.12.3. Speaking of which, Tim just approved that so it will be going live as soon as all the tests pass and I get the release drafted up.

So just let me know when the tests are good on this branch and I'll give my approval also.

rileyjmurray commented 1 month ago

@sserita, @coreyostrove all tests pass, including new tests that I wrote in test_forwardsim.py. Merge when ready!

sandialabs / pyGSTi