<gh_stars>0
import logging
import warnings
from collections import OrderedDict

import numpy as np
from nengo.exceptions import SimulationError, ValidationError

from nengo_loihi.builder.discretize import (
    LEARN_FRAC,
    Q_BITS,
    U_BITS,
    decay_int,
    learn_overflow_bits,
    overflow_signed,
    scale_pes_errors,
    shift,
)
from nengo_loihi.probe import LoihiProbe, ProbeFilter

logger = logging.getLogger(__name__)


class EmulatorInterface:
    """Software emulator for Loihi chip behaviour.

    Parameters
    ----------
    model : Model
        Model specification that will be simulated.
    seed : int, optional (Default: None)
        A seed for all stochastic operations done in this simulator.
    """

    strict = False

    def __init__(self, model, seed=None):
        self.closed = True

        if seed is None:
            seed = np.random.randint(2**31 - 1)
        self.seed = seed
        logger.debug("EmulatorInterface seed: %d", seed)
        self.rng = np.random.RandomState(self.seed)

        self.block_info = BlockInfo(model.blocks)
        self.inputs = list(model.inputs)
        logger.debug("EmulatorInterface dtype: %s", self.block_info.dtype)

        self.compartment = CompartmentState(self.block_info, strict=self.strict)
        self.synapses = SynapseState(
            self.block_info,
            pes_error_scale=getattr(model, "pes_error_scale", 1.0),
            strict=self.strict,
        )
        self.axons = AxonState(self.block_info)
        self.probes = ProbeState(self.block_info, list(model.probes), model.dt)

        self.t = 0
        self.closed = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def clear_probes(self):
        # Probe data in `self.probes` is cleared after it is retrieved by a
        # `self.collect_probe_output` call, so nothing needs to happen here.
        # This function remains so that `nengo_loihi.Simulator` can call `clear_probes`
        # on all of its internal simulators without checking their types.
        pass

    def close(self):
        self.closed = True

        # remove references to states to free memory (except probes)
        self.block_info = None
        self.inputs = None
        self.compartment = None
        self.synapses = None
        self.axons = None

    def collect_probe_output(self, probe):
        """Collect and clear output from a particular probe"""
        return self.probes.collect_probe_output(probe)

    def chip2host(self, probes_receivers):
        increment = None
        for probe, receiver in probes_receivers.items():
            inc = self.probes.send(self.t, probe, receiver)
            increment = inc if increment is None else increment
            assert increment == inc

    def host2chip(self, spikes, errors):
        for spike_input, t, spike_idxs in spikes:
            spike_input.add_spikes(t, spike_idxs)

        self.synapses.update_pes_errors(errors)

    def run_steps(self, steps):
        """Simulate for the given number of ``dt`` steps.

        Parameters
        ----------
        steps : int
            Number of steps to run the simulation for.
        """
        for _ in range(steps):
            self.step()

    def step(self):
        """Advance the simulation by 1 step (``dt`` seconds)."""
        self.t += 1
        self.compartment.advance_input()
        self.synapses.inject_current(
            self.t, self.inputs, self.axons, self.compartment.spiked
        )
        self.synapses.update_input(self.compartment.input)
        self.synapses.update_traces(self.t, self.rng)
        self.synapses.update_weights(self.t, self.rng)
        self.compartment.update(self.rng)
        self.probes.update(self.t, self.compartment)


class BlockInfo:
    """Provide information about all the LoihiBlocks in the model.

    Attributes
    ----------
    dtype : dtype
        Datatype of the blocks. Either ``np.float32`` if the blocks are not
        discretized or ``np.int32`` if they are. All blocks are the same.
    blocks : list of LoihiBlock
        List of all the blocks in the model.
    n_compartments : int
        Total number of compartments across all blocks.
    slices : dict of {LoihiBlock: slice}
        Maps each block to a slice for that block's compartments with
        respect to all compartments. Used to slice into any array storing
        data across all compartments.
    """

    def __init__(self, blocks):
        self.blocks = list(blocks)
        self.slices = OrderedDict()

        assert self.dtype in (np.float32, np.int32)

        start_ix = end_ix = 0
        for block in self.blocks:
            end_ix += block.n_neurons
            self.slices[block] = slice(start_ix, end_ix)
            assert block.compartment.vth.dtype == self.dtype
            assert block.compartment.bias.dtype == self.dtype
            start_ix = end_ix

        self.n_compartments = end_ix

    @property
    def dtype(self):
        return self.blocks[0].compartment.vth.dtype


class IterableState:
    """Base class for aspects of the emulator state.

    This class takes the name of a LoihiBlock attribute as the
    ``block_key`` and maps these objects to their parent blocks and slices.

    Attributes
    ----------
    dtype : dtype
        Datatype of the state elements (given by the BlockInfo datatype).
    block_map : dict of {item: block}
        Maps an item (determined by ``block_key``) to the block
        it belongs to.
    n_compartments : int
        The total number of neuron compartments (given by BlockInfo).
    slices : dict of {item: slice}
        Maps an item to the ``block_info.slice`` for the block
        it belongs to.
    strict : bool (Default: True)
        Whether "undesired" chip effects (ex. overflow) raise errors (``True``)
        or whether they only raise warnings (``False``).
    """

    def __init__(self, block_info, block_key, strict=True):
        self.n_compartments = block_info.n_compartments
        self.dtype = block_info.dtype
        self.strict = strict

        blocks_items = list(self._blocks_items(block_info.blocks, block_key))
        self.block_map = OrderedDict((item, block) for block, item in blocks_items)
        self.slices = OrderedDict(
            (item, block_info.slices[block]) for block, item in blocks_items
        )

    @staticmethod
    def _blocks_items(blocks, block_key):
        for block in blocks:
            if block_key == "compartment":
                # one item per block
                yield block, getattr(block, block_key)
            else:
                # multiple items per block (attribute is iterable)
                for item in getattr(block, block_key):
                    yield block, item

    def error(self, msg):
        if self.strict:
            raise SimulationError(msg)
        else:
            warnings.warn(msg)

    def items(self):
        return self.slices.items()


