codedistance 0.0.1__py3-none-any.whl

codedistance/NHow.py

@@ -0,0 +1,804 @@
+import numpy as np
+import numba as nb
+from .common import *
+
+#####################################################
+## Ring Functions for Calculating Howell Matrix Form
+#####################################################
+
+@nb.jit(nb.int16[:,:](nb.int16,nb.int16,nb.int16))
+def Gcdex_jit(a,b,N):
+    '''Extended GCD: Return g,s,t,u,v such that: as + bt = g where g = gcd(a,b); AND au + bv = 0'''
+    a = a % N
+    b = b % N
+    s = nb.int16(0)
+    old_s = nb.int16(1)
+    t = nb.int16(1)
+    old_t = nb.int16(0)
+    r = b
+    old_r = a
+    while r != nb.int16(0):
+        quotient = old_r // r
+        (old_r, r) = (r, old_r - quotient * r)
+        (old_s, s) = (s, old_s - quotient * s)
+        (old_t, t) = (t, old_t - quotient * t)
+    p = np.sign(t * old_s - s * old_t)
+    u = p * s
+    v = p * t
+    # g = old_r
+    s = old_s
+    t = old_t
+    return np.array([[s,t],[u,v]],dtype=nb.int16)
+
+@nb.jit(nb.int16(nb.int16,nb.int16))
+def Ann_jit(a,N):
+    '''Annihilator of a modulo N: Return u such that a u mod N = 0. Return 0 if a is a unit. Return 1 if a == 0'''
+    a = a % N
+    if a == 0:
+        return 1
+    u = N // np.gcd(a,N)
+    return u % N
+
+@nb.jit(nb.int16(nb.int16,nb.int16))
+def Split_jit(a,N):
+    a = a % N
+    if N == nb.int16(0):
+        return nb.int16(0)
+    if a == nb.int16(0):
+        return nb.int16(1)
+    r = int(np.ceil(np.log2(np.log2(N)))) if N > 1 else nb.int16(1)
+    for i in range(r):
+        a = a*a % N
+    return N // np.gcd(a,N)
+
+@nb.jit(nb.int16(nb.int16,nb.int16,nb.int16))
+def Stab_jit(a,b,N):
+    ### return c such that GCD(a + bc, N) = GCD(a,b) modulo N
+    a = nb.int16(a % N)
+    b = nb.int16(b % N)
+    g = np.gcd(np.gcd(a,b),N)
+    c = Split_jit(a//g,N//g)
+    return c % N
+
+@nb.jit(nb.int16(nb.int16,nb.int16,nb.int16))
+def Div_jit(a,b,N):
+    '''Return c such that bc = a mod N or None if no such c exists'''
+    a = a % N
+    b = b % N
+    if b < 1:
+        return nb.int16(-1)
+    g = np.gcd(b,N)
+    if a % g == nb.int16(0):
+        r = a % b
+        while r > nb.int16(0):
+            a += N
+            r = a % b
+        return a // b % N
+    return nb.int16(-1)
+
+@nb.jit(nb.int16(nb.int16,nb.int16))
+def Unit_jit(a,N):
+    '''Return a unit c such that ac = gcd(a,N)  mod N.'''
+    a = a % N
+    if a == nb.int16(0):
+        return nb.int16(1)
+    g = np.gcd(a,N)
+    s = Div_jit(g,a,N)
+    if g == nb.int16(1):
+        return s
+    d = Stab_jit(s,N//g,N)
+    c = (s + d * N // g) % N
+    return c
+
+
+
+#####################################################
+## RREF Modulo 2
+#####################################################
+
+
+@nb.jit(nb.types.Tuple((nb.int8[:,:],nb.int64[:]))(nb.int8[:,:],nb.int64,nb.int64,nb.int64,nb.int64))
+def HowZ2Set(A,tB,nB,nC,r0):
+    if np.sum(A) == 0:
+        return A,np.arange(0,dtype=nb.int64)
+    m,n = A.shape
+    pivots = []
+    B = []
+    # B = [set() for i in range(m)]
+    # rowData,colData = np.nonzero(A)
+    # for ix in range(len(rowData)):
+    #     B[rowData[ix]].add(colData[ix])
+    for myrow in A:
+        s = set()
+        for i in range(n):
+            if myrow[i]==1:
+                s.add(i)
+        B.append(s)
+    r = r0
+    for j in range(nC):
+        for t in range(tB):
+            ## c is the column of B we are currently looking at
+            c = j + t * nB
+            iList = [i for i in range(r,m) if c in B[i]]
+            if len(iList) > 0:
+                i = iList.pop(0)
+                pivots.append(c)
+                ## found j: if j > r, swap row j with row r
+                if i > r:  
+                    ## swap using bitflips - more elegant than array indexing
+                    B[r] ^= B[i]
+                    B[i] ^= B[r] 
+                    B[r] ^= B[i]
+                ## eliminate non-zero entries in column c apart from row r
+                for i in [i for i in range(r) if c in B[i] ] + iList:     
+                    B[i] ^= B[r]
+                r +=1
+    C = np.zeros((m,n),dtype=nb.int8)
+    for i in range(m):
+        for j in B[i]:
+            C[i][j] = 1
+    return C,np.array(pivots,dtype=nb.int64)
+
+def KerZ2(A,cols,r0,verbose=False):
+    '''fast kernel method - test'''
+    # if RREF(A) is of form I|D up to col perms then ker(A) = D.T|I
+    A = np.array(A,dtype=np.int8)
+    B, pivots = RREFZ2(A,cols,r0)
+    w = np.sum(B, axis=-1)
+    B = B[w > 0]
+    m, n = B.shape
+    k = n - m
+    nonpivots = invRange(n,pivots)
+    C = np.zeros((k,n),dtype=np.int8)
+    for i in range(k):
+        C[i,nonpivots[i]] = 1
+    C[:,pivots] = B[:,nonpivots].T
+    if verbose:
+        return B, pivots, C
+    return C
+
+@nb.njit(nb.types.Tuple((nb.int8[:,:],nb.int64[:]))(nb.int8[:,:],nb.int64[:],nb.int64))
+def RREFZ2(A,cols,r0):
+    '''RREF using cols as indices to perform reductions'''
+    pivots = []
+    B = A.copy()
+    if np.sum(B) > 0:
+        m = len(B)
+        r = r0
+        for c in cols:
+            iList = [i for i in range(r,m) if B[i,c] > 0]
+            if len(iList) > 0:
+                i = iList.pop(0)
+                pivots.append(c)
+                ## found j: if j > r, swap row j with row r
+                if i > r:  
+                    ## swap using bitflips - more elegant than array indexing
+                    B[r] ^= B[i]
+                    B[i] ^= B[r] 
+                    B[r] ^= B[i]
+                ## eliminate non-zero entries in column c apart from row r
+                for i in [i for i in range(r) if B[i,c] > 0] + iList:     
+                    B[i] ^= B[r]
+                r += 1
+    return B,np.array(pivots,dtype=nb.int64)
+
+@nb.jit(nb.types.Tuple((nb.int8[:,:],nb.int64[:]))(nb.int8[:,:],nb.int64,nb.int64,nb.int64,nb.int64))
+def HowZ2(A,tB,nB,nC,r0):
+    pivots = []
+    B = A.copy()
+    if np.sum(A) == 0:
+        return B,np.array(pivots,dtype=nb.int64)
+    m = len(B)
+    r = r0
+
+    for j in range(nC):
+        for t in range(tB):
+            ## c is the column of B we are currently looking at
+            c = j + t * nB
+            iList = [i for i in range(r,m) if B[i,c] > 0]
+            if len(iList) > 0:
+                i = iList.pop(0)
+                pivots.append(c)
+                ## found j: if j > r, swap row j with row r
+                if i > r:  
+                    ## swap using bitflips - more elegant than array indexing
+                    # B[r] = B[r] ^ B[i]
+                    # B[i] = B[r] ^ B[i]
+                    # B[r] = B[r] ^ B[i]
+                    B[r] ^= B[i]
+                    B[i] ^= B[r] 
+                    B[r] ^= B[i]
+                ## eliminate non-zero entries in column c apart from row r
+                for i in [i for i in range(r) if B[i,c] > 0] + iList:     
+                    # B[i] = B[i] ^ B[r]
+                    B[i] ^= B[r]
+                r +=1
+    return B,np.array(pivots,dtype=nb.int64)
+
+@nb.jit(nb.int8[:,:](nb.int8[:,:]))
+def binSpanNb(A):
+    r,n = A.shape
+    temp = np.zeros((1 << r, n),dtype=nb.int8)
+    j = 1
+    for i in range(r):
+        k = j << 1
+        temp[j:k] = temp[:j] ^ A[i]
+        j = k
+    return temp
+
+def SpanBin(A):
+    '''Span of binary matrix A - seems faster than jit version...'''
+    r,n = A.shape
+    temp = np.zeros((1 << r, n),dtype=np.int8)
+    j = 1
+    for i in range(r):
+        k = j << 1
+        temp[j:k] = temp[:j] ^ A[i]
+        j = k
+    return temp
+
+
+#####################################################
+## Calculate Howell Matrix form modulo N
+#####################################################
+
+@nb.jit(nb.types.Tuple((nb.int16[:,:],nb.types.List(nb.int64)))(nb.int16[:,:],nb.int16,nb.int64,nb.int64,nb.int64,nb.int64))
+def HowZN(A,N,tB,nB,nC,r0):
+    '''Return Howell basis of A mod N plus row operations to convert to this form'''
+    pivots = []
+    if np.sum(A) == 0:
+        return A,pivots
+    m,n = A.shape
+    B = [a for a in A]
+    N = nb.int16(N)
+    r = r0
+
+    ## c is the column of B we are currently looking at
+    for mc in range(nC):
+        for t in range(tB):
+            c = mc + t * nB
+            ## find j such that B[j][c] > 0
+            jList = [j for j in range(r,m) if B[j][c] > 0]
+            if len(jList) > 0:
+                j = jList.pop(0)
+                pivots.append(mc)
+                ## found j: if j > r, swap row j with row r
+                if j > r:  
+                    B[j],B[r] = B[r],B[j]
+                ## Multiplying by x ensures that B[r][c] is a minimal representative
+                b = B[r][c]
+                x = Unit_jit(b,N)
+                if(x > 1):
+                    B[r] = np.mod(B[r] * x, N)
+                ## eliminate entries in column c below row r
+                for j in jList:
+                    a, b = B[r][c],B[j][c]
+                    C = Gcdex_jit(a,b,N)
+                    Br = np.mod(B[r] * C[0,0] + B[j] * C[0,1], N)
+                    Bj = np.mod(B[r] * C[1,0] + B[j] * C[1,1], N)
+                    B[r] = Br
+                    B[j] = Bj
+                ## ensure entries in column c above row r are less than B[r][c]
+                b = B[r][c]
+                for j in range(r):
+                    if B[j][c] >= b:
+                        x = nb.int16(B[j][c] // b)
+                        B[j] = np.mod(B[j] - B[r] * x,N)
+                ## Multiplying by x = Ann(b) eliminates b = B[r][c], but rest of the row may be non-zero
+                ## If x > 0 then b is a zero divisor and we add a row
+                ## If x == 0, b is a unit and we move to the next value of l
+                x = Ann_jit(b,N)
+                if x > 0:
+                    B.append(np.mod(B[r] * x,N))
+                    m += 1
+                r +=1
+    temp = np.empty((m,n),dtype=nb.int16)
+    for i in range (m):
+        temp[i] = B[i]
+    return temp,pivots
+
+# @nb.jit
+def blockDims(n,nA=0,tB=1,nC=-1):
+    nA = min(n,nA)
+    nB = (n - nA) // tB
+    if nC < 0 or nC > nB:
+        nC = nB 
+    return nA,nB,nC
+
+def getH(A,N,nA=0,tB=1,nC=-1,r0=0,retPivots=False):
+    '''Return Howell matrix form modulo N:
+    A: input matrix
+    N: linear algebra modulo N
+    nA: number of cols appended to right - not subject to row reduction
+    tB: number of blocks in the matrix
+    nC: number of columns to reduce
+    r0: starting row for reduction'''
+    ## nB: number of columns in each block
+    m,n = A.shape
+    nA,nB,nC = blockDims(n,nA,tB,nC)
+    if N==2:
+        A = np.array(A,dtype=np.int8)
+        H,pivots = HowZ2(A,tB,nB,nC,r0)
+        # H,pivots = HowZ2Set(A,tB,nB,nC,r0)
+    else:
+        A = np.array(A,dtype=np.int16)
+        H,pivots = HowZN(A,N,tB,nB,nC,r0)
+    w = np.sum(H,axis=-1)
+    ix = [i for i in range(len(H)) if i < r0 or w[i] >0]
+    H = H[ix]
+    return (H, list(pivots)) if retPivots else H
+    
+def getHU(A,N,nA=0,tB=1,nC=-1,r0=0):
+    '''Return Howell matrix form modulo N plus transformation matrix U such that H = U @ A mod N'''
+    m,n = A.shape
+    B = np.hstack([A,ZMatI(m)])
+    nA += m
+    HU = getH(B,N,nA,tB,nC,r0)
+    return HU[:,:n],HU[:,n:]
+
+def getK(A,N,nA=0,tB=1,nC=-1,retPivots=False):
+    '''Return Kernel K such that K @ A.T = 0 mod N'''
+    ## Transpose so nA,tB,nC,r0 not applicable
+    H, U = getHU(A.T,N)
+    ix = np.sum(H,axis=-1) == 0
+    K = U[ix,:]
+    ## r0 not applicable...
+    return getH(K,N,nA,tB,nC,retPivots=retPivots)
+
+
+def HowRes(A,B,N,tB=1):
+    '''Return R, V such that  B = R + V @ A mod N for A,B matrices, unknowns R, V'''
+    return solveH(A,B,N,tB)
+
+def HowResU(A,B,N,tB=1):
+    '''Return R, V such that  B = R + V @ A mod N for A,B matrices, unknowns R, V'''
+    R,V,H,U,K = solveHU(A,B,N,tB)
+    return R,V
+
+def solveHU(A,B,N,tB=1):
+    '''Solve B = R + (V + <K>) @ A mod N for A,B matrices, unknowns R, V, K
+    A: m x n
+    B: r x n
+    R: r x n
+    V: r x m
+    K: k x m'''
+    B1D = len(B.shape) == 1
+    A, B = ZMat2D(A), ZMat2D(B)
+    m,n = A.shape
+    r,n1 = B.shape
+    if n != n1:
+        print(func_name(), 'A,B incompatible shape')
+    ## Make matrix of form [[I,B,0],[0,A,I]]
+    BA = np.hstack([B,ZMatZeros((r,m))])
+    BA = ZMatVstack([BA,np.hstack([A,ZMatI(m)])])
+    ## Howell form - only consider first m2 + n1 columns
+    HBA = getH(BA,N,nA=m,tB=tB,r0=r)
+    ## How results in matrix of form [[I,R,V],[0,H,U],[0,0,K]]
+    ## R: top middle block of HBA
+    R = HBA[:r,:n]
+    ## if B is 1D, return 1D result
+    if B1D:
+        R = R[0]
+    ## V: top right block of HBA - negative to reverse residue calc
+    V = np.mod(-HBA[:r,n:],N)
+    ## H: bottom middle block of HBA
+    H = HBA[r:,:n]
+    ## U: bottom right block of HBA
+    U = HBA[r:,n:]
+    ix = np.sum(H,axis=-1) > 0
+    H,U,K = H[ix,:],U[ix,:],U[np.logical_not(ix), :]
+    return R,V,H,U,K
+
+def solveH(A,B,N,tB=1):
+    '''Solve B = R + V @ A mod N for unknown R, H = getH(A,N)
+    A: m x n
+    B: r x n
+    R: r x n'''
+    B1D = len(np.shape(B)) == 1
+    A, B = ZMat2D(A), ZMat2D(B)
+    m,n = A.shape
+    r,n1 = B.shape
+    if n != n1:
+        print(func_name(), 'A,B incompatible shape')
+    ## Make matrix of form [[B],[A]]
+    BA = ZMatVstack([B,A])
+    ## Howell form - only consider first m2 + n1 columns
+    HBA = getH(BA,N,tB=tB,r0=r)
+    ## How results in matrix of form [[I,R,V],[0,H,U],[0,0,K]]
+    ## R: top block
+    R = HBA[:r,:]
+    ## H: lower block
+    H = HBA[r:,:]
+    ## if B is 1D, return 1D result
+    if B1D:
+        R = R[0]
+    return R,H
+
+# @nb.jit
+def HowPivots(A,nA=0,tB=1,nC=-1,r0=0):
+    '''Leading indices of matrix in Howell/RREF form
+    - Examine first nC columns
+    - Matrix has nBlocks blocks A = [A_0|A_1|..|A_n-1]'''
+    temp = []
+    if np.sum(A) == 0:
+        return temp
+    m,n = A.shape
+    nA,nB,nC = blockDims(n,nA,tB,nC)
+    r,c = r0,0
+    for r in range(m):
+        while c < nC and np.all([A[r,c + t * nB] == 0 for t in range(tB)]):
+            c+=1
+        if c < nC:
+            temp.append(c)
+        else:
+            return temp
+    return temp
+
+# def ZMatBlockSum(A,nA=0,tB=1,nC=-1,N=2):
+#     '''Sum block of Zmat
+#     - up to nC
+#     - tB - number of blocks
+#     - if A = [A_0|A_1|..|A_n] return A_0 + 2 A_1 + ... + 2^n A_n'''
+#     nB,nC = blockDims(A,nA,tB,nC)
+#     ix = ZMat(range(nC))
+#     S = A.take(indices=ix, axis=-1)
+#     for t in range(1,tB):
+#         ix += nB
+#         S += A.take(indices=ix, axis=-1) * (N ** t)
+#     return S
+
+# @nb.jit
+# @nb.jit(nb.types.Tuple((nb.int8[:,:],nb.int64[:]))(nb.int8[:,:],nb.int64,nb.int64,nb.int64,nb.int64))
+
+
+
+def indepL(L,pList,nA,tB,retWeights=False):
+    w = np.sum(L[:,-nA:],axis=-1)
+    L = L[w > 0,:]
+    probs,weights = ZMatProbW(L,pList=pList,nA=nA,tB=tB,N=2)
+    ix = argsort(probs)
+    probs = probs[ix]
+    weights = weights[ix]
+    L = L[ix,:]
+    H = L[:,-nA:]
+    H,U = getHU(H,2)
+    K = U[np.sum(H,axis=-1) == 0,:]
+    K,LI = getH(K,2,retPivots=True)
+    ix = invRange(len(L),LI)
+    L = L[ix,:]
+    probs = probs[ix]
+    weights = weights[ix]
+    if retWeights:
+        return L,probs,weights
+    else:
+        return L,probs
+
+
+def uniqL(L,pList,nA,tB,pMax = 1e6):
+    H = L[:,-nA:]
+    ix = np.sum(H,axis=-1)
+    L = L[ix>0]
+    # L = ZMat({tuple(x) for x in L})
+    L = np.unique(L,axis=0)
+    probs = ZMatProb(L,pList,nA,tB)
+    if len(probs) > pMax:
+        ix = argsort(probs)
+        ix = ix[-pMax:]
+        L = L[ix]
+        probs = probs[ix]
+    return L,probs
+
+# @nb.jit
+def ZMatTake(A,ix):
+    '''Take indices ix form A along axis -1 '''
+    B = A[:,ix].copy()
+    return B
+
+    
+def XZhad(E, ix=None):
+    '''Apply Hadamard to qubits according to index ix for Pauli operator list E'''
+    m,n = E.shape
+    nB = n//2
+    if ix is None:
+        ix = np.arange(nB)
+    ixN = ix + nB
+    ix1 = np.arange(n)
+    ix1[ix] = ixN
+    ix1[ixN] = ix
+    return ZMatPermuteCols(E,ix1)
+
+def XZCX(E, a,b):
+    '''Apply CX(a->b) for Pauli operator list E'''
+    m,n = E.shape
+    nB = n//2
+    E1 = E.copy()
+    E1[:,b] ^= E[:,a]
+    E1[:,a+nB] ^= E[:,b+nB]
+    return E1
+
+def XZCZ(E, a,b):
+    '''Apply CZ(a b) for Pauli operator list E'''
+    m,n = E.shape
+    nB = n//2
+    E1 = E.copy()
+    E1[:,a+nB] ^= E[:,b]
+    if a != b:
+        ## if a==b, we apply S; otherwise it's CZ
+        E1[:,b+nB] ^= E[:,a]
+    return E1
+
+# @nb.jit
+def ZMatBlockSum(A,nA=0,tB=1,nC=-1,N=2):
+    '''Sum block of Zmat
+    - up to nC
+    - tB - number of blocks
+    - if A = [A_0|A_1|..|A_n] return A_0 + 2 A_1 + ... + 2^n A_n'''
+    m,n = A.shape
+    # print(len(A))
+    nA,nB,nC = blockDims(n,nA,tB,nC)
+    ix = np.arange(nC)
+    S = ZMatTake(A,ix)
+    for t in np.arange(1,tB):
+        ix += nB
+        S += ZMatTake(A,ix) * (N ** t)
+    return S
+
+def ZMatBlockSum(A,nA=0,tB=1,nC=-1,N=2):
+    '''Sum block of Zmat
+    - up to nC
+    - tB - number of blocks
+    - if A = [A_0|A_1|..|A_n] return A_0 + 2 A_1 + ... + 2^n A_n'''
+    m,n = A.shape
+    # print(len(A))
+    nA,nB,nC = blockDims(n,nA,tB,nC)
+    S = np.empty((m,nC),dtype=int)
+    c1 = 0
+    c2 = c1 + nC
+    S[:,:] = A[:,c1:c2]
+    c1 = c2
+    for t in range(1,tB):
+        c2 = c1 + nC
+        if N == 2:
+            S += (A[:,c1:c2] << t)
+        else:
+            S += A[:,c1:c2] * (N ** t)
+        c1 = c2
+    return S
+
+def ZMatPermuteCols(A,ix,nA=0,tB=1,nC=-1):
+    m,n = A.shape
+    nA,nB,nC = blockDims(n,nA=nA,tB=tB,nC=nC)
+    ix = list(ix) + invRange(nB,ix)
+    ix = ZMat(ix)
+    ind = np.arange(n)
+    for t in range(tB):
+        ind[t*nB:(t+1)*nB] = ix
+        ix += nB
+    return ZMatTake(A,ind)
+
+
+
+
+
+####################################
+## Matrix Multiplication - avoid overflow for large matrices
+####################################
+
+# @nb.jit (nb.int16[:,:](nb.int16[:,:],nb.int16[:,:],nb.int64))
+# def matMulZN(A,B,N):
+#     m1,n1 = A.shape
+#     m2,n2 = B.shape
+#     C = np.zeros((m1,n2),dtype=nb.int16)
+#     m = min(n1,m2)
+#     for i in range(m1):
+#         for j in range(n2):
+#             temp = nb.int16(0)
+#             for k in range(m):
+#                 temp = np.mod(temp + (A[i,k] * B[k,j]),N)
+#             C[i,j] = temp
+#     return C
+
+def matMulZN(A,B,N):
+    return np.mod(A @ B, N)
+
+@nb.jit (nb.int8[:,:](nb.int8[:,:],nb.int8[:,:]))
+def matMulZ2(A,B):
+    m1,n1 = A.shape
+    m2,n2 = B.shape
+    C = np.zeros((m1,n2),dtype=nb.int8)
+    m = min(n1,m2)
+    for i in range(m1):
+        for j in range(n2):
+            temp = nb.int8(0)
+            for k in range(m):
+                temp = temp ^ (A[i,k] & B[k,j])
+            C[i,j] = temp
+    return C
+
+def matMul(A,B,N):
+    '''Multiply two integer matrices modulo N'''
+    A = ZMat2D(A)
+    B = ZMat2D(B)
+    if N==2:
+        A = np.array(A,dtype=np.int8)
+        B = np.array(B,dtype=np.int8)
+        return matMulZ2(A,B)
+    else:
+        A = np.array(A,dtype=np.int16)
+        B = np.array(B,dtype=np.int16)
+        return matMulZN(A,B,N)
+    return np.mod(A @ B, N)
+
+def mod1(A):
+    '''Replace values > 0 in A with 1'''
+    return np.clip(A,0,1)
+    A = ZMat(A)
+    A[A>0] = 1
+    return A
+
+def RemoveZeroRows(A,N=False):
+    '''Remove any zero rows from integer matrix A'''
+    # A = ZMat(A)
+    w = np.sum(A, axis=-1)
+    return A[w > 0]
+
+def restoreShape(A,s):
+    if s!=A.shape:
+        return np.reshape(A,s)
+    return A
+
+
+def ZMatWeight(A,nA=0,tB=1,nC=-1):
+    s = A.shape
+    if len(s) != 2:
+        return restoreShape(np.sum(mod1(ZMatBlockSum(ZMat2D(A),nA,tB,nC,1)),axis=-1),s[:-1])
+    return np.sum(mod1(ZMatBlockSum(A,nA,tB,nC,1)),axis=-1)
+
+def pListDefault(tB=1,pI=0.7):
+    pLen = (1 << tB) - 1
+    p = (1-pI)/pLen
+    return np.array([pI] + [p] * pLen)
+
+def ZMatProbW(A,pList=None,nA=0,tB=1,nC=-1,N=2):
+    '''Return probabilities and weight of each row of A
+    '''
+    if len(A) == 0:
+        return [],[]
+    if pList is None:
+        pList = pListDefault(tB)
+    pList = np.array(pList) 
+    W = ZMatBlockSum(A,tB=tB,N=2,nA=nA)
+    s = np.prod(pList[W],axis=-1)
+    W = np.sum(mod1(W),axis=-1)
+    return s, W
+
+def ZMatProb(A,pList,nA=0,tB=1,nC=-1,N=2):
+    s, w =  ZMatProbW(A,pList,nA=0,tB=1,nC=-1,N=2)
+    return s
+
+def lowWeightGens(L,S=None,N=2,tB=1,pList=None,retProb=False):
+    '''Get set of lowest weight generators of logicals L
+    S is a list of stabilisers
+    tB is number of blocks
+    N is precision
+    pList is list of probabilities'''
+    # Default probabilities
+    if pList is None:
+        pList = pListDefault(tB=tB)
+    w = ZMatProb(L,pList,tB=tB,N=N)
+    ## order L by increasing prob/decreasing weight
+    ix = argsort(w)
+    L = L[ix]
+    w = w[ix]
+    LS = L if S is None else ZMatVstack([L,S])
+    ## RREF plus transformation matrix
+    H, U = getHU(LS,N,tB=tB)
+    ## K is a list of linear combinations of rows which result in zero
+    ix = np.sum(H,axis=-1) == 0
+    K = U[ix,:]
+    ## RREF - so combinations of lowest weight rows are to the RHS
+    K, LI = getH(K,N,retPivots=True)
+    # K = K[:,:len(L)]
+    ## get leading entries of K - these correspond to rows we should exclude 
+    ## as adding them to the set results in a redundancy
+    # LI = HowPivots(K,nC=len(L))
+    ix = invRange(len(L),LI)
+    L = L[ix,:]
+    w = w[ix]
+    return (L,w) if retProb else L
+
+##################################################
+####            Operations on Spans           ####
+##################################################
+
+def affineIntersection(A1,o1,A2,o2,N,C=False):
+    '''Calculate intersection of two affine spaces
+    U1 = o1 + <A1> mod N
+    U2 = o2 + <A2> mod N'''
+    if C is not False:
+        A,o = C
+        tocheck = np.mod(o + A,N)
+        tocheck = ZMatVstack([[o],tocheck])
+        for v in tocheck:
+            v1 = np.mod(v - o1,N)
+            b,u = HowRes(A1,v1,N)
+            if not isZero(b):
+                print(v1, 'Not in span of A1')
+                return False
+            v2 = np.mod(v - o2,N)
+            b,u = HowRes(A2,v2,N)
+            if not isZero(b):
+                print(v2, 'Not in span of A2')
+                return False
+        return True
+    o = affineIntercept(A1,o1,A2,o2,N)
+    if o is False:
+        return False
+    ## new affine space is intersection
+    A = nsIntersection([A1,A2],N)
+    ## residue mod A to simplify result
+    R, V = HowRes(A, o, N)
+    return A,R[0]
+
+def affineIntercept(A1,o1,A2,o2,N,C=False):
+    '''Calculate intersection of two affine spaces
+    U1 = o1 + <A1> mod N
+    U2 = o2 + <A2> mod N
+    return o which is in both U1 and U2 or False'''
+    if C is not False:
+        ## check if C-o1 is in <A1>
+        R,V = HowRes(A1,np.mod(C-o1,N),N)
+        if not isZero(R):
+            return False
+        ## check if C-o2 is in <A2>
+        R,V = HowRes(A2,np.mod(C-o2,N),N)
+        if not isZero(R):
+            return False
+        return True
+    A = ZMatVstack([A1, A2])
+    ## find solution o1-o2 = (v1|v2) @ (A1//A2) = v1@A1 + v2@A2 mod N <=> o1 - v1@A1 = o2 + v2@A2
+    R, v1v2 = HowResU(A, np.mod(o1-o2,N), N)
+    if not isZero(R):
+        ## there is no solution if residue is non-zero
+        return False
+    ## extract v1 from v1v2 - corresponds to first len(A1) entries of the vector
+    v1 = v1v2[:,:len(A1)]
+    v1A1 = matMul(v1, A1, N)[0]
+    ## note minus sign
+    return np.mod(o1 - v1A1,N)
+
+def nsIntersection(Alist,N,C=False):
+    '''Intersection of multiple rowspans AList modulo N.'''
+    if C is not False:
+        for A in Alist:
+            ## check if C in <A> for each A in AList
+            R,V = HowRes(A,C,N)
+            if not isZero(R):
+                return False
+        return True
+    if len(Alist) == 0:
+        return False
+    A = Alist[0]
+    for B in Alist[1:]:
+        if len(B) < len(A):
+            A,B = B,A
+        AB = ZMatVstack([A,B])
+        Kt = getK(AB.T,N)
+        A0 = ZMatVstack([A,ZMatZeros(B.shape)])
+        C = matMul(Kt, A0,N)
+        A = getH(C,N)
+    return A
+
+def nsUnion(Alist,N,C=False):
+    '''Union of multiple rowspaces Alist modulo N'''
+    if C is not False:
+        ## check if A in C for all A in Alist
+        for A in Alist:
+            R,V = HowRes(C,A,N)
+            if not isZero(R):
+                return False
+        return True
+    if len(Alist) == 0:
+        return False
+    return getH(ZMatVstack(Alist),N) 
+