@@ -23,11 +23,9 @@ from jax import numpy as jnp, random, jit
2323import numpy as np
2424
2525from ngclearn.utils.model_utils import scanner
26- from ngcsimlib.compilers import compile_command, wrap_command
2726from ngcsimlib.context import Context
28- from ngcsimlib.commands import Command
27+ from ngcsimlib.compilers.process import Process
2928# # import model-specific mechanisms
30- from ngclearn.operations import summation
3129from ngclearn.components.neurons.spiking.adExCell import AdExCell
3230
3331# # create seeding keys (JAX-style)
@@ -43,19 +41,20 @@ w0 = 0. ## initial recovery value (for its reset condition)
4341
4442# # create simple system with only one AdEx
4543with Context(" Model" ) as model:
46- cell = AdExCell(" z0 " , n_units = 1 , tau_m = 15 ., resist_m = 1 ., tau_w = 400 .,
47- v_sharpness = 2 .,
48- intrinsic_mem_thr = - 55 ., v_thr = 5 ., v_rest = - 72 ., v_reset = - 75 .,
49- a = 0.1 , b = 0.75 ,
50- v0 = v0, w0 = w0, integration_type = " euler " , key = subkeys[ 0 ] )
44+ cell = AdExCell(
45+ " z0 " , n_units = 1 , tau_m = 15 ., resist_m = 1 ., tau_w = 400 ., v_sharpness = 2 .,
46+ intrinsic_mem_thr = - 55 ., v_thr = 5 ., v_rest = - 72 ., v_reset = - 75 ., a = 0.1 ,
47+ b = 0.75 , v0 = v0, w0 = w0, integration_type = " euler " , key = subkeys[ 0 ]
48+ )
5149
5250 # # create and compile core simulation commands
53- reset_cmd, reset_args = model.compile_by_key(cell, compile_key = " reset" )
54- model.add_command(wrap_command(jit(model.reset)), name = " reset" )
55- advance_cmd, advance_args = model.compile_by_key(cell,
56- compile_key = " advance_state" )
57- model.add_command(wrap_command(jit(model.advance_state)), name = " advance" )
51+ advance_process = (Process()
52+ >> cell.advance_state)
53+ model.wrap_and_add_command(jit(advance_process.pure), name = " advance" )
5854
55+ reset_process = (Process()
56+ >> cell.reset)
57+ model.wrap_and_add_command(jit(reset_process.pure), name = " reset" )
5958
6059 # # set up non-compiled utility commands
6160 @Context.dynamicCommand
0 commit comments