Faster Layout and MDC Store Creation

[coreyostrove]

Building off of the changes in PR #445 as part of an effort at reducing the overhead associated with running GST analysis, this PR further speeds up the construction of COPA layouts and MDC stores/objective functions. Both in general, but even more so in the context of iterative GST estimation.

Note: This branch is forked off of 'feature-faster-circuit-primitives', and so some of these commits are from that branch and will get cleaned up whenever #445 gets merged into develop (in case you notice duplicates).

Here is a summary of the changes:

1. Caching for layout creation: One of the main bottlenecks in layout creation is performing a variety of (relatively expensive) preprocessing steps on circuits. In the context of iterative GST we almost always use a nested circuit structure wherein for each new iteration we keep all of the previous circuits and add a few new ones. New layouts are constructed at each iteration, and previously we started from scratch each time. Now prior to starting the iterative estimation we precompute and construct a cache for the more expensive circuit structures used in layout generation, significantly reducing the amount of repeated calculations.
2. Caching for MDC store construction: Similar story to layout creation. Previously one of the bottlenecks in MDC store construction was computing the number of possible outcomes for a circuit (including POVMs and potentially instruments). This resulted in repeated calls to a method called expand_instruments_and_separate_povm which is quite expensive. Like layouts we construct a fresh MDC store at each iteration, so by caching this information we significantly reduce the amount of repeated calculations.
3. The expand_instruments_and_separate_povm method has been refactored, moving it from being a method of Circuits which takes as input a model into being a method of OpModels which takes as input a circuit. This method was a strange fit for the Circuit class, as it relies pretty heavily on knowledge of the internal details of a model, and from an abstraction hierarchy level this was also a out of character thing for a Circuit to be asked to do. It makes much more sense for an OpModel (the implementation effectively assumed you'd only be calling it on an OpModel given the attributes and methods it assumed were implemented, so now that is enforced), and OpModel is also where the conceptually similar complete_circuit(s) and split_circuit(s) methods live. In addition to the refactor, I have also added a new bulk_expand_instruments_and_separate_povm which more efficiently performs this operation on a list of circuits.
4. Miscellaneous performance updates in service of speeding up layout and MDC store construction.

I look forward to any and all feedback.

P.S. I didn't mention how these changes affect the bottom line on runtime yet. For single-qubit GST with the XYI gate set, the full TP parameterization, and a maximum depth of 128 (as before, I haven't yet tested this on the 2Q case, or for other experiment designs, so ymmv) overall we have:

• MatrixCOPALayout creation is ~4X faster
• MapCOPALayout creation is ~2X faster
• ModelDataSetCircuitStore/TimeIndependentMDCObjectiveFunction creation is ~5X faster.

Those changes are measured relative to the runtimes following the changes made on #445. Combined with #445 this gives an overall reduction in end-to-end runtime (for the 1Q test case described above) of ~50% when using the MatrixForwardSimulator, and ~15% when using the MapForwardSimulator (proportionally less of the overall runtime is spent performing these subroutines with map, so this difference makes sense).

[rileyjmurray]

It looks like this branch came off of the feature-faster-circuit-primitives branch. That's making Git think this PR contains all the material from PR #445. I'll hold off on a full review of this PR until (1) #445 is merged and (2) this PR is updated, as needed, in order for Git to see the minimal set of changes.

That said, I do have three actionable comments. Two comments are given in-line. Here's a broader comment.

Right now the changes introduce free-functions like create_matrix_copa_layout_circuit_cache and create_map_copa_layout_circuit_cache. I'd prefer that these were static class functions, and that the name class signifier (_matrix_, _map__, etc..) be dropped from the function names. From there, you could invoke the function just with mld.sim.create_copa_layout_circuit_cache(...).

If you'd like you could make create_copa_layout_circuit_cache a method in the ForwardSimulator base class. If that method were to raise a NotImplementedError by default then you could replace some branching if-statements with

try:
    precomp_layour_circuit_cache = mdl.sim.create_copa_layout_circuit_cache(...)
except NotImplementedError:
    precomp_layour_circuit_cache = None

[coreyostrove]

Thanks for the feedback and suggestions, @rileyjmurray. I have incorporated your feedback and refactored the cache creation methods into static methods of the corresponding forward simulator classes. I also addressed the two inline comments, so I have marked those as resolved.

[coreyostrove]

Question for folks, but particularly aimed at @enielse. A comparable line to this one for computing the quantity named observed_outcomes appears here, but in the map layout code this queries the outcomes property of DataSetRow instead of the unique_outcomes property as is done here. I think this should be unique_outcomes in both cases (I don't see a reason why it wouldn't also be for map), but I wanted to run this by you before I made that change.

[sserita]

I don't see any reason why map cannot also use unique_outcomes instead of outcomes, but am basing that just on looking through this code while reviewing. Expert @enielse opinion still welcome.

[rileyjmurray]

This looks good overall, but the current changes have broken unit tests for some forward simulators. I don't know why all of the failures are happening. One of the failures is because TorchForwardSimulator uses Circuit.expand_instruments_and_separate_povm.

[coreyostrove]

