From 68a9ed7ded78555f8bcaf0929d90f58906cded62 Mon Sep 17 00:00:00 2001 From: Omar Nassar Date: Sun, 30 Mar 2025 20:02:47 +0200 Subject: [PATCH 1/3] improve speed of reduced set models Reduced Set models don't use numba for enhanced performance. This uses numba njit wrapper to compile the deriviative functions during runtime, and provide considerable performance gain. Since numba currently doesn't fully support jit-ed classes, derivative functions have to move outside the class scope. To improved readability, they were moved to a new file: _dfun_stefanescu_jirsa.py. To test the changes a new unit test module was added: models_benchmark.py. Here, we compute the derivative multiple times within a time window with/without numba and compare the throughput. Around 5x performance gains were observed. We also verify that numba doesn't change derivative values. --- .../models/_dfun_stefanescu_jirsa.py | 82 +++++++++++++++++++ .../tvb/simulator/models/stefanescu_jirsa.py | 73 ++++------------- .../library/simulator/models_benchmark.py | 78 ++++++++++++++++++ 3 files changed, 177 insertions(+), 56 deletions(-) create mode 100644 tvb_library/tvb/simulator/models/_dfun_stefanescu_jirsa.py create mode 100644 tvb_library/tvb/tests/library/simulator/models_benchmark.py diff --git a/tvb_library/tvb/simulator/models/_dfun_stefanescu_jirsa.py b/tvb_library/tvb/simulator/models/_dfun_stefanescu_jirsa.py new file mode 100644 index 0000000000..453ff8336a --- /dev/null +++ b/tvb_library/tvb/simulator/models/_dfun_stefanescu_jirsa.py @@ -0,0 +1,82 @@ +import numba as nb +import numpy as np + + +def dfun_fitzhughnaguma(state_variables, coupling, local_coupling, tau, e_i, K11, Aik, K12, Bik, IE_i, + b, m_i, K21, Cik, II_i, f_i, n_i): + xi = state_variables[0, :] + eta = state_variables[1, :] + alpha = state_variables[2, :] + beta = state_variables[3, :] + derivative = np.empty_like(state_variables) + + # Sum the activity from the modes + c_0 = coupling[0, :].sum(axis=1)[:, np.newaxis] + + # Compute derivatives + derivative[0] = ( + tau * (xi - e_i * xi ** 3 / 3.0 - eta) + + K11 * (np.dot(xi, Aik) - xi) + - K12 * (np.dot(alpha, Bik) - xi) + + tau * (IE_i + c_0 + local_coupling * xi) + ) + + derivative[1] = (xi - b * eta + m_i) / tau + + derivative[2] = ( + tau * (alpha - f_i * alpha ** 3 / 3.0 - beta) + + K21 * (np.dot(xi, Cik) - alpha) + + tau * (II_i + c_0 + local_coupling * xi) + ) + + derivative[3] = (alpha - b * beta + n_i) / tau + + return derivative + +_dfun_fitzhughnaguma_numba = nb.njit(dfun_fitzhughnaguma) +def dfun_fitzhughnaguma_numba(state_variables, coupling, local_coupling, tau, e_i, K11, Aik, K12, Bik, IE_i, + b, m_i, K21, Cik, II_i, f_i, n_i): + return _dfun_fitzhughnaguma_numba(state_variables, coupling, local_coupling, + tau, e_i, K11, Aik, K12, Bik, IE_i, b, + m_i, K21, Cik, II_i, f_i, n_i) + +def dfun_hindmarshrose(state_variables, coupling, local_coupling, r, a_i, b_i, c_i, d_i, s, K11, A_ik, K12, B_ik, IE_i, m_i, K21, C_ik, II_i, e_i, f_i, h_i, p_i, n_i): + xi = state_variables[0, :] + eta = state_variables[1, :] + tau = state_variables[2, :] + alpha = state_variables[3, :] + beta = state_variables[4, :] + gamma = state_variables[5, :] + derivative = np.empty_like(state_variables) + + c_0 = coupling[0, :].sum(axis=1)[:, np.newaxis] + # c_1 = coupling[1, :] + + derivative[0] = (eta - a_i * xi ** 3 + b_i * xi ** 2 - tau + + K11 * (np.dot(xi, A_ik) - xi) - + K12 * (np.dot(alpha, B_ik) - xi) + + IE_i + c_0 + local_coupling * xi) + + derivative[1] = c_i - d_i * xi ** 2 - eta + + derivative[2] = r * s * xi - r * tau - m_i + + derivative[3] = (beta - e_i * alpha ** 3 + f_i * alpha ** 2 - gamma + + K21 * (np.dot(xi, C_ik) - alpha) + + II_i + c_0 + local_coupling * xi) + + derivative[4] = h_i - p_i * alpha ** 2 - beta + + derivative[5] = r * s * alpha - r * gamma - n_i + + return derivative + +_dfun_hindmarshrose_numba = nb.njit(dfun_hindmarshrose) +@nb.njit +def dfun_hindmarshrose_numba(state_variables, coupling, local_coupling, r, a_i, b_i, c_i, d_i, s, K11, A_ik, K12, B_ik, IE_i, m_i, K21, C_ik, II_i, e_i, f_i, h_i, p_i, n_i): + return _dfun_hindmarshrose_numba(state_variables, coupling, local_coupling, + r, a_i, b_i, c_i, d_i, s, K11, A_ik, + K12, B_ik, IE_i, m_i, K21, C_ik, II_i, + e_i, f_i, h_i, p_i, n_i) + + diff --git a/tvb_library/tvb/simulator/models/stefanescu_jirsa.py b/tvb_library/tvb/simulator/models/stefanescu_jirsa.py index df235e1e2d..3609825862 100644 --- a/tvb_library/tvb/simulator/models/stefanescu_jirsa.py +++ b/tvb_library/tvb/simulator/models/stefanescu_jirsa.py @@ -31,6 +31,7 @@ import numpy from scipy.integrate import trapezoid as scipy_integrate_trapz from scipy.stats import norm as scipy_stats_norm +from ._dfun_stefanescu_jirsa import dfun_fitzhughnaguma, dfun_fitzhughnaguma_numba, dfun_hindmarshrose, dfun_hindmarshrose_numba from .base import Model from tvb.basic.neotraits.api import NArray, Final, List, Range @@ -106,6 +107,7 @@ class ReducedSetFitzHughNagumo(ReducedSetBase): """ # Define traited attributes for this model, these represent possible kwargs. + use_numba = True tau = NArray( label=r":math:`\tau`", default=numpy.array([3.0]), @@ -213,32 +215,14 @@ def dfun(self, state_variables, coupling, local_coupling=0.0): \dot{\beta}_i &= \frac{1}{c}\left(\alpha_i-b\beta_i+n_i\right) """ - - xi = state_variables[0, :] - eta = state_variables[1, :] - alpha = state_variables[2, :] - beta = state_variables[3, :] - derivative = numpy.empty_like(state_variables) - # sum the activity from the modes - c_0 = coupling[0, :].sum(axis=1)[:, numpy.newaxis] - - # TODO: generalize coupling variables to a matrix form - # c_1 = coupling[1, :] # this cv represents alpha - - derivative[0] = (self.tau * (xi - self.e_i * xi ** 3 / 3.0 - eta) + - self.K11 * (numpy.dot(xi, self.Aik) - xi) - - self.K12 * (numpy.dot(alpha, self.Bik) - xi) + - self.tau * (self.IE_i + c_0 + local_coupling * xi)) - - derivative[1] = (xi - self.b * eta + self.m_i) / self.tau - - derivative[2] = (self.tau * (alpha - self.f_i * alpha ** 3 / 3.0 - beta) + - self.K21 * (numpy.dot(xi, self.Cik) - alpha) + - self.tau * (self.II_i + c_0 + local_coupling * xi)) - - derivative[3] = (alpha - self.b * beta + self.n_i) / self.tau - - return derivative + dfun_to_use = dfun_fitzhughnaguma + if self.use_numba: + dfun_to_use = dfun_fitzhughnaguma_numba + return dfun_to_use( + state_variables, coupling, local_coupling, + self.tau, self.e_i, self.K11, self.Aik, self.K12, self.Bik, self.IE_i, + self.b, self.m_i, self.K21, self.Cik, self.II_i, self.f_i, self.n_i + ) def update_derived_parameters(self): """ @@ -484,6 +468,7 @@ class ReducedSetHindmarshRose(ReducedSetBase): II_i = None m_i = None n_i = None + use_numba = True def dfun(self, state_variables, coupling, local_coupling=0.0): r""" @@ -509,36 +494,12 @@ def dfun(self, state_variables, coupling, local_coupling=0.0): \dot{\gamma}_i &= rs \alpha_i - r \gamma_i - n_i """ - - xi = state_variables[0, :] - eta = state_variables[1, :] - tau = state_variables[2, :] - alpha = state_variables[3, :] - beta = state_variables[4, :] - gamma = state_variables[5, :] - derivative = numpy.empty_like(state_variables) - - c_0 = coupling[0, :].sum(axis=1)[:, numpy.newaxis] - # c_1 = coupling[1, :] - - derivative[0] = (eta - self.a_i * xi ** 3 + self.b_i * xi ** 2 - tau + - self.K11 * (numpy.dot(xi, self.A_ik) - xi) - - self.K12 * (numpy.dot(alpha, self.B_ik) - xi) + - self.IE_i + c_0 + local_coupling * xi) - - derivative[1] = self.c_i - self.d_i * xi ** 2 - eta - - derivative[2] = self.r * self.s * xi - self.r * tau - self.m_i - - derivative[3] = (beta - self.e_i * alpha ** 3 + self.f_i * alpha ** 2 - gamma + - self.K21 * (numpy.dot(xi, self.C_ik) - alpha) + - self.II_i + c_0 + local_coupling * xi) - - derivative[4] = self.h_i - self.p_i * alpha ** 2 - beta - - derivative[5] = self.r * self.s * alpha - self.r * gamma - self.n_i - - return derivative + dfun_to_use = dfun_hindmarshrose + if self.use_numba: + dfun_to_use = dfun_hindmarshrose_numba + return dfun_to_use( + state_variables, coupling, local_coupling, + self.r, self.a_i, self.b_i, self.c_i, self.d_i, self.s, self.K11, self.A_ik, self.K12, self.B_ik, self.IE_i, self.m_i, self.K21, self.C_ik, self.II_i, self.e_i, self.f_i, self.h_i, self.p_i, self.n_i) def update_derived_parameters(self, corrected_d_p=True): """ diff --git a/tvb_library/tvb/tests/library/simulator/models_benchmark.py b/tvb_library/tvb/tests/library/simulator/models_benchmark.py new file mode 100644 index 0000000000..1b3958a681 --- /dev/null +++ b/tvb_library/tvb/tests/library/simulator/models_benchmark.py @@ -0,0 +1,78 @@ +import numpy as np +import time +from tvb.tests.library.base_testcase import BaseTestCase +from tvb.simulator.models.stefanescu_jirsa import ReducedSetFitzHughNagumo +from tvb.simulator.models.stefanescu_jirsa import ReducedSetHindmarshRose +from tvb.simulator.models.base import Model + + +class TestBenchmarkModels(BaseTestCase): + """Test cases that benchmark the performance of models' implementation + """ + + def randn_state_for_model(self, model: Model, n_node): + shape = (model.nvar, n_node, model.number_of_modes) + state = np.random.randn(*shape) + return state + + + def zero_coupling_for_model(self, model: Model, n_node): + n_cvar = len(model.cvar) + shape = (n_cvar, n_node, model.number_of_modes) + coupling = np.zeros(shape) + return coupling + + def eps_for_model(self, model: Model, n_node, time_limit=0.5, state=None, coupling=None): + model.configure() + if state is None: + state = self.randn_state_for_model(model, n_node) + if coupling is None: + coupling = self.zero_coupling_for_model(model, n_node) + # throw one away in case of initialization + model.dfun(state, coupling) + # start timing + tic = time.time() + n_eval = 0 + while (time.time() - tic) < time_limit: + model.dfun(state, coupling) + n_eval += 1 + toc = time.time() + return n_eval / (toc - tic) + + def test_reduced_set_fitz_hugh_nagumo_when_using_numba_improve_performance(self): + # create, & initialize numba & no-numba models + rs_fhn_model = ReducedSetFitzHughNagumo() + n_node = 10 + state = self.randn_state_for_model(rs_fhn_model, n_node) + coupling = self.zero_coupling_for_model(rs_fhn_model, n_node) + + rs_fhn_model.use_numba = False + no_numba_performance = self.eps_for_model(rs_fhn_model, 10, state=state, coupling=coupling) + no_numba_dfun_derivative = rs_fhn_model.dfun(state, coupling) + rs_fhn_model.use_numba = True + numba_performance = self.eps_for_model(rs_fhn_model, 10, state=state, coupling=coupling) + numba_dfun_derivative = rs_fhn_model.dfun(state, coupling) + speedup = numba_performance / no_numba_performance + print(f"speedup: {speedup}") + print(f"no_numba_performance: {no_numba_performance}, numba_performance: {numba_performance}") + assert speedup > 1 + assert np.allclose(no_numba_dfun_derivative, numba_dfun_derivative) + + def test_reduced_set_hindmarsh_rose_when_using_numba_improve_performance(self): + # create, & initialize numba & no-numba models + rs_hr_model = ReducedSetHindmarshRose() + n_node = 10 + state = self.randn_state_for_model(rs_hr_model, n_node) + coupling = self.zero_coupling_for_model(rs_hr_model, n_node) + + rs_hr_model.use_numba = False + no_numba_performance = self.eps_for_model(rs_hr_model, 10, state=state, coupling=coupling) + no_numba_dfun_derivative = rs_hr_model.dfun(state, coupling) + rs_hr_model.use_numba = True + numba_performance = self.eps_for_model(rs_hr_model, 10, state=state, coupling=coupling) + numba_dfun_derivative = rs_hr_model.dfun(state, coupling) + speedup = numba_performance / no_numba_performance + print(f"speedup: {speedup}") + print(f"no_numba_performance: {no_numba_performance}, numba_performance: {numba_performance}") + assert speedup > 1 + assert np.allclose(no_numba_dfun_derivative, numba_dfun_derivative) From d32512d0107cc8a464ec8cf3ad0e69b34aa4e2cc Mon Sep 17 00:00:00 2001 From: Omar Nassar Date: Mon, 31 Mar 2025 22:36:01 +0200 Subject: [PATCH 2/3] use random coupling to ensure the coupling expressions inside the numba dfun are tested --- .../tvb/tests/library/simulator/models_benchmark.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tvb_library/tvb/tests/library/simulator/models_benchmark.py b/tvb_library/tvb/tests/library/simulator/models_benchmark.py index 1b3958a681..99d12ca7e8 100644 --- a/tvb_library/tvb/tests/library/simulator/models_benchmark.py +++ b/tvb_library/tvb/tests/library/simulator/models_benchmark.py @@ -16,10 +16,10 @@ def randn_state_for_model(self, model: Model, n_node): return state - def zero_coupling_for_model(self, model: Model, n_node): + def randn_coupling_for_model(self, model: Model, n_node): n_cvar = len(model.cvar) shape = (n_cvar, n_node, model.number_of_modes) - coupling = np.zeros(shape) + coupling = np.random.randn(*shape) return coupling def eps_for_model(self, model: Model, n_node, time_limit=0.5, state=None, coupling=None): @@ -27,7 +27,7 @@ def eps_for_model(self, model: Model, n_node, time_limit=0.5, state=None, coupli if state is None: state = self.randn_state_for_model(model, n_node) if coupling is None: - coupling = self.zero_coupling_for_model(model, n_node) + coupling = self.randn_coupling_for_model(model, n_node) # throw one away in case of initialization model.dfun(state, coupling) # start timing @@ -44,7 +44,7 @@ def test_reduced_set_fitz_hugh_nagumo_when_using_numba_improve_performance(self) rs_fhn_model = ReducedSetFitzHughNagumo() n_node = 10 state = self.randn_state_for_model(rs_fhn_model, n_node) - coupling = self.zero_coupling_for_model(rs_fhn_model, n_node) + coupling = self.randn_coupling_for_model(rs_fhn_model, n_node) rs_fhn_model.use_numba = False no_numba_performance = self.eps_for_model(rs_fhn_model, 10, state=state, coupling=coupling) @@ -63,7 +63,7 @@ def test_reduced_set_hindmarsh_rose_when_using_numba_improve_performance(self): rs_hr_model = ReducedSetHindmarshRose() n_node = 10 state = self.randn_state_for_model(rs_hr_model, n_node) - coupling = self.zero_coupling_for_model(rs_hr_model, n_node) + coupling = self.randn_coupling_for_model(rs_hr_model, n_node) rs_hr_model.use_numba = False no_numba_performance = self.eps_for_model(rs_hr_model, 10, state=state, coupling=coupling) From cb09a91bd94f2c95635d565f218a7c81166bf765 Mon Sep 17 00:00:00 2001 From: Omar Nassar Date: Sun, 13 Apr 2025 16:23:03 +0200 Subject: [PATCH 3/3] rename unit test names to shorter, more consice ones --- tvb_library/tvb/tests/library/simulator/models_benchmark.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tvb_library/tvb/tests/library/simulator/models_benchmark.py b/tvb_library/tvb/tests/library/simulator/models_benchmark.py index 99d12ca7e8..bd3ad9fced 100644 --- a/tvb_library/tvb/tests/library/simulator/models_benchmark.py +++ b/tvb_library/tvb/tests/library/simulator/models_benchmark.py @@ -39,7 +39,7 @@ def eps_for_model(self, model: Model, n_node, time_limit=0.5, state=None, coupli toc = time.time() return n_eval / (toc - tic) - def test_reduced_set_fitz_hugh_nagumo_when_using_numba_improve_performance(self): + def test_rsfhn_numba(self): # create, & initialize numba & no-numba models rs_fhn_model = ReducedSetFitzHughNagumo() n_node = 10 @@ -58,7 +58,7 @@ def test_reduced_set_fitz_hugh_nagumo_when_using_numba_improve_performance(self) assert speedup > 1 assert np.allclose(no_numba_dfun_derivative, numba_dfun_derivative) - def test_reduced_set_hindmarsh_rose_when_using_numba_improve_performance(self): + def test_rshr_numba(self): # create, & initialize numba & no-numba models rs_hr_model = ReducedSetHindmarshRose() n_node = 10