"""
Module providing `NumpyCodeObject`.
"""
import sys
from collections.abc import Iterable
import numpy as np
from brian2.core.base import BrianObjectException
from brian2.core.preferences import prefs, BrianPreference
from brian2.core.variables import (DynamicArrayVariable, ArrayVariable,
AuxiliaryVariable, Subexpression)
from brian2.core.functions import Function
from ...codeobject import CodeObject, constant_or_scalar, check_compiler_kwds
from ...templates import Templater
from ...generators.numpy_generator import NumpyCodeGenerator
from ...targets import codegen_targets
__all__ = ['NumpyCodeObject']
# Preferences
prefs.register_preferences(
'codegen.runtime.numpy',
'Numpy runtime codegen preferences',
discard_units = BrianPreference(
default=False,
docs="""
Whether to change the namespace of user-specifed functions to remove
units.
"""
)
)
[docs]class LazyArange(Iterable):
"""
A class that can be used as a `~numpy.arange` replacement (with an implied
step size of 1) but does not actually create an array of values until
necessary. It is somewhat similar to the ``range()`` function in Python 3,
but does not use a generator. It is tailored to a special use case, the
``_vectorisation_idx`` variable in numpy templates, and not meant for
general use. The ``_vectorisation_idx`` is used for stateless function
calls such as ``rand()`` and for the numpy codegen target determines the
number of values produced by such a call. This will often be the number of
neurons or synapses, and this class avoids creating a new array of that size
at every code object call when all that is needed is the *length* of the
array.
Examples
--------
>>> from brian2.codegen.runtime.numpy_rt.numpy_rt import LazyArange
>>> ar = LazyArange(10)
>>> len(ar)
10
>>> len(ar[:5])
5
>>> type(ar[:5])
<class 'brian2.codegen.runtime.numpy_rt.numpy_rt.LazyArange'>
>>> ar[5]
5
>>> for value in ar[3:7]:
... print(value)
...
3
4
5
6
>>> len(ar[np.array([1, 2, 3])])
3
"""
def __init__(self, stop, start=0, indices=None):
self.start = start
self.stop = stop
self.indices = indices
def __len__(self):
if self.indices is None:
return self.stop - self.start
else:
return len(self.indices)
def __getitem__(self, item):
if isinstance(item, slice):
if self.indices is None:
start, stop, step = item.start, item.stop, item.step
if step not in [None, 1]:
raise NotImplementedError("Step should be 1")
if start is None:
start = 0
if stop is None:
stop = len(self)
return LazyArange(start=self.start+start,
stop=min([self.start+stop, self.stop]))
else:
raise NotImplementedError("Cannot slice LazyArange with indices")
elif isinstance(item, np.ndarray):
if item.dtype == np.dtype(bool):
item = np.nonzero(item)[0] # convert boolean array into integers
if len(item) == 0:
return np.array([], dtype=np.int32)
if np.min(item) < 0 or np.max(item) > len(self):
raise IndexError("Indexing array contains out-of-bounds values")
return LazyArange(start=self.start, stop=self.stop, indices=item)
elif isinstance(item, int):
if self.indices is None:
index = self.start + item
if index >= self.stop:
raise IndexError(index)
return index
else:
return self.indices[item]
else:
raise TypeError("Can only index with integer, numpy array, or slice.")
def __iter__(self):
if self.indices is None:
return iter(np.arange(self.start, self.stop))
else:
return iter(self.indices)
# Allow conversion to a proper array with np.array(...)
def __array__(self, dtype=None):
if self.indices is None:
return np.arange(self.start, self.stop)
else:
return self.indices + self.start
# Allow basic arithmetics (used when shifting stuff for subgroups)
def __add__(self, other):
if isinstance(other, int):
return LazyArange(start=self.start + other, stop=self.stop + other)
else:
return NotImplemented
def __radd__(self, other):
return self.__add__(other)
def __sub__(self, other):
if isinstance(other, int):
return LazyArange(start=self.start - other, stop=self.stop - other)
else:
return NotImplemented
[docs]class NumpyCodeObject(CodeObject):
"""
Execute code using Numpy
Default for Brian because it works on all platforms.
"""
templater = Templater('brian2.codegen.runtime.numpy_rt', '.py_',
env_globals={'constant_or_scalar': constant_or_scalar})
generator_class = NumpyCodeGenerator
class_name = 'numpy'
def __init__(self, owner, code, variables, variable_indices,
template_name, template_source, compiler_kwds,
name='numpy_code_object*'):
check_compiler_kwds(compiler_kwds, [],
'numpy')
from brian2.devices.device import get_device
self.device = get_device()
self.namespace = {'_owner': owner,
# TODO: This should maybe go somewhere else
'logical_not': np.logical_not}
CodeObject.__init__(self, owner, code, variables, variable_indices,
template_name, template_source,
compiler_kwds=compiler_kwds, name=name)
self.variables_to_namespace()
[docs] @classmethod
def is_available(cls):
# no test necessary for numpy
return True
[docs] def variables_to_namespace(self):
# Variables can refer to values that are either constant (e.g. dt)
# or change every timestep (e.g. t). We add the values of the
# constant variables here and add the names of non-constant variables
# to a list
# A list containing tuples of name and a function giving the value
self.nonconstant_values = []
for name, var in self.variables.items():
if isinstance(var, (AuxiliaryVariable, Subexpression)):
continue
try:
if not hasattr(var, 'get_value'):
raise TypeError()
value = var.get_value()
except TypeError:
# Either a dummy Variable without a value or a Function object
if isinstance(var, Function):
impl = var.implementations[self.__class__].get_code(self.owner)
self.namespace[name] = impl
else:
self.namespace[name] = var
continue
if isinstance(var, ArrayVariable):
self.namespace[self.generator_class.get_array_name(var)] = value
if var.scalar and var.constant:
self.namespace[name] = value[0]
else:
self.namespace[name] = value
if isinstance(var, DynamicArrayVariable):
dyn_array_name = self.generator_class.get_array_name(var,
access_data=False)
self.namespace[dyn_array_name] = self.device.get_value(var,
access_data=False)
# Also provide the Variable object itself in the namespace (can be
# necessary for resize operations, for example)
self.namespace[f"_var_{name}"] = var
# There is one type of objects that we have to inject into the
# namespace with their current value at each time step: dynamic
# arrays that change in size during runs (i.e. not synapses but
# e.g. the structures used in monitors)
if (isinstance(var, DynamicArrayVariable) and
var.needs_reference_update):
self.nonconstant_values.append((self.generator_class.get_array_name(var,
self.variables),
var.get_value))
[docs] def update_namespace(self):
# update the values of the non-constant values in the namespace
for name, func in self.nonconstant_values:
self.namespace[name] = func()
[docs] def compile_block(self, block):
code = getattr(self.code, block, '').strip()
if not code or 'EMPTY_CODE_BLOCK' in code:
return None
return compile(code, '(string)', 'exec')
[docs] def run_block(self, block):
compiled_code = self.compiled_code[block]
if not compiled_code:
return
try:
exec(compiled_code, self.namespace)
except Exception as exc:
code = getattr(self.code, block)
message = (f"An exception occured during the execution of the "
f"'{block}' block of code object {self.name}.\n")
lines = code.split('\n')
message += 'The error was raised in the following line:\n'
_, _, tb = sys.exc_info()
tb = tb.tb_next # Line in the code object's code
message += f"{lines[tb.tb_lineno - 1]}\n"
raise BrianObjectException(message, self.owner) from exc
# output variables should land in the variable name _return_values
if '_return_values' in self.namespace:
return self.namespace['_return_values']
codegen_targets.add(NumpyCodeObject)