bsb-nest 6.0.0a5__py3-none-any.whl

bsb_nest/__init__.py

@@ -0,0 +1,27 @@
+"""
+NEST simulation adapter for the BSB framework.
+"""
+
+from bsb import SimulationBackendPlugin
+
+from .adapter import NestAdapter
+from .devices import (
+    DCGenerator,
+    Multimeter,
+    PoissonGenerator,
+    SinusoidalPoissonGenerator,
+    SpikeRecorder,
+)
+from .simulation import NestSimulation
+
+__plugin__ = SimulationBackendPlugin(Simulation=NestSimulation, Adapter=NestAdapter)
+
+__all__ = [
+    "DCGenerator",
+    "Multimeter",
+    "NestAdapter",
+    "NestSimulation",
+    "PoissonGenerator",
+    "SinusoidalPoissonGenerator",
+    "SpikeRecorder",
+]

bsb_nest/adapter.py

@@ -0,0 +1,203 @@
+import contextlib
+import sys
+import typing
+
+import nest
+from bsb import (
+    AdapterError,
+    AdapterProgress,
+    SimulationData,
+    SimulationResult,
+    SimulatorAdapter,
+    report,
+    warn,
+)
+from neo import SpikeTrain
+from tqdm import tqdm
+
+from .exceptions import KernelWarning, NestConnectError, NestModelError, NestModuleError
+
+if typing.TYPE_CHECKING:
+    from .simulation import NestSimulation
+
+
+class NestResult(SimulationResult):
+    def record(self, nc, **annotations):
+        recorder = nest.Create("spike_recorder", params={"record_to": "memory"})
+        nest.Connect(nc, recorder)
+
+        def flush(segment):
+            events = recorder.events[0]
+
+            segment.spiketrains.append(
+                SpikeTrain(
+                    events["times"],
+                    array_annotations={"senders": events["senders"]},
+                    t_stop=nest.biological_time,
+                    units="ms",
+                    **annotations,
+                )
+            )
+
+        self.create_recorder(flush)
+
+
+class NestAdapter(SimulatorAdapter):
+    def __init__(self, comm=None):
+        super().__init__(comm=comm)
+        self.loaded_modules = set()
+
+    def simulate(self, *simulations, post_prepare=None):
+        try:
+            self.reset_kernel()
+            return super().simulate(*simulations, post_prepare=post_prepare)
+        finally:
+            self.reset_kernel()
+
+    def prepare(self, simulation):
+        """
+        Prepare the simulation environment in NEST.
+
+        This method initializes internal data structures and performs all
+        setup steps required before running the simulation:
+
+        - Loads and installs required NEST modules.
+        - Applies simulation-level settings (e.g., resolution, verbosity, seed).
+        - Creates neuron populations based on cell models.
+        - Establishes connectivity between neurons using connection models.
+        - Instantiates devices (e.g., recorders, stimuli) used in the simulation.
+
+        If any error occurs during preparation, the corresponding internal state
+        is cleaned up to avoid partial setups.
+
+        :param simulation: The simulation configuration to prepare.
+        :type simulation: NestSimulation
+        :returns: The prepared simulation data associated with the given simulation.
+        :rtype: bsb.simulation.adapter.SimulationData
+        """
+        self.simdata[simulation] = SimulationData(
+            simulation, result=NestResult(simulation)
+        )
+        try:
+            report("Installing  NEST modules...", level=2)
+            self.load_modules(simulation)
+            self.set_settings(simulation)
+            report("Creating neurons...", level=2)
+            self.create_neurons(simulation)
+            report("Creating connections...", level=2)
+            self.connect_neurons(simulation)
+            report("Creating devices...", level=2)
+            self.create_devices(simulation)
+            return self.simdata[simulation]
+        except Exception:
+            del self.simdata[simulation]
+            raise
+
+    def reset_kernel(self):
+        nest.ResetKernel()
+        # Reset which modules we should consider explicitly loaded by the user
+        # to appropriately warn them when they load them twice.
+        self.loaded_modules = set()
+
+    def run(self, *simulations):
+        unprepared = [sim for sim in simulations if sim not in self.simdata]
+        if unprepared:
+            raise AdapterError(f"Unprepared for simulations: {', '.join(unprepared)}")
+        report("Simulating...", level=2)
+        duration = max(sim.duration for sim in simulations)
+        progress = AdapterProgress(duration)
+        try:
+            with nest.RunManager():
+                for oi, i in progress.steps(step=1):
+                    nest.Run(i - oi)
+                    progress.tick(i)
+        finally:
+            results = [self.simdata[sim].result for sim in simulations]
+            for sim in simulations:
+                del self.simdata[sim]
+        progress.complete()
+        report("Simulation done.", level=2)
+        return results
+
+    def load_modules(self, simulation):
+        for module in simulation.modules:
+            try:
+                nest.Install(module)
+                self.loaded_modules.add(module)
+            except Exception as e:
+                if e.errorname == "DynamicModuleManagementError":
+                    if "loaded already" in e.message:
+                        # Modules stay loaded in between `ResetKernel` calls.
+                        # If the module is not in the `loaded_modules` set, then
+                        # it's the first time this `reset`/`prepare` cycle,
+                        # and there is no user-side issue.
+                        if module in self.loaded_modules:
+                            warn(f"Already loaded '{module}'.", KernelWarning)
+                    elif "file not found" in e.message:
+                        raise NestModuleError(f"Module {module} not found") from None
+                    else:
+                        raise
+                else:
+                    raise
+
+    def create_neurons(self, simulation):
+        """
+        Create a population of nodes in the NEST simulator based on the cell model
+        configurations.
+        """
+        simdata = self.simdata[simulation]
+        for cell_model in simulation.cell_models.values():
+            simdata.populations[cell_model] = cell_model.create_population(simdata)
+
+    def connect_neurons(self, simulation):
+        """
+        Connect the cells in NEST according to the connection model configurations
+        """
+        simdata = self.simdata[simulation]
+        iter = simulation.connection_models.values()
+        if self.comm.get_rank() == 0:
+            iter = tqdm(iter, desc="", file=sys.stdout)
+        for connection_model in iter:
+            with contextlib.suppress(AttributeError):
+                # Only rank 0 should report progress bar
+                iter.set_description(connection_model.name)
+            cs = simulation.scaffold.get_connectivity_set(
+                connection_model.tag or connection_model.name
+            )
+            try:
+                pre_nodes = simdata.populations[simulation.get_model_of(cs.pre_type)]
+            except KeyError:
+                raise NestModelError(f"No model found for {cs.pre_type}") from None
+            try:
+                post_nodes = simdata.populations[simulation.get_model_of(cs.post_type)]
+            except KeyError:
+                raise NestModelError(f"No model found for {cs.post_type}") from None
+            try:
+                simdata.connections[connection_model] = (
+                    connection_model.create_connections(
+                        simdata, pre_nodes, post_nodes, cs, self.comm
+                    )
+                )
+            except Exception:
+                raise NestConnectError(
+                    f"{connection_model} error during connect."
+                ) from None
+
+    def create_devices(self, simulation):
+        simdata = self.simdata[simulation]
+        for device_model in simulation.devices.values():
+            device_model.implement(self, simulation, simdata)
+
+    def set_settings(self, simulation: "NestSimulation"):
+        nest.set_verbosity(simulation.verbosity)
+        nest.resolution = simulation.resolution
+        nest.overwrite_files = True
+        if simulation.seed is not None:
+            nest.rng_seed = simulation.seed
+
+    def check_comm(self):
+        if nest.NumProcesses() != self.comm.get_size():
+            raise RuntimeError(
+                f"NEST is managing {nest.NumProcesses()} processes, but "
+                f"{self.comm.get_size()} were detected. Please check your MPI setup."
+            )

bsb_nest/cell.py

@@ -0,0 +1,32 @@
+import nest
+from bsb import CellModel, config
+
+from .distributions import NestRandomDistribution, nest_parameter
+
+
+@config.node
+class NestCell(CellModel):
+    model = config.attr(type=str, default="iaf_psc_alpha")
+    """Importable reference to the NEST model describing the cell type."""
+    constants = config.dict(type=nest_parameter())
+    """Dictionary of the constants values to assign to the cell model."""
+
+    def create_population(self, simdata):
+        n = len(simdata.placement[self])
+        population = nest.Create(self.model, n) if n else nest.NodeCollection([])
+        self.set_constants(population)
+        self.set_parameters(population, simdata)
+        return population
+
+    def set_constants(self, population):
+        population.set(
+            {
+                k: (v() if isinstance(v, NestRandomDistribution) else v)
+                for k, v in self.constants.items()
+            }
+        )
+
+    def set_parameters(self, population, simdata):
+        ps = simdata.placement[self]
+        for param in self.parameters:
+            population.set(param.name, param.get_value(ps))

bsb_nest/connection.py

@@ -0,0 +1,193 @@
+import functools
+import sys
+
+import nest
+import numpy as np
+import psutil
+from bsb import ConnectionModel, compose_nodes, config, types
+from tqdm import tqdm
+
+from .distributions import nest_parameter
+from .exceptions import NestConnectError
+
+
+@config.node
+class NestSynapseSettings:
+    """
+    Class interfacing a NEST synapse model.
+    """
+
+    model = config.attr(type=str, default="static_synapse")
+    """Importable reference to the NEST model describing the synapse type."""
+    weight = config.attr(type=float, required=True)
+    """Weight of the connection between the presynaptic and the postsynaptic cells."""
+    delay = config.attr(type=float, required=True)
+    """Delay of the transmission between the presynaptic and the postsynaptic cells."""
+    receptor_type = config.attr(type=int)
+    """Index of the postsynaptic receptor to target."""
+    constants = config.catch_all(type=nest_parameter())
+    """Dictionary of the constants values to assign to the synapse model."""
+
+
+@config.node
+class NestConnectionSettings:
+    """
+    Class interfacing a NEST connection rule.
+    """
+
+    rule = config.attr(type=str)
+    """Importable reference to the Nest connection rule used to connect the cells."""
+    constants = config.catch_all(type=types.any_())
+    """Dictionary of parameters to assign to the connection rule."""
+
+
+class LazySynapseCollection:
+    def __init__(self, pre, post):
+        self._pre = pre
+        self._post = post
+
+    def __len__(self):
+        return self.collection.__len__()
+
+    def __str__(self):
+        return self.collection.__str__()
+
+    def __iter__(self):
+        return iter(self.collection)
+
+    def __getattr__(self, attr):
+        return getattr(self.collection, attr)
+
+    @functools.cached_property
+    def collection(self):
+        return nest.GetConnections(self._pre, self._post)
+
+
+@config.dynamic(attr_name="model_strategy", required=False)
+class NestConnection(compose_nodes(NestConnectionSettings, ConnectionModel)):
+    """
+    Class interfacing a NEST connection, including its connection rule and synaptic
+    parameters.
+    """
+
+    model_strategy: str
+    """
+    Specifies the strategy used by the connection model for synapse creation and
+    management.
+    """
+
+    synapse = config.attr(type=NestSynapseSettings, required=True)
+    """Nest synapse model with its parameters."""
+
+    def create_connections(self, simdata, pre_nodes, post_nodes, cs, comm):
+        import nest
+
+        syn_spec = self.get_syn_spec()
+        if syn_spec["synapse_model"] not in nest.Models(mtype="synapses"):
+            raise NestConnectError(
+                f"Unknown synapse model '{syn_spec['synapse_model']}'."
+            )
+        if self.rule is not None:
+            nest.Connect(pre_nodes, post_nodes, self.get_conn_spec(), syn_spec)
+        else:
+            comm.barrier()
+            for pre_locs, post_locs in self.predict_mem_iterator(
+                pre_nodes, post_nodes, cs, comm
+            ):
+                comm.barrier()
+                if len(pre_locs) == 0 or len(post_locs) == 0:
+                    continue
+                cell_pairs, multiplicity = np.unique(
+                    np.column_stack((pre_locs[:, 0], post_locs[:, 0])),
+                    return_counts=True,
+                    axis=0,
+                )
+                prel = pre_nodes.tolist()
+                postl = post_nodes.tolist()
+                ssw = {**syn_spec}
+                bw = syn_spec["weight"]
+                ssw["weight"] = [bw * m for m in multiplicity]
+                ssw["delay"] = [syn_spec["delay"]] * len(ssw["weight"])
+                nest.Connect(
+                    [prel[x] for x in cell_pairs[:, 0]],
+                    [postl[x] for x in cell_pairs[:, 1]],
+                    "one_to_one",
+                    ssw,
+                    return_synapsecollection=False,
+                )
+            comm.barrier()
+        return LazySynapseCollection(pre_nodes, post_nodes)
+
+    def predict_mem_iterator(self, pre_nodes, post_nodes, cs, comm):
+        avmem = psutil.virtual_memory().available
+        predicted_all_mem = (
+            len(pre_nodes) * 8 * 2 + len(post_nodes) * 8 * 2 + len(cs) * 6 * 8 * (16 + 2)
+        ) * comm.get_size()
+        n_chunks = len(cs.get_local_chunks("out"))
+        predicted_local_mem = (predicted_all_mem / n_chunks) if n_chunks > 0 else 0.0
+        if predicted_local_mem > avmem / 2:
+            # Iterate block-by-block
+            return self.block_iterator(cs, comm)
+        elif predicted_all_mem > avmem / 2:
+            # Iterate local hyperblocks
+            return self.local_iterator(cs, comm)
+        else:
+            # Iterate all
+            return (cs.load_connections().as_globals().all(),)
+
+    def block_iterator(self, cs, comm):
+        locals = cs.get_local_chunks("out")
+
+        def block_iter():
+            iter = locals
+            if comm.get_rank() == 0:
+                iter = tqdm(iter, desc="hyperblocks", file=sys.stdout)
+            for local in iter:
+                inner_iter = cs.load_connections().as_globals().from_(local)
+                if comm.get_rank() == 0:
+                    yield from tqdm(
+                        inner_iter,
+                        desc="blocks",
+                        total=len(cs.get_global_chunks("out", local)),
+                        file=sys.stdout,
+                        leave=False,
+                    )
+                else:
+                    yield from inner_iter
+
+        return block_iter()
+
+    def local_iterator(self, cs, comm):
+        iter = cs.get_local_chunks("out")
+        if comm.get_rank() == 0:
+            iter = tqdm(iter, desc="hyperblocks", file=sys.stdout)
+        yield from (
+            cs.load_connections().as_globals().from_(local).all() for local in iter
+        )
+
+    def get_connectivity_set(self):
+        if self.tag is not None:
+            return self.scaffold.get_connectivity_set(self.tag)
+        else:
+            return self.connection_model
+
+    def get_conn_spec(self):
+        return {
+            "rule": self.rule,
+            **self.constants,
+        }
+
+    def get_syn_spec(self):
+        return {
+            **{
+                label: value
+                for attr, label in (
+                    ("model", "synapse_model"),
+                    ["weight"] * 2,
+                    ["delay"] * 2,
+                    ["receptor_type"] * 2,
+                )
+                if (value := getattr(self.synapse, attr)) is not None
+            },
+            **self.synapse.constants,
+        }

