17844297625

Committed 18 Sep 2025 11:52PM UTC coverage: 89.605% (+0.4%) from 89.243%

Build # 17844297625

Build Type

push

github

Committed by

avivajpeyi

Commit Message

Add varma dataset

Run Details

200 of 211 branches covered (94.79%)

Branch coverage included in aggregate %.

82 of 86 new or added lines in 1 file covered. (95.35%)

1705 of 1915 relevant lines covered (89.03%)

1.78 hits per line

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

95.76

/src/log_psplines/example_datasets/varma_data.py

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import rfft
from typing import Optional


class VARMAData:
    """
    Simulate Vector Autoregressive Moving Average (VARMA) processes and compute related spectral properties.
    """

    def __init__(
            self,
            n_samples: int = 1024,
            sigma: np.ndarray = np.array([[1.0, 0.9], [0.9, 1.0]]),
            var_coeffs: np.ndarray = np.array([[[0.5, 0.0], [0.0, -0.3]], [[0.0, 0.0], [0.0, -0.5]]]),
            vma_coeffs: np.ndarray = np.array([[[1.0, 0.0], [0.0, 1.0]]]),
            seed: int = None,
    ):
        """
        Initialize the SimVARMA class.

        Args:
            n_samples (int): Number of samples to generate.
            var_coeffs (np.ndarray): VAR coefficient array.
            vma_coeffs (np.ndarray): VMA coefficient array.
            sigma (np.ndarray): Covariance matrix or scalar variance.
        """
        self.n_samples = n_samples
        self.var_coeffs = var_coeffs
        self.vma_coeffs = vma_coeffs
        self.sigma = sigma
        self.dim = vma_coeffs.shape[1]
        self.psd_scaling = 1
        self.n_freq_samples = n_samples // 2

        self.fs = 2 * np.pi
        self.freq = (
                np.linspace(0, 0.5, self.n_freq_samples, endpoint=False)[1:] * self.fs
        )
        self.data = None  # set in "resimulate"
        self.periodogram = None  # set in "resimulate"
        self.welch_psd = None  # set in "resimulate"
        self.welch_f = None  # set in "resimulate"
        self.resimulate(seed=seed)

        self.psd = _calculate_true_varma_psd(
            self.n_freq_samples,
            self.dim,
            self.var_coeffs,
            self.vma_coeffs,
            self.sigma,
        )

    def resimulate(self, seed=None):
        """
        Simulate VARMA process.

        Args:
            seed (int, optional): Random seed for reproducibility.

        Returns:
            np.ndarray: Simulated VARMA process data.
        """
        if seed is not None:
            np.random.seed(seed)

        lag_ma = self.vma_coeffs.shape[0]
        lag_ar = self.var_coeffs.shape[0]

        if self.sigma.shape[0] == 1:
            cov_matrix = np.identity(self.dim) * self.sigma
        else:
            cov_matrix = self.sigma

        x_init = np.zeros((lag_ar + 1, self.dim))
        x = np.empty((self.n_samples + 101, self.dim))
        x[:] = np.nan
        x[: lag_ar + 1] = x_init
        epsilon = np.random.multivariate_normal(
            np.zeros(self.dim), cov_matrix, size=[lag_ma]
        )

        for i in range(lag_ar + 1, x.shape[0]):
            epsilon = np.concatenate(
                [
                    np.random.multivariate_normal(
                        np.zeros(self.dim), cov_matrix, size=[1]
                    ),
                    epsilon[:-1],
                ]
            )
            x[i] = np.sum(
                np.matmul(
                    self.var_coeffs,
                    x[i - 1: i - lag_ar - 1: -1][..., np.newaxis],
                ),
                axis=(0, -1),
            ) + np.sum(
                np.matmul(self.vma_coeffs, epsilon[..., np.newaxis]),
                axis=(0, -1),
            )

        self.data = x[101:]

    def plot(self, axs=None, fname: Optional[str] = None):
        """
        Matrix plot: diagonal is the PSD, below diagonal is real CSD, above diagonal is imag CSD.
        Plots both the true PSD (from self.psd) and the periodogram (from self.data).
        """

        if self.data is None:
            raise ValueError("No data to plot. Run resimulate first.")
        n_freq = self.n_freq_samples - 1
        dim = self.dim
        # Compute periodogram and CSD
        data = self.data
        N = data.shape[0]
        # Remove mean
        data = data - np.mean(data, axis=0)
        # FFT
        fft_data = rfft(data, axis=0)[:n_freq]
        periodogram = np.empty((n_freq, dim, dim), dtype=np.complex128)
        for i in range(dim):
            for j in range(dim):
                periodogram[:, i, j] = fft_data[:, i] * np.conj(fft_data[:, j]) / N
        # True PSD (self.psd): shape (n_freq, dim, dim)
        true_psd = self.psd  # shape (n_freq, dim, dim)
        freq = self.freq  # shape (n_freq,)
        # Setup axes
        if axs is None:
            fig, axs = plt.subplots(dim, dim, figsize=(4 * dim, 4 * dim), sharex=True)
        else:
            fig = axs[0, 0].figure

        data_kwgs = dict(alpha=0.3, lw=2, zorder=-10, color='k')
        true_kwgs = dict(lw=1, zorder=10, color='k')
        for i in range(dim):
            for j in range(dim):
                ax = axs[i, j]
                if i == j:
                    # Diagonal: PSD
                    ax.plot(freq, true_psd[:, i, i].real, label="True PSD", **true_kwgs)
                    ax.plot(freq, periodogram[:, i, i].real, label="Periodogram", **data_kwgs)
                    ax.set_title(f"PSD: channel {i + 1}")
                elif i > j:
                    # Below diagonal: real CSD
                    ax.plot(freq, true_psd[:, i, j].real, label="True Re(CSD)",  **true_kwgs)
                    ax.plot(freq, periodogram[:, i, j].real, label="Periodogram Re(CSD)", **data_kwgs)
                    ax.set_title(f"Re(CSD): {i + 1},{j + 1}")
                else:
                    # Above diagonal: imag CSD
                    ax.plot(freq, true_psd[:, i, j].imag, label="True Im(CSD)", **true_kwgs)
                    ax.plot(freq, periodogram[:, i, j].imag, label="Periodogram Im(CSD)", **data_kwgs)
                    ax.set_title(f"Im(CSD): {i + 1},{j + 1}")
                if i == dim - 1:
                    ax.set_xlabel("Frequency (rad)")
                if j == 0:
                    ax.set_ylabel("Power / CSD")
                ax.legend(fontsize=8)
        fig.tight_layout()
        if fname:
            fig.savefig(fname, bbox_inches="tight")
        return axs