class CompartmentState(IterableState):
    """State representing the Compartments of all blocks."""

    MAX_DELAY = 1  # delay not yet implemented

    def __init__(self, block_info, strict=True):
        super().__init__(block_info, "compartment", strict=strict)

        # Initialize NumPy arrays to store compartment-related data
        self.input = np.zeros((self.MAX_DELAY, self.n_compartments), dtype=self.dtype)
        self.current = np.zeros(self.n_compartments, dtype=self.dtype)
        self.voltage = np.zeros(self.n_compartments, dtype=self.dtype)
        self.spiked = np.zeros(self.n_compartments, dtype=bool)
        self.spike_count = np.zeros(self.n_compartments, dtype=np.int32)
        self.ref_count = np.zeros(self.n_compartments, dtype=np.int32)

        self.decay_u = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.decay_v = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.scale_u = np.ones(self.n_compartments, dtype=self.dtype)
        self.scale_v = np.ones(self.n_compartments, dtype=self.dtype)

        self.vth = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.vmin = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.vmax = np.full(self.n_compartments, np.nan, dtype=self.dtype)

        self.bias = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.ref = np.full(self.n_compartments, np.nan, dtype=self.dtype)

        # Fill in arrays with parameters from CompartmentSegments
        for compartment, sl in self.items():
            self.decay_u[sl] = compartment.decay_u
            self.decay_v[sl] = compartment.decay_v
            if compartment.scale_u:
                self.scale_u[sl] = compartment.decay_u
            if compartment.scale_v:
                self.scale_v[sl] = compartment.decay_v
            self.vth[sl] = compartment.vth
            self.vmin[sl] = compartment.vmin
            self.vmax[sl] = compartment.vmax
            self.bias[sl] = compartment.bias
            self.ref[sl] = compartment.refract_delay

        assert not np.any(np.isnan(self.decay_u))
        assert not np.any(np.isnan(self.decay_v))
        assert not np.any(np.isnan(self.vth))
        assert not np.any(np.isnan(self.vmin))
        assert not np.any(np.isnan(self.vmax))
        assert not np.any(np.isnan(self.bias))
        assert not np.any(np.isnan(self.ref))

        if self.dtype == np.int32:
            assert (self.scale_u == 1).all()
            assert (self.scale_v == 1).all()
            self._decay_current = lambda x, u: decay_int(x, self.decay_u, offset=1) + u
            self._decay_voltage = lambda x, u: decay_int(x, self.decay_v) + u

            def overflow(x, bits, name=None):
                _, o = overflow_signed(x, bits=bits, out=x)
                if np.any(o):
                    self.error("Overflow" + (" in %s" % name if name else ""))

        elif self.dtype == np.float32:

            def decay_float(x, u, d, s):
                return (1 - d) * x + s * u

            self._decay_current = lambda x, u: decay_float(
                x, u, d=self.decay_u, s=self.scale_u
            )
            self._decay_voltage = lambda x, u: decay_float(
                x, u, d=self.decay_v, s=self.scale_v
            )

            def overflow(x, bits, name=None):
                pass  # do not do overflow in floating point

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        self._overflow = overflow

        self.noise = NoiseState(block_info)

    def advance_input(self):
        self.input[:-1] = self.input[1:]
        self.input[-1] = 0

    def update(self, rng):
        noise = self.noise.sample(rng)
        q0 = self.input[0, :]
        q0[~(self.noise.target_u)] += noise[~(self.noise.target_u)]
        self._overflow(q0, bits=Q_BITS, name="q0")

        self.current[:] = self._decay_current(self.current, q0)
        self._overflow(self.current, bits=U_BITS, name="current")
        u2 = self.current + self.bias
        u2[self.noise.target_u] += noise[self.noise.target_u]
        self._overflow(u2, bits=U_BITS, name="u2")

        self.voltage[:] = self._decay_voltage(self.voltage, u2)
        # We have not been able to create V overflow on the chip, so we do
        # not include it here. See github.com/nengo/nengo-loihi/issues/130
        # self.overflow(self.v, bits=V_BIT, name="V")

        np.clip(self.voltage, self.vmin, self.vmax, out=self.voltage)
        self.voltage[self.ref_count > 0] = 0
        # TODO^: don't zero voltage in case neuron is saving overshoot

        self.spiked[:] = self.voltage > self.vth
        self.voltage[self.spiked] = 0
        self.ref_count[self.spiked] = self.ref[self.spiked]
        # decrement ref_count
        np.clip(self.ref_count - 1, 0, None, out=self.ref_count)

        self.spike_count[self.spiked] += 1


class NoiseState(IterableState):
    """State representing the noise parameters for all compartments."""

    def __init__(self, block_info):
        super().__init__(block_info, "compartment")
        self.enabled = np.full(self.n_compartments, np.nan, dtype=bool)
        self.exp = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.mant_offset = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.target_u = np.full(self.n_compartments, np.nan, dtype=bool)

        # Fill in arrays with parameters from Compartment
        for compartment, sl in self.items():
            self.enabled[sl] = compartment.enable_noise
            self.exp[sl] = compartment.noise_exp
            self.mant_offset[sl] = compartment.noise_offset
            self.target_u[sl] = compartment.noise_at_membrane

        if self.dtype == np.int32:
            # TODO: if we could do this mult with shifts, it'd be faster, but
            # numpy has no function taking a vector of positive/negative shifts
            self.mult = np.where(self.enabled, 2.0 ** (self.exp - 7), 0)
            self.mant_offset *= 64

            def uniform(rng, n=self.n_compartments):
                return rng.randint(-127, 128, size=n, dtype=np.int32)

        elif self.dtype == np.float32:
            self.mult = np.where(self.enabled, 10.0**self.exp, 0)

            def uniform(rng, n=self.n_compartments):
                return rng.uniform(-1, 1, size=n).astype(np.float32)

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        assert not np.any(np.isnan(self.enabled))
        assert not np.any(np.isnan(self.exp))
        assert not np.any(np.isnan(self.mant_offset))
        assert not np.any(np.isnan(self.target_u))
        assert not np.any(np.isnan(self.mult))

        self._uniform = uniform

    def sample(self, rng):
        x = self._uniform(rng)
        return ((x + self.mant_offset) * self.mult).astype(self.dtype)


