Refactor Schottky and Theta classes for improved functionality and performance

Author: mmasdeu

darmonpoints/schottky.py

@@ -27,7 +27,7 @@
 from .divisors import Divisors, DivisorsElement
 from .meromorphic import *
 from .theta import *
-from .util import muted
+from .util import muted, our_sqrt
 
 infinity = Infinity
 
@@ -100,17 +100,20 @@ def uniq(lst):
 
 
 @muted
-@cached_function
 def find_eigenvector_matrix(g):
+    K = g.parent().base_ring()
     t = g.trace()
     n = g.determinant()
-    delta = (t * t - 4 * n).sqrt()
-    vaps = sorted([(t + delta) / 2, (t - delta) / 2], key=lambda o: o.valuation())
-    veps = sum(((g - l).right_kernel().basis() for l in vaps), [])
+    delta = (t * t - K(4) * n).sqrt()
+    vaps = sorted([(t + delta) / K(2), (t - delta) / K(2)], key=lambda o: o.valuation())
+    assert len(vaps) == 2
+    veps = []
+    for l in vaps:
+        B = g - l
+        veps.append((B[0,1], -B[0,0])) # The right kernel of a rank-1 2x2 matrix
     return g.matrix_space()(veps).transpose()
 
 
-@cached_function
 def find_parameter(g, ball, pi=None, check=True):
     if pi is None:
         pi = g.parent().base_ring().uniformizer()
@@ -210,6 +213,8 @@ def theta_naive(self, m, a=ZZ(0), b=ZZ(1), z=None, gamma=None, **kwargs):
         K = self.K
         if z is Infinity:
             return K(1)
+        if isinstance(z, Divisors): # z is a divisor!
+            return prod(self.theta_naive(m, a, b, P, gamma, **kwargs)**n for P, n in z)
         if gamma is not None:
             if b != ZZ(1):
                 raise ValueError("If gamma is specified, then b should not")
@@ -225,9 +230,9 @@ def theta_naive(self, m, a=ZZ(0), b=ZZ(1), z=None, gamma=None, **kwargs):
                 can_stop = True
             for _, g in self.enumerate_group_elements(i):
                 ga = act(g, a)
-                tmp1 = K(1) if ga is Infinity else z - ga
+                tmp1 = K(1) if ga is Infinity or z == ga else z - ga
                 gb = act(g, b)
-                tmp2 = K(1) if gb is Infinity else z - gb
+                tmp2 = K(1) if gb is Infinity or z == gb else z - gb
                 new_num *= tmp1
                 new_den *= tmp2
                 try:
@@ -353,7 +358,7 @@ def construct_tree(self, level):
         return self.NJtree
 
     @cached_method
-    @muted
+    # @muted # DEBUG
     def fundamental_domain(self, level=1, check=True):
         while True:
             level += 1
@@ -364,7 +369,7 @@ def fundamental_domain(self, level=1, check=True):
             except ValueError as e:
                 verbose(str(e))
 
-    @muted
+    # @muted # DEBUG
     def _fundamental_domain(self, i0=None, check=True):
         tree = self.NJtree
         if i0 is None:
@@ -546,13 +551,14 @@ def test_fundamental_domain(self):
         gens = self._generators
         test_fundamental_domain(gens, balls)
 
-    def balls(self):
+    def balls(self, generators=None):
         try:
             return self._balls
         except AttributeError:
             pass
         K = self.base_ring()
-        generators = self._generators
+        if generators is None:
+            generators = self._generators
         G = PreSchottkyGroup(K, generators)
         gp = G.group()
         good_gens, good_balls, _ = G.fundamental_domain()
@@ -608,22 +614,18 @@ def a_point(self):
                 return x0
         raise RuntimeError("Reached maximum iterations (100)")
 
-    def find_point(self, gamma, eps=1, idx=None):
+    def find_point(self, gamma, eps=1, idx=None, check=True):
         balls = self.balls()
         gens = self.gens_extended()
         if idx is not None:
             B = balls[-idx]
             if idx > 0:
                 if self._generators[idx - 1] != gamma:
                     B1 = next(balls[-i] for i, g in gens if g == gamma)
-                    # print(f'!! {B = } but {B1 = }')
-                    # print(balls)
                     B = B1
             else:
                 if self._generators[-idx - 1] ** -1 != gamma:
                     B1 = next(balls[-i] for i, g in gens if g == gamma)
