API接口: NeurAI 包#

Subpackages#

Submodules#

neurai.config.set_platform(platform='cpu')#

Sets the status for JAX platform.

Parameters:

platform (str, optional) – A string specifying the platform. Default is ‘cpu’.

Returns:

None

neurai.config.set_simulate_status(enable_simulate_jit=True, multi_thread_connect=True, dt=0.1, save_ir_collector=False, use_brainmodel=True)#

Sets the status for compiler settings.

Parameters:

  • enable_simulate_jit (bool, optional) – A boolean flag indicating if Just-In-Time (JIT) compilation is enabled. Default is True.

  • multi_thread_connect (bool, optional) – A boolean flag indicating if multi-threading is enabled. Default is False. If set to True, due to the creation of multiple threads, different PRNGKey will be generated when creating synapses, which could potentially lead to variations in the results of multiple simulation runs.

  • save_ir_collector (bool) – A boolean flag indicating if save the simulate network in ir.

Returns:

None

class neurai.monitor.CallbackRecorder(callback=None)#

Bases: Recorder

Record data using callback

Parameters:

callback (callable, optional) – Data is processed through callback functions, by default None

class neurai.monitor.Event(step, key, value)#

Bases: object

Stores events that log parameters

Parameters:

  • step (float) – The simulation current time step.

  • key (str) – The variable name

  • value (jnp.ndarray) – The variable value

class neurai.monitor.FileRecorder(file_path='./', over_write=False)#

Bases: Recorder

Record data using file

Parameters:

  • file_path (str, optional) – Data record address, by default ‘./’

  • over_write (bool, optional) – Whether to rewrite the file, by default False

class neurai.monitor.MemoryRecorder#

Bases: Recorder

Record data using memory

class neurai.monitor.Monitor(monitors=None, recorder=None)#

Bases: object

Abstract base class used to build new monitors.

Parameters:

  • recorder (Recorder,) – Record data to different locations.

  • monitors (List[MonitorConf],) – Configure the parameters or variables that need to be monitored.

class neurai.monitor.MonitorBS(monitors=None, recorder=None, start=0, stop=inf, interval=1, maxlen=10000, mode='latest')#

Bases: Monitor

Class for monitoring brain simulation models.

Parameters:

  • recorder (Recorder,) – Record data to different locations.

  • monitors (MonitorConf,) – Configure the parameters or variables that need to be monitored.

  • start (float, optional) – The start time step for monitoring, by default 0.

  • stop (float, optional) – The stop time step for monitoring, by default np.inf.

  • interval (int, optional) – The interval between two successive monitoring, by default 1.

  • maxlen (int , optional) – The maximum number of monitored data,by default 10000.

  • mode (str , optional) – The data record mode ,by default ‘oldest’.

simulate_init(param, sim_t, dt)#

Initializes the monitoring process by creating a dictionary for storing the monitored data.

Parameters:

  • sim_t (float) – The total simulation time.

  • dt (float) – The simulation step size.

Return type:

None

step_record(param, mon, step, step_logs=None)#

Records the monitored data at the end of each time step.

Parameters:

  • param (dict) – The dictionary of model parameters.

  • mon (dict) – The dictionary of parameters the need to be recorded to.

  • step (int) – The current number of steps.

  • step_logs (dict, optional) – The log files for each time step, by default None.

Return type:

dict

class neurai.monitor.MonitorConf(layer_name=None, var_name=None, grad=None, slices=None)#

Bases: object

Configures the variables that need to be monitored.

Parameters:

  • var_name (str) – Name of the variable to monitor.

  • layer_name (str, optional) – Optional name of the layer the variable belongs to, by default None

  • grad (str, optional) – Optional gradient associated with the variable, by default None

  • slices (range, optional) – Monitor the slice information of variables, by default None

class neurai.monitor.Recorder(callback=None, file_path='./', over_write=False)#

Bases: object

Record data to different locations

Parameters:

  • callback (callable, optional) – Data is processed through callback functions, by default None

  • file_path (str, optional) – Data record address, by default ‘./’

  • over_write (bool, optional) – Whether to rewrite the file, by default False

Module contents#