sgn-drift 0.1.0__py3-none-any.whl

sgndrift/psd/estimators.py

@@ -0,0 +1,150 @@
+"""Core PSD estimation logic classes (Math only).
+Refactored to remove invalid boundary zeroing and enforce input shape.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+from scipy.special import loggamma
+from sympy import EulerGamma
+
+EULERGAMMA = float(EulerGamma.evalf())
+
+
+@dataclass
+class BaseEstimator(ABC):
+    """Base class for PSD estimation logic."""
+
+    size: int
+    normalization: float = 1.0
+
+    # Internal State
+    n_samples: int = field(init=False, default=0)
+    current_psd: np.ndarray = field(init=False, repr=False, default=None)
+
+    def __post_init__(self):
+        # Initialize with ones to avoid divide-by-zero
+        self.current_psd = np.ones(self.size)
+
+    def _check_shape(self, data: np.ndarray) -> None:
+        """Validate input data shape matches estimator configuration."""
+        if data.shape[-1] != self.size:
+            raise ValueError(
+                f"Input data size {data.shape[-1]} does not match estimator size {self.size}"
+            )
+
+    @abstractmethod
+    def update(self, data: np.ndarray) -> None:
+        """Update state with new frequency-domain data."""
+        pass
+
+    def get_psd(self) -> np.ndarray:
+        return self.current_psd
+
+
+@dataclass
+class MGMEstimator(BaseEstimator):
+    """
+    Median-Geometric-Mean Estimator (Standard LIGO).
+    """
+
+    n_median: int = 7
+    n_average: int = 64
+
+    # Internal State
+    history: deque = field(init=False, repr=False, default=None)
+    geo_mean_log: Optional[np.ndarray] = field(init=False, repr=False, default=None)
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.history = deque(maxlen=self.n_median)
+
+    @staticmethod
+    def _median_bias(nn):
+        """XLALMedianBias"""
+        ans = 1.0
+        n = (nn - 1) // 2
+        for i in range(1, n + 1):
+            ans -= 1.0 / (2 * i)
+            ans += 1.0 / (2 * i + 1)
+        return ans
+
+    @staticmethod
+    def _log_median_bias_geometric(nn):
+        """XLALLogMedianBiasGeometric"""
+        return np.log(MGMEstimator._median_bias(nn)) - nn * (
+            loggamma(1.0 / nn) - np.log(nn)
+        )
+
+    def update(self, data: np.ndarray) -> None:
+        self._check_shape(data)
+
+        if np.iscomplexobj(data):
+            power = np.abs(data) ** 2
+        else:
+            power = data
+
+        self.history.append(power)
+
+        if self.n_samples == 0:
+            self.geo_mean_log = np.log(power)
+            self.n_samples += 1
+        else:
+            self.n_samples = min(self.n_samples + 1, self.n_average)
+
+            bias = self._log_median_bias_geometric(len(self.history))
+
+            # Match Legacy: use sort and integer index
+            stacked = np.array(self.history)
+            sorted_bins = np.sort(stacked, axis=0)
+            idx = len(self.history) // 2
+            log_bin_median = np.log(sorted_bins[idx])
+
+            self.geo_mean_log = (
+                self.geo_mean_log * (self.n_samples - 1) + log_bin_median - bias
+            ) / self.n_samples
+
+        self.current_psd = np.exp(self.geo_mean_log + EULERGAMMA) * self.normalization
+
+    def set_reference(self, psd: np.ndarray, weight: int):
+        self._check_shape(psd)
+
+        raw = psd / self.normalization
+        # Avoid log(0)
+        raw = np.where(raw > 0, raw, 1e-300)
+
+        self.history.clear()
+        for _ in range(self.n_median):
+            self.history.append(raw)
+
+        self.geo_mean_log = np.log(raw) - EULERGAMMA
+        self.n_samples = min(weight, self.n_average)
+        self.current_psd = psd.copy()
+
+
+@dataclass
+class RecursiveEstimator(BaseEstimator):
+    """
+    Exponential Moving Average Estimator.
+    """
+
+    alpha: float = 0.1
+    _initialized: bool = field(init=False, default=False)
+
+    def update(self, data: np.ndarray) -> None:
+        self._check_shape(data)
+
+        power = (
+            np.abs(data) ** 2 if np.iscomplexobj(data) else data
+        ) * self.normalization
+
+        if not self._initialized:
+            self.current_psd = power
+            self._initialized = True
+        else:
+            self.current_psd = (1 - self.alpha) * self.current_psd + self.alpha * power

sgndrift/sinks/drift_sink.py

@@ -0,0 +1,154 @@
+"""
+Sinks for Drift Events.
+"""
+
+from __future__ import annotations
+
+import csv
+import os
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, Optional, TextIO
+
+from sgn.base import SinkPad
+from sgnts.base import EventFrame, TSFrame, TSSink
+from sgndrift.transforms.drift import DriftEvent
+
+
+@dataclass
+class DriftCSVSink(TSSink):
+    """
+    Writes DriftEvent data to a CSV file.
+    Inherits from TSSink to integrate with sgn-ts pipelines.
+    """
+
+    filename: str = "drift.csv"
+
+    # Mark 'in' as unaligned to prevent Audioadapter creation for discrete events
+    static_unaligned_sink_pads: ClassVar[list[str]] = ["in"]
+
+    # Internal state
+    _file: Optional[TextIO] = field(init=False, repr=False, default=None)
+    _writer: Any = field(init=False, repr=False, default=None)
+
+    def __post_init__(self):
+        # Force all input pads to be unaligned to prevent Audioadapter creation.
+        # This is necessary because EventFrames are discrete and lack sample rates.
+        # We set this before super().__post_init__() so TimeSeriesMixin uses it.
+        self.unaligned = list(self.sink_pad_names)
+        super().__post_init__()
+
+    def configure(self) -> None:
+        """Configure input frame types to expect EventFrame."""
+        for name in self.sink_pad_names:
+            self.input_frame_types[name] = EventFrame
+
+    @property
+    def min_latest(self) -> int:
+        """
+        Override min_latest to handle the case where all inputs are unaligned.
+        Base implementation crashes if self.inbufs is empty.
+        """
+        if not self.inbufs:
+            latest_offsets = []
+            for pad in self.unaligned_sink_pads:
+                frame = self.unaligned_data.get(pad)
+                if frame and hasattr(frame, "data") and frame.data:
+                    # Assuming frame.data is list of buffers
+                    latest_offsets.append(frame.data[-1].noffset)
+            return max(latest_offsets) if latest_offsets else 0
+        return super().min_latest
+
+    @property
+    def earliest(self) -> int:
+        """
+        Override earliest to handle the case where all inputs are unaligned.
+        """
+        if not self.inbufs:
+            earliest_offsets = []
+            for pad in self.unaligned_sink_pads:
+                frame = self.unaligned_data.get(pad)
+                if frame and hasattr(frame, "data") and frame.data:
+                    earliest_offsets.append(frame.data[0].offset)
+            return min(earliest_offsets) if earliest_offsets else 0
+        return super().earliest
+
+    def _align(self) -> None:
+        """
+        Override alignment logic.
+        Since input is unaligned, base class _align() would fail.
+        We simply check if unaligned data is present.
+        """
+        # Assume alignment is satisfied if we have data on the first pad
+        # For multiple pads, we might want to check all, but TSSink usually has one.
+        if not self.sink_pads:
+            self._is_aligned = False
+            return
+
+        sink_pad = self.sink_pads[0]
+        if self.unaligned_data.get(sink_pad) is not None:
+            self._is_aligned = True
+        else:
+            self._is_aligned = False
+
+    def process(self, input_frames: dict[SinkPad, TSFrame]) -> None:
+        """
+        Process incoming frames and write to CSV.
+        TSSink.internal() calls this method with frames collected from all pads.
+        """
+        # We assume a single sink pad named "in"
+        # Since we configured the pad to expect EventFrame, input_frames contains EventFrames.
+        if not self.sink_pads:
+            return
+
+        sink_pad = self.sink_pads[0]
+        frame = input_frames.get(sink_pad)
+
+        if frame is None:
+            return
+
+        if frame.EOS:
+            self.mark_eos(sink_pad)
+
+        if frame.is_gap:
+            return
+
+        # Check for data
+        if not hasattr(frame, "data") or not frame.data:
+            return
+
+        for buf in frame.data:
+            if not hasattr(buf, "data") or not buf.data:
+                continue
+
+            for event in buf.data:
+                if not isinstance(event, DriftEvent):
+                    continue
+
+                row = {"time": event.epoch}
+                row.update(event.data)
+
+                if self._file is None:
+                    self._open_file(row.keys())
+
+                self._writer.writerow(row)
+
+        if self._file:
+            self._file.flush()
+
+    def _open_file(self, keys):
+        exists = os.path.exists(self.filename)
+        self._file = open(self.filename, "a", newline="")
+        # Ensure deterministic column order with 'time' first
+        data_keys = sorted([k for k in keys if k != "time"])
+        fieldnames = ["time"] + data_keys
+        self._writer = csv.DictWriter(self._file, fieldnames=fieldnames)
+        if not exists:
+            self._writer.writeheader()
+
+    def cleanup(self):
+        if self._file:
+            self._file.close()
+            self._file = None
+
+    def __del__(self):
+        self.cleanup()

sgndrift/transforms/drift.py

@@ -0,0 +1,145 @@
+"""
+Geometric Diagnostics: Elements for tracking the manifold velocity of detector noise.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import ClassVar, Dict, Tuple
+
+import numpy as np
+
+from sgn.base import SourcePad
+from sgndrift.transforms.psd import PSDEvent
+from sgndrift.psd.drift import calculate_fisher_velocity
+from sgnts.base import EventBuffer, EventFrame, TSTransform
+
+
+@dataclass
+class DriftEvent:
+    """
+    Container for Fisher Information Velocity (Drift) data.
+    """
+
+    epoch: float
+    data: Dict[str, float]
+
+
+@dataclass
+class FisherVelocity(TSTransform):
+    """
+    Computes Fisher Information Velocity (Geometric Drift) between consecutive PSDs.
+
+    Wraps sgndrift.psd.drift.calculate_fisher_velocity.
+
+    Inputs:
+        EventFrame containing PSDEvent objects.
+
+    Outputs:
+        EventFrame containing DriftEvent objects.
+    """
+
+    # Mark 'in' as unaligned to prevent TimeSeriesMixin from creating an Audioadapter.
+    static_unaligned_sink_pads: ClassVar[list[str]] = ["in"]
+
+    bands: Dict[str, Tuple[float, float]] = field(default_factory=dict)
+
+    _prev_data: np.ndarray = field(init=False, repr=False, default=None)
+    _prev_epoch: float = field(init=False, repr=False, default=None)
+
+    def configure(self) -> None:
+        """Configure element-specific attributes."""
+        # Inform the element that it handles EventFrames
+        for name in self.sink_pad_names:
+            self.input_frame_types[name] = EventFrame
+        for name in self.source_pad_names:
+            self.output_frame_types[name] = EventFrame
+
+    @property
+    def min_latest(self) -> int:
+        if not self.inbufs:
+            latest_offsets = []
+            for pad in self.unaligned_sink_pads:
+                frame = self.unaligned_data.get(pad)
+                if frame and frame.data:
+                    latest_offsets.append(frame.data[-1].noffset)
+            return max(latest_offsets) if latest_offsets else 0
+        return super().min_latest
+
+    @property
+    def earliest(self) -> int:
+        if not self.inbufs:
+            earliest_offsets = []
+            for pad in self.unaligned_sink_pads:
+                frame = self.unaligned_data.get(pad)
+                if frame and frame.data:
+                    earliest_offsets.append(frame.data[0].offset)
+            return min(earliest_offsets) if earliest_offsets else 0
+        return super().earliest
+
+    def _align(self) -> None:
+        sink_pad = self.sink_pads[0]
+        if self.unaligned_data.get(sink_pad) is not None:
+            self._is_aligned = True
+        else:
+            self._is_aligned = False
+
+    def new(self, pad: SourcePad) -> EventFrame:
+        sink_pad = self.sink_pads[0]
+        in_frame = self.unaligned_data.get(sink_pad)
+        self.unaligned_data[sink_pad] = None
+
+        if in_frame is None or in_frame.is_gap:
+            return EventFrame(is_gap=True, EOS=in_frame.EOS if in_frame else False)
+
+        if not hasattr(in_frame, "data") or not in_frame.data:
+            return EventFrame(is_gap=True, EOS=in_frame.EOS)
+        if not in_frame.data[0].data:
+            return EventFrame(is_gap=True, EOS=in_frame.EOS)
+
+        psd_event = in_frame.data[0].data[0]
+
+        if not isinstance(psd_event, PSDEvent):
+            return EventFrame(is_gap=True, EOS=in_frame.EOS)
+
+        current_data = psd_event.data
+        current_epoch = psd_event.epoch
+        freqs = psd_event.frequencies
+        df = psd_event.delta_f
+
+        drift_results = {}
+
+        # Only calculate if we have history
+        if self._prev_data is not None:
+            dt = current_epoch - self._prev_epoch
+            if dt > 0:
+                drift_results = calculate_fisher_velocity(
+                    current_psd=current_data,
+                    previous_psd=self._prev_data,
+                    dt=dt,
+                    frequencies=freqs,
+                    delta_f=df,
+                    bands=self.bands,
+                )
+
+        # Update History
+        self._prev_data = current_data.copy()
+        self._prev_epoch = current_epoch
+
+        # Handle startup transient (return zeros instead of empty)
+        if not drift_results:
+            bands_keys = self.bands.keys() if self.bands else ["total"]
+            drift_results = {k: 0.0 for k in bands_keys}
+
+        out_event = DriftEvent(epoch=current_epoch, data=drift_results)
+
+        buf = in_frame.data[0]
+        ts = buf.offset
+        dur = (
+            buf.duration if hasattr(buf, "duration") and buf.duration else 1_000_000_000
+        )
+        te = ts + dur
+
+        out_buf = EventBuffer.from_span(ts, te, [out_event])
+
+        return EventFrame(data=[out_buf], EOS=in_frame.EOS)

