[BBPBGLIB-1184] Fix CoreNEURON reports restore functionality

neurodamus/core/coreneuron_configuration.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from ._utils import run_only_rank0
 from . import NeurodamusCore as Nd
+from .configuration import ConfigurationError
 from ..report import get_section_index
 
 
@@ -75,6 +76,7 @@ class _CoreNEURONConfig(object):
     """
     sim_config_file = "sim.conf"
     report_config_file = "report.conf"
+    restore_path = None
     output_root = "output"
     datadir = f"{output_root}/coreneuron_input"
     default_cell_permute = 0
@@ -85,6 +87,47 @@ class _CoreNEURONConfig(object):
     def instantiate_artificial_cell(self):
         self.artificial_cell_object = Nd.CoreNEURONArtificialCell()
 
+    @run_only_rank0
+    def update_tstop(self, report_name, nodeset_name, tstop):
+        # Try current directory first
+        report_conf = Path(self.output_root) / self.report_config_file
+        if not report_conf.exists():
+            # Try one level up from restore_path
+            parent_report_conf = Path(self.restore_path) / ".." / self.report_config_file
+            if parent_report_conf.exists():
+                # Copy the file to current location
+                report_conf.write_bytes(parent_report_conf.read_bytes())
+            else:
+                raise ConfigurationError(f"Report config file not found in {report_conf} "
+                                         f"or {parent_report_conf}")
+
+        # Read all content
+        with report_conf.open('rb') as f:
+            lines = f.readlines()
+
+        # Find and update the matching line
+        found = False
+        for i, line in enumerate(lines):
+            try:
+                parts = line.decode().split()
+                # Report name and target name must match in order to update the tstop
+                if parts[0:2] == [report_name, nodeset_name]:
+                    parts[9] = f"{tstop:.6f}"
+                    lines[i] = (' '.join(parts) + '\n').encode()
+                    found = True
+                    break
+            except (UnicodeDecodeError, IndexError):
+                # Ignore lines that cannot be decoded (binary data)
+                continue
+
+        if not found:
+            raise ConfigurationError(f"Report '{report_name}' with target '{nodeset_name}' "
+                                     "not matching any report in the 'save' execution")
+
+        # Write back
+        with report_conf.open('wb') as f:
+            f.writelines(lines)
+
     @run_only_rank0
     def write_report_config(
             self, report_name, target_name, report_type, report_variable,

neurodamus/node.py

@@ -279,6 +279,8 @@ def __init__(self, config_file, options=None):
                 # Instantiate the CoreNEURON artificial cell object which is used to fill up
                 # the empty ranks. This need to be done before the circuit is finitialized
                 CoreConfig.instantiate_artificial_cell()
+                if SimConfig.restore_coreneuron:
+                    CoreConfig.restore_path = SimConfig.restore
             self._run_conf = SimConfig.run_conf
             self._target_manager = TargetManager(self._run_conf)
             self._target_spec = TargetSpec(self._run_conf.get("CircuitTarget"))
@@ -825,7 +827,7 @@ def enable_reports(self):
             pop_offsets_alias = self._circuits.get_population_offsets()
         else:
             pop_offsets_alias = CircuitManager.read_population_offsets()
-        if SimConfig.use_coreneuron:
+        if SimConfig.use_coreneuron and not SimConfig.restore_coreneuron:
             CoreConfig.write_report_count(len(reports_conf))
 
         for rep_name, rep_conf in reports_conf.items():
@@ -839,7 +841,7 @@ def enable_reports(self):
                 continue
 
             if SimConfig.use_coreneuron and MPI.rank == 0:
-                if not self._report_write_coreneuron_config(rep_name, rep_conf, target, rep_params):
+                if not self._report_write_coreneuron_config(rep_conf, target, rep_params):
                     n_errors += 1
                     continue
 
@@ -940,9 +942,15 @@ def _report_build_params(self, rep_name, rep_conf, target, pop_offsets_alias_pop
         )
 
     #
-    def _report_write_coreneuron_config(self, rep_name, rep_conf, target, rep_params):
+    def _report_write_coreneuron_config(self, rep_conf, target, rep_params):
         target_spec = TargetSpec(rep_conf["Target"])
 
+        # For restore case with no change in reporting, we can directly update the end time.
+        # Note: If different reports are needed during restore, this workflow needs to be adapted.
+        if SimConfig.restore_coreneuron:
+            CoreConfig.update_tstop(rep_params.name, target_spec.name, rep_params.end)
+            return True
+
         # for sonata config, compute target_type from user inputs
         if "Sections" in rep_conf and "Compartments" in rep_conf:
             def _compute_corenrn_target_type(section_type, compartment_type):
@@ -963,8 +971,7 @@ def _compute_corenrn_target_type(section_type, compartment_type):
 
         reporton_comma_separated = ",".join(rep_params.report_on.split())
         core_report_params = (
-            (os.path.basename(rep_conf.get("FileName", rep_name)),
-                target_spec.name, rep_params.rep_type, reporton_comma_separated)
+            (rep_params.name, target_spec.name, rep_params.rep_type, reporton_comma_separated)
             + rep_params[3:5] + (target_type,) + rep_params[5:8]
             + (target.get_gids(), SimConfig.corenrn_buff_size)
         )
@@ -1008,8 +1015,10 @@ def _report_setup(self, report, rep_conf, target, rep_type):
                 report.add_synapse_report(cell, point, spgid, pop_name, pop_offset)
 
     def _reports_init(self, pop_offsets_alias):
-        pop_offsets = pop_offsets_alias[0]
+        if SimConfig.restore_coreneuron:
+            return
 
+        pop_offsets = pop_offsets_alias[0]
         if SimConfig.use_coreneuron:
             # write spike populations
             if hasattr(CoreConfig, "write_population_count"):