sage/period_computation.sage

@@ -1,5 +1,9 @@
 import logging
-from sage.all import sage_eval, PolynomialRing, QQ, SR, log
+from sage.all import sage_eval, PolynomialRing, QQ, SR, log, var
+
+# Logger
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG)
 
 class PFOperator:
     """
@@ -23,7 +27,7 @@ class PFOperator:
     def __init__(self, operator_string: str, no_variables: int = 1):
         self.no_variables = no_variables
         self.operator = self._operator_from_string(operator_string)
-        logging.info("Initialised PFOperator: %s", self.operator)
+        logger.info("Initialised PFOperator: %s", self.operator)
 
 class Period:
     """
@@ -181,21 +185,65 @@ class Period:
             if coefficients is not None:
                 raise ValueError("Provide either coefficients or period_string, not both.")
             self.coefficients = self._period_from_string(period_string)
-            logging.debug("Using string for initialisation.")
+            self.period_string = period_string
+            logger.debug("Using string for initialisation.")
         elif coefficients is None:
             self.coefficients = {}
         else:
             self.coefficients = coefficients
             self.period_string = self._period_to_string()
-            logging.debug("Using coefficients for initialisation.")
+            logger.debug("Using coefficients for initialisation.")
 
         if order is None:
             self.order = self._max_z_degree()
-            logging.debug("Order not provided, using maximum z-degree: %d", self.order)
+            logger.debug("Order not provided, using maximum z-degree: %d", self.order)
         else:
             self.coefficients = self._truncate_coefficients(order)
             self.order = order
             self.period_string = self._period_to_string()
-            logging.debug("Truncated coefficients to order %d.", self.order)
+            logger.debug("Truncated coefficients to order %d.", self.order)
 
-        logging.info("Initialised Period in %d variables at order %d.", self.no_variables, self.order)
+        logger.info("Initialised Period in %d variables at order %d.", self.no_variables, self.order)
+
+
+class PeriodAnsatz(Period):
+    """
+    A class for period Ansätze, characterised by number of variables and their z- and log-multi-degrees.
+    """
+
+    def _multi_indices(self, total_degree: int) -> list:
+        # All no_variables-tuples of non-negative integers whose entries sum to at most total_degree.
+        def helper(remaining_vars, remaining_degree):
+            if remaining_vars == 0:
+                yield ()
+                return
+            for first in range(remaining_degree + 1):
+                for rest in helper(remaining_vars - 1, remaining_degree - first):
+                    yield (first,) + rest
+
+        return list(helper(self.no_variables, total_degree))
+
+    def __init__(self, no_variables: int, z_degree: int, log_degree: int):
+        self.no_variables = no_variables
+        self.z_degree = z_degree
+        self.log_degree = log_degree
+
+        log_indices = self._multi_indices(log_degree)
+        z_indices = self._multi_indices(z_degree)
+
+        # Every (log_index, z_index) pair allowed by the degree bounds gets its own
+        # fresh unknown, to be solved for once a PFOperator is applied to the Ansatz.
+        coefficients = {
+            log_index: {
+                z_index: var("a_" + "_".join(str(x) for x in log_index + z_index))
+                for z_index in z_indices
+            }
+            for log_index in log_indices
+        }
+
+        super().__init__(no_variables=no_variables, coefficients=coefficients, order=z_degree)
+        self.expansion_coefficients = self.coefficients
+        logger.info(
+            "Initialised PeriodAnsatz with %d unknown coefficient(s).",
+            sum(len(z_dict) for z_dict in self.expansion_coefficients.values()),
+        )