sgndrift/transforms/psd.py

@@ -0,0 +1,190 @@
+"""
+SGN Elements for PSD Estimation.
+wraps sgnligo.psd.estimators logic into TSTransform elements.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+import scipy.signal
+from sgn.base import SourcePad
+from sgnts.base import (
+    AdapterConfig,
+    EventBuffer,
+    EventFrame,
+    Offset,
+    TSTransform,
+)
+
+from sgndrift.psd.estimators import BaseEstimator, MGMEstimator, RecursiveEstimator
+
+# Optional LAL import for conversion methods
+try:
+    import lal
+except ImportError:
+    lal = None
+
+
+@dataclass
+class PSDEvent:
+    """
+    Container for a PSD estimate event (Pure Python/NumPy).
+    Decoupled from LAL to ensure stability in non-LAL environments.
+    """
+
+    data: np.ndarray
+    frequencies: np.ndarray
+    epoch: float
+    delta_f: float
+
+    def to_lal(self) -> Optional[object]:
+        """
+        Convert to LAL REAL8FrequencySeries if LAL is available.
+        Uses standard 'strain^2 s' unit definition.
+        """
+        if lal is None:
+            return None
+
+        try:
+            # Standard unit construction used in sgnligo.psd.psd
+            unit = lal.Unit("strain^2 s")
+
+            series = lal.CreateREAL8FrequencySeries(
+                "psd",
+                lal.LIGOTimeGPS(self.epoch),
+                0.0,
+                self.delta_f,
+                unit,
+                len(self.data),
+            )
+            series.data.data = self.data
+            return series
+        except Exception:
+            return None
+
+
+@dataclass
+class PSDEstimator(TSTransform):
+    """
+    Base TSTransform for PSD Estimation.
+    Outputs EventFrame containing PSDEvent objects.
+    """
+
+    fft_length: float = 4.0
+    overlap: float = 0.5
+    sample_rate: int = 16384
+    window_type: str = "hann"
+
+    # Internal state
+    _estimator: BaseEstimator = field(init=False, repr=False, default=None)
+    _window: np.ndarray = field(init=False, repr=False, default=None)
+    _freqs: np.ndarray = field(init=False, repr=False, default=None)
+    _norm_factor: float = field(init=False, repr=False, default=1.0)
+    _delta_f: float = field(init=False, repr=False, default=0.0)
+
+    def __post_init__(self):
+        n_samples = int(self.fft_length * self.sample_rate)
+        stride = int(n_samples * (1 - self.overlap))
+        overlap_samples = n_samples - stride
+
+        self.adapter_config = AdapterConfig()
+        self.adapter_config.stride = Offset.fromsamples(stride, self.sample_rate)
+        self.adapter_config.overlap = (
+            0,
+            Offset.fromsamples(overlap_samples, self.sample_rate),
+        )
+        self.adapter_config.skip_gaps = True
+
+        super().__post_init__()
+
+        self._window = scipy.signal.get_window(self.window_type, n_samples)
+        s2 = np.sum(self._window**2)
+        self._norm_factor = 2.0 / (self.sample_rate * s2)
+
+        self._freqs = np.fft.rfftfreq(n_samples, d=1 / self.sample_rate)
+        self._delta_f = self._freqs[1] - self._freqs[0]
+
+        self._init_estimator(len(self._freqs))
+
+    def _init_estimator(self, size: int):
+        raise NotImplementedError
+
+    def new(self, pad: SourcePad) -> EventFrame:
+        in_frame = self.preparedframes[self.sink_pads[0]]
+
+        if in_frame.is_gap or not in_frame.buffers:
+            return EventFrame(is_gap=True, EOS=in_frame.EOS)
+
+        buf = in_frame.buffers[0]
+        data = buf.data
+
+        if len(data) != len(self._window):
+            return EventFrame(is_gap=True, EOS=in_frame.EOS)
+
+        # 1. Compute FFT
+        windowed = data * self._window
+        fft_data = np.fft.rfft(windowed)
+
+        # 2. Update Estimator
+        self._estimator.update(fft_data)
+        psd_data = self._estimator.get_psd().copy()
+
+        # 3. Create Output Event
+        # Calculate timestamps in nanoseconds for EventBuffer
+        ts = Offset.tons(buf.offset)
+        # Duration is derived from buffer length
+        duration_samples = len(data)
+        duration_offset = Offset.fromsamples(duration_samples, self.sample_rate)
+        te = Offset.tons(buf.offset + duration_offset)
+
+        # Epoch for PSD metadata (start of window)
+        epoch = Offset.tosec(buf.offset)
+
+        event = PSDEvent(
+            data=psd_data, frequencies=self._freqs, epoch=epoch, delta_f=self._delta_f
+        )
+
+        # Use factory method to avoid constructor signature issues
+        out_buf = EventBuffer.from_span(ts, te, [event])
+
+        meta = in_frame.metadata.copy() if in_frame.metadata else {}
+
+        lal_obj = event.to_lal()
+        if lal_obj:
+            meta["psd"] = lal_obj
+
+        meta["psd_numpy"] = psd_data
+        meta["psd_freqs"] = self._freqs
+
+        return EventFrame(data=[out_buf], metadata=meta, EOS=in_frame.EOS)
+
+
+@dataclass
+class RecursivePSD(PSDEstimator):
+    """Fast, Low-Latency PSD Estimator."""
+
+    alpha: float = 0.1
+
+    def _init_estimator(self, size: int):
+        self._estimator = RecursiveEstimator(
+            size=size, normalization=self._norm_factor, alpha=self.alpha
+        )
+
+
+@dataclass
+class MGMPSD(PSDEstimator):
+    """Standard Median-Geometric-Mean PSD Estimator."""
+
+    n_median: int = 7
+    n_average: int = 64
+
+    def _init_estimator(self, size: int):
+        self._estimator = MGMEstimator(
+            size=size,
+            normalization=self._norm_factor,
+            n_median=self.n_median,
+            n_average=self.n_average,
+        )