17394129984

Committed 02 Sep 2025 05:25AM UTC coverage: 90.419% (+9.4%) from 81.047%

Build # 17394129984

Build Type

push

github

Committed by

avivajpeyi

Commit Message

run precommits

Run Details

169 of 180 branches covered (93.89%)

Branch coverage included in aggregate %.

145 of 166 new or added lines in 15 files covered. (87.35%)

62 existing lines in 11 files now uncovered.

1492 of 1657 relevant lines covered (90.04%)

1.8 hits per line

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

90.32

/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)[1:]
        self.times = np.arange(self.n) / self.fs
        self.psd_theoretical = self._compute_theoretical_psd()

        # convert to Timeseries and Periodogram datatypes
        self.ts = Timeseries(t=self.times, y=self.ts, std=self.sigma)
        self.periodogram = self.ts.to_periodogram()
        self.welch_psd = self._compute_welch_psd()

    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_welch_psd(self) -> Periodogram:
        nperseg = min(256, self.n // 4)  # Use a segment length of 256 or n//4
        freqs, Pxx = welch(
            self.ts.y,
            fs=self.fs,
            nperseg=nperseg,
            noverlap=0.5,
            scaling="density",
            detrend="constant",
        )
        return Periodogram(freqs, Pxx)

    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)[1:]	2✔