class SynapseState(IterableState):
    """State representing all synapses.

    Attributes
    ----------
    pes_error_scale : float
        Scaling for the errors of PES learning rules.
    pes_errors : {Synapse: ndarray(n_neurons / 2)}
        Maps synapse to PES learning rule errors for those synapses.
    spikes_in : {Synapse: list}
        Maps synapse to a queue of input spikes targeting those synapses.
    traces : {Synapse: ndarray(Synapse.n_axons)}
        Maps synapse to trace values for each of their axons.
    trace_spikes : {Synapse: set}
        Maps synapse to a queue of input spikes waiting to be added to those
        synapse traces.
    """

    def __init__(self, block_info, pes_error_scale=1.0, strict=True):  # noqa: C901
        super().__init__(block_info, "synapses", strict=strict)

        self.pes_error_scale = pes_error_scale

        self.spikes_in = OrderedDict()
        self.traces = OrderedDict()
        self.trace_spikes = OrderedDict()
        self.pes_errors = OrderedDict()
        for synapse in self.slices:
            n = synapse.n_axons
            self.spikes_in[synapse] = []

            if synapse.learning:
                self.traces[synapse] = np.zeros(n, dtype=self.dtype)
                self.trace_spikes[synapse] = set()
                self.pes_errors[synapse] = np.zeros(
                    self.block_map[synapse].n_neurons // 2, dtype=self.dtype
                )
                # ^ Currently, PES learning only happens on Nodes, where we
                # have pairs of on/off neurons. Therefore, the number of error
                # dimensions is half the number of neurons.

        if self.dtype == np.int32:

            def stochastic_round(
                x, dtype=self.dtype, rng=None, clip=None, name="values"
            ):
                x_sign = np.sign(x).astype(dtype)
                x_frac, x_int = np.modf(np.abs(x))
                p = rng.rand(*x.shape)
                y = x_int.astype(dtype) + (x_frac > p)
                if clip is not None:
                    q = y > clip
                    if np.any(q):
                        warnings.warn("Clipping %s" % name)
                    y[q] = clip
                return x_sign * y

            def trace_round(x, rng=None):
                return stochastic_round(x, rng=rng, clip=127, name="synapse trace")

            def weight_update(synapse, delta_ws, rng=None):
                synapse_cfg = synapse.synapse_cfg
                wgt_exp = synapse_cfg.real_weight_exp
                shift_bits = synapse_cfg.shift_bits
                overflow = learn_overflow_bits(n_factors=2)
                for w, delta_w in zip(synapse.weights, delta_ws):
                    product = shift(
                        delta_w * synapse._lr_int,
                        LEARN_FRAC + synapse._lr_exp - overflow,
                    )
                    learn_w = shift(w, LEARN_FRAC - wgt_exp) + product
                    learn_w[:] = stochastic_round(
                        learn_w * 2 ** (-LEARN_FRAC - shift_bits),
                        clip=2 ** (8 - shift_bits) - 1,
                        rng=rng,
                        name="learning weights",
                    )
                    w[:] = np.left_shift(learn_w, wgt_exp + shift_bits)

        elif self.dtype == np.float32:

            def trace_round(x, rng=None):
                return x  # no rounding

            def weight_update(synapse, delta_ws, rng=None):
                for w, delta_w in zip(synapse.weights, delta_ws):
                    w += synapse.learning_rate * delta_w

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        self._trace_round = trace_round
        self._weight_update = weight_update

    def inject_current(self, t, spike_inputs, all_axons, spiked):
        # --- clear spikes going in to each synapse
        for spike_queue in self.spikes_in.values():
            spike_queue.clear()

        # --- inputs pass spikes to synapses
        if t >= 2:  # input spikes take one time-step to arrive
            for spike_input in spike_inputs:
                compartment_idxs = spike_input.spike_idxs(t - 1)
                for axon in spike_input.axons:
                    spikes = axon.map_spikes(compartment_idxs)
                    self.spikes_in[axon.target].extend(
                        s for s in spikes if s is not None
                    )

        # --- axons pass spikes to synapses
        for axon, a_idx in all_axons.items():
            compartment_idxs = spiked[a_idx].nonzero()[0]
            spikes = axon.map_spikes(compartment_idxs)
            self.spikes_in[axon.target].extend(s for s in spikes if s is not None)

    def update_input(self, input):
        for synapse, s_slice in self.items():
            qb = input[:, s_slice]

            for spike in self.spikes_in[synapse]:
                base = synapse.axon_compartment_base(spike.axon_idx)
                if base is None:
                    continue

                weights, indices = synapse.axon_weights_indices(
                    spike.axon_idx, atom=spike.atom
                )
                qb[0, base + indices] += weights

    def update_pes_errors(self, errors):
        # TODO: these are sent every timestep, but learning only happens every
        # `tepoch * 2**learn_k` timesteps (see Synapse). Need to average.
        for pes_errors in self.pes_errors.values():
            pes_errors[:] = 0

        for synapse, _, e in errors:
            pes_errors = self.pes_errors[synapse]
            assert pes_errors.shape == e.shape
            pes_errors += scale_pes_errors(e, scale=self.pes_error_scale)

    def update_weights(self, t, rng):
        for synapse, pes_error in self.pes_errors.items():
            if t % synapse.learn_epoch == 0:
                trace = self.traces[synapse]
                e = np.hstack([-pes_error, pes_error])
                delta_w = np.outer(trace, e)
                self._weight_update(synapse, delta_w, rng=rng)

    def update_traces(self, t, rng):
        for synapse in self.traces:
            trace_spikes = self.trace_spikes.get(synapse, None)
            if trace_spikes is not None:
                for spike in self.spikes_in[synapse]:
                    if spike.axon_idx in trace_spikes:
                        self.error("Synaptic trace spikes lost")
                    trace_spikes.add(spike.axon_idx)

            trace = self.traces.get(synapse, None)
            if trace is not None and t % synapse.train_epoch == 0:
                tau = synapse.tracing_tau
                decay = np.exp(-synapse.train_epoch / tau)
                trace1 = decay * trace
                trace1[list(trace_spikes)] += synapse.tracing_mag
                trace[:] = self._trace_round(trace1, rng=rng)
                trace_spikes.clear()


