DeRhamComputation/as_covers/group.sage

class group:
    def __init__(self, name, elts, one, mult, inv, gens):
        self.name = name
        self.elts = elts
        self.one = one
        self.mult = mult
        self.inv = inv
        self.order = len(self.elts)
        self.gens = gens
        
    def __repr__(self):
        return self.name
    
    def elt(self, a_tuple):
        return group_elt(a_tuple, self)
    
    def ONE(self):
        return self.elt(self.one)
    
    def GENS(self):
        return [self.elt(aa) for aa in self.gens]
    
class group_elt:
    def __init__(self, as_tuple, group):
        self.group = group
        self.as_tuple = as_tuple
        
    def __repr__(self):
        return str(self.as_tuple)
    
    def __mul__(self, other):
        result_as_tuple = self.group.mult(self.as_tuple, other.as_tuple)
        return group_elt(result_as_tuple, self.group)
    
    def __neg__(self):
        result_as_tuple = self.group.inv(self.as_tuple)
        return group_elt(result_as_tuple, self.group)
    
    def __pow__(self, m):
        if m == 0:
            return self.group.ONE()
        if m == 1:
            return self
        if m == 2:
            return self*self
        if m < 0:
            return -(self^(-m))
        if m%2 == 1:
            return (self^(m//2))^2*self
        return (self^(m//2))^2
    
    def __eq__(self, other):
        return self.as_tuple == other.as_tuple
    
    
def cyclic_gp(p, n):
    name = "cyclic group of order " + str(p) + "^" + str(n)
    elts = [i for i in range(p^n)]
    one = 0
    mult = lambda i1, i2: (i1 + i2) % (p ** n)
    inv = lambda i: (-i) % (p ** n)
    gens = [1]
    gp = group(name, elts, one, mult, inv, gens)
    return gp

def elementary_gp(p, n):
    name = "(Z/" + str(p) + ")" + "^" + str(n)
    pr = [list(GF(p)) for _ in range(n)]
    from itertools import product
    elts = []
    for a in product(*pr):
        elts += [a]
    one = elts[0]
    mult = lambda i1, i2: [(i1[j] + i2[j]) % p for j in range(n)]
    inv = lambda i: [(-i[j]) % p for j in range(n)]
    gens = []
    for i in range(n):
        e = n*[0]
        e[i] = 1
        gens += [tuple(e)]
    gp = group(name, elts, one, mult, inv, gens)
    return gp



def heisenberg(p):
    name = "Heisenberg group E(" + str(p) + "^3)"
    elts = [(i, j, k) for i in range(p) for j in range(p) for k in range(p)]
    one = 0
    mult = lambda elt1, elt2 : ((elt1[0] + elt2[0])%p, (elt1[1] + elt2[1])%p, (-elt1[1]*elt2[0] + elt1[2] + elt2[2])%p)
    inv = lambda elt : (p-elt[0], p-elt[1], (p - elt[2] - (p-elt[0])*(p-elt[1]))%p)
    gens = [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
    gp = group(name, elts, one, mult, inv, gens)
    return gp