FOGI model serialization and some FOGI model optimization

pygsti/baseobjs/errorgenbasis.py

@@ -16,11 +16,12 @@
 from pygsti.baseobjs import Basis as _Basis
 from pygsti.baseobjs.errorgenlabel import GlobalElementaryErrorgenLabel as _GlobalElementaryErrorgenLabel,\
 LocalElementaryErrorgenLabel as _LocalElementaryErrorgenLabel
-
+from pygsti.baseobjs.nicelyserializable import NicelySerializable as _NicelySerializable
 from pygsti.tools import optools as _ot
+from pygsti.baseobjs.statespace import StateSpace as _StateSpace
 
 
-class ElementaryErrorgenBasis(object):
+class ElementaryErrorgenBasis(_NicelySerializable):
     """
     A basis for error-generator space defined by a set of elementary error generators.
 
@@ -49,6 +50,7 @@ def __len__(self):
         Number of elementary errorgen elements in this basis.
         """
         return len(self.labels)
+
 
 #helper function for checking label types.
 def _all_elements_same_type(lst):
@@ -88,6 +90,7 @@ def __init__(self, state_space, labels, basis_1q=None):
             comprise the basis element labels for the values of the
             `ElementaryErrorgenLabels` in `labels`.
         """
+        super().__init__()
         labels = tuple(labels)
 
         #add an assertion that the labels are ElementaryErrorgenLabels and that all of the labels are the same type.
@@ -116,6 +119,34 @@ def __init__(self, state_space, labels, basis_1q=None):
         self._cached_dual_matrices = None
         self._cached_supports = None
 
+    def __eq__(self, other):
+        """Compare self with other. Return true only if they are identical, including the order of their labels.
+
+        Args:
+            other (ExplicitElementaryErrorgenBasis): Error generator basis to compare against
+
+        Returns:
+            Boolean: True if they are identical, False otherwise
+        """
+        if not isinstance(other, ExplicitElementaryErrorgenBasis):
+            return False
+        return self.state_space.__eq__(other.state_space) and [label.__str__() for label in self.labels] == [label.__str__() for label in other.labels] and self._basis_1q.__eq__(other._basis_1q)
+
+    def _to_nice_serialization(self):
+        state = super()._to_nice_serialization()
+        state.update({'state_space' : self.state_space._to_nice_serialization(),
+                      'labels' : [label.__str__() for label in self.labels],
+                      'label_type' : 'global' if isinstance(self.labels[0], _GlobalElementaryErrorgenLabel) else 'local',
+                      '_basis_1q' : self._basis_1q if isinstance(self._basis_1q, str) else self._basis_1q._to_nice_serialization()
+        })
+        return state
+    @classmethod
+    def from_nice_serialization(cls, state):
+        if state['label_type'] == 'global':
+            cast = _GlobalElementaryErrorgenLabel.cast
+        else:
+            cast = _LocalElementaryErrorgenLabel.cast
+        return cls(_StateSpace.from_nice_serialization(state['state_space']), [cast(label) for label in state['labels']], state['_basis_1q'] if isinstance(state['_basis_1q'], str) else _Basis.from_nice_serialization(state['_basis_1q']))
     @property
     def labels(self):
         return self._labels
@@ -259,7 +290,7 @@ def union(self, other_basis):
         #assert that these two bases have compatible label types.
         msg = 'Incompatible `ElementaryErrrogenLabel` types, the two `ElementaryErrorgenBasis` should have the same label type.'
         assert type(self._labels[0]) == type(other_basis.labels[0]), msg
-        
+
         #Get the union of the two bases labels.
         union_labels = set(self._labels) | set(other_basis.labels)
         union_state_space = self.state_space.union(other_basis.state_space)
@@ -516,7 +547,7 @@ def __init__(self, basis_1q, state_space, elementary_errorgen_types=('H', 'S', '
             `GlobalElementaryErrorgenLabel` and `LocalElementaryErrorgenLabel`,
             respectively.
         """
-
+        super().__init__()
         if isinstance(basis_1q, _Basis):
             self._basis_1q = basis_1q
         elif isinstance(basis_1q, str):
@@ -589,6 +620,24 @@ def __init__(self, basis_1q, state_space, elementary_errorgen_types=('H', 'S', '
         #                      (IXX,XXI) #   on right, loop over all possible choices of at least one, an at most m,
         #                      (IXX,XXX) #    nontrivial indices to place within the m nontriv left indices (1 & 2 here)
 
+    def _to_nice_serialization(self):
+
+        state = super()._to_nice_serialization()
+        state.update({'basis_1q' : self._basis_1q.to_nice_serialization(),
+                      'state_space' : self.state_space.to_nice_serialization(),
+                      'elementary_errorgen_types': self._elementary_errorgen_types,
+                      'max_weights' : self.max_weights,
+                      'sslbl_overlap' : self._sslbl_overlap,
+                      'default_label_type': self._default_lbl_typ 
+        })
+        return state
+
+    @classmethod
+    def from_nice_serialization(cls, state):
+
+        return cls(_Basis.from_nice_serialization(state['basis_1q']), _StateSpace.from_nice_serialization(state['state_space']), \
+                   state['elementary_errorgen_types'], max_weights=state['max_weights'], sslbl_overlap=state['sslbl_overlap'], \
+                    default_label_type=state['default_label_type'])
     def __len__(self):
         """ Number of elementary errorgen elements in this basis """
         return self._offsets[self._elementary_errorgen_types[-1]]['END']
@@ -741,6 +790,7 @@ def label_index(self, label, ok_if_missing=False, identity_label='I'):
         identity_label : str, optional (default 'I')
             An optional string specifying the label used to denote the identity in basis element labels.
         """
+
         if isinstance(label, _LocalElementaryErrorgenLabel):
             label = _GlobalElementaryErrorgenLabel.cast(label, self.sslbls, identity_label=identity_label)
 
@@ -762,8 +812,10 @@ def label_index(self, label, ok_if_missing=False, identity_label='I'):
         elif eetype in ('C', 'A'):
             assert(len(trivial_bel) == 1)  # assumes this is a single character
             nontrivial_inds = [i for i, letter in enumerate(bels[0]) if letter != trivial_bel]
-            left_support = tuple([self.sslbls[i] for i in nontrivial_inds])
-
+            left_support = tuple([label.sslbls[i] for i in nontrivial_inds])
+            #left_support = tuple([self.sslbls[i] for i in nontrivial_inds])
+            #This line above is supposed to be a bugfix, I want to verify this with
+            #Corey before removing this comment
             if ok_if_missing and (support, left_support) not in self._offsets[eetype]:
                 return None
             base = self._offsets[eetype][(support, left_support)]
@@ -772,7 +824,6 @@ def label_index(self, label, ok_if_missing=False, identity_label='I'):
                 support, left_support, eetype, [trivial_bel], nontrivial_bels))}
         else:
             raise ValueError("Invalid elementary errorgen type: %s" % str(eetype))
-
         return base + indices[label]
 
     def create_subbasis(self, sslbl_overlap, retain_max_weights=True):
@@ -829,4 +880,4 @@ def difference(self, other_basis):
         other_basis : `ElementaryErrorgenBasis`
             `ElementaryErrorgenBasis` to construct the difference with.
         """
-        return self.to_explicit_basis().difference(other_basis)
+        return self.to_explicit_basis().difference(other_basis)

pygsti/baseobjs/errorgenspace.py

@@ -11,26 +11,44 @@
 #***************************************************************************************************
 
 import numpy as _np
-
+from pygsti.baseobjs.nicelyserializable import NicelySerializable as _NicelySerializable
 from pygsti.tools import matrixtools as _mt
+from pygsti.baseobjs.errorgenbasis import ExplicitElementaryErrorgenBasis
 
+class ErrorgenSpace(_NicelySerializable):
 
-class ErrorgenSpace(object):
-    """
-    A vector space of error generators, spanned by some basis.
-
-    This object collects the information needed to specify a space
-    within the space of all error generators.
-    """
 
     def __init__(self, vectors, basis):
+        """
+        A vector space of error generators, spanned by some basis.
+
+        This object collects the information needed to specify a space
+        within the space of all error generators.
+
+        Parameters
+        ----------
+        vectors : numpy array
+            List of vectors that span the space
+
+            elemgen_basis : ElementaryErrorgenBasis
+                The elementary error generator basis that define the entries of self.vectors 
+        """
+        super().__init__()
         self.vectors = vectors
         self.elemgen_basis = basis
         #Question: have multiple bases or a single one?
         #self._vectors = [] if (items is None) else items  # list of (basis, vectors_mx) pairs
         # map sslbls => (vectors, basis) where basis.sslbls == sslbls
         # or basis => vectors if bases can hash well(?)
-
+    def _to_nice_serialization(self):
+        state = super()._to_nice_serialization()
+        state.update({'vectors' : self._encodemx(self.vectors),
+                      'basis': self.elemgen_basis._to_nice_serialization()
+        })
+        return state
+    @classmethod
+    def from_nice_serialization(cls, state):
+        return cls(cls._decodemx(state['vectors']), ExplicitElementaryErrorgenBasis.from_nice_serialization(state['basis']))
     def intersection(self, other_space, free_on_unspecified_space=False, use_nice_nullspace=False):
         """
         TODO: docstring
@@ -75,6 +93,18 @@ def intersection(self, other_space, free_on_unspecified_space=False, use_nice_nu
 
         return ErrorgenSpace(intersection_vecs, common_basis)
 
+    def __eq__(self, other):
+        """Compare self with other. Return true only if they are identical, including the order of their vectors.
+
+        Args:
+            other (ErrorgenSpace): Error generator space to compare against
+
+        Returns:
+            Boolean: True if they are identical, False otherwise
+        """
+        if not isinstance(other, ErrorgenSpace):
+            return False
+        return _np.allclose(self.vectors, other.vectors) and self.elemgen_basis.__eq__(other.elemgen_basis)
     def union(self, other_space):
         """
         TODO: docstring
@@ -96,4 +126,4 @@ def normalize(self, norm_order=2):
         """
         for j in range(self.vectors.shape[1]):
             sign = +1 if max(self.vectors[:, j]) >= -min(self.vectors[:, j]) else -1
-            self.vectors[:, j] /= sign * _np.linalg.norm(self.vectors[:, j], ord=norm_order)
+            self.vectors[:, j] /= sign * _np.linalg.norm(self.vectors[:, j], ord=norm_order)

pygsti/baseobjs/statespace.py

@@ -418,7 +418,7 @@ def create_subspace(self, labels):
         StateSpace
         """
         # Default, generic, implementation constructs an explicit state space
-        labels = set(labels)
+        labels = sorted(set(labels))
         sub_tpb_labels = []
         sub_tpb_udims = []
         sub_tpb_types = []
@@ -1389,7 +1389,6 @@ def __str__(self):
             ['*'.join(["%s(%d%s)" % (lbl, self.label_dims[lbl], 'c' if (self.label_types[lbl] == 'C') else '')
                        for lbl in tpb]) for tpb in self._labels])
 
-
 def default_space_for_dim(dim):
     """
     Create a state space for a given superoperator dimension.

pygsti/models/explicitmodel.py

@@ -23,6 +23,7 @@
 from pygsti.models.memberdict import OrderedMemberDict as _OrderedMemberDict
 from pygsti.models.layerrules import LayerRules as _LayerRules
 from pygsti.models.modelparaminterposer import ModelParamsInterposer as _ModelParamsInterposer
+from pygsti.models.fogistore import FirstOrderGaugeInvariantStore as _FirstOrderGaugeInvariantStore
 from pygsti.models.gaugegroup import GaugeGroup as _GaugeGroup
 from pygsti.forwardsims.forwardsim import ForwardSimulator as _FSim
 from pygsti.forwardsims import matrixforwardsim as _matrixfwdsim
@@ -1603,7 +1604,9 @@ def _to_nice_serialization(self):
                       'default_gauge_group': (self.default_gauge_group.to_nice_serialization()
                                               if (self.default_gauge_group is not None) else None),
                       'parameter_interposer': (self._param_interposer.to_nice_serialization()
-                                               if (self._param_interposer is not None) else None)
+                                               if (self._param_interposer is not None) else None),
+                      'fogi_store': (self.fogi_store.to_nice_serialization()
+                                               if (self.fogi_store is not None) else None)
                       })
 
         mmgraph = self.create_modelmember_graph()
@@ -1619,6 +1622,8 @@ def _from_nice_serialization(cls, state):
             if (state['default_gauge_group'] is not None) else None
         param_interposer = _ModelParamsInterposer.from_nice_serialization(state['parameter_interposer']) \
             if (state['parameter_interposer'] is not None) else None
+        fogi_store = _FirstOrderGaugeInvariantStore.from_nice_serialization(state['fogi_store']) \
+            if (state['fogi_store'] is not None) else None
         param_labels = state.get('parameter_labels', None)
         param_bounds = state.get('parameter_bounds', None)
 
@@ -1637,6 +1642,7 @@ def _from_nice_serialization(cls, state):
         mdl._clean_paramvec()
         mdl.default_gauge_group = default_gauge_group
         mdl.param_interposer = param_interposer
+        mdl.fogi_store = fogi_store
 
         Np = len(mdl._paramlbls)  # _clean_paramvec sets up ._paramlbls so its length == # of params
         if param_labels and len(param_labels) == Np: