From f2574e106680add41967d76b4c2e4733fe43f3d6 Mon Sep 17 00:00:00 2001
From: Jochem Smit <jhsmit@gmail.com>
Date: Tue, 26 Nov 2024 09:52:16 +0100
Subject: [PATCH] refactor and test process photon data method

---
 dont_fret/process.py         |  6 +++---
 dont_fret/web/datamanager.py |  6 ++++--
 dont_fret/web/methods.py     |  4 ++--
 tests/test_models.py         | 14 ++++++++++++++
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/dont_fret/process.py b/dont_fret/process.py
index 5bc775a..2a4297a 100644
--- a/dont_fret/process.py
+++ b/dont_fret/process.py
@@ -9,7 +9,7 @@
 from tqdm.auto import tqdm
 
 from dont_fret.config import cfg
-from dont_fret.config.config import BurstColor, DontFRETConfig
+from dont_fret.config.config import BurstColor
 from dont_fret.fileIO import PhotonFile
 from dont_fret.models import Bursts, PhotonData
 
@@ -73,10 +73,10 @@ def batch_search_and_save(
             f.result()
 
 
-def full_search(
+def process_photon_data(
     photon_data: PhotonData, burst_colors: list[BurstColor], hooks: dict = cfg.hooks
 ) -> Bursts:
-    """search and alex / fret cde"""
+    """search and apply hooks"""
     bursts = photon_data.burst_search(burst_colors)
 
     # Apply hooks
diff --git a/dont_fret/web/datamanager.py b/dont_fret/web/datamanager.py
index 4127f9d..cfd90a8 100644
--- a/dont_fret/web/datamanager.py
+++ b/dont_fret/web/datamanager.py
@@ -14,7 +14,7 @@
 from dont_fret.config.config import BurstColor, DontFRETConfig, cfg
 from dont_fret.fileIO import PhotonFile
 from dont_fret.models import Bursts, PhotonData
-from dont_fret.process import full_search
+from dont_fret.process import process_photon_data
 from dont_fret.web.methods import get_duration, get_info, make_burst_dataframe
 from dont_fret.web.models import BurstNode, PhotonNode
 
@@ -78,7 +78,9 @@ async def get_bursts(
 
             try:
                 photon_data = await self.get_photons(photon_node)
-                bursts = await self.run(full_search, photon_data, burst_colors, self.cfg.hooks)
+                bursts = await self.run(
+                    process_photon_data, photon_data, burst_colors, self.cfg.hooks
+                )
 
                 future.set_result(bursts)
             except Exception as e:
diff --git a/dont_fret/web/methods.py b/dont_fret/web/methods.py
index 433b344..0cb95c2 100644
--- a/dont_fret/web/methods.py
+++ b/dont_fret/web/methods.py
@@ -15,7 +15,7 @@
 from dont_fret.config.config import BurstColor
 from dont_fret.fileIO import PhotonFile
 from dont_fret.models import Bursts, PhotonData
-from dont_fret.process import full_search
+from dont_fret.process import process_photon_data
 from dont_fret.web.models import BurstFilterItem, BurstNode, FRETNode, PhotonNode
 
 
@@ -58,7 +58,7 @@ def make_burst_nodes(
 
     hooks = hooks or {}
     for name, burst_colors in burst_settings.items():
-        bursts = [full_search(photon_data, burst_colors, hooks) for photon_data in photons]
+        bursts = [process_photon_data(photon_data, burst_colors, hooks) for photon_data in photons]
         # bursts = [photons.burst_search(burst_colors) for photons in photons]
         # if alex_2cde:
         #     bursts = [b.alex_2cde(photons) for b, photons in zip(bursts, photons)]
diff --git a/tests/test_models.py b/tests/test_models.py
index 5f35f20..f951e9e 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -8,6 +8,7 @@
 from dont_fret.config.config import BurstColor
 from dont_fret.fileIO import PhotonFile
 from dont_fret.models import BinnedPhotonData, Bursts, PhotonData
+from dont_fret.process import process_photon_data
 
 cwd = Path(__file__).parent
 input_data_dir = cwd / "test_data" / "input"
@@ -80,6 +81,8 @@ def test_load_save_bursts(dcbs_bursts: Bursts, tmp_path: Path):
     assert dcbs_bursts.burst_data.equals(bursts_load.burst_data)
     assert dcbs_bursts.photon_data.equals(bursts_load.photon_data)
     assert dcbs_bursts.cfg == bursts_load.cfg
+    assert dcbs_bursts.cfg
+    assert bursts_load.cfg
     assert asdict(dcbs_bursts.cfg) == asdict(bursts_load.cfg)
     assert dcbs_bursts.metadata == bursts_load.metadata
 
@@ -109,6 +112,17 @@ def test_burst_search(ph_ds1: PhotonData):
         assert (df_ref["time_max"] == df_test["timestamps_max"] * ph_ds1.timestamps_unit).all()
 
 
+def test_process_photon_data(ph_ds1: PhotonData):
+    hooks = {
+        "alex_2cde": {},
+        "fret_2cde": {},
+    }
+
+    bursts = process_photon_data(ph_ds1, APBS_TEST, hooks=hooks)
+    assert "alex_2cde" in bursts.burst_data.columns
+    assert "fret_2cde" in bursts.burst_data.columns
+
+
 def test_binning(ph_ds1):
     trace = BinnedPhotonData(ph_ds1, bounds=(0, ph_ds1.tmax))
     assert len(trace) == 81439