class AxonState(IterableState):
    """State representing all (output) Axons."""

    def __init__(self, block_info):
        super().__init__(block_info, "axons")


class ProbeState:
    """State representing all probes.

    Attributes
    ----------
    dt : float
        Time constant of the Emulator.
    filters : {nengo_loihi.Probe: function}
        Maps Probes to the filtering function for that probe.
    filter_pos : {nengo_loihi.Probe: int}
        Maps Probes to the position of their filter in the data.
    block_probes : {nengo_loihi.Probe: slice}
        Maps Probes to the BlockInfo slice for the block they are probing.
    input_probes : {nengo_loihi.Probe: SpikeInput}
        Maps Probes to the SpikeInput that they are probing.
    """

    def __init__(self, block_info, probes, dt):
        self.dt = dt

        self.probes = OrderedDict()
        for probe in probes:
            block_slices = OrderedDict(
                [(block, block_info.slices[block]) for block in probe.target]
            )
            self.probes[probe] = block_slices

        self.outputs = {}
        for probe, block_slices in self.probes.items():
            self.outputs[probe] = [[] for block in block_slices]

        self.probe_filter = ProbeFilter(dt=dt)

    def collect_probe_output(self, probe):
        """Collect and clear output from a particular probe"""
        assert isinstance(probe, LoihiProbe)
        out = probe.weight_outputs(self.outputs[probe])
        out = self.probe_filter(probe, out)
        self.clear_probe(probe)
        return out

    def clear_probe(self, probe):
        for block_output in self.outputs[probe]:
            block_output.clear()

    def send(self, t, probe, receiver):
        """Send probed data to the receiver node.

        Returns
        -------
        n_send_steps : int
            The number of steps sent to the receiver.
        """
        # we don't currently filter here, so make sure the probe isn't expecting it
        assert probe.synapse is None, "Filtering should be done on host-side connection"

        outputs = self.outputs[probe]
        n_send_steps = len(outputs[0])
        if n_send_steps > 0:
            assert t >= n_send_steps, (t, n_send_steps)
            t0 = t - n_send_steps

            x = probe.weight_outputs(outputs)
            for j, xx in enumerate(x):
                receiver.receive(self.dt * (t0 + j + 2), xx)

        self.clear_probe(probe)
        return n_send_steps

    def update(self, t, compartment):
        for probe, block_slices in self.probes.items():
            assert hasattr(compartment, probe.key)
            target_attr = getattr(compartment, probe.key)

            for k, (block, block_slice) in enumerate(block_slices.items()):
                output = target_attr[block_slice][probe.slice[k]]
                if output.base is not None:
                    # if not already copied by the probe slice, then copy
                    output = output.copy()

                self.outputs[probe][k].append(output)

import logging
import warnings
from collections import OrderedDict

import numpy as np
from nengo.exceptions import SimulationError, ValidationError

from nengo_loihi.builder.discretize import (
    LEARN_FRAC,
    Q_BITS,
    U_BITS,
    decay_int,
    learn_overflow_bits,
    overflow_signed,
    scale_pes_errors,
    shift,
)
from nengo_loihi.probe import LoihiProbe, ProbeFilter

logger = logging.getLogger(__name__)


class EmulatorInterface:
    """Software emulator for Loihi chip behaviour.

    Parameters
    ----------
    model : Model
        Model specification that will be simulated.
    seed : int, optional (Default: None)
        A seed for all stochastic operations done in this simulator.
    """

    strict = False

    def __init__(self, model, seed=None):
        self.closed = True

        if seed is None:
            seed = np.random.randint(2**31 - 1)
        self.seed = seed
        logger.debug("EmulatorInterface seed: %d", seed)
        self.rng = np.random.RandomState(self.seed)

        self.block_info = BlockInfo(model.blocks)
        self.inputs = list(model.inputs)
        logger.debug("EmulatorInterface dtype: %s", self.block_info.dtype)

        self.compartment = CompartmentState(self.block_info, strict=self.strict)
        self.synapses = SynapseState(
            self.block_info,
            pes_error_scale=getattr(model, "pes_error_scale", 1.0),
            strict=self.strict,
        )
        self.axons = AxonState(self.block_info)
        self.probes = ProbeState(self.block_info, list(model.probes), model.dt)

        self.t = 0
        self.closed = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def clear_probes(self):
        # Probe data in `self.probes` is cleared after it is retrieved by a
        # `self.collect_probe_output` call, so nothing needs to happen here.
        # This function remains so that `nengo_loihi.Simulator` can call `clear_probes`
        # on all of its internal simulators without checking their types.
        pass

    def close(self):
        self.closed = True

        # remove references to states to free memory (except probes)
        self.block_info = None
        self.inputs = None
        self.compartment = None
        self.synapses = None
        self.axons = None

    def collect_probe_output(self, probe):
        """Collect and clear output from a particular probe"""
        return self.probes.collect_probe_output(probe)

    def chip2host(self, probes_receivers):
        increment = None
        for probe, receiver in probes_receivers.items():
            inc = self.probes.send(self.t, probe, receiver)
            increment = inc if increment is None else increment
            assert increment == inc

    def host2chip(self, spikes, errors):
        for spike_input, t, spike_idxs in spikes:
            spike_input.add_spikes(t, spike_idxs)

        self.synapses.update_pes_errors(errors)

    def run_steps(self, steps):
        """Simulate for the given number of ``dt`` steps.

        Parameters
        ----------
        steps : int
            Number of steps to run the simulation for.
        """
        for _ in range(steps):
            self.step()

    def step(self):
        """Advance the simulation by 1 step (``dt`` seconds)."""
        self.t += 1
        self.compartment.advance_input()
        self.synapses.inject_current(
            self.t, self.inputs, self.axons, self.compartment.spiked
        )
        self.synapses.update_input(self.compartment.input)
        self.synapses.update_traces(self.t, self.rng)
        self.synapses.update_weights(self.t, self.rng)
        self.compartment.update(self.rng)
        self.probes.update(self.t, self.compartment)


