Add support for compartment reports (soma + named section) as well as spike reports

bluecellulab/cell/core.py

@@ -525,6 +525,10 @@
         nc.record(spike_vec)
         self.recordings[f"spike_detector_{location}_{threshold}"] = spike_vec
 
+    def is_recording_spikes(self, location: str, threshold: float) -> bool:
+        key = f"spike_detector_{location}_{threshold}"
+        return key in self.recordings
+
     def get_recorded_spikes(self, location: str, threshold: float = -30) -> list[float]:
         """Get recorded spikes in the current cell.
 
@@ -756,6 +760,18 @@
         """Get a vector of AIS voltage."""
         return self.get_recording('self.axonal[1](0.5)._ref_v')
 
+    def add_variable_recording(self, variable: str, section, segx):
+        if variable == "v":
+            self.add_voltage_recording(section, segx)
+        else:
+            raise ValueError(f"Unsupported variable for recording: {variable}")
+
+    def get_variable_recording(self, variable: str, section, segx) -> np.ndarray:
+        if variable == "v":
+            return self.get_voltage_recording(section=section, segx=segx)
+        else:
+            raise ValueError(f"Unsupported variable '{variable}'")
+
     @property
     def n_segments(self) -> int:
         """Get the number of segments in the cell."""

bluecellulab/circuit/config/sonata_simulation_config.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 from __future__ import annotations
 from functools import lru_cache
+import json
+import logging
 from pathlib import Path
 from typing import Optional
 
@@ -21,6 +23,8 @@
 
 from bluepysnap import Simulation as SnapSimulation
 
+logger = logging.getLogger(__name__)
+
 
 class SonataSimulationConfig:
     """Sonata implementation of SimulationConfig protocol."""
@@ -74,9 +78,42 @@
             result.append(ConnectionOverrides.from_sonata(conn_entry))
         return result
 
+    @lru_cache(maxsize=1)
+    def get_compartment_sets(self) -> dict[str, dict]:
+        filepath = self.impl.config.get("compartment_sets_file")
+        if not filepath:
+            raise ValueError("No 'compartment_sets_file' entry found in SONATA config.")
+        with open(filepath, 'r') as f:
+            return json.load(f)
+
+    @lru_cache(maxsize=1)
+    def get_node_sets(self) -> dict[str, dict]:
+        filepath = self.impl.circuit.config.get("node_sets_file")
+        if not filepath:
+            raise ValueError("No 'node_sets_file' entry found in SONATA config.")
+        with open(filepath, 'r') as f:
+            return json.load(f)
+
+    @lru_cache(maxsize=1)
+    def get_report_entries(self) -> dict[str, dict]:
+        """Returns the 'reports' dictionary from the SONATA simulation config.
+
+        Each key is a report name, and the value is its configuration.
+        """
+        reports = self.impl.config.get("reports", {})
+        if not isinstance(reports, dict):
+            raise ValueError("Invalid format for 'reports' in SONATA config.")
+        return reports
+
     def connection_entries(self) -> list[ConnectionOverrides]:
         return self._connection_entries() + self._connection_overrides
 
+    def report_file_path(self, report_cfg: dict, report_key: str) -> Path:
+        """Resolve the full path for the report output file."""
+        output_dir = Path(self.output_root_path)
+        file_name = report_cfg.get("file_name", f"{report_key}.h5")
+        return output_dir / file_name
+
     @property
     def base_seed(self) -> int:
         return self.impl.run.random_seed
@@ -135,7 +172,13 @@
 
     @property
     def output_root_path(self) -> str:
-        return self.impl.config["output"]["output_dir"]
+        return self.impl.config.get("output", {}).get("output_dir", "output")
+
+    @property
+    def spikes_file_path(self) -> Path:
+        output_dir = Path(self.output_root_path)
+        spikes_file = self.impl.config.get("output", {}).get("spikes_file", "spikes.h5")
+        return output_dir / spikes_file
 
     @property
     def extracellular_calcium(self) -> Optional[float]:

bluecellulab/circuit/iotools.py

@@ -17,7 +17,6 @@
 from pathlib import Path
 import logging
 
-import bluepy
 import numpy as np
 
 from bluecellulab.circuit.node_id import CellId
@@ -28,6 +27,7 @@
 def parse_outdat(path: str | Path) -> dict[CellId, np.ndarray]:
     """Parse the replay spiketrains in a out.dat formatted file pointed to by
     path."""
+    import bluepy
     spikes = bluepy.impl.spike_report.SpikeReport.load(path).get()
     # convert Series to DataFrame with 2 columns for `groupby` operation
     spike_df = spikes.to_frame().reset_index()

bluecellulab/circuit_simulation.py

@@ -17,10 +17,12 @@
 
 from __future__ import annotations
 from collections.abc import Iterable
+import os
 from pathlib import Path
 from typing import Optional
 import logging
 
+from collections import defaultdict
 import neuron
 import numpy as np
 import pandas as pd
@@ -45,6 +47,7 @@
 from bluecellulab.importer import load_mod_files
 from bluecellulab.rngsettings import RNGSettings
 from bluecellulab.simulation.neuron_globals import NeuronGlobals
+from bluecellulab.simulation.report import configure_all_reports, write_compartment_report, write_sonata_spikes
 from bluecellulab.stimulus.circuit_stimulus_definitions import Noise, OrnsteinUhlenbeck, RelativeOrnsteinUhlenbeck, RelativeShotNoise, ShotNoise
 import bluecellulab.stimulus.circuit_stimulus_definitions as circuit_stimulus_definitions
 from bluecellulab.exceptions import BluecellulabError
@@ -301,6 +304,16 @@
                 add_linear_stimuli=add_linear_stimuli
             )
 
+        configure_all_reports(
+            cells=self.cells,
+            simulation_config=self.circuit_access.config
+        )
+
+        # add spike recordings
+        for cell in self.cells.values():
+            if not cell.is_recording_spikes("soma", threshold=self.spike_threshold):
+                cell.start_recording_spikes(None, location="soma", threshold=self.spike_threshold)
+
     def _add_stimuli(self, add_noise_stimuli=False,
                      add_hyperpolarizing_stimuli=False,
                      add_relativelinear_stimuli=False,
@@ -458,13 +471,26 @@
     @staticmethod
     def merge_pre_spike_trains(*train_dicts) -> dict[CellId, np.ndarray]:
         """Merge presynaptic spike train dicts."""
-        filtered_dicts = [d for d in train_dicts if d not in [None, {}, [], ()]]
+        filtered_dicts = [d for d in train_dicts if isinstance(d, dict) and d]
+
+        if not filtered_dicts:
+            logger.warning("merge_pre_spike_trains: No presynaptic spike trains found.")
+            return {}
 
         all_keys = set().union(*[d.keys() for d in filtered_dicts])
-        return {
-            k: np.sort(np.concatenate([d[k] for d in filtered_dicts if k in d]))
-            for k in all_keys
-        }
+        result = {}
+
+        for k in all_keys:
+            valid_arrays = []
+            for d in filtered_dicts:
+                if k in d:
+                    val = d[k]
+                    if isinstance(val, (np.ndarray, list)) and len(val) > 0:
+                        valid_arrays.append(np.asarray(val))
+            if valid_arrays:
+                result[k] = np.sort(np.concatenate(valid_arrays))
+
+        return result
 
     def _add_connections(
             self,
@@ -646,6 +672,8 @@
             forward_skip_value=forward_skip_value,
             show_progress=show_progress)
 
+        self.write_reports()
+
     def get_mainsim_voltage_trace(
             self, cell_id: int | tuple[str, int], t_start=None, t_stop=None, t_step=None
     ) -> np.ndarray:
@@ -779,3 +807,76 @@
                                  record_dt=cell_kwargs['record_dt'],
                                  template_format=cell_kwargs['template_format'],
                                  emodel_properties=cell_kwargs['emodel_properties'])
+
+    def write_reports(self):
+        """Write all reports defined in the simulation config."""
+        report_entries = self.circuit_access.config.get_report_entries()
+
+        for report_name, report_cfg in report_entries.items():
+            report_type = report_cfg.get("type", "compartment")
+            section = report_cfg.get("sections")
+
+            if report_type != "compartment":
+                raise NotImplementedError(f"Report type '{report_type}' is not supported.")
+
+            output_path = self.circuit_access.config.report_file_path(report_cfg, report_name)
+            if section == "compartment_set":
+                if report_cfg.get("cells") is not None:
+                    raise ValueError(
+                        "Report config error: 'cells' must not be set when using 'compartment_set' sections."
+                    )
+                compartment_sets = self.circuit_access.config.get_compartment_sets()
+                write_compartment_report(
+                    report_name=report_name,
+                    output_path=output_path,
+                    cells=self.cells,
+                    report_cfg=report_cfg,
+                    source_sets=compartment_sets,
+                    source_type="compartment_set"
+                )
+
+            else:
+                node_sets = self.circuit_access.config.get_node_sets()
+                if report_cfg.get("compartments") not in ("center", "all"):
+                    raise ValueError(
+                        f"Unsupported 'compartments' value '{report_cfg.get('compartments')}' "
+                        "for node-based section recording (must be 'center' or 'all')."
+                    )
+                write_compartment_report(
+                    report_name=report_name,
+                    output_path=output_path,
+                    cells=self.cells,
+                    report_cfg=report_cfg,
+                    source_sets=node_sets,
+                    source_type="node_set"
+                )
+
+        self.write_spike_report()
+
+    def write_spike_report(self):
+        """Collect and write in-memory recorded spike times to a SONATA HDF5
+        file, grouped by population as required by the SONATA specification."""
+        output_path = self.circuit_access.config.spikes_file_path
+
+        if os.path.exists(output_path):
+            os.remove(output_path)
+
+        # Group spikes per population
+        spikes_by_population = defaultdict(dict)
+        for gid, cell in self.cells.items():
+            pop = getattr(gid, 'population_name', None)
+            if pop is None:
+                continue
+            try:
+                cell_spikes = cell.get_recorded_spikes(location="soma", threshold=self.spike_threshold)
+                if cell_spikes is not None:
+                    spikes_by_population[pop][gid.id] = list(cell_spikes)
+            except AttributeError:
+                continue
+
+        # Ensure we at least create empty groups for all known populations
+        all_populations = set(getattr(gid, 'population_name', None) for gid in self.cells.keys())
+
+        for pop in all_populations:
+            spikes = spikes_by_population.get(pop, {})  # May be empty
+            write_sonata_spikes(output_path, spikes, pop)