This file contains a class Trignomio that allows evaluating and making operations with trignomios, like adding, multiplying, etc A trignomio is similar to a polynomial, but uses trigonometric functions

trignomios.py

"""
This file contains a class Trignomio that allows evaluating and
making operations with trignomios, like adding, multiplying, etc

A trignomio is similar to a polynomial, but uses trignometic functions

Example
-------
>>> trig = Trignomio([1, 2, 3])
>>> print(trig)
1 + 2 * sin(x) + 3 * cos(x)
>>> trig = Trignomio([1, 2, 3, 4])
>>> print(trig)
1 + 2 * sin(x) + 3 * cos(x) + 4 * sin(2*x)
"""

from __future__ import annotations

from typing import Tuple, Union

import numpy as np

Parameter = Union[int, float]
Scalar = Union[int, float]


class Trignomio:
    """
    Defines a trignometric version of a polynomial

    Meaning, if a polynomial can be defined as
    P(x) = a0 + a1 * x + a2 * x^2 + ... + ap * x^p

    We can define a trignometric version:
    T(x) = a0 + a1 * sin(w*x) + a2 * cos(w*x) + ...
           + a_{2p-1} * sin(p*w*x) + a_{2p} * cos(p*w*x)

    Example
    -------
    >>> trig = Trignomio([1, 0, 1])
    >>> print(trig)
    1 + cos(x)
    >>> trig.eval(0)
    2.0
    >>> trig.eval(np.pi/2)
    1.0
    >>> trig = Trignomio([1, 2, 0, -3])
    >>> print(trig)
    1 + 2 * sin(x) - 3*sin(2*x)

    """

    def __init__(self, values: Tuple[Scalar], omega: Scalar = 1):
        values = tuple(values)
        for val in values:
            float(val)
        if not len(values) % 2:
            values = values + (0 * values[0],)
        self.__values = values
        self.__omega = omega

    @property
    def pmax(self) -> int:
        """
        Returns the maximum harmonic value of the trignomio

        Example
        -------
        >>> trig = Trignomio([1, 0, 1])  # 1 + cos(x)
        >>> trig.pmax
        1
        >>> trig = Trignomio([1, 0, 1, 2])  # 1 + cos(x) + sin(2*x)
        >>> trig.pmax
        2
        """
        return len(self.__values) // 2

    @property
    def omega(self) -> Scalar:
        """
        Returns the harmonic of the trignomio

        Example
        -------
        >>> trig = Trignomio([1, 0, 1])  # 1 + cos(x)
        >>> trig.omega
        1
        >>> trig = Trignomio([1, 0, 1], pi)  # 1 + cos(pi*x)
        >>> trig.omega
        3.14159265
        """
        return self.__omega

    def __iter__(self):
        yield from self.__values

    def __getitem__(self, index):
        return self.__values[index]

    def __neg__(self) -> Trignomio:
        return self.__class__(tuple(-coef for coef in self), self.omega)

    def __add__(self, other: Union[Scalar, Trignomio]):
        if not isinstance(other, Trignomio):
            coefs = list(self)
            coefs[0] += other
            return self.__class__(coefs, self.omega)
        values = [0] * (1 + 2 * max(self.pmax, other.pmax))
        for i, val in enumerate(self):
            values[i] += val
        for i, val in enumerate(other):
            values[i] += val
        return self.__class__(values, self.__omega)

    def __sub__(self, other: Union[Scalar, Trignomio]) -> Trignomio:
        return self.__add__(-other)

    def __mul__(self, other: Union[Scalar, Trignomio]) -> Trignomio:
        if not isinstance(other, Trignomio):
            coefs = tuple(other * coef for coef in self)
            return self.__class__(coefs, self.omega)
        values = [0] * (1 + 2 * (self.pmax + other.pmax))
        for i, vali in enumerate(self):
            p = (i + 1) // 2
            for j, valj in enumerate(other):
                q = (j + 1) // 2
                s, d = p + q, abs(p - q)
                temp = vali * valj / 2
                if (i % 2) ^ (j % 2):
                    if p != q:
                        sind = -1 if (i > j) ^ (i % 2) else 1
                        values[2 * d - 1] += temp * sind
                    values[2 * s - 1] += temp
                else:
                    coss = -1 if i % 2 else 1
                    values[2 * d] += temp
                    values[2 * s] += temp * coss
        return self.__class__(values, self.omega)

    def __truediv__(self, other: Scalar) -> Trignomio:
        coefs = tuple(coef / other for coef in self)
        return self.__class__(coefs, self.omega)

    def __rsub__(self, other: Scalar) -> Trignomio:
        return (-self).__add__(other)

    def __rmul__(self, other: Scalar) -> Trignomio:
        return self.__mul__(other)

    def __radd__(self, other: Scalar) -> Trignomio:
        return self.__add__(other)

    def eval(self, node: Parameter, derivate: int = 0):
        """
        Evaluates the trignomio at given node

        Example
        -------
        >>> trig = Trignomio([1, 2])
        >>> print(trig)
        1 + 2*sin(x)
        >>> trig.eval(0)
        1.0
        >>> trig.eval(pi/2)
        3.0
        >>> trig.eval(0, 0)
        2.0
        >>> trig.eval(pi/2, 2)
        -2.0
        """
        if derivate != 0:
            raise NotImplementedError
        w = self.omega
        results = 0
        for i, value in enumerate(self):
            p = (i + 1) // 2
            base = np.sin(p * w * node) if i % 2 else np.cos(p * w * node)
            results += value * base
        return results

    def shift(self, amount: Parameter) -> Trignomio:
        """
        Evaluates the trignomio at given node

        Example
        -------
        >>> old_trig = Trignomio([1, 2])
        >>> print(old_trig)
        1 + 2*sin(x)
        >>> new_trig = old_trig.shift(np.pi/2)
        >>> print(new_trig)
        1 - 2*cos(x)
        """
        newcoefs = [0] * (1 + 2 * self.pmax)
        newcoefs[0] += self[0]
        for p in range(1, self.pmax + 1):
            sinpwa = np.sin(p * self.omega * amount)
            cospwa = np.cos(p * self.omega * amount)
            cofs, cofc = self[2 * p - 1], self[2 * p]
            newcoefs[2 * p - 1] = cofs * cospwa + cofc * sinpwa
            newcoefs[2 * p] = -cofs * sinpwa + cofc * cospwa
        return self.__class__(newcoefs, self.omega)

    def __call__(self, nodes):
        return self.eval(nodes, 0)

    def __str__(self) -> str:
        msgs = []
        flag = False
        for i, coef in enumerate(self):
            if coef == 0:
                continue
            if coef < 0:
                msg = "- "
            elif flag:
                msg = "+ "
            else:
                msg = ""
            flag = True
            p = (i + 1) // 2
            coef = abs(coef)
            if coef != 1 or i == 0:
                msg += str(coef)
            if i > 0:
                if coef != 1:
                    msg += " * "
                msg += "sin(" if i % 2 else "cos("
                if p > 1:
                    msg += f"{p}*"
                msg += "w*t)"
            msgs.append(msg)
        return " ".join(msgs)

    def __repr__(self) -> str:
        return str(self)


def test_random_trignomials():
    """
    Function to test if the trignomials coefficients
    are correctly computed
    """
    ntests = 100
    maxpmax = 6
    tsample = np.linspace(-1, 1, 17)
    for _ in range(ntests):
        dega, degb = np.random.randint(0, maxpmax + 1, 2)
        coefsa = np.random.uniform(-1, 1, dega + 1)
        coefsb = np.random.uniform(-1, 1, degb + 1)
        polya = Trignomio(coefsa)
        polyb = Trignomio(coefsb)
        polyc = polya + polyb
        polyd = polya * polyb
        polye = polya.shift(1)
        valuesa = polya(tsample)
        valuesb = polyb(tsample)
        valuesc = polyc(tsample)
        valuesd = polyd(tsample)
        valuese = polye(1 + tsample)

        np.testing.assert_allclose(valuesa + valuesb, valuesc)
        np.testing.assert_allclose(valuesa * valuesb, valuesd)
        np.testing.assert_allclose(valuese, valuesa)


if __name__ == "__main__":
    test_random_trignomials()