New gauge group, plus minor changes to error generator classes

pygsti/modelmembers/operations/embeddederrorgen.py

@@ -14,6 +14,7 @@
 import warnings as _warnings
 
 from pygsti.modelmembers.operations.embeddedop import EmbeddedOp as _EmbeddedOp
+from pygsti.modelmembers.operations.lindbladerrorgen import LindbladErrorgen as _LinbladErrorGen
 
 
 # Idea:
@@ -50,7 +51,7 @@ class EmbeddedErrorgen(_EmbeddedOp):
         of the EmbeddedErrorgen.
     """
 
-    def __init__(self, state_space, target_labels, errgen_to_embed):
+    def __init__(self, state_space, target_labels, errgen_to_embed: _LinbladErrorGen):
         _EmbeddedOp.__init__(self, state_space, target_labels, errgen_to_embed)
 
         # set "API" error-generator members (to interface properly w/other objects)

pygsti/modelmembers/operations/experrorgenop.py

@@ -19,11 +19,11 @@
 import scipy.sparse.linalg as _spsl
 
 from pygsti.modelmembers.operations.linearop import LinearOperator as _LinearOperator
-from pygsti.modelmembers.operations.lindbladerrorgen import LindbladParameterization as _LindbladParameterization
+from pygsti.modelmembers.operations.lindbladerrorgen import LindbladErrorgen as _LindbladErrorgen
+from pygsti.modelmembers.operations.embeddederrorgen import EmbeddedErrorgen as _EmbeddedErrorGen
 from pygsti.modelmembers import modelmember as _modelmember, term as _term
 from pygsti.modelmembers.errorgencontainer import ErrorGeneratorContainer as _ErrorGeneratorContainer
 from pygsti.baseobjs.polynomial import Polynomial as _Polynomial
-from pygsti.tools import matrixtools as _mt
 
 IMAG_TOL = 1e-7  # tolerance for imaginary part being considered zero
 MAX_EXPONENT = _np.log(_np.finfo('d').max) - 10.0  # so that exp(.) doesn't overflow
@@ -44,7 +44,7 @@ class ExpErrorgenOp(_LinearOperator, _ErrorGeneratorContainer):
         operator is `exp(L)`.
     """
 
-    def __init__(self, errorgen):
+    def __init__(self, errorgen : _LindbladErrorgen | _EmbeddedErrorGen):
         # Extract superop dimension from 'errorgen'
         state_space = errorgen.state_space
         self.errorgen = errorgen  # don't copy (allow object reuse)

pygsti/modelmembers/operations/fulltpop.py

@@ -35,15 +35,16 @@ class FullTPOp(_DenseOperator, _Torchable):
     An operation matrix that is fully parameterized except for
     the first row, which is frozen to be [1 0 ... 0] so that the action
     of the operation, when interpreted in the Pauli or Gell-Mann basis, is
-    trace preserving (TP).
+    trace preserving (TP). Bases other than Pauli or Gell-Mann are supported
+    only if their first element is the identity matrix.
 
     Parameters
     ----------
     m : array_like or LinearOperator
         a square 2D array-like or LinearOperator object representing the operation action.
         The shape of m sets the dimension of the operation.
 
-    basis : Basis or {'pp','gm','std'} or None
+    basis : Basis or {'pp','gm'} or None
         The basis used to construct the Hilbert-Schmidt space representation
         of this state as a super-operator.  If None, certain functionality,
         such as access to Kraus operators, will be unavailable.
@@ -71,8 +72,12 @@ def __init__(self, m, basis=None, evotype="default", state_space=None):
         _DenseOperator.__init__(self, mx, basis, evotype, state_space)
         assert(self._rep.base.flags['C_CONTIGUOUS'] and self._rep.base.flags['OWNDATA'])
         assert(isinstance(self._ptr, _ProtectedArray))
-        self._paramlbls = _np.array(["MxElement %d,%d" % (i, j) for i in range(1, self.dim) for j in range(self.dim)],
-                                    dtype=object)
+        self._paramlbls = _np.array(
+            ["MxElement %d,%d" % (i, j) for i in range(1, self.dim) for j in range(self.dim)], dtype=object
+        )
+        if self._basis is not None:
+            assert self._basis.first_element_is_identity  # type: ignore
+        return
 
     @property
     def _ptr(self):

pygsti/modelmembers/operations/lindbladcoefficients.py

