LogPSplinePSD 0.0.3__py3-none-any.whl

log_psplines/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.0.3'
+__version_tuple__ = version_tuple = (0, 0, 3)

log_psplines/arviz_utils.py

@@ -0,0 +1,142 @@
+import os
+from typing import List
+
+import arviz as az
+import numpy as np
+from xarray import DataArray, Dataset
+
+from .psplines import LogPSplines
+
+
+def get_weights(
+    idata: az.InferenceData,
+    thin: int = 10,
+) -> np.ndarray:
+    """
+    Extract weight samples from arviz InferenceData.
+
+    Parameters
+    ----------
+    idata : az.InferenceData
+        Inference data containing weight samples
+    thin : int
+        Thinning factor
+
+    Returns
+    -------
+    jnp.ndarray
+        Weight samples, shape (n_samples_thinned, n_weights)
+    """
+    # Get weight samples and flatten chains
+    weight_samples = (
+        idata.posterior.weights.values
+    )  # (chains, draws, n_weights)
+    weight_samples = weight_samples.reshape(
+        -1, weight_samples.shape[-1]
+    )  # (chains*draws, n_weights)
+
+    # Thin samples
+    return weight_samples[::thin]
+
+
+def get_psd_samples_arviz(
+    idata: az.InferenceData, spline_model: LogPSplines, thin: int = 10
+) -> np.ndarray:
+    """
+    Extract PSD samples from arviz InferenceData.
+
+    Parameters
+    ----------
+    idata : az.InferenceData
+        Inference data containing weight samples
+    spline_model : LogPSplines
+        Spline model for reconstruction
+    thin : int
+        Thinning factor
+
+    Returns
+    -------
+    jnp.ndarray
+        PSD samples, shape (n_samples_thinned, n_frequencies)
+    """
+    # Get weight samples and flatten chains
+    weight_samples = get_weights(idata, thin=thin)
+
+    # Compute PSD samples
+    psd_samples = []
+    for weights in weight_samples:
+        ln_spline = spline_model.basis.T @ weights
+        ln_psd = ln_spline + spline_model.log_parametric_model
+        psd_samples.append(np.exp(ln_psd))
+
+    return np.array(psd_samples)
+
+
+def _make_dataset_from_dict(data_dict, coords=None):
+    dataset_vars = {}
+    for k, v in data_dict.items():
+        if (
+            isinstance(v, tuple)
+            and len(v) == 2
+            and isinstance(v[0], (list, str))
+        ):
+            dims, data = v
+            dataset_vars[k] = DataArray(data, dims=dims)
+        else:
+            dataset_vars[k] = DataArray(v)
+    return Dataset(dataset_vars, coords=coords)
+
+
+def compare_runs(
+    run1: az.InferenceData,
+    run2: az.InferenceData,
+    labels: List[str],
+    outdir: str,
+) -> Dataset:
+    """
+    Compare two InferenceData runs and return a Dataset with differences.
+
+    Parameters
+    ----------
+    run1 : az.InferenceData
+        First run to compare
+    run2 : az.InferenceData
+        Second run to compare
+
+    Returns
+    -------
+    Dataset
+        Dataset containing the differences between the two runs
+    """
+    import matplotlib.pyplot as plt
+
+    os.makedirs(outdir, exist_ok=True)
+
+    # Ensure both runs have the same variables
+    common_vars = set(run1.posterior.data_vars) & set(run2.posterior.data_vars)
+    if not common_vars:
+        raise ValueError("No common variables found in the two runs.")
+
+    # Plot density
+    fig = az.plot_density(
+        [run1.posterior, run2.posterior],
+        data_labels=labels,
+        shade=0.2,
+        hdi_prob=0.94,
+    )
+    plt.suptitle("Density Comparison", fontsize=14)
+    plt.tight_layout()
+    plt.savefig(f"{outdir}/density_comparison.png")
+    plt.close()
+
+    # Get summaries
+    summary1 = az.summary(run1)
+    summary2 = az.summary(run2)
+
+    # Compute difference in summaries
+    common_vars = summary1.index.intersection(summary2.index)
+    diff = summary1.loc[common_vars] - summary2.loc[common_vars]
+    diff.to_csv(f"{outdir}/summary_diff.csv")
+
+    print("Summary Differences:")
+    print(diff)

