Generates a smoothly varying standard normal time series.

kld_noise_generator.py

"""Generates a smoothly varying standard normal time series."""

import numpy as np
import scipy.linalg
import torch


class KLDNoiseGenerator(torch.nn.Module):
    """Generates a smoothly varying standard normal time series.

    See https://en.wikipedia.org/wiki/Langevin_dynamics, and

    A statistical method of estimation and simulation for systems of stochastic differential
    equations (Shoji and Ozaki, Biometrika (1998), 85, 1, pp. 240-243)

    Args:
        initial_L (torch.Tensor): Initial value of the time series.
        dt (float): Time step.
        gamma (float, optional): Damping factor. Defaults to 2.0 (critical damping).
        initial_V (torch.Tensor, optional): Initial value of the velocity.
    """

    def __init__(self, initial_L, dt, gamma=2.0, initial_V=None):
        super().__init__()
        initial_V = torch.randn_like(initial_L) if initial_V is None else initial_V
        self.register_buffer("state", torch.stack((initial_L, initial_V), dim=-1))
        A, B = self.compute_A_and_B(dt, gamma)
        self.register_buffer("A", A.to(self.state))
        self.register_buffer("B", B.to(self.state))

    @staticmethod
    def compute_A_and_B(dt, gamma):
        J = np.array([[0.0, 1.0], [-1.0, -gamma]])
        sigma_sigma_T = np.array([[0.0, 0.0], [0.0, 2 * gamma]])
        A = scipy.linalg.expm(J * dt)
        Q = A @ sigma_sigma_T @ A.T - sigma_sigma_T
        V = scipy.linalg.solve_continuous_lyapunov(J, Q)
        B = scipy.linalg.cholesky(V, lower=True)
        return torch.from_numpy(A), torch.from_numpy(B)

    def forward(self):
        L = self.state[..., 0].clone(memory_format=torch.contiguous_format)
        noise = torch.randn_like(self.state)
        self.state.copy_(self.state @ self.A.T + noise @ self.B.T)
        return L


def main():
    """Test and demo."""

    # Test that mean and std are ~1 after many steps.
    x = torch.randn(10000)
    std, mean = torch.std_mean(x)
    print(f"initial mean: {mean:.6f}, initial std: {std:.6f}")
    gen = KLDNoiseGenerator(x, dt=0.2, gamma=2.0)
    for _ in range(1000):
        x = gen()
    std, mean = torch.std_mean(x)
    print(f"final mean: {mean:.6f}, final std: {std:.6f}")

    # Demo
    import matplotlib.pyplot as plt

    x = torch.randn(4)
    xs = []
    gen = KLDNoiseGenerator(x, dt=0.01, gamma=2.0)
    for _ in range(1001):
        xs.append(gen())
    xs = torch.stack(xs)
    ts = torch.arange(xs.shape[0]) * 0.01
    plt.plot(ts, xs)
    plt.xlabel("time")
    plt.ylabel("value")
    plt.grid()
    plt.savefig("kld_noise.png")


if __name__ == "__main__":
    main()