• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

nz-gravity / LogPSplinePSD / 16407223521

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

push

github

avivajpeyi
fix: add benchmarking fix for cli

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
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] = a1*x[t-1] + a2*x[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 - a1*e^{-i*2πf/fs} - a2*e^{-i*2πf*2/fs} - ... - ap*e^{-i*2πf*p/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()
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc