From 85f69f58846af94e547c785e5e83d42413c68f6f Mon Sep 17 00:00:00 2001
From: Alexis Arnaudon <alexis.arnaudon@epfl.ch>
Date: Mon, 13 May 2024 10:10:57 +0200
Subject: [PATCH] Feat: Add python dynamical feature registration (#391)

---
 efel/api.py         |  6 ++++++
 tests/test_basic.py | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/efel/api.py b/efel/api.py
index b4063d67..aefda482 100644
--- a/efel/api.py
+++ b/efel/api.py
@@ -411,6 +411,12 @@ def get_mean_feature_values(
     return featureDicts  # type: ignore
 
 
+def register_feature(feature_function: Callable):
+    """Register a new feature."""
+    pyfeatures.all_pyfeatures.append(feature_function.__name__)
+    setattr(pyfeatures, feature_function.__name__, feature_function)
+
+
 reset()
 
 
diff --git a/tests/test_basic.py b/tests/test_basic.py
index d7e158e6..1c7c88af 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -40,6 +40,7 @@
 from efel.api import get_feature_values
 from efel.api import set_setting
 from efel.api import get_distance
+from efel.api import register_feature
 
 
 testdata_dir = Path(__file__).parent / 'testdata'
@@ -4522,3 +4523,29 @@ def test_spontaneous_firing():
     numpy.testing.assert_allclose(ap_fall_time, ap_fall_time_expected)
     numpy.testing.assert_allclose(ap_rise_rate, ap_rise_rate_expected)
     numpy.testing.assert_allclose(ap_fall_rate, ap_fall_rate_expected)
+
+
+def test_register_feature():
+    """basic: Test register feature"""
+    import efel
+    efel.reset()
+
+    trace = {}
+    trace['T'] = numpy.arange(0, 100, 0.1)
+    trace['V'] = numpy.ones(len(trace['T'])) * -80.0
+    trace['stim_start'] = [25]
+    trace['stim_end'] = [75]
+
+    def new_feature():
+        v = efel.pyfeatures.get_cpp_feature("voltage")
+        stim_start = efel.pyfeatures.pyfeatures._get_cpp_data("stim_start")
+        return numpy.array([v[0], stim_start])
+
+    efel.register_feature(new_feature)
+    result = efel.get_feature_values([trace], ["new_feature"])[0]["new_feature"]
+    assert result[0] == -80.0
+    assert result[1] == 25.0
+
+    # remove it to keep other tests untouched
+    del efel.pyfeatures.new_feature
+    del efel.pyfeatures.all_pyfeatures[-1]