having consistent sampling options across ssim methods #44 (#52)

* consistent sampling options across ssim methods

* lint fix

* lint fix

* Updated type hint in _sample_array function to ndarray from Sequence

* lint fix

* bug fix

* bug fix

---------

Co-authored-by: Kiliç Ilkan Fabrice <[email protected]>

Author: ilkilic

bluecellulab/circuit/simulation_access.py

@@ -15,17 +15,11 @@
 
 from __future__ import annotations
 from pathlib import Path
-from platform import python_version_tuple
 from typing import Optional, Protocol
 
 from bluecellulab.circuit.config import SimulationConfig, SonataSimulationConfig
 from bluecellulab.exceptions import ExtraDependencyMissingError
 
-if python_version_tuple() < ('3', '9'):
-    from typing import Sequence
-else:
-    from collections.abc import Sequence
-
 from bluecellulab import BLUEPY_AVAILABLE
 if BLUEPY_AVAILABLE:
     import bluepy
@@ -38,23 +32,22 @@
 from bluecellulab.circuit.config import BluepySimulationConfig
 
 
-def _sample_array(arr: Sequence, t_step: float, sim_t_step: float) -> Sequence:
+def _sample_array(arr: np.ndarray, ratio: float) -> np.ndarray:
     """Sample an array at a given time step.
 
     Args:
         arr: Array to sample.
-        t_step: User specified time step to sample at.
-        sim_t_step: Time step used in the main simulation.
+        ratio: The ratio between the time step used for sampling and the time step used in the simulation
+                (the time step should be a multiple of the simulation time step)
 
     Returns:
         Array sampled at the given time step.
     """
-    ratio = t_step / sim_t_step
-    if t_step == sim_t_step:
+    if ratio == 1:
         return arr
     elif not np.isclose(ratio, round(ratio)):
-        raise ValueError(
-            f"Time step {t_step} is not a multiple of the simulation time step {sim_t_step}.")
+        raise ValueError("The ratio is not close to a whole number. "
+                         "The time step should be a multiple of the simulation time step.")
     return arr[::round(ratio)]
 
 
@@ -105,15 +98,17 @@ def get_soma_voltage(
             .to_numpy()
         )
         if t_step is not None:
-            arr = _sample_array(arr, t_step, self._config._soma_report_dt)
+            ratio = t_step / self._config._soma_report_dt
+            arr = _sample_array(arr, ratio)
         return arr
 
     def get_soma_time_trace(self, t_step: Optional[float] = None) -> np.ndarray:
         """Retrieve the time trace from the main simulation."""
         report = self.impl.report('soma')
         arr = report.get_gid(report.gids[0]).index.to_numpy()
         if t_step is not None:
-            arr = _sample_array(arr, t_step, self._config._soma_report_dt)
+            ratio = t_step / self._config._soma_report_dt
+            arr = _sample_array(arr, ratio)
         return arr
 
     def get_spikes(self) -> dict[CellId, np.ndarray]:
