Filtering Lorenz-63 With An Ensemble Kalman Filter

import time

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from jax import jit, random

from cuthbert import filter as run_filter
from cuthbert.enkf import ensemble_kalman_filter
from cuthbert.gaussian import taylor
from cuthbert.gaussian.types import LinearizedKalmanFilterState
from cuthbert.smc import particle_filter
from cuthbertlib.resampling import adaptive, systematic
from cuthbertlib.stats.multivariate_normal import logpdf
from cuthbertlib.types import LogConditionalDensity

jax.config.update("jax_enable_x64", True)

# Lorenz-63 parameters
lorenz_sigma = 10.0
lorenz_rho = 28.0
lorenz_beta = 8.0 / 3.0

# Discretization and noise
dt_inner = 0.01
n_inner_steps = 35
dt = dt_inner * n_inner_steps
diff_std = 1.0  # Diffusion noise standard deviation
obs_std = 0.3   # Observation noise standard deviation

x_dim = 3
y_dim = 1


def lorenz_drift(x):
    """Lorenz-63 drift function."""
    return jnp.array(
        [
            lorenz_sigma * (x[1] - x[0]),
            x[0] * (lorenz_rho - x[2]) - x[1],
            x[0] * x[1] - lorenz_beta * x[2],
        ]
    )


def lorenz_step(x):
    """Advance one assimilation interval with Heun integration."""

    def body(_, x_curr):
        k1 = lorenz_drift(x_curr)
        x_pred = x_curr + dt_inner * k1
        k2 = lorenz_drift(x_pred)
        return x_curr + 0.5 * dt_inner * (k1 + k2)

    return jax.lax.fori_loop(0, n_inner_steps, body, x)


# Noise covariances
Q = (diff_std**2 * dt) * jnp.eye(x_dim)
chol_Q = jnp.linalg.cholesky(Q)
R = (obs_std**2) * jnp.eye(y_dim)
chol_R = jnp.linalg.cholesky(R)

# Observation model: observe only x_1
H = jnp.array([[1.0, 0.0, 0.0]])
d_obs = jnp.zeros(y_dim)

# Simulate ground truth
num_time_steps = 200
key = random.key(0)

# Spin up to reach the attractor
x = jnp.array([1.0, 1.0, 1.0])
for _ in range(1_000):
    key, dyn_key = random.split(key)
    x = lorenz_step(x) + chol_Q @ random.normal(dyn_key, (x_dim,))

# Initial distribution centered near the attractor
m0 = x
P0 = 2.0 * jnp.eye(x_dim)
chol_P0 = jnp.linalg.cholesky(P0)

# Now simulate with observations
xs, ys = [], []
for t in range(num_time_steps):
    key, dyn_key, obs_key = random.split(key, 3)
    x = lorenz_step(x) + chol_Q @ random.normal(dyn_key, (x_dim,))
    y = H @ x + d_obs + chol_R @ random.normal(obs_key, (y_dim,))
    xs.append(x)
    ys.append(y)

true_states = jnp.stack(xs)
ys = jnp.stack(ys)
times = jnp.arange(num_time_steps) * dt

model_inputs = jnp.arange(num_time_steps + 1)

def get_init_log_density(model_inputs):
    def init_log_density(x):
        return logpdf(x, m0, chol_P0, nan_support=False)

    return init_log_density, m0


def get_dynamics_log_density(
    state: LinearizedKalmanFilterState, model_inputs
) -> tuple[LogConditionalDensity, jnp.ndarray, jnp.ndarray]:
    lin_point = state.mean

    def dynamics_log_density(x_prev, x):
        return logpdf(x, lorenz_step(x_prev), chol_Q, nan_support=False)

    return dynamics_log_density, lin_point, lorenz_step(lin_point)


def get_observation_func(
    state: LinearizedKalmanFilterState, model_inputs
) -> tuple[LogConditionalDensity, jnp.ndarray, jnp.ndarray]:
    idx = model_inputs - 1

    def obs_log_density(x, y):
        return logpdf(y, H @ x + d_obs, chol_R, nan_support=False)

    return obs_log_density, state.mean, ys[idx]


ekf = taylor.build_filter(
    get_init_log_density,
    get_dynamics_log_density,
    get_observation_func,
    associative=False,
)

jitted_filter = jit(run_filter, static_argnames=["filter_obj"])

n_timing = 20
ekf_states = jitted_filter(ekf, model_inputs)  # warm up
jax.block_until_ready(ekf_states)
_times = []
for _ in range(n_timing):
    _t0 = time.perf_counter()
    jax.block_until_ready(jitted_filter(ekf, model_inputs))
    _times.append(time.perf_counter() - _t0)
ekf_time = float(jnp.median(jnp.array(_times)))

ekf_means = ekf_states.mean
ekf_chol_covs = ekf_states.chol_cov

enkf = ensemble_kalman_filter.build_filter(
    init_sample=lambda key, mi: m0 + chol_P0 @ random.normal(key, m0.shape),
    get_dynamics=lambda mi: lambda x, key: lorenz_step(x)
    + chol_Q @ random.normal(key, (x_dim,)),
    get_observations=lambda mi: (lambda x: H @ x + d_obs, chol_R, ys[mi - 1]),
    n_particles=25,
    inflation=0.05,
    perturbed_obs=True,
)

key, enkf_key = random.split(key)
enkf_states = jitted_filter(enkf, model_inputs, key=enkf_key)  # warm up
jax.block_until_ready(enkf_states)
_times = []
for _ in range(n_timing):
    _t0 = time.perf_counter()
    jax.block_until_ready(jitted_filter(enkf, model_inputs, key=enkf_key))
    _times.append(time.perf_counter() - _t0)
enkf_time = float(jnp.median(jnp.array(_times)))

enkf_means = enkf_states.mean

adaptive_systematic = adaptive.ess_decorator(systematic.resampling, 0.5)

pf = particle_filter.build_filter(
    init_sample=lambda key, mi: m0 + chol_P0 @ random.normal(key, (x_dim,)),
    propagate_sample=lambda key, state, mi: lorenz_step(state)
    + chol_Q @ random.normal(key, (x_dim,)),
    log_potential=lambda s_prev, s, mi: logpdf(
        ys[mi - 1], H @ s + d_obs, chol_R, nan_support=False
    ),
    n_filter_particles=125,
    resampling_fn=adaptive_systematic,
)

key, pf_key = random.split(key)
pf_states = jitted_filter(pf, model_inputs, key=pf_key)  # warm up
jax.block_until_ready(pf_states)
_times = []
for _ in range(n_timing):
    _t0 = time.perf_counter()
    jax.block_until_ready(jitted_filter(pf, model_inputs, key=pf_key))
    _times.append(time.perf_counter() - _t0)
pf_time = float(jnp.median(jnp.array(_times)))

pf_weights = jax.nn.softmax(pf_states.log_weights, axis=-1)
pf_means = jnp.sum(pf_states.particles * pf_weights[..., None], axis=-2)