ALLEN INSTITUTE

# Source code for bmtk.simulator.filternet.lgnmodel.transferfunction

    from sympy.utilities.lambdify import lambdify
    from sympy import Matrix
    import sympy.parsing.sympy_parser as symp
    import sympy.abc
    import numpy as np

    class ScalarTransferFunction(object):
        def __init__(self, transfer_function_string, symbol=sympy.abc.s):
            self.symbol = symbol
            self.transfer_function_string = transfer_function_string
            # replacing sympy.Heaviside() with np.heaviside() for better performance
            modules = [{'Heaviside': lambda x, y=0.5: np.heaviside(x, 0.5)}, 'numpy', 'sympy']
            self.closure = lambdify(self.symbol, symp.parse_expr(self.transfer_function_string), modules=modules)

        def __call__(self, s):
            return self.closure(s)

        def to_dict(self):
            return {'class': (__name__, self.__class__.__name__), 'function': self.transfer_function_string}

        def imshow(self, rates, times=None, show=True):
            import matplotlib.pyplot as plt

            vals = [self(rate) for rate in rates]
            times = np.linspace(0.0, 1.0, len(rates)) if times is None else times

            fig, ax1 = plt.subplots()
            ax1.set_xlabel('time (seconds)')
            ax1.set_ylabel(str(self.symbol), color='b')
            ax1.plot(times, rates, '--b')
            ax1.tick_params(axis='y', labelcolor='b')

            ax2 = ax1.twinx()
            ax2.set_ylabel('transform', color='r')
            ax2.plot(times, vals, 'r')
            ax2.tick_params(axis='y', labelcolor='r')

            plt.title(self.transfer_function_string)
            if show:
                plt.show()

    class MultiTransferFunction(object):
        def __init__(self, symbol_tuple, transfer_function_string):
            self.symbol_tuple = symbol_tuple
            self.transfer_function_string = transfer_function_string
            modules = [{'Heaviside': lambda x, y=0.5: np.heaviside(x, 0.5)}, 'numpy', 'sympy']
            self.closure = lambdify(self.symbol_tuple,symp.parse_expr(self.transfer_function_string), modules=modules)

        def __call__(self, *s):
            if isinstance(s[0], (float,)):
                return self.closure(*s)
            else:
                return np.array(list(map(lambda x: self.closure(*x), zip(*s))))

        def to_dict(self):
            return {'class': (__name__, self.__class__.__name__), 'function': self.transfer_function_string}