log_psplines/datatypes.py

@@ -0,0 +1,83 @@
+import dataclasses
+
+import jax.numpy as jnp
+
+
+@dataclasses.dataclass
+class Timeseries:
+    t: jnp.ndarray
+    y: jnp.ndarray
+    std: float = 1.0
+
+    @property
+    def n(self):
+        return len(self.t)
+
+    @property
+    def fs(self) -> float:
+        """Sampling frequency computed from the time array."""
+        return float(1 / (self.t[1] - self.t[0]))
+
+    def to_periodogram(self) -> "Periodogram":
+        """Compute the one-sided periodogram of the timeseries."""
+        freq = jnp.fft.rfftfreq(len(self.y), d=1 / self.fs)
+        power = jnp.abs(jnp.fft.rfft(self.y)) ** 2 / len(self.y)
+        return Periodogram(freq[1:], power[1:])
+
+    def standardise(self):
+        """Standardise the timeseries to have zero mean and unit variance."""
+        self.std = float(jnp.std(self.y))
+        y = (self.y - jnp.mean(self.y)) / self.std
+        return Timeseries(self.t, y, self.std)
+
+    def __repr__(self):
+        return f"Timeseries(n={len(self.t)}, std={self.std:.3f}, fs={self.fs:.3f})"
+
+
+@dataclasses.dataclass
+class Periodogram:
+    freqs: jnp.ndarray
+    power: jnp.ndarray
+    filtered: bool = False
+
+    def __post_init__(self):
+        # assert no nans
+        if jnp.isnan(self.freqs).any() or jnp.isnan(self.power).any():
+
+            raise ValueError("Frequency or power contains NaN values.")
+
+    @property
+    def n(self):
+        return len(self.freqs)
+
+    @property
+    def fs(self) -> float:
+        """Sampling frequency computed from the frequency array."""
+        return float(2 * self.freqs[-1])
+
+    def highpass(self, min_freq: float) -> "Periodogram":
+        """Return a new Periodogram with frequencies above a threshold."""
+        mask = self.freqs > min_freq
+        return Periodogram(self.freqs[mask], self.power[mask], filtered=True)
+
+    def to_timeseries(self) -> "Timeseries":
+        """Compute the inverse FFT of the periodogram."""
+        y = jnp.fft.irfft(self.power, n=2 * (self.n - 1))
+        t = jnp.linspace(0, 1 / self.fs, len(y))
+        return Timeseries(t, y)
+
+    def __mul__(self, other):
+        return Periodogram(self.freqs, self.power * other)
+
+    def __truediv__(self, other):
+        return Periodogram(self.freqs, self.power / other)
+
+    def __repr__(self):
+        return f"Periodogram(n={self.n}, fs={self.fs:.3f}, filtered={self.filtered})"
+
+
+def compute_welsch_psd(
+    freqs: jnp.ndarray, power: jnp.ndarray, alpha: float = 2.0
+) -> jnp.ndarray:
+    """Compute the Welsch power spectral density of a periodogram."""
+    return power / (1 + (freqs / alpha) ** 2)

log_psplines/example_datasets/__init__.py

@@ -0,0 +1 @@
+from .ar_data import ARData

log_psplines/example_datasets/ar_data.py

