Update Arbor backend to work with Arbor 0.10.0

.github/workflows/full-test.yml

@@ -58,7 +58,7 @@ jobs:
     - name: Install Arbor
       if: startsWith(matrix.os, 'ubuntu')
       run: |
-        python -m pip install arbor==0.9.0 libNeuroML morphio
+        python -m pip install arbor==0.10.0 libNeuroML morphio
     - name: Install NESTML
       if: startsWith(matrix.os, 'ubuntu')
       run: |

Makefile

@@ -97,7 +97,7 @@ $(NEST_STAMP): $(NEST_VENV_STAMP) $(NEST_SRC_UNPACKED)
 	    $(CURDIR)/pyNN/nest/extensions
 	cd $(NEST_BUILD_DIR)/pynn_extensions && make install
 	$(NEST_PIP) install \
-	    "neuron>=9.0.0" nrnutils "arbor==0.9.0" \
+	    "neuron>=9.0.0" nrnutils "arbor==0.10.0" \
 	    brian2 libNeuroML scipy matplotlib Cheetah3 h5py Jinja2 \
 	    pytest pytest-xdist pytest-cov flake8 morphio nestml
 	$(NEST_PIP) install -e .

pyNN/arbor/cells.py

@@ -2,12 +2,21 @@
 from collections import defaultdict
 from lazyarray import larray
 import arbor
+from arbor import units as U
 from neuroml import Point3DWithDiam
 from ..morphology import Morphology, NeuroMLMorphology, MorphIOMorphology, IonChannelDistribution
 from ..models import BaseCellType
 from ..parameters import ParameterSpace
 
 
+# Units for the parameters of current-source mechanisms (e.g. iclamp).
+ELECTRODE_PARAM_UNITS = {
+    "tstart": U.ms,
+    "duration": U.ms,
+    "current": U.nA,
+}
+
+
 def convert_point(p3d: Point3DWithDiam) -> arbor.mpoint:
     return arbor.mpoint(p3d.x, p3d.y, p3d.z, p3d.diameter/2)
 
@@ -64,7 +73,7 @@ def _build_tree(self, i):
                 if segment.name not in self.labels:
                     self.labels[segment.name] = f"(segment {i})"
         elif isinstance(std_morphology, MorphIOMorphology):
-            tree = arbor.load_swc_neuron(std_morphology.morphology_file, raw=True)
+            tree = arbor.load_swc_neuron(std_morphology.morphology_file).segment_tree
         else:
             raise ValueError("{} not supported as a neuron morphology".format(type(std_morphology)))
 
@@ -89,18 +98,22 @@ def _build_decor(self, i):
         decor = arbor.decor()
         # Set the default properties of the cell (this overrides the model defaults).
         decor.set_property(
-            cm=self.parameters["cm"][i] * 0.01,  # µF/cm² -> F/m²
-            rL=self.parameters["Ra"][i] * 1,     # Ω·cm
-            Vm=self.initial_values["v"][i]
+            cm=self.parameters["cm"][i] * U.uF / U.cm2,
+            rL=self.parameters["Ra"][i] * U.Ohm * U.cm,
+            Vm=self.initial_values["v"][i] * U.mV
         )
         if not self.parameters["ionic_species"]._evaluated:
             self.parameters["ionic_species"].evaluate(simplify=True)
         for ion_name, ionic_species in self.parameters["ionic_species"].items():
             assert ion_name == ionic_species.ion_name
-            decor.set_ion(ion_name,
-                          int_con=ionic_species.internal_concentration,
-                          ext_con=ionic_species.external_concentration,
-                          rev_pot=ionic_species.reversal_potential)  # method="nernst/na")
+            ion_kwargs = {}
+            if ionic_species.internal_concentration is not None:
+                ion_kwargs["int_con"] = ionic_species.internal_concentration * U.mM
+            if ionic_species.external_concentration is not None:
+                ion_kwargs["ext_con"] = ionic_species.external_concentration * U.mM
+            if ionic_species.reversal_potential is not None:
+                ion_kwargs["rev_pot"] = ionic_species.reversal_potential * U.mV
+            decor.set_ion(ion_name, **ion_kwargs)  # method="nernst/na")
         for native_name, region_params in mechanism_parameters.items():
             for region, params in region_params.items():
                 if native_name == "hh":
@@ -124,12 +137,16 @@ def _build_decor(self, i):
             location_generator = current_source["location_generator"]
             mechanism = getattr(arbor, current_source["model_name"])
             for locset, label in location_generator.generate_locations(morph, label=f"{current_source['model_name']}_label"):
-                params = current_source["parameters"].evaluate(simplify=True)
+                params = current_source["parameters"].evaluate(simplify=True).as_dict()
+                params = {
+                    name: value * ELECTRODE_PARAM_UNITS[name] if name in ELECTRODE_PARAM_UNITS else value
+                    for name, value in params.items()
+                }
                 mech = mechanism(**params)
                 decor.place(locset, mech, label)
 
         # add spike source
-        decor.place('"root"', arbor.threshold_detector(-10), "detector")
+        decor.place('"root"', arbor.threshold_detector(-10 * U.mV), "detector")
         # todo: allow user to choose location and threshold value
 
         policy = arbor.cv_policy_max_extent(10.0)

pyNN/arbor/populations.py

