Repository: brian-team/brian2modelfitting (language: python)

Issue:
Refactor `Inferencer`/`Fitter`
Both classes have a lot in common (everything related to the Brian simulation), and should inherit from a common superclass to share the code. Another possible approach would be to put all the shared code in a separate class that neither `Fitter` nor `Inferencer` inherits from, and delegate to an object of that class (e.g. everything related to the simulation could be in a `BrianSimulator` class, and both `Fitter` and `Inferencer` would do something like `self.brian_simulator = BrianSimulator(...)` in their `__init__` and then call `self.brian_simulator...` for all simulation-related tasks.) Actually, there is already a `Simulator` class (with implementations for runtime and standalone) – maybe some more code could simply move into this class.


Discussion / hints:

Fix patch:
diff --git a/brian2modelfitting/base.py b/brian2modelfitting/base.py
new file mode 100644
index 0000000..4a20e15
--- /dev/null
+++ b/brian2modelfitting/base.py
@@ -0,0 +1,272 @@
+"""
+Basic functions such as input handling, shared between fitting and inference classes.
+"""
+from collections.abc import Mapping
+import numbers
+
+from brian2 import get_logger
+from brian2.core.namespace import get_local_namespace
+from brian2.core.functions import Function
+from brian2.devices import RuntimeDevice, get_device, device
+from brian2.devices.cpp_standalone.device import CPPStandaloneDevice
+from brian2.units.fundamentalunits import (Quantity,
+                                           DIMENSIONLESS,
+                                           fail_for_dimension_mismatch)
+
+from .simulator import RuntimeSimulator, CPPStandaloneSimulator
+
+logger = get_logger(__name__)
+
+
+def get_full_namespace(additional_namespace, level=0):
+    # Get the local namespace with all the values that could be relevant
+    # in principle -- by filtering things out, we avoid circular loops
+    namespace = {key: value
+                 for key, value in get_local_namespace(level=level + 1).items()
+                 if (not key.startswith('_') and
+                     isinstance(value, (Quantity, numbers.Number, Function)))}
+    namespace.update(additional_namespace)
+
+    return namespace
+
+
+def handle_input_args(input_dict_or_arr, input_var, model):
+    """
+    Handle the input argument of the fit/inference methods.
+
+    Parameters
+    ----------
+    input_dict_or_arr : dict or `np.ndarray`, optional
+        A dictionary mapping the name of the input variable to the input.
+        Note that only a single input is currently supported.
+        When used together with the ``input_var`` argument (deprecated),
+        this can also be a single array.
+    input_var : str, optional
+        The name of the input variable. Deprecated, use the ``input_dict``
+        argument instead.
+    model : `~brian2.equations.equations.Equations`
+        The equations of the model to fit/infer.
+
+    Returns
+    -------
+    input_arr, input_var : (`~numpy.ndarray`, str)
+        The input array and the name of the input variable.
+    """
+    # Support deprecated legacy syntax of input_var + input or the new
+    # syntax with a dictionary as input
+    if input_var is not None:
+        logger.warn("Use the 'input' argument with a dictionary instead "
+                    "of giving the name as 'input_var'",
+                    name_suffix='deprecated_input_var')
+        if isinstance(input_dict_or_arr, Mapping) and input_var not in input_dict_or_arr:
+            raise ValueError("Name given as 'input_var' and key in "
+                             "'input' dictionary do not match.")
+    else:
+        if not isinstance(input_dict_or_arr, Mapping):
+            raise TypeError("'input' argument has to be a dictionary "
+                            "mapping the name of the input variable to the "
+                            "input.")
+        if len(input_dict_or_arr) > 1:
+            raise NotImplementedError("Only a single input is currently "
+                                      "supported.")
+        input_var = list(input_dict_or_arr.keys())[0]
+
+    if isinstance(input_dict_or_arr, Mapping):
+        input_arr = input_dict_or_arr[input_var]
+    else:
+        input_arr = input_dict_or_arr
+
+    if input_var != "spikes" and input_var not in model.identifiers:
+        raise NameError(f"{input_var} is not an identifier in the model")
+    return input_arr, input_var
+
+
+def handle_output_args(output_dict, output_vars, model):
+    """
+    Handle the output argument of the fit/inference methods.
+
+    Parameters
+    ----------
+    output_dict : dict, optional
+        A dictionary mapping the name of the output variables to the output.
+    output_vars : list of str, optional
+        The name of the output variables. Deprecated, use the ``output_dict``
+        argument instead.
+    model : `~brian2.equations.equations.Equations`
+        The equations of the model to fit/infer.
+
+    Returns
+    -------
+    outputs, output_var : (list, str)
+        The output list and the name of the output variables.
+    """
+    # Support deprecated legacy syntax of output_var + output or the new
+    # syntax with a dictionary as output
+    if output_vars is not None:
+        logger.warn("Use the 'output' argument with a dictionary instead "
+                    "of giving the name as 'output_var'",
+                    name_suffix='deprecated_output_var')
+        if isinstance(output_vars, str):
+            output_vars = [output_vars]
+
+        if isinstance(output_dict, Mapping):
+            if set(output_vars) != set(output_dict.keys()):
+                raise ValueError("Names given as 'output_var' and keys "
+                                 "in 'output' dictionary do not match.")
+            output_list = list(output_dict.values())
+        elif not isinstance(output_dict, list):
+            output_list = [output_dict]
+        else:
+            output_list = output_dict
+    else:
+        if not isinstance(output_dict, Mapping):
+            raise TypeError("'output' argument has to be a dictionary "
+                            "mapping the name of the input variable to the "
+                            "input.")
+        output_vars = list(output_dict.keys())
+        output_list = list(output_dict.values())
+
+    for o_var in output_vars:
+        if o_var != 'spikes' and o_var not in model.names:
+            raise NameError(f"{o_var} is not a model variable")
+
+    return output_list, output_vars
+
+
+def output_dims(output_list, output_var, model):
+    """
+    Verify the output dimensions.
+
+    Parameters
+    ----------
+    output_list : list of `~numpy.ndarray`
+        List of output values, for each variable in ``output_var``.
+    output_var : list of str
+        List of output variable names.
+    model : `~brian2.equations.equations.Equations`
+        The equations of the model to fit/infer.
+
+    Returns
+    -------
+    dims : list of `~brian2.units.fundamentalunits.Dimension`
+        The dimensions of the output variables.
+    """
+    dims = []
+    for o_var, out in zip(output_var, output_list):
+        if o_var == 'spikes':
+            dims.append(DIMENSIONLESS)
+        else:
+            dims.append(model[o_var].dim)
+            fail_for_dimension_mismatch(out, dims[-1],
+                                        'The provided target values '
+                                        '("output") need to have the same '
+                                        'units as the variable '
+                                        '{}'.format(o_var))
+    return dims
+
+
+def input_equations(input_var, input_dim):
+    """
+    Define the input variable for the equations of the model.
+
+    Parameters
+    ----------
+    input_arr : `~numpy.ndarray`
+        The input array.
+    input_var : str
+        The name of the input variable.
+
+    Returns
+    -------
+    input_eqs : str
+        The equations setting the input variable.
+    """
+    input_dim = '1' if input_dim is DIMENSIONLESS else repr(input_dim)
+    input_eqs = "{} = input_var(t, i % n_traces) : {}".format(input_var,
+                                                              input_dim)
+    return input_eqs
+
+
+def output_equations(output_var, output_dim):
+    """
+    Make the output variables available within the model equations. This
+    can be useful for approaches that couple the system to the target values.
+
+    Parameters
+    ----------
+    output_var : list of str
+        The names of the output variables.
+    output_dim : list of `~brian2.units.fundamentalunits.Dimension`
+        The dimensions of the output variables.
+
+    Returns
+    -------
+    output_eqs : str
+        The equations setting a variable named ``<output_variable>_target``.
+    """
+    counter = 0
+    output_eqs = []
+    for o_var, o_dim in zip(output_var, output_dim):
+        if o_var != 'spikes':
+            counter += 1
+            # For approaches that couple the system to the target values,
+            # provide a convenient variable
+            output_expr = f'output_var_{counter}(t, i % n_traces)'
+            output_dim = ('1' if o_dim is DIMENSIONLESS
+                          else repr(o_dim))
+            output_eqs.append("{}_target = {} : {}".format(o_var,
+                                                           output_expr,
+                                                           output_dim))
+    return "\n".join(output_eqs)
+
+
+def handle_param_init(param_init, model):
+    """
+    Verify that the parameters given in ``param_init`` are valid.
+
+    Parameters
+    ----------
+    param_init : dict or None
+        A dictionary mapping the name of the parameters to their initial value.
+    model : `~brian2.equations.equations.Equations`
+        The equations of the model to fit/infer.
+
+    Returns
+    -------
+    param_init : dict
+        The dictionary mapping the name of the parameters to their initial values.
+    """
+    if not param_init:
+        param_init = {}
+    for param, val in param_init.items():
+        if not (param in model.diff_eq_names or
+                param in model.parameter_names):
+            raise ValueError(f"{param} is not a model variable or a "
+                             "parameter in the model")
+    return param_init
+
+
+def setup_fit():
+    """
+    Function sets up simulator in one of the two available modes: runtime
+    or standalone. The `.Simulator` that will be used depends on the currently
+    set `.Device`. In the case of `.CPPStandaloneDevice`, the device will also
+    be reset if it has already run a simulation.
+
+    Returns
+    -------
+    simulator : `.Simulator`
+    """
+    simulators = {
+        CPPStandaloneDevice: CPPStandaloneSimulator(),
+        RuntimeDevice: RuntimeSimulator()
+    }
+    if isinstance(get_device(), CPPStandaloneDevice):
+        if device.has_been_run is True:
+            build_options = dict(device.build_options)
+            get_device().reinit()
+            get_device().activate(**build_options)
+    simulator = [sim for dev, sim in simulators.items()
+                 if isinstance(get_device(), dev)]
+    assert len(simulator) == 1, f"Found {len(simulator)} simulators for device {get_device().__class__.__name__}"
+    return simulator[0]
diff --git a/brian2modelfitting/fitter.py b/brian2modelfitting/fitter.py
index 961d6cb..884df41 100644
--- a/brian2modelfitting/fitter.py
+++ b/brian2modelfitting/fitter.py
@@ -5,23 +5,26 @@
 from typing import Sequence, Mapping
 
 import sympy