This looks good overall, but the current changes have broken unit tests for some forward simulators. I don't know why all of the failures are happening. One of the failures is because TorchForwardSimulator uses Circuit.expand_instruments_and_separate_povm.

TIL that you can copy a link to a specific line of the github actions log, neat.

Yup, this makes sense. One of the changes made on this branch was moving that method from the Circuit class to the OpModel class. This branch was forked off of a version of develop before the torch forward simulator code was merged in, so I'll need to make some minor updates.

[sserita]

Excited to get this in and get those speedups in layout creation! First, I have a few comments: one which I think will error in certain edge cases, and others that are optional but may cut down on branching to increase maintainability.

[sserita]

Are we checking this because only matrix and map forward sims have this defined? If so, thoughts on a try/catch on NotImplementedError instead so that if create_[...]_cache is defined on another forward sim, we don't have to remember and come back to change this check?

[coreyostrove]

This staticmethod is implemented in the ForwardSimulator base class, but raises a NotImplementedError. Though that might eventually change (it would be possible in principle to implement caching for the base COPA layout class' cache creation, but I haven't done so yet).

I've been spending too much time the past few months on SWE youtube, and one of the SWE design philosophies du jour that I've been enamored with as of late is 'Locality of Behavior' (LoB, as described here). While I'm not a purist by any means, in this instance I think it is easier (for someone reading this for the first time, or the first time in a while) to understand this portion of the code by explicitly signaling for which forward simulators the caching will happen, rather than the try/except paradigm which would require looking into the source code for the various forward simulator classes to identify when it won't happen.

That said, you're right there is a trade-off in doing so, which is needing to remember to come back to this and explicitly enable caching in protocols which support it. (Realistically it'd most likely be one of either me, you or Riley messing around with pyGSTi internals at this level, so hopefully we'd remember).

[rileyjmurray]

I hadn't heard of LoB before, but I like it.

[sserita]

What does this mean? In the <10% of cases that you expect to run into, will the cache/resulting layout still be created correctly?

[coreyostrove]

Good question, it took me a minute to remember what I was thinking about here. (It wasn't loaded into cache, ba dum dum tsss...)

The thing that happens in the init, that isn't happening here, are the following calls:

unique_circuits, to_unique = self._compute_unique_circuits(circuits)
aliases = circuits.op_label_aliases if isinstance(circuits, _CircuitList) else None
ds_circuits = _lt.apply_aliases_to_circuits(unique_circuits, aliases)

This allows for the possibility of aliasing in the keys of the DataSet. This aliasing is the part that I didn't implement in the code shown below. The thing that would happen in the <10% of cases I was alluding to, i.e. the cases where someone is using aliasing (as far as I am aware this is not a widely used feature) is that ds_row would be None some amount (maybe all) of the time, and we'd pass a list with None entries into bulk_expand_instruments_and_separate_povm. This is allowed (and is in fact the default), but it would result in expanded circuits being created for every POVM and instrument effect, instead of just those for which outcomes were observed. The resulting layout wouldn't be wrong, but would be less memory efficient than it could otherwise be, as it would have space allocated for probabilities we won't be calculating. It would also make the bulk_expand_instruments_and_separate_povm less efficient than it otherwise could be.

All of that said, I am not entirely sure why I didn't implement the aliasing for the DataSet keys at the time. The answer is probably something along the lines of: "I wasn't sure how the aliasing code worked and couldn't be arsed to figure it out at the time." Anyhow, I see no reason not to add this back in (I'm pretty sure I can just essentially copy and paste the lines above into the cache creation routine), so I'll give it a crack.

[sserita]

I don't see any reason why map cannot also use unique_outcomes instead of outcomes, but am basing that just on looking through this code while reviewing. Expert @enielse opinion still welcome.

[coreyostrove]

Thanks for the careful feedback, @sserita!

I have just pushed some changes addressing this feedback, to summarize the main changes:

1. I've added back in alias support during cache creation to better match the original implementations in the various COPA layout inits. (Inspired by one of your questions about a comment I had left).
2. I have replace some of the branching statements for cache existence with an empty dictionary + get based version where appropriate. There are a few places where the behavior in the case where the cache isn't present is meaningfully different than when invoking the fallback behavior for missing values (because they are using bulk methods), and in those instances I haven't made this change.

There is also a miscellaneous typo fix in the RB unit tests that I accidentally introduced with the IRB branch that I'm throwing in for good measure, in case you notice the bizarre inclusion.

[sserita]

The aliasing looks good. I just have one question left on the RB test change.

[sserita]

I'm confused about why this change is needed. Aren't we in test/unit/protocols, so ../../test_packages is the right way to get to those files?

[coreyostrove]

I thought so originally myself, but I noticed when running unit tests locally last night that a new test_packages directory was being created at the base level (i.e. at the same level as test and pygsti) and populated with these files. In retrospect I suppose I could have equivalently removed one of the sets of '..'. Not sure why this is the case though, tbh, maybe something to do with how pytest sets the cwd?

[coreyostrove]

It is also possible this was related to which directory I was launching pytest from, so if you have a venv handy to run these tests on and have a different experience that'd be good to know.