@@ -87,7 +87,17 @@ def arbor_cell_description(self, gid):
             for key, value in params.items():
                 if isinstance(value, Sequence):
                     params[key] = value.value
-            schedule = self.celltype.arbor_schedule(**params)
+            schedule_units = getattr(self.celltype, "arbor_schedule_units", {})
+            schedule_params = {}
+            for key, value in params.items():
+                unit = schedule_units.get(key)
+                if unit is None or value is None:
+                    schedule_params[key] = value
+                elif isinstance(value, (list, tuple, np.ndarray)):
+                    schedule_params[key] = [float(v) * unit for v in value]
+                else:
+                    schedule_params[key] = value * unit
+            schedule = self.celltype.arbor_schedule(**schedule_params)
             return arbor.spike_source_cell("spike-source", schedule)
         else:
             args = self._arbor_cell_description[index]

pyNN/arbor/recording.py

@@ -5,6 +5,7 @@
 from collections import defaultdict
 import numpy as np
 import arbor
+from arbor import units as U
 
 from .. import recording
 from . import simulator
@@ -58,18 +59,24 @@ def _localize_variables(self, variables, locations):
         return resolved_variables
 
     def _set_arbor_sim(self, arbor_sim):
+        # Since Arbor 0.10.0, probes are addressed by (gid, tag) rather than by
+        # a positional cell_member index. The tag assigned here (a per-gid probe
+        # counter, in the same iteration order as _get_arbor_probes) must match
+        # the tag given to the corresponding probe there.
         self.handles = defaultdict(list)
         probe_indices = defaultdict(int)
         for variable in self.recorded:
             if variable.name != "spikes":
                 for cell in self.recorded[variable]:
-                    probeset_id = arbor.cell_member(cell.gid, probe_indices[cell.gid])
+                    tag = str(probe_indices[cell.gid])
                     probe_indices[cell.gid] += 1
-                    handle = arbor_sim.sample(probeset_id, arbor.regular_schedule(self.sampling_interval))
+                    handle = arbor_sim.sample(
+                        cell.gid, tag, arbor.regular_schedule(self.sampling_interval * U.ms))
                     self.handles[variable].append(handle)
 
     def _get_arbor_probes(self, gid):
         probes = []
+        probe_index = 0
         for variable in self.recorded:
             if variable.location is None:
                 pass
@@ -79,12 +86,15 @@ def _get_arbor_probes(self, gid):
             if gid in [cell.gid for cell in self.recorded[variable]]:
                 if variable.name == "spikes":
                     continue
-                elif variable.name == "v":
-                    probe = arbor.cable_probe_membrane_voltage(locset)
+                # Tag must match the one assigned in _set_arbor_sim (per-gid index).
+                tag = str(probe_index)
+                probe_index += 1
+                if variable.name == "v":
+                    probe = arbor.cable_probe_membrane_voltage(locset, tag)
                 else:
                     mech_name, state_name = variable.name.split(".")
                     arbor_model = mech_name  # to do: find_arbor_model(mech_name)
-                    probe = arbor.cable_probe_density_state(locset, arbor_model, state_name)
+                    probe = arbor.cable_probe_density_state(locset, arbor_model, state_name, tag)
                 probes.append(probe)
         return probes
 

pyNN/arbor/simulator.py

@@ -13,6 +13,7 @@
 except ImportError:
     pass
 import arbor
+from arbor import units as U
 from .. import common
 from ..core import find
 
@@ -196,7 +197,7 @@ def run(self, simtime):
             for recorder in self.recorders:
                 recorder._set_arbor_sim(self.arbor_sim)
         self.t += simtime
-        self.arbor_sim.run(self.t, self.dt)
+        self.arbor_sim.run(self.t * U.ms, self.dt * U.ms)
         self.running = True
 
     def run_until(self, tstop):

pyNN/arbor/standardmodels.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 import arbor
+from arbor import units as U
 
 from ..standardmodels import cells, ion_channels, synapses, electrodes, receptors, build_translations
 from ..parameters import ParameterSpace, IonicSpecies
@@ -32,6 +33,7 @@ class SpikeSourcePoisson(cells.SpikeSourcePoisson):
     # todo: manage "seed"
     arbor_cell_kind = arbor.cell_kind.spike_source
     arbor_schedule = arbor.poisson_schedule
+    arbor_schedule_units = {"tstart": U.ms, "freq": U.Hz, "tstop": U.ms}
 
 
 class SpikeSourceArray(cells.SpikeSourceArray):
@@ -42,6 +44,8 @@ class SpikeSourceArray(cells.SpikeSourceArray):
     )
     arbor_cell_kind = arbor.cell_kind.spike_source
     arbor_schedule = arbor.explicit_schedule
+    # Since Arbor 0.10.0, schedule parameters must be unit-typed.
+    arbor_schedule_units = {"times": U.ms}
 
 
 class BaseCurrentSource(object):

pyproject.toml

@@ -42,7 +42,7 @@ sonata = ["h5py"]
 morphologies = ["morphio"]
 neuron = ["neuron", "nrnutils"]
 brian2 = ["brian2"]
-arbor = ["arbor==0.9.0", "libNeuroML", "morphio"]
+arbor = ["arbor==0.10.0", "libNeuroML", "morphio"]
 spiNNaker = ["spyNNaker"]
 neuroml = ["libNeuroML"]
 nestml = ["nestml"]