polygen.py

import numpy as np
from numpy.polynomial import Polynomial
from dist import Dist, ConstDist

class Poly:
    
    def __init__(self, degree, dist):
        self.degree = degree
        if isinstance(dist, list):
            if len(dist) != degree + 1:
                raise ValueError("If a list of distributions is given, it must have length degree+1")
            self.dists = dist
        elif isinstance(dist, Dist):
            self.dists = [dist] * (degree + 1)
        else:
            raise ValueError("dist must be a Dist instance or list of Dist instances")

    def sample_coeffs(self):
        """Return an array of sampled coefficients (real or complex)."""
        coeffs = np.array([d.sample()[0] for d in self.dists])
        if np.iscomplexobj(coeffs):
            coeffs = coeffs.astype(np.complex128)
        else:
            coeffs = coeffs.astype(np.float64)
        return coeffs

    def sample(self):
        """Sample a random polynomial with coefficients from the specified distributions."""
        coeffs = self.sample_coeffs()
        return Polynomial(coeffs)

    def __add__(self, other):
        if not isinstance(other, Poly):
            raise TypeError("Polynomial addition only supported between Poly objects.")
        deg = max(self.degree, other.degree)
        new_dists = []
        for i in range(deg + 1):
            a = self.dists[i] if i <= self.degree else ConstDist(0)
            b = other.dists[i] if i <= other.degree else ConstDist(0)
            new_dists.append(a + b)
        return Poly(deg, new_dists)

    def __sub__(self, other):
        if not isinstance(other, Poly):
            raise TypeError("Polynomial subtraction only supported between Poly objects.")
        deg = max(self.degree, other.degree)
        new_dists = []
        for i in range(deg + 1):
            a = self.dists[i] if i <= self.degree else ConstDist(0)
            b = other.dists[i] if i <= other.degree else ConstDist(0)
            new_dists.append(a - b)
        return Poly(deg, new_dists)

    def __mul__(self, other):
        if not isinstance(other, Poly):
            raise TypeError("Polynomial multiplication only supported between Poly objects.")
        deg = self.degree + other.degree
        new_dists = [ConstDist(0)] * (deg + 1)
        for i in range(self.degree + 1):
            for j in range(other.degree + 1):
                new_dists[i + j] = new_dists[i + j] + (self.dists[i] * other.dists[j])
        return Poly(deg, new_dists)

    def compose(self, other):
        """Composition P(Q(x)) as a random polynomial."""
        if not isinstance(other, Poly):
            raise TypeError("Polynomial composition only supported between Poly objects.")
        # Initialize as zero polynomial
        result = Poly(0, ConstDist(0))
        base = Poly(0, ConstDist(1))  # Q(x)^0
        for coeff in self.dists:
            result = result + (Poly(other.degree * base.degree, [coeff]) * base)
            base = base * other
        return result

    def __repr__(self):
        return f"Poly(degree={self.degree}, dists={self.dists})"