class BlockInfo:
    """Provide information about all the LoihiBlocks in the model.

    Attributes
    ----------
    dtype : dtype
        Datatype of the blocks. Either ``np.float32`` if the blocks are not
        discretized or ``np.int32`` if they are. All blocks are the same.
    blocks : list of LoihiBlock
        List of all the blocks in the model.
    n_compartments : int
        Total number of compartments across all blocks.
    slices : dict of {LoihiBlock: slice}
        Maps each block to a slice for that block's compartments with
        respect to all compartments. Used to slice into any array storing
        data across all compartments.
    """

    def __init__(self, blocks):
        self.blocks = list(blocks)
        self.slices = OrderedDict()

        assert self.dtype in (np.float32, np.int32)

        start_ix = end_ix = 0
        for block in self.blocks:
            end_ix += block.n_neurons
            self.slices[block] = slice(start_ix, end_ix)
            assert block.compartment.vth.dtype == self.dtype
            assert block.compartment.bias.dtype == self.dtype
            start_ix = end_ix

        self.n_compartments = end_ix

    @property
    def dtype(self):
        return self.blocks[0].compartment.vth.dtype


class IterableState:
    """Base class for aspects of the emulator state.

    This class takes the name of a LoihiBlock attribute as the
    ``block_key`` and maps these objects to their parent blocks and slices.

    Attributes
    ----------
    dtype : dtype
        Datatype of the state elements (given by the BlockInfo datatype).
    block_map : dict of {item: block}
        Maps an item (determined by ``block_key``) to the block
        it belongs to.
    n_compartments : int
        The total number of neuron compartments (given by BlockInfo).
    slices : dict of {item: slice}
        Maps an item to the ``block_info.slice`` for the block
        it belongs to.
    strict : bool (Default: True)
        Whether "undesired" chip effects (ex. overflow) raise errors (``True``)
        or whether they only raise warnings (``False``).
    """

    def __init__(self, block_info, block_key, strict=True):
        self.n_compartments = block_info.n_compartments
        self.dtype = block_info.dtype
        self.strict = strict

        blocks_items = list(self._blocks_items(block_info.blocks, block_key))
        self.block_map = OrderedDict((item, block) for block, item in blocks_items)
        self.slices = OrderedDict(
            (item, block_info.slices[block]) for block, item in blocks_items
        )

    @staticmethod
    def _blocks_items(blocks, block_key):
        for block in blocks:
            if block_key == "compartment":
                # one item per block
                yield block, getattr(block, block_key)
            else:
                # multiple items per block (attribute is iterable)
                for item in getattr(block, block_key):
                    yield block, item

    def error(self, msg):
        if self.strict:
            raise SimulationError(msg)
        else:
            warnings.warn(msg)

    def items(self):
        return self.slices.items()


