[BBPBGLIB-854] Implement a CLI option to run coreneuron in direct mod…

…e without writing model data to disk (#100)

## Context
This PR implements a new CLI option "--coreneuron-direct-mode" that
enables to run the coreneuron simulation using direct memory transfer
from neuron, without writing intermediate model data to disk.

## Scope

- When "--coreneuron-direct-mode", we set the Coreneuron
parameter`coreneuron.file_mode=False`.
- And Coreneuron writes the element reports but not the spike report. So
we still create `sim.conf`, `report.conf` and an empty folder
`coreneuron_input`. The spike report is created explicitly in neurodamus
via `sonata_spikes()` just like the Neuron runs.
- Enable fast membrane current calculation `Nd.cvode.use_fast_imem(1)`
for `i_membrane` and `lfp` reports.

## Testing
New test file `test_coreneuron_directmode.py`.
The blueconfig pipeline tests with "--coreneuron-direct-mode" enabled
have been done manually with the blueconfig branch (https://bbpgitlab.epfl.ch/hpc/sim/blueconfigs/-/merge_requests/114).


## Review
* [x] PR description is complete
* [x] Coding style (imports, function length, New functions, classes or
files) are good
* [x] Unit/Scientific test added
* [ ] Updated Readme, in-code, developer documentation

neurodamus/commands.py

@@ -61,6 +61,8 @@ def neurodamus(args=None):
         --model-stats           Show model stats in CoreNEURON simulations [default: False]
         --dry-run               Dry-run simulation to estimate memory usage [default: False]
         --num-target-ranks=<number>  Number of ranks to target for dry-run load balancing
+        --coreneuron-direct-mode     Run CoreNeuron in direct memory mode transfered from Neuron,
+                                     without writing model data to disk.
     """
     options = docopt_sanitize(docopt(neurodamus.__doc__, args))
     config_file = options.pop("ConfigFile")

neurodamus/core/configuration.py

@@ -76,6 +76,7 @@ class CliOptions(ConfigT):
     simulator = None
     dry_run = False
     num_target_ranks = None
+    coreneuron_direct_mode = False
 
     # Restricted Functionality support, mostly for testing
 
@@ -235,6 +236,7 @@ class _SimConfig(object):
     spike_threshold = -30
     dry_run = False
     num_target_ranks = None
+    coreneuron_direct_mode = False
 
     _validators = []
     _requisitors = []
@@ -1005,6 +1007,27 @@ def _spikes_sort_order(config: _SimConfig, run_conf):
                                  "BBP supports 'none' and 'by_time'")
 
 
+@SimConfig.validator
+def _coreneuron_direct_mode(config: _SimConfig, run_conf):
+    user_config = config.cli_options
+    direct_mode = user_config.coreneuron_direct_mode
+    if direct_mode:
+        if config.use_neuron:
+            raise ConfigurationError("--coreneuron-direct-mode is not valid for NEURON")
+        if config.modelbuilding_steps > 1:
+            logging.warning("--coreneuron-direct-mode not valid for multi-cyle model building, "
+                            "continue with file mode")
+            direct_mode = False
+        if config.save or config.restore:
+            logging.warning("--coreneuron-direct-mode not valid for save/restore, "
+                            "continue with file mode")
+            direct_mode = False
+
+    if direct_mode:
+        logging.info("Run CORENEURON direct mode without writing model data to disk")
+    config.coreneuron_direct_mode = direct_mode
+
+
 def get_debug_cell_gid(cli_options):
     gid = cli_options.get("dump_cell_state") if cli_options else None
     if gid:

neurodamus/core/coreneuron_configuration.py

@@ -171,13 +171,13 @@ def write_spike_filename(self, filename):
             fp.write(filename)
             fp.write("\n")
 
-    def psolve_core(self, save_path=None, restore_path=None):
+    def psolve_core(self, save_path=None, restore_path=None, coreneuron_direct_mode=False):
         from neuron import coreneuron
         from . import NeurodamusCore as Nd
 
         Nd.cvode.cache_efficient(1)
         coreneuron.enable = True
-        coreneuron.file_mode = True
+        coreneuron.file_mode = not coreneuron_direct_mode
         coreneuron.sim_config = f"{self.output_root}/{self.sim_config_file}"
         if save_path:
             coreneuron.save_path = save_path

neurodamus/node.py

@@ -283,7 +283,7 @@ def __init__(self, config_file, options=None):
             self._run_conf = SimConfig.run_conf
             self._target_manager = TargetManager(self._run_conf)
             self._target_spec = TargetSpec(self._run_conf.get("CircuitTarget"))
-            if SimConfig.use_neuron:
+            if SimConfig.use_neuron or SimConfig.coreneuron_direct_mode:
                 self._sonatareport_helper = Nd.SonataReportHelper(Nd.dt, True)
             self._base_circuit: CircuitConfig = SimConfig.base_circuit
             self._extra_circuits = SimConfig.extra_circuits
@@ -858,6 +858,11 @@ def enable_reports(self):
             has_gids = len(self._circuits.global_manager.get_final_gids()) > 0
             report = Report(*rep_params, SimConfig.use_coreneuron) if has_gids else None
 
+            # With coreneuron direct mode, enable fast membrane current calculation for i_membrane
+            if SimConfig.coreneuron_direct_mode and \
+                    "i_membrane" in rep_params.report_on or rep_params.rep_type == "lfp":
+                Nd.cvode.use_fast_imem(1)
+
             if not SimConfig.use_coreneuron or rep_params.rep_type == "Synapse":
                 try:
                     self._report_setup(report, rep_conf, target, rep_params.rep_type)
@@ -1094,7 +1099,7 @@ def sim_init(self, corenrn_gen=None, **sim_opts):
         if corenrn_gen:
             self._sim_corenrn_write_config()
 
-        if SimConfig.use_neuron:
+        if SimConfig.use_neuron or SimConfig.coreneuron_direct_mode:
             self._sim_init_neuron()
 
         if ospath.isfile("debug_gids.txt"):
@@ -1212,9 +1217,11 @@ def _coreneuron_ensure_all_ranks_have_gids(self, corenrn_data):
                 dummyfile.write("%s\n0\n" % coredata_version)
 
     # -
-    def _sim_corenrn_configure_datadir(self, corenrn_restore):
+    def _sim_corenrn_configure_datadir(self, corenrn_restore, coreneuron_direct_mode):
         corenrn_datadir = SimConfig.coreneuron_datadir
         os.makedirs(corenrn_datadir, exist_ok=True)
+        if coreneuron_direct_mode:
+            return corenrn_datadir
         corenrn_datadir_shm = SHMUtil.get_datadir_shm(corenrn_datadir)
 
         # Clean-up any previous simulations in the same output directory
@@ -1273,14 +1280,16 @@ def _sim_corenrn_configure_datadir(self, corenrn_restore):
     @timeit(name="corewrite")
     def _sim_corenrn_write_config(self, corenrn_restore=False):
         log_stage("Dataset generation for CoreNEURON")
-        CoreConfig.datadir = self._sim_corenrn_configure_datadir(corenrn_restore)
+        CoreConfig.datadir = self._sim_corenrn_configure_datadir(corenrn_restore,
+                                                                 SimConfig.coreneuron_direct_mode)
         fwd_skip = self._run_conf.get("ForwardSkip", 0) if not corenrn_restore else 0
 
         if not corenrn_restore:
             CompartmentMapping(self._circuits.global_manager).register_mapping()
-            with self._coreneuron_ensure_all_ranks_have_gids(CoreConfig.datadir):
-                self._pc.nrnbbcore_write(CoreConfig.datadir)
-                MPI.barrier()  # wait for all ranks to finish corenrn data generation
+            if not SimConfig.coreneuron_direct_mode:
+                with self._coreneuron_ensure_all_ranks_have_gids(CoreConfig.datadir):
+                    self._pc.nrnbbcore_write(CoreConfig.datadir)
+                    MPI.barrier()  # wait for all ranks to finish corenrn data generation
 
         CoreConfig.write_sim_config(
             Nd.tstop,
@@ -1312,8 +1321,11 @@ def run_all(self):
             self.sonata_spikes()
         if SimConfig.use_coreneuron:
             print_mem_usage()
-            self.clear_model(avoid_clearing_queues=False)
+            if not SimConfig.coreneuron_direct_mode:
+                self.clear_model(avoid_clearing_queues=False)
             self._run_coreneuron()
+            if SimConfig.coreneuron_direct_mode:
+                self.sonata_spikes()
         return timings
 
     # -
@@ -1329,7 +1341,8 @@ def _run_coreneuron(self):
         logging.info("Launching simulation with CoreNEURON")
         CoreConfig.psolve_core(
             getattr(SimConfig, "save", None),
-            getattr(SimConfig, "restore", None)
+            getattr(SimConfig, "restore", None),
+            SimConfig.coreneuron_direct_mode
         )
 
     #

tests/integration-e2e/test_coreneuron_directmode.py

@@ -0,0 +1,37 @@
+import os
+import numpy.testing as npt
+import numpy as np
+
+
+def test_coreneuron_no_write_model(USECASE3):
+    from libsonata import SpikeReader, ElementReportReader
+    from neurodamus import Neurodamus
+    from neurodamus.core.configuration import SimConfig
+    nd = Neurodamus(
+        str(USECASE3 / "simulation_sonata_coreneuron.json"),
+        keep_build=True,
+        coreneuron_direct_mode=True
+    )
+    nd.run()
+    coreneuron_data = SimConfig.coreneuron_datadir
+    assert not next(os.scandir(coreneuron_data), None), f"{coreneuron_data} should be empty."
+
+    spikes_path = os.path.join(SimConfig.output_root, nd._run_conf.get("SpikesFile"))
+    spikes_reader = SpikeReader(spikes_path)
+    pop_A = spikes_reader["NodeA"]
+    pop_B = spikes_reader["NodeB"]
+    spike_dict = pop_A.get_dict()
+    npt.assert_allclose(spike_dict["timestamps"][:10], np.array([0.2, 0.3, 0.3, 2.5, 3.4,
+                                                                4.2, 5.5, 7., 7.4, 8.6]))
+    npt.assert_allclose(spike_dict["node_ids"][:10], np.array([0, 1, 2, 0, 1, 2, 0, 0, 1, 2]))
+    assert not pop_B.get()
+
+    soma_reader = ElementReportReader(SimConfig.reports.get("soma_report").get('FileName'))
+    soma_A = soma_reader["NodeA"]
+    soma_B = soma_reader["NodeB"]
+    data_A = soma_A.get(tstop=0.5)
+    data_B = soma_B.get(tstop=0.5)
+    npt.assert_allclose(data_A.data, np.array([[-75.], [-39.78627], [-14.380434], [15.3370695],
+                                               [1.7240616], [-13.333434]]))
+    npt.assert_allclose(data_B.data, np.array([[-75.], [-75.00682], [-75.010414], [-75.0118],
+                                               [-75.01173], [-75.010635]]))