Avoid accessing non-existent state_space field of some Basis classes (e.g., TensorProdBasis) (sandialabs#842)

Fixes a bug in deserializing models that use a TensorProdBasis, and a similar bug in leakage/gaugeopt.py.

Also cleans up whitespace and adds type annotations.

Author: rileyjmurray

pygsti/baseobjs/errorgenbasis.py

@@ -63,6 +63,7 @@ def _all_elements_same_type(lst):
             return False
     return True
 
+
 class ExplicitElementaryErrorgenBasis(ElementaryErrorgenBasis):
     """
     This basis object contains the information  necessary for building, 
@@ -329,6 +330,7 @@ def difference(self, other_basis):
         difference_state_space = self.state_space
         return ExplicitElementaryErrorgenBasis(difference_state_space, sorted(difference_labels, key=lambda label: label.__str__()), self._basis_1q)
 
+
 class CompleteElementaryErrorgenBasis(ElementaryErrorgenBasis):
     """
     This basis object contains the information  necessary for building, 
@@ -398,7 +400,6 @@ def _count_uptriangle_labels_for_support(cls, support, left_support, type_str, t
 
         return cnt
 
-
     @classmethod
     def _create_ordered_labels(cls, type_str, basis_1q, state_space,
                                max_weight=None, sslbl_overlap=None,
@@ -880,4 +881,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/leakage/gaugeopt.py

@@ -12,15 +12,16 @@
 import copy
 from typing import Dict, List, Optional, TYPE_CHECKING, Union
 
-from pygsti.baseobjs.basis import Basis
+from pygsti.baseobjs.basis import Basis, BuiltinBasis
 
 if TYPE_CHECKING:
     from pygsti.models import ExplicitOpModel
     from pygsti.protocols.gst import GSTGaugeOptSuite, ModelEstimateResults
 
 
-def _direct_sum_unitary_group(subspace_bases, full_basis, triviality_flags=None,
-                              level_partition=None):
+def _direct_sum_unitary_group(subspace_bases: list[BuiltinBasis], full_basis: Basis,
+        triviality_flags=None, level_partition=None
+    ):
     """
     Build a gauge group that acts as an independent unitary on each summand of a
     direct-sum decomposition H = H₀ ⨁ H₁ ⨁ ... of Hilbert space.
@@ -186,7 +187,7 @@ def lagoified_gopparams_dicts(gopparams_dicts: List[Dict]) -> List[Dict]:
         # ^ We use subspace-restricted loss functions that only care about mismatches
         #   between an estimate and a target when restricted to the computational subspace.
         #
-        gg = UnitaryGaugeGroup(tm.basis.state_space, tm.basis)
+        gg = UnitaryGaugeGroup(tm.state_space, tm.basis)
         inner_dict['gauge_group'] = gg
         inner_dict['_gaugeGroupEl'] = gg.compute_element(gg.initial_params)
         # ^ We insist on the unitary gauge group because other common gauge groups

pygsti/modelmembers/states/densestate.py

@@ -142,7 +142,10 @@ def _from_memoized_dict(cls, mm_dict, serial_memo):
             if 'Basis object has unexpected dimension' in se and len(serial_memo) > 0:
                 member = list(serial_memo.values())[0]
                 basis = member.parent.basis
-                state_space = basis.state_space
+                # Use the parent model's state_space directly rather than basis.state_space:
+                # composite bases (e.g. TensorProdBasis, DirectSumBasis) don't define state_space,
+                # while every Model is guaranteed to have one.
+                state_space = member.parent.state_space
                 return cls(vec, basis, mm_dict['evotype'], state_space)
             raise e
 

test/unit/algorithms/test_gaugeopt_correctness.py

@@ -424,6 +424,36 @@ def test_non_coordinate_subspace_rejected(self):
             _leakage_direct_sum_group(basis)
 
 
+class LagoifiedGopparamsDictsTester(BaseCase):
+    """
+    Regression test for pygsti.leakage.gaugeopt.lagoified_gopparams_dicts.
+
+    It used to build its unitary gauge group from `tm.basis.state_space`, which raises
+    AttributeError whenever `tm.basis` is a composite basis (e.g. TensorProdBasis) --
+    the common case for any multi-subsystem model, leakage or not. It should instead use
+    `tm.state_space`, which every Model has regardless of its basis's type.
+    """
+
+    def test_tensor_prod_basis_target_model_does_not_crash(self):
+        from pygsti.baseobjs import ExplicitStateSpace
+        from pygsti.baseobjs.basis import Basis, TensorProdBasis
+        from pygsti.leakage.gaugeopt import lagoified_gopparams_dicts
+        from pygsti.models.gaugegroup import DirectSumUnitaryGroup
+
+        # A qubit tensored with a leaky qutrit: `basis` is a TensorProdBasis, which (unlike
+        # BuiltinBasis) has no `state_space` attribute of its own.
+        basis = TensorProdBasis((Basis.cast('pp', 4), Basis.cast('l2p1', 9)))
+        self.assertFalse(hasattr(basis, 'state_space'))
+        state_space = ExplicitStateSpace(['Q0', 'Q1'], [2, 3])
+        target_model = ExplicitOpModel(state_space, basis)
+
+        gopparams_dicts = lagoified_gopparams_dicts([{'target_model': target_model}])
+
+        self.assertEqual(len(gopparams_dicts), 2)
+        self.assertIsInstance(gopparams_dicts[0]['gauge_group'], UnitaryGaugeGroup)
+        self.assertIsInstance(gopparams_dicts[1]['gauge_group'], DirectSumUnitaryGroup)
+
+
 class FindPerfectGauge_DirectSumGaugeGroup4LevelTester(BaseCase):
     """
     Gauge-recovery test for the generalized direct-sum gauge group on a 4-level