Merge pull request #171 from shixnya/develop

FilterNet transfer function change
AllenInstitute · Apr 21, 2021 · 6f868d9 · 6f868d9
2 parents 224bcb6 + 6c36bce
commit 6f868d9
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/bmtk/simulator/filternet/lgnmodel/transferfunction.py b/bmtk/simulator/filternet/lgnmodel/transferfunction.py
@@ -8,7 +8,9 @@ class ScalarTransferFunction(object):
     def __init__(self, transfer_function_string, symbol=sympy.abc.s):
         self.symbol = symbol
         self.transfer_function_string = transfer_function_string
-        self.closure = lambdify(self.symbol, symp.parse_expr(self.transfer_function_string), modules=['sympy'])
+        # replacing sympy.Heaviside() with np.heaviside() for better performance
+        modules = [{'Heaviside': lambda x: 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)
@@ -42,7 +44,8 @@ class MultiTransferFunction(object):
     def __init__(self, symbol_tuple, transfer_function_string):
         self.symbol_tuple = symbol_tuple
         self.transfer_function_string = transfer_function_string
-        self.closure = lambdify(self.symbol_tuple, symp.parse_expr(self.transfer_function_string), modules=['sympy'])
+        modules = [{'Heaviside': lambda x: 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,)):