core_operations.py

"""@package docstring
B-ASIC Core Operations Module.
TODO: More info.
"""

from numbers import Number
from typing import Optional
from numpy import conjugate, sqrt, abs as np_abs

from b_asic.port import SignalSourceProvider, InputPort, OutputPort
from b_asic.operation import AbstractOperation
from b_asic.graph_component import Name, TypeName


class Constant(AbstractOperation):
    """Constant value operation.
    TODO: More info.
    """

    def __init__(self, value: Number = 0, name: Name = ""):
        super().__init__(input_count = 0, output_count = 1, name = name)
        self.set_param("value", value)

    @property
    def type_name(self) -> TypeName:
        return "c"

    def evaluate(self):
        return self.param("value")
    
    @property
    def value(self) -> Number:
        """TODO: docstring"""
        return self.param("value")

    @value.setter
    def value(self, value: Number):
        """TODO: docstring"""
        return self.set_param("value", value)


class Addition(AbstractOperation):
    """Binary addition operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "add"

    def evaluate(self, a, b):
        return a + b


class Subtraction(AbstractOperation):
    """Binary subtraction operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "sub"

    def evaluate(self, a, b):
        return a - b


class Multiplication(AbstractOperation):
    """Binary multiplication operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "mul"

    def evaluate(self, a, b):
        return a * b


class Division(AbstractOperation):
    """Binary division operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "div"

    def evaluate(self, a, b):
        return a / b


class SquareRoot(AbstractOperation):
    """Unary square root operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])

    @property
    def type_name(self) -> TypeName:
        return "sqrt"

    def evaluate(self, a):
        return sqrt(complex(a))


class ComplexConjugate(AbstractOperation):
    """Unary complex conjugate operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])

    @property
    def type_name(self) -> TypeName:
        return "conj"

    def evaluate(self, a):
        return conjugate(a)


class Max(AbstractOperation):
    """Binary max operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "max"

    def evaluate(self, a, b):
        assert not isinstance(a, complex) and not isinstance(b, complex), \
            ("core_operations.Max does not support complex numbers.")
        return a if a > b else b


class Min(AbstractOperation):
    """Binary min operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 1, name = name, input_sources = [src0, src1])

    @property
    def type_name(self) -> TypeName:
        return "min"

    def evaluate(self, a, b):
        assert not isinstance(a, complex) and not isinstance(b, complex), \
            ("core_operations.Min does not support complex numbers.")
        return a if a < b else b


class Absolute(AbstractOperation):
    """Unary absolute value operation.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])

    @property
    def type_name(self) -> TypeName:
        return "abs"

    def evaluate(self, a):
        return np_abs(a)


class ConstantMultiplication(AbstractOperation):
    """Unary constant multiplication operation.
    TODO: More info.
    """

    def __init__(self, value: Number, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])
        self.set_param("value", value)

    @property
    def type_name(self) -> TypeName:
        return "cmul"

    def evaluate(self, a):
        return a * self.param("value")


class ConstantAddition(AbstractOperation):
    """Unary constant addition operation.
    TODO: More info.
    """

    def __init__(self, value: Number, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])
        self.set_param("value", value)

    @property
    def type_name(self) -> TypeName:
        return "cadd"

    def evaluate(self, a):
        return a + self.param("value")


class ConstantSubtraction(AbstractOperation):
    """Unary constant subtraction operation.
    TODO: More info.
    """

    def __init__(self, value: Number, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])
        self.set_param("value", value)

    @property
    def type_name(self) -> TypeName:
        return "csub"

    def evaluate(self, a):
        return a - self.param("value")


class ConstantDivision(AbstractOperation):
    """Unary constant division operation.
    TODO: More info.
    """

    def __init__(self, value: Number, src0: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 1, output_count = 1, name = name, input_sources = [src0])
        self.set_param("value", value)

    @property
    def type_name(self) -> TypeName:
        return "cdiv"

    def evaluate(self, a):
        return a / self.param("value")

class Butterfly(AbstractOperation):
    """Butterfly operation that returns two outputs.
    The first output is a + b and the second output is a - b.
    TODO: More info.
    """

    def __init__(self, src0: Optional[SignalSourceProvider] = None, src1: Optional[SignalSourceProvider] = None, name: Name = ""):
        super().__init__(input_count = 2, output_count = 2, name = name, input_sources = [src0, src1])

    def evaluate(self, a, b):
        return a + b, a - b

    @property
    def type_name(self) -> TypeName:
        return "bfly"