8 files changed, 638 insertions, 0 deletions
diff --git a/mdpoly/__init__.py b/mdpoly/__init__.py
new file mode 100644
index 0000000..b08f191
--- /dev/null
+++ b/mdpoly/__init__.py
@@ -0,0 +1,55 @@
+"""
+A Library to represent multivariate polynomials
+
+"""
+# internal classes imported with underscore because they should not be exposed
+# to the end users
+from .abc import (Shape as _Shape)
+
+from .algebra import (PolyRingAlgebra as _PolyRingAlgebra,
+                      MatrixAlgebra as  _MatrixAlgebra)
+
+from .leaves import (Const as _Const,
+                     Var as _Var,
+                     Param as _Param)
+
+from typing import Self, Sequence
+
+
+Shape = _Shape
+
+
+class _FromHelpers:
+    @classmethod
+    def from_names(cls, comma_separated_names: str, strip: bool =True) -> Sequence[Self]:
+        """ Generate scalar variables from comma separated list of names """
+        names = comma_separated_names.split(",")
+        if strip:
+            names = map(str.strip, names)
+        yield from map(cls, names)
+
+
+class Constant(_Const, _PolyRingAlgebra, _FromHelpers):
+    """ Constant values """
+
+
+class Variable(_Var, _PolyRingAlgebra, _FromHelpers):
+    """ Polynomial Variable """
+
+
+class Parameter(_Param, _PolyRingAlgebra):
+    """ Parameter that can be substituted """
+
+
+class MatrixConstant(_Const, _MatrixAlgebra):
+    """ Matrix constant """
+    
+
+class MatrixVariable(_Var, _MatrixAlgebra):
+    """ Matrix Polynomial Variable """
+
+
+class MatrixParameter(_Param, _PolyRingAlgebra):
+    """ Matrix Parameter """
+
+
diff --git a/mdpoly/abc.py b/mdpoly/abc.py
new file mode 100644
index 0000000..53b16a5
--- /dev/null
+++ b/mdpoly/abc.py
@@ -0,0 +1,75 @@
+""" Abstract Base Classes of MDPoly
+
+Defines:
+  - Tree data structure to represent expressions
+  - Representations of polynomials
+"""
+from .types import Number, Shape, MatrixIndex, PolyIndex
+
+from typing import Self, Sequence, Protocol, runtime_checkable
+from abc import abstractmethod
+
+
+@runtime_checkable
+class Leaf(Protocol):
+    """
+    Leaf of the binary tree.
+    """
+    name: str
+    shape: Shape
+
+
+@runtime_checkable
+class Expr(Protocol):
+    """
+    Binary tree to represent a mathematical expression.
+    """
+    left: Self | Leaf
+    right: Self | Leaf
+    shape: Shape
+
+    def __repr__(self):
+        name = self.__class__.__qualname__
+        return f"{name}(left={self.left}, right={self.right})"
+
+    def leaves(self) -> Sequence[Leaf]:
+        if isinstance(self.left, Leaf):
+            yield self.left
+        else:
+            yield from self.left.leaves()
+
+        if isinstance(self.right, Leaf):
+            yield self.right
+        else:
+            yield from self.right.leaves()
+    
+    def __iter__(self) -> Sequence[Self | Leaf]:
+        return (self.left, self.right)
+
+
+class Repr(Protocol):
+    """
+    Representation of a polynomial expression.
+    """
+    @abstractmethod
+    def at(self, entry: MatrixIndex, term: PolyIndex) -> Number:
+        """ Access polynomial entry """
+
+    @abstractmethod
+    def set(self, entry: MatrixIndex, term: PolyIndex, value: Number) -> None:
+        """ Set value of polynomial entry """
+
+    @abstractmethod
+    def entries(self) -> Sequence[MatrixIndex]:
+        """ Return indices to non-zero entries of the matrix """
+
+    @abstractmethod
+    def terms(self, entry: MatrixIndex) -> Sequence[PolyIndex]:
+        """ Return indices to non-zero terms in the polynomial at the given
+        matrix entry """
+
+    def __iter__(self) -> Sequence[tuple[MatrixIndex, PolyIndex, Number]]:
+        """ Iterate over non-zero entries of the representations """
+        for entry in self.entries():
+            for term in self.terms(entry):
+                yield entry, term, self.at(entry, term)
diff --git a/mdpoly/algebra.py b/mdpoly/algebra.py
new file mode 100644
index 0000000..8a8e961
--- /dev/null
+++ b/mdpoly/algebra.py
@@ -0,0 +1,235 @@
+"""
+Algebraic Structures for Expressions
+
+"""
+from .abc import Expr
+from .leaves import Nothing, Const, Param
+from .errors import AlgebraicError
+# from .state import State
+# from .representations import HasSparseRepr, SparseRepr
+
+from typing import Protocol, TypeVar, Any, runtime_checkable
+from functools import wraps
+from enum import Enum
+from abc import abstractmethod
+
+
+def binary_operator(cls: Expr) -> Expr:
+    __init_cls__ = cls.__init__
+    @wraps(cls.__init__)
+    def __init_binary__(self, left, right, *args, **kwargs):
+        __init_cls__(self, *args, **kwargs)
+        self.left, self.right = left, right
+
+    cls.__init__ = __init_binary__
+    return cls
+
+
+def unary_operator(cls: Expr) -> Expr:
+    __init_cls__ = cls.__init__
+    @wraps(cls.__init__)
+    def __init_unary__(self, left, *args, **kwargs):
+        __init_cls__(self, *args, **kwargs)
+        self.left, self.right = left, Nothing
+
+    cls.__init__ = __init_unary__
+    return cls
+
+
+T = TypeVar("T")
+
+@runtime_checkable
+class AlgebraicStructure(Protocol):
+    """
+    Provides method to enforce algebraic closure of operations
+    """
+    class Algebra(Enum):
+        """ Types of algebras """
+        poly_ring = 0
+        matrix_ring = 1
+
+    _algebra: Algebra
+    _parameter: type
+    _constant: type
+
+    @classmethod
+    def _is_constant(cls: T, other: T) -> bool:
+        return isinstance(other, cls._constant)
+
+    @classmethod
+    def _is_paramter(cls: T, other: T) -> bool:
+        return isinstance(other, cls._parameter)
+
+    @classmethod
+    def _is_const_or_param(cls: T, other: T) -> bool:
+        return cls._is_constant(other) or cls._is_parameter(other)
+
+    @classmethod
+    def _assert_same_algebra(cls: T, other: T) -> None:
+        if not cls._algebra == other._algebra:
+            if not cls._is_constant(other):
+                raise AlgebraicError("Cannot perform operation between types from "
+                                     f"different algebraic structures {cls} ({cls._algebra}) "
+                                     f"and {type(other)} ({other._algebra})")
+
+    @classmethod
+    def _wrap_if_constant(cls: T, other: Any):
+        if isinstance(other, AlgebraicStructure):
+            cls._assert_same_algebra(other)
+            return other
+
+        if isinstance(other, cls._constant):
+            return other
+
+        return cls._constant(name="", value=other)
+
+    # FIXME: add shape rules in algebra
+    # @classmethod
+    # def _assert_same_shape(cls: T, other: T) -> None:
+    #     if not cls.shape == other.shape:
+    #         raise AlgebraicError("Cannot perform algebraic operations between "
+    #                              f"objects with different shapes {cls.shape} and {other.shape}")
+
+
+
+class ReducibleExpr(Protocol):
+    @abstractmethod
+    def reduce(self) -> Expr: ...
+
+
+# ┏━┓┏━┓╻  ╻ ╻┏┓╻┏━┓┏┳┓╻┏━┓╻     ┏━┓╻  ┏━╸┏━╸┏┓ ┏━┓┏━┓
+# ┣━┛┃ ┃┃  ┗┳┛┃┗┫┃ ┃┃┃┃┃┣━┫┃     ┣━┫┃  ┃╺┓┣╸ ┣┻┓┣┳┛┣━┫
+# ╹  ┗━┛┗━╸ ╹ ╹ ╹┗━┛╹ ╹╹╹ ╹┗━╸   ╹ ╹┗━╸┗━┛┗━╸┗━┛╹┗╸╹ ╹
+
+class PolyRingAlgebra(AlgebraicStructure, Protocol):
+    """
+    Endows with the algebraic structure of a polynomial ring.
+    """
+    _algebra = AlgebraicStructure.Algebra.poly_ring
+    _parameter = Param
+    _constant = Const
+
+    def __add__(self, other):
+        other = self._wrap_if_constant(other)
+        return Add(self, other)
+
+    def __sub__(self, other):
+        other = self._wrap_if_constant(other)
+        return Subtract(self, other)
+
+    def __neg__(self, other):
+        return Multiply(self._constant(-1), other)
+
+    def __mul__(self, other):
+        other = self._wrap_if_constant(other)
+        return Multiply(self, other)
+
+    def __rmul__(self, other):
+        other = self._wrap_if_constant(other)
+        return Multiply(other, self)
+
+    def __truediv__(self, other):
+        other = self._wrap_if_constant(other)
+        if not self._is_const_or_param(other):
+            raise AlgebraicError("Cannot divide by variables in polynomial ring.")
+        return Multiply(Const(value=1/other.value), self)
+
+    def __pow__(self, other):
+        other = self._wrap_if_constant(other)
+        if not self._is_const_or_param(other):
+            raise AlgebraicError(f"Cannot raise to powers of type {type(other)} in "
+                                 "polynomial ring. Only constants and parameters are allowed.")
+        return Exponentiate(self, other)
+
+
+@binary_operator
+class Add(Expr, PolyRingAlgebra):
+    """ Sum operator """
+
+
+@binary_operator
+class Subtract(Expr, ReducibleExpr, PolyRingAlgebra):
+    """ Subtraction operation """
+    def reduce(self) -> Expr:
+        return Add(self.left, Multiply(self._constant(value=-1), self.right))
+
+
+@binary_operator
+class Multiply(Expr, PolyRingAlgebra):
+    """ Multiplication operator """
+
+
+@binary_operator
+class Exponentiate(Expr, PolyRingAlgebra):
+    """ Exponentiation operator """
+
+
+@unary_operator
+class Diff(Expr, PolyRingAlgebra):
+    """ Total differentiation """
+
+
+# ┏┳┓┏━┓╺┳╸┏━┓╻╻ ╻   ┏━┓╻  ┏━╸┏━╸┏┓ ┏━┓┏━┓
+# ┃┃┃┣━┫ ┃ ┣┳┛┃┏╋┛   ┣━┫┃  ┃╺┓┣╸ ┣┻┓┣┳┛┣━┫
+# ╹ ╹╹ ╹ ╹ ╹┗╸╹╹ ╹   ╹ ╹┗━╸┗━┛┗━╸┗━┛╹┗╸╹ ╹
+
+class MatrixAlgebra(AlgebraicStructure, Protocol):
+    """
+    Endows with the algebraic structure of a matrix ring.
+    """
+    _algebra = AlgebraicStructure.Algebra.matrix_ring
+    # FIXME: consider MatParam or something like that?
+    _parameter = Param
+    _constant = Const
+
+    def __add__(self, other):
+        other = self._wrap_if_constant(other)
+        return MatAdd(self, other)
+
+    def __sub___(self, other):
+        other = self._wrap_if_constant(other)
+        return MatAdd(self, ScalarMul(Const(value=-1), other))
+
+    def __mul__(self, scalar):
+        scalar = self._wrap_if_constant(scalar)
+        return ScalarMul(scalar, self) 
+
+    def __rmul__(self, scalar):
+        scalar = self._wrap_if_constant(scalar)
+        return ScalarMul(scalar, self)
+
+    def __matmul__(self, other):
+        other = self._wrap_if_constant(other)
+        return MatMul(self, other)
+
+    def __rmatmul(self, other):
+        other = self._wrap_if_constant(other)
+        return MatMul(other, self)
+
+    def __truediv__(self, scalar):
+        scalar = self._wrap_if_constant(scalar)
+        raise NotImplementedError
+
+    def transpose(self):
+        raise NotImplementedError
+
+    # Shorthands & syntactic sugar
+
+    @property
+    def T(self):
+        return self.transpose()
+
+
+@binary_operator
+class MatAdd(Expr, MatrixAlgebra):
+    """ Matrix Addition """
+
+
+@binary_operator
+class ScalarMul(Expr, MatrixAlgebra):
+    """ Matrix-Scalar Multiplication """
+
+
+@binary_operator
+class MatMul(Expr, MatrixAlgebra):
+    """ Matrix Multiplication """
diff --git a/mdpoly/errors.py b/mdpoly/errors.py
new file mode 100644
index 0000000..65290b3
--- /dev/null
+++ b/mdpoly/errors.py
@@ -0,0 +1,17 @@
+
+class AlgebraicError(Exception):
+    """
+    This is raised whenever an invalid mathematical operation is performed.
+    See more in mdpoly/ops.py
+    """
+
+class InvalidShape(Exception):
+    """
+    This is raised whenever there is an invalid shape
+    """
+
+class IndexError(Exception):
+    """
+    This is raised whenever there is an invalid index
+    """
+
diff --git a/mdpoly/leaves.py b/mdpoly/leaves.py
new file mode 100644
index 0000000..4dd487d
--- /dev/null
+++ b/mdpoly/leaves.py
@@ -0,0 +1,41 @@
+from .abc import Leaf
+from .types import Number, Shape
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Nothing(Leaf):
+    """
+    A leaf that is nothing. This is a placeholder (eg. for unary operators).
+    """
+    name: str = "<nothing>"
+    shape: Shape = Shape(0, 0)
+
+
+@dataclass(frozen=True)
+class Const(Leaf):
+    """
+    A constant
+    """
+    value: Number
+    name: str = ""
+    shape: Shape = Shape.scalar()
+
+
+@dataclass(frozen=True)
+class Var(Leaf):
+    """
+    Polynomial variable
+    """
+    name: str
+    shape: Shape = Shape.scalar()
+
+
+@dataclass(frozen=True)
+class Param(Leaf):
+    """
+    A parameter
+    """
+    name: str
+    shape: Shape = Shape.scalar()
diff --git a/mdpoly/representations.py b/mdpoly/representations.py
new file mode 100644
index 0000000..de03838
--- /dev/null
+++ b/mdpoly/representations.py
@@ -0,0 +1,121 @@
+from .abc import Repr 
+from .types import Number, Shape, MatrixIndex, PolyIndex
+
+from typing import Protocol, Sequence
+from abc import abstractmethod
+
+import numpy as np
+import numpy.typing as npt
+
+
+class SparseRepr(Repr):
+    """
+    Sparse representation of polynomial
+    """
+    data: dict[MatrixIndex, dict[PolyIndex, Number]]
+
+    def __init__(self):
+        self.data = {}
+
+    def at(self, entry: MatrixIndex, term: PolyIndex) -> Number:
+        """ Access polynomial entry """
+        if entry in self.data.keys():
+            if term in self.data[entry].keys():
+                return self.data[entry][term]
+        
+        return 0
+
+    def set(self, entry: MatrixIndex, term: PolyIndex, value: Number) -> None:
+        """ Set value of polynomial entry """
+        if entry not in self.data.keys():
+            self.data[entry] = {}
+
+        self.data[entry][term] = value
+
+    def entries(self) -> Sequence[MatrixIndex]:
+        """ Return indices to non-zero entries of the matrix """
+        yield from self.data.keys()
+
+    def terms(self, entry: MatrixIndex) -> Sequence[PolyIndex]:
+        """ Return indices to non-zero terms in the polynomial at the given
+        matrix entry """
+        if entry in self.data.keys():
+            yield from self.data[entry]
+
+        return []
+
+
+class HasSparseRepr(Protocol):
+    @abstractmethod
+    def to_sparse(self) -> SparseRepr: ...
+
+
+class SparseMatrixRepr(Repr):
+    """
+    Sparse matrix representation of a polynomial.
+
+    Collect the monomials in a vector x (the "basis") and store a matrix Q and
+    such that the quadratic form x^T Q x gives the polynomial.
+    """
+
+    def at(self, entry: MatrixIndex, term: PolyIndex) -> Number:
+        """ Access polynomial entry """
+
+    def set(self, entry: MatrixIndex, term: PolyIndex, value: Number) -> None:
+        """ Set value of polynomial entry """
+
+    def entries(self) -> Sequence[MatrixIndex]:
+        """ Return indices to non-zero entries of the matrix """
+
+    def terms(self, entry: MatrixIndex) -> Sequence[PolyIndex]:
+        """ Return indices to non-zero terms in the polynomial at the given
+        matrix entry """
+
+    def basis(self):
+        raise NotImplementedError
+
+    def matrix(self):
+        raise NotImplementedError
+
+
+class HasSparseMatrixRepr(Protocol):
+    @abstractmethod
+    def to_sparse_matrix(self) -> SparseMatrixRepr: ...
+
+
+class DenseRepr(Repr):
+    """
+    Dense representation of a polynomial
+    """
+    data: npt.NDArray[Number]
+
+    def __init__(self, shape: Shape, dtype:type =float):
+        self.data = np.zeros(shape, dtype)
+
+    def at(self, entry: MatrixIndex, term: PolyIndex) -> Number:
+        """ Access polynomial entry """
+        raise NotImplementedError
+
+    def set(self, entry: MatrixIndex, term: PolyIndex, value: Number) -> None:
+        """ Set value of polynomial entry """
+        raise NotImplementedError
+
+    def entries(self) -> Sequence[MatrixIndex]:
+        """ Return indices to non-zero entries of the matrix """
+        raise NotImplementedError
+
+    def terms(self, entry: MatrixIndex) -> Sequence[PolyIndex]:
+        """ Return indices to non-zero terms in the polynomial at the given
+        matrix entry """
+        raise NotImplementedError
+
+
+class HasDenseRepr(Protocol):
+    @abstractmethod
+    def to_dense(self) -> DenseRepr: ...
+
+
+class DenseMatrixRepr(Repr):
+    """
+    Dense matrix representation of a polynomials
+    """
diff --git a/mdpoly/state.py b/mdpoly/state.py
new file mode 100644
index 0000000..56cd7e0
--- /dev/null
+++ b/mdpoly/state.py
@@ -0,0 +1,32 @@
+from .types import Number
+from .leaves import Var, Param
+
+from typing import NewType
+
+Index = NewType('Index', int)
+
+class State:
+    variables: dict[Var, Index]
+    parameters: dict[Param, Number]
+
+    def __init__(self, variables, parameters={}):
+        self._last_index = 0
+        self.variables = variables
+        self.parameters = parameters
+
+    def _make_index(self) -> Index:
+        """ Make a new index """
+        self._last_index += 1
+        return self._last_index
+
+    def index(self, var: Var) -> Index:
+        """ Get the index for a variable """
+        if var not in self.variables.keys():
+            new_index = self._make_index()
+            self.variables[var] = new_index
+            return new_index
+
+        return self.variables[var]
+        
+
+
diff --git a/mdpoly/types.py b/mdpoly/types.py
new file mode 100644
index 0000000..14847fa
--- /dev/null
+++ b/mdpoly/types.py
@@ -0,0 +1,62 @@
+from .errors import InvalidShape
+
+from typing import Self, NamedTuple, Iterable
+
+Number = int | float
+
+
+def filtertype(t: type, seq: Iterable):
+    yield from filter(lambda x: isinstance(x, t), seq)
+
+
+class Shape(NamedTuple):
+    """
+    Describes the shape of a mathematical object.
+    """
+    rows: int
+    cols: int
+
+    @classmethod
+    def infer(cls) -> Self:
+        return cls(-1, -1)
+
+    @classmethod
+    def scalar(cls) -> Self:
+        return cls(1, 1)
+    
+    @classmethod
+    def row(cls, n: int) -> Self:
+        if n <= 0:
+            raise InvalidShape("Row vector must have dimension >= 1")
+        return cls(1, n)
+
+    @classmethod
+    def column(cls, n: int) -> Self:
+        if n <= 0:
+            raise InvalidShape("Column vector must have dimension >= 1")
+        return cls(n, 1)
+
+
+class MatrixIndex(NamedTuple):
+    """
+    Tuple to index an element of a matrix or vector.
+    """
+    row: int
+    col: int
+
+    @classmethod
+    @property
+    def infer(cls, row=-1, col=-1) -> Self:
+        return cls(row, col)
+
+
+class PolyTermIndex(NamedTuple):
+    term: int
+    power: Number
+
+
+class PolyIndex:
+    """
+    Tuple to index a coefficient of a polynomial.
+    """
+    indices: tuple[PolyTermIndex]