pyGSTi 0.9.12.1__cp310-cp310-win_amd64.whl → 0.9.13__cp310-cp310-win_amd64.whl

pygsti/forwardsims/torchfwdsim.py

@@ -0,0 +1,265 @@
+"""
+Defines a ForwardSimulator class called "TorchForwardSimulator" that can leverage the automatic
+differentation features of PyTorch.
+
+This file also defines two helper classes: StatelessCircuit and StatelessModel.
+
+See also: pyGSTi/modelmembers/torchable.py.
+"""
+#***************************************************************************************************
+# Copyright 2024, National Technology & Engineering Solutions of Sandia, LLC (NTESS).
+# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
+# in this software.
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License.  You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
+#***************************************************************************************************
+
+
+from __future__ import annotations
+from typing import Tuple, Optional, Dict, TYPE_CHECKING
+if TYPE_CHECKING:
+    from pygsti.baseobjs.label import Label
+    from pygsti.models.explicitmodel import ExplicitOpModel
+    from pygsti.circuits.circuit import SeparatePOVMCircuit
+    from pygsti.layouts.copalayout import CircuitOutcomeProbabilityArrayLayout
+    import torch
+
+import warnings as warnings
+from pygsti.modelmembers.torchable import Torchable
+from pygsti.forwardsims.forwardsim import ForwardSimulator
+
+try:
+    import torch
+    TORCH_ENABLED = True
+except ImportError:
+    TORCH_ENABLED = False
+    pass
+
+
+class StatelessCircuit:
+    """
+    Helper data structure for specifying a quantum circuit (consisting of prep,
+    applying a sequence of gates, and applying a POVM to the output of the last gate).
+    """
+
+    def __init__(self, spc: SeparatePOVMCircuit):
+        self.prep_label = spc.circuit_without_povm[0]
+        self.op_labels  = spc.circuit_without_povm[1:]
+        self.povm_label = spc.povm_label
+        self.outcome_probs_dim = len(spc.effect_labels)
+        # ^ This definition of outcome_probs_dim will need to be changed if/when
+        #   we extend any Instrument class to be Torchable.
+        return
+
+
+class StatelessModel:
+    """
+    A container for the information in an ExplicitOpModel that's "stateless" in the sense of
+    object-oriented programming:
+    
+        * A list of StatelessCircuits
+        * Metadata for parameterized ModelMembers
+
+    StatelessModels have instance functions to facilitate computation of (differentable!)
+    circuit outcome probabilities.
+
+    Design notes
+    ------------
+    Much of this functionality could be packed into the TorchForwardSimulator class.
+    Keeping it separate from TorchForwardSimulator helps clarify that it uses none of
+    the sophiciated machinery in TorchForwardSimulator's base class.
+    """
+
+    def __init__(self, model: ExplicitOpModel, layout: CircuitOutcomeProbabilityArrayLayout):
+        circuits = []
+        self.outcome_probs_dim = 0
+        #TODO: Refactor this to use the bulk_expand_instruments_and_separate_povm codepath
+        for _, circuit, outcomes in layout.iter_unique_circuits():
+            expanded_circuits = model.expand_instruments_and_separate_povm(circuit, outcomes)
+            if len(expanded_circuits) > 1:
+                raise NotImplementedError("I don't know what to do with this.")
+            spc = next(iter(expanded_circuits))
+            c = StatelessCircuit(spc)
+            circuits.append(c)
+            self.outcome_probs_dim += c.outcome_probs_dim
+        self.circuits = circuits
+
+        # We need to verify assumptions on what layout.iter_unique_circuits() returns.
+        # Looking at the implementation of that function, the assumptions can be
+        # framed in terms of the "layout._element_indicies" dict.
+        eind = layout._element_indices
+        assert isinstance(eind, dict)
+        items = iter(eind.items())
+        k_prev, v_prev = next(items)
+        assert k_prev == 0
+        assert v_prev.start == 0
+        for k, v in items:
+            assert k == k_prev + 1
+            assert v.start == v_prev.stop
+            k_prev = k
+            v_prev = v
+        assert self.outcome_probs_dim == v_prev.stop
+
+        self.param_metadata = []
+        for lbl, obj in model._iter_parameterized_objs():
+            assert isinstance(obj, Torchable), f"{type(obj)} does not subclass {Torchable}."
+            param_type = type(obj)
+            param_data = (lbl, param_type) + (obj.stateless_data(),)
+            self.param_metadata.append(param_data)
+        self.params_dim = None 
+        # ^ That's set in get_free_params.
+
+        self.default_to_reverse_ad = None
+        # ^ That'll be set to a boolean the next time that get_free_params is called.
+        return
+    
+    def get_free_params(self, model: ExplicitOpModel) -> Tuple[torch.Tensor]:
+        """
+        Return a tuple of Tensors that encode the states of the provided model's ModelMembers
+        (where "state" in meant the sense of  object-oriented programming).
+        
+        We compare the labels of the input model's ModelMembers to those of the model provided
+        to StatelessModel.__init__(...). We raise an error if an inconsistency is detected.
+        """
+        free_params = []
+        prev_idx = 0
+        for i, (lbl, obj) in enumerate(model._iter_parameterized_objs()):
+            gpind = obj.gpindices_as_array()
+            vec = obj.to_vector()
+            vec_size = vec.size
+            vec = torch.from_numpy(vec)
+            assert gpind[0] == prev_idx and gpind[-1] == prev_idx + vec_size - 1
+            # ^ We should have gpind = (prev_idx, prev_idx + 1, ..., prev_idx + vec.size - 1).
+            #   That assert checks a cheap necessary condition that this holds.
+            prev_idx += vec_size
+            if self.param_metadata[i][0] != lbl:
+                message = """
+                The model passed to get_free_params has a qualitatively different structure from
+                the model used to construct this StatelessModel. Specifically, the two models have
+                qualitative differences in the output of "model._iter_parameterized_objs()".
+
+                The presence of this structral difference essentially gaurantees that a subsequent
+                call to get_torch_bases would silently fail, so we're forced to raise an error here.
+                """
+                raise ValueError(message)
+            free_params.append(vec)
+        self.params_dim = prev_idx
+        self.default_to_reverse_ad = self.outcome_probs_dim < self.params_dim
+        return tuple(free_params)
+
+    def get_torch_bases(self, free_params: Tuple[torch.Tensor]) -> Dict[Label, torch.Tensor]:
+        """
+        Take data of the kind produced by get_free_params and format it in the way required by
+        circuit_probs_from_torch_bases.
+
+        Note
+        ----
+        If you want to use the returned dict to build a PyTorch Tensor that supports the 
+        .backward() method, then you need to make sure that fp.requires_grad is True for all
+        fp in free_params. This can be done by calling fp._requires_grad(True) before calling
+        this function.
+        """
+        assert len(free_params) == len(self.param_metadata)
+         # ^ A sanity check that we're being called with the correct number of arguments.
+        torch_bases = dict()
+        for i, val in enumerate(free_params):
+
+            label, type_handle, stateless_data = self.param_metadata[i]
+            param_t = type_handle.torch_base(stateless_data, val)
+            torch_bases[label] = param_t
+        
+        return torch_bases
+
+    def circuit_probs_from_torch_bases(self, torch_bases: Dict[Label, torch.Tensor]) -> torch.Tensor:
+        """
+        Compute the circuit outcome probabilities that result when all of this StatelessModel's
+        StatelessCircuits are run with data in torch_bases.
+    
+        Return the results as a single (vectorized) torch Tensor.
+        """
+        probs = []
+        for c in self.circuits:
+            superket = torch_bases[c.prep_label]
+            superops = [torch_bases[ol] for ol in c.op_labels]
+            povm_mat = torch_bases[c.povm_label]
+            for superop in superops:
+                superket = superop @ superket
+            circuit_probs = povm_mat @ superket
+            probs.append(circuit_probs)
+        probs = torch.concat(probs)
+        return probs
+    
+    def circuit_probs_from_free_params(self, *free_params: Tuple[torch.Tensor], enable_backward=False) -> torch.Tensor:
+        """
+        This is the basic function we expose to pytorch for automatic differentiation. It returns the circuit
+        outcome probabilities resulting when the states of ModelMembers associated with this StatelessModel
+        are set based on free_params. 
+
+        If you want to call PyTorch's .backward() on the returned Tensor (or a function of that Tensor), then
+        you should set enable_backward=True. Keep the default value of enable_backward=False in all other
+        situations, including when using PyTorch's jacrev function.
+        """
+        if enable_backward:
+            for fp in free_params:
+                fp._requires_grad(True)
+        torch_bases = self.get_torch_bases(free_params)
+        probs = self.circuit_probs_from_torch_bases(torch_bases)
+        return probs
+
+
+class TorchForwardSimulator(ForwardSimulator):
+    """
+    A forward simulator that leverages automatic differentiation in PyTorch.
+    """
+
+    ENABLED = TORCH_ENABLED
+
+    def __init__(self, model : Optional[ExplicitOpModel] = None):
+        if not self.ENABLED:
+            raise RuntimeError('PyTorch could not be imported.')
+        self.model = model
+        super(ForwardSimulator, self).__init__(model)
+
+    def _bulk_fill_probs(self, array_to_fill, layout, split_model = None) -> None:
+        if split_model is None:
+            slm = StatelessModel(self.model, layout)
+            free_params = slm.get_free_params(self.model)
+            torch_bases = slm.get_torch_bases(free_params)
+        else:
+            slm, torch_bases = split_model
+
+        probs = slm.circuit_probs_from_torch_bases(torch_bases)
+        array_to_fill[:slm.outcome_probs_dim] = probs.cpu().detach().numpy().ravel()
+        return
+
+    def _bulk_fill_dprobs(self, array_to_fill, layout, pr_array_to_fill) -> None:
+        slm = StatelessModel(self.model, layout)
+        # ^ TODO: figure out how to safely recycle StatelessModel objects from one
+        #   call to another. The current implementation is wasteful if we need to 
+        #   compute many jacobians without structural changes to layout or self.model.
+        free_params = slm.get_free_params(self.model)
+    
+        if pr_array_to_fill is not None:
+            torch_bases = slm.get_torch_bases(free_params)
+            splitm = (slm, torch_bases)
+            self._bulk_fill_probs(pr_array_to_fill, layout, splitm)
+
+        argnums = tuple(range(len(slm.param_metadata)))
+        if slm.default_to_reverse_ad:
+            # Then slm.circuit_probs_from_free_params will automatically construct the
+            # torch_base dict to support reverse-mode AD.
+            J_func = torch.func.jacrev(slm.circuit_probs_from_free_params, argnums=argnums)
+        else:
+            # Then slm.circuit_probs_from_free_params will automatically skip the extra
+            # steps needed for torch_base to support reverse-mode AD.
+            J_func = torch.func.jacfwd(slm.circuit_probs_from_free_params, argnums=argnums)
+        # ^ Note that this _bulk_fill_dprobs function doesn't accept parameters that
+        #   could be used to override the default behavior of the StatelessModel. If we
+        #   have a need to override the default in the future then we'd need to override
+        #   the ForwardSimulator function(s) that call self._bulk_fill_dprobs(...).
+
+        J_val = J_func(*free_params)
+        J_val = torch.column_stack(J_val)
+        array_to_fill[:] = J_val.cpu().detach().numpy()
+        return

