From ea67d36a31d5e65240d5fcbc734cbcd2b3c6e1c0 Mon Sep 17 00:00:00 2001
From: Nick Koskelo <koskelo2@illinois.edu>
Date: Fri, 25 Jul 2025 16:05:51 -0700
Subject: [PATCH 1/2] Add type annotations for the basis files and correcting
 hypothesis significance levels in conducting multiple  tests.

---
 pygsti/baseobjs/basis.py             | 201 ++++++++++++++-------------
 pygsti/baseobjs/basisconstructors.py |  22 +--
 pygsti/tools/hypothesis.py           |  27 +++-
 3 files changed, 137 insertions(+), 113 deletions(-)

diff --git a/pygsti/baseobjs/basis.py b/pygsti/baseobjs/basis.py
index b195a3ef1..9d4061790 100644
--- a/pygsti/baseobjs/basis.py
+++ b/pygsti/baseobjs/basis.py
@@ -9,12 +9,21 @@
 # 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
+# Used to allow certain type annotations. (e.g. class method returning instance of the class).
 
 import copy as _copy
 import itertools as _itertools
 import warnings as _warnings
 from functools import lru_cache
-from typing import Union
+from typing import (
+    Union,
+    TypeAlias,
+    Literal,
+    Optional,
+    Sequence,
+    Iterable,
+)
 
 import numpy as _np
 import scipy.sparse as _sps
@@ -94,14 +103,14 @@ class Basis(_NicelySerializable):
     matrices.  The `size` of this basis is `2`, as there are two elements
     (and two vector elements).  Since vector elements are the length-4
     flattened Pauli matrices, the dimension (`dim`) is `4`.  Since the
-    elements are 2x2 Pauli matrices, the `elshape` is `(2,2)`.
+    elements are 2x2 Pauli matrices, the `elshape` is `(2, 2)`.
 
     As another example consider a basis which spans all the diagonal
     2x2 matrices.  The elements of this basis are the two matrix units
-    with a 1 in the (0,0) or (1,1) location.  The vector elements,
-    however, are the length-2 [1,0] and [0,1] vectors obtained by extracting
+    with a 1 in the (0, 0) or (1, 1) location.  The vector elements,
+    however, are the length-2 [1, 0] and [0, 1] vectors obtained by extracting
     just the diagonal entries from each basis element.  Thus, for this
-    basis, `size=2`, `dim=2`, and `elshape=(2,2)` - so the dimension is
+    basis, `size=2`, `dim=2`, and `elshape=(2, 2)` - so the dimension is
     not just the product of `elshape` entries (equivalently, `elsize`).
 
     Parameters
@@ -142,7 +151,7 @@ class Basis(_NicelySerializable):
     elshape : int
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
 
     elndim : int
@@ -159,7 +168,7 @@ class Basis(_NicelySerializable):
     # could be static methods.
 
     @classmethod
-    def cast_from_name_and_statespace(cls, name: str, state_space: _StateSpace, sparse=None):
+    def cast_from_name_and_statespace(cls, name: str, state_space: _StateSpace, sparse: Optional[bool] = None) -> Basis:
         tpbBases = []
         block_labels = state_space.tensor_product_blocks_labels
         if len(block_labels) == 1 and len(block_labels[0]) == 1:
@@ -175,15 +184,15 @@ def cast_from_name_and_statespace(cls, name: str, state_space: _StateSpace, spar
                     tpbBases.append(BuiltinBasis(nm, state_space.label_dimension(tpbLabels[0]), sparse))
                 else:
                     tpbBases.append(TensorProdBasis([
-                        BuiltinBasis(name if (state_space.label_type(l) == 'Q') else 'cl',
-                                        state_space.label_dimension(l), sparse) for l in tpbLabels]))
+                        BuiltinBasis(name if (state_space.label_type(lbl) == 'Q') else 'cl',
+                                        state_space.label_dimension(lbl), sparse) for lbl in tpbLabels]))
         if len(tpbBases) == 1:
             return tpbBases[0]
         else:
             return DirectSumBasis(tpbBases)
 
     @classmethod
-    def cast_from_name_and_dims(cls, name: str, dim: Union[int,list,tuple], sparse=None):
+    def cast_from_name_and_dims(cls, name: str, dim: Union[int, list, tuple], sparse: Optional[bool] = None) -> Basis:
         if isinstance(dim, (list, tuple)):  # list/tuple of block dimensions
             tpbBases = []
             for tpbDim in dim:
@@ -199,48 +208,48 @@ def cast_from_name_and_dims(cls, name: str, dim: Union[int,list,tuple], sparse=N
                 return DirectSumBasis(tpbBases)
         else:
             return BuiltinBasis(name, dim, sparse)
-        
+
     @classmethod
-    def cast_from_basis(cls, basis, dim=None, sparse=None):
+    def cast_from_basis(cls, basis: Basis, dim=None, sparse: Optional[bool] = None) -> Basis:
         #then just check to make sure consistent with `dim` & `sparse`
         if dim is not None:
             if isinstance(dim, _StateSpace):
                 state_space = dim
                 if hasattr(basis, 'state_space'):  # TODO - should *all* basis objects have a state_space?
-                    assert(state_space.is_compatible_with(basis.state_space)), \
+                    assert (state_space.is_compatible_with(basis.state_space)), \
                         "Basis object has incompatible state space: %s != %s" % (str(state_space),
                                                                                     str(basis.state_space))
             else:  # assume dim is an integer
-                assert(dim == basis.dim or dim == basis.elsize), \
+                assert (dim == basis.dim or dim == basis.elsize), \
                     "Basis object has unexpected dimension: %d != %d or %d" % (dim, basis.dim, basis.elsize)
         if sparse is not None:
             basis = basis.with_sparsity(sparse)
         return basis
 
     @classmethod
-    def cast_from_arrays(cls, arrays, dim=None, sparse=None):
+    def cast_from_arrays(cls, arrays: _np.ndarray, dim=None, sparse: Optional[bool] = None) -> Basis:
         b = ExplicitBasis(arrays, sparse=sparse)
         if dim is not None:
-            assert(dim == b.dim), "Created explicit basis has unexpected dimension: %d vs %d" % (dim, b.dim)
+            assert (dim == b.dim), "Created explicit basis has unexpected dimension: %d vs %d" % (dim, b.dim)
         if sparse is not None:
-            assert(sparse == b.sparse), "Basis object has unexpected sparsity: %s" % (b.sparse)
+            assert (sparse == b.sparse), "Basis object has unexpected sparsity: %s" % (b.sparse)
         return b
 
     @classmethod
-    def cast(cls, arg, dim=None, sparse=None):
-        #print("DB: CAST = ",arg,dim)
+    def cast(cls, arg, dim=None, sparse: Optional[bool] = None) -> Basis:
+        #print("DB: CAST = ", arg, dim)
         if isinstance(arg, Basis):
             return cls.cast_from_basis(arg, dim, sparse)
         if isinstance(arg, str):
             if isinstance(dim, _StateSpace):
                 return cls.cast_from_name_and_statespace(arg, dim, sparse)
             return cls.cast_from_name_and_dims(arg, dim, sparse)
-        if (arg is None) or (hasattr(arg,'__len__') and len(arg) == 0):
+        if (arg is None) or (hasattr(arg, '__len__') and len(arg) == 0):
             return ExplicitBasis([], [], "*Empty*", "Empty (0-element) basis", False, sparse)
             # ^ The original implementation would return this value under two conditions.
-            #   Either arg was None, or isinstance(arg,(tuple,list,ndarray)) and len(arg) == 0.
+            #   Either arg was None, or isinstance(arg, (tuple, list, ndarray)) and len(arg) == 0.
             #   We're just slightly relaxing the type requirement by using this check instead.
-        
+
         # At this point, original behavior would check that arg is a tuple, list, or ndarray.
         # Instead, we'll just require that arg[0] is well-defined. This is enough to discern
         # between the two cases we can still support.
@@ -252,8 +261,7 @@ def cast(cls, arg, dim=None, sparse=None):
 
         raise ValueError("Can't cast %s to be a basis!" % str(type(arg)))
 
-
-    def __init__(self, name, longname, real, sparse):
+    def __init__(self, name: str, longname: str, real: bool, sparse: bool):
         super().__init__()
         self.name = name
         self.longname = longname
@@ -283,7 +291,7 @@ def elshape(self):
         """
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
         """
         # shape of "natural" elements - size may be > self.dim (to display naturally)
@@ -322,7 +330,7 @@ def first_element_is_identity(self):
         d = self.elshape[0]
         return _np.allclose(self.elements[0], _np.identity(d) * (_np.linalg.norm(self.elements[0]) / _np.sqrt(d)))
 
-    def is_simple(self):
+    def is_simple(self) -> bool:
         """
         Whether the flattened-element vector space is the *same* space as the space this basis's vectors belong to.
 
@@ -332,7 +340,7 @@ def is_simple(self):
         """
         return self.elsize == self.dim
 
-    def is_complete(self):
+    def is_complete(self) -> bool:
         """
         Whether this is a complete basis, i.e. this basis's vectors span the entire space that they live in.
 
@@ -342,7 +350,7 @@ def is_complete(self):
         """
         return self.dim == self.size
 
-    def is_partial(self):
+    def is_partial(self) -> bool:
         """
         The negative of :meth:`is_complete`, effectively "is_incomplete".
 
@@ -368,7 +376,7 @@ def vector_elements(self):
             # Use flatten (rather than ravel) to ensure a copy is made.
             return [el.flatten() for el in self.elements]
 
-    def copy(self):
+    def copy(self) -> Basis:
         """
         Make a copy of this Basis object.
 
@@ -378,7 +386,7 @@ def copy(self):
         """
         return _copy.deepcopy(self)
 
-    def with_sparsity(self, desired_sparsity):
+    def with_sparsity(self, desired_sparsity: bool) -> Basis:
         """
         Returns either this basis or a copy of it with the desired sparsity.
 
@@ -400,7 +408,7 @@ def with_sparsity(self, desired_sparsity):
         else:
             return self._copy_with_toggled_sparsity()
 
-    def _copy_with_toggled_sparsity(self):
+    def _copy_with_toggled_sparsity(self) -> Basis:
         raise NotImplementedError("Derived classes should implement this!")
 
     def __str__(self):
@@ -418,7 +426,7 @@ def __len__(self):
     def __eq__(self, other):
         return self.is_equivalent(other, sparseness_must_match=True)
 
-    def is_equivalent(self, other, sparseness_must_match=True):
+    def is_equivalent(self, other, sparseness_must_match: bool = True) -> bool:
         """
         Tests whether this basis is equal to another basis, optionally ignoring sparseness.
 
@@ -454,7 +462,7 @@ def is_equivalent(self, other, sparseness_must_match=True):
                 return _np.array_equal(self.elements, other.elements)
             else:
                 return _np.array_equal(self.elements, other)
-    
+
     @lru_cache(maxsize=4)
     def create_transform_matrix(self, to_basis):
         """
@@ -528,7 +536,7 @@ def is_normalized(self):
         """
         if self.elndim == 2:
             for i, mx in enumerate(self.elements):
-                t = _np.linalg.norm(mx) # == sqrt(tr(mx mx))
+                t = _np.linalg.norm(mx)  # == sqrt(tr(mx mx))
                 if not _np.isclose(t, 1.0): return False
             return True
         elif self.elndim == 1:
@@ -575,10 +583,10 @@ def from_std_transform_matrix(self):
             if self.is_complete():
                 return _spsl.inv(self.to_std_transform_matrix.tocsc()).tocsr()
             else:
-                assert(self.size < self.dim), "Basis seems to be overcomplete: size > dimension!"
+                assert (self.size < self.dim), "Basis seems to be overcomplete: size > dimension!"
                 # we'd need to construct a different pseudo-inverse if the above assert fails
 
-                A = self.to_std_transform_matrix  # shape (dim,size) - should have indep *cols*
+                A = self.to_std_transform_matrix  # shape (dim, size) - should have indep *cols*
                 Adag = A.getH()        # shape (size, dim)
                 invAdagA = _spsl.inv(Adag.tocsr().dot(A.tocsc())).tocsr()
                 return invAdagA.dot(Adag.tocsc())
@@ -586,10 +594,10 @@ def from_std_transform_matrix(self):
             if self.is_complete():
                 return _inv(self.to_std_transform_matrix)
             else:
-                assert(self.size < self.dim), "Basis seems to be overcomplete: size > dimension!"
+                assert (self.size < self.dim), "Basis seems to be overcomplete: size > dimension!"
                 # we'd need to construct a different pseudo-inverse if the above assert fails
 
-                A = self.to_std_transform_matrix  # shape (dim,size) - should have indep *cols*
+                A = self.to_std_transform_matrix  # shape (dim, size) - should have indep *cols*
                 Adag = A.transpose().conjugate()  # shape (size, dim)
                 return _np.dot(_inv(_np.dot(Adag, A)), Adag)
 
@@ -616,7 +624,7 @@ def to_elementstd_transform_matrix(self):
         # *is* a standard representation of the vector space this basis or partial-basis
         # acts upon (this is *not* true for direct-sum bases, where the flattened
         # elements represent vectors in a larger "embedding" space (w/larger dim than actual space).
-        assert(self.is_simple()), "Incorrectly using a simple-assuming implementation of to_elementstd_transform_matrix"
+        assert (self.is_simple()), "Incorrectly using a simple-assuming implementation of to_elementstd_transform_matrix"
         return self.to_std_transform_matrix
 
     @property
@@ -661,7 +669,7 @@ def create_equivalent(self, builtin_basis_name):
         Basis
         """
         #This default implementation assumes that this basis is simple.
-        assert(self.is_simple()), "Incorrectly using a simple-assuming implementation of create_equivalent()"
+        assert (self.is_simple()), "Incorrectly using a simple-assuming implementation of create_equivalent()"
         return BuiltinBasis(builtin_basis_name, self.dim, sparse=self.sparse)
 
     #TODO: figure out if we actually need the return value from this function to
@@ -692,11 +700,11 @@ def create_simple_equivalent(self, builtin_basis_name=None):
         Basis
         """
         #This default implementation assumes that this basis is simple.
-        assert(self.is_simple()), "Incorrectly using a simple-assuming implementation of create_simple_equivalent()"
+        assert (self.is_simple()), "Incorrectly using a simple-assuming implementation of create_simple_equivalent()"
         if builtin_basis_name is None: return self.copy()
         else: return self.create_equivalent(builtin_basis_name)
 
-    def is_compatible_with_state_space(self, state_space):
+    def is_compatible_with_state_space(self, state_space: _StateSpace) -> bool:
         """
         Checks whether this basis is compatible with a given state space.
 
@@ -758,7 +766,7 @@ class LazyBasis(Basis):
         The basis labels
     """
 
-    def __init__(self, name, longname, real, sparse):
+    def __init__(self, name, longname, real, sparse: bool):
         """
         Creates a new LazyBasis.  Parameters are the same as those to
         :meth:`Basis.__init__`.
@@ -870,7 +878,9 @@ class ExplicitBasis(Basis):
     """
     Count = 0  # The number of custom bases, used for serialized naming
 
-    def __init__(self, elements, labels=None, name=None, longname=None, real=False, sparse=None, vector_elements=None):
+    def __init__(self, elements: _np.ndarray, labels: Optional[list] = None, name: Optional[str] = None,
+                 longname: Optional[str] = None, real: Optional[bool] = False, sparse: Optional[bool] = None,
+                 vector_elements: Optional[_np.ndarray] = None):
         '''
         Create a new ExplicitBasis.
 
@@ -935,13 +945,13 @@ def __init__(self, elements, labels=None, name=None, longname=None, real=False,
                         el = _np.array(el)  # try to convert to a numpy array
 
                 if elshape is None: elshape = el.shape
-                else: assert(elshape == el.shape), "Inconsistent element shapes!"
+                else: assert (elshape == el.shape), "Inconsistent element shapes!"
                 self.elements.append(el)
             dim = int(_np.prod(elshape))
         self.ellookup = {lbl: el for lbl, el in zip(self.labels, self.elements)}  # fast by-label element lookup
 
         if vector_elements is not None:
-            assert(len(vector_elements) == size), "Must have the same number of `elements` and `vector_elements`"
+            assert (len(vector_elements) == size), "Must have the same number of `elements` and `vector_elements`"
             if sparse:
                 self._vector_elements = [(el if _sps.issparse(el) else _sps.lil_matrix(el)) for el in vector_elements]
             else:
@@ -970,7 +980,7 @@ def _to_nice_serialization(self):
         return state
 
     @classmethod
-    def _from_nice_serialization(cls, state):
+    def _from_nice_serialization(cls, state) -> ExplicitBasis:
         vels = [cls._decodemx(vel) for vel in state['vector_elements']] \
             if (state.get('vector_elements', None) is not None) else None
         return cls([cls._decodemx(el) for el in state['elements']],
@@ -1000,7 +1010,7 @@ def elshape(self):
         """
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
         """
         # shape of "natural" elements - size may be > self.dim (to display naturally)
@@ -1021,7 +1031,7 @@ def vector_elements(self):
         else:
             return Basis.vector_elements.fget(self)  # call base class get-property fn
 
-    def _copy_with_toggled_sparsity(self):
+    def _copy_with_toggled_sparsity(self) -> ExplicitBasis:
         return ExplicitBasis(self.elements, self.labels, self.name, self.longname, self.real, not self.sparse,
                              self._vector_elements)
 
@@ -1060,15 +1070,15 @@ class BuiltinBasis(LazyBasis):
         or dense numpy arrays (the default).
     """
 
-    def __init__(self, name, dim_or_statespace, sparse=False):
+    def __init__(self, name, dim_or_statespace, sparse: Optional[bool] = False):
         from pygsti.baseobjs import statespace as _statespace
-        assert(name in _basis_constructor_dict), "Unknown builtin basis name '%s'!" % name
+        assert (name in _basis_constructor_dict), "Unknown builtin basis name '%s'!" % name
         if sparse is None: sparse = False  # choose dense matrices by default (when sparsity is "unspecified")
 
         if isinstance(dim_or_statespace, _statespace.StateSpace):
             self.state_space = dim_or_statespace
         elif name == 'cl':  # HACK for now, until we figure out better classical state spaces
-            self.state_space = _statespace.ExplicitStateSpace([('L%d' % i,) for i in range(dim_or_statespace)])
+            self.state_space = _statespace.ExplicitStateSpace([('L%d' % i, ) for i in range(dim_or_statespace)])
         elif name == "sv":
             # A state vector can have any shape. It does not need to be a perfect square root.
             self.state_space = _statespace.default_space_for_udim(dim_or_statespace)
@@ -1084,7 +1094,7 @@ def __init__(self, name, dim_or_statespace, sparse=False):
         #precompute some properties
         self._size, self._dim, self._elshape = _basis_constructor_dict[self.name].sizes(dim=self._get_dimension_to_pass_to_constructor(), sparse=self.sparse)
         #Check that sparse is True only when elements are *matrices*
-        assert(not self.sparse or len(self._elshape) == 2), "`sparse == True` is only allowed for *matrix*-valued bases!"
+        assert (not self.sparse or len(self._elshape) == 2), "`sparse == True` is only allowed for *matrix*-valued bases!"
 
     def _to_nice_serialization(self):
         state = super()._to_nice_serialization()
@@ -1095,18 +1105,17 @@ def _to_nice_serialization(self):
         return state
 
     @classmethod
-    def _from_nice_serialization(cls, state):
+    def _from_nice_serialization(cls, state) -> BuiltinBasis:
         statespace = _StateSpace.from_nice_serialization(state['state_space'])
         return cls(state['name'], statespace, state['sparse'])
-    
-    def _get_dimension_to_pass_to_constructor(self):
+
+    def _get_dimension_to_pass_to_constructor(self) -> int:
         """
         A basis in the state-vector (name= 'sv') case will correspond to a basis of vectors of length d.
         This means that it will be operated on by matrices of shape (d \times d).
         """
         return self.state_space.udim if self.name == "sv" else self.state_space.dim
 
-
     @property
     def dim(self):
         """
@@ -1128,7 +1137,7 @@ def elshape(self):
         """
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
         """
         return self._elshape
@@ -1148,18 +1157,18 @@ def _lazy_build_elements(self):
 
         cargs = {'dim': self._get_dimension_to_pass_to_constructor(), 'sparse': self.sparse}
         self._elements = _np.array(f(**cargs))  # a list of (dense) mxs or vectors -> ndarray (possibly sparse in future?)
-        assert(len(self._elements) == self.size), "Logic error: wrong number of elements were created!"
+        assert (len(self._elements) == self.size), "Logic error: wrong number of elements were created!"
 
     def _lazy_build_labels(self):
         f = _basis_constructor_dict[self.name].labeler
         cargs = {'dim': self._get_dimension_to_pass_to_constructor(), 'sparse': self.sparse}
         self._labels = f(**cargs)
-        assert(len(self._labels) == self.size)
+        assert (len(self._labels) == self.size)
 
-    def _copy_with_toggled_sparsity(self):
+    def _copy_with_toggled_sparsity(self) -> BuiltinBasis:
         return BuiltinBasis(self.name, self.state_space, not self.sparse)
 
-    def is_equivalent(self, other, sparseness_must_match=True):
+    def is_equivalent(self, other, sparseness_must_match: bool = True) -> bool:
         """
         Tests whether this basis is equal to another basis, optionally ignoring sparseness.
 
@@ -1202,11 +1211,11 @@ class DirectSumBasis(LazyBasis):
     component_bases : iterable
         A list of the component bases.  Each list elements may be either
         a Basis object or a tuple of arguments to :func:`Basis.cast`,
-        e.g. `('pp',4)`.
+        e.g. `('pp', 4)`.
 
     name : str, optional
         The name of this basis.  If `None`, the names of the component bases
-        joined with "+" is used.
+        joined with " + " is used.
 
     longname : str, optional
         A longer description of this basis.  If `None`, then a long name is
@@ -1218,7 +1227,7 @@ class DirectSumBasis(LazyBasis):
         The "vectors" of this basis, always 1D (sparse or dense) arrays.
     """
 
-    def __init__(self, component_bases, name=None, longname=None):
+    def __init__(self, component_bases, name: Optional[str] = None, longname: Optional[str] = None):
         '''
         Create a new DirectSumBasis - a basis for a space that is the direct-sum
         of the spaces spanned by other "component" bases.
@@ -1228,17 +1237,17 @@ def __init__(self, component_bases, name=None, longname=None):
         component_bases : iterable
             A list of the component bases.  Each list elements may be either
             a Basis object or a tuple of arguments to :func:`Basis.cast`,
-            e.g. `('pp',4)`.
+            e.g. `('pp', 4)`.
 
         name : str, optional
             The name of this basis.  If `None`, the names of the component bases
-            joined with "+" is used.
+            joined with " + " is used.
 
         longname : str, optional
             A longer description of this basis.  If `None`, then a long name is
             automatically generated.
         '''
-        assert(len(component_bases) > 0), "Must supply at least one component basis"
+        assert (len(component_bases) > 0), "Must supply at least one component basis"
 
         self._component_bases = []
         self._vector_elements = None  # vectorized elements: 1D arrays
@@ -1247,19 +1256,19 @@ def __init__(self, component_bases, name=None, longname=None):
             if isinstance(compbasis, Basis):
                 self._component_bases.append(compbasis)
             else:
-                #compbasis can be a list/tuple of args to Basis.cast, e.g. ('pp',2)
+                #compbasis can be a list/tuple of args to Basis.cast, e.g. ('pp', 2)
                 self._component_bases.append(Basis.cast(*compbasis))
 
         if name is None:
-            name = "+".join([c.name for c in self._component_bases])
+            name = " + ".join([c.name for c in self._component_bases])
         if longname is None:
             longname = "Direct-sum basis with components " + ", ".join(
                 [c.name for c in self._component_bases])
 
         real = all([c.real for c in self._component_bases])
         sparse = all([c.sparse for c in self._component_bases])
-        assert(all([c.real == real for c in self._component_bases])), "Inconsistent `real` value among component bases!"
-        assert(all([c.sparse == sparse for c in self._component_bases])), "Inconsistent sparsity among component bases!"
+        assert (all([c.real == real for c in self._component_bases])), "Inconsistent `real` value among component bases!"
+        assert (all([c.sparse == sparse for c in self._component_bases])), "Inconsistent sparsity among component bases!"
 
         #precompute various basis properties. can add more as they are deemed frequently accessed.
         self._dim = sum([c.dim for c in self._component_bases])
@@ -1276,7 +1285,7 @@ def _to_nice_serialization(self):
         return state
 
     @classmethod
-    def _from_nice_serialization(cls, state):
+    def _from_nice_serialization(cls, state) -> DirectSumBasis:
         component_bases = [Basis.from_nice_serialization(b) for b in state['component_bases']]
         return cls(component_bases, state['name'], state['longname'])
 
@@ -1306,16 +1315,16 @@ def elshape(self):
         """
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
         """
         elndim = len(self._component_bases[0].elshape)
-        assert(all([len(c.elshape) == elndim for c in self._component_bases])
-               ), "Inconsistent element ndims among component bases!"
+        assert (all([len(c.elshape) == elndim for c in self._component_bases])), \
+            "Inconsistent element ndims among component bases!"
         return tuple([sum([c.elshape[k] for c in self._component_bases]) for k in range(elndim)])
 
     def __hash__(self):
-        return hash((self.name,)+tuple((hash(comp) for comp in self._component_bases)))
+        return hash((self.name, ) + tuple((hash(comp) for comp in self._component_bases)))
 
     def _lazy_build_vector_elements(self):
         if self.sparse:
@@ -1326,7 +1335,7 @@ def _lazy_build_vector_elements(self):
         i, start = 0, 0
         for compbasis in self._component_bases:
             for lbl, vel in zip(compbasis.labels, compbasis.vector_elements):
-                assert(_sps.issparse(vel) == self.sparse), "Inconsistent sparsity!"
+                assert (_sps.issparse(vel) == self.sparse), "Inconsistent sparsity!"
                 if self.sparse:
                     mx = _sps.lil_matrix((self.dim, 1), dtype='complex')
                     mx[start:start + compbasis.dim, 0] = vel
@@ -1336,7 +1345,7 @@ def _lazy_build_vector_elements(self):
                 i += 1
             start += compbasis.dim
 
-        assert(i == self.size)
+        assert (i == self.size)
         self._vector_elements = compMxs
 
     def _lazy_build_elements(self):
@@ -1457,7 +1466,7 @@ def to_elementstd_transform_matrix(self):
             elements are 4x4 matrices) and `size` is the size of this basis (its
             number of vectors).
         """
-        assert(not self.sparse), "to_elementstd_transform_matrix not implemented for sparse mode"
+        assert (not self.sparse), "to_elementstd_transform_matrix not implemented for sparse mode"
         expanddim = self.elsize  # == _np.prod(self.elshape)
         if self.sparse:
             toSimpleStd = _sps.lil_matrix((expanddim, self.size), dtype='complex')
@@ -1472,7 +1481,7 @@ def to_elementstd_transform_matrix(self):
             toSimpleStd[:, i] = vel
         return toSimpleStd
 
-    def create_equivalent(self, builtin_basis_name):
+    def create_equivalent(self, builtin_basis_name) -> DirectSumBasis:
         """
         Create an equivalent basis with components of type `builtin_basis_name`.
 
@@ -1493,7 +1502,7 @@ def create_equivalent(self, builtin_basis_name):
         equiv_components = [c.create_equivalent(builtin_basis_name) for c in self._component_bases]
         return DirectSumBasis(equiv_components)
 
-    def create_simple_equivalent(self, builtin_basis_name=None):
+    def create_simple_equivalent(self, builtin_basis_name=None) -> BuiltinBasis:
         """
         Create a basis of type `builtin_basis_name` whose elements are compatible with this basis.
 
@@ -1542,7 +1551,7 @@ class TensorProdBasis(LazyBasis):
     component_bases : iterable
         A list of the component bases.  Each list elements may be either
         a Basis object or a tuple of arguments to :func:`Basis.cast`,
-        e.g. `('pp',4)`.
+        e.g. `('pp', 4)`.
 
     name : str, optional
         The name of this basis.  If `None`, the names of the component bases
@@ -1553,7 +1562,7 @@ class TensorProdBasis(LazyBasis):
         automatically generated.
     """
 
-    def __init__(self, component_bases, name=None, longname=None):
+    def __init__(self, component_bases, name: Optional[str] = None, longname: Optional[str] = None):
         '''
         Create a new TensorProdBasis whose elements are the tensor products
         of the elements of a set of "component" bases.
@@ -1563,7 +1572,7 @@ def __init__(self, component_bases, name=None, longname=None):
         component_bases : iterable
             A list of the component bases.  Each list elements may be either
             a Basis object or a tuple of arguments to :func:`Basis.cast`,
-            e.g. `('pp',4)`.
+            e.g. `('pp', 4)`.
 
         name : str, optional
             The name of this basis.  If `None`, the names of the component bases
@@ -1573,14 +1582,14 @@ def __init__(self, component_bases, name=None, longname=None):
             A longer description of this basis.  If `None`, then a long name is
             automatically generated.
         '''
-        assert(len(component_bases) > 0), "Must supply at least one component basis"
+        assert (len(component_bases) > 0), "Must supply at least one component basis"
 
         self._component_bases = []
         for compbasis in component_bases:
             if isinstance(compbasis, Basis):
                 self._component_bases.append(compbasis)
             else:
-                #compbasis can be a list/tuple of args to Basis.cast, e.g. ('pp',2)
+                #compbasis can be a list/tuple of args to Basis.cast, e.g. ('pp', 2)
                 self._component_bases.append(Basis.cast(*compbasis))
 
         if name is None:
@@ -1591,8 +1600,10 @@ def __init__(self, component_bases, name=None, longname=None):
 
         real = all([c.real for c in self._component_bases])
         sparse = all([c.sparse for c in self._component_bases])
-        #assert(all([c.real == real for c in self._component_bases])), "Inconsistent `real` value among component bases!"
-        assert(all([c.sparse == sparse for c in self._component_bases])), "Inconsistent sparsity among component bases!"
+        #assert (all([c.real == real for c
+        #            in self._component_bases])), "Inconsistent `real` value among component bases!"
+        assert (all([c.sparse == sparse for c
+                     in self._component_bases])), "Inconsistent sparsity among component bases!"
 
         #precompute certain properties. Can add more as deemed frequently accessed.
         self._dim = int(_np.prod([c.dim for c in self._component_bases]))
@@ -1601,7 +1612,7 @@ def __init__(self, component_bases, name=None, longname=None):
         # basis are in their "natrual" representation where it makes sense to take tensor products.  For
         # example, a direct-sum basis may hold elements in a compact way that violate this... but I'm not sure if they
         # do and this needs to be checked.  For now, we could just disable this overly-restrictive assert:
-        assert(all([c.is_simple() for c in self._component_bases])), \
+        assert (all([c.is_simple() for c in self._component_bases])), \
             "Components of a tensor product basis must be *simple* (have vector-dimension == size of elements)"
         # because we use the natural representation to take tensor (kronecker) products.
         # Note: this assertion also means dim == product(component_elsizes) == elsize, so basis is *simple*
@@ -1617,7 +1628,7 @@ def _to_nice_serialization(self):
         return state
 
     @classmethod
-    def _from_nice_serialization(cls, state):
+    def _from_nice_serialization(cls, state) -> TensorProdBasis:
         component_bases = [Basis.from_nice_serialization(b) for b in state['component_bases']]
         return cls(component_bases, state['name'], state['longname'])
 
@@ -1647,7 +1658,7 @@ def elshape(self):
         """
         The shape of each element.  Typically either a length-1 or length-2
         tuple, corresponding to vector or matrix elements, respectively.
-        Note that *vector elements* always have shape `(dim,)` (or `(dim,1)`
+        Note that *vector elements* always have shape `(dim, )` (or `(dim, 1)`
         in the sparse case).
         """
         elndim = max([c.elndim for c in self._component_bases])
@@ -1659,14 +1670,14 @@ def elshape(self):
         return tuple(elshape)
 
     def __hash__(self):
-        return hash((self.name,) + tuple((hash(comp) for comp in self._component_bases)))
+        return hash((self.name, ) + tuple((hash(comp) for comp in self._component_bases)))
 
     def _lazy_build_elements(self):
         #LAZY building of elements (in case we never need them)
         if self.sparse:
             compMxs = [None] * self.size
         else:
-            compMxs = _np.zeros((self.size,) + self.elshape, 'complex')
+            compMxs = _np.zeros((self.size, ) + self.elshape, 'complex')
 
         #Take kronecker product of *natural* reps of component-basis elements
         # then reshape to vectors at the end.  This requires that the vector-
@@ -1691,7 +1702,7 @@ def _lazy_build_labels(self):
         for i, factor_lbls in enumerate(_itertools.product(*comp_lbls)):
             self._labels.append(''.join(factor_lbls))
 
-    def _copy_with_toggled_sparsity(self):
+    def _copy_with_toggled_sparsity(self) -> TensorProdBasis:
         return TensorProdBasis([cb._copy_with_toggled_sparsity() for cb in self._component_bases],
                                self.name, self.longname)
 
diff --git a/pygsti/baseobjs/basisconstructors.py b/pygsti/baseobjs/basisconstructors.py
index 7edff58c4..5edd81f47 100644
--- a/pygsti/baseobjs/basisconstructors.py
+++ b/pygsti/baseobjs/basisconstructors.py
@@ -94,7 +94,7 @@ class MatrixBasisConstructor(object):
         real components.
     """
 
-    def __init__(self, longname, matrixgen_fn, labelgen_fn, real, first_element_is_identity):
+    def __init__(self, longname: str, matrixgen_fn, labelgen_fn, real: bool, first_element_is_identity):
         """
         Create a new MatrixBasisConstructor:
 
@@ -123,7 +123,7 @@ def __init__(self, longname, matrixgen_fn, labelgen_fn, real, first_element_is_i
         self.real = real
         self.first_element_is_identity = first_element_is_identity
 
-    def matrix_dim(self, dim):
+    def matrix_dim(self, dim: int) -> int:
         """
         Helper function that converts a *vector-space* dimension `dim` to matrix-dimension by taking a sqrt.
 
@@ -140,7 +140,7 @@ def matrix_dim(self, dim):
         assert(d**2 == dim), "Matrix bases can only have dimension = perfect square (not %d)!" % dim
         return d
 
-    def labeler(self, dim, sparse):
+    def labeler(self, dim, sparse: bool) -> list[str]:
         """
         Get the labels of a basis to be constructed.
 
@@ -158,7 +158,7 @@ def labeler(self, dim, sparse):
         """
         return self.labelgen_fn(self.matrix_dim(dim))
 
-    def constructor(self, dim, sparse):
+    def constructor(self, dim, sparse: bool):
         """
         Get the elements of a basis to be constructed.
 
@@ -178,7 +178,7 @@ def constructor(self, dim, sparse):
         if sparse: els = [_sps.csr_matrix(el) for el in els]
         return els
 
-    def sizes(self, dim, sparse):
+    def sizes(self, dim, sparse: bool) -> tuple[int, int, tuple[int, int]]:
         """
         Get some relevant sizes/dimensions for constructing a basis.
 
@@ -222,7 +222,7 @@ class DiagonalMatrixBasisConstructor(MatrixBasisConstructor):
     a VectorBasisConstructor, but element are diagonal matrices rather than vectors)
     """
 
-    def constructor(self, dim, sparse):
+    def constructor(self, dim, sparse: bool):
         """
         Get the elements of a basis to be constructed.
 
@@ -245,7 +245,7 @@ def constructor(self, dim, sparse):
         if sparse: els = [_sps.csr_matrix(el) for el in els]
         return els
 
-    def sizes(self, dim, sparse):
+    def sizes(self, dim, sparse: bool):
         """
         Get some relevant sizes/dimensions for constructing a basis.
 
@@ -285,7 +285,7 @@ class SingleElementMatrixBasisConstructor(MatrixBasisConstructor):
     A constructor for a basis containing just a single element (e.g. the identity).
     """
 
-    def sizes(self, dim, sparse):
+    def sizes(self, dim, sparse: bool):
         """
         See docstring for :class:`MatrixBasisConstructor`
 
@@ -365,7 +365,7 @@ def __init__(self, longname, vectorgen_fn, labelgen_fn, real):
         self.real = real
         self.first_element_is_identity = False  # only applies to matrix bases
 
-    def labeler(self, dim, sparse):
+    def labeler(self, dim, sparse: bool):
         """
         Get the labels of a basis to be constructed.
 
@@ -383,7 +383,7 @@ def labeler(self, dim, sparse):
         """
         return self.labelgen_fn(dim)
 
-    def constructor(self, dim, sparse):
+    def constructor(self, dim, sparse: bool):
         """
         Get the elements of a basis to be constructed.
 
@@ -403,7 +403,7 @@ def constructor(self, dim, sparse):
         assert(not sparse), "Sparse vector bases not supported (yet)"
         return els
 
-    def sizes(self, dim, sparse):
+    def sizes(self, dim, sparse: bool):
         """
         Get some relevant sizes/dimensions for constructing a basis.
 
diff --git a/pygsti/tools/hypothesis.py b/pygsti/tools/hypothesis.py
index 62643dd99..2660f99ef 100644
--- a/pygsti/tools/hypothesis.py
+++ b/pygsti/tools/hypothesis.py
@@ -10,10 +10,19 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 
+from typing import (
+    Optional,
+    Sequence,
+    Literal,
+    TypeAlias,
+)
 import numpy as _np
 
 
-def bonferroni_correction(significance, numtests):
+CORRECTION_METHOD_T: TypeAlias = Literal["bonferroni", "sidak"]
+
+
+def bonferroni_correction(significance: float, numtests: int) -> float:
     """
     Calculates the standard Bonferroni correction.
 
@@ -38,11 +47,14 @@ def bonferroni_correction(significance, numtests):
     return local_significance
 
 
-def sidak_correction(significance, numtests):
+def sidak_correction(significance: float, numtests: int) -> float:
     """
     Sidak correction.
 
-    TODO: docstring - better explanaition
+    Convert a significance level for an individual test repeated independently, numtest times
+    into a significance level for the whole family of tests.
+
+    https://en.wikipedia.org/wiki/sidak_correction
 
     Parameters
     ----------
@@ -60,8 +72,8 @@ def sidak_correction(significance, numtests):
     return adjusted_significance
 
 
-def generalized_bonferroni_correction(significance, weights, numtests=None,
-                                      nested_method='bonferroni', tol=1e-10):
+def generalized_bonferroni_correction(significance: float, weights: _np.ndarray, numtests: Optional[int] = None,
+                                      nested_method: CORRECTION_METHOD_T = 'bonferroni', tol: float = 1e-10) -> float:
     """
     Generalized Bonferroni correction.
 
@@ -88,14 +100,15 @@ def generalized_bonferroni_correction(significance, weights, numtests=None,
     float
     """
     weights = _np.array(weights)
-    assert(_np.abs(_np.sum(weights) - 1.) < tol), "Invalid weighting! The weights must add up to 1."
+    assert (_np.allclose(_np.abs(weights), weights)), "Invalid weighting! The weights must be non-negative."
+    assert (_np.abs(_np.sum(weights) - 1.) < tol), "Invalid weighting! The weights must add up to 1."
 
     adjusted_significance = _np.zeros(len(weights), float)
     adjusted_significance = significance * weights
 
     if numtests is not None:
 
-        assert(len(numtests) == len(weights)), "The number of tests must be specified for each weight!"
+        assert (len(numtests) == len(weights)), "The number of tests must be specified for each weight!"
         for i in range(0, len(weights)):
 
             if nested_method == 'bonferroni':

From f724a194ea67e207a191f8b9ebcdf078fdfc9e58 Mon Sep 17 00:00:00 2001
From: Nick Koskelo <koskelo2@illinois.edu>
Date: Fri, 1 Aug 2025 10:26:34 -0700
Subject: [PATCH 2/2] TypeAlias is from python 3.10 but we need to also support
 python 3.9.

---
 pygsti/baseobjs/basis.py   | 8 ++------
 pygsti/tools/hypothesis.py | 8 ++------
 2 files changed, 4 insertions(+), 12 deletions(-)

diff --git a/pygsti/baseobjs/basis.py b/pygsti/baseobjs/basis.py
index 9d4061790..d84f203eb 100644
--- a/pygsti/baseobjs/basis.py
+++ b/pygsti/baseobjs/basis.py
@@ -18,11 +18,7 @@
 from functools import lru_cache
 from typing import (
     Union,
-    TypeAlias,
-    Literal,
     Optional,
-    Sequence,
-    Iterable,
 )
 
 import numpy as _np
@@ -218,7 +214,7 @@ def cast_from_basis(cls, basis: Basis, dim=None, sparse: Optional[bool] = None)
                 if hasattr(basis, 'state_space'):  # TODO - should *all* basis objects have a state_space?
                     assert (state_space.is_compatible_with(basis.state_space)), \
                         "Basis object has incompatible state space: %s != %s" % (str(state_space),
-                                                                                    str(basis.state_space))
+                                                                                str(basis.state_space))
             else:  # assume dim is an integer
                 assert (dim == basis.dim or dim == basis.elsize), \
                     "Basis object has unexpected dimension: %d != %d or %d" % (dim, basis.dim, basis.elsize)
@@ -624,7 +620,7 @@ def to_elementstd_transform_matrix(self):
         # *is* a standard representation of the vector space this basis or partial-basis
         # acts upon (this is *not* true for direct-sum bases, where the flattened
         # elements represent vectors in a larger "embedding" space (w/larger dim than actual space).
-        assert (self.is_simple()), "Incorrectly using a simple-assuming implementation of to_elementstd_transform_matrix"
+        assert self.is_simple(), "Incorrectly using a simple-assuming implementation of to_elementstd_transform_matrix"
         return self.to_std_transform_matrix
 
     @property
diff --git a/pygsti/tools/hypothesis.py b/pygsti/tools/hypothesis.py
index 2660f99ef..8e1c17170 100644
--- a/pygsti/tools/hypothesis.py
+++ b/pygsti/tools/hypothesis.py
@@ -12,16 +12,11 @@
 
 from typing import (
     Optional,
-    Sequence,
     Literal,
-    TypeAlias,
 )
 import numpy as _np
 
 
-CORRECTION_METHOD_T: TypeAlias = Literal["bonferroni", "sidak"]
-
-
 def bonferroni_correction(significance: float, numtests: int) -> float:
     """
     Calculates the standard Bonferroni correction.
@@ -73,7 +68,8 @@ def sidak_correction(significance: float, numtests: int) -> float:
 
 
 def generalized_bonferroni_correction(significance: float, weights: _np.ndarray, numtests: Optional[int] = None,
-                                      nested_method: CORRECTION_METHOD_T = 'bonferroni', tol: float = 1e-10) -> float:
+                                      nested_method: Literal["bonferroni", "sidak"] = 'bonferroni',
+                                      tol: float = 1e-10) -> float:
     """
     Generalized Bonferroni correction.