@@ -528,7 +528,7 @@ def set_elementary_errorgens(self, elementary_errorgens, on_missing='ignore', tr
                     raise ValueError("Missing entry for %s in dictionary of elementary errorgens." % str(eeg_lbl))
             flat_data[i] = val
 
-        self.block_data[(slice(None, None),) * self.block_data.ndim] = flat_data.reshape(self.block_data.shape)
+        self.block_data[:] = flat_data.reshape(self.block_data.shape)
         self._truncate_block_data(truncate)
 
         #set a flag to indicate that the coefficients (as returned by elementary_errorgens)
@@ -652,7 +652,7 @@ def _block_data_to_params(self, truncate=False):
             in the case of `'other'` blocks.
         """
         if truncate is False:
-            ttol = -1e-15  # (was -1e-12) # truncation tolerance
+            ttol = -1e-14  # (was -1e-12) # truncation tolerance
         elif truncate is True:
             ttol = -_np.inf
         else:
@@ -725,25 +725,24 @@ def _block_data_to_params(self, truncate=False):
                 nonzero_block_data = perm_block_data[0:num_nonzero, 0:num_nonzero]
                 assert(_np.isclose(_np.linalg.norm(self.block_data), _np.linalg.norm(nonzero_block_data)))
 
-                #evals, U = _np.linalg.eigh(nonzero_block_data)  # works too (assert hermiticity above)
-                evals, U = _np.linalg.eig(nonzero_block_data)
+                evals, U = _np.linalg.eigh(nonzero_block_data)
 
                 assert(all([ev > ttol for ev in evals])), \
                     ("Lindblad coefficients are not CPTP (truncate == %s)! (largest neg = %g)"
-                     % (str(truncate), min(evals.real)))
+                     % (str(truncate), min(evals)))
 
                 if ttol < 0:  # if we're truncating and assert above allows *negative* eigenvalues
                     #push any slightly negative evals of other_projs positive so that
                     # the Cholesky decomp will work.
-                    Ui = _np.linalg.inv(U)
+                    Ui = U.T.conj()
                     pos_evals = evals.clip(1e-16, None)
-                    nonzero_block_data = _np.dot(U, _np.dot(_np.diag(pos_evals), Ui))
+                    nonzero_block_data = U @ (pos_evals[:, None] * Ui)
                     try:
                         nonzero_Lmx = _np.linalg.cholesky(nonzero_block_data)
                         # if Lmx not postitive definite, try again with 1e-12 (same lines as above)
-                    except _np.linalg.LinAlgError:                         # pragma: no cover
+                    except _np.linalg.LinAlgError:                          # pragma: no cover
                         pos_evals = evals.clip(1e-12, 1e100)                # pragma: no cover
-                        nonzero_block_data = _np.dot(U, _np.dot(_np.diag(pos_evals), Ui))  # pragma: no cover
+                        nonzero_block_data = U @ (pos_evals[:, None] * Ui)  # pragma: no cover
                         nonzero_Lmx = _np.linalg.cholesky(nonzero_block_data)
                 else:  # truncate == False or == 0 case
                     nonzero_Lmx = _np.linalg.cholesky(nonzero_block_data)
@@ -753,15 +752,15 @@ def _block_data_to_params(self, truncate=False):
                 Lmx = perm.T @ perm_Lmx @ perm
 
                 for i in range(num_bels):
-                    assert(_np.linalg.norm(_np.imag(Lmx[i, i])) < IMAG_TOL)
+                    assert(_np.abs(Lmx[i, i].imag) < IMAG_TOL)
                     params[i, i] = Lmx[i, i].real
                     for j in range(i):
                         params[i, j] = Lmx[i, j].real
                         params[j, i] = Lmx[i, j].imag
 
             elif self._param_mode == "elements":  # params mx stores block_data (hermitian) directly
                 for i in range(num_bels):
-                    assert(_np.linalg.norm(_np.imag(self.block_data[i, i])) < IMAG_TOL)
+                    assert(_np.abs(self.block_data[i, i].imag) < IMAG_TOL)
                     params[i, i] = self.block_data[i, i].real
                     for j in range(i):
                         params[i, j] = self.block_data[i, j].real
@@ -1257,4 +1256,4 @@ def __str__(self):
 
 @lru_cache(maxsize=16)
 def cached_diag_indices(n):
-    return _np.diag_indices(n)
+    return _np.diag_indices(n)

pygsti/modelmembers/operations/lindbladerrorgen.py

@@ -465,9 +465,9 @@ def from_elementary_errorgens(cls, elementary_errorgens, parameterization='auto'
             if parameterization == "auto" else LindbladParameterization.cast(parameterization)
 
         eegs_by_typ = {
-            'ham': {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type == 'H'},
-            'other_diagonal': {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type == 'S'},
-            'other': {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type != 'H'}
+            'ham':            {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type == 'H' },
+            'other_diagonal': {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type == 'S' },
+            'other':          {eeglbl: v for eeglbl, v in elementary_errorgens.items() if eeglbl.errorgen_type != 'H' }
         }
 
         blocks = []

pygsti/models/gaugegroup.py

@@ -10,14 +10,16 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 
+from typing import Tuple
 import numpy as _np
+import scipy.linalg as _la
 
-from pygsti.baseobjs import StateSpace as _StateSpace
+from pygsti.baseobjs import StateSpace as _StateSpace, statespace as _statespace
 from pygsti.modelmembers import operations as _op
-from pygsti.baseobjs import statespace as _statespace
-from pygsti.baseobjs.basis import Basis as _Basis
+from pygsti.baseobjs.basis import Basis as _Basis, BuiltinBasis as _BuiltinBasis
 from pygsti.baseobjs.nicelyserializable import NicelySerializable as _NicelySerializable
 from pygsti.evotypes.evotype import Evotype as _Evotype
+from pygsti.tools.optools import superop_to_unitary, unitary_to_superop
 
 
 class GaugeGroup(_NicelySerializable):
@@ -441,6 +443,10 @@ def __init__(self, operation):
         self._inv_matrix = None
         GaugeGroupElement.__init__(self)
 
+    @property
+    def operation(self):
+        return self._operation
+
     @property
     def transform_matrix(self):
         """
@@ -750,6 +756,16 @@ def __init__(self, operation):
         TPGaugeGroupElement.__init__(self, operation)
 
 
+class UnitaryGaugeGroupElement(OpGaugeGroupElement):
+
+    def __init__(self, operation: _op.ExpErrorgenOp):
+        OpGaugeGroupElement.__init__(self, operation)
+
+    @property
+    def operation(self) -> _op.ExpErrorgenOp:
+        return self._operation # type: ignore
+
+
 class UnitaryGaugeGroup(OpGaugeGroupWithBasis):
     """
     A gauge group consisting of unitary gauge-transform matrices.
@@ -781,24 +797,8 @@ def __init__(self, state_space, basis, evotype='default'):
         operation = _op.ExpErrorgenOp(errgen)
         OpGaugeGroupWithBasis.__init__(self, operation, UnitaryGaugeGroupElement, "Unitary", basis)
 
-
-class UnitaryGaugeGroupElement(OpGaugeGroupElement):
-    """
-    Element of a :class:`UnitaryGaugeGroup`
-
-    Parameters
-    ----------
-    operation : LinearOperator
-        The operation to base this element on. It provides both parameterization
-        information and the gauge transformation matrix itself.
-    """
-
-    def __init__(self, operation):
-        """
-        Creates a new gauge group element based on `operation`, which
-        is assumed to have the correct parameterization.
-        """
-        OpGaugeGroupElement.__init__(self, operation)
+    def compute_element(self, param_vec) -> UnitaryGaugeGroupElement:
+        return super().compute_element(param_vec)
 
 
 class SpamGaugeGroup(OpGaugeGroup):
@@ -1094,3 +1094,150 @@ def _to_nice_serialization(self):
     @classmethod
     def _from_nice_serialization(cls, state):  # memo holds already de-serialized objects
         return cls(state['operation_dimension'])
+
+
+class DirectSumUnitaryGroup(GaugeGroup):
+    """
+    A subgroup of the unitary group, where the unitary operators in the group all have a
+    shared block-diagonal structure.
+
+    Example setting where this is useful: 
+        The system's Hilbert space is naturally expressed as a direct sum, H = U ⨁ V,
+        and we want gauge optimization to preserve the natural separation between U and V.
+    """
+
+    def __init__(self, subgroups: Tuple[UnitaryGaugeGroup | TrivialGaugeGroup, ...], basis, name="Direct sum gauge group"):
+        self.subgroups = subgroups
+        self.basis = basis
+        self._param_dims = _np.array([sg.num_params for sg in subgroups])
+        self._num_params = _np.sum(self._param_dims)
+        super().__init__(name)
+
+    @property
+    def num_params(self):
+        return self._num_params
+
+    def compute_element(self, param_vec):
+        assert param_vec.size == self.num_params
+        subelements = []
+        offset = 0
+        for pd, sg in zip(self._param_dims, self.subgroups):
+            sub_param_vec = param_vec[offset:offset + pd]
+            sub_element = sg.compute_element(sub_param_vec)
+            subelements.append(sub_element)
+            offset += pd
+        out = DirectSumUnitaryGroupElement(tuple(subelements), self.basis)
+        return out
+
+    @property
+    def initial_params(self):
+        paramvecs = [sg.initial_params for sg in self.subgroups]
+        return _np.concatenate(paramvecs)
+
+    def _to_nice_serialization(self):
+        state = super()._to_nice_serialization()
+        state.update({
+            'subgroups': tuple([se.to_nice_serialization() for se in self.subgroups]),
+            'state_space_dimension': int(self.basis.dim),
+            'basis': self.basis if isinstance(self.basis, str) else self.basis.to_nice_serialization(),
+            'name': self.name
+        })
+        return state
+
+    @classmethod
+    def _from_nice_serialization(cls, state):
+        subgroups = []
+        for sg_dict in state['subgroups']:
+            if sg_dict['class'] == 'UnitaryGaugeGroup':
+                sg = UnitaryGaugeGroup.from_nice_serialization(sg_dict)
+            else:
+                sg = TrivialGaugeGroup.from_nice_serialization(sg_dict)
+            subgroups.append(sg)
+        subgroups = tuple(subgroups)
+        basis = state['basis'] if isinstance(state['basis'], str) else _Basis.from_nice_serialization(state['basis'])
+        name = state['name']
+        return cls(subgroups, basis, name)
+
+
+class DirectSumUnitaryGroupElement(GaugeGroupElement):
+
+    def __init__(self, subelements: Tuple[UnitaryGaugeGroupElement | TrivialGaugeGroupElement, ...], basis):
+        self.subelements = subelements
+        self.basis = basis
+        self._update_matrices()
+        self._vectorized_op_dim = self.basis.dim ** 2
+        self._num_params = _np.sum([se.num_params for se in self.subelements])
+        super().__init__()
+
+    @property
+    def transform_matrix(self):
+        return self._m
+
+    @property
+    def transform_matrix_inverse(self):
+        return self._invm
+
+    def deriv_wrt_params(self, wrt_filter=None):
+        raise NotImplementedError()
+
+    def to_vector(self):
+        v = _np.concatenate([se.to_vector() for se in self.subelements])
+        return v
+
+    def from_vector(self, v):
+        assert v.size == self.num_params
+        offset = 0
+        for se in self.subelements:
+            block = v[offset:(offset + se.num_params)]
+            se.from_vector(block)
+            offset += se.num_params
+        self._update_matrices()
+        return
+
+    def _update_matrices(self):
+        u_blocks, num_params = [], []
+        for se in self.subelements:
+            if isinstance(se, TrivialGaugeGroupElement):
+                se_ss = _statespace.default_space_for_dim(se.transform_matrix.shape[0])
+                se_basis = _BuiltinBasis('std', se_ss.dim)
+                u_blocks.append(_np.eye(se_ss.udim))
+            else:
+                u_abstract_op = se.operation
+                se_basis = u_abstract_op.errorgen.matrix_basis
+                u = superop_to_unitary(se.transform_matrix, se_basis)
+                u_blocks.append(u)
+            num_params.append(se.num_params)
+        u = _la.block_diag(*u_blocks)
+        self._u = u
+        self._m = unitary_to_superop(self._u, self.basis)
+        self._invm = unitary_to_superop(self._u.T.conj(), self.basis)
+        return
+
+    @property
+    def num_params(self):
+        return self._num_params
+
+    def inverse(self):
+        return InverseGaugeGroupElement(self)
+
+    def _to_nice_serialization(self):
+        state = super()._to_nice_serialization()
+        state.update({
+            'class': self.__class__.__name__,
+            'subelements': tuple([se.to_nice_serialization() for se in self.subelements]),
+            'basis': self.basis if isinstance(self.basis, str) else self.basis.to_nice_serialization()
+        })
+        return state
+
+    @classmethod
+    def _from_nice_serialization(cls, state):
+        subelements = []
+        for se_dict in state['subelements']:
+            if se_dict['class'] == 'UnitaryGaugeGroupElement':
+                se = UnitaryGaugeGroupElement.from_nice_serialization(se_dict)
+            else:
+                se = TrivialGaugeGroupElement.from_nice_serialization(se_dict)
+            subelements.append(se)
+        subelements = tuple(subelements)
+        basis = state['basis'] if isinstance(state['basis'], str) else _Basis.from_nice_serialization(state['basis'])
+        return cls(subelements, basis)