def _calculate_true_varma_psd(
        n_samples: int,
        dim: int,
        var_coeffs: np.ndarray,
        vma_coeffs: np.ndarray,
        sigma: np.ndarray,
) -> np.ndarray:
    """
    Calculate the spectral matrix for given frequencies.

    Args:
        n_samples int: Number of samples to generate for the true PSD (up to 0.5).
        var_coeffs (np.ndarray): VAR coefficient array.
        vma_coeffs (np.ndarray): VMA coefficient array.
        sigma (np.ndarray): Covariance matrix or scalar variance.

    Returns:
        np.ndarray: VARMA spectral matrix (PSD) for freq from 0 to 0.5.
    """
    freq = np.linspace(0, 0.5, n_samples, endpoint=False)[1:]
    spec_matrix = np.apply_along_axis(
        lambda f: _calculate_spec_matrix_helper(
            f, dim, var_coeffs, vma_coeffs, sigma
        ),
        axis=1,
        arr=freq.reshape(-1, 1),
    )
    return spec_matrix / (2 * np.pi)


def _calculate_spec_matrix_helper(f, dim, var_coeffs, vma_coeffs, sigma):
    """
    Helper function to calculate spectral matrix for a single frequency.

    Args:
        f (float): Single frequency value.
        var_coeffs (np.ndarray): VAR coefficient array.
        vma_coeffs (np.ndarray): VMA coefficient array.
        sigma (np.ndarray): Covariance matrix or scalar variance.

    Returns:
        np.ndarray: Calculated spectral matrix for the given frequency.
    """
    if sigma.shape[0] == 1:
        cov_matrix = np.identity(dim) * sigma
    else:
        cov_matrix = sigma

    k_ar = np.arange(1, var_coeffs.shape[0] + 1)
    A_f_re_ar = np.sum(
        var_coeffs * np.cos(np.pi * 2 * k_ar * f)[:, np.newaxis, np.newaxis],
        axis=0,
    )
    A_f_im_ar = -np.sum(
        var_coeffs * np.sin(np.pi * 2 * k_ar * f)[:, np.newaxis, np.newaxis],
        axis=0,
    )
    A_f_ar = A_f_re_ar + 1j * A_f_im_ar
    A_bar_f_ar = np.identity(dim) - A_f_ar
    H_f_ar = np.linalg.inv(A_bar_f_ar)

    k_ma = np.arange(vma_coeffs.shape[0])
    A_f_re_ma = np.sum(
        vma_coeffs * np.cos(np.pi * 2 * k_ma * f)[:, np.newaxis, np.newaxis],
        axis=0,
    )
    A_f_im_ma = -np.sum(
        vma_coeffs * np.sin(np.pi * 2 * k_ma * f)[:, np.newaxis, np.newaxis],
        axis=0,
    )
    A_f_ma = A_f_re_ma + 1j * A_f_im_ma
    A_bar_f_ma = A_f_ma
    H_f_ma = A_bar_f_ma

    spec_mat = H_f_ar @ H_f_ma @ cov_matrix @ H_f_ma.conj().T @ H_f_ar.conj().T
    return spec_mat

