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)