From ecb03f8ec536d368f868e6403ea7f92b4b0e9801 Mon Sep 17 00:00:00 2001
From: anilbey <anil.tuncel@epfl.ch>
Date: Tue, 22 Feb 2022 14:27:11 +0100
Subject: [PATCH] thalamus: don't require bpo_threshold_current_dep's current
 (#41)

* remove unused fitness_calculator argument

* consider depolarized threshold current being None when generating currents

* add test for get_thalamus_stim_currents' exception

* pycodestyle fix
---
 emodelrunner/protocols/create_protocols.py   | 10 ++++++++--
 emodelrunner/protocols/thalamus_protocols.py |  3 ---
 tests/unit_tests/test_protocols.py           | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/emodelrunner/protocols/create_protocols.py b/emodelrunner/protocols/create_protocols.py
index e82208c..a6bc092 100644
--- a/emodelrunner/protocols/create_protocols.py
+++ b/emodelrunner/protocols/create_protocols.py
@@ -158,13 +158,19 @@ def get_thalamus_stim_currents(self, responses, mtype, dt):
             mtype (str): mtype to index the responses
             dt (float): timestep of the generated currents (ms)
 
+        Raises:
+            KeyError: when "dep" holding or threshold current is missing
+
         Returns:
             dict: the generated currents in a dict.
         """
         thres_i_hyp = responses[f"{mtype}.bpo_threshold_current_hyp"]
-        thres_i_dep = responses[f"{mtype}.bpo_threshold_current_dep"]
         holding_i_hyp = responses[f"{mtype}.bpo_holding_current_hyp"]
-        holding_i_dep = responses[f"{mtype}.bpo_holding_current_dep"]
+        try:
+            thres_i_dep = responses[f"{mtype}.bpo_threshold_current_dep"]
+            holding_i_dep = responses[f"{mtype}.bpo_holding_current_dep"]
+        except KeyError:
+            thres_i_dep = holding_i_dep = None
         currents = {}
         for protocol in self.protocols.protocols:
             currents.update(
diff --git a/emodelrunner/protocols/thalamus_protocols.py b/emodelrunner/protocols/thalamus_protocols.py
index 883c75d..9981b48 100644
--- a/emodelrunner/protocols/thalamus_protocols.py
+++ b/emodelrunner/protocols/thalamus_protocols.py
@@ -51,7 +51,6 @@ def __init__(
         thdetect_protocol_hyp=None,
         other_protocols=None,
         pre_protocols=None,
-        fitness_calculator=None,
     ):
         """Constructor."""
         # pylint: disable=too-many-arguments
@@ -76,8 +75,6 @@ def __init__(
 
         self.pre_protocols = pre_protocols
 
-        self.fitness_calculator = fitness_calculator
-
     def subprotocols(self):
         """Return all the subprotocols contained in this protocol, is recursive."""
         subprotocols = collections.OrderedDict({self.name: self})
diff --git a/tests/unit_tests/test_protocols.py b/tests/unit_tests/test_protocols.py
index 9999fe0..77f3ada 100644
--- a/tests/unit_tests/test_protocols.py
+++ b/tests/unit_tests/test_protocols.py
@@ -15,6 +15,7 @@
 # limitations under the License.
 
 from pathlib import Path
+from types import SimpleNamespace
 
 from emodelrunner.load import (
     load_config,
@@ -108,6 +109,23 @@ def test_using_thalamus_protocols(self):
             }
             assert prot_obj_names - thalamus_recipe_protocol_keys == set()
 
+    def test_thalamus_stim_currents_exception(self):
+        """Tests the exception case in get_thalamus_stim_currents."""
+        mtype = "test_mtype"
+        responses = {
+            f"{mtype}.bpo_threshold_current_hyp": 0.1,
+            f"{mtype}.bpo_holding_current_hyp": 0.3,
+            f"{mtype}.bpo_threshold_current_dep": 0.4,
+        }
+
+        mock_obj = SimpleNamespace(
+            protocols=[SimpleNamespace(generate_current=(lambda *args: {"args": args}))]
+        )
+
+        thal_protocols = ProtocolBuilder(protocols=mock_obj)
+        currents = thal_protocols.get_thalamus_stim_currents(responses, mtype, dt=0.025)
+        assert currents["args"] == (0.1, None, 0.3, None, 0.025)
+
 
 class TestProtocolParser:
     """Tests for the ProtocolParser class."""