bsb_nest/device.py

@@ -0,0 +1,149 @@
+import abc
+import warnings
+
+import nest
+from bsb import DeviceModel, Targetting, config, refs, types
+
+
+@config.node
+class NestRule:
+    """
+    Interface to connect a device directly through the NEST interface.
+    """
+
+    rule = config.attr(type=str, required=True)
+    """Connection rule to connect """
+    constants = config.catch_all(type=types.any_())
+    """Dictionary of parameters for the targetting rule."""
+    cell_models = config.reflist(refs.sim_cell_model_ref)
+    """Reference to the Nest cell model to target with the Device"""
+
+
+@config.dynamic(attr_name="device", auto_classmap=True, default="external")
+class NestDevice(DeviceModel):
+    device: str
+    """Name of the NEST device model (e.g., "spike_generator", "poisson_generator")."""
+    weight = config.attr(type=float, required=True)
+    """weight of the connection between the device and its target"""
+    delay = config.attr(type=float, required=True)
+    """delay of the transmission between the device and its target"""
+    targetting = config.attr(
+        type=types.or_(Targetting, NestRule), default=dict, call_default=True
+    )
+    """Targets of the device, which should be either a population or a nest rule"""
+    receptor_type = config.attr(type=int, required=False, default=0)
+    """Integer ID of the postsynaptic target receptor"""
+
+    def get_dict_targets(
+        self,
+        adapter,
+        simulation,
+        simdata,
+    ) -> dict:
+        """
+        Get a dictionary from a target group to its NEST Collection
+        for each target group of the device.
+
+        :param bsb_nest.adapter.NestAdapter adapter: Nest adapter instance
+        :param bsb_nest.simulation.NestSimulation simulation: Nest simulation instance
+        :param bsb.simulation.adapter.SimulationData simdata: Simulation data instance
+        :return: dictionary of device target group to NEST Collection
+        :rtype: dict
+        """
+        if isinstance(self.targetting, Targetting):
+            node_collector = self.targetting.get_targets(adapter, simulation, simdata)
+        else:
+            node_collector = {
+                model: simdata.populations[model][targets]
+                for model, targets in simdata.populations.items()
+                if not self.targetting.cell_models or model in self.targetting.cell_models
+            }
+        return node_collector
+
+    @staticmethod
+    def _flatten_nodes_ids(dict_targets):
+        return sum(dict_targets.values(), start=nest.NodeCollection())
+
+    @staticmethod
+    def _invert_targets_dict(dict_targets):
+        return {elem: k.name for k, v in dict_targets.items() for elem in v.tolist()}
+
+    def get_target_nodes(
+        self,
+        adapter,
+        simulation,
+        simdata,
+    ):
+        """
+        Get the NEST Collection of the targets of the device.
+
+        :param bsb_nest.adapter.NestAdapter adapter:
+        :param bsb_nest.simulation.NestSimulation simulation: Nest simulation instance
+        :param bsb.simulation.adapter.SimulationData simdata: Simulation data instance
+        :return: Flattened NEST collection with all the targets of the device
+        """
+        targets_dict = self.get_dict_targets(adapter, simulation, simdata)
+        return self._flatten_nodes_ids(targets_dict)
+
+    def connect_to_nodes(self, device, nodes):
+        if len(nodes) == 0:
+            warnings.warn(f"{self.name} has no targets", stacklevel=2)
+        else:
+            try:
+                nest.Connect(
+                    device,
+                    nodes,
+                    syn_spec={
+                        "weight": self.weight,
+                        "delay": self.delay,
+                        "receptor_type": self.receptor_type,
+                    },
+                )
+
+            except Exception as e:
+                if "does not send output" not in str(e):
+                    raise
+                nest.Connect(
+                    nodes,
+                    device,
+                    syn_spec={"weight": self.weight, "delay": self.delay},
+                )
+
+    def register_device(self, simdata, device):
+        simdata.devices[self] = device
+        return device
+
+    @abc.abstractmethod
+    def implement(
+        self,
+        adapter,
+        simulation,
+        simdata,
+    ):
+        """
+        Create, connect and register the Nest device.
+
+        :param bsb_nest.adapter.NestAdapter adapter:
+        :param bsb_nest.simulation.NestSimulation simulation: Nest simulation instance
+        :param bsb.simulation.adapter.SimulationData simdata: Simulation data instance
+        """
+        pass
+
+
+@config.node
+class ExtNestDevice(NestDevice, classmap_entry="external"):
+    """
+    Class interfacing Nest devices.
+    """
+
+    nest_model = config.attr(type=str, required=True)
+    """Importable reference to the NEST model describing the device type."""
+    constants = config.dict(type=types.or_(types.number(), str))
+    """Dictionary of the constants values to assign to the device model."""
+
+    def implement(self, adapter, simulation, simdata):
+        simdata.devices[self] = device = nest.Create(
+            self.nest_model, params=self.constants
+        )
+        nodes = self.get_target_nodes(adapter, simulation, simdata)
+        self.connect_to_nodes(device, nodes)