93	self.times = np.arange(self.n) / self.fs	2✔
94	self.psd_theoretical = self._compute_theoretical_psd()	2✔
95
96	# convert to Timeseries and Periodogram datatypes
97	self.ts = Timeseries(t=self.times, y=self.ts, std=self.sigma)	2✔
98	self.periodogram = self.ts.to_periodogram()	2✔
99	self.welch_psd = self._compute_welch_psd()	2✔
100
101	def __repr__(self):
102	return f"ARData(order={self.order}, n={self.n})"
103
104	def _generate_timeseries(self) -> np.ndarray:	2✔
105	"""
106	Generate an AR(p) time series of length n using the recursion
107
108	x[t] = a1x[t-1] + a2x[t-2] + ... + ap*x[t-p] + noise[t],
109
110	where noise[t] ~ Normal(0, sigma^2). For t < 0, we assume x[t] = 0.
111
112	Returns
113	-------
114	ts : np.ndarray
115	Simulated AR(p) time series of length n.
116	"""
117	n = self.n + 100 # Generate extra samples to avoid edge effects	2✔
118	rng = np.random.default_rng(self.seed)	2✔
119	x = np.zeros(n, dtype=float)	2✔
120	noise = rng.normal(loc=0.0, scale=self.sigma, size=n)	2✔
121
122	# Iterate from t = p .. n-1
123	for t in range(self.order, self.n):	2✔
124	past_terms = 0.0	2✔
125	# sum over a_k * x[t-k-1]
126	for k, a_k in enumerate(self.ar_coefs, start=1):	2✔
127	past_terms += a_k * x[t - k]	2✔
128	x[t] = past_terms + noise[t]	2✔
129
130	# Return only the last n samples
131	return x[self.order : self.order + self.n]	2✔
132
133	def _compute_theoretical_psd(self) -> np.ndarray:	2✔
134	"""
135	Compute the theoretical one‐sided PSD (power per Hz) of the AR(p) process:
136
137	S_theory(f) = (sigma^2 / fs) / \| 1 - a1e^{-i2πf/fs} - a2e^{-i2πf2/fs} - ... - ape^{-i2πfp/fs} \|^2
138
139	evaluated at freqs = [0, 1, 2, ..., fs/2].
140
141	Returns
142	-------
143	psd_th : np.ndarray
144	One‐sided theoretical PSD of length n//2 + 1.
145	"""
146	# digital‐frequency omega = 2π (f / fs)
147	omega = 2 * np.pi * self.freqs / self.fs	2✔
148
149	# Form the denominator polynomial: 1 - sum_{k=1}^p a_k e^{-i k omega}
150	# We compute numerator = sigma^2 / fs, denominator=\|...\|^2
151	denom = np.ones_like(omega, dtype=complex)	2✔
152	for k, a_k in enumerate(self.ar_coefs, start=1):	2✔
153	denom -= a_k * np.exp(-1j * k * omega)	2✔
154	denom_mag2 = np.abs(denom) ** 2	2✔
155
156	psd_th = (self.sigma**2 / self.fs) / denom_mag2	2✔
157	return psd_th.real * 2 # should already be float	2✔
158
159	def _compute_welch_psd(self) -> Periodogram:	2✔
160	nperseg = min(256, self.n // 4) # Use a segment length of 256 or n//4	2✔
161	freqs, Pxx = welch(	2✔
162	self.ts.y,
163	fs=self.fs,
164	nperseg=nperseg,
165	noverlap=0.5,
166	scaling="density",
167	detrend="constant",
168	)
169	return Periodogram(freqs, Pxx)	2✔
170
171	def plot(	2✔
172	self,
173	ax: Optional[plt.Axes] = None,
174	*,
175	show_legend: bool = True,
176	periodogram_kwargs: Optional[dict] = None,
177	theoretical_kwargs: Optional[dict] = None,
178	) -> plt.Axes:
179	"""
180	Plot the one‐sided raw periodogram and the theoretical PSD
181	on the same axes (log–log).
182
183	Parameters
184	----------
185	ax : Optional[plt.Axes]
186	If provided, plot onto this Axes object. Otherwise, create a new figure/axes.
187	show_legend : bool, default=True
188	Whether to display a legend.
189	periodogram_kwargs : Optional[dict], default=None
190	Additional kwargs to pass to plt.semilogy when plotting the periodogram.
191	theoretical_kwargs : Optional[dict], default=None
192	Additional kwargs to pass to plt.semilogy when plotting the theoretical PSD.
193
194	Returns
195	-------
196	ax : plt.Axes
197	The Axes object containing the plot.
198	"""
199	if ax is None:	2✔
UNCOV 200	fig, ax = plt.subplots(figsize=(8, 4))	×
201
202	# Default plotting styles
203	p_kwargs = {"label": "Raw Periodogram", "alpha": 0.6, "linewidth": 1.0}	2✔
204	t_kwargs = {	2✔
205	"label": "Theoretical PSD",
206	"linestyle": "--",
207	"color": "C1",
208	"linewidth": 2.0,
209	}
210
211	if periodogram_kwargs is not None:	2✔
UNCOV 212	p_kwargs.update(periodogram_kwargs)	×
213	if theoretical_kwargs is not None:	2✔
UNCOV 214	t_kwargs.update(theoretical_kwargs)	×
215
216	# Plot raw periodogram
217	ax.semilogy(self.freqs, self.periodogram.power, **p_kwargs)	2✔
218
219	# Plot theoretical PSD
220	ax.semilogy(self.freqs, self.psd_theoretical, **t_kwargs)	2✔
221
222	try:	2✔
223	f, Pxx = welch(	2✔
224	self.ts.y,
225	fs=self.fs,
226	nperseg=min(256, len(self.ts.y) // 4),
227	scaling="density",
228	)
229	ax.semilogy(f, Pxx, label="Welch PSD", color="C3", linestyle=":")	2✔
UNCOV 230	except Exception as e:	×
UNCOV 231	pass	×
232
233	ax.set_xlabel("Frequency [Hz]")	2✔
234	ax.set_ylabel("PSD [power/Hz]")	2✔
235	ax.set_title(	2✔
236	f"AR({self.order}) Process: Periodogram vs Theoretical PSD"
237	)
238
239	if show_legend:	2✔
240	ax.legend()	2✔
241
242	ax.grid(True, which="both", ls=":", alpha=0.5)	2✔
243	return ax	2✔
244
245
246	# Example usage:
247	if __name__ == "__main__":
248	# --- Simulate AR(2) over 8 seconds at 1024 Hz ---
249	ar2 = ARData(
250	ar_coefs=[0.9, -0.5], duration=8.0, fs=1024.0, sigma=1.0, seed=42
251	)
252	fig, ax = plt.subplots(figsize=(8, 4))
253	ar2.plot(ax=ax)
254	plt.show()
255
256	# --- Simulate AR(4) over 4 seconds at 2048 Hz ---
257	# e.g. coefficients [0.5, -0.3, 0.1, -0.05]
258	ar4 = ARData(
259	ar_coefs=[0.5, -0.3, 0.1, -0.05], duration=4.0, fs=2048.0, sigma=1.0
260	)
261	fig2, ax2 = plt.subplots(figsize=(8, 4))
262	ar4.plot(
263	ax=ax2,
264	periodogram_kwargs={"color": "C2"},
265	theoretical_kwargs={"color": "k", "linestyle": "-."},
266	)
267	plt.show()

nz-gravity / LogPSplinePSD / 17394129984

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