extended_tf.py

# extended_tf.py - Extended definition of the discrete
# transfer function implemented in scipy.signal.
# Supports arithmetical operations between transfer function
# and feedback loop computation
#
# Code author: [Alessio Russo - alessior@kth.se]
# Last update: 10th January 2020, by alessior@kth.se
#
# Copyright (c) [2021] Alessio Russo [alessior@kth.se]. All rights reserved.
# This file is part of PythonVRFT.
# PythonVRFT is free software: you can redistribute it and/or modify
# it under the terms of the MIT License. You should have received a copy of
# the MIT License along with PythonVRFT.
# If not, see <https://opensource.org/licenses/MIT>.
#

from __future__ import division

import numpy as np
import scipy.signal as scipysig
from scipy.signal.ltisys import TransferFunction as TransFun
from numpy import polymul, polyadd


class ExtendedTF(scipysig.ltisys.TransferFunctionDiscrete):
    """
    Extended definition of the discrete transfer function implemented in scipy.signal.
    Supports arithmetical operations between transfer function and feedback loop
    computation
    """

    def __init__(self, num: np.ndarray, den: np.ndarray, dt: float):
        self._dt = dt
        super().__init__(num, den, dt=dt)

    def __neg__(self):
        return ExtendedTF(-self.num, self.den, dt=self._dt)

    def __floordiv__(self, other):
        # can't make sense of integer division right now
        return NotImplemented

    def __mul__(self, other):
        if type(other) in [int, float]:
            return ExtendedTF(self.num*other, self.den, dt=self._dt)
        elif type(other) in [TransFun, ExtendedTF]:
            numer = polymul(self.num, other.num)
            denom = polymul(self.den, other.den)
            return ExtendedTF(numer, denom, dt=self._dt)

    def __truediv__(self, other):
        if type(other) in [int, float]:
            return ExtendedTF(self.num,self.den*other, dt=self._dt)
        if type(other) in [TransFun, ExtendedTF]:
            numer = polymul(self.num, other.den)
            denom = polymul(self.den, other.num)
            return ExtendedTF(numer, denom, dt=self._dt)

    def __rtruediv__(self, other):
        if type(other) in [int, float]:
            return ExtendedTF(other*self.den, self.num, dt=self._dt)
        if type(other) in [TransFun, ExtendedTF]:
            numer = polymul(self.den, other.num)
            denom = polymul(self.num, other.den)
            return ExtendedTF(numer, denom, dt=self._dt)

    def __add__(self,other):
        if type(other) in [int, float]:
            return ExtendedTF(polyadd(self.num, self.den*other), self.den, dt=self._dt)
        if type(other) in [TransFun, type(self)]:
            if len(self.den) == len(other.den) and np.all(self.den == other.den):
                numer = polyadd(self.num, other.num)
                denom = self.den
            else:
                numer = polyadd(polymul(self.num, other.den), polymul(self.den, other.num))
                denom = polymul(self.den, other.den)
            return ExtendedTF(numer, denom, dt=self._dt)

    def __sub__(self, other):
        if type(other) in [int, float]:
            return ExtendedTF(polyadd(self.num, -self.den*other), self.den, dt=self._dt)
        if type(other) in [TransFun, type(self)]:
            if len(self.den) == len(other.den) and np.all(self.den == other.den):
                numer = polyadd(self.num, -other.num)
                denom = self.den
            else:
                numer = polyadd(polymul(self.num, other.den), -polymul(self.den, other.num))
                denom = polymul(self.den, other.den)
            return ExtendedTF(numer, denom, dt=self._dt)

    def __rsub__(self, other):
        if type(other) in [int, float]:
            return ExtendedTF(polyadd(-self.num, self.den*other), self.den, dt=self._dt)
        if type(other) in [TransFun, type(self)]:
            if len(self.den) == len(other.den) and np.all(self.den == other.den):
                numer = polyadd(self.num, -other.num)
                denom = self.den
            else:
                numer = polyadd(polymul(self.num, other.den), -polymul(self.den, other.num))
                denom = polymul(self.den, other.den)
            return ExtendedTF(numer, denom, dt=self._dt)

    def feedback(self):
        """ Computes T(z)/(1+T(z)) """
        num = self.num
        den = self.den
        den = polyadd(num, den)
        self = ExtendedTF(num, den, dt=self.dt)
        return self

    __rmul__ = __mul__
    __radd__ = __add__