def __init__(self, size, monitors=None, name=None, show_code=False):
# name
# -----
if name is None:
name = ''
else:
name = '_' + name
global _NeuGroup_NO
_NeuGroup_NO += 1
name = f'NG{_NeuGroup_NO}{name}'
# size
# ----
if isinstance(size, (list, tuple)):
if len(size) <= 0:
raise errors.ModelDefError(
'size must be int, or a tuple/list of int.')
if not isinstance(size[0], int):
raise errors.ModelDefError(
'size must be int, or a tuple/list of int.')
size = tuple(size)
elif isinstance(size, int):
size = (size, )
else:
raise errors.ModelDefError(
'size must be int, or a tuple/list of int.')
self.size = size
# initialize
# ----------
super(NeuGroup, self).__init__(steps={'update': self.update},
monitors=monitors,
name=name,
show_code=show_code)
def __init__(self, steps, monitors=None, name=None, host=None, show_code=False):
# host of the data
# ----------------
if host is None:
host = self
self.host = host
# model
# -----
if callable(steps):
self.steps = OrderedDict([(steps.__name__, steps)])
elif isinstance(steps, (list, tuple)) and callable(steps[0]):
self.steps = OrderedDict([(step.__name__, step) for step in steps])
elif isinstance(steps, dict):
self.steps = steps
else:
raise errors.ModelDefError(f'Unknown model type: {type(steps)}. Currently, BrainPy '
f'only supports: function, list/tuple/dict of functions.')
# name
# ----
if name is None:
global _DynamicSystem_NO
name = f'DS{_DynamicSystem_NO}'
_DynamicSystem_NO += 1
if not name.isidentifier():
raise errors.ModelUseError(f'"{name}" isn\'t a valid identifier according to Python '
f'language definition. Please choose another name.')
self.name = name
# monitors
# ---------
if monitors is None:
monitors = []
self.mon = Monitor(target=self, variables=monitors)
# runner
# -------
self.driver = backend.get_node_driver()(pop=self)
# run function
# ------------
self.run_func = None
# others
# ---
self.show_code = show_code
if self.target_backend is None:
raise errors.ModelDefError('Must define "target_backend".')
if isinstance(self.target_backend, str):
self._target_backend = (self.target_backend, )
elif isinstance(self.target_backend, (tuple, list)):
if not isinstance(self.target_backend[0], str):
raise errors.ModelDefError('"target_backend" must be a list/tuple of string.')
self._target_backend = tuple(self.target_backend)
else:
raise errors.ModelDefError(f'Unknown setting of "target_backend": {self.target_backend}')
def get_args(f):
"""Get the function arguments.
Parameters
----------
f : callable
The function.
Returns
-------
args : tuple
The variable names, the other arguments, and the original args.
"""
# 1. get the function arguments
parameters = inspect.signature(f).parameters
arguments = []
for name, par in parameters.items():
if par.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD:
arguments.append(par.name)
elif par.kind is inspect.Parameter.KEYWORD_ONLY:
arguments.append(par.name)
elif par.kind is inspect.Parameter.VAR_POSITIONAL:
raise errors.ModelDefError(
'Step function do not support positional parameters, e.g., *args'
)
elif par.kind is inspect.Parameter.POSITIONAL_ONLY:
raise errors.ModelDefError(
'Step function do not support positional only parameters, e.g., /'
)
elif par.kind is inspect.Parameter.VAR_KEYWORD:
raise errors.ModelDefError(
f'Step function do not support dict of keyword arguments: {str(par)}'
)
else:
raise errors.ModelDefError(f'Unknown argument type: {par.kind}')
# 2. check the function arguments
class_kw = None
if len(arguments) > 0 and arguments[0] in backend.CLASS_KEYWORDS:
class_kw = arguments[0]
arguments = arguments[1:]
for a in arguments:
if a in backend.CLASS_KEYWORDS:
raise errors.DiffEqError(f'Class keywords "{a}" must be defined '
f'as the first argument.')
return class_kw, arguments
def build(self, inputs, inputs_is_formatted=False, return_code=True, mon_length=0, show_code=False):
"""Build the object for running.
Parameters
----------
inputs : list, tuple, optional
The object inputs.
inputs_is_formatted : bool
Whether the "inputs" is formatted.
return_code : bool
Whether return the formatted codes.
mon_length : int
The monitor length.
Returns
-------
calls : list, tuple
The code lines to call step functions.
"""
if (self._target_backend[0] != 'general') and (backend.get_backend_name() not in self._target_backend):
raise errors.ModelDefError(f'The model {self.name} is target to run on {self._target_backend}, '
f'but currently the selected backend is {backend.get_backend_name()}')
if not inputs_is_formatted:
inputs = utils.format_pop_level_inputs(inputs, self, mon_length)
return self.driver.build(formatted_inputs=inputs,
mon_length=mon_length,
return_code=return_code,
show_code=(self.show_code or show_code))
def register_constant_delay(self, key, size, delay_time):
if not hasattr(self, 'constant_delays'):
self.constant_delays = {}
if key in self.constant_delays:
raise errors.ModelDefError(
f'"{key}" has been registered as an constant delay.')
self.constant_delays[key] = delays.ConstantDelay(size, delay_time)
return self.constant_delays[key]
def __init__(self, target, variables):
self.target = target
for mon_var in variables:
if not hasattr(target, mon_var):
raise errors.ModelDefError(
f"Item {mon_var} isn't defined in model {target}, "
f"so it can not be monitored.")
item_names = []
mon_indices = []
item_content = {}
if variables is not None:
if isinstance(variables, (list, tuple)):
for mon_var in variables:
if isinstance(mon_var, str):
var_data = getattr(target, mon_var)
mon_key = mon_var
mon_idx = None
mon_shape = ops.shape(var_data)
elif isinstance(mon_var, (tuple, list)):
mon_key = mon_var[0]
var_data = getattr(target, mon_key)
mon_idx = mon_var[1]
mon_shape = ops.shape(mon_idx) # TODO: matrix index
else:
raise errors.ModelUseError(
f'Unknown monitor item: {str(mon_var)}')
item_names.append(mon_key)
mon_indices.append(mon_idx)
dtype = var_data.dtype if hasattr(var_data,
'dtype') else None
item_content[mon_var] = ops.zeros((1, ) + mon_shape,
dtype=dtype)
elif isinstance(variables, dict):
for k, v in variables.items():
item_names.append(k)
mon_indices.append(v)
if v is None:
shape = ops.shape(getattr(target, k))
else:
shape = ops.shape(v)
val_data = getattr(target, k)
dtype = val_data.dtype if hasattr(val_data,
'dtype') else None
item_content[k] = ops.zeros((1, ) + shape, dtype=dtype)
else:
raise errors.ModelUseError(
f'Unknown monitors type: {type(variables)}')
self.ts = None
self.item_names = item_names
self.item_indices = mon_indices
self.item_contents = item_content
self.num_item = len(item_content)
def get_steps_func(self, show_code=False):
for func_name, step in self.steps.items():
class_args, arguments = utils.get_args(step)
host_name = self.host.name
calls = []
for arg in arguments:
if hasattr(self.host, arg):
calls.append(f'{host_name}.{arg}')
elif arg in backend.SYSTEM_KEYWORDS:
calls.append(arg)
else:
raise errors.ModelDefError(
f'Step function "{func_name}" of {self.host} '
f'define an unknown argument "{arg}" which is not '
f'an attribute of {self.host} nor the system keywords '
f'{backend.SYSTEM_KEYWORDS}.')
self.formatted_funcs[func_name] = {
'func': step,
'scope': {
host_name: self.host
},
'call': [f'{host_name}.{func_name}({", ".join(calls)})']
}
def transform_integrals_to_model(integrals):
from brainpy.integrators import sympy_analysis
if callable(integrals):
integrals = [integrals]
all_scope = dict()
all_variables = set()
all_parameters = set()
analyzers = []
for integral in integrals:
# integral function
if Dispatcher is not None and isinstance(integral, Dispatcher):
integral = integral.py_func
else:
integral = integral
# original function
f = integral.origin_f
if Dispatcher is not None and isinstance(f, Dispatcher):
f = f.py_func
func_name = f.__name__
# code scope
closure_vars = inspect.getclosurevars(f)
code_scope = dict(closure_vars.nonlocals)
code_scope.update(dict(closure_vars.globals))
# separate variables
analysis = ast_analysis.separate_variables(f)
variables_for_returns = analysis['variables_for_returns']
expressions_for_returns = analysis['expressions_for_returns']
for vi, (key, vars) in enumerate(variables_for_returns.items()):
variables = []
for v in vars:
if len(v) > 1:
raise ValueError(
'Cannot analyze multi-assignment code line.')
variables.append(v[0])
expressions = expressions_for_returns[key]
var_name = integral.variables[vi]
DE = sympy_analysis.SingleDiffEq(var_name=var_name,
variables=variables,
expressions=expressions,
derivative_expr=key,
scope=code_scope,
func_name=func_name)
analyzers.append(DE)
# others
for var in integral.variables:
if var in all_variables:
raise errors.ModelDefError(
f'Variable {var} has been defined before. Cannot group '
f'this integral as a dynamic system.')
all_variables.add(var)
all_parameters.update(integral.parameters)
all_scope.update(code_scope)
return DynamicModel(integrals=integrals,
analyzers=analyzers,
variables=list(all_variables),
parameters=list(all_parameters),
scopes=all_scope)
def class2func(cls_func, host, func_name=None, show_code=False):
"""Transform the function in a class into the ordinary function which is
compatible with the Numba JIT compilation.
Parameters
----------
cls_func : function
The function of the instantiated class.
func_name : str
The function name. If not given, it will get the function by `cls_func.__name__`.
show_code : bool
Whether show the code.
Returns
-------
new_func : function
The transformed function.
"""
class_arg, arguments = utils.get_args(cls_func)
func_name = cls_func.__name__ if func_name is None else func_name
host_name = host.name
# arguments 1
calls = []
for arg in arguments:
if hasattr(host, arg):
calls.append(f'{host_name}.{arg}')
elif arg in backend.SYSTEM_KEYWORDS:
calls.append(arg)
else:
raise errors.ModelDefError(
f'Step function "{func_name}" of {host} '
f'define an unknown argument "{arg}" which is not '
f'an attribute of {host} nor the system keywords '
f'{backend.SYSTEM_KEYWORDS}.')
# analysis
analyzed_results = analyze_step_func(host=host, f=cls_func)
delay_call = analyzed_results['delay_call']
# code_string = analyzed_results['code_string']
main_code = analyzed_results['code_string']
code_scope = analyzed_results['code_scope']
self_data_in_right = analyzed_results['self_data_in_right']
self_data_without_index_in_left = analyzed_results[
'self_data_without_index_in_left']
self_data_with_index_in_left = analyzed_results[
'self_data_with_index_in_left']
# main_code = get_func_body_code(code_string)
num_indent = get_num_indent(main_code)
data_need_pass = sorted(
list(set(self_data_in_right + self_data_with_index_in_left)))
data_need_return = self_data_without_index_in_left
# check delay
replaces_early = {}
replaces_later = {}
if len(delay_call) > 0:
for delay_ in delay_call.values():
# delay_ = dict(type=calls[-1],
# args=args,
# keywords=keywords,
# kws_append=kws_append,
# func=func,
# org_call=org_call,
# rep_call=rep_call,
# data_need_pass=data_need_pass)
if delay_['type'] == 'push':
if len(delay_['args'] + delay_['keywords']) == 2:
func = numba.njit(delay.push_type2)
elif len(delay_['args'] + delay_['keywords']) == 1:
func = numba.njit(delay.push_type1)
else:
raise ValueError(f'Unknown delay push. {delay_}')
else:
if len(delay_['args'] + delay_['keywords']) == 1:
func = numba.njit(delay.pull_type1)
elif len(delay_['args'] + delay_['keywords']) == 0:
func = numba.njit(delay.pull_type0)
else:
raise ValueError(f'Unknown delay pull. {delay_}')
delay_call_name = delay_['func']
data_need_pass.remove(delay_call_name)
data_need_pass.extend(delay_['data_need_pass'])
replaces_early[delay_['org_call']] = delay_['rep_call']
replaces_later[delay_call_name] = delay_call_name.replace('.', '_')
code_scope[delay_call_name.replace('.', '_')] = func
for target, dest in replaces_early.items():
main_code = main_code.replace(target, dest)
# main_code = tools.word_replace(main_code, replaces_early)
# arguments 2: data need pass
new_args = arguments + []
for data in sorted(set(data_need_pass)):
splits = data.split('.')
replaces_later[data] = data.replace('.', '_')
obj = host
for attr in splits[1:]:
obj = getattr(obj, attr)
if callable(obj):
code_scope[data.replace('.', '_')] = obj
continue
new_args.append(data.replace('.', '_'))
calls.append('.'.join([host_name] + splits[1:]))
# data need return
assigns = []
returns = []
for data in data_need_return:
splits = data.split('.')
assigns.append('.'.join([host_name] + splits[1:]))
returns.append(data.replace('.', '_'))
replaces_later[data] = data.replace('.', '_')
# code scope
code_scope[host_name] = host
# codes
header = f'def new_{func_name}({", ".join(new_args)}):\n'
main_code = header + tools.indent(main_code, spaces_per_tab=2)
if len(returns):
main_code += f'\n{" " * num_indent + " "}return {", ".join(returns)}'
main_code = tools.word_replace(main_code, replaces_later)
if show_code:
print(main_code)
print(code_scope)
print()
# recompile
exec(compile(main_code, '', 'exec'), code_scope)
func = code_scope[f'new_{func_name}']
func = numba.jit(**NUMBA_PROFILE)(func)
return func, calls, assigns