class CompartmentState(IterableState):
    """State representing the Compartments of all blocks."""

    MAX_DELAY = 1  # delay not yet implemented

    def __init__(self, block_info, strict=True):
        super().__init__(block_info, "compartment", strict=strict)

        # Initialize NumPy arrays to store compartment-related data
        self.input = np.zeros((self.MAX_DELAY, self.n_compartments), dtype=self.dtype)
        self.current = np.zeros(self.n_compartments, dtype=self.dtype)
        self.voltage = np.zeros(self.n_compartments, dtype=self.dtype)
        self.spiked = np.zeros(self.n_compartments, dtype=bool)
        self.spike_count = np.zeros(self.n_compartments, dtype=np.int32)
        self.ref_count = np.zeros(self.n_compartments, dtype=np.int32)

        self.decay_u = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.decay_v = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.scale_u = np.ones(self.n_compartments, dtype=self.dtype)
        self.scale_v = np.ones(self.n_compartments, dtype=self.dtype)

        self.vth = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.vmin = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.vmax = np.full(self.n_compartments, np.nan, dtype=self.dtype)

        self.bias = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.ref = np.full(self.n_compartments, np.nan, dtype=self.dtype)

        # Fill in arrays with parameters from CompartmentSegments
        for compartment, sl in self.items():
            self.decay_u[sl] = compartment.decay_u
            self.decay_v[sl] = compartment.decay_v
            if compartment.scale_u:
                self.scale_u[sl] = compartment.decay_u
            if compartment.scale_v:
                self.scale_v[sl] = compartment.decay_v
            self.vth[sl] = compartment.vth
            self.vmin[sl] = compartment.vmin
            self.vmax[sl] = compartment.vmax
            self.bias[sl] = compartment.bias
            self.ref[sl] = compartment.refract_delay

        assert not np.any(np.isnan(self.decay_u))
        assert not np.any(np.isnan(self.decay_v))
        assert not np.any(np.isnan(self.vth))
        assert not np.any(np.isnan(self.vmin))
        assert not np.any(np.isnan(self.vmax))
        assert not np.any(np.isnan(self.bias))
        assert not np.any(np.isnan(self.ref))

        if self.dtype == np.int32:
            assert (self.scale_u == 1).all()
            assert (self.scale_v == 1).all()
            self._decay_current = lambda x, u: decay_int(x, self.decay_u, offset=1) + u
            self._decay_voltage = lambda x, u: decay_int(x, self.decay_v) + u

            def overflow(x, bits, name=None):
                _, o = overflow_signed(x, bits=bits, out=x)
                if np.any(o):
                    self.error("Overflow" + (" in %s" % name if name else ""))

        elif self.dtype == np.float32:

            def decay_float(x, u, d, s):
                return (1 - d) * x + s * u

            self._decay_current = lambda x, u: decay_float(
                x, u, d=self.decay_u, s=self.scale_u
            )
            self._decay_voltage = lambda x, u: decay_float(
                x, u, d=self.decay_v, s=self.scale_v
            )

            def overflow(x, bits, name=None):
                pass  # do not do overflow in floating point

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        self._overflow = overflow

        self.noise = NoiseState(block_info)

    def advance_input(self):
        self.input[:-1] = self.input[1:]
        self.input[-1] = 0

    def update(self, rng):
        noise = self.noise.sample(rng)
        q0 = self.input[0, :]
        q0[~(self.noise.target_u)] += noise[~(self.noise.target_u)]
        self._overflow(q0, bits=Q_BITS, name="q0")

        self.current[:] = self._decay_current(self.current, q0)
        self._overflow(self.current, bits=U_BITS, name="current")
        u2 = self.current + self.bias
        u2[self.noise.target_u] += noise[self.noise.target_u]
        self._overflow(u2, bits=U_BITS, name="u2")

        self.voltage[:] = self._decay_voltage(self.voltage, u2)
        # We have not been able to create V overflow on the chip, so we do
        # not include it here. See github.com/nengo/nengo-loihi/issues/130
        # self.overflow(self.v, bits=V_BIT, name="V")

        np.clip(self.voltage, self.vmin, self.vmax, out=self.voltage)
        self.voltage[self.ref_count > 0] = 0
        # TODO^: don't zero voltage in case neuron is saving overshoot

        self.spiked[:] = self.voltage > self.vth
        self.voltage[self.spiked] = 0
        self.ref_count[self.spiked] = self.ref[self.spiked]
        # decrement ref_count
        np.clip(self.ref_count - 1, 0, None, out=self.ref_count)

        self.spike_count[self.spiked] += 1


class NoiseState(IterableState):
    """State representing the noise parameters for all compartments."""

    def __init__(self, block_info):
        super().__init__(block_info, "compartment")
        self.enabled = np.full(self.n_compartments, np.nan, dtype=bool)
        self.exp = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.mant_offset = np.full(self.n_compartments, np.nan, dtype=self.dtype)
        self.target_u = np.full(self.n_compartments, np.nan, dtype=bool)

        # Fill in arrays with parameters from Compartment
        for compartment, sl in self.items():
            self.enabled[sl] = compartment.enable_noise
            self.exp[sl] = compartment.noise_exp
            self.mant_offset[sl] = compartment.noise_offset
            self.target_u[sl] = compartment.noise_at_membrane

        if self.dtype == np.int32:
            # TODO: if we could do this mult with shifts, it'd be faster, but
            # numpy has no function taking a vector of positive/negative shifts
            self.mult = np.where(self.enabled, 2.0 ** (self.exp - 7), 0)
            self.mant_offset *= 64

            def uniform(rng, n=self.n_compartments):
                return rng.randint(-127, 128, size=n, dtype=np.int32)

        elif self.dtype == np.float32:
            self.mult = np.where(self.enabled, 10.0**self.exp, 0)

            def uniform(rng, n=self.n_compartments):
                return rng.uniform(-1, 1, size=n).astype(np.float32)

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        assert not np.any(np.isnan(self.enabled))
        assert not np.any(np.isnan(self.exp))
        assert not np.any(np.isnan(self.mant_offset))
        assert not np.any(np.isnan(self.target_u))
        assert not np.any(np.isnan(self.mult))

        self._uniform = uniform

    def sample(self, rng):
        x = self._uniform(rng)
        return ((x + self.mant_offset) * self.mult).astype(self.dtype)


