def _test_cfun2(self, cfun): "This variant tests per node cfun kernel." sim = self._prep_sim(cfun) # prep & invoke kernel template = '''import numpy as np <%include file="nb-coupling2.py.mako"/> def kernel(t, cx, weights, state, di): for i in range(weights.shape[0]): % for cvar, cterm in zip(sim.model.cvar, sim.model.coupling_terms): cx[${loop.index},i] = cx_${cterm}(t, i, weights, state[${cvar}], di) % endfor ''' kernel = NbBackend().build_py_func(template, dict(sim=sim, np=np), print_source=True) fill = np.r_[:sim.history.buffer.size] fill = np.reshape(fill, sim.history.buffer.shape[:-1]) sim.history.buffer[..., 0] = fill sim.current_state[:] = fill[0, :, :, None] # [current_state, oldest, next oldest, ..., current - 1] buf = sim.history.buffer[..., 0] # (horizon, svar, node) horizon, nsvar, nnode = buf.shape # switch to Jan's history layout here, time increasing with index buf2 = np.zeros((nsvar, nnode, horizon + 10)) buf2[:, :, :horizon - 1] = np.transpose(buf[1:], (1, 2, 0)) buf2[:, :, horizon - 1] = buf[0] weights = sim.connectivity.weights.astype('f') cX = np.zeros((nsvar, nnode)) kernel(horizon - 1, cX, weights, buf2, sim.connectivity.idelays) # do comparison (t, y), = sim.run() np.testing.assert_allclose(cX, y[0, :, :, 0], 1e-5, 1e-6)
def _test_cfun(self, cfun): "Test a Python cfun template." sim = self._prep_sim(cfun) # prep & invoke kernel template = f'''import numpy as np <%include file="nb-coupling.py.mako"/> ''' kernel = NbBackend().build_py_func(template, dict(sim=sim, np=np), name='coupling', print_source=True) fill = np.r_[:sim.history.buffer.size] fill = np.reshape(fill, sim.history.buffer.shape[:-1]) sim.history.buffer[..., 0] = fill sim.current_state[:] = fill[0, :, :, None] buf = sim.history.buffer[..., 0] # kernel has history in reverse order except 1st element 🤕 rbuf = np.concatenate((buf[0:1], buf[1:][::-1]), axis=0) state = np.transpose(rbuf, (1, 0, 2)).astype('f') weights = sim.connectivity.weights.astype('f') cX = np.zeros_like(state[:, 0]) kernel(cX, weights, state, sim.connectivity.idelays) # do comparison (t, y), = sim.run() np.testing.assert_allclose(cX, y[0, :, :, 0], 1e-5, 1e-6)
def _eval_cg(self, integrator_, state, weights_): class sim: integrator = integrator_ connectivity = Connectivity.from_file() model = MontbrioPazoRoxin() sim.connectivity.speed = np.r_[np.inf] sim.connectivity.configure() sim.integrator.configure() sim.connectivity.set_idelays(sim.integrator.dt) template = ''' import numpy as np import numba as nb cx_Coupling_Term_r = nb.njit(lambda t, i, w, r: np.dot(w[i],r[:,t])) cx_Coupling_Term_V = nb.njit(lambda t, i, w, V: np.dot(w[i],V[:,t])) dx_r = nb.njit(lambda r,V,Coupling_Term_r,Coupling_Term_V,parmat: -r*Coupling_Term_r**2/76) dx_V = nb.njit(lambda r,V,Coupling_Term_r,Coupling_Term_V,parmat: -V*Coupling_Term_V**2/76) <%include file="nb-integrate.py.mako" /> ''' integrate = NbBackend().build_py_func(template, dict(sim=sim, np=np), name='integrate', print_source=True) parmat = np.zeros(0) np.random.seed(42) args = state, weights_, parmat if isinstance(sim.integrator, IntegratorStochastic): print(sim.integrator.noise.nsig.shape) args = args + (sim.integrator.noise.nsig, ) integrate(0, *args) return state
def _test_dfun(self, model_): "Test a Numba cfun template." class sim: # dummy sim model = model_ template = '''import numpy as np <% svars = ', '.join(sim.model.state_variables) cvars = ', '.join(sim.model.coupling_terms) %> <%include file="nb-dfuns.py.mako"/> def kernel(dx, state, cx, parmat): for i in range(state.shape[1]): ${svars} = state[:, i] % for cvar in sim.model.coupling_terms: ${cvar} = cx[${loop.index}, i] % endfor % for svar in sim.model.state_variables: dx[${loop.index},i] = dx_${svar}(${svars}, ${cvars}, parmat[:,i] if parmat.size else None) % endfor ''' kernel = NbBackend().build_py_func(template, dict(sim=sim, np=np), print_source=True) cX = np.random.rand(2, 128) / 10.0 dX = np.zeros_like(cX) state = np.random.rand(2, 128) parmat = sim.model.spatial_parameter_matrix kernel(dX, state, cX, parmat) drh, dVh = dX dr, dV = sim.model.dfun(state, cX) np.testing.assert_allclose(drh, dr) np.testing.assert_allclose(dVh, dV)
def test_nb_mako_100ms(benchmark): sim = make_sim(100.0) template = '<%include file="nb-sim.py.mako"/>' content = dict(sim=sim, np=np, debug_nojit=False) kernel = NbBackend().build_py_func(template, content, print_source=True, name='run_sim') benchmark(lambda: kernel(sim))
def _test_mpr(self, integrator, delays=False): sim = self._create_sim( integrator, inhom_mmpr=True, delays=delays, run_sim=False ) template = '<%include file="nb-sim.py.mako"/>' content = dict(sim=sim, np=np, debug_nojit=True) kernel = NbBackend().build_py_func( template, content, print_source=True, name='run_sim', ) np.random.seed(42) state = kernel(sim) # (nsvar, nnode, horizon + nstep) yh = np.transpose(state[:,:,sim.connectivity.horizon:], (2,0,1)) (_, y), = sim.run() self._check_match(y, yh)