"""
This module defines the `StateUpdateMethod` class that acts as a base class for
all stateupdaters and allows to register stateupdaters so that it is able to
return a suitable stateupdater object for a given set of equations. This is used
for example in `NeuronGroup` when no state updater is given explicitly.
"""
import time
from abc import ABCMeta, abstractmethod
from collections.abc import Iterable
from brian2.utils.caching import cached
from brian2.utils.logger import get_logger
__all__ = ["StateUpdateMethod"]
logger = get_logger(__name__)
[docs]
class UnsupportedEquationsException(Exception):
pass
[docs]
class StateUpdateMethod(metaclass=ABCMeta):
stateupdaters = dict()
[docs]
@abstractmethod
def __call__(self, equations, variables=None, method_options=None):
"""
Generate abstract code from equations. The method also gets the
the variables because some state updaters have to check whether
variable names reflect other state variables (which can change from
timestep to timestep) or are external values (which stay constant during
a run) For convenience, this arguments are optional -- this allows to
directly see what code a state updater generates for a set of equations
by simply writing ``euler(eqs)``, for example.
Parameters
----------
equations : `Equations`
The model equations.
variables : dict, optional
The `Variable` objects for the model variables.
method_options : dict, optional
Additional options specific to the state updater.
Returns
-------
code : str
The abstract code performing a state update step.
"""
pass
[docs]
@staticmethod
def register(name, stateupdater):
"""
Register a state updater. Registered state updaters can be referred to
via their name.
Parameters
----------
name : str
A short name for the state updater (e.g. `'euler'`)
stateupdater : `StateUpdaterMethod`
The state updater object, e.g. an `ExplicitStateUpdater`.
"""
# only deal with lower case names -- we don't want to have 'Euler' and
# 'euler', for example
name = name.lower()
if name in StateUpdateMethod.stateupdaters:
raise ValueError(
f"A stateupdater with the name '{name}' has already been registered"
)
if not isinstance(stateupdater, StateUpdateMethod):
raise ValueError(
f"Given stateupdater of type {type(stateupdater)} does "
"not seem to be a valid stateupdater."
)
StateUpdateMethod.stateupdaters[name] = stateupdater
[docs]
@staticmethod
@cached
def apply_stateupdater(
equations, variables, method, method_options=None, group_name=None
):
"""
apply_stateupdater(equations, variables, method, method_options=None, group_name=None)
Applies a given state updater to equations. If a `method` is given, the
state updater with the given name is used or if is a callable, then it
is used directly. If a `method` is a list of names, all the
methods will be tried until one that doesn't raise an
`UnsupportedEquationsException` is found.
Parameters
----------
equations : `Equations`
The model equations.
variables : `dict`
The dictionary of `Variable` objects, describing the internal
model variables.
method : {callable, str, list of str}
A callable usable as a state updater, the name of a registered
state updater or a list of names of state updaters.
Returns
-------
abstract_code : str
The code integrating the given equations.
"""
if isinstance(method, Iterable) and not isinstance(method, str):
the_method = None
start_time = time.time()
for one_method in method:
try:
one_method_start_time = time.time()
code = StateUpdateMethod.apply_stateupdater(
equations, variables, one_method, group_name=group_name
)
the_method = one_method
one_method_time = time.time() - one_method_start_time
break
except UnsupportedEquationsException:
pass
except TypeError:
raise TypeError(
"Each element in the list of methods has "
"to be a string or a callable, got "
f"{type(one_method)}."
)
total_time = time.time() - start_time
if the_method is None:
raise ValueError(
"No stateupdater that is suitable for the "
"given equations has been found."
)
# If only one method was tried
if method[0] == the_method:
timing = f"took {one_method_time:.2f}s"
else:
timing = (
f"took {one_method_time:.2f}s, trying other methods took "
f"{total_time - one_method_time:.2f}s"
)
if group_name is not None:
msg_text = (
"No numerical integration method specified for group "
f"'{group_name}', using method '{the_method}' ({timing})."
)
else:
msg_text = (
"No numerical integration method specified, "
f"using method '{the_method}' ({timing})."
)
logger.info(msg_text, "method_choice")
else:
if callable(method):
# if this is a standard state updater, i.e. if it has a
# can_integrate method, check this method and raise a warning if it
# claims not to be applicable.
stateupdater = method
method = getattr(
stateupdater, "__name__", repr(stateupdater)
) # For logging, get a nicer name
elif isinstance(method, str):
method = method.lower() # normalize name to lower case
stateupdater = StateUpdateMethod.stateupdaters.get(method, None)
if stateupdater is None:
raise ValueError(
"No state updater with the name '{method}' is known."
)
else:
raise TypeError(
"method argument has to be a string, a "
"callable, or an iterable of such objects. "
f"Got {type(method)}"
)
start_time = time.time()
code = stateupdater(equations, variables, method_options)
method_time = time.time() - start_time
timing = "took %.2fs" % method_time
if group_name is not None:
logger.debug(
f"Group {group_name}: using numerical integration "
f"method {method} ({timing})",
"method_choice",
)
else:
logger.debug(
f"Using numerical integration method: {method} f({{timing}})",
"method_choice",
)
return code