class SynapseState(IterableState):
    """State representing all synapses.

    Attributes
    ----------
    pes_error_scale : float
        Scaling for the errors of PES learning rules.
    pes_errors : {Synapse: ndarray(n_neurons / 2)}
        Maps synapse to PES learning rule errors for those synapses.
    spikes_in : {Synapse: list}
        Maps synapse to a queue of input spikes targeting those synapses.
    traces : {Synapse: ndarray(Synapse.n_axons)}
        Maps synapse to trace values for each of their axons.
    trace_spikes : {Synapse: set}
        Maps synapse to a queue of input spikes waiting to be added to those
        synapse traces.
    """

    def __init__(self, block_info, pes_error_scale=1.0, strict=True):  # noqa: C901
        super().__init__(block_info, "synapses", strict=strict)

        self.pes_error_scale = pes_error_scale

        self.spikes_in = OrderedDict()
        self.traces = OrderedDict()
        self.trace_spikes = OrderedDict()
        self.pes_errors = OrderedDict()
        for synapse in self.slices:
            n = synapse.n_axons
            self.spikes_in[synapse] = []

            if synapse.learning:
                self.traces[synapse] = np.zeros(n, dtype=self.dtype)
                self.trace_spikes[synapse] = set()
                self.pes_errors[synapse] = np.zeros(
                    self.block_map[synapse].n_neurons // 2, dtype=self.dtype
                )
                # ^ Currently, PES learning only happens on Nodes, where we
                # have pairs of on/off neurons. Therefore, the number of error
                # dimensions is half the number of neurons.

        if self.dtype == np.int32:

            def stochastic_round(
                x, dtype=self.dtype, rng=None, clip=None, name="values"
            ):
                x_sign = np.sign(x).astype(dtype)
                x_frac, x_int = np.modf(np.abs(x))
                p = rng.rand(*x.shape)
                y = x_int.astype(dtype) + (x_frac > p)
                if clip is not None:
                    q = y > clip
                    if np.any(q):
                        warnings.warn("Clipping %s" % name)
                    y[q] = clip
                return x_sign * y

            def trace_round(x, rng=None):
                return stochastic_round(x, rng=rng, clip=127, name="synapse trace")

            def weight_update(synapse, delta_ws, rng=None):
                synapse_cfg = synapse.synapse_cfg
                wgt_exp = synapse_cfg.real_weight_exp
                shift_bits = synapse_cfg.shift_bits
                overflow = learn_overflow_bits(n_factors=2)
                for w, delta_w in zip(synapse.weights, delta_ws):
                    product = shift(
                        delta_w * synapse._lr_int,
                        LEARN_FRAC + synapse._lr_exp - overflow,
                    )
                    learn_w = shift(w, LEARN_FRAC - wgt_exp) + product
                    learn_w[:] = stochastic_round(
                        learn_w * 2 ** (-LEARN_FRAC - shift_bits),
                        clip=2 ** (8 - shift_bits) - 1,
                        rng=rng,
                        name="learning weights",
                    )
                    w[:] = np.left_shift(learn_w, wgt_exp + shift_bits)

        elif self.dtype == np.float32:

            def trace_round(x, rng=None):
                return x  # no rounding

            def weight_update(synapse, delta_ws, rng=None):
                for w, delta_w in zip(synapse.weights, delta_ws):
                    w += synapse.learning_rate * delta_w

        else:
            raise ValidationError(
                "dtype %r not supported" % self.dtype, attr="dtype", obj=block_info
            )

        self._trace_round = trace_round
        self._weight_update = weight_update

    def inject_current(self, t, spike_inputs, all_axons, spiked):
        # --- clear spikes going in to each synapse
        for spike_queue in self.spikes_in.values():
            spike_queue.clear()

        # --- inputs pass spikes to synapses
        if t >= 2:  # input spikes take one time-step to arrive
            for spike_input in spike_inputs:
                compartment_idxs = spike_input.spike_idxs(t - 1)
                for axon in spike_input.axons:
                    spikes = axon.map_spikes(compartment_idxs)
                    self.spikes_in[axon.target].extend(
                        s for s in spikes if s is not None
                    )

        # --- axons pass spikes to synapses
        for axon, a_idx in all_axons.items():
            compartment_idxs = spiked[a_idx].nonzero()[0]
            spikes = axon.map_spikes(compartment_idxs)
            self.spikes_in[axon.target].extend(s for s in spikes if s is not None)

    def update_input(self, input):
        for synapse, s_slice in self.items():
            qb = input[:, s_slice]

            for spike in self.spikes_in[synapse]:
                base = synapse.axon_compartment_base(spike.axon_idx)
                if base is None:
                    continue

                weights, indices = synapse.axon_weights_indices(
                    spike.axon_idx, atom=spike.atom
                )
                qb[0, base + indices] += weights

    def update_pes_errors(self, errors):
        # TODO: these are sent every timestep, but learning only happens every
        # `tepoch * 2**learn_k` timesteps (see Synapse). Need to average.
        for pes_errors in self.pes_errors.values():
            pes_errors[:] = 0

        for synapse, _, e in errors:
            pes_errors = self.pes_errors[synapse]
            assert pes_errors.shape == e.shape
            pes_errors += scale_pes_errors(e, scale=self.pes_error_scale)

    def update_weights(self, t, rng):
        for synapse, pes_error in self.pes_errors.items():
            if t % synapse.learn_epoch == 0:
                trace = self.traces[synapse]
                e = np.hstack([-pes_error, pes_error])
                delta_w = np.outer(trace, e)
                self._weight_update(synapse, delta_w, rng=rng)

    def update_traces(self, t, rng):
        for synapse in self.traces:
            trace_spikes = self.trace_spikes.get(synapse, None)
            if trace_spikes is not None:
                for spike in self.spikes_in[synapse]:
                    if spike.axon_idx in trace_spikes:
                        self.error("Synaptic trace spikes lost")
                    trace_spikes.add(spike.axon_idx)

            trace = self.traces.get(synapse, None)
            if trace is not None and t % synapse.train_epoch == 0:
                tau = synapse.tracing_tau
                decay = np.exp(-synapse.train_epoch / tau)
                trace1 = decay * trace
                trace1[list(trace_spikes)] += synapse.tracing_mag
                trace[:] = self._trace_round(trace1, rng=rng)
                trace_spikes.clear()


class AxonState(IterableState):
    """State representing all (output) Axons."""

    def __init__(self, block_info):
        super().__init__(block_info, "axons")


