Unify report writing across single-rank and MPI runs using recording sites

bluecellulab/cell/core.py

@@ -19,7 +19,7 @@
 
 from pathlib import Path
 import queue
-from typing import List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 from typing_extensions import deprecated
 
 import neuron
@@ -46,7 +46,7 @@
 from bluecellulab.stimulus.circuit_stimulus_definitions import SynapseReplay
 from bluecellulab.synapse import SynapseFactory, Synapse
 from bluecellulab.synapse.synapse_types import SynapseID
-from bluecellulab.type_aliases import HocObjectType, NeuronSection, SectionMapping
+from bluecellulab.type_aliases import HocObjectType, NeuronSection, ReportSite, SectionMapping
 from bluecellulab.cell.section_tools import currents_vars, section_to_variable_recording_str
 
 logger = logging.getLogger(__name__)
@@ -129,6 +129,7 @@ def __init__(self,
         neuron.h.finitialize()
 
         self.recordings: dict[str, HocObjectType] = {}
+        self.report_sites: dict[str, list[dict]] = {}
         self.synapses: dict[SynapseID, Synapse] = {}
         self.connections: dict[SynapseID, bluecellulab.Connection] = {}
 
@@ -1012,47 +1013,66 @@ def resolve_segments_from_config(self, report_cfg) -> List[Tuple[NeuronSection,
                     targets.append((sec, sec_name, seg.x))
         return targets
 
-    def configure_recording(self, recording_sites, variable_name, report_name):
-        """Configure recording of a variable on a single cell.
-
-        This function sets up the recording of the specified variable (e.g., membrane voltage)
-        in the target cell, for each resolved segment.
+    def configure_recording(self,
+                            recording_sites: Iterable[tuple[NeuronSection | None, str, float]],
+                            variable_name: str,
+                            report_name: str
+                            ) -> list[tuple[ReportSite, str]]:
+        """Attach NEURON recordings for a variable at the given sites and
+        return the recording names created.
 
         Parameters
         ----------
-        cell : Any
-            The cell object on which to configure recordings.
-
-        recording_sites : list of tuples
-            List of tuples (section, section_name, segment) where:
-            - section is the section object in the cell.
-            - section_name is the name of the section.
-            - segment is the Neuron segment index (0-1).
-
+        recording_sites : iterable
+            (section, section_name, segx) tuples describing recording locations.
         variable_name : str
-            The name of the variable to record (e.g., "v" for membrane voltage).
-
+            Variable to record (e.g. "v", "ina", "kca.gkca").
         report_name : str
-            The name of the report (used in logging).
+            Report identifier (for logging).
+
+        Returns
+        -------
+        list[tuple[ReportSite, str]]
+            (site, rec_name) pairs for successfully configured recordings.
         """
         node_id = self.cell_id.id
+        configured: list[tuple[ReportSite, str]] = []
+
+        for site in recording_sites:
+            sec, sec_name, seg = site
+            report_site = ReportSite(sec, sec_name, float(seg))
 
-        for sec, sec_name, seg in recording_sites:
             try:
-                self.add_variable_recording(variable=variable_name, section=sec, segx=seg)
+                section_obj = self.soma if sec is None else sec
+                rec_name = section_to_variable_recording_str(section_obj, float(seg), variable_name)
+
+                if rec_name not in self.recordings:
+                    self.add_variable_recording(
+                        variable=variable_name,
+                        section=None if sec is None else sec,
+                        segx=float(seg),
+                    )
+
+                configured.append((report_site, rec_name))
+
                 logger.info(
                     f"Recording '{variable_name}' at {sec_name}({seg}) on GID {node_id} for report '{report_name}'"
                 )
+
             except AttributeError:
                 logger.warning(
-                    f"Recording for variable '{variable_name}' is not implemented in Cell."
+                    "Recording for variable '%s' is not implemented at %s(%s) on GID %s for report '%s'",
+                    variable_name, sec_name, seg, node_id, report_name,
                 )
-                return
+
             except Exception as e:
                 logger.warning(
-                    f"Failed to record '{variable_name}' at {sec_name}({seg}) on GID {node_id} for report '{report_name}': {e}"
+                    f"Failed to record '{variable_name}' at {sec_name}({seg}) on GID {node_id} "
+                    f"for report '{report_name}': {e}"
                 )
 
+        return configured
+
     def add_currents_recordings(
         self,
         section,

bluecellulab/circuit/circuit_access/sonata_circuit_access.py

@@ -29,7 +29,7 @@
 from bluecellulab.circuit import CellId, SynapseProperty
 from bluecellulab.circuit.config import SimulationConfig
 from bluecellulab.circuit.synapse_properties import SynapseProperties
-from bluecellulab.circuit.config import SimulationConfig, SonataSimulationConfig
+from bluecellulab.circuit.config import SonataSimulationConfig
 from bluecellulab.circuit.synapse_properties import (
     properties_from_snap,
     properties_to_snap,
@@ -301,7 +301,7 @@ def morph_filepath(self, cell_id: CellId) -> str:
         node_population = self._circuit.nodes[cell_id.population_name]
         try:  # if asc defined in alternate morphology
             return str(node_population.morph.get_filepath(cell_id.id, extension="asc"))
-        except BluepySnapError as e:
+        except BluepySnapError:
             logger.debug(f"No asc morphology found for {cell_id}, trying swc.")
             return str(node_population.morph.get_filepath(cell_id.id))
 

bluecellulab/circuit_simulation.py

@@ -22,13 +22,12 @@
 import logging
 import warnings
 
-from bluecellulab.reports.utils import configure_all_reports
+from bluecellulab.reports.utils import prepare_recordings_for_reports
 import neuron
 import numpy as np
 import pandas as pd
 from pydantic.types import NonNegativeInt
 from typing_extensions import deprecated
-from typing import Optional
 
 import bluecellulab
 from bluecellulab.cell import CellDict
@@ -334,7 +333,7 @@ def instantiate_gids(
                 add_linear_stimuli=add_linear_stimuli
             )
 
-        configure_all_reports(
+        self.recording_index, self.sites_index = prepare_recordings_for_reports(
             cells=self.cells,
             simulation_config=self.circuit_access.config
         )

bluecellulab/reports/manager.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Dict
+from typing import Any, Optional, Dict
+
+from bluecellulab.circuit.node_id import CellId
 from bluecellulab.reports.writers import get_writer
 from bluecellulab.reports.utils import SUPPORTED_REPORT_TYPES, extract_spikes_from_cells  # helper you already have / write
 
@@ -30,31 +32,36 @@ def __init__(self, config, sim_dt: float):
 
     def write_all(
         self,
-        cells_or_traces: Dict,
-        spikes_by_pop: Optional[Dict[str, Dict[int, list]]] = None,
+        cells: Dict[CellId, Any],
+        spikes_by_pop: Optional[Dict[str, Dict[int, list[float]]]] = None,
     ):
-        """Write all configured reports (compartment and spike) in SONATA
-        format.
+        """Write all configured SONATA reports (compartment and spike).
+
+        `cells` maps CellId to live Cell objects or recording proxies.
+        For compartment reports each entry must provide:
+            - ``report_sites``: ``{report_name: [site_dict, ...]}``
+            - ``get_recording(rec_name)`` → recorded trace
+
+        If ``spikes_by_pop`` is not provided, spike times are obtained from the
+        cells via ``get_recorded_spikes(location=..., threshold=...)``.
 
         Parameters
         ----------
-        cells_or_traces : dict
-            A dictionary mapping (population, gid) to either:
-            - Cell objects with recorded data (used in single-process simulations), or
-            - Precomputed trace dictionaries, e.g., {"voltage": ndarray}, typically gathered across ranks in parallel runs.
-
-        spikes_by_pop : dict, optional
-            A precomputed dictionary of spike times by population.
-            If not provided, spike times are extracted from `cells_or_traces`.
-
-        Notes
-        -----
-        In parallel simulations, you must gather all traces and spikes to rank 0 and pass them here.
-        """
-        self._write_voltage_reports(cells_or_traces)
-        self._write_spike_report(spikes_by_pop or extract_spikes_from_cells(cells_or_traces, location=self.cfg.spike_location, threshold=self.cfg.spike_threshold))
+        cells : Dict[CellId, Any]
+            Cell objects or proxies exposing recordings and report topology.
 
-    def _write_voltage_reports(self, cells_or_traces):
+        spikes_by_pop : dict[str, dict[int, list[float]]], optional
+            Precomputed spikes ``{population: {gid: [times...]}}``. If omitted,
+            spikes are extracted from the cells.
+        """
+        self._write_compartment_reports(cells)
+        self._write_spike_report(
+            spikes_by_pop or extract_spikes_from_cells(
+                cells, location=self.cfg.spike_location, threshold=self.cfg.spike_threshold
+            )
+        )
+
+    def _write_compartment_reports(self, cells):
         for name, rcfg in self.cfg.get_report_entries().items():
             if rcfg.get("type") not in SUPPORTED_REPORT_TYPES:
                 continue
@@ -83,9 +90,9 @@ def _write_voltage_reports(self, cells_or_traces):
 
             out_path = self.cfg.report_file_path(rcfg, name)
             writer = get_writer("compartment")(rcfg, out_path, self.dt)
-            writer.write(cells_or_traces, self.cfg.tstart, self.cfg.tstop)
+            writer.write(cells, self.cfg.tstart, self.cfg.tstop)
 
-    def _write_spike_report(self, spikes_by_pop):
+    def _write_spike_report(self, spikes_by_pop: Dict[str, Dict[int, list[float]]]):
         out_path = self.cfg.spikes_file_path
         writer = get_writer("spikes")({}, out_path, self.dt)
         writer.write(spikes_by_pop)