pygsti/forwardsims/weakforwardsim.py

@@ -55,7 +55,7 @@ class WeakForwardSimulator(_ForwardSimulator):
         circuit : Circuit
             A tuple-like object of *simplified* gates (e.g. may include
             instrument elements like 'Imyinst_0') generated by
-            Circuit.expand_instruments_and_separate_povm()
+            OpModel.expand_instruments_and_separate_povm()
 
         resource_alloc: ResourceAlloc
             Currently not used
@@ -77,7 +77,7 @@ class WeakForwardSimulator(_ForwardSimulator):
         assert(resource_alloc is None), "WeakForwardSimulator cannot use a resource_alloc for one shot."
 
         #prep_label, op_labels, povm_label = self.model.split_circuit(spc_circuit)
-        spc_dict = circuit.expand_instruments_and_separate_povm(self.model,
+        spc_dict = self.model.expand_instruments_and_separate_povm(circuit,
                                                                 observed_outcomes=None)  # FUTURE: observed outcomes?
         assert(len(spc_dict) == 1), "Circuits with instruments are not supported by weak forward simulator (yet)"
         spc = next(iter(spc_dict.keys()))  # first & only SeparatePOVMCircuit

pygsti/io/__init__.py

@@ -12,8 +12,7 @@ pyGSTi Input/Output Python Package
 
 # Import the most important/useful routines of each module into
 # the package namespace
-#from .legacyio import enable_no_cython_unpickling
-#from .legacyio import enable_old_object_unpickling  # , disable_old_object_unpickling
+
 from .readers import *
 from .metadir import *
 from .stdinput import *

pygsti/io/mongodb.py

@@ -158,7 +158,6 @@ def read_auxtree_from_mongodb_doc(mongodb, doc, auxfile_types_member='auxfile_ty
 
 
 def _load_auxdoc_member(mongodb, member_name, typ, metadata, quick_load):
-    from pymongo import ASCENDING, DESCENDING
     subtypes = typ.split(':')
     cur_typ = subtypes[0]
     next_typ = ':'.join(subtypes[1:])
@@ -809,7 +808,6 @@ def remove_auxtree_from_mongodb(mongodb, collection_name, doc_id, auxfile_types_
 
 
 def _remove_auxdoc_member(mongodb, member_name, typ, metadata, session, recursive):
-    from pymongo import ASCENDING, DESCENDING
     subtypes = typ.split(':')
     cur_typ = subtypes[0]
     next_typ = ':'.join(subtypes[1:])

pygsti/io/stdinput.py

@@ -128,7 +128,7 @@ class StdInputParser(object):
             else:
                 circuit = _Circuit._fastinit(layer_tuple, line_lbls, editable=False,
                                              name='', stringrep=s, occurrence=occurrence_id,
-                                             compilable_layer_indices=compilable_indices)
+                                             compilable_layer_indices_tup=compilable_indices)
 
             if self.use_global_parse_cache:
                 _global_parse_cache[create_subcircuits][s] = circuit
@@ -164,10 +164,10 @@ class StdInputParser(object):
         occurrence_id: int or None
             The "occurence id" - an integer following a second '@' symbol that identifies a particular
             copy of this circuit.
-        compilable_indices : tuple or None
+        compilable_indices : tuple
             A tuple of layer indices (into `label_tuple`) marking the layers that can be "compiled",
             and are *not* followed by a barrier so they can be compiled with following layers.  This
-            is non-`None` only when there are explicit markers within the circuit string indicating
+            is non-empty only when there are explicit markers within the circuit string indicating
             the presence or absence of barriers.
         """
         if lookup is None:
@@ -175,8 +175,8 @@ class StdInputParser(object):
         self._circuit_parser.lookup = lookup
         circuit_tuple, circuit_labels, occurrence_id, compilable_indices = \
             self._circuit_parser.parse(s, create_subcircuits)
-        # print "DB: result = ",result
-        # print "DB: stack = ",self.exprStack
+        compilable_indices = compilable_indices if compilable_indices is not None else ()
+
         return circuit_tuple, circuit_labels, occurrence_id, compilable_indices
 
     def parse_dataline(self, s, lookup=None, expected_counts=-1, create_subcircuits=True,
@@ -326,14 +326,9 @@ class StdInputParser(object):
                         self.parse_circuit_raw(line, {}, create_subcircuits)
                     if parsed_line_lbls is None:
                         parsed_line_lbls = line_labels  # default to the passed-in argument
-                        #nlines = num_lines
-                    #else: nlines = None  # b/c we've got a valid line_lbls
                     circuit = _Circuit._fastinit(layer_lbls, parsed_line_lbls, editable=False,
                                                  name='', stringrep=line.strip(), occurrence=occurrence_id,
-                                                 compilable_layer_indices=compilable_indices)
-                    #circuit = _Circuit(layer_lbls, stringrep=line.strip(),
-                    #                        line_labels=parsed_line_lbls, num_lines=nlines,
-                    #                        expand_subcircuits=False, check=False, occurrence=occurrence_id)
+                                                 compilable_layer_indices_tup=compilable_indices)
                     ##Note: never expand subcircuits since parse_circuit_raw already does this w/create_subcircuits arg
                 circuit_list.append(circuit)
         return circuit_list
@@ -507,11 +502,6 @@ class StdInputParser(object):
 
         last_circuit = last_commentDict = None
 
-        #REMOVE DEBUG
-        #from mpi4py import MPI
-        #comm = MPI.COMM_WORLD
-        #debug_circuit_elements = 0; debug_test_simple_dict = {}; circuit_bytes = 0; sizeof_bytes = 0
-
         with open(filename, 'r') as inputfile:
             for (iLine, line) in enumerate(inputfile):
                 if iLine % nSkip == 0 or iLine + 1 == nLines: display_progress(iLine + 1, nLines, filename)
@@ -593,11 +583,6 @@ class StdInputParser(object):
                                     warnings.append("Dataline for circuit '%s' has zero counts and will be ignored" % s)
                                 continue  # skip lines in dataset file with zero counts (no experiments done)
                             else:
-                                #if not bBad:
-                                #    s = circuitStr if len(circuitStr) < 40 else circuitStr[0:37] + "..."
-                                #    warnings.append("Dataline for circuit '%s' has zero counts." % s)
-                                # don't make a fuss if we don't ignore the lines (needed for
-                                # fill_in_empty_dataset_with_fake_data).
                                 pass
 
                         #Call this low-level function for performance, so need to construct outcome *index* arrays above
@@ -639,7 +624,6 @@ class StdInputParser(object):
                         else:
                             raise ValueError("Invalid circuit data-line prefix: '%s'" % parts[0])
 
-        #REMOVE print("Rank %d DONE load loop.  circuit bytes = %g" % (comm.rank, circuit_bytes))
         if looking_for in ("circuit_data", "circuit_data_or_line") and current_item:
             #add final circuit info (no blank line at end of file)
             dataset.add_raw_series_data(current_item['circuit'], current_item.get('outcomes', []),

pygsti/layouts/copalayout.py

@@ -190,24 +190,26 @@ class CircuitOutcomeProbabilityArrayLayout(_NicelySerializable):
         if unique_circuits is None and to_unique is None:
             unique_circuits, to_unique = self._compute_unique_circuits(circuits)
         self._unique_circuits = unique_circuits
-        self._unique_circuit_index = _collections.OrderedDict(
-            [(c, i) for i, c in enumerate(self._unique_circuits)])  # original circuits => unique circuit indices
+        self._unique_circuit_index = {c:i for i, c in enumerate(self._unique_circuits)}  # original circuits => unique circuit indices
         self._to_unique = to_unique  # original indices => unique circuit indices
         self._unique_complete_circuits = unique_complete_circuits  # Note: can be None
         self._param_dimensions = param_dimensions
         self._resource_alloc = _ResourceAllocation.cast(resource_alloc)
 
-        max_element_index = max(_it.chain(*[[ei for ei, _ in pairs] for pairs in elindex_outcome_tuples.values()])) \
-            if len(elindex_outcome_tuples) > 0 else -1  # -1 makes _size = 0 below
-        indices = set(i for tuples in elindex_outcome_tuples.values() for i, o in tuples)
+        indices = [i for tuples in elindex_outcome_tuples.values() for i, _ in tuples]
+        max_element_index = max(indices) if len(elindex_outcome_tuples) > 0 else -1  # -1 makes _size = 0 below
+        indices = set(indices)
+        
+        
         self._size = max_element_index + 1
         assert(len(indices) == self._size), \
             "Inconsistency: %d distinct indices but max index + 1 is %d!" % (len(indices), self._size)
 
-        self._outcomes = _collections.OrderedDict()
-        self._element_indices = _collections.OrderedDict()
+        self._outcomes = dict()
+        self._element_indices = dict()
+        sort_idx_func = lambda x: x[0]
         for i_unique, tuples in elindex_outcome_tuples.items():
-            sorted_tuples = sorted(tuples, key=lambda x: x[0])  # sort by element index
+            sorted_tuples = sorted(tuples, key=sort_idx_func)  # sort by element index
             elindices, outcomes = zip(*sorted_tuples)  # sorted by elindex so we make slices whenever possible
             self._outcomes[i_unique] = tuple(outcomes)
             self._element_indices[i_unique] = _slct.list_to_slice(elindices, array_ok=True)
@@ -598,10 +600,6 @@ class CircuitOutcomeProbabilityArrayLayout(_NicelySerializable):
         """
         jtj[:] = _np.dot(j.T, j)
 
-    #Not needed
-    #def allocate_jtj_shared_mem_buf(self):
-    #    return _np.empty((self._param_dimensions[0], self._param_dimensions[0]), 'd'), None
-
     def memory_estimate(self, array_type, dtype='d'):
         """
         Memory required to allocate an array of a given type (in bytes).

pygsti/layouts/distlayout.py

@@ -360,9 +360,6 @@ class DistributableCOPALayout(_CircuitOutcomeProbabilityArrayLayout):
             to_send = 0  # default = contribute nothing to MPI.SUM below
 
             if i in atoms_dict:
-                #print("DB (%d): updating elindex_outcome_tuples w/Atom %d:\n%s"
-                #      % (rank, i, "\n".join(["%d: %s" % (indx, str(tups))
-                #                             for indx, tups in atoms_dict[i].elindex_outcome_tuples.items()])))
                 if start is None:
                     start = stop = offset
                 assert(stop == offset)  # This should be checked by _assert_sequential(myAtomIndices) above
@@ -810,43 +807,6 @@ class DistributableCOPALayout(_CircuitOutcomeProbabilityArrayLayout):
         super().__init__(local_circuits, local_unique_circuits, local_to_unique, local_elindex_outcome_tuples,
                          local_unique_complete_circuits, param_dimensions, resource_alloc)
 
-        #DEBUG LAYOUT PRINTING
-        #def cnt_str(cnt):
-        #    if cnt is None: return "%4s" % '-'
-        #    return "%4d" % cnt
-        #def slc_str(slc):
-        #    if slc is None: return "%14s" % '--'
-        #    return "%3d->%3d (%3d)" % (slc.start, slc.stop, slc.stop - slc.start) \
-        #        if isinstance(slc, slice) else "%14s" % str(slc)
-        #shm = bool(resource_alloc.host_comm is not None)  # shared mem?
-        #if rank == 0:
-        #    print("%11s %-14s %-14s %-14s   %-14s %-4s %-14s %-4s %-14s %-4s" % (
-        #        '#', 'g-elements', 'g-params', 'g-param2s',
-        #        'h-elements','tot', 'h-params','tot', 'h-params2','tot'),
-        #          flush=True)
-        #resource_alloc.comm.barrier()
-        #for r in range(resource_alloc.comm.size):
-        #    if r == rank:
-        #        my_desc = ("%3d (%2d.%2d)" % (rank, resource_alloc.host_index, resource_alloc.host_comm.rank)) \
-        #                  if shm else ("%11d" % rank)
-        #        print(my_desc, slc_str(self.global_element_slice), slc_str(self.global_param_slice),
-        #              slc_str(self.global_param2_slice), ' ',
-        #              slc_str(self.host_element_slice), cnt_str(self.host_num_elements),
-        #              slc_str(self.host_param_slice), cnt_str(self.host_num_params),
-        #              slc_str(self.host_param2_slice), cnt_str(self.host_num_params2),  flush=True)
-        #    resource_alloc.comm.barrier()
-        #
-        #if rank == 0:
-        #    print("%11s %-14s %-14s %-4s" % ('#', 'g-pfine', 'h-pfine', 'tot'), flush=True)
-        #resource_alloc.comm.barrier()
-        #for r in range(resource_alloc.comm.size):
-        #    if r == rank:
-        #        my_desc = ("%3d (%2d.%2d)" % (rank, resource_alloc.host_index, resource_alloc.host_comm.rank)) \
-        #                  if shm else ("%11d" % rank)
-        #        print(my_desc, slc_str(self.global_param_fine_slice), slc_str(self.host_param_fine_slice),
-        #              cnt_str(self.host_num_params_fine), flush=True)
-        #    resource_alloc.comm.barrier()
-
     @property
     def max_atom_elements(self):
         """ The most elements owned by a single atom. """