add dataset for re pon and poff only

funkelab · Jul 26, 2024 · 2f44147 · 2f44147
1 parent 3459e1f
commit 2f44147
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/sparkle_stats/datasets/__init__.py b/sparkle_stats/datasets/__init__.py
@@ -1,4 +1,5 @@
 from .p_on_p_off_only_dataset import PonPoffOnlyDataset
+from .re_pon_poff_dataset import RePonPoffDataset
 from .zarr_dataset import ZarrDataset
 from .zarr_intensity_only_dataset import ZarrIntensityOnlyDataset
 from .zarr_state_dataset import ZarrStateOnlyDataset
@@ -8,4 +9,5 @@
     "ZarrStateOnlyDataset",
     "ZarrIntensityOnlyDataset",
     "PonPoffOnlyDataset",
+    "RePonPoffDataset",
 ]
diff --git a/sparkle_stats/datasets/re_pon_poff_dataset.py b/sparkle_stats/datasets/re_pon_poff_dataset.py
@@ -0,0 +1,12 @@
+from .zarr_intensity_only_dataset import ZarrIntensityOnlyDataset
+
+
+class RePonPoffDataset(ZarrIntensityOnlyDataset):
+    def __getitem__(self, item):
+        trace, parameters = super().__getitem__(item)
+        parameters = parameters[:, [0, -2, -1]]
+        return trace, parameters
+
+    @property
+    def output_classes(self):
+        return 3