16407223521

Committed 21 Jul 2025 02:31AM UTC coverage: 80.997% (+3.1%) from 77.939%

Build # 16407223521

Build Type

push

github

Committed by

avivajpeyi

Commit Message

fix: add benchmarking fix for cli

Run Details

96 of 123 branches covered (78.05%)

Branch coverage included in aggregate %.

143 of 151 new or added lines in 4 files covered. (94.7%)

7 existing lines in 2 files now uncovered.

1106 of 1361 relevant lines covered (81.26%)

1.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.8

/src/log_psplines/example_datasets/ar_data.py

from typing import Optional, Sequence

import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import welch

from ..datatypes import Periodogram, Timeseries


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

        self.ts = self._generate_timeseries()
        self.freqs = np.fft.rfftfreq(self.n, d=1.0 / self.fs)
        self.times = np.arange(self.n) / self.fs
        self.psd_theoretical = self._compute_theoretical_psd()
        self.periodogram = self._compute_periodogram()

        # convert to Timeseries and Periodogram datatypes
        self.ts = Timeseries(t=self.times, y=self.ts, std=self.sigma)
        self.periodogram = Periodogram(
            freqs=self.freqs, power=self.periodogram, filtered=False
        )

    def __repr__(self):
        return f"ARData(order={self.order}, n={self.n})"

    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.
        """
        n = self.n + 100  # Generate extra samples to avoid edge effects
        rng = np.random.default_rng(self.seed)
        x = np.zeros(n, dtype=float)
        noise = rng.normal(loc=0.0, scale=self.sigma, size=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 only the last n samples
        return x[self.order : self.order + self.n]

    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 * 2  # 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.power, **p_kwargs)

        # Plot theoretical PSD
        ax.semilogy(self.freqs, self.psd_theoretical, **t_kwargs)

        try:
            f, Pxx = welch(
                self.ts.y,
                fs=self.fs,
                nperseg=min(256, len(self.ts.y) // 4),
                scaling="density",
            )
            ax.semilogy(f, Pxx, label="Welch PSD", color="C3", linestyle=":")
        except Exception as e:
            pass

        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()

1	from typing import Optional, Sequence	2✔
2
3	import matplotlib.pyplot as plt	2✔
4	import numpy as np	2✔
5	from scipy.signal import welch	2✔
6
7	from ..datatypes import Periodogram, Timeseries	2✔
8
9
10	class ARData:	2✔
11	"""
12	A class to simulate an AR(p) process (for p up to 4, or higher) and
13	compute its theoretical PSD as well as the raw periodogram.
14
15	Attributes
16	----------
17	ar_coefs : np.ndarray
18	1D array of AR coefficients [a1, a2, ..., ap].
19	order : int
20	Order p of the AR process.
21	sigma : float
22	Standard deviation of the white‐noise driving the AR process.
23	fs : float
24	Sampling frequency [Hz].
25	duration : float
26	Total duration of the time series [s].
27	n : int
28	Number of samples = int(duration * fs).
29	seed : Optional[int]
30	Seed for the random number generator (if given).
31	ts : np.ndarray
32	Simulated time‐domain AR(p) series of length n.
33	freqs : np.ndarray
34	One‐sided frequency axis (length n//2 + 1).
35	psd_theoretical : np.ndarray
36	Theoretical one‐sided PSD (power per Hz) sampled at freqs.
37	periodogram : np.ndarray
38	One‐sided raw periodogram (power per Hz) from the simulated ts.
39	"""
40
41	def __init__(	2✔
42	self,
43	order: int,
44	duration: float,
45	fs: float,
46	sigma: float = 1.0,
47	seed: Optional[int] = None,
48	ar_coefs: Sequence[float] = None,
49	) -> None:
50	"""
51	Parameters
52	----------
53	ar_coefs : Sequence[float]
54	Coefficients [a1, a2, ..., ap] for an AR(p) model.
55	For example, for AR(2) with x[t] = a1 x[t-1] + a2 x[t-2] + noise,
56	pass ar_coefs=[a1, a2].
57	duration : float
58	Total length of the time series in seconds.
59	fs : float
60	Sampling frequency in Hz.
61	sigma : float, default=1.0
62	Standard deviation of the white‐noise innovations.
63	seed : Optional[int], default=None
64	Seed for the random number generator (if you want reproducible draws).
65	"""
66	self.order = order	2✔
67
68	if ar_coefs is None:	2✔
69	if order == 1:	2✔
70	ar_coefs = [0.9]	2✔
71	elif order == 2:	2✔
72	ar_coefs = [1.45, -0.9025]	2✔
73	elif order == 3:	2✔
74	ar_coefs = [0.9, -0.8, 0.7]	2✔
75	elif order == 4:	2✔
76	ar_coefs = [0.9, -0.8, 0.7, -0.6]	2✔
77	elif order == 5:	×
78	ar_coefs = [1, -2.2137, 2.9403, -2.1697, 0.9606]	×
79
80	else:
81	assert len(self.ar_coefs) == order	×
82
83	self.ar_coefs = np.asarray(ar_coefs, dtype=float)	2✔
84	self.order = len(self.ar_coefs)	2✔
85	self.sigma = float(sigma)	2✔
86	self.fs = float(fs)	2✔
87	self.duration = float(duration)	2✔
88	self.n = int(self.duration * self.fs)	2✔
89	self.seed = seed	2✔
90
91	self.ts = self._generate_timeseries()	2✔
92	self.freqs = np.fft.rfftfreq(self.n, d=1.0 / self.fs)	2✔
93	self.times = np.arange(self.n) / self.fs	2✔
94	self.psd_theoretical = self._compute_theoretical_psd()	2✔
95	self.periodogram = self._compute_periodogram()	2✔
96
97	# convert to Timeseries and Periodogram datatypes
98	self.ts = Timeseries(t=self.times, y=self.ts, std=self.sigma)	2✔
99	self.periodogram = Periodogram(	2✔
100	freqs=self.freqs, power=self.periodogram, filtered=False
101	)
102
103	def __repr__(self):
104	return f"ARData(order={self.order}, n={self.n})"
105
106	def _generate_timeseries(self) -> np.ndarray:	2✔
107	"""
108	Generate an AR(p) time series of length n using the recursion
109
110	x[t] = a1x[t-1] + a2x[t-2] + ... + ap*x[t-p] + noise[t],
111
112	where noise[t] ~ Normal(0, sigma^2). For t < 0, we assume x[t] = 0.
113
114	Returns
115	-------
116	ts : np.ndarray
117	Simulated AR(p) time series of length n.
118	"""
119	n = self.n + 100 # Generate extra samples to avoid edge effects	2✔
120	rng = np.random.default_rng(self.seed)	2✔
121	x = np.zeros(n, dtype=float)	2✔
122	noise = rng.normal(loc=0.0, scale=self.sigma, size=n)	2✔
123
124	# Iterate from t = p .. n-1
125	for t in range(self.order, self.n):	2✔
126	past_terms = 0.0	2✔
127	# sum over a_k * x[t-k-1]
128	for k, a_k in enumerate(self.ar_coefs, start=1):	2✔
129	past_terms += a_k * x[t - k]	2✔
130	x[t] = past_terms + noise[t]	2✔
131
132	# Return only the last n samples
133	return x[self.order : self.order + self.n]	2✔
134
135	def _compute_theoretical_psd(self) -> np.ndarray:	2✔
136	"""
137	Compute the theoretical one‐sided PSD (power per Hz) of the AR(p) process:
138
139	S_theory(f) = (sigma^2 / fs) / \| 1 - a1e^{-i2πf/fs} - a2e^{-i2πf2/fs} - ... - ape^{-i2πfp/fs} \|^2
140
141	evaluated at freqs = [0, 1, 2, ..., fs/2].
142
143	Returns
144	-------
145	psd_th : np.ndarray
146	One‐sided theoretical PSD of length n//2 + 1.
147	"""
148	# digital‐frequency omega = 2π (f / fs)
149	omega = 2 * np.pi * self.freqs / self.fs	2✔
150
151	# Form the denominator polynomial: 1 - sum_{k=1}^p a_k e^{-i k omega}
152	# We compute numerator = sigma^2 / fs, denominator=\|...\|^2
153	denom = np.ones_like(omega, dtype=complex)	2✔
154	for k, a_k in enumerate(self.ar_coefs, start=1):	2✔
155	denom -= a_k * np.exp(-1j * k * omega)	2✔
156	denom_mag2 = np.abs(denom) ** 2	2✔
157
158	psd_th = (self.sigma**2 / self.fs) / denom_mag2	2✔
159	return psd_th.real * 2 # should already be float	2✔
160
161	def _compute_periodogram(self) -> np.ndarray:	2✔
162	"""
163	Compute the one‐sided raw periodogram of the simulated time series:
164
165	Pxx(f_k) = (1 / (n * fs)) * \|H(f_k)\|^2,
166	then double all bins except DC (k=0) and Nyquist (k=n/2) if n is even.
167
168	Returns
169	-------
170	pxx : np.ndarray
171	One‐sided periodogram (power per Hz) of length n//2 + 1.
172	"""
173	# 1) Full FFT
174	H_full = np.fft.fft(self.ts)	2✔
175
176	# 2) Compute \|H\|^2 and normalize by (n * fs) → gives power per Hz
177	Pxx_full = (1.0 / (self.n * self.fs)) * np.abs(H_full) ** 2	2✔
178
179	# 3) Keep only the first (n//2 + 1) bins for real‐input one‐sided PSD
180	Pxx_one = Pxx_full[: self.n // 2 + 1]	2✔
181
182	# 4) Double all interior bins (1 .. n//2-1) to account for negative frequencies
183	if self.n % 2 == 0:	2✔
184	# n even → Nyquist is index n/2 and should NOT be doubled
185	Pxx_one[1:-1] *= 2.0	2✔
186	else:
187	# n odd → last index is floor(n/2), which is still not doubled
UNCOV 188	Pxx_one[1:] *= 2.0	×
189
190	return Pxx_one	2✔
191
192	def plot(	2✔
193	self,
194	ax: Optional[plt.Axes] = None,
195	*,
196	show_legend: bool = True,
197	periodogram_kwargs: Optional[dict] = None,
198	theoretical_kwargs: Optional[dict] = None,
199	) -> plt.Axes:
200	"""
201	Plot the one‐sided raw periodogram and the theoretical PSD
202	on the same axes (log–log).
203
204	Parameters
205	----------
206	ax : Optional[plt.Axes]
207	If provided, plot onto this Axes object. Otherwise, create a new figure/axes.
208	show_legend : bool, default=True
209	Whether to display a legend.
210	periodogram_kwargs : Optional[dict], default=None
211	Additional kwargs to pass to plt.semilogy when plotting the periodogram.
212	theoretical_kwargs : Optional[dict], default=None
213	Additional kwargs to pass to plt.semilogy when plotting the theoretical PSD.
214
215	Returns
216	-------
217	ax : plt.Axes
218	The Axes object containing the plot.
219	"""
220	if ax is None:	2✔
UNCOV 221	fig, ax = plt.subplots(figsize=(8, 4))	×
222
223	# Default plotting styles
224	p_kwargs = {"label": "Raw Periodogram", "alpha": 0.6, "linewidth": 1.0}	2✔
225	t_kwargs = {	2✔
226	"label": "Theoretical PSD",
227	"linestyle": "--",
228	"color": "C1",
229	"linewidth": 2.0,
230	}
231
232	if periodogram_kwargs is not None:	2✔
UNCOV 233	p_kwargs.update(periodogram_kwargs)	×
234	if theoretical_kwargs is not None:	2✔
UNCOV 235	t_kwargs.update(theoretical_kwargs)	×
236
237	# Plot raw periodogram
238	ax.semilogy(self.freqs, self.periodogram.power, **p_kwargs)	2✔
239
240	# Plot theoretical PSD
241	ax.semilogy(self.freqs, self.psd_theoretical, **t_kwargs)	2✔
242
243	try:	2✔
244	f, Pxx = welch(	2✔
245	self.ts.y,
246	fs=self.fs,
247	nperseg=min(256, len(self.ts.y) // 4),
248	scaling="density",
249	)
250	ax.semilogy(f, Pxx, label="Welch PSD", color="C3", linestyle=":")	2✔
UNCOV 251	except Exception as e:	×
UNCOV 252	pass	×
253
254	ax.set_xlabel("Frequency [Hz]")	2✔
255	ax.set_ylabel("PSD [power/Hz]")	2✔
256	ax.set_title(	2✔
257	f"AR({self.order}) Process: Periodogram vs Theoretical PSD"
258	)
259
260	if show_legend:	2✔
261	ax.legend()	2✔
262
263	ax.grid(True, which="both", ls=":", alpha=0.5)	2✔
264	return ax	2✔
265
266
267	# Example usage:
268	if __name__ == "__main__":
269	# --- Simulate AR(2) over 8 seconds at 1024 Hz ---
270	ar2 = ARData(
271	ar_coefs=[0.9, -0.5], duration=8.0, fs=1024.0, sigma=1.0, seed=42
272	)
273	fig, ax = plt.subplots(figsize=(8, 4))
274	ar2.plot(ax=ax)
275	plt.show()
276
277	# --- Simulate AR(4) over 4 seconds at 2048 Hz ---
278	# e.g. coefficients [0.5, -0.3, 0.1, -0.05]
279	ar4 = ARData(
280	ar_coefs=[0.5, -0.3, 0.1, -0.05], duration=4.0, fs=2048.0, sigma=1.0
281	)
282	fig2, ax2 = plt.subplots(figsize=(8, 4))
283	ar4.plot(
284	ax=ax2,
285	periodogram_kwargs={"color": "C2"},
286	theoretical_kwargs={"color": "k", "linestyle": "-."},
287	)
288	plt.show()

nz-gravity / LogPSplinePSD / 16407223521

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous