From 88d51a5dd671ecdb480e4708346172759268f819 Mon Sep 17 00:00:00 2001
From: "max.isi" <max.isi@ligo.org>
Date: Wed, 16 Oct 2024 10:48:17 -0400
Subject: [PATCH] data wip

---
 src/jimgw/single_event/data.py | 90 ++++++++++++++++++++--------------
 1 file changed, 54 insertions(+), 36 deletions(-)

diff --git a/src/jimgw/single_event/data.py b/src/jimgw/single_event/data.py
index c21917b6..c79102ce 100644
--- a/src/jimgw/single_event/data.py
+++ b/src/jimgw/single_event/data.py
@@ -28,48 +28,66 @@
 
 class Data(ABC):
     """
-    Base class for all data.
+    Base class for all data. The time domain data are considered the primary
+    entitiy; the Fourier domain data are derived from an FFT after applying a
+    window. The structure is set up so that :attr:`td` and :attr:`fd` are always
+    Fourier conjugates of each other: the one-sided Fourier series is complete
+    up to the Nyquist frequency
 
     """
-
     name: str
 
-    fd: Float[Array, " n_sample"]
-    td: Float[Array, " n_sample//2"]  # fix this
-    psd: Float[Array, " n_sample//2"]  # fix this
-    frequencies: Float[Array, " n_sample//2"]
-    times: Float[Array, " n_sample"]
+    td: Float[Array, " n_time"]
+    fd: Float[Array, " n_time // 2 + 1"]
+
+    epoch: float
+    delta_t: float
+
+    window: Float[Array, " n_time"]
 
     @property
     def duration(self) -> Float:
         """Duration of the data in seconds."""
-        if len(self.frequencies) == 0:
-            return 0
-        return 1 / (self.frequencies[1] - self.frequencies[0])
-    
+        return self.n_time * self.delta_t
+
+    @property
+    def n_time(self) -> int:
+        """Number of time samples."""
+        return len(self.td)
+
+    @property
+    def n_freq(self) -> int:
+        """Number of frequency samples."""
+        return self.n_time // 2 + 1
+
+    @property
+    def times(self) -> Float[Array, " n_time"]:
+        """Times of the data in seconds."""
+        return jnp.arange(self.n_time) * self.delta_t + self.epoch
+
     @property
-    def delta_t(self) -> Float:
-        """Sampling interval of the data in seconds."""
-        return self.times[1] - self.times[0]
-    
-    def 
-
-    def load_psd_from_data(self, data: TimeSeries, **kws) -> None:
-        seglen = self.duration
-        self.psd = data.psd(fftlength=seglen).value
-
-    def compute_psd_from_gwosc(self,
-                    start_time: Float | None = None,
-                    end_time: Float | None = None,
-                    off_source: bool = True,
-                    **kws) -> None:
-        if data is None:
-            if pad:
-                # pull more data to compute a PSD
-                
-        # n = len(data)
-        # delta_t = 1.0
-        # data = jnp.fft.rfft(data * tukey(n, tukey_alpha)) * delta_t
-        # freq = jnp.fft.rfftfreq(n, delta_t)
-        # return jnp.abs(data) ** 2 / delta_t
-        raise NotImplementedError
+    def frequencies(self) -> Float[Array, " n_time // 2 + 1"]:
+        """Frequencies of the data in Hz."""
+        return jnp.fft.rfftfreq(self.n_time, self.delta_t)
+
+    def __init__(self, td: Float[Array, " n_time"],
+                 delta_t: float = 1.,
+                 epoch: float = 0.,
+                 **kws) -> None:
+        self.td = td
+        self.delta_t = delta_t
+        self.epoch = epoch
+        self.window = kws.get("window", jnp.ones_like(self.td))
+
+    def set_tukey_window(self, alpha: float = 0.4):
+        self.window = jnp.array(tukey(self.n_time, alpha))
+
+    def fft(self, **kws) -> None:
+        if "window" in kws:
+            self.window = kws["window"]
+        self.fd = jnp.fft.rfft(self.td * self.window) * self.delta_t
+
+    def frequency_slice(self, f_min: float, f_max: float) -> \
+            Float[Array, " n_sample"]:
+        f = self.frequencies
+        return self.fd[(f >= f_min) & (f <= f_max)]