@@ -140,14 +135,16 @@ def get_soma_voltage(
         report = self.impl.reports["soma"].filter(cell_id.id, t_start, t_end)
         arr = report.report[cell_id.population_name][cell_id.id].values
         if t_step is not None:
-            arr = _sample_array(arr, t_step, self.impl.dt)
+            ratio = t_step / self.impl.dt
+            arr = _sample_array(arr, ratio)
         return arr
 
     def get_soma_time_trace(self, t_step: Optional[float] = None) -> np.ndarray:
         report = self.impl.reports["soma"]
         arr = report.filter().report.index.values
         if t_step is not None:
-            arr = _sample_array(arr, t_step, self.impl.dt)
+            ratio = t_step / self.impl.dt
+            arr = _sample_array(arr, ratio)
         return arr
 
     def get_spikes(self) -> dict[CellId, np.ndarray]:

bluecellulab/ssim.py

@@ -38,7 +38,7 @@
 from bluecellulab.circuit.config import SimulationConfig
 from bluecellulab.circuit.format import determine_circuit_format, CircuitFormat
 from bluecellulab.circuit.node_id import create_cell_id, create_cell_ids
-from bluecellulab.circuit.simulation_access import BluepySimulationAccess, SimulationAccess, SonataSimulationAccess
+from bluecellulab.circuit.simulation_access import BluepySimulationAccess, SimulationAccess, SonataSimulationAccess, _sample_array
 from bluecellulab.stimuli import Pattern
 import bluecellulab.stimuli as stimuli
 from bluecellulab.exceptions import BluecellulabError
@@ -656,9 +656,18 @@ def get_mainsim_voltage_trace(
         cell_id = create_cell_id(cell_id)
         return self.simulation_access.get_soma_voltage(cell_id, t_start, t_stop, t_step)
 
-    def get_mainsim_time_trace(self) -> np.ndarray:
-        """Get the time trace from the main simulation."""
-        return self.simulation_access.get_soma_time_trace()
+    def get_mainsim_time_trace(self, t_step=None) -> np.ndarray:
+        """Get the time trace from the main simulation.
+
+        Parameters
+        -----------
+        t_step: time step (should be a multiple of report time step T;
+        equals T by default)
+
+        Returns:
+            One dimentional np.ndarray to represent the times.
+        """
+        return self.simulation_access.get_soma_time_trace(t_step)
 
     def get_time(self) -> np.ndarray:
         """Get the time vector for the recordings, contains negative times.
@@ -668,17 +677,56 @@ def get_time(self) -> np.ndarray:
         first_key = next(iter(self.cells))
         return self.cells[first_key].get_time()
 
-    def get_time_trace(self) -> np.ndarray:
-        """Get the time vector for the recordings, negative times removed."""
+    def get_time_trace(self, t_step=None) -> np.ndarray:
+        """Get the time vector for the recordings, negative times removed.
+
+        Parameters
+        -----------
+        t_step: time step (should be a multiple of report time step T;
+        equals T by default)
+
+        Returns:
+            One dimentional np.ndarray to represent the times.
+        """
         time = self.get_time()
-        return time[np.where(time >= 0.0)]
+        time = time[np.where(time >= 0.0)]
 
-    def get_voltage_trace(self, cell_id: int | tuple[str, int]) -> np.ndarray:
-        """Get the voltage vector for the cell_id, negative times removed."""
+        if t_step is not None:
+            ratio = t_step / self.dt
+            time = _sample_array(time, ratio)
+        return time
+
+    def get_voltage_trace(
+            self, cell_id: int | tuple[str, int], t_start=None, t_stop=None, t_step=None
+    ) -> np.ndarray:
+        """Get the voltage vector for the cell_id, negative times removed.
+
+        Parameters
+        -----------
+        cell_id: cell id of interest.
+        t_start, t_stop: time range of interest,
+        report time range is used by default.
+        t_step: time step (should be a multiple of report time step T;
+        equals T by default)
+
+        Returns:
+            One dimentional np.ndarray to represent the voltages.
+        """
         cell_id = create_cell_id(cell_id)
         time = self.get_time()
         voltage = self.cells[cell_id].get_soma_voltage()
-        return voltage[np.where(time >= 0.0)]
+
+        if t_start is None or t_start < 0:
+            t_start = 0
+        if t_stop is None:
+            t_stop = np.inf
+
+        voltage = voltage[np.where((time >= t_start) & (time <= t_stop))]
+
+        if t_step is not None:
+            ratio = t_step / self.dt
+            voltage = _sample_array(voltage, ratio)
+        return voltage
 
     def delete(self):
         """Delete ssim and all of its attributes.

tests/test_circuit/test_simulation_access.py

@@ -25,20 +25,20 @@ def test_sample_array():
     arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
     t_step = 0.1
     sim_t_step = 0.1
-    assert np.array_equal(_sample_array(arr, t_step, sim_t_step), arr)
+    assert np.array_equal(_sample_array(arr, t_step / sim_t_step), arr)
     t_step = 0.2
     sim_t_step = 0.1
-    assert np.array_equal(_sample_array(arr, t_step, sim_t_step), [1, 3, 5, 7, 9])
+    assert np.array_equal(_sample_array(arr, t_step / sim_t_step), [1, 3, 5, 7, 9])
     t_step = 0.3
     sim_t_step = 0.1
-    assert np.array_equal(_sample_array(arr, t_step, sim_t_step), [1, 4, 7, 10])
+    assert np.array_equal(_sample_array(arr, t_step / sim_t_step), [1, 4, 7, 10])
     t_step = 0.4
     sim_t_step = 0.1
-    assert np.array_equal(_sample_array(arr, t_step, sim_t_step), [1, 5, 9])
+    assert np.array_equal(_sample_array(arr, t_step / sim_t_step), [1, 5, 9])
     t_step = 0.1
     sim_t_step = 0.2
     with pytest.raises(ValueError):
-        _sample_array(arr, t_step, sim_t_step)
+        _sample_array(arr, t_step / sim_t_step)
 
 
 class TestSonataSimulationAccess:

tests/test_ssim_sonata.py

@@ -1,7 +1,6 @@
 """Testing SSim with SONATA simulations."""
 
 from pathlib import Path
-from bluecellulab.circuit.simulation_access import _sample_array
 
 import numpy as np
 import pytest
@@ -34,11 +33,8 @@ def test_sim_quick_scx_sonata(input_type):
     sim.run(t_stop)
 
     # Get the voltage trace
-    time = sim.get_time_trace()
-    voltage = sim.get_voltage_trace(cell_id)
-    # Sampling
-    voltage = _sample_array(voltage, 1.0, 0.025)
-    time = _sample_array(time, 1.0, 0.025)
+    time = sim.get_time_trace(1)
+    voltage = sim.get_voltage_trace(cell_id, 0, t_stop, 1)
     voltage = voltage[:len(voltage) - 1]  # remove last point, mainsim produces 1 less
     time = time[:len(time) - 1]  # remove last point, mainsim produces 1 less
     mainsim_voltage = sim.get_mainsim_voltage_trace(cell_id)
@@ -69,11 +65,8 @@ def test_sim_quick_scx_sonata_multicircuit(input_type):
     t_stop = 20.0
     sim.run(t_stop)
     for cell_id in cell_ids:
-        voltage = sim.get_voltage_trace(cell_id)
-        # Sampling
-        voltage = _sample_array(voltage, 1.0, 0.025)
-        # remove the last one
-        voltage = voltage[:len(voltage) - 1]
+        voltage = sim.get_voltage_trace(cell_id, 0, t_stop, 1)
+        voltage = voltage[:len(voltage) - 1]  # remove last point, mainsim produces 1 less
         mainsim_voltage = sim.get_mainsim_voltage_trace(cell_id)
         voltage_diff = voltage - mainsim_voltage
         rms_error = np.sqrt(np.mean(voltage_diff ** 2))