From 729e100e3531b34fa0a67436fdb27d55e67c1f8e Mon Sep 17 00:00:00 2001
From: kushalkolar <kushalkolar@gmail.com>
Date: Fri, 6 Oct 2023 02:29:17 -0400
Subject: [PATCH 1/2] playing with pynapple

---
 mesmerize_core/caiman_extensions/_pynapple.py | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 mesmerize_core/caiman_extensions/_pynapple.py

diff --git a/mesmerize_core/caiman_extensions/_pynapple.py b/mesmerize_core/caiman_extensions/_pynapple.py
new file mode 100644
index 0000000..d090409
--- /dev/null
+++ b/mesmerize_core/caiman_extensions/_pynapple.py
@@ -0,0 +1,37 @@
+import numpy as np
+import pandas as pd
+import pynapple as nap
+from caiman.source_extraction.cnmf import CNMF
+
+
+from ._utils import validate
+from .cnmf import CNMFExtensions
+
+
+@pd.api.extensions.register_series_accessor("nap")
+class PynappleExtension:
+    def __init__(self, series: pd.Series):
+        self._series = series
+
+    @validate("cnmf")
+    def get_tsd_frame(self):
+        cnmf_obj: CNMF = self._series.cnmf.get_output()
+
+        framerate = cnmf_obj.params.data["fr"]
+
+        n_frames = cnmf_obj.estimates.C[1]
+
+        duration_seconds = n_frames / framerate
+
+        timestamps = np.linspace(0, duration_seconds, n_frames)
+
+        n_components = cnmf_obj.estimates.C.shape[0]
+
+        tsdframe = nap.TsdFrame(
+            t=timestamps,
+            d=cnmf_obj.estimates.C.T,
+            time_units="s",
+            columns=list(map(range(n_components)))
+        )
+
+        return tsdframe

From 2856837968d27c949dc78f5b38b15b49bf58bfd4 Mon Sep 17 00:00:00 2001
From: kushalkolar <kushalkolar@gmail.com>
Date: Fri, 6 Oct 2023 02:38:01 -0400
Subject: [PATCH 2/2] fix typos

---
 mesmerize_core/caiman_extensions/_pynapple.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mesmerize_core/caiman_extensions/_pynapple.py b/mesmerize_core/caiman_extensions/_pynapple.py
index d090409..5ceaae8 100644
--- a/mesmerize_core/caiman_extensions/_pynapple.py
+++ b/mesmerize_core/caiman_extensions/_pynapple.py
@@ -19,7 +19,7 @@ def get_tsd_frame(self):
 
         framerate = cnmf_obj.params.data["fr"]
 
-        n_frames = cnmf_obj.estimates.C[1]
+        n_frames = cnmf_obj.estimates.C.shape[1]
 
         duration_seconds = n_frames / framerate
 
@@ -31,7 +31,7 @@ def get_tsd_frame(self):
             t=timestamps,
             d=cnmf_obj.estimates.C.T,
             time_units="s",
-            columns=list(map(range(n_components)))
+            columns=list(map(str, range(n_components)))
         )
 
         return tsdframe