Merge pull request #605: Add type annotation for on_space parameter

Add type annotation for on_space parameter. Also add an optional enum, but keep use of string literals throughout current code.

Author: rileyjmurray

pygsti/__init__.py

@@ -7,6 +7,9 @@
 # http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
 #***************************************************************************************************
 """ A Python implementation of LinearOperator Set Tomography """
+from .pgtypes import (
+    SpaceT
+)
 
 from . import baseobjs
 from . import algorithms as alg
@@ -19,6 +22,7 @@
 from . import protocols
 from . import report as rpt
 from . import serialization
+from . import enums
 
 # Import the most important/useful routines of each module/sub-package
 # into the package namespace
@@ -30,6 +34,7 @@
 from pygsti.tools.gatetools import *  # *_qubit_gate fns
 from .drivers import *
 from .tools import *
+
 # NUMPY BUG FIX (imported from tools)
 from pygsti.baseobjs._compatibility import _numpy14einsumfix
 

pygsti/algorithms/gaugeopt.py

@@ -517,8 +517,8 @@ def _jacobian_fn(gauge_group_el):
                 # d(op_term) = S_inv * (-dS * S_inv * G * S + G * dS) = S_inv * (-dS * G' + G * dS)
                 #   Note: (S_inv * G * S) is G' (transformed G)
                 wt = item_weights.get(lbl, opWeight)
-                left = -1 * _np.dot(dS, mdl_post.operations[lbl].to_dense(on_space='minimal'))  # shape (n,d1,d2)
-                right = _np.swapaxes(_np.dot(G.to_dense(on_space='minimal'), dS), 0, 1)  # shape (d1,n,d2) -> (n,d1,d2)
+                left = -1 * _np.dot(dS, mdl_post.operations[lbl].to_dense('minimal'))  # shape (n,d1,d2)
+                right = _np.swapaxes(_np.dot(G.to_dense('minimal'), dS), 0, 1)  # shape (d1,n,d2) -> (n,d1,d2)
                 result = _np.swapaxes(_np.dot(S_inv, left + right), 1, 2)  # shape (d1, d2, n)
                 result = result.reshape((d**2, n))  # must copy b/c non-contiguous
                 my_jacMx[start:start + d**2] = wt * result
@@ -530,8 +530,8 @@ def _jacobian_fn(gauge_group_el):
                 wt = item_weights.get(ilbl, opWeight)
                 for lbl, G in Inst.items():
                     # same calculation as for operation terms
-                    left = -1 * _np.dot(dS, mdl_post.instruments[ilbl][lbl].to_dense(on_space='minimal'))  # (n,d1,d2)
-                    right = _np.swapaxes(_np.dot(G.to_dense(on_space='minimal'), dS), 0, 1)  # (d1,n,d2) -> (n,d1,d2)
+                    left = -1 * _np.dot(dS, mdl_post.instruments[ilbl][lbl].to_dense('minimal'))  # (n,d1,d2)
+                    right = _np.swapaxes(_np.dot(G.to_dense('minimal'), dS), 0, 1)  # (d1,n,d2) -> (n,d1,d2)
                     result = _np.swapaxes(_np.dot(S_inv, left + right), 1, 2)  # shape (d1, d2, n)
                     result = result.reshape((d**2, n))  # must copy b/c non-contiguous
                     my_jacMx[start:start + d**2] = wt * result
@@ -544,7 +544,7 @@ def _jacobian_fn(gauge_group_el):
                 #   Note: (S_inv * rho) is transformed rho
                 wt = item_weights.get(lbl, spamWeight)
                 Sinv_dS = _np.dot(S_inv, dS)  # shape (d1,n,d2)
-                result = -1 * _np.dot(Sinv_dS, rho.to_dense(on_space='minimal'))  # shape (d,n)
+                result = -1 * _np.dot(Sinv_dS, rho.to_dense('minimal'))  # shape (d,n)
                 my_jacMx[start:start + d] = wt * result
                 start += d
 
@@ -554,7 +554,7 @@ def _jacobian_fn(gauge_group_el):
                 for lbl, E in povm.items():
                     # d(ET_term) = E.T * dS
                     wt = item_weights.get(povmlbl + "_" + lbl, spamWeight)