-from numpy import (ones, array, arange, concatenate, mean, argmin, nanargmin,
-                   reshape, zeros, sqrt, ndarray, broadcast_to, sum, cumsum,
-                   hstack, tile, repeat)
+import numpy as np
 from scipy.optimize import least_squares
 
 from brian2.parsing.sympytools import sympy_to_str, str_to_sympy
 from brian2.units.fundamentalunits import DIMENSIONLESS, get_dimensions, fail_for_dimension_mismatch
-from brian2.utils.stringtools import get_identifiers
 
-from brian2 import (NeuronGroup, defaultclock, get_device, Network,
-                    StateMonitor, SpikeMonitor, second, get_local_namespace,
+from brian2 import (NeuronGroup, defaultclock, second,
                     Quantity, get_logger, Expression, ms)
 from brian2.input import TimedArray
-from brian2.equations.equations import Equations, SUBEXPRESSION, SingleEquation
-from brian2.devices import device, RuntimeDevice
-from brian2.devices.cpp_standalone.device import CPPStandaloneDevice
-from brian2.core.functions import Function
+from brian2.equations.equations import Equations, SingleEquation
+
+from .base import (get_full_namespace,
+                   handle_input_args,
+                   handle_output_args,
+                   handle_param_init,
+                   input_equations,
+                   output_equations,
+                   output_dims, setup_fit)
+
+from .sensitivity import get_sensitivity_equations, get_sensitivity_init
 from .simulator import RuntimeSimulator, CPPStandaloneSimulator
 from .metric import Metric, SpikeMetric, TraceMetric, MSEMetric, GammaFactor, normalize_weights
 from .optimizer import Optimizer, NevergradOptimizer
@@ -41,7 +44,7 @@ def get_param_dic(params, param_names, n_traces, n_samples):
 
     for param_combination in params:
         for param_name, param_value in param_combination.items():
-            d[param_name].extend(repeat(param_value, n_traces))
+            d[param_name].extend(np.repeat(param_value, n_traces))
 
     return d
 
@@ -59,172 +62,11 @@ def get_spikes(monitor, n_samples, n_traces):
         sample_spikes = []
         for trace in range(n_traces):
             i += 1
-            sample_spikes.append(array(spike_trains[i], copy=False))
+            sample_spikes.append(np.array(spike_trains[i], copy=False))
         spikes.append(sample_spikes)
     return spikes
 
 
-def get_full_namespace(additional_namespace, level=0):
-    # Get the local namespace with all the values that could be relevant
-    # in principle -- by filtering things out, we avoid circular loops
-    namespace = {key: value
-                 for key, value in get_local_namespace(level=level + 1).items()
-                 if (not key.startswith('_') and
-                     isinstance(value, (Quantity, numbers.Number, Function)))}
-    namespace.update(additional_namespace)
-
-    return namespace
-
-
-def setup_fit():
-    """
-    Function sets up simulator in one of the two available modes: runtime
-    or standalone. The `.Simulator` that will be used depends on the currently
-    set `.Device`. In the case of `.CPPStandaloneDevice`, the device will also
-    be reset if it has already run a simulation.
-
-    Returns
-    -------
-    simulator : `.Simulator`
-    """
-    simulators = {
-        CPPStandaloneDevice: CPPStandaloneSimulator(),
-        RuntimeDevice: RuntimeSimulator()
-    }
-    if isinstance(get_device(), CPPStandaloneDevice):
-        if device.has_been_run is True:
-            build_options = dict(device.build_options)
-            get_device().reinit()
-            get_device().activate(**build_options)
-    simulator = [sim for dev, sim in simulators.items()
-                 if isinstance(get_device(), dev)]
-    assert len(simulator) == 1, f"Found {len(simulator)} simulators for device {get_device().__class__.__name__}"
-    return simulator[0]
-
-
-def get_sensitivity_equations(group, parameters, namespace=None, level=1,
-                              optimize=True):
-    """
-    Get equations for sensitivity variables.
-
-    Parameters
-    ----------
-    group : `NeuronGroup`
-        The group of neurons that will be simulated.
-    parameters : list of str
-        Names of the parameters that are fit.
-    namespace : dict, optional
-        The namespace to use.
-    level : `int`, optional
-        How much farther to go down in the stack to find the namespace.
-    optimize : bool, optional
-        Whether to remove sensitivity variables from the equations that do
-        not evolve if initialized to zero (e.g. ``dS_x_y/dt = -S_x_y/tau``
-        would be removed). This avoids unnecessary computation but will fail
-        in the rare case that such a sensitivity variable needs to be
-        initialized to a non-zero value. Defaults to ``True``.
-
-    Returns
-    -------
-    sensitivity_eqs : `Equations`
-        The equations for the sensitivity variables.
-    """
-    if namespace is None:
-        namespace = get_local_namespace(level)
-        namespace.update(group.namespace)
-
-    eqs = group.equations
-    diff_eqs = eqs.get_substituted_expressions(group.variables)
-    diff_eq_names = [name for name, _ in diff_eqs]
-
-    system = sympy.Matrix([str_to_sympy(diff_eq[1].code)
-                           for diff_eq in diff_eqs])
-    J = system.jacobian([str_to_sympy(d) for d in diff_eq_names])
-
-    sensitivity = []
-    sensitivity_names = []
-    for parameter in parameters:
-        F = system.jacobian([str_to_sympy(parameter)])
-        names = [str_to_sympy(f'S_{diff_eq_name}_{parameter}')
-                 for diff_eq_name in diff_eq_names]
-        sensitivity.append(J * sympy.Matrix(names) + F)
-        sensitivity_names.append(names)
-
-    new_eqs = []
-    for names, sensitivity_eqs, param in zip(sensitivity_names, sensitivity, parameters):
-        for name, eq, orig_var in zip(names, sensitivity_eqs, diff_eq_names):
-            unit = eqs[orig_var].dim / group.variables[param].dim
-            unit = repr(unit) if not unit.is_dimensionless else '1'
-            if optimize:
-                # Check if the equation stays at zero if initialized at zero
-                zeroed = eq.subs(name, sympy.S.Zero)
-                if zeroed == sympy.S.Zero:
-                    # No need to include equation as differential equation
-                    if unit == '1':
-                        new_eqs.append(f'{sympy_to_str(name)} = 0 : {unit}')
-                    else:
-                        new_eqs.append(f'{sympy_to_str(name)} = 0*{unit} : {unit}')
-                    continue
-            rhs = sympy_to_str(eq)
-            if rhs == '0':  # avoid unit mismatch
-                rhs = f'0*{unit}/second'
-            new_eqs.append('d{lhs}/dt = {rhs} : {unit}'.format(lhs=sympy_to_str(name),
-                                                               rhs=rhs,
-                                                               unit=unit))
-    new_eqs = Equations('\n'.join(new_eqs))
-    return new_eqs
-
-
-def get_sensitivity_init(group, parameters, param_init):
-    """
-    Calculate the initial values for the sensitivity parameters (necessary if
-    initial values are functions of parameters).
-
-    Parameters
-    ----------
-    group : `NeuronGroup`
-        The group of neurons that will be simulated.
-    parameters : list of str
-        Names of the parameters that are fit.
-    param_init : dict
-        The dictionary with expressions to initialize the model variables.
-
-    Returns
-    -------
-    sensitivity_init : dict
-        Dictionary of expressions to initialize the sensitivity
-        parameters.
-    """
-    sensitivity_dict = {}
-    for var_name, expr in param_init.items():
-        if not isinstance(expr, str):
-            continue
-        identifiers = get_identifiers(expr)
-        for identifier in identifiers:
-            if (identifier in group.variables
-                    and getattr(group.variables[identifier],
-                                'type', None) == SUBEXPRESSION):
-                raise NotImplementedError('Initializations that refer to a '
-                                          'subexpression are currently not '
-                                          'supported')
-            sympy_expr = str_to_sympy(expr)
-            for parameter in parameters:
-                diffed = sympy_expr.diff(str_to_sympy(parameter))
-                if diffed != sympy.S.Zero:
-                    if getattr(group.variables[parameter],
-                               'type', None) == SUBEXPRESSION:
-                        raise NotImplementedError('Sensitivity '
-                                                  f'S_{var_name}_{parameter} '
-                                                  'is initialized to a non-zero '
-                                                  'value, but it has been '
-                                                  'removed from the equations. '
-                                                  'Set optimize=False to avoid '
-                                                  'this.')
-                    init_expr = sympy_to_str(diffed)
-                    sensitivity_dict[f'S_{var_name}_{parameter}'] = init_expr
-    return sensitivity_dict
-
-
 class Fitter(metaclass=abc.ABCMeta):
     """
     Base Fitter class for model fitting applications.
@@ -296,69 +138,28 @@ def __init__(self, dt, model, input, output, n_samples, input_var=None,
 
         if isinstance(model, str):
             model = Equations(model)
+        self.model = model
 
-        # Support deprecated legacy syntax of input_var + input or the new
-        # syntax with a dictionary as input
-        if input_var is not None:
-            logger.warn('Use the \'input\' argument with a dictionary instead '
-                        'of giving the name as \'input_var\'',
-                        name_suffix='deprecated_input_var')
-            if isinstance(input, Mapping) and input_var not in input:
-                raise ValueError('Name given as \'input_var\' and key in '
-                                 '\'input\' dictionary do not match.')
-        else:
-            if not isinstance(input, Mapping):
-                raise TypeError('\'input\' argument has to be a dictionary '
-                                'mapping the name of the input variable to the '
-                                'input.')
-            if len(input) > 1:
-                raise NotImplementedError('Only a single input is currently '
-                                          'supported.')
-            input_var = list(input.keys())[0]
-
-        if isinstance(input, Mapping):
-            input = input[input_var]
-
-        if input_var != 'spikes' and input_var not in model.identifiers:
-            raise NameError(f"{input_var} is not an identifier in the model")
-
-        # Support deprecated legacy syntax of output_var + input or the new
-        # syntax with a dictionary as output
-        if output_var is not None:
-            logger.warn('Use the \'output\' argument with a dictionary instead '
-                        'of giving the name as \'output_var\'',
-                        name_suffix='deprecated_output_var')
-            if isinstance(output_var, str):
-                output_var = [output_var]
-
-            if isinstance(output, Mapping):
-                if set(output_var) != set(output.keys()):
-                    raise ValueError('Names given as \'output_var\' and keys '
-                                     'in \'output\' dictionary do not match.')
-            elif not isinstance(output, list):
-                output = [output]
-        else:
-            if not isinstance(output, Mapping):
-                raise TypeError('\'output\' argument has to be a dictionary '
-                                'mapping the name of the input variable to the '
-                                'input.')
-            output_var = list(output.keys())
-            output = list(output.values())
+        input_arr, input_var = handle_input_args(input, input_var, model)
+        self.input = input_arr
+        input_eqs = input_equations(input_var, get_dimensions(input_arr))
+        self.model += input_eqs
 
-        for o_var in output_var:
-            if o_var != 'spikes' and o_var not in model.names:
-                raise NameError(f"{o_var} is not a model variable")
+        output, output_var = handle_output_args(output, output_var, model)
 
         self.output_var = output_var
         self.output = output
-        self.output_ = [array(o, copy=False) for o in output]
+        self.output_ = [np.array(o, copy=False) for o in output]
+        self.output_dim = output_dims(output, output_var, model)
+
+        self.model += output_equations(output_var, self.output_dim)
 
         self.dt = dt
 
         self.simulator = None
 
         self.parameter_names = model.parameter_names
-        self.n_traces, n_steps = input.shape
+        self.n_traces, n_steps = input_arr.shape
         self.duration = n_steps * dt
         # Sample size requested by user
         self._n_samples = n_samples
@@ -370,44 +171,10 @@ def __init__(self, dt, model, input, output, n_samples, input_var=None,
         self.refractory = refractory
         self.penalty = penalty
 
-        self.input = input
-
-        self.output_dim = []
-        for o_var, out in zip(self.output_var, output):
-            if o_var == 'spikes':
-                self.output_dim.append(DIMENSIONLESS)
-            else:
-                self.output_dim.append(model[o_var].dim)
-                fail_for_dimension_mismatch(out, self.output_dim[-1],
-                                            'The provided target values '
-                                            '("output") need to have the same '
-                                            'units as the variable '
-                                            '{}'.format(o_var))
-        self.model = model
-
         self.use_units = use_units
         self.iteration = 0
-        input_dim = get_dimensions(input)
-        input_dim = '1' if input_dim is DIMENSIONLESS else repr(input_dim)
-        input_eqs = "{} = input_var(t, i % n_traces) : {}".format(input_var,
-                                                                  input_dim)
-        self.model += input_eqs
 
-        counter = 0
-        for o_var, o_dim in zip(self.output_var, self.output_dim):
-            if o_var != 'spikes':
-                counter += 1
-                # For approaches that couple the system to the target values,
-                # provide a convenient variable
-                output_expr = f'output_var_{counter}(t, i % n_traces)'
-                output_dim = ('1' if o_dim is DIMENSIONLESS
-                              else repr(o_dim))
-                output_eqs = "{}_target = {} : {}".format(o_var,
-                                                          output_expr,
-                                                          output_dim)
-                self.model += output_eqs
-
-        input_traces = TimedArray(input.transpose(), dt=dt)
+        input_traces = TimedArray(input_arr.transpose(), dt=dt)
         self.input_traces = input_traces
 
         # initialization of attributes used later
@@ -420,60 +187,13 @@ def __init__(self, dt, model, input, output, n_samples, input_var=None,
         self.optimizer = None
         self.metric = None
 
-        if not param_init:
-            param_init = {}
-        for param, val in param_init.items():
-            if not (param in self.model.diff_eq_names or
-                    param in self.model.parameter_names):
-                raise ValueError("%s is not a model variable or a "
-                                 "parameter in the model" % param)
-        self.param_init = param_init
+        self.param_init = handle_param_init(param_init, model)
+
 
     @property
     def n_neurons(self):
         return self.n_traces * self.n_samples
 
-    def setup_simulator(self, network_name, n_neurons, output_var, param_init,
-                        calc_gradient=False, optimize=True, online_error=False,
-                        level=1):
-        simulator = setup_fit()
-
-        namespace = get_full_namespace({'input_var': self.input_traces,
-                                        'n_traces': self.n_traces},
-                                       level=level+1)
-        if hasattr(self, 't_start'):  # OnlineTraceFitter
-            namespace['t_start'] = self.t_start
-        counter = 0
-        for o_var, out in zip(self.output_var, self.output):
-            if self.output_var != 'spikes':
-                counter += 1
-                namespace[f'output_var_{counter}'] = TimedArray(out.transpose(),
-                                                                dt=self.dt)
-        neurons = self.setup_neuron_group(n_neurons, namespace,
-                                          calc_gradient=calc_gradient,
-                                          optimize=optimize,
-                                          online_error=online_error)
-        network = Network(neurons)
-        if isinstance(output_var, str):
-            output_var = [output_var]
-        if 'spikes' in output_var:
-            network.add(SpikeMonitor(neurons, name='spikemonitor'))
-
-        record_vars = [v for v in output_var if v != 'spikes']
-        if calc_gradient:
-            record_vars.extend([f'S_{out_var}_{p}'
-                                for p in self.parameter_names
-                                for out_var in self.output_var])
-        if len(record_vars):
-            network.add(StateMonitor(neurons, record_vars, record=True,
-                                     name='statemonitor', dt=self.dt))
-
-        if calc_gradient:
-            param_init = dict(param_init)
-            param_init.update(get_sensitivity_init(neurons, self.parameter_names,
-                                                   param_init))
-        simulator.initialize(network, param_init, name=network_name)
-        return simulator
 
     def setup_neuron_group(self, n_neurons, namespace, calc_gradient=False,
                            optimize=True, online_error=False, name='neurons'):
@@ -605,8 +325,8 @@ def optimization_iter(self, optimizer, metric, penalty):
         self.simulator.run(self.duration, d_param, self.parameter_names,
                            iteration=self.iteration)
 
-        raw_errors = array(self.calc_errors(metric))
-        errors = sum(raw_errors, axis=0)
+        raw_errors = np.array(self.calc_errors(metric))
+        errors = np.sum(raw_errors, axis=0)
 
         if penalty is not None:
             error_penalty = getattr(self.simulator.neurons, penalty + '_')
@@ -624,7 +344,7 @@ def optimization_iter(self, optimizer, metric, penalty):
 
             errors += error_penalty
         self._objective_errors_normed.extend(raw_errors.T.tolist())
-        unnormalized_raw_errors = array([metric.revert_normalization(err)
+        unnormalized_raw_errors = np.array([metric.revert_normalization(err)
                                          for metric, err in zip(metric, raw_errors)]).T.tolist()
         self._objective_errors.extend(unnormalized_raw_errors)
         optimizer.tell(parameters, errors)
@@ -742,11 +462,21 @@ def fit(self, optimizer, metric=None, n_rounds=1, callback='text',
         if (restart or
                 self.simulator is None or
                 self.simulator.current_net != 'fit'):
-            self.simulator = self.setup_simulator('fit', self.n_neurons,
-                                                  output_var=self.output_var,
-                                                  online_error=online_error,
-                                                  param_init=self.param_init,
-                                                  level=level+1)
+
+            namespace = self.create_namespace(level=level+1)
+            self.simulator = setup_fit()  # for standalone, it is important to do this before creating the neurons
+            neurons = self.setup_neuron_group(self.n_neurons,
+                                              namespace,
+                                              online_error=online_error)
+            record_vars = [v for v in self.output_var if v != 'spikes']
+            record_spikes = 'spikes' in self.output_var
+
+            self.simulator.setup('fit',
+                                 neurons=neurons,
+                                 record_vars=record_vars,
+                                 record_spikes=record_spikes,
+                                 param_init=self.param_init,
+                                 dt=self.dt)
 
         # Run Optimization Loop
         for index in range(n_rounds):
@@ -754,7 +484,7 @@ def fit(self, optimizer, metric=None, n_rounds=1, callback='text',
                                                                      self.metric,
                                                                      penalty)
             self.iteration += 1
-            best_idx = nanargmin(self.optimizer.errors)
+            best_idx = np.nanargmin(self.optimizer.errors)
             self._best_error = self.optimizer.errors[best_idx]
             self._best_objective_errors_normed = tuple(self._objective_errors_normed[best_idx])
             self._best_objective_errors = tuple(self._objective_errors[best_idx])
@@ -807,6 +537,18 @@ def fit(self, optimizer, metric=None, n_rounds=1, callback='text',
 
         return self.best_params, self.best_error
 
+    def create_namespace(self, level):
+        namespace = get_full_namespace({'input_var': self.input_traces,
+                                        'n_traces': self.n_traces},
+                                       level=level + 1)
+        counter = 0
+        for o_var, out in zip(self.output_var, self.output):
+            if self.output_var != 'spikes':
+                counter += 1
+                namespace[f'output_var_{counter}'] = TimedArray(out.transpose(),
+                                                                dt=self.dt)
+        return namespace
+
     def _verify_metric_argument(self, metric, metric_class=Metric):
         if isinstance(metric, Metric):
             metric = {varname: metric
@@ -925,7 +667,7 @@ def results(self, format='list', use_units=None):
         params = self.optimizer.tested_parameters
         if use_units:
             error_dim = self.metric[0].get_dimensions(self.output_dim[0])
-            errors = Quantity(array(self.optimizer.errors).flatten(),
+            errors = Quantity(np.array(self.optimizer.errors).flatten(),
                               dim=error_dim)
             if len(self.output_var) > 1:
                 # Add additional information for the raw errors
@@ -942,13 +684,13 @@ def results(self, format='list', use_units=None):
                               in enumerate(zip(self.output_var, self.metric, self.output_dim))
                               }
         else:
-            errors = array(array(self.optimizer.errors).flatten())
+            errors = np.array(np.array(self.optimizer.errors).flatten())
 
         dim = self.model.dimensions
 
         if format == 'list':
             res_list = []
-            for j in arange(0, len(params)):
+            for j in np.arange(0, len(params)):
                 temp_data = params[j]
                 if use_units:
                     res_dict = {n: Quantity(temp_data[n], dim=dim[n]) for n in names}
@@ -980,7 +722,7 @@ def results(self, format='list', use_units=None):
                 res_dict = {param_name: Quantity(param_value, dim=dim[param_name])
                             for param_name, param_value in temp_dict.items()}
             else:
-                res_dict = {param_name: array(param_value)
+                res_dict = {param_name: np.array(param_value)
                             for param_name, param_value in temp_dict.items()}
 
             res_dict['error'] = errors
@@ -992,17 +734,17 @@ def results(self, format='list', use_units=None):
                         res_dict['objective_errors'] = {output_var: raw_errors[output_var]
                                                         for output_var in self.output_var}
                     else:
-                        res_dict['objective_errors_normalized'] = {output_var: array([raw_error_normed[idx]
+                        res_dict['objective_errors_normalized'] = {output_var: np.array([raw_error_normed[idx]
                                                                                     for raw_error_normed in self._objective_errors_normed])
                                                                 for idx, output_var in enumerate(self.output_var)}
-                        res_dict['objective_errors'] = {output_var: array([raw_error[idx]
+                        res_dict['objective_errors'] = {output_var: np.array([raw_error[idx]
                                                                         for raw_error in self._objective_errors])
                                                         for idx, output_var in enumerate(self.output_var)}
                 else:
                     for idx, output_var in enumerate(self.output_var):
-                        res_dict['normalized_error_' + output_var] = array([raw_error_normed[idx]
+                        res_dict['normalized_error_' + output_var] = np.array([raw_error_normed[idx]
                                                                             for raw_error_normed in self._objective_errors_normed])
-                        res_dict['error_' + output_var] = array([raw_error[idx]
+                        res_dict['error_' + output_var] = np.array([raw_error[idx]
                                                                  for raw_error in self._objective_errors])
             if format == 'dict':
                 return res_dict
@@ -1061,10 +803,18 @@ def generate(self, output_var=None, params=None, param_init=None,
             output_var = self.output_var
         elif isinstance(output_var, str):
             output_var = [output_var]
-        self.simulator = self.setup_simulator('generate', self.n_traces,
-                                              output_var=output_var,
-                                              param_init=param_init,
-                                              level=level+1)
+        namespace = self.create_namespace(level=level+1)
+        self.simulator = setup_fit()  # For standalone, it is important to do this before the neuron group is created
+        neurons = self.setup_neuron_group(self.n_traces,
+                                          namespace=namespace)
+        record_vars = [v for v in output_var if v != 'spikes']
+        record_spikes = 'spikes' in output_var
+        self.simulator.setup("generate",
+                             neurons=neurons,
+                             record_vars=record_vars,
+                             param_init=param_init,
+                             dt=self.dt,
+                             record_spikes=record_spikes)
         param_dic = get_param_dic([params],
                                   self.parameter_names, self.n_traces, 1)
         self.simulator.run(self.duration, param_dic, self.parameter_names,
@@ -1136,7 +886,7 @@ def calc_errors(self, metric):
             traces = getattr(self.simulator.networks['fit']['statemonitor'],
                              o_var+'_')
             # Reshape traces for easier calculation of error
-            traces = reshape(traces, (traces.shape[0]//self.n_traces,
+            traces = np.reshape(traces, (traces.shape[0]//self.n_traces,
                                       self.n_traces,
                                       -1))
             errors = m.calc(traces, o, self.dt)
@@ -1293,12 +1043,32 @@ def refine(self, params=None, metric=None,
             min_bounds.append(min_bound)
             max_bounds.append(max_bound)
 
-        self.simulator = self.setup_simulator('refine', self.n_traces,
-                                              output_var=self.output_var,
-                                              param_init=self.param_init,
-                                              calc_gradient=calc_gradient,
-                                              optimize=optimize,
-                                              level=level+1)
+        self.simulator = setup_fit()   # for standalone, it is important to do this before creating the neurons
+
+        namespace = self.create_namespace(level=level+1)
+        neurons = self.setup_neuron_group(self.n_traces,
+                                          namespace=namespace,
+                                          calc_gradient=calc_gradient,
+                                          optimize=optimize)
+
+        record_vars = [v for v in self.output_var if v != 'spikes']
+        if calc_gradient:
+            record_vars.extend([f'S_{out_var}_{p}'
+                                for p in self.parameter_names
+                                for out_var in self.output_var])
+        record_spikes = 'spikes' in self.output_var
+        if calc_gradient:
+            param_init = dict(self.param_init)
+            param_init.update(get_sensitivity_init(neurons, self.parameter_names,
+                                                   param_init))
+        else:
+            param_init = self.param_init
+        self.simulator.setup('refine',
+                             neurons=neurons,
+                             record_vars=record_vars,
+                             record_spikes=record_spikes,
+                             param_init=param_init,
+                             dt=self.dt)
 
         t_start_steps = [int(round(t_s / self.dt)) if t_w is None else 0
                          for t_s, t_w in zip(t_start, t_weights)]
@@ -1329,19 +1099,19 @@ def _calc_error(x):
                 if t_w is None:
                     residual = trace[:, t_s_steps:] - out[:, t_s_steps:]
                 else:
-                    residual = (trace - out) * sqrt(t_w)
+                    residual = (trace - out) * np.sqrt(t_w)
                 one_residual.append((residual*norm).flatten())
 
             output_len = [output[:, t_s_steps:].size
                           for output, t_s_steps in zip(self.output,
                                                        t_start_steps)]
-            end_idx = cumsum(output_len)
-            start_idx = hstack([0, end_idx[:-1]])
-            error = tuple([mean(r**2) for r in one_residual])
-            combined_error = sum(array(error))
+            end_idx = np.cumsum(output_len)
+            start_idx = np.hstack([0, end_idx[:-1]])
+            error = tuple([np.mean(r**2) for r in one_residual])
+            combined_error = np.sum(np.array(error))
             errors.append(error)
             combined_errors.append(combined_error)
-            best_idx = argmin(combined_errors)
+            best_idx = np.argmin(combined_errors)
 
             if self.use_units:
                 norm_dim = get_dimensions(normalization[0]) ** 2
@@ -1370,7 +1140,7 @@ def _calc_error(x):
                                                self.output_dim,
                                                normalization)])
             else:
-                all_errors = array(combined_errors)
+                all_errors = np.array(combined_errors)
                 params = {p: float(val) for p, val
                           in zip(self.parameter_names, x)}
                 best_raw_error_normed = errors[best_idx]
@@ -1387,7 +1157,7 @@ def _calc_error(x):
             callback_func(params, errors,
                           best_params, best_error, n_evaluations[0], additional_info)
             n_evaluations[0] += 1
-            return array(hstack(one_residual))
+            return np.array(np.hstack(one_residual))
 
         def _calc_gradient(params):
             residuals = []
@@ -1402,10 +1172,10 @@ def _calc_gradient(params):
                     if t_w is None:
                         residual = trace[:, t_s_steps:]
                     else:
-                        residual = trace * sqrt(t_w)
+                        residual = trace * np.sqrt(t_w)
                     one_residual.append((residual*norm).flatten())
-                residuals.append(array(hstack(one_residual)))
-            gradient = array(residuals)
+                residuals.append(np.array(np.hstack(one_residual)))
+            gradient = np.array(residuals)
             return gradient.T
 
         assert 'jac' not in kwds
@@ -1451,7 +1221,7 @@ def __init__(self, model, input, output, dt, reset, threshold,
                          method=method, param_init=param_init,
                          penalty=penalty, use_units=use_units)
         self.output = [Quantity(o) for o in output]
-        self.output_ = [array(o) for o in output]
+        self.output_ = [np.array(o) for o in output]
 
         if param_init:
             for param, val in param_init.items():
@@ -1511,7 +1281,7 @@ def __init__(self, model, input_var, input, output_var, output, dt,
                          param_init=param_init, penalty=penalty)
 
         self.output = [Quantity(output)]
-        self.output_ = [array(output)]
+        self.output_ = [np.array(output)]
 
         if output_var not in self.model.names:
             raise NameError("%s is not a model variable" % output_var)
@@ -1537,6 +1307,11 @@ def __init__(self, model, input_var, input, output_var, output, dt,
 
         self.simulator = None
 
+    def create_namespace(self, level):
+        namespace = super().create_namespace(level=level+1)
+        namespace['t_start'] = self.t_start
+        return namespace
+
     def fit(self, optimizer, n_rounds=1, callback='text',
             restart=False, start_iteration=None, penalty=None,
             level=0, **params):
@@ -1554,8 +1329,8 @@ def fit(self, optimizer, n_rounds=1, callback='text',
     def calc_errors(self, metric=None):
         """Calculates error in online fashion.To be used inside optim_iter."""
         errors = self.simulator.neurons.total_error/int((self.duration-self.t_start)/defaultclock.dt)
-        errors = mean(errors.reshape((self.n_samples, self.n_traces)), axis=1)
-        return [array(errors)]
+        errors = np.mean(errors.reshape((self.n_samples, self.n_traces)), axis=1)
+        return [np.array(errors)]
 
     def generate_traces(self, params=None, param_init=None, level=0):
         """Generates traces for best fit of parameters and all inputs"""
diff --git a/brian2modelfitting/inferencer.py b/brian2modelfitting/inferencer.py
index 7110995..43487d6 100644
--- a/brian2modelfitting/inferencer.py
+++ b/brian2modelfitting/inferencer.py
@@ -28,15 +28,13 @@
 except ImportError:
     sbi = None
     torch = None
-# from sbi.utils.get_nn_models import (posterior_nn,
-#                                      likelihood_nn,
-#                                      classifier_nn)
-# from sbi.utils.torchutils import BoxUniform
-# from sbi.inference.posteriors.direct_posterior import DirectPosterior
-# import sbi.analysis
-# import sbi.inference
-# import torch
 
+from .base import (handle_input_args,
+                   handle_output_args,
+                   handle_param_init,
+                   input_equations,
+                   output_equations,
+                   output_dims)
 from .simulator import RuntimeSimulator, CPPStandaloneSimulator
 from .utils import tqdm
 
@@ -235,80 +233,34 @@ def __init__(self, dt, model, input, output, features=None, method=None,
         # model equations
         if isinstance(model, str):
             model = Equations(model)
-        else:
-            raise TypeError('Equations must be appropriately formatted.')
+        self.model = model
 
         # input data traces
-        if not isinstance(input, Mapping):
-            raise TypeError('`input` argument must be a dictionary mapping'
-                            ' the name of the input variable and `input`.')
-        if len(input) > 1:
-            raise NotImplementedError('Only a single input is supported.')
-        input_var = list(input.keys())[0]
-        input = input[input_var]
-        if input_var not in model.identifiers:
-            raise NameError(f'{input_var} is not an identifier in the model.')
-
-        # output data traces
-        if not isinstance(output, Mapping):
-            raise TypeError('`output` argument must be a dictionary mapping'
-                            ' the name of the output variable and `output`.')
-        output_var = list(output.keys())
-        output = list(output.values())
-        for o_var in output_var:
-            if o_var != 'spikes' and o_var not in model.names:
-                raise NameError(f'{o_var} is not a model variable')
+        input_arr, input_var = handle_input_args(input, None, model)
+        self.input = input_arr
+        input_eqs = input_equations(input_var, get_dimensions(input_arr))
+        self.model += input_eqs
+
+        output, output_var = handle_output_args(output, None, model)
+
         self.output_var = output_var
         self.output = output
+        self.output_ = [np.array(o, copy=False) for o in output]
+        self.output_dim = output_dims(output, output_var, model)
+
+        self.model += output_equations(output_var, self.output_dim)
 
         # create variable for parameter names
         self.param_names = sorted(model.parameter_names)
 
         # set the simulation time for a given time scale
-        self.n_traces, n_steps = input.shape
+        self.n_traces, n_steps = input_arr.shape
         self.sim_time = dt * n_steps
 
-        # handle multiple output variables
-        self.output_dim = []
-        for o_var, out in zip(self.output_var, self.output):
-            if o_var == 'spikes':
-                self.output_dim.append(DIMENSIONLESS)
-            else:
-                self.output_dim.append(model[o_var].dim)
-                fail_for_dimension_mismatch(out, self.output_dim[-1],
-                                            'The provided target values must'
-                                            ' have the same units as the'
-                                            f' variable {o_var}')
-
-        # add input to equations
-        self.model = model
-        input_dim = get_dimensions(input)
-        input_dim = '1' if input_dim is DIMENSIONLESS else repr(input_dim)
-        input_eqs = f'{input_var} = input_var(t, i % n_traces) : {input_dim}'
-        self.model += input_eqs
-
-        # add output to equations
-        counter = 0
-        for o_var, o_dim in zip(self.output_var, self.output_dim):
-            if o_var != 'spikes':
-                counter += 1
-                output_expr = f'output_var_{counter}(t, i % n_traces)'
-                output_dim = ('1' if o_dim is DIMENSIONLESS else repr(o_dim))
-                output_eqs = f'{o_var}_target = {output_expr} : {output_dim}'
-                self.model += output_eqs
-
         # create ``TimedArray`` object for input w.r.t. a given time scale
-        self.input_traces = TimedArray(input.transpose(), dt=self.dt)
-
-        # handle initial values for the ODE system
-        if not param_init:
-            param_init = {}
-        for param in param_init.keys():
-            if not (param in self.model.diff_eq_names or
-                    param in self.model.parameter_names):
-                raise ValueError(f'{param} is not a model variable or a'
-                                 ' parameter in the model')
-        self.param_init = param_init
+        self.input_traces = TimedArray(input_arr.transpose(), dt=self.dt)
+
+        self.param_init = handle_param_init(param_init, model)
 
         # handle the rest of optional parameters for the ``NeuronGroup`` class
         self.method = method
diff --git a/brian2modelfitting/metric.py b/brian2modelfitting/metric.py
index 54a93af..bc8bf14 100644
--- a/brian2modelfitting/metric.py
+++ b/brian2modelfitting/metric.py
@@ -1,14 +1,16 @@
+import multiprocessing
 import warnings
 import abc
-from collections import defaultdict
+
+import numpy as np
 
 try:
     import efel
 except ImportError:
     efel = None
 from itertools import repeat
-from brian2 import Hz, second, Quantity, ms, us, get_dimensions, mV
-from brian2.units.fundamentalunits import check_units, in_unit, DIMENSIONLESS
+from brian2 import second, Quantity, ms, get_dimensions, mV
+from brian2.units.fundamentalunits import check_units, DIMENSIONLESS
 from numpy import (array, sum, abs, amin, digitize, rint, arange, inf, NaN,
                    clip, mean)
 
@@ -101,19 +103,53 @@ def get_gamma_factor(model, data, delta, time, dt, rate_correction=True):
     return clip(rate_term - gamma, 0, inf)
 
 
-def calc_eFEL(traces, inp_times, feat_list, dt):
+def calc_eFEL(traces, inp_times, feat_list, dt, parallel_processes=0):
+    """
+    Calculate eFEL features for each trace.
+
+    Parameters
+    ----------
+    traces: `~numpy.ndarray`
+        2D array of shape ``(n_traces, time steps)`` with the traces.
+    inp_times: list of list
+        List of input times for each trace.
+    feat_list : list of str
+        List of eFEL features to calculate.
+    dt : `~brian2.units.fundamentalunits.Quantity`
+        The length of a single time step.
+    parallel_processes : int, optional
+        The number of parallel processes to use for the calculation. If 0,
+        no parallelization is used. If -1, the number of processes is
+        automatically determined. If >0, the specified number of processes
+        is used. Defaults to 0.
+
+    Returns
+    -------
+    features: list of dict
+        List of dictionaries with the features for each trace.
+    """
     _check_efel()
     out_traces = []
     for i, trace in enumerate(traces):
         time = arange(0, len(trace))*dt/ms
+        assert isinstance(time, np.ndarray)
         temp_trace = {}
         temp_trace['T'] = time
-        temp_trace['V'] = trace/mV
+        temp_trace['V'] = trace/float(mV)  # efel expects mV
         temp_trace['stim_start'] = [inp_times[i][0]]
         temp_trace['stim_end'] = [inp_times[i][1]]
         out_traces.append(temp_trace)
 
-    results = efel.getFeatureValues(out_traces, feat_list)
+    if parallel_processes == 0:
+        parallel_map = map
+    elif parallel_processes > 0:
+        parallel_map = multiprocessing.Pool(parallel_processes).map
+    elif parallel_processes == -1:
+        parallel_map = multiprocessing.Pool().map
+    else:
+        raise ValueError(f"Invalid value for 'parallel_processes': "
+                         f"{parallel_processes}")
+    results = efel.getFeatureValues(out_traces, feat_list, parallel_map=parallel_map)
     return results
 
 
@@ -337,6 +373,7 @@ def calc(self, model_traces, data_traces, dt):
             ``(n_samples, )``.
         """
         start_steps = int(round(self.t_start/dt))
+
         if self.t_weights is not None:
             features = self.get_features(model_traces * float(self.normalization),
                                          data_traces * float(self.normalization),
@@ -485,8 +522,43 @@ def revert_normalization(self, error):
 
 
 class FeatureMetric(TraceMetric):
+    """
+    Calculate the error between goal and calculated output using the
+    features returned by the eFEL package.
+
+    Parameters
+    ----------
+    stim_times: list of lists of `~brian2.units.fundamentalunits.Quantity`
+        The start and end times of the stimuli. The list has to have the same
+        length as the number of traces, and each element of the list has to be
+        a list or tuple with two elements, the start and end time of the stimulus.
+    feat_list: list of strings
+        The names of the features that should be calculated. You can get a list
+        of all available features from the ``eFEL`` package via
+        ``efel.api.getFeatureNames()``. For a description of the features, see
+        `https://efel.readthedocs.io/en/latest/eFeatures.html`_.
+    weights: dict, optional
+        A dictionary with the feature names as keys and the weights as values.
+        The weights are used to combine the individual features into a single
+        error value. The default is to use a weight of 1 for each feature.
+    combine: callable, optional
+        A function that combines the feature values for the model and the data
+        trace into a single error value. The default is to use the absolute
+        difference between the two values.
+    t_start: `~brian2.units.fundamentalunits.Quantity`, optional
+        Start of time window considered for calculating the fit error.
+    normalization : float, optional
+        A normalization term that will be used rescale results before
+        handing them to the optimization algorithm.
+    parallel_processes: int, optional
+        The number of parallel processes to use for calculating the features.
+        If 0, the features are calculated in the main process. If >0, the
+        features are calculated in parallel processes. A special value of -1
+        means to use as many processes as there are CPUs on the system (as
+        returned by ``os.cpu_count()``). Defaults to 0.
+    """
     def __init__(self, stim_times, feat_list, weights=None, combine=None,
-                 t_start=0*second, normalization=1.):
+                 t_start=0*second, normalization=1., parallel_processes=0):
         _check_efel()
         super(FeatureMetric, self).__init__(t_start=t_start,
                                             normalization=normalization)
@@ -511,6 +583,7 @@ def combine(x, y):
             raise TypeError("Weights has to be a dictionary!")
 
         self.weights = weights
+        self.parallel_processes = parallel_processes
 
     def check_values(self, feat_list):
         """Removes all the None values and checks for array features"""
@@ -542,16 +615,20 @@ def get_features(self, traces, output, dt):
                 raise ValueError("Specify the stim_times variable of appropiate "
                                  "size (same as number of traces or 1).")
 
-        out_feat = calc_eFEL(output, self.stim_times, self.feat_list, dt)
+        # TODO: This could be cached
+        out_feat = calc_eFEL(output, self.stim_times, self.feat_list, dt, parallel_processes=self.parallel_processes)
         self.check_values(out_feat)
 
+        # calculate features across all samples and traces in a single call
+        model_feat = calc_eFEL(traces.reshape(-1, traces.shape[-1]),
+                               np.tile(self.stim_times, (n_samples, 1)),
+                               self.feat_list, dt, parallel_processes=self.parallel_processes)
+        self.check_values(model_feat)
+
         features = []
-        for one_sample in traces:
-            sample_feat = calc_eFEL(one_sample, self.stim_times,
-                                    self.feat_list, dt)
-            self.check_values(sample_feat)
+        for sample_idx in range(n_samples):
             sample_features = []
-            for one_trace_feat, one_out in zip(sample_feat, out_feat):
+            for one_trace_feat, one_out in zip(model_feat[sample_idx*n_traces:(sample_idx + 1)*n_traces], out_feat):
                 sample_features.append(self.feat_to_err(one_trace_feat,
                                                         one_out))
             # Convert the list of dictionaries to a dictionary of lists
diff --git a/brian2modelfitting/optimizer.py b/brian2modelfitting/optimizer.py
index e163a02..e155147 100644
--- a/brian2modelfitting/optimizer.py
+++ b/brian2modelfitting/optimizer.py
@@ -153,8 +153,8 @@ def __init__(self,  method='DE', use_nevergrad_recommendation=False,
         super(Optimizer, self).__init__()
 
         if method not in registry:
-            raise AssertionError("Unknown to Nevergrad optimization method:"
-                                 + method)
+            raise AssertionError("Unknown to Nevergrad optimization method: "
+                                 + method + "\nAvailable methods: " + ", ".join(sorted(s for s in registry.keys())))
         self.tested_parameters = []
         self.errors = []
         self.method = method
diff --git a/brian2modelfitting/sensitivity.py b/brian2modelfitting/sensitivity.py
new file mode 100644
index 0000000..0c25464
--- /dev/null
+++ b/brian2modelfitting/sensitivity.py
@@ -0,0 +1,132 @@
+"""
+Functions to determine sensitivity equations and initial values.
+"""
+import sympy
+
+from brian2.core.namespace import get_local_namespace
+from brian2.equations.equations import Equations, SUBEXPRESSION
+from brian2.parsing.sympytools import str_to_sympy, sympy_to_str
+from brian2.utils.stringtools import get_identifiers
+
+
+def get_sensitivity_equations(group, parameters, namespace=None, level=1,
+                              optimize=True):
+    """
+    Get equations for sensitivity variables.
+
+    Parameters
+    ----------
+    group : `NeuronGroup`
+        The group of neurons that will be simulated.
+    parameters : list of str
+        Names of the parameters that are fit.
+    namespace : dict, optional
+        The namespace to use.
+    level : `int`, optional
+        How much farther to go down in the stack to find the namespace.
+    optimize : bool, optional
+        Whether to remove sensitivity variables from the equations that do
+        not evolve if initialized to zero (e.g. ``dS_x_y/dt = -S_x_y/tau``
+        would be removed). This avoids unnecessary computation but will fail
+        in the rare case that such a sensitivity variable needs to be
+        initialized to a non-zero value. Defaults to ``True``.
+
+    Returns
+    -------
+    sensitivity_eqs : `Equations`
+        The equations for the sensitivity variables.
+    """
+    if namespace is None:
+        namespace = get_local_namespace(level)
+        namespace.update(group.namespace)
+
+    eqs = group.equations
+    diff_eqs = eqs.get_substituted_expressions(group.variables)
+    diff_eq_names = [name for name, _ in diff_eqs]
+
+    system = sympy.Matrix([str_to_sympy(diff_eq[1].code)
+                           for diff_eq in diff_eqs])
+    J = system.jacobian([str_to_sympy(d) for d in diff_eq_names])
+
+    sensitivity = []
+    sensitivity_names = []
+    for parameter in parameters:
+        F = system.jacobian([str_to_sympy(parameter)])
+        names = [str_to_sympy(f'S_{diff_eq_name}_{parameter}')
+                 for diff_eq_name in diff_eq_names]
+        sensitivity.append(J * sympy.Matrix(names) + F)
+        sensitivity_names.append(names)
+
+    new_eqs = []
+    for names, sensitivity_eqs, param in zip(sensitivity_names, sensitivity, parameters):
+        for name, eq, orig_var in zip(names, sensitivity_eqs, diff_eq_names):
+            unit = eqs[orig_var].dim / group.variables[param].dim
+            unit = repr(unit) if not unit.is_dimensionless else '1'
+            if optimize:
+                # Check if the equation stays at zero if initialized at zero
+                zeroed = eq.subs(name, sympy.S.Zero)
+                if zeroed == sympy.S.Zero:
+                    # No need to include equation as differential equation
+                    if unit == '1':
+                        new_eqs.append(f'{sympy_to_str(name)} = 0 : {unit}')
+                    else:
+                        new_eqs.append(f'{sympy_to_str(name)} = 0*{unit} : {unit}')
+                    continue
+            rhs = sympy_to_str(eq)
+            if rhs == '0':  # avoid unit mismatch
+                rhs = f'0*{unit}/second'
+            new_eqs.append('d{lhs}/dt = {rhs} : {unit}'.format(lhs=sympy_to_str(name),
+                                                               rhs=rhs,
+                                                               unit=unit))
+    new_eqs = Equations('\n'.join(new_eqs))
+    return new_eqs
+
+
+def get_sensitivity_init(group, parameters, param_init):
+    """
+    Calculate the initial values for the sensitivity parameters (necessary if
+    initial values are functions of parameters).
+
+    Parameters
+    ----------
+    group : `NeuronGroup`
+        The group of neurons that will be simulated.
+    parameters : list of str
+        Names of the parameters that are fit.
+    param_init : dict
+        The dictionary with expressions to initialize the model variables.
+
+    Returns
+    -------
+    sensitivity_init : dict
+        Dictionary of expressions to initialize the sensitivity
+        parameters.
+    """
+    sensitivity_dict = {}
+    for var_name, expr in param_init.items():
+        if not isinstance(expr, str):
+            continue
+        identifiers = get_identifiers(expr)
+        for identifier in identifiers:
+            if (identifier in group.variables
+                    and getattr(group.variables[identifier],
+                                'type', None) == SUBEXPRESSION):
+                raise NotImplementedError('Initializations that refer to a '
+                                          'subexpression are currently not '
+                                          'supported')
+            sympy_expr = str_to_sympy(expr)
+            for parameter in parameters:
+                diffed = sympy_expr.diff(str_to_sympy(parameter))
+                if diffed != sympy.S.Zero:
+                    if getattr(group.variables[parameter],
+                               'type', None) == SUBEXPRESSION:
+                        raise NotImplementedError('Sensitivity '
+                                                  f'S_{var_name}_{parameter} '
+                                                  'is initialized to a non-zero '
+                                                  'value, but it has been '
+                                                  'removed from the equations. '
+                                                  'Set optimize=False to avoid '
+                                                  'this.')
+                    init_expr = sympy_to_str(diffed)
+                    sensitivity_dict[f'S_{var_name}_{parameter}'] = init_expr
+    return sensitivity_dict
diff --git a/brian2modelfitting/simulator.py b/brian2modelfitting/simulator.py
index 852b329..ee1edaf 100644
--- a/brian2modelfitting/simulator.py
+++ b/brian2modelfitting/simulator.py
@@ -1,7 +1,7 @@
 import os
 import abc
 from numpy import atleast_1d
-from brian2 import device, Network
+from brian2 import device, Network, StateMonitor, SpikeMonitor
 
 
 def initialize_parameter(variableview, value):
@@ -60,6 +60,35 @@ def __init__(self):
     statemonitor = property(lambda self: self.networks[self.current_net]['statemonitor'])
     spikemonitor = property(lambda self: self.networks[self.current_net]['spikemonitor'])
 
+    def setup(self, network_name, neurons, record_vars, param_init, dt, record_spikes=False):
+        """
+        Basic setup for the simulation.
+
+        Parameters
+        ----------
+        network_name : str
+            Name of the network
+        neurons : `~brian2.groups.neurongroup.NeuronGroup`
+            The ``NeuronGroup`` to be simulated
+        record_vars : list[str]
+            Names of the variables to be recorded
+        param_init : dict
+            Dictionary of initial values for the parameters
+        dt : `~brian2.units.fundamentalunits.Quantity`
+            Time step of the simulation
+        record_spikes : bool, optional
+            Whether to record spikes or not (default: False)
+        """
+        network = Network(neurons)
+        if record_spikes:
+            network.add(SpikeMonitor(neurons, name='spikemonitor'))
+
+        if len(record_vars):
+            network.add(StateMonitor(neurons, record_vars, record=True,
+                                     name='statemonitor', dt=dt))
+
+        self.initialize(network, param_init, name=network_name)
+
     def initialize(self, network, var_init, name='fit'):
         """
         Prepares the simulation for running


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index d747368..885376a 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -27,7 +27,7 @@ jobs:
     - name: Install brian2modelfitting
       run: |
         python -m pip install --upgrade pip wheel
-        python -m pip install flake8 pytest-coverage coveralls
+        python -m pip install flake8 pytest-coverage pytest-timeout coveralls
         python -m pip install ".[all]"
     - name: Update to latest Brian development version
       run: python -m pip install -i https://test.pypi.org/simple/ --pre --upgrade Brian2
@@ -40,7 +40,7 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
-        pytest --cov=brian2modelfitting
+        pytest --timeout=60 --cov=brian2modelfitting
     - name: Upload coverage to coveralls
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/brian2modelfitting/tests/test_metric.py b/brian2modelfitting/tests/test_metric.py
index 17273e0..216b0bf 100644
--- a/brian2modelfitting/tests/test_metric.py
+++ b/brian2modelfitting/tests/test_metric.py
@@ -216,14 +216,22 @@ def test_get_errors_gamma():
     assert_almost_equal(errors, [0., 0.])
 
 
-def test_calc_EFL():
+@pytest.mark.parametrize("parallel_processes", [0, 2, -1, -2])  # only testing that it works at all
+def test_calc_EFL(parallel_processes):
     # "voltage traces" that are constant at -70*mV, -60mV, -50mV, -40mV for
     # 50ms each.
     dt = 1*ms
     voltage = np.ones((2, 200))*np.repeat([-70, -60, -50, -40], 50)*mV
     # Note that calcEFL takes times in ms
     inp_times = [[99, 150], [49, 150]]
-    results = calc_eFEL(voltage, inp_times, ['voltage_base'], dt=dt)
+    if parallel_processes == -2:  # invalid value
+        with pytest.raises(ValueError):
+            calc_eFEL(voltage, inp_times, ['voltage_base'], dt=dt,
+                      parallel_processes=parallel_processes)
+        return
+    results = calc_eFEL(voltage, inp_times, ['voltage_base'], dt=dt,
+                        parallel_processes=parallel_processes)
+
     assert len(results) == 2
     assert all(res.keys() == {'voltage_base'} for res in results)
     assert_almost_equal(results[0]['voltage_base'], -60)
diff --git a/brian2modelfitting/tests/test_modelfitting_tracefitter.py b/brian2modelfitting/tests/test_modelfitting_tracefitter.py
index de6fca6..a2ccefe 100644
--- a/brian2modelfitting/tests/test_modelfitting_tracefitter.py
+++ b/brian2modelfitting/tests/test_modelfitting_tracefitter.py
@@ -355,7 +355,7 @@ def test_fitter_fit_methods(method):
                      output=output_traces,
                      n_samples=30)
     # Skip a few methods that seem to hang due to multi-threading deadlocks (?) or simply take very long
-    skip = ['BO', 'ParaPortfolio', 'BAR', 'MultiBFGS', 'MultiCobyla', 'MultiSQP', 'NgIohRW']
+    skip = ['BO', 'ParaPortfolio', 'BAR', 'MultiBFGS', 'MultiCobyla', 'MultiSQP', 'NgIohRW', 'F3SQPCMA']
     if any(s in method for s in skip):
         pytest.skip(f'Skipping method {method}')