1	import matplotlib.pyplot as plt	2✔
2	import numpy as np
3	import matplotlib.pyplot as plt	2✔
4	from numpy.fft import rfft	2✔
5	from typing import Optional	2✔
6
7
8	class VARMAData:	2✔
9	"""
10	Simulate Vector Autoregressive Moving Average (VARMA) processes and compute related spectral properties.
11	"""
12
13	def __init__(	2✔
14	self,
15	n_samples: int = 1024,
16	sigma: np.ndarray = np.array([[1.0, 0.9], [0.9, 1.0]]),
17	var_coeffs: np.ndarray = np.array([[[0.5, 0.0], [0.0, -0.3]], [[0.0, 0.0], [0.0, -0.5]]]),
18	vma_coeffs: np.ndarray = np.array([[[1.0, 0.0], [0.0, 1.0]]]),
19	seed: int = None,
20	):
21	"""
22	Initialize the SimVARMA class.
23
24	Args:
25	n_samples (int): Number of samples to generate.
26	var_coeffs (np.ndarray): VAR coefficient array.
27	vma_coeffs (np.ndarray): VMA coefficient array.
28	sigma (np.ndarray): Covariance matrix or scalar variance.
29	"""
30	self.n_samples = n_samples
31	self.var_coeffs = var_coeffs
32	self.vma_coeffs = vma_coeffs	2✔
33	self.sigma = sigma	2✔
34	self.dim = vma_coeffs.shape[1]	2✔
35	self.psd_scaling = 1	2✔
36	self.n_freq_samples = n_samples // 2	2✔
37		2✔
38	self.fs = 2 * np.pi	2✔
39	self.freq = (
40	np.linspace(0, 0.5, self.n_freq_samples, endpoint=False)[1:] * self.fs	2✔
41	)	2✔
42	self.data = None # set in "resimulate"
43	self.periodogram = None # set in "resimulate"
44	self.welch_psd = None # set in "resimulate"
45	self.welch_f = None # set in "resimulate"	2✔
46	self.resimulate(seed=seed)	2✔
47		2✔
48	self.psd = _calculate_true_varma_psd(	2✔
49	self.n_freq_samples,	2✔
50	self.dim,
51	self.var_coeffs,	2✔
52	self.vma_coeffs,
53	self.sigma,
54	)
55
56	def resimulate(self, seed=None):
57	"""
58	Simulate VARMA process.
59		2✔
60	Args:
61	seed (int, optional): Random seed for reproducibility.
62
63	Returns:
64	np.ndarray: Simulated VARMA process data.
65	"""
66	if seed is not None:
67	np.random.seed(seed)
68
69	lag_ma = self.vma_coeffs.shape[0]	2✔
NEW 70	lag_ar = self.var_coeffs.shape[0]	×
71
72	if self.sigma.shape[0] == 1:	2✔
73	cov_matrix = np.identity(self.dim) * self.sigma	2✔
74	else:
75	cov_matrix = self.sigma	2✔
NEW 76		×
77	x_init = np.zeros((lag_ar + 1, self.dim))
78	x = np.empty((self.n_samples + 101, self.dim))	2✔
79	x[:] = np.nan
80	x[: lag_ar + 1] = x_init	2✔
81	epsilon = np.random.multivariate_normal(	2✔
82	np.zeros(self.dim), cov_matrix, size=[lag_ma]	2✔
83	)	2✔
84		2✔
85	for i in range(lag_ar + 1, x.shape[0]):
86	epsilon = np.concatenate(
87	[
88	np.random.multivariate_normal(	2✔
89	np.zeros(self.dim), cov_matrix, size=[1]	2✔
90	),
91	epsilon[:-1],
92	]
93	)
94	x[i] = np.sum(
95	np.matmul(
96	self.var_coeffs,
97	x[i - 1: i - lag_ar - 1: -1][..., np.newaxis],	2✔
98	),
99	axis=(0, -1),
100	) + np.sum(
101	np.matmul(self.vma_coeffs, epsilon[..., np.newaxis]),
102	axis=(0, -1),
103	)
104
105	self.data = x[101:]
106
107	def plot(self, axs=None, fname: Optional[str] = None):
108	"""	2✔
109	Matrix plot: diagonal is the PSD, below diagonal is real CSD, above diagonal is imag CSD.
110	Plots both the true PSD (from self.psd) and the periodogram (from self.data).	2✔
111	"""
112
113	if self.data is None:
114	raise ValueError("No data to plot. Run resimulate first.")
115	n_freq = self.n_freq_samples - 1
116	dim = self.dim	2✔
NEW 117	# Compute periodogram and CSD	×
118	data = self.data	2✔
119	N = data.shape[0]	2✔
120	# Remove mean
121	data = data - np.mean(data, axis=0)	2✔
122	# FFT	2✔
123	fft_data = rfft(data, axis=0)[:n_freq]
124	periodogram = np.empty((n_freq, dim, dim), dtype=np.complex128)	2✔
125	for i in range(dim):
126	for j in range(dim):	2✔
127	periodogram[:, i, j] = fft_data[:, i] * np.conj(fft_data[:, j]) / N	2✔
128	# True PSD (self.psd): shape (n_freq, dim, dim)	2✔
129	true_psd = self.psd # shape (n_freq, dim, dim)	2✔
130	freq = self.freq # shape (n_freq,)	2✔
131	# Setup axes
132	if axs is None:
133	fig, axs = plt.subplots(dim, dim, figsize=(4 * dim, 4 * dim), sharex=True)
134	else:	2✔
135	fig = axs[0, 0].figure	2✔
136
137	data_kwgs = dict(alpha=0.3, lw=2, zorder=-10, color='k')	2✔
138	true_kwgs = dict(lw=1, zorder=10, color='k')	2✔
139	for i in range(dim):
140	for j in range(dim):
141	ax = axs[i, j]
NEW 142	if i == j:	×
143	# Diagonal: PSD
144	ax.plot(freq, true_psd[:, i, i].real, label="True PSD", **true_kwgs)	2✔
145	ax.plot(freq, periodogram[:, i, i].real, label="Periodogram", **data_kwgs)	2✔
146	ax.set_title(f"PSD: channel {i + 1}")	2✔
147	elif i > j:	2✔
148	# Below diagonal: real CSD	2✔
149	ax.plot(freq, true_psd[:, i, j].real, label="True Re(CSD)", **true_kwgs)	2✔
150	ax.plot(freq, periodogram[:, i, j].real, label="Periodogram Re(CSD)", **data_kwgs)
151	ax.set_title(f"Re(CSD): {i + 1},{j + 1}")	2✔
152	else:
153	# Above diagonal: imag CSD
154	ax.plot(freq, true_psd[:, i, j].imag, label="True Im(CSD)", **true_kwgs)
155	ax.plot(freq, periodogram[:, i, j].imag, label="Periodogram Im(CSD)", **data_kwgs)
156	ax.set_title(f"Im(CSD): {i + 1},{j + 1}")
157	if i == dim - 1:	2✔
158	ax.set_xlabel("Frequency (rad)")
159	if j == 0:
160	ax.set_ylabel("Power / CSD")
161	ax.legend(fontsize=8)
162	fig.tight_layout()
163	if fname:	2✔
164	fig.savefig(fname, bbox_inches="tight")	2✔
165	return axs
166		2✔
167
168	def _calculate_true_varma_psd(
169	n_samples: int,
170	dim: int,
171	var_coeffs: np.ndarray,
172	vma_coeffs: np.ndarray,	2✔
173	sigma: np.ndarray,
174	) -> np.ndarray:
175	"""
176	Calculate the spectral matrix for given frequencies.
177
178	Args:	2✔
179	n_samples int: Number of samples to generate for the true PSD (up to 0.5).
180	var_coeffs (np.ndarray): VAR coefficient array.
181	vma_coeffs (np.ndarray): VMA coefficient array.	2✔
182	sigma (np.ndarray): Covariance matrix or scalar variance.
183
184	Returns:
185	np.ndarray: VARMA spectral matrix (PSD) for freq from 0 to 0.5.
186	"""
187	freq = np.linspace(0, 0.5, n_samples, endpoint=False)[1:]	2✔
188	spec_matrix = np.apply_along_axis(
189	lambda f: _calculate_spec_matrix_helper(
190	f, dim, var_coeffs, vma_coeffs, sigma
191	),
192	axis=1,
193	arr=freq.reshape(-1, 1),	2✔
194	)	2✔
195	return spec_matrix / (2 * np.pi)	2✔
196		2✔
197		2✔
198	def _calculate_spec_matrix_helper(f, dim, var_coeffs, vma_coeffs, sigma):	2✔
199	"""	2✔
200	Helper function to calculate spectral matrix for a single frequency.	2✔
201		2✔
202	Args:	2✔
203	f (float): Single frequency value.
204	var_coeffs (np.ndarray): VAR coefficient array.
205	vma_coeffs (np.ndarray): VMA coefficient array.	2✔
206	sigma (np.ndarray): Covariance matrix or scalar variance.
207
208	Returns:
209	np.ndarray: Calculated spectral matrix for the given frequency.
210	"""
211	if sigma.shape[0] == 1:
212	cov_matrix = np.identity(dim) * sigma
213	else:
214	cov_matrix = sigma
215
216	k_ar = np.arange(1, var_coeffs.shape[0] + 1)
217	A_f_re_ar = np.sum(
218	var_coeffs * np.cos(np.pi * 2 * k_ar * f)[:, np.newaxis, np.newaxis],
219	axis=0,
220	)
221	A_f_im_ar = -np.sum(
222	var_coeffs * np.sin(np.pi * 2 * k_ar * f)[:, np.newaxis, np.newaxis],
223	axis=0,
224	)	2✔
225	A_f_ar = A_f_re_ar + 1j * A_f_im_ar	2✔
226	A_bar_f_ar = np.identity(dim) - A_f_ar
227	H_f_ar = np.linalg.inv(A_bar_f_ar)
228
229	k_ma = np.arange(vma_coeffs.shape[0])
230	A_f_re_ma = np.sum(
231	vma_coeffs * np.cos(np.pi * 2 * k_ma * f)[:, np.newaxis, np.newaxis],
232	axis=0,	2✔
233	)
234	A_f_im_ma = -np.sum(
235	vma_coeffs * np.sin(np.pi * 2 * k_ma * f)[:, np.newaxis, np.newaxis],	2✔
236	axis=0,
237	)
238	A_f_ma = A_f_re_ma + 1j * A_f_im_ma
239	A_bar_f_ma = A_f_ma
240	H_f_ma = A_bar_f_ma
241
242	spec_mat = H_f_ar @ H_f_ma @ cov_matrix @ H_f_ma.conj().T @ H_f_ar.conj().T
243	return spec_mat

nz-gravity / LogPSplinePSD / 17844297625

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