-                    result = _np.dot(E.to_dense(on_space='minimal')[None, :], dS).T  # shape (1,n,d2).T => (d2,n,1)
+                    result = _np.dot(E.to_dense('minimal')[None, :], dS).T  # shape (1,n,d2).T => (d2,n,1)
                     my_jacMx[start:start + d] = wt * result.squeeze(2)  # (d2,n)
                     start += d
 
@@ -851,7 +851,7 @@ def _spam_penalty_jac_fill(spam_penalty_vec_grad_to_fill, mdl_pre, mdl_post,
 
         #get sgn(denMx) == d(|denMx|_Tr)/d(denMx) in std basis
         # dmDim = denMx.shape[0]
-        denMx = _tools.vec_to_stdmx(prepvec.to_dense(on_space='minimal')[:, None], op_basis)
+        denMx = _tools.vec_to_stdmx(prepvec.to_dense('minimal')[:, None], op_basis)
         assert(_np.linalg.norm(denMx - denMx.T.conjugate()) < 1e-4), \
             "denMx should be Hermitian!"
 
@@ -865,7 +865,7 @@ def _spam_penalty_jac_fill(spam_penalty_vec_grad_to_fill, mdl_pre, mdl_post,
         # get d(prepvec')/dp = d(S_inv * prepvec)/dp in op_basis [shape == (n,dim)]
         #                    = (-S_inv*dS*S_inv) * prepvec = -S_inv*dS * prepvec'
         Sinv_dS = _np.dot(S_inv, dS)  # shape (d1,n,d2)
-        dVdp = -1 * _np.dot(Sinv_dS, prepvec.to_dense(on_space='minimal')[:, None]).squeeze(2)  # shape (d,n,1) => (d,n)
+        dVdp = -1 * _np.dot(Sinv_dS, prepvec.to_dense('minimal')[:, None]).squeeze(2)  # shape (d,n,1) => (d,n)
         assert(dVdp.shape == (d, n))
 
         # denMx = sum( spamvec[i] * Bmx[i] )
@@ -890,7 +890,7 @@ def _spam_penalty_jac_fill(spam_penalty_vec_grad_to_fill, mdl_pre, mdl_post,
         for lbl, effectvec in povm.items():
 
             #get sgn(EMx) == d(|EMx|_Tr)/d(EMx) in std basis
-            EMx = _tools.vec_to_stdmx(effectvec.to_dense(on_space='minimal')[:, None], op_basis)
+            EMx = _tools.vec_to_stdmx(effectvec.to_dense('minimal')[:, None], op_basis)
             # dmDim = EMx.shape[0]
             assert(_np.linalg.norm(EMx - EMx.T.conjugate()) < 1e-4), \
                 "denMx should be Hermitian!"
@@ -905,7 +905,7 @@ def _spam_penalty_jac_fill(spam_penalty_vec_grad_to_fill, mdl_pre, mdl_post,
             # get d(effectvec')/dp = [d(effectvec.T * S)/dp].T in op_basis [shape == (n,dim)]
             #                      = [effectvec.T * dS].T
             #  OR = dS.T * effectvec
-            pre_effectvec = mdl_pre.povms[povmlbl][lbl].to_dense(on_space='minimal')[:, None]
+            pre_effectvec = mdl_pre.povms[povmlbl][lbl].to_dense('minimal')[:, None]
             dVdp = _np.dot(pre_effectvec.T, dS).squeeze(0).T
             # shape = (1,d) * (n, d1,d2) = (1,n,d2) => (n,d2) => (d2,n)
             assert(dVdp.shape == (d, n))

pygsti/circuits/cloudcircuitconstruction.py

@@ -27,11 +27,8 @@
 from pygsti.circuits.circuitstructure import GermFiducialPairPlaquette as _GermFiducialPairPlaquette, \
     PlaquetteGridCircuitStructure as _PlaquetteGridCircuitStructure
 
-from pygsti.tools import basistools as _bt
-from pygsti.tools import internalgates as _itgs
 from pygsti.tools import listtools as _lt
 from pygsti.tools import mpitools as _mpit
-from pygsti.tools import optools as _ot
 from pygsti.tools import slicetools as _slct
 from pygsti.tools.legacytools import deprecate as _deprecated_fn
 

pygsti/enums/__init__.py

@@ -0,0 +1,6 @@
+"""
+These enums are intended to make it simplier to call functions with fixed strings expected as input.
+"""
+
+
+from .convertspaceenum import SpaceConversionType

pygsti/enums/convertspaceenum.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class SpaceConversionType(Enum):
+    Minimal = 'minimal'
+    Hilbert = 'Hilbert'
+    HilbertSchmidt = 'HilbertSchmidt'

pygsti/evotypes/chp/opreps.py

@@ -17,6 +17,7 @@
 from .. import basereps as _basereps
 from pygsti.baseobjs.statespace import StateSpace as _StateSpace
 from ...tools import internalgates as _itgs
+from pygsti import SpaceT
 
 
 class OpRep(_basereps.OpRep):
@@ -49,7 +50,7 @@ def adjoint_acton_random(self, state, rand_state):
     def _chp_ops(self, seed_or_state=None):
         return self.base_chp_ops
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         try:
             str_ops = str(self._chp_ops())
         except Exception:

pygsti/evotypes/densitymx_slow/effectreps.py

@@ -15,7 +15,7 @@
 # import functools as _functools
 from pygsti.baseobjs.statespace import StateSpace as _StateSpace
 from ...tools import matrixtools as _mt
-
+from pygsti import SpaceT
 
 class EffectRep:
     """Any representation of an "effect" in the sense of a POVM."""
@@ -44,7 +44,7 @@ def probability(self, state):
         # can assume state is a StateRep and self.state_rep is
         return _np.dot(self.state_rep.data, state.data)  # not vdot b/c *real* data
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         return self.state_rep.to_dense(on_space)
 
 
@@ -79,7 +79,7 @@ def probability(self, state):
         Edense = self.to_dense('HilbertSchmidt', scratch)
         return _np.dot(Edense, state.data)  # not vdot b/c data is *real*
 
-    def to_dense(self, on_space, outvec=None):
+    def to_dense(self, on_space: SpaceT, outvec=None):
         if on_space not in ('minimal', 'HilbertSchmidt'):
             raise ValueError("'densitymx' evotype cannot produce Hilbert-space ops!")
         return _mt.zvals_int64_to_dense(self.zvals_int, self.nfactors, outvec, False, self.abs_elval)
@@ -105,7 +105,7 @@ def __init__(self, povm_factors, effect_labels, state_space):
         super(EffectRepTensorProduct, self).__init__(state_space)
         self.factor_effects_have_changed()
 
-    def to_dense(self, on_space, outvec=None):
+    def to_dense(self, on_space: SpaceT, outvec=None):
 
         if on_space not in ('minimal', 'HilbertSchmidt'):
             raise ValueError("'densitymx' evotype cannot produce Hilbert-space ops!")

pygsti/evotypes/densitymx_slow/opreps.py

@@ -23,7 +23,7 @@
 from ...tools import lindbladtools as _lbt
 from ...tools import matrixtools as _mt
 from ...tools import optools as _ot
-
+from pygsti import SpaceT
 
 class OpRep:
     """
@@ -51,12 +51,12 @@ def aslinearoperator(self):
         def mv(v):
             if v.ndim == 2 and v.shape[1] == 1: v = v[:, 0]
             in_state = _StateRepDense(_np.ascontiguousarray(v, 'd'), self.state_space, None)
-            return self.acton(in_state).to_dense(on_space='HilbertSchmidt')
+            return self.acton(in_state).to_dense(on_space="HilbertSchmidt")
 
         def rmv(v):
             if v.ndim == 2 and v.shape[1] == 1: v = v[:, 0]
             in_state = _StateRepDense(_np.ascontiguousarray(v, 'd'), self.state_space, None)
-            return self.adjoint_acton(in_state).to_dense(on_space='HilbertSchmidt')
+            return self.adjoint_acton(in_state).to_dense(on_space="HilbertSchmidt")
         return LinearOperator((self.dim, self.dim), matvec=mv, rmatvec=rmv)  # transpose, adjoint, dot, matmat?
 
 
@@ -80,7 +80,7 @@ def __init__(self, mx, basis, state_space):
     def base_has_changed(self):
         pass
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         if on_space not in ('minimal', 'HilbertSchmidt'):
             raise ValueError("'densitymx_slow' evotype cannot produce Hilbert-space ops!")
         return self.base
@@ -115,7 +115,7 @@ def __init__(self, mx, basis, state_space):
     def base_has_changed(self):
         self.superop_base[:, :] = _ot.unitary_to_superop(self.base, self.basis)
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         if on_space in ('minimal', 'HilbertSchmidt'):
             return self.to_dense_superop()
         else:  # 'Hilbert'
@@ -167,7 +167,7 @@ def adjoint_acton(self, state):
         Aadj = self.A.conjugate(copy=True).transpose()
         return _StateRepDense(Aadj.dot(state.data), state.state_space, None)
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         if on_space not in ('minimal', 'HilbertSchmidt'):
             raise ValueError("'densitymx_slow' evotype cannot produce Hilbert-space ops!")
         return self.A.toarray()
@@ -212,7 +212,7 @@ def __str__(self):
     def copy(self):
         return OpRepKraus(self.basis, list(self.kraus_reps), self.state_space)
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         assert(on_space in ('minimal', 'HilbertSchmidt')), \
             'Can only compute OpRepKraus.to_dense on HilbertSchmidt space!'
         return sum([rep.to_dense(on_space) for rep in self.kraus_reps])
@@ -245,7 +245,7 @@ def copy(self):
     def update_unitary_rates(self, rates):
         self.unitary_rates[:] = rates
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         assert(on_space in ('minimal', 'HilbertSchmidt'))  # below code only works in this case
         return sum([rate * rep.to_dense(on_space) for rate, rep in zip(self.unitary_rates, self.unitary_reps)])
 

pygsti/evotypes/densitymx_slow/statereps.py

@@ -17,7 +17,7 @@
 from pygsti.baseobjs.statespace import StateSpace as _StateSpace
 from ...tools import basistools as _bt
 from ...tools import optools as _ot
-
+from pygsti import SpaceT
 try:
     from ...tools import fastcalc as _fastcalc
 except ImportError:
@@ -51,7 +51,7 @@ def actionable_staterep(self):
         # a probability/amplitude by POVM effect reps.
         return self  # for most classes, the rep itself is actionable
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         if on_space not in ('minimal', 'HilbertSchmidt'):
             raise ValueError("'densitymx' evotype cannot produce Hilbert-space ops!")
         return self.data
@@ -136,7 +136,7 @@ class StateRepComposed(StateRep):
     def __init__(self, state_rep, op_rep, state_space):
         self.state_rep = state_rep
         self.op_rep = op_rep
-        super(StateRepComposed, self).__init__(state_rep.to_dense('HilbertSchmidt'), state_space)
+        super(StateRepComposed, self).__init__(state_rep.to_dense("HilbertSchmidt"), state_space)
         self.reps_have_changed()
 
     def reps_have_changed(self):
@@ -158,9 +158,9 @@ def reps_have_changed(self):
         if len(self.factor_reps) == 0:
             vec = _np.empty(0, 'd')
         else:
-            vec = self.factor_reps[0].to_dense('HilbertSchmidt')
+            vec = self.factor_reps[0].to_dense("HilbertSchmidt")
             for i in range(1, len(self.factor_reps)):
-                vec = _np.kron(vec, self.factor_reps[i].to_dense('HilbertSchmidt'))
+                vec = _np.kron(vec, self.factor_reps[i].to_dense("HilbertSchmidt"))
         self.data[:] = vec
 
     def __reduce__(self):

pygsti/evotypes/stabilizer_slow/effectreps.py

@@ -13,7 +13,7 @@
 from .. import basereps as _basereps
 from pygsti.baseobjs.statespace import StateSpace as _StateSpace
 from ...tools import matrixtools as _mt
-
+from pygsti import SpaceT
 
 class EffectRep(_basereps.EffectRep):
     def __init__(self, state_space):
@@ -29,7 +29,7 @@ def probability(self, state):
     def amplitude(self, state):
         return state.sframe.extract_amplitude(self.zvals)
 
-    def to_dense(self, on_space):
+    def to_dense(self, on_space: SpaceT):
         return _mt.zvals_to_dense(self.zvals, superket=bool(on_space not in ('minimal', 'Hilbert')))
 
 
@@ -46,7 +46,7 @@ def __str__(self):
         s = "Stabilizer effect vector for %d qubits with outcome %s" % (nQubits, str(self.zvals))
         return s
 
-    def to_dense(self, on_space, outvec=None):
+    def to_dense(self, on_space: SpaceT, outvec=None):
         return _mt.zvals_to_dense(self.zvals, superket=bool(on_space not in ('minimal', 'Hilbert')))