@@ -0,0 +1,270 @@
+from typing import Optional, Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class ARData:
+    """
+    A class to simulate an AR(p) process (for p up to 4, or higher) and
+    compute its theoretical PSD as well as the raw periodogram.
+
+    Attributes
+    ----------
+    ar_coefs : np.ndarray
+        1D array of AR coefficients [a1, a2, ..., ap].
+    order : int
+        Order p of the AR process.
+    sigma : float
+        Standard deviation of the white‐noise driving the AR process.
+    fs : float
+        Sampling frequency [Hz].
+    duration : float
+        Total duration of the time series [s].
+    n : int
+        Number of samples = int(duration * fs).
+    seed : Optional[int]
+        Seed for the random number generator (if given).
+    ts : np.ndarray
+        Simulated time‐domain AR(p) series of length n.
+    freqs : np.ndarray
+        One‐sided frequency axis (length n//2 + 1).
+    psd_theoretical : np.ndarray
+        Theoretical one‐sided PSD (power per Hz) sampled at freqs.
+    periodogram : np.ndarray
+        One‐sided raw periodogram (power per Hz) from the simulated ts.
+    """
+
+    def __init__(
+        self,
+        order: int,
+        duration: float,
+        fs: float,
+        sigma: float = 1.0,
+        seed: Optional[int] = None,
+        ar_coefs: Sequence[float] = None,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        ar_coefs : Sequence[float]
+            Coefficients [a1, a2, ..., ap] for an AR(p) model.
+            For example, for AR(2) with x[t] = a1 x[t-1] + a2 x[t-2] + noise,
+            pass ar_coefs=[a1, a2].
+        duration : float
+            Total length of the time series in seconds.
+        fs : float
+            Sampling frequency in Hz.
+        sigma : float, default=1.0
+            Standard deviation of the white‐noise innovations.
+        seed : Optional[int], default=None
+            Seed for the random number generator (if you want reproducible draws).
+        """
+        self.order = order
+
+        if ar_coefs is None:
+            if order == 1:
+                ar_coefs = [0.9]
+            elif order == 2:
+                ar_coefs = [1.45, -0.9025]
+            elif order == 3:
+                ar_coefs = [0.9, -0.8, 0.7]
+            elif order == 4:
+                ar_coefs = [0.9, -0.8, 0.7, -0.6]
+            elif order == 5:
+                ar_coefs = [1, -2.2137, 2.9403, -2.1697, 0.9606]
+
+        else:
+            assert len(self.ar_coefs) == order
+
+        self.ar_coefs = np.asarray(ar_coefs, dtype=float)
+        self.order = len(self.ar_coefs)
+        self.sigma = float(sigma)
+        self.fs = float(fs)
+        self.duration = float(duration)
+        self.n = int(self.duration * self.fs)
+        self.seed = seed
+
+        # 1) Simulate the AR(p) time series
+        self.ts = self._generate_timeseries()
+
+        # 2) Build the one‐sided frequency axis
+        #    rfftfreq returns [0, 1, 2, ..., fs/2] with n//2 + 1 points
+        self.freqs = np.fft.rfftfreq(self.n, d=1.0 / self.fs)
+
+        # 3) Compute theoretical PSD on that frequency grid
+        self.psd_theoretical = self._compute_theoretical_psd()
+
+        # 4) Compute the one‐sided raw periodogram (power per Hz)
+        self.periodogram = self._compute_periodogram()
+
+    def _generate_timeseries(self) -> np.ndarray:
+        """
+        Generate an AR(p) time series of length n using the recursion
+
+            x[t] = a1*x[t-1] + a2*x[t-2] + ... + ap*x[t-p] + noise[t],
+
+        where noise[t] ~ Normal(0, sigma^2).  For t < 0, we assume x[t] = 0.
+
+        Returns
+        -------
+        ts : np.ndarray
+            Simulated AR(p) time series of length n.
+        """
+        rng = np.random.default_rng(self.seed)
+        x = np.zeros(self.n, dtype=float)
+        noise = rng.normal(loc=0.0, scale=self.sigma, size=self.n)
+
+        # Iterate from t = p .. n-1
+        for t in range(self.order, self.n):
+            past_terms = 0.0
+            # sum over a_k * x[t-k-1]
+            for k, a_k in enumerate(self.ar_coefs, start=1):
+                past_terms += a_k * x[t - k]
+            x[t] = past_terms + noise[t]
+
+        return x
+
+    def _compute_theoretical_psd(self) -> np.ndarray:
+        """
+        Compute the theoretical one‐sided PSD (power per Hz) of the AR(p) process:
+
+            S_theory(f) = (sigma^2 / fs) / | 1 - a1*e^{-i*2πf/fs} - a2*e^{-i*2πf*2/fs} - ... - ap*e^{-i*2πf*p/fs} |^2
+
+        evaluated at freqs = [0, 1, 2, ..., fs/2].
+
+        Returns
+        -------
+        psd_th : np.ndarray
+            One‐sided theoretical PSD of length n//2 + 1.
+        """
+        # digital‐frequency omega = 2π (f / fs)
+        omega = 2 * np.pi * self.freqs / self.fs
+
+        # Form the denominator polynomial: 1 - sum_{k=1}^p a_k e^{-i k omega}
+        # We compute numerator = sigma^2 / fs, denominator=|...|^2
+        denom = np.ones_like(omega, dtype=complex)
+        for k, a_k in enumerate(self.ar_coefs, start=1):
+            denom -= a_k * np.exp(-1j * k * omega)
+        denom_mag2 = np.abs(denom) ** 2
+
+        psd_th = (self.sigma**2 / self.fs) / denom_mag2
+        return psd_th.real  # should already be float
+
+    def _compute_periodogram(self) -> np.ndarray:
+        """
+        Compute the one‐sided raw periodogram of the simulated time series:
+
+            Pxx(f_k) = (1 / (n * fs)) * |H(f_k)|^2,
+            then double all bins except DC (k=0) and Nyquist (k=n/2) if n is even.
+
+        Returns
+        -------
+        pxx : np.ndarray
+            One‐sided periodogram (power per Hz) of length n//2 + 1.
+        """
+        # 1) Full FFT
+        H_full = np.fft.fft(self.ts)
+
+        # 2) Compute |H|^2 and normalize by (n * fs) → gives power per Hz
+        Pxx_full = (1.0 / (self.n * self.fs)) * np.abs(H_full) ** 2
+
+        # 3) Keep only the first (n//2 + 1) bins for real‐input one‐sided PSD
+        Pxx_one = Pxx_full[: self.n // 2 + 1]
+
+        # 4) Double all interior bins (1 .. n//2-1) to account for negative frequencies
+        if self.n % 2 == 0:
+            # n even → Nyquist is index n/2 and should NOT be doubled
+            Pxx_one[1:-1] *= 2.0
+        else:
+            # n odd → last index is floor(n/2), which is still not doubled
+            Pxx_one[1:] *= 2.0
+
+        return Pxx_one
+
+    def plot(
+        self,
+        ax: Optional[plt.Axes] = None,
+        *,
+        show_legend: bool = True,
+        periodogram_kwargs: Optional[dict] = None,
+        theoretical_kwargs: Optional[dict] = None,
+    ) -> plt.Axes:
+        """
+        Plot the one‐sided raw periodogram and the theoretical PSD
+        on the same axes (log–log).
+
+        Parameters
+        ----------
+        ax : Optional[plt.Axes]
+            If provided, plot onto this Axes object. Otherwise, create a new figure/axes.
+        show_legend : bool, default=True
+            Whether to display a legend.
+        periodogram_kwargs : Optional[dict], default=None
+            Additional kwargs to pass to plt.semilogy when plotting the periodogram.
+        theoretical_kwargs : Optional[dict], default=None
+            Additional kwargs to pass to plt.semilogy when plotting the theoretical PSD.
+
+        Returns
+        -------
+        ax : plt.Axes
+            The Axes object containing the plot.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 4))
+
+        # Default plotting styles
+        p_kwargs = {"label": "Raw Periodogram", "alpha": 0.6, "linewidth": 1.0}
+        t_kwargs = {
+            "label": "Theoretical PSD",
+            "linestyle": "--",
+            "color": "C1",
+            "linewidth": 2.0,
+        }
+
+        if periodogram_kwargs is not None:
+            p_kwargs.update(periodogram_kwargs)
+        if theoretical_kwargs is not None:
+            t_kwargs.update(theoretical_kwargs)
+
+        # Plot raw periodogram
+        ax.semilogy(self.freqs, self.periodogram, **p_kwargs)
+
+        # Plot theoretical PSD
+        ax.semilogy(self.freqs, self.psd_theoretical, **t_kwargs)
+
+        ax.set_xlabel("Frequency [Hz]")
+        ax.set_ylabel("PSD [power/Hz]")
+        ax.set_title(
+            f"AR({self.order}) Process: Periodogram vs Theoretical PSD"
+        )
+
+        if show_legend:
+            ax.legend()
+
+        ax.grid(True, which="both", ls=":", alpha=0.5)
+        return ax
+
+
+# Example usage:
+if __name__ == "__main__":
+    # --- Simulate AR(2) over 8 seconds at 1024 Hz ---
+    ar2 = ARData(
+        ar_coefs=[0.9, -0.5], duration=8.0, fs=1024.0, sigma=1.0, seed=42
+    )
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ar2.plot(ax=ax)
+    plt.show()
+
+    # --- Simulate AR(4) over 4 seconds at 2048 Hz ---
+    # e.g. coefficients [0.5, -0.3, 0.1, -0.05]
+    ar4 = ARData(
+        ar_coefs=[0.5, -0.3, 0.1, -0.05], duration=4.0, fs=2048.0, sigma=1.0
+    )
+    fig2, ax2 = plt.subplots(figsize=(8, 4))
+    ar4.plot(
+        ax=ax2,
+        periodogram_kwargs={"color": "C2"},
+        theoretical_kwargs={"color": "k", "linestyle": "-."},
+    )
+    plt.show()

