Simulation Example

Simulation Example#

Here we demonstrate how to use this package to estimate the multivariate PSD of VAR(2) simulated data.

! pip install sgvb_psd -q 
import numpy as np
import matplotlib.pyplot as plt

from sgvb_psd.utils.sim_varma import SimVARMA

np.random.seed(0)

var2 = SimVARMA(
    n_samples=2 ** 13,
    var_coeffs=np.array(
        [[[0.5, 0.0], [0.0, -0.3]], [[0.0, 0.0], [0.0, -0.5]]]
    ),
    vma_coeffs=np.array([[[1.0, 0.0], [0.0, 1.0]]]),
    sigma=np.array([[1.0, 0.9], [0.9, 1.0]]),
)
var2

VAR(2) Data

from sgvb_psd.psd_estimator import PSDEstimator

optim = PSDEstimator(
    x=var2.data,
    N_theta=50,
    nchunks=4,
    ntrain_map=1000,
    max_hyperparm_eval=1,
    fs=2 * np.pi,
)
_ = optim.run(lr=0.003);
_ = optim.plot(
    true_psd=[var2.psd, var2.freq],
    off_symlog=False,
    xlims=[0, np.pi],
)
plt.savefig(f"var2_psd.png")
plt.close()

_ = optim.plot_coherence(true_psd=[var2.psd, var2.freq], labels="XY")
plt.savefig(f"var2_coh.png")
plt.close()

_ = optim.plot_vi_losses()
plt.savefig(f"var2_vi_losses.png")
plt.close()

VAR(2) PSD

VAR(2) Coherence

VAR(2) VI Losses