[Python] How to compute the power spectral density of a vector-valued process without...

I'm simulating a 2D Ornstein-Uhlenbeck process (Langevin equation for velocity), and I'm interested in computing the power spectral density (PSD) of the vector-valued velocity process.

Following the Wiener–Khinchin theorem, I compute the velocity autocorrelation function (VACF) defined as:

$$ R(\tau) = \frac{1}{T - \tau} \int_0^{T - \tau} \langle \mathbf{v}(t) \cdot \mathbf{v}(t + \tau) \rangle , dt $$

Then I apply the Fast Fourier Transform (FFT) to obtain the power spectral density (PSD). However, since I only compute the VACF for positive lags $\tau \geq 0$, I need to mirror the VACF manually to make it symmetric before applying np.fft.fft.

That mirroring step feels like a workaround, and I'm wondering:

• Is there a correct or more natural way to compute the PSD of a vector-valued process without having to explicitly mirror the VACF?
• Or alternatively, is there a shortcut that computes the PSD directly without computing the VACF at all?

Below is my working example. I calculate the VACF as the average dot product over time for all trajectories, and then mirror it to compute the PSD via FFT.

1. Simulation of 2D Langevin trajectories

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

np.random.seed(0)
n = 10_000 # Number of time steps
dt = .1 # Time step size (delta t)
dim = 2 # Dimension of the velocity vector (2D: x and y)
N = 10 # Number of independent trajectories (ensemble size)
gamma = 1 # Damping coefficient in the Langevin equation
sigma = 1 # Noise strength

dW = np.random.normal(scale=sigma * np.sqrt(dt), size=(n, dim, N)) Gaussian noise (Wiener increments)
v = np.zeros((n, dim, N)) # Velocity - shape (time, dimension, trajectories)
r = np.zeros((n, dim, N)) # Position - shape (time, dimension, trajectories)

# Integrate the Langevin equation using the Euler-Maruyama method
for i in range(n - 1)
v[i + 1] = v - gamma * v * dt + dW # Update velocity: damping + noise
r[i + 1] = r + v * dt # Update position: Euler integration of velocity



2. VACF and PSD computation

# Computes the VACF: ⟨v(t) · v(t + τ)⟩
def manual_vacf(v, lag):
vacf = [] # VACF for each trajectory
for i in range(v.shape[-1]): # Loop over trajectories
v_ = v[..., i] # Select trajectory i (shape: [time, dim])
vacf_ = np.empty(lag) # VACF of this trajectory
for lag_ in range(1, lag + 1): # Loop over lag values
v1, v2 = v_[:-lag_], v_[lag_:] # Vectors at t and t+lag
# Remove mean and element wise product
v1v2 = (v1 - v1.mean(axis=0)) * (v2 - v2.mean(axis=0))
v1_dot_v2 = np.sum(v1v2, axis=1) # Dot product v(t) · v(t + τ)
vacf_[lag_ - 1] = np.mean(v1_dot_v2) # Average over all time points
vacf.append(vacf_) # Store VACF of this trajectory
return np.mean(vacf, axis=0) # Average VACF over all trajectories

# Mirrors the VACF and computes the PSD using the Wiener–Khinchin theorem
def psd_wiener_khinchin(vacf, lag, dt):
vacf_mirror = np.hstack([np.flip(vacf), vacf]) # Extend VACF to negative lags
ft = np.fft.fft(vacf_mirror)[:lag] # FFT and keep only non-negative frequencies
freq = np.fft.fftfreq(2 * lag, dt)[:lag] # Frequency values
psd = np.abs(ft) * dt # PSD (magnitude × dt)
return psd, freq

# === vacf ===
lag_time = 4
lag = int(lag_time / dt)
tau = dt * np.arange(lag)
vacf_manual = manual_vacf(v, lag)
tau_theory = dt * np.arange(-lag, lag)
vacf_theory = theoretical_vacf(tau_theory, gamma, sigma)

# === psd ===
psd_manual, freq_manual = psd_wiener_khinchin(vacf_manual, lag, dt)
psd_theory = theoretical_psd(freq_manual, gamma, sigma)


3. Plotting

fig, axs = plt.subplots(1, 3, figsize=(10, 2))

# first 3 simulated trajectories
for i in range(3):
axs[0].plot(*r[..., i].T)
axs[0].axis('equal')
axs[0].set_xlabel('x')
axs[0].set_ylabel('y')

# vacf
axs[1].plot(tau, vacf_manual)
axs[1].plot(tau_theory, vacf_theory)
axs[1].set_xlim(-lag_time, lag_time)
axs[1].set_xlabel('lag time')
axs[1].set_ylabel('VACF')

# psd
axs[2].plot(freq_manual, psd_manual, label='simulation')
axs[2].plot(freq_manual, psd_theory, label='theory')
axs[2].set_yscale('log')
axs[2].set_xscale('log')
axs[2].set_xlabel('frequency')
axs[2].set_ylabel('PSD')
axs[2].legend()

plt.show()




Continue reading...