log_psplines/example_datasets/lvk_data.py

@@ -0,0 +1,177 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.signal as signal
+from gwpy.timeseries import TimeSeries
+
+
+@dataclass
+class LVKData:
+    """
+    A dataclass for downloading, loading, and computing PSDs of gravitational-wave strain data.
+
+    Upon initialization, the PSDs for all overlapping segments are computed immediately.
+    """
+
+    strain: TimeSeries
+    duration: int
+    segment_duration: int
+    segment_overlap: float
+    min_freq: Optional[float] = None
+    max_freq: Optional[float] = None
+
+    # Fields computed in __post_init__; not passed at construction
+    fs: float = field(init=False)
+    n: int = field(init=False)
+    nperseg: int = field(init=False)
+    noverlap: int = field(init=False)
+    step: int = field(init=False)
+    n_segments: int = field(init=False)
+    freqs: np.ndarray = field(init=False)
+    psds: np.ndarray = field(init=False)
+    median_psd: np.ndarray = field(init=False)
+
+    def __post_init__(self):
+        # Sampling info
+        self.fs = float(self.strain.sample_rate.value)
+        self.n = len(self.strain)
+
+        # Number of samples per segment and overlap in samples
+        self.nperseg = int(self.fs * self.segment_duration)
+        self.noverlap = int(self.nperseg * self.segment_overlap)
+        self.step = self.nperseg - self.noverlap
+
+        # Compute number of segments
+        self.n_segments = (self.n - self.noverlap) // self.step
+
+        # Extract raw numpy array of strain values
+        data = self.strain.value
+
+        # Build strided array of shape (n_segments, nperseg)
+        shape = (self.n_segments, self.nperseg)
+        strides = (self.step * data.strides[-1], data.strides[-1])
+        segments = np.lib.stride_tricks.as_strided(
+            data, shape=shape, strides=strides
+        )
+
+        # Compute one-sided PSD for each segment
+        freqs_full, psd_full = signal.welch(
+            segments,
+            fs=self.fs,
+            nperseg=self.nperseg,
+            noverlap=self.noverlap,
+            axis=-1,
+            return_onesided=True,
+            scaling="density",
+        )
+
+        # Apply frequency mask if requested
+        freq_mask = np.ones_like(freqs_full, dtype=bool)
+        if self.min_freq is not None:
+            freq_mask &= freqs_full >= self.min_freq
+        if self.max_freq is not None:
+            freq_mask &= freqs_full <= self.max_freq
+
+        self.freqs = freqs_full[freq_mask]
+        self.psds = psd_full[:, freq_mask]
+        self.median_psd = np.median(self.psds, axis=0)
+
+    @classmethod
+    def download(
+        cls,
+        detector: str = "H1",
+        gps_start: int = 1126259462,
+        duration: int = 1024,
+        channel: Optional[str] = None,
+    ) -> TimeSeries:
+        """
+        Download open strain data from GWOSC for a given detector and GPS range.
+        """
+        print(
+            f"Downloading {detector} data [{gps_start} - {gps_start + duration}]..."
+        )
+        strain = TimeSeries.fetch_open_data(
+            detector, gps_start, gps_start + duration
+        )
+        if channel:
+            strain.channel = channel
+        return strain
+
+    @classmethod
+    def load(
+        cls,
+        detector: str = "H1",
+        gps_start: int = 1126259462,
+        duration: int = 1024,
+        segment_duration: int = 4,
+        segment_overlap: float = 0.5,
+        min_freq: Optional[float] = None,
+        max_freq: Optional[float] = None,
+        cache_file: str = "strain_cache.gwf",
+        channel: str = "H1:GWOSC-STRAIN",
+    ) -> "LVKData":
+        """
+        Load strain data from cache or download if needed, then compute PSDs.
+        """
+        if os.path.exists(cache_file):
+            try:
+                strain = TimeSeries.read(cache_file)
+                print(f"Loaded cached strain from '{cache_file}'")
+            except Exception as e:
+                print(
+                    f"Failed to read cache '{cache_file}': {e}. Redownloading..."
+                )
+                os.remove(cache_file)
+                strain = cls.download(
+                    detector, gps_start, duration, channel=channel
+                )
+                strain.write(cache_file)
+                print(f"Cached new strain to '{cache_file}'")
+        else:
+            strain = cls.download(
+                detector, gps_start, duration, channel=channel
+            )
+            strain.write(cache_file)
+            print(f"Cached strain to '{cache_file}'")
+
+        return cls(
+            strain=strain,
+            duration=duration,
+            segment_duration=segment_duration,
+            segment_overlap=segment_overlap,
+            min_freq=min_freq,
+            max_freq=max_freq,
+        )
+
+    def compute_median_psd(
+        self, n_segments: Optional[int] = None
+    ) -> np.ndarray:
+        """
+        Return the median PSD computed over the first `n_segments` segments.
+        """
+        if n_segments is None:
+            n_segments = self.n_segments
+        if n_segments > self.n_segments:
+            raise ValueError(
+                "n_segments exceeds available number of segments."
+            )
+        return np.median(self.psds[:n_segments, :], axis=0)
+
+    def plot_psd(self) -> plt.Figure:
+        """
+        Plot all individual-segment PSDs in gray and the median PSD in red.
+        """
+        fig, ax = plt.subplots(figsize=(8, 5))
+        ax.loglog(self.freqs, self.psds.T, color="gray", alpha=0.3)
+        ax.loglog(
+            self.freqs, self.median_psd, color="r", lw=2, label="Median PSD"
+        )
+        ax.set_xlabel("Frequency (Hz)")
+        ax.set_ylabel("PSD [strain^2/Hz]")
+        ax.set_title(f"PSD: {self.strain.channel}")
+        ax.grid(True, which="both", ls=":")
+        ax.legend()
+        return fig