Make new constructors Variable.from_powers and Variable.make_combinations

1 files changed, 59 insertions, 14 deletions

diff --git a/mdpoly/algebra.py b/mdpoly/algebra.py
index 580c2e2..b98af01 100644
--- a/mdpoly/algebra.py
+++ b/mdpoly/algebra.py
@@ -7,15 +7,17 @@ from .errors import AlgebraicError, InvalidShape
 from .state import State
 from .index import Shape, MatrixIndex, PolyIndex, Number
 
-from typing import cast
+from typing import cast, Iterable
 from functools import reduce
-from itertools import product
+from itertools import product, chain, combinations_with_replacement
+from math import prod
 from abc import abstractmethod
 
 import operator
 
 
 class BinaryOp(Expr):
+    """ Binary Operator """
     def __init__(self, left, right):
         self._left = left
         self._right = right
@@ -38,6 +40,7 @@ class BinaryOp(Expr):
 
 
 class UnaryOp(Expr):
+    """ Unary Operator """
     def __init__(self, left, right=None):
         self._inner = left
 
@@ -82,10 +85,20 @@ class Reducible(Expr):
         """ Reduce the expression to its basic elements """
 
 
+# TODO: review this idea before implementing
+# class Simplifiable(Expr):
+#     """ Simplifiable Expression. """
+# 
+#     @abstractmethod
+#     def simplify(self) -> Expr:
+#         """ Simplify the expression """
+
+
 # ┏━┓┏━┓╻  ╻ ╻┏┓╻┏━┓┏┳┓╻┏━┓╻     ┏━┓╻  ┏━╸┏━╸┏┓ ┏━┓┏━┓
 # ┣━┛┃ ┃┃  ┗┳┛┃┗┫┃ ┃┃┃┃┃┣━┫┃     ┣━┫┃  ┃╺┓┣╸ ┣┻┓┣┳┛┣━┫
 # ╹  ┗━┛┗━╸ ╹ ╹ ╹┗━┛╹ ╹╹╹ ╹┗━╸   ╹ ╹┗━╸┗━┛┗━╸┗━┛╹┗╸╹ ╹
 
+
 class PolyRingExpr(Expr):
     r""" Endows with the algebraic structure of a polynomial ring.
 
@@ -93,6 +106,8 @@ class PolyRingExpr(Expr):
     polynomials are scalars.
     """
 
+    # -- Properties ---
+
     @property
     def algebra(self) -> Algebra:
         return Algebra.poly_ring
@@ -105,24 +120,46 @@ class PolyRingExpr(Expr):
                                 f"shapes {self.left.shape} and {self.right.shape}.")
         return self.left.shape
 
+    # --- Helper methods for construction ---
+
+    @classmethod
+    def from_powers(cls, variable: PolyRingExpr, exponents: Iterable[int]) -> PolyRingExpr:
+        """ Given a variable, say :math:`x`, and a list of exponents ``[0, 3, 5]``
+        returns a polynomial :math:`x^5 + x^3 + 1`.
+        """
+        nonzero_exponents = filter(lambda e: e != 0, exponents)
+        constant_term = PolyConst(1 if 0 in exponents else 0)
+        # FIXME: remove annoying 0
+        return sum((variable ** e for e in nonzero_exponents), constant_term)
+
+
+    @classmethod
+    def make_combinations(cls, variables: Iterable[Var], max_degree: int) -> Iterable[PolyRingExpr]:
+        """ Make a list of """
+        variables = chain(variables, (PolyConst(1),))
+        for comb in combinations_with_replacement(variables, max_degree):
+            yield prod(comb)
+
+    # --- Operator Overloading ---
+
     def __add__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolyAdd(self, other)
 
     def __radd__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolyAdd(other, self)
 
     def __sub__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolySub(self, other)
 
     def __rsub__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolyAdd(other, self)
 
     def __neg__(self):
@@ -130,12 +167,12 @@ class PolyRingExpr(Expr):
 
     def __mul__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolyMul(self, other)
 
     def __rmul__(self, other):
         self._assert_same_algebra(self, other)
-        other = self._wrap(Number, Const, other)
+        other = self._wrap(Number, PolyConst, other)
         return PolyMul(other, self)
 
     def __matmul__(self, other):
@@ -145,8 +182,8 @@ class PolyRingExpr(Expr):
         self.__rmatmul__(other)
 
     def __truediv__(self, other):
-        other = self._wrap(Number, Const, other)
-        if not isinstance(other, Const | Param):
+        other = self._wrap(Number, PolyConst, other)
+        if not isinstance(other, PolyConst | PolyParam):
             raise AlgebraicError("Cannot divide by variables in polynomial ring.")
 
         return PolyMul(self, other)
@@ -155,13 +192,18 @@ class PolyRingExpr(Expr):
         raise AlgebraicError("Cannot perform right division in polynomial ring.")
 
     def __pow__(self, other):
-        other = self._wrap(Number, Const, other)
-        if not isinstance(other, Const | Param):
+        other = self._wrap(Number, PolyConst, other)
+        if not isinstance(other, PolyConst | PolyParam):
              raise AlgebraicError(f"Cannot raise to powers of type {type(other)} in "
                                   "polynomial ring. Only constants and parameters are allowed.")
         return PolyExp(self, other)
 
 
+class PolyConst(Const, PolyRingExpr): ...
+class PolyVar(Const, PolyRingExpr): ...
+class PolyParam(Param, PolyRingExpr): ...
+
+
 class PolyAdd(BinaryOp, PolyRingExpr):
     """ Addition operator between scalar polynomials. """
 
@@ -322,6 +364,9 @@ class MatrixExpr(Expr):
     def algebra(self) -> Algebra:
         return Algebra.matrix_ring
 
+    def at(self, entry: MatrixIndex, scalar_type: type) -> Expr:
+        raise NotImplementedError
+
     def __add__(self, other):
         self._assert_same_algebra(self, other)
         other = self._wrap(Number, MatConst, other)
@@ -361,12 +406,12 @@ class MatrixExpr(Expr):
         scalar = self._wrap_if_constant(scalar)
         raise NotImplementedError
 
-    def transpose(self):
+    def transpose(self) -> MatrixExpr:
         """ Matrix transposition. """
         return MatTranspose(self)
 
     @property
-    def T(self):
+    def T(self) -> MatrixExpr:
         """ Shorthand for :py:meth:`mdpoly.algebra.MatrixAlgebra.transpose`. """
         return self.transpose()