-                    # print(f'!! {B = } but {B1 = }')
-                    # print(balls)
                     B = B1
         else:
             B = next(balls[-i] for i, g in gens if g == gamma)
@@ -632,11 +634,10 @@ def find_point(self, gamma, eps=1, idx=None):
         except TypeError:
             raise RuntimeError(
                 "Fundamental domain has balls with fractional radius. \
-            Try with a quadratic (ramified) extension."
+Try with a quadratic (ramified) extension."
             )
-        test = lambda x: self.in_fundamental_domain(x, strict=False)
-
-        try:
+        if check:
+            test = lambda x: self.in_fundamental_domain(x, strict=False)
             if not test(ans):
                 raise RuntimeError(
                     f"{ans = } should be in fundamental domain, but it is not."
@@ -646,11 +647,6 @@ def find_point(self, gamma, eps=1, idx=None):
                 raise RuntimeError(
                     f"gamma * ans = {ga} should be in fundamental domain, but it is not."
                 )
-        except RuntimeError:
-            new_idx = -idx if idx is not None else None
-            ginv = gamma**-1
-            ans = self.find_point(ginv, eps, new_idx)
-            return act(ginv, ans)
         return ans
 
     def to_fundamental_domain(self, x):
@@ -731,7 +727,7 @@ def find_equivalent_divisor(self, D):
             if e == 0:
                 continue
             zz = self.find_point(g, idx=i + 1)
-            ans -= Div([(e, zz), (-e, act(g, zz))])
+            ans += Div([(-e, zz), (e, act(g, zz))])
         return ans
 
     def theta(self, prec, a, b=None, **kwargs):
@@ -767,9 +763,7 @@ def theta(self, prec, a, b=None, **kwargs):
         s = kwargs.pop("s", None)
         if s is not None:
             D0 += s * D0
-        recursive = kwargs.get("recursive", True)
-        if recursive:
-            D0 = self.find_equivalent_divisor(D0)
+        D0 = self.find_equivalent_divisor(D0)
         ans = ThetaOC(self, a=D0, b=None, prec=prec, base_ring=K, **kwargs)
         z = kwargs.pop("z", None)
         improve = kwargs.pop("improve", True)
@@ -788,32 +782,37 @@ def u_function(self, gamma, prec, a=None, **kwargs):
             if z is None:
                 kwargs.pop("z", None)
                 return lambda z: prod(
-                    self._u_function(gens[abs(i) - 1], prec, a=a)(z) ** ZZ(sgn(i))
+                    self._u_function(abs(i) - 1, prec, a=a)(z) ** ZZ(sgn(i))
                     for i in wd
                 )
             return prod(
-                self._u_function(gens[abs(i) - 1], prec, a=a)(z) ** ZZ(sgn(i))
+                self._u_function(abs(i) - 1, prec, a=a)(z) ** ZZ(sgn(i))
                 for i in wd
             )
+        naive = kwargs.get('naive', None)
         if z is None:
-            return lambda z: self._u_function(gens[wd[0] - 1], prec, a=a)(z)
+            return lambda z: self._u_function(wd[0] - 1, prec, a=a, naive=naive)(z)
         else:
-            return self._u_function(gens[wd[0] - 1], prec, a=a)(z)
+            return self._u_function(wd[0] - 1, prec, a=a, naive=naive)(z)
 
     @cached_method
-    def _u_function(self, gamma, prec, a):
+    def _u_function(self, i, prec, a, naive=False):
         r"""
         Compute u_gamma
         """
-        (i,) = self.word_problem(gamma)
-        assert i > 0
+        gamma = self._generators[i]
         if a is None:
-            a = self.a_point()  # self.find_point(gamma, idx=g)
-        a = self.base_ring()(1) * a
-        K = a.parent()
+            a = self.find_point(gamma, idx=i) # self.a_point()
+            K = a.parent()
+        else:
+            a = self.base_ring()(1) * a
+            K = a.parent()            
         D = self.find_equivalent_divisor(Divisors(K)([(1, a), (-1, act(gamma, a))]))
-        ans = ThetaOC(self, a=D, b=None, prec=prec, base_ring=K)
-        ans = ans.improve(prec)
+        if naive:
+            return lambda z : self.theta_naive(prec, z=z, a=a, gamma=gamma)
+        else:
+            ans = ThetaOC(self, a=D, b=None, prec=prec, base_ring=K)
+            ans = ans.improve(prec)
         return ans
 
     @cached_method

darmonpoints/theta.py

@@ -100,8 +100,8 @@ def __init__(self, G, a=None, b=None, prec=None, **kwargs):
         self.K = K
         self.G = G
         self.p = G.pi
-        generators = G.generators()
-        balls = G.balls()
+        _ = G.generators()
+        _ = G.balls()
         if prec is None:
             try:
                 self.prec = K.precision_cap()
@@ -123,20 +123,21 @@ def __init__(self, G, a=None, b=None, prec=None, **kwargs):
         self.b = b
         self.m = 1
         self.z = K["z"].gen()
-
+        self.MM = MeromorphicFunctions(self.K, self.p, self.prec)
         params = G.parameters
         gens_ext = G.gens_extended()
-        # self.val will contain the 0 and 1 terms
-        D1 = sum((g * D for (_, g), _ in zip(gens_ext, params)), self.Div([]))
-        self.val = D + D1
-        # self.val = self.z.parent()(1)
-        # self.val *= prod((self.z - P) ** n for P, n in (D + D1) if P != Infinity)
-        self.MM = MeromorphicFunctions(self.K, self.p, self.prec)
-        self.Fnlist = [{}]
-        D1dict = {i: g * D for (i, g), tau in zip(gens_ext, params)}
-        for (i, g), tau in zip(gens_ext, params):
-            gD = sum(g * val for j, val in D1dict.items() if j != -i)
-            self.Fnlist[0][i] = self.MM(gD, tau)
+        
+        # Corresponding to words of length exactly 1
+        D1dict = {i : g * D for i, g in gens_ext}
+
+        # Corresponding to words of length exactly 2
+        # D2dict = {i : sum(g * val) for j, val in D1dict.items() if j != i for i, g in gens_ext }
+        
+        # self.val will contain the 0, 1 and 2 terms
+        self.val = sum(D1dict.values(), D)
+        self.Fnlist = [{ i :
+            self.MM(sum(g * val for j, val in D1dict.items() if j != -i), tau)
+                for (i, g), tau in zip(gens_ext, params)}]
 
     def improve(self, m):
         gens_ext = self.G.gens_extended()
@@ -151,12 +152,13 @@ def improve(self, m):
                     )
         for it in range(m):
             if self.m >= self.prec:
-                self.Fnlist = [
+                if len(self.Fnlist) > 0:
+                    self.Fnlist = [
                     {
                         ky: sum((F[ky] for F in self.Fnlist[1:]), self.Fnlist[0][ky])
                         for ky in self.Fnlist[0]
                     }
-                ]
+                    ]
                 return self
             tmp = {}
             for (i, gi), tau in zip(gens_ext, params):
@@ -177,9 +179,6 @@ def improve(self, m):
         ]
         return self
 
-    def improve_one(self):
-        return self.improve(1)
-
     def _repr_(self):
         a, b = self.a, self.b
         if b is None:
@@ -221,23 +220,18 @@ def __call__(self, z, **kwargs):
         return self.evaluate(z, **kwargs)
 
     def evaluate(self, z, **kwargs):
-        # recursive = kwargs.get("recursive", True)
         if not isinstance(z, DivisorsElement):
             z = self.Div([(1, z)])
         if z.degree() != 0:
-            # inf, wd = self.G.to_fundamental_domain(Infinity)
             z -= self.Div([(z.degree(), Infinity)])
         z = self.G.find_equivalent_divisor(z)
-        ans = prod(eval_rat(self.val, P) ** n for P, n in z)
-        newans = prod(
-            F(P) ** n for FF in self.Fnlist for ky, F in FF.items() for P, n in z
-        )
-        ans *= newans
-        return ans
+        ans0 = prod(eval_rat(self.val, P) ** n for P, n in z)
+        ans1 = prod(F(z) for FF in self.Fnlist for F in FF.values())        
+        return ans0 * ans1
 
     def eval_derivative(self, z, return_value=False):
-        assert not isinstance(z, DivisorsElement)
-        # z = self.Div([(1,z)])
+        if isinstance(z, DivisorsElement):
+            raise NotImplementedError
         z0, wd = self.G.to_fundamental_domain(z)
         gens = (
             [None]