bsb_nest/devices/__init__.py

@@ -0,0 +1,13 @@
+from .dc_generator import DCGenerator
+from .multimeter import Multimeter
+from .poisson_generator import PoissonGenerator
+from .sinusoidal_poisson_generator import SinusoidalPoissonGenerator
+from .spike_recorder import SpikeRecorder
+
+__all__ = [
+    "DCGenerator",
+    "Multimeter",
+    "PoissonGenerator",
+    "SinusoidalPoissonGenerator",
+    "SpikeRecorder",
+]

bsb_nest/devices/dc_generator.py

@@ -0,0 +1,23 @@
+import nest
+from bsb import config
+
+from ..device import NestDevice
+
+
+@config.node
+class DCGenerator(NestDevice, classmap_entry="dc_generator"):
+    amplitude = config.attr(type=float, required=True)
+    """Current amplitude of the dc generator"""
+    start = config.attr(type=float, required=False, default=0.0)
+    """Activation time in ms"""
+    stop = config.attr(type=float, required=False, default=None)
+    """Deactivation time in ms. 
+        If not specified, generator will last until the end of the simulation."""
+
+    def implement(self, adapter, simulation, simdata):
+        nodes = self.get_target_nodes(adapter, simulation, simdata)
+        params = {"amplitude": self.amplitude, "start": self.start}
+        if self.stop is not None and self.stop > self.start:
+            params["stop"] = self.stop
+        device = self.register_device(simdata, nest.Create("dc_generator", params=params))
+        self.connect_to_nodes(device, nodes)