Controls the passing of NEURON global params to processes (#165)

* control the passing of NEURON global params to processes

* a more general solution in IsolatedProcess

* no need to explicitly pass neuron global args

* future annotations add import

bluecellulab/__init__.py

@@ -7,7 +7,7 @@
 except ImportError:
     BLUEPY_AVAILABLE = False
 
-from .importer import *  # NOQA
+from bluecellulab.importer import import_hoc
 from .verbosity import *
 from .cell import Cell, create_ball_stick  # NOQA
 from .circuit import EmodelProperties
@@ -19,7 +19,7 @@
 from .simulation import Simulation  # NOQA
 from .rngsettings import RNGSettings  # NOQA
 from .circuit_simulation import CircuitSimulation, CircuitSimulation  # NOQA
-
+import neuron
 
 from .simulation.neuron_globals import NeuronGlobals
 

bluecellulab/analysis/inject_sequence.py

@@ -7,9 +7,9 @@
 import numpy as np
 from bluecellulab.cell.core import Cell
 from bluecellulab.cell.template import TemplateParams
+from bluecellulab.simulation.parallel import IsolatedProcess
 from bluecellulab.simulation.simulation import Simulation
 from bluecellulab.stimulus.factory import Stimulus, StimulusFactory
-from bluecellulab.utils import IsolatedProcess
 
 
 class StimulusName(Enum):

bluecellulab/circuit/circuit_access/bluepy_circuit_access.py

@@ -19,8 +19,9 @@
 from pathlib import Path
 from typing import Optional
 
+import neuron
 import pandas as pd
-from bluecellulab import circuit, neuron
+from bluecellulab import circuit
 from bluecellulab.circuit.circuit_access import EmodelProperties
 from bluecellulab.circuit.config import BluepySimulationConfig
 from bluecellulab.circuit.config.definition import SimulationConfig

bluecellulab/circuit/circuit_access/sonata_circuit_access.py

@@ -22,8 +22,9 @@
 from bluepysnap.circuit_ids import CircuitNodeId, CircuitEdgeIds
 from bluepysnap.exceptions import BluepySnapError
 from bluepysnap import Circuit as SnapCircuit
+import neuron
 import pandas as pd
-from bluecellulab import circuit, neuron
+from bluecellulab import circuit
 from bluecellulab.circuit.circuit_access.definition import EmodelProperties
 from bluecellulab.circuit import CellId, SynapseProperty
 from bluecellulab.circuit.config import SimulationConfig

bluecellulab/simulation/neuron_globals.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Module that handles the global NEURON parameters."""
 
+from typing import NamedTuple
 import neuron
 from bluecellulab.circuit.config.sections import Conditions, MechanismConditions
 from bluecellulab.exceptions import error_context
@@ -62,6 +63,11 @@ def set_minis_single_vesicle_values(mech_conditions: MechanismConditions) -> Non
             )
 
 
+class NeuronGlobalParams(NamedTuple):
+    temperature: float
+    v_init: float
+
+
 class NeuronGlobals:
     _instance = None
 
@@ -96,3 +102,10 @@ def v_init(self):
     def v_init(self, value):
         self._v_init = value
         neuron.h.v_init = value
+
+    def export_params(self) -> NeuronGlobalParams:
+        return NeuronGlobalParams(self.temperature, self.v_init)
+
+    def load_params(self, params: NeuronGlobalParams) -> None:
+        self.temperature = params.temperature
+        self.v_init = params.v_init

bluecellulab/simulation/parallel.py

@@ -0,0 +1,36 @@
+"""Controlled simulations in parallel."""
+
+from __future__ import annotations
+from multiprocessing.pool import Pool
+from bluecellulab.simulation.neuron_globals import NeuronGlobals
+
+
+class IsolatedProcess(Pool):
+    """Multiprocessing Pool that restricts a worker to run max 1 process.
+
+    Use this when running isolated NEURON simulations. Running 2 NEURON
+    simulations on a single process is to be avoided. Required global
+    NEURON simulation parameters will automatically be passed to each
+    worker.
+    """
+
+    def __init__(self, processes: int | None = 1):
+        """Initialize the IsolatedProcess pool.
+
+        Args:
+            processes: The number of processes to use for running the simulations.
+                       If set to None, then the number returned by os.cpu_count() is used.
+        """
+        neuron_global_params = NeuronGlobals.get_instance().export_params()
+        super().__init__(
+            processes=processes,
+            initializer=self.init_worker,
+            initargs=(neuron_global_params,),
+            maxtasksperchild=1,
+        )
+
+    @staticmethod
+    def init_worker(neuron_global_params):
+        """Load global parameters for the NEURON environment in each worker
+        process."""
+        NeuronGlobals.get_instance().load_params(neuron_global_params)

bluecellulab/tools.py

@@ -25,7 +25,8 @@
 import bluecellulab
 from bluecellulab.circuit.circuit_access import EmodelProperties
 from bluecellulab.exceptions import UnsteadyCellError
-from bluecellulab.utils import CaptureOutput, IsolatedProcess
+from bluecellulab.simulation.parallel import IsolatedProcess
+from bluecellulab.utils import CaptureOutput
 
 logger = logging.getLogger(__name__)
 

bluecellulab/utils.py

@@ -1,19 +1,21 @@
 """Utility functions used within BlueCellulab."""
+
 from __future__ import annotations
 import contextlib
 import io
 import json
-from multiprocessing.pool import Pool
 
 import numpy as np
 
 
 def run_once(func):
     """A decorator to ensure a function is only called once."""
+
     def wrapper(*args, **kwargs):
         if not wrapper.has_run:
             wrapper.has_run = True
             return func(*args, **kwargs)
+
     wrapper.has_run = False
     return wrapper
 
@@ -32,29 +34,25 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 class NumpyEncoder(json.JSONEncoder):
     def default(self, obj):
-        if isinstance(obj, (np.int_, np.intc, np.intp, np.int8,
-                            np.int16, np.int32, np.int64, np.uint8,
-                            np.uint16, np.uint32, np.uint64)):
+        if isinstance(
+            obj,
+            (
+                np.int_,
+                np.intc,
+                np.intp,
+                np.int8,
+                np.int16,
+                np.int32,
+                np.int64,
+                np.uint8,
+                np.uint16,
+                np.uint32,
+                np.uint64,
+            ),
+        ):
             return int(obj)
-        elif isinstance(obj, (np.float_, np.float16, np.float32,
-                              np.float64)):
+        elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)):
             return float(obj)
         elif isinstance(obj, np.ndarray):
             return obj.tolist()
         return json.JSONEncoder.default(self, obj)
-
-
-class IsolatedProcess(Pool):
-    """Multiprocessing Pool that restricts a worker to run max 1 process.
-
-    Use this when running isolated NEURON simulations. Running 2 NEURON
-    simulations on a single process is to be avoided.
-    """
-    def __init__(self, processes: int | None = 1):
-        """Initialize the IsolatedProcess pool.
-
-        Args:
-            processes: The number of processes to use for running the stimuli.
-            If set to None, then the number returned by os.cpu_count() is used.
-        """
-        super().__init__(processes=processes, maxtasksperchild=1)

examples/3-bpap/bpap.ipynb

@@ -65,12 +65,13 @@
     "from pathlib import Path\n",
     "from itertools import islice\n",
     "\n",
+    "import neuron\n",
     "import numpy as np\n",
     "from matplotlib.pyplot import get_cmap\n",
     "import matplotlib.pyplot as plt\n",
     "from scipy.optimize import curve_fit\n",
     "\n",
-    "from bluecellulab import Cell, neuron\n",
+    "from bluecellulab import Cell\n",
     "from bluecellulab.simulation import Simulation\n",
     "from bluecellulab.stimulus.circuit_stimulus_definitions import Hyperpolarizing\n",
     "from bluecellulab.circuit.circuit_access import EmodelProperties"

examples/4-epsp/epsp.ipynb

@@ -65,11 +65,12 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "from bluecellulab import Cell, neuron\n",
+    "from bluecellulab import Cell\n",
     "from bluecellulab.simulation import Simulation\n",
     "from bluecellulab.simulation.neuron_globals import NeuronGlobals\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
+    "import neuron\n",
     "import numpy as np\n",
     "import seaborn as sns\n",
     "\n",

tests/test_simulation/test_neuron_globals.py

@@ -4,7 +4,7 @@
 import neuron
 
 from bluecellulab.circuit.config.sections import ConditionEntry, Conditions, MechanismConditions
-from bluecellulab.simulation.neuron_globals import NeuronGlobals, set_global_condition_parameters, set_init_depleted_values, set_minis_single_vesicle_values
+from bluecellulab.simulation.neuron_globals import NeuronGlobalParams, NeuronGlobals, set_global_condition_parameters, set_init_depleted_values, set_minis_single_vesicle_values
 
 
 @mock.patch("neuron.h")
@@ -62,6 +62,25 @@ def test_neuron_globals():
     NeuronGlobals.get_instance().v_init = -70.0
     assert neuron.h.v_init == -70.0
 
+    # set back to default
+    NeuronGlobals.get_instance().temperature = 34.0
+    NeuronGlobals.get_instance().v_init = -65.0
+
     # exception initiating singleton
     with pytest.raises(RuntimeError):
         NeuronGlobals()
+
+
+def test_neuron_global_params():
+    """Unit test for NeuronGlobalParams."""
+    params = NeuronGlobals.get_instance().export_params()
+    assert params.temperature == 34.0
+    assert params.v_init == -65.0
+
+    altered_params = NeuronGlobalParams(temperature=25.0, v_init=-95.0)
+    NeuronGlobals.get_instance().load_params(altered_params)
+    assert neuron.h.celsius == 25.0
+    assert neuron.h.v_init == -95.0
+
+    # set back to default
+    NeuronGlobals.get_instance().load_params(params)

tests/test_simulation/test_parallel.py

@@ -0,0 +1,8 @@
+from bluecellulab.simulation.parallel import IsolatedProcess
+
+
+def test_isolated_process():
+    """Test to ensure isolated process keeps its properties."""
+    runner = IsolatedProcess()
+    assert runner._processes == 1
+    assert runner._maxtasksperchild == 1

tests/test_utils.py

@@ -1,7 +1,7 @@
 import json
 
 import numpy as np
-from bluecellulab.utils import CaptureOutput, IsolatedProcess, NumpyEncoder, run_once
+from bluecellulab.utils import CaptureOutput, NumpyEncoder, run_once
 
 
 # Decorated function for testing
@@ -56,10 +56,3 @@ def test_numpy_encoder():
         json.dumps(np.array([True, False, True]), cls=NumpyEncoder)
         == "[true, false, true]"
     )
-
-
-def test_isolated_process():
-    """Test to ensure isolated process keeps its properties."""
-    runner = IsolatedProcess()
-    assert runner._processes == 1
-    assert runner._maxtasksperchild == 1