class ProbeState:
    """State representing all probes.

    Attributes
    ----------
    dt : float
        Time constant of the Emulator.
    filters : {nengo_loihi.Probe: function}
        Maps Probes to the filtering function for that probe.
    filter_pos : {nengo_loihi.Probe: int}
        Maps Probes to the position of their filter in the data.
    block_probes : {nengo_loihi.Probe: slice}
        Maps Probes to the BlockInfo slice for the block they are probing.
    input_probes : {nengo_loihi.Probe: SpikeInput}
        Maps Probes to the SpikeInput that they are probing.
    """

    def __init__(self, block_info, probes, dt):
        self.dt = dt

        self.probes = OrderedDict()
        for probe in probes:
            block_slices = OrderedDict(
                [(block, block_info.slices[block]) for block in probe.target]
            )
            self.probes[probe] = block_slices

        self.outputs = {}
        for probe, block_slices in self.probes.items():
            self.outputs[probe] = [[] for block in block_slices]

        self.probe_filter = ProbeFilter(dt=dt)

    def collect_probe_output(self, probe):
        """Collect and clear output from a particular probe"""
        assert isinstance(probe, LoihiProbe)
        out = probe.weight_outputs(self.outputs[probe])
        out = self.probe_filter(probe, out)
        self.clear_probe(probe)
        return out

    def clear_probe(self, probe):
        for block_output in self.outputs[probe]:
            block_output.clear()

    def send(self, t, probe, receiver):
        """Send probed data to the receiver node.

        Returns
        -------
        n_send_steps : int
            The number of steps sent to the receiver.
        """
        # we don't currently filter here, so make sure the probe isn't expecting it
        assert probe.synapse is None, "Filtering should be done on host-side connection"

        outputs = self.outputs[probe]
        n_send_steps = len(outputs[0])
        if n_send_steps > 0:
            assert t >= n_send_steps, (t, n_send_steps)
            t0 = t - n_send_steps

            x = probe.weight_outputs(outputs)
            for j, xx in enumerate(x):
                receiver.receive(self.dt * (t0 + j + 2), xx)

        self.clear_probe(probe)
        return n_send_steps

    def update(self, t, compartment):
        for probe, block_slices in self.probes.items():
            assert hasattr(compartment, probe.key)
            target_attr = getattr(compartment, probe.key)

            for k, (block, block_slice) in enumerate(block_slices.items()):
                output = target_attr[block_slice][probe.slice[k]]
                if output.base is not None:
                    # if not already copied by the probe slice, then copy
                    output = output.copy()

                self.outputs[probe][k].append(output)
Software emulator for Loihi chip behaviour.      Parameters     ----------     model : Model         Model specification that will be simulated.     seed : int, optional (Default: None)         A seed for all stochastic operations done in this simulator. # Probe data in `self.probes` is cleared after it is retrieved by a # `self.collect_probe_output` call, so nothing needs to happen here. # This function remains so that `nengo_loihi.Simulator` can call `clear_probes` # on all of its internal simulators without checking their types. # remove references to states to free memory (except probes) Collect and clear output from a particular probe Simulate for the given number of ``dt`` steps.          Parameters         ----------         steps : int             Number of steps to run the simulation for. Advance the simulation by 1 step (``dt`` seconds). Provide information about all the LoihiBlocks in the model.      Attributes     ----------     dtype : dtype         Datatype of the blocks. Either ``np.float32`` if the blocks are not         discretized or ``np.int32`` if they are. All blocks are the same.     blocks : list of LoihiBlock         List of all the blocks in the model.     n_compartments : int         Total number of compartments across all blocks.     slices : dict of {LoihiBlock: slice}         Maps each block to a slice for that block's compartments with         respect to all compartments. Used to slice into any array storing         data across all compartments. Base class for aspects of the emulator state.      This class takes the name of a LoihiBlock attribute as the     ``block_key`` and maps these objects to their parent blocks and slices.      Attributes     ----------     dtype : dtype         Datatype of the state elements (given by the BlockInfo datatype).     block_map : dict of {item: block}         Maps an item (determined by ``block_key``) to the block         it belongs to.     n_compartments : int         The total number of neuron compartments (given by BlockInfo).     slices : dict of {item: slice}         Maps an item to the ``block_info.slice`` for the block         it belongs to.     strict : bool (Default: True)         Whether "undesired" chip effects (ex. overflow) raise errors (``True``)         or whether they only raise warnings (``False``). # one item per block # multiple items per block (attribute is iterable) State representing the Compartments of all blocks. # delay not yet implemented # Initialize NumPy arrays to store compartment-related data # Fill in arrays with parameters from CompartmentSegments # do not do overflow in floating point # We have not been able to create V overflow on the chip, so we do # not include it here. See github.com/nengo/nengo-loihi/issues/130 # self.overflow(self.v, bits=V_BIT, name="V") # TODO^: don't zero voltage in case neuron is saving overshoot # decrement ref_count State representing the noise parameters for all compartments. # Fill in arrays with parameters from Compartment # TODO: if we could do this mult with shifts, it'd be faster, but # numpy has no function taking a vector of positive/negative shifts State representing all synapses.      Attributes     ----------     pes_error_scale : float         Scaling for the errors of PES learning rules.     pes_errors : {Synapse: ndarray(n_neurons / 2)}         Maps synapse to PES learning rule errors for those synapses.     spikes_in : {Synapse: list}         Maps synapse to a queue of input spikes targeting those synapses.     traces : {Synapse: ndarray(Synapse.n_axons)}         Maps synapse to trace values for each of their axons.     trace_spikes : {Synapse: set}         Maps synapse to a queue of input spikes waiting to be added to those         synapse traces. # noqa: C901 # ^ Currently, PES learning only happens on Nodes, where we # have pairs of on/off neurons. Therefore, the number of error # dimensions is half the number of neurons. # no rounding # --- clear spikes going in to each synapse # --- inputs pass spikes to synapses # input spikes take one time-step to arrive # --- axons pass spikes to synapses # TODO: these are sent every timestep, but learning only happens every # `tepoch * 2**learn_k` timesteps (see Synapse). Need to average. State representing all (output) Axons. State representing all probes.      Attributes     ----------     dt : float         Time constant of the Emulator.     filters : {nengo_loihi.Probe: function}         Maps Probes to the filtering function for that probe.     filter_pos : {nengo_loihi.Probe: int}         Maps Probes to the position of their filter in the data.     block_probes : {nengo_loihi.Probe: slice}         Maps Probes to the BlockInfo slice for the block they are probing.     input_probes : {nengo_loihi.Probe: SpikeInput}         Maps Probes to the SpikeInput that they are probing. Collect and clear output from a particular probe Send probed data to the receiver node.          Returns         -------         n_send_steps : int             The number of steps sent to the receiver. # we don't currently filter here, so make sure the probe isn't expecting it # if not already copied by the probe slice, then copy