From e09f906c3955d58be99d1ec6adf903121967f7cd Mon Sep 17 00:00:00 2001
From: ilkilic <10600022+ilkilic@users.noreply.github.com>
Date: Thu, 16 Jan 2025 15:20:08 +0100
Subject: [PATCH] Add IV curve plot (#3)

* add iv-curve

* add unit test

* add calculate_rheobase function in tools + unit tests
---
 .gitignore                           |  2 +
 bluecellulab/analysis/analysis.py    | 63 ++++++++++++++++++++++
 bluecellulab/analysis/plotting.py    | 25 +++++++++
 bluecellulab/tools.py                | 74 ++++++++++++++++++++++++++
 tests/test_analysis/test_analysis.py | 78 ++++++++++++++++++++++++++++
 tests/test_tools.py                  | 64 ++++++++++++++++++++++-
 6 files changed, 305 insertions(+), 1 deletion(-)
 create mode 100644 bluecellulab/analysis/analysis.py
 create mode 100644 bluecellulab/analysis/plotting.py
 create mode 100644 tests/test_analysis/test_analysis.py

diff --git a/.gitignore b/.gitignore
index 8abc17c..c3bce7e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 /modlib
 **/x86_64
+**/arm64
 /.tox
 *.pyc
 *.egg-info
@@ -26,3 +27,4 @@ nrnivmodl.log
 *.btr
 *.whl
 .ipynb_checkpoints
+cADpyr_L2TPC_dendrogram.png
\ No newline at end of file
diff --git a/bluecellulab/analysis/analysis.py b/bluecellulab/analysis/analysis.py
new file mode 100644
index 0000000..beb3412
--- /dev/null
+++ b/bluecellulab/analysis/analysis.py
@@ -0,0 +1,63 @@
+"""Module for analyzing cell simulation results."""
+try:
+    import efel
+except ImportError:
+    efel = None
+import numpy as np
+
+from bluecellulab.stimulus import StimulusFactory
+from bluecellulab.tools import calculate_rheobase
+from bluecellulab.analysis.inject_sequence import run_stimulus
+from bluecellulab.analysis.plotting import plot_iv_curve
+
+
+def compute_plot_iv_curve(cell, stim_start=100.0, duration=500.0, post_delay=100.0, threshold_voltage=-30, nb_bins=11):
+    """Compute and plot the IV curve from a given cell by injecting a
+    predefined range of currents.
+
+    Args:
+        cell (bluecellulab.cell.Cell): The initialized cell model.
+        stim_start (float): Start time for current injection (in ms). Default is 100.0 ms.
+        duration (float): Duration of current injection (in ms). Default is 500.0 ms.
+        post_delay (float): Delay after the stimulation ends (in ms). Default is 100.0 ms.
+        nb_bins (int): Number of current injection levels. Default is 11.
+
+    Returns:
+        tuple: A tuple containing:
+            - list_amp (np.ndarray): The predefined injected step current levels (nA).
+            - steady_states (np.ndarray): The corresponding steady-state voltages (mV).
+    """
+    rheobase = calculate_rheobase(cell)
+
+    list_amp = np.linspace(rheobase - 2, rheobase - 0.1, nb_bins)  # [nA]
+
+    steps = []
+    times = []
+    voltages = []
+    # inject step current and record voltage response
+    for amp in list_amp:
+        stim_factory = StimulusFactory(dt=0.1)
+        step_stimulus = stim_factory.step(pre_delay=stim_start, duration=duration, post_delay=post_delay, amplitude=amp)
+        recording = run_stimulus(cell.template_params, step_stimulus, section="soma[0]", segment=0.5)
+        steps.append(step_stimulus)
+        times.append(recording.time)
+        voltages.append(recording.voltage)
+
+    steady_states = []
+    # compute steady state response
+    efel.set_setting('Threshold', threshold_voltage)
+    for voltage, t in zip(voltages, times):
+        trace = {
+            'T': t,
+            'V': voltage,
+            'stim_start': [stim_start],
+            'stim_end': [stim_start + duration]
+        }
+        features_results = efel.get_feature_values([trace], ['steady_state_voltage_stimend'])
+        steady_state = features_results[0]['steady_state_voltage_stimend']
+        steady_states.append(steady_state)
+
+    # plot I-V curve
+    plot_iv_curve(list_amp, steady_states)
+
+    return np.array(list_amp), np.array(steady_states)
diff --git a/bluecellulab/analysis/plotting.py b/bluecellulab/analysis/plotting.py
new file mode 100644
index 0000000..fcca88c
--- /dev/null
+++ b/bluecellulab/analysis/plotting.py
@@ -0,0 +1,25 @@
+"""Module for plotting analysis results of cell simulations."""
+
+import matplotlib.pyplot as plt
+
+
+def plot_iv_curve(currents, voltages):
+    """Plots the IV curve.
+
+    Args:
+        currents (iterable): The injected current levels (nA).
+        voltages (iterable): The corresponding steady-state voltages (mV).
+    Raises:
+        ValueError: If the lengths of currents and voltages do not match.
+    """
+    if len(currents) != len(voltages):
+        raise ValueError("currents and voltages must have the same length")
+
+    # Plotting
+    plt.figure(figsize=(10, 6))
+    plt.plot(voltages, currents, marker='o', linestyle='-', color='b')
+    plt.title("I-V Curve")
+    plt.ylabel("Injected current [nA]")
+    plt.xlabel("Steady state voltage [mV]")
+    plt.tight_layout()
+    plt.show()
diff --git a/bluecellulab/tools.py b/bluecellulab/tools.py
index 5ea43d2..e0a99f5 100644
--- a/bluecellulab/tools.py
+++ b/bluecellulab/tools.py
@@ -23,6 +23,7 @@
 import numpy as np
 
 import bluecellulab
+from bluecellulab.cell import Cell
 from bluecellulab.circuit.circuit_access import EmodelProperties
 from bluecellulab.exceptions import UnsteadyCellError
 from bluecellulab.simulation.parallel import IsolatedProcess
@@ -351,3 +352,76 @@ def check_empty_topology() -> bool:
         neuron.h.topology()
 
     return stdout == ['', '']
+
+
+def calculate_max_thresh_current(cell: Cell, threshold_voltage: float = -30.0) -> float:
+    """Calculate the upper bound threshold current.
+
+    Args:
+        cell (bluecellulab.cell.Cell): The initialized cell model.
+        threshold_voltage (float, optional): Voltage threshold for spike detection. Default is -30.0 mV.
+
+    Returns:
+        float: The upper bound threshold current.
+    """
+    # Calculate resting membrane potential (rmp)
+    rmp = calculate_SS_voltage(
+        template_path=cell.template_params.template_filepath,
+        morphology_path=cell.template_params.morph_filepath,
+        template_format=cell.template_params.template_format,
+        emodel_properties=cell.template_params.emodel_properties,
+        step_level=0.0,
+    )
+
+    # Calculate input resistance (rin)
+    rin = calculate_input_resistance(
+        template_path=cell.template_params.template_filepath,
+        morphology_path=cell.template_params.morph_filepath,
+        template_format=cell.template_params.template_format,
+        emodel_properties=cell.template_params.emodel_properties,
+    )
+
+    # Calculate upperbound threshold current
+    upperbound_threshold_current = (threshold_voltage - rmp) / rin
+    upperbound_threshold_current = np.min([upperbound_threshold_current, 2.0])
+
+    return upperbound_threshold_current
+
+
+def calculate_rheobase(cell: Cell,
+                       threshold_voltage: float = -30.0,
+                       threshold_search_stim_start: float = 300.0,
+                       threshold_search_stim_stop: float = 1000.0) -> float:
+    """Calculate the rheobase by first computing the upper bound threshold
+    current.
+
+    Args:
+        cell (bluecellulab.cell.Cell): The initialized cell model.
+        threshold_voltage (float, optional): Voltage threshold for spike detection. Default is -30.0 mV.
+        threshold_search_stim_start (float, optional): Start time for threshold search stimulation (in ms). Default is 300.0 ms.
+        threshold_search_stim_stop (float, optional): Stop time for threshold search stimulation (in ms). Default is 1000.0 ms.
+
+    Returns:
+        float: The rheobase current.
+    """
+    if cell.template_params.emodel_properties is None:
+        raise ValueError("emodel_properties cannot be None")
+
+    # Calculate upper bound threshold current
+    upperbound_threshold_current = calculate_max_thresh_current(cell, threshold_voltage)
+
+    # Compute rheobase
+    rheobase = search_threshold_current(
+        template_name=cell.template_params.template_filepath,
+        morphology_path=cell.template_params.morph_filepath,
+        template_format=cell.template_params.template_format,
+        emodel_properties=cell.template_params.emodel_properties,
+        hyp_level=cell.template_params.emodel_properties.holding_current,
+        inj_start=threshold_search_stim_start,
+        inj_stop=threshold_search_stim_stop,
+        min_current=cell.template_params.emodel_properties.holding_current,
+        max_current=upperbound_threshold_current,
+        current_precision=0.005,
+    )
+
+    return rheobase
diff --git a/tests/test_analysis/test_analysis.py b/tests/test_analysis/test_analysis.py
new file mode 100644
index 0000000..265eadb
--- /dev/null
+++ b/tests/test_analysis/test_analysis.py
@@ -0,0 +1,78 @@
+"""Unit tests for the analysis module."""
+
+from unittest.mock import MagicMock, patch
+import numpy as np
+import pytest
+from bluecellulab.analysis.analysis import compute_plot_iv_curve
+from pathlib import Path
+from bluecellulab.cell import Cell
+from bluecellulab.circuit.circuit_access import EmodelProperties
+
+
+parent_dir = Path(__file__).resolve().parent.parent
+
+
+class MockRecording:
+    def __init__(self):
+        self.time = [1, 2, 3]
+        self.voltage = [-70, -55, -40]
+
+
+@pytest.fixture
+def mock_run_stimulus():
+    return MagicMock(return_value=MockRecording())
+
+
+@pytest.fixture
+def mock_search_threshold_current():
+    return MagicMock(return_value=0.1)
+
+
+@pytest.fixture
+def mock_steady_state_voltage_stimend():
+    return MagicMock(return_value=-65)
+
+
+@pytest.fixture
+def mock_efel():
+    efel_mock = MagicMock()
+    efel_mock.getFeatureValues.return_value = [{'steady_state_voltage_stimend': [-65]}]
+    return efel_mock
+
+
+@pytest.fixture
+def mock_cell():
+    emodel_properties = EmodelProperties(
+        threshold_current=1.1433533430099487,
+        holding_current=1.4146618843078613,
+        AIS_scaler=1.4561502933502197,
+        soma_scaler=1.0
+    )
+    cell = Cell(
+        f"{parent_dir}/examples/circuit_sonata_quick_scx/components/hoc/cADpyr_L2TPC.hoc",
+        f"{parent_dir}/examples/circuit_sonata_quick_scx/components/morphologies/asc/rr110330_C3_idA.asc",
+        template_format="v6",
+        emodel_properties=emodel_properties
+    )
+    return cell
+
+
+def test_plot_iv_curve(mock_cell, mock_run_stimulus, mock_search_threshold_current, mock_efel):
+    """Test the plot_iv_curve function."""
+    with patch('bluecellulab.cell.core', mock_cell), \
+         patch('bluecellulab.analysis.analysis.run_stimulus', mock_run_stimulus), \
+         patch('bluecellulab.tools.search_threshold_current', mock_search_threshold_current), \
+         patch('bluecellulab.analysis.analysis.efel', mock_efel):
+
+        stim_start = 100.0
+        duration = 500.0
+        post_delay = 100.0
+        threshold_voltage = -30
+        nb_bins = 11
+
+        list_amp, steady_states = compute_plot_iv_curve(mock_cell, stim_start, duration, post_delay, threshold_voltage, nb_bins)
+
+        assert isinstance(list_amp, np.ndarray)
+        assert isinstance(steady_states, np.ndarray)
+        assert len(list_amp) == nb_bins
+        assert len(steady_states) == nb_bins
diff --git a/tests/test_tools.py b/tests/test_tools.py
index b7c45e3..7571029 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -17,17 +17,52 @@
 
 import numpy as np
 import pytest
+from unittest.mock import MagicMock, patch
 
+from bluecellulab.cell import Cell
 from bluecellulab import CircuitSimulation
 from bluecellulab.cell.ballstick import create_ball_stick
 from bluecellulab.circuit.circuit_access import EmodelProperties
 from bluecellulab.circuit.node_id import create_cell_id
 from bluecellulab.exceptions import UnsteadyCellError
-from bluecellulab.tools import calculate_SS_voltage, calculate_SS_voltage_subprocess, calculate_input_resistance, detect_hyp_current, detect_spike, detect_spike_step, detect_spike_step_subprocess, holding_current, holding_current_subprocess, search_threshold_current, template_accepts_cvode, check_empty_topology
+from bluecellulab.tools import calculate_SS_voltage, calculate_SS_voltage_subprocess, calculate_input_resistance, detect_hyp_current, detect_spike, detect_spike_step, detect_spike_step_subprocess, holding_current, holding_current_subprocess, search_threshold_current, template_accepts_cvode, check_empty_topology, calculate_max_thresh_current, calculate_rheobase
+
 
 script_dir = Path(__file__).parent
 
 
+@pytest.fixture
+def mock_cell():
+    emodel_properties = EmodelProperties(
+        threshold_current=1.1433533430099487,
+        holding_current=1.4146618843078613,
+        AIS_scaler=1.4561502933502197,
+        soma_scaler=1.0
+    )
+    cell = Cell(
+        f"{script_dir}/examples/circuit_sonata_quick_scx/components/hoc/cADpyr_L2TPC.hoc",
+        f"{script_dir}/examples/circuit_sonata_quick_scx/components/morphologies/asc/rr110330_C3_idA.asc",
+        template_format="v6",
+        emodel_properties=emodel_properties
+    )
+    return cell
+
+
+@pytest.fixture
+def mock_calculate_SS_voltage():
+    return MagicMock(return_value=-65.0)
+
+
+@pytest.fixture
+def mock_calculate_input_resistance():
+    return MagicMock(return_value=100.0)
+
+
+@pytest.fixture
+def mock_search_threshold_current():
+    return MagicMock(return_value=0.1)
+
+
 @pytest.mark.v5
 def test_calculate_SS_voltage_subprocess():
     """Tools: Test calculate_SS_voltage"""
@@ -241,3 +276,30 @@ def test_check_empty_topology():
     assert check_empty_topology() is True
     cell = create_ball_stick()
     assert check_empty_topology() is False
+
+
+def test_calculate_max_thresh_current(mock_cell, mock_calculate_SS_voltage, mock_calculate_input_resistance):
+    """Test the calculate_max_thresh_current function."""
+    with patch('bluecellulab.tools.calculate_SS_voltage', mock_calculate_SS_voltage), \
+         patch('bluecellulab.tools.calculate_input_resistance', mock_calculate_input_resistance):
+
+        threshold_voltage = -30.0
+        upperbound_threshold_current = calculate_max_thresh_current(mock_cell, threshold_voltage)
+
+        assert upperbound_threshold_current == (threshold_voltage - mock_calculate_SS_voltage.return_value) / mock_calculate_input_resistance.return_value
+        assert upperbound_threshold_current == 0.35
+
+
+def test_calculate_rheobase(mock_cell, mock_calculate_SS_voltage, mock_calculate_input_resistance, mock_search_threshold_current):
+    """Test the calculate_rheobase function."""
+    with patch('bluecellulab.tools.calculate_SS_voltage', mock_calculate_SS_voltage), \
+         patch('bluecellulab.tools.calculate_input_resistance', mock_calculate_input_resistance), \
+         patch('bluecellulab.tools.search_threshold_current', mock_search_threshold_current):
+
+        threshold_voltage = -30.0
+        threshold_search_stim_start = 300.0
+        threshold_search_stim_stop = 1000.0
+
+        rheobase = calculate_rheobase(mock_cell, threshold_voltage, threshold_search_stim_start, threshold_search_stim_stop)
+
+        assert rheobase == mock_search_threshold_current.return_value