Implementing moving origin of ideal

sage/period_computation.sage

@@ -1,6 +1,6 @@
 import copy
 import logging
-from sage.all import sage_eval, PolynomialRing, QQ, SR, log, matrix, prod, var
+from sage.all import sage_eval, PolynomialRing, QQ, SR, function, log, matrix, prod, solve, var
 
 load("sage/util.py")
 
@@ -71,6 +71,139 @@ class PFOperator:
         logger.debug("Simplified PFOperator to: %s", simplified.operator_string)
         return simplified
 
+    def change_coordinates(self, *new_coordinates: str) -> "PFOperator":
+        """
+        Rewrites this operator in new coordinates w0, ..., w{n-1} := new_coordinates(z0, ...,
+        z{n-1}), where new_coordinates is given as n expressions in the *old* coordinates
+        z0, ..., z{n-1}; the inverse change of coordinates needed to do this is solved for
+        automatically (the two aren't independent - one determines the other). The new
+        origin w = 0 is the point z0, ..., z{n-1} at which new_coordinates vanishes.
+
+        E.g. pf_operator.change_coordinates("1 - 5*z0", "1/z1") moves the origin to
+        z0 = 1/5, z1 = infinity.
+        """
+        n = self.no_variables
+        if len(new_coordinates) != n:
+            raise ValueError(
+                "Expected %d new coordinate expression(s), got %d." % (n, len(new_coordinates))
+            )
+
+        z_gens = [var('z%d' % i) for i in range(n)]
+        parse_locals = {str(z): z for z in z_gens}
+        try:
+            new_coords = [SR(sage_eval(expr, locals=parse_locals)) for expr in new_coordinates]
+        except Exception as e:
+            raise ValueError("Invalid coordinate expression: %s" % e)
+
+        w_gens = [var('w%d' % i) for i in range(n)]
+        equations = [w_gens[i] == new_coords[i] for i in range(n)]
+        solutions = solve(equations, z_gens, solution_dict=True)
+        if not solutions:
+            raise ValueError("Could not invert the given change of coordinates %s." % (new_coordinates,))
+        inverse_coords = [solutions[0][z] for z in z_gens]
+
+        Y = function('Y')(*w_gens)
+        y = Y.subs({w_gens[i]: new_coords[i] for i in range(n)})
+
+        def apply_theta(expr, i):
+            return z_gens[i] * expr.diff(z_gens[i])
+
+        total = SR(0)
+        for coeff, monomial in self.operator:
+            exponents = monomial.exponents()[0]
+            z_exponents = exponents[:n]
+            theta_exponents = exponents[n:]
+
+            term = y
+            for i, power in enumerate(theta_exponents):
+                for _ in range(power):
+                    term = apply_theta(term, i)
+
+            z_monomial = prod(z_gens[i] ** z_exponents[i] for i in range(n))
+            total += SR(coeff) * z_monomial * term
+
+        total = total.subs({z_gens[i]: inverse_coords[i] for i in range(n)})
+        total = total.simplify_full().expand()
+
+        total_theta_degree = max(sum(monomial.exponents()[0][n:]) for _, monomial in self.operator)
+        theta_ring = PolynomialRing(QQ, n, ['theta%d' % i for i in range(n)])
+        theta_gens = theta_ring.gens()
+
+        # Peel off the coefficient of each derivative order k, highest total order first (so
+        # a coefficient can never still contain an as-yet-unextracted higher derivative as a
+        # factor), and rewrite w^k*D^k Y as a falling-factorial polynomial in theta applied
+        # to Y - the Euler-operator identity theta*(theta-1)*...*(theta-k+1) = w^k*d^k/dw^k,
+        # taken variable by variable. Each term's coefficient is left as a general rational
+        # function of w for now; only once every term has been collected is the whole
+        # operator scaled by their common denominator (see below).
+        remainder = total
+        terms = []
+        for k in sorted(_multi_indices(n, total_theta_degree), key=lambda k: -sum(k)):
+            if sum(k) == 0:
+                dterm = Y
+            else:
+                args = [arg for i in range(n) if k[i] for arg in (w_gens[i], k[i])]
+                dterm = Y.diff(*args)
+
+            coeff = remainder.coefficient(dterm)
+            if coeff == 0:
+                continue
+            remainder -= coeff * dterm
+
+            falling_factorial = theta_ring(1)
+            for i in range(n):
+                for j in range(k[i]):
+                    falling_factorial *= (theta_gens[i] - j)
+
+            w_monomial = prod(w_gens[i] ** k[i] for i in range(n))
+            terms.append(((coeff / w_monomial).simplify_rational(), falling_factorial))
+
+        remainder = remainder.simplify_full()
+        if remainder != 0:
+            raise ValueError("Residual nonzero after change of coordinates: %s" % remainder)
+        if not terms:
+            raise ValueError("Change of coordinates produced the zero operator.")
+
+        w_ring = PolynomialRing(QQ, n, ['w%d' % i for i in range(n)])
+        common_denominator = w_ring(1)
+        for rational_coeff, _ in terms:
+            common_denominator = common_denominator.lcm(w_ring(rational_coeff.denominator()))
+
+        result_terms = {}
+        for rational_coeff, falling_factorial in terms:
+            scaled = w_ring((rational_coeff * SR(common_denominator)).simplify_rational())
+            for w_coeff, w_monomial in scaled:
+                p = w_monomial.exponents()[0]
+                for theta_coeff, theta_monomial in falling_factorial:
+                    q = theta_monomial.exponents()[0]
+                    key = (p, q)
+                    result_terms[key] = result_terms.get(key, QQ(0)) + w_coeff * theta_coeff
+
+        content = gcd([QQ(coeff) for coeff in result_terms.values() if coeff != 0])
+        if content not in (0, 1):
+            result_terms = {key: coeff / content for key, coeff in result_terms.items()}
+
+        def _format_monomial(coeff, p, q):
+            factors = [str(QQ(coeff))]
+            for i in range(n):
+                if p[i]:
+                    factors.append("z%d^%d" % (i, p[i]) if p[i] != 1 else "z%d" % i)
+            for i in range(n):
+                if q[i]:
+                    factors.append("theta%d^%d" % (i, q[i]) if q[i] != 1 else "theta%d" % i)
+            return "*".join(factors)
+
+        terms_string = [
+            _format_monomial(coeff, p, q) for (p, q), coeff in result_terms.items() if coeff != 0
+        ]
+        if not terms_string:
+            raise ValueError("Change of coordinates produced the zero operator.")
+        operator_string = " + ".join(terms_string)
+
+        new_operator = PFOperator(operator_string, no_variables=n)
+        logger.debug("Changed coordinates of PFOperator to: %s", new_operator.operator_string)
+        return new_operator
+
 class PFOperatorAnsatz(PFOperator):
     """
     A class for Picard-Fuchs operator Ansätze, characterised by number of variables and their z- and theta-multi-degrees.
@@ -122,7 +255,17 @@ class PFIdeal:
     def remove_operator(self, pf_operator: PFOperator):
         self.operators.remove(pf_operator)
         logger.info("Removed PFOperator from PFIdeal: %s", pf_operator.operator_string)
-    
+
+    def change_coordinates(self, *new_coordinates: str) -> "PFIdeal":
+        """
+        Moves the origin of the whole system of differential equations to a new point, by
+        returning a new PFIdeal whose operators are each rewritten in new_coordinates; see
+        PFOperator.change_coordinates for what new_coordinates means.
+        """
+        new_operators = [pf_operator.change_coordinates(*new_coordinates) for pf_operator in self.operators]
+        logger.info("Changed coordinates of PFIdeal with %d operator(s) to %s.", len(new_operators), new_coordinates)
+        return PFIdeal(new_operators)
+
     def find_power_series_solution(self, indicials: list, order: int) -> list:
         """
         Finds power series solutions (no logs) to this PFIdeal at given indicial exponents/order.