Part 5 (NumPyro-free): Filter optimization vs joint optimization¶

In the original Part 5, we used Stochastic Variational Inference (SVI) to compare two strategies:

• optimize over parameters while a filter marginalizes latent states, and
• optimize over parameters together with an explicit latent state path.

In this notebook, we mirror that comparison without NumPyro PPL. The Bayesian guide/posterior objects disappear, but the optimization structure remains:

• Filter optimization: minimize the negative marginal log-likelihood with Filter.
• Joint optimization: minimize the negative joint log density with dsx.log_prob, optimizing both rho and state_path_params directly.

Defining and Sampling From the Model¶

We return to defining and generating data from our running LTI example.

For data generation we use the pure-JAX dsx.simulate(...) entry point, so even the synthetic-data step is NumPyro-free.

In [1]:

import jax.numpy as jnp
import jax.random as jr
import numpyro
import numpyro.distributions as dist
import dynestyx as dsx
from dynestyx import DynamicalModel

# for convenience, we can define "fixed" things in the model outside of it.
# this is not required, but it helps keep the model clean.
state_dim = 2
observation_dim = 1
control_dim = 1

# Create the known matrices B, C
B = jnp.eye(state_dim, control_dim)
C = jnp.eye(observation_dim, state_dim)

# create the initial condition as a distribution
initial_condition = dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim))


def lti_model(
    sigma_obs=0.1,
    sigma_process=0.1,
    obs_times=None,
    obs_values=None,
    ctrl_times=None,
    ctrl_values=None,
    predict_times=None,
):
    # sample the unknown parameter
    rho = numpyro.sample("rho", dist.Uniform(-0.5, 0.5))
    A = jnp.array([[0, 0.3], [rho, -0.2]])

    # create the state evolution as a callable mapping to a distribution
    # Crucially, this depends on A, which depends on rho, which is unknown.
    # Thus, the state evolution MUST be defined within `lti_model`, not outside.
    state_evolution = lambda x, u, t_now, t_next: dist.MultivariateNormal(
        A @ x + B @ u, sigma_process**2 * jnp.eye(state_dim)
    )

    # create the observation model as a callable mapping to a distribution
    observation_model = lambda x, u, t: dist.MultivariateNormal(
        C @ x, sigma_obs**2 * jnp.eye(observation_dim)
    )

    # create the dynamical model
    dynamics = DynamicalModel(
        control_dim=control_dim,
        initial_condition=initial_condition,
        state_evolution=state_evolution,
        observation_model=observation_model,
    )

    # sample from the dynamical model
    return dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
        ctrl_times=ctrl_times,
        ctrl_values=ctrl_values,
        predict_times=predict_times,
    )

import jax.numpy as jnp import jax.random as jr import numpyro import numpyro.distributions as dist import dynestyx as dsx from dynestyx import DynamicalModel # for convenience, we can define "fixed" things in the model outside of it. # this is not required, but it helps keep the model clean. state_dim = 2 observation_dim = 1 control_dim = 1 # Create the known matrices B, C B = jnp.eye(state_dim, control_dim) C = jnp.eye(observation_dim, state_dim) # create the initial condition as a distribution initial_condition = dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)) def lti_model( sigma_obs=0.1, sigma_process=0.1, obs_times=None, obs_values=None, ctrl_times=None, ctrl_values=None, predict_times=None, ): # sample the unknown parameter rho = numpyro.sample("rho", dist.Uniform(-0.5, 0.5)) A = jnp.array([[0, 0.3], [rho, -0.2]]) # create the state evolution as a callable mapping to a distribution # Crucially, this depends on A, which depends on rho, which is unknown. # Thus, the state evolution MUST be defined within `lti_model`, not outside. state_evolution = lambda x, u, t_now, t_next: dist.MultivariateNormal( A @ x + B @ u, sigma_process**2 * jnp.eye(state_dim) ) # create the observation model as a callable mapping to a distribution observation_model = lambda x, u, t: dist.MultivariateNormal( C @ x, sigma_obs**2 * jnp.eye(observation_dim) ) # create the dynamical model dynamics = DynamicalModel( control_dim=control_dim, initial_condition=initial_condition, state_evolution=state_evolution, observation_model=observation_model, ) # sample from the dynamical model return dsx.sample( "f", dynamics, obs_times=obs_times, obs_values=obs_values, ctrl_times=ctrl_times, ctrl_values=ctrl_values, predict_times=predict_times, )

In [2]:

# create a synthetic control sequence as i.i.d. Gaussians
obs_times = jnp.arange(0.0, 100.0, 1.0)  # T=100 steps
ctrl_times = obs_times  # same times for controls
ctrl_values = jr.normal(jr.PRNGKey(0), (len(ctrl_times), control_dim))

rho_true = 0.3


def make_data(sigma_obs=0.1, sigma_process=0.1):
    A = jnp.array([[0, 0.3], [rho_true, -0.2]])
    # Use a structured discrete-time model so dsx.simulate(...) can auto-route
    # without a NumPyro handler stack.
    dynamics = dsx.LTI_discrete(
        A=A,
        Q=sigma_process**2 * jnp.eye(state_dim),
        H=C,
        R=sigma_obs**2 * jnp.eye(observation_dim),
        B=B,
        initial_mean=jnp.zeros(state_dim),
        initial_cov=jnp.eye(state_dim),
    )
    result = dsx.simulate(
        dynamics,
        rng_key=jr.PRNGKey(1),
        predict_times=obs_times,
        ctrl_times=ctrl_times,
        ctrl_values=ctrl_values,
    )
    print("make_data shapes:", result.times.shape, result.observations.shape)
    obs_values = result.observations[0]
    return obs_times, obs_values, ctrl_times, ctrl_values


obs_times, obs_values, ctrl_times, ctrl_values = make_data(sigma_obs=0.1, sigma_process=0.1)

# create a synthetic control sequence as i.i.d. Gaussians obs_times = jnp.arange(0.0, 100.0, 1.0) # T=100 steps ctrl_times = obs_times # same times for controls ctrl_values = jr.normal(jr.PRNGKey(0), (len(ctrl_times), control_dim)) rho_true = 0.3 def make_data(sigma_obs=0.1, sigma_process=0.1): A = jnp.array([[0, 0.3], [rho_true, -0.2]]) # Use a structured discrete-time model so dsx.simulate(...) can auto-route # without a NumPyro handler stack. dynamics = dsx.LTI_discrete( A=A, Q=sigma_process**2 * jnp.eye(state_dim), H=C, R=sigma_obs**2 * jnp.eye(observation_dim), B=B, initial_mean=jnp.zeros(state_dim), initial_cov=jnp.eye(state_dim), ) result = dsx.simulate( dynamics, rng_key=jr.PRNGKey(1), predict_times=obs_times, ctrl_times=ctrl_times, ctrl_values=ctrl_values, ) print("make_data shapes:", result.times.shape, result.observations.shape) obs_values = result.observations[0] return obs_times, obs_values, ctrl_times, ctrl_values obs_times, obs_values, ctrl_times, ctrl_values = make_data(sigma_obs=0.1, sigma_process=0.1)

make_data shapes: (1, 100) (1, 100, 1)

Shape convention note: dsx.simulate(...) returns arrays with a leading n_simulations axis (size 1 by default). In this notebook we index that axis explicitly via [0, ...] when extracting the single synthetic trajectory.

Define the dynamics, NumPyro-free¶

In the original notebook, the unknown parameter $\rho$ was a latent NumPyro sample site with a uniform prior. Without NumPyro there is no notion of a prior: $\rho$ becomes a plain argument, and the model is just a factory returning a DynamicalModel.

One further change from the original: the exact (cuthbert) Kalman filter we use below requires the model to be declared in structured linear-Gaussian form --- LinearGaussianStateEvolution / LinearGaussianObservation rather than anonymous callables --- so we build the same LTI system with the dsx.LTI_discrete factory (exactly as dynestyx's own tests do).

In [3]:

def make_dynamics(rho, sigma_obs=0.1, sigma_process=0.1):
    A = jnp.array([[0, 0.3], [rho, -0.2]])
    return dsx.LTI_discrete(
        A=A,
        Q=sigma_process**2 * jnp.eye(state_dim),
        H=C,
        R=sigma_obs**2 * jnp.eye(observation_dim),
        B=B,
        initial_mean=jnp.zeros(state_dim),
        initial_cov=jnp.eye(state_dim),
    )

def make_dynamics(rho, sigma_obs=0.1, sigma_process=0.1): A = jnp.array([[0, 0.3], [rho, -0.2]]) return dsx.LTI_discrete( A=A, Q=sigma_process**2 * jnp.eye(state_dim), H=C, R=sigma_obs**2 * jnp.eye(observation_dim), B=B, initial_mean=jnp.zeros(state_dim), initial_cov=jnp.eye(state_dim), )

Filter optimization and joint optimization¶

We now mirror the two-strand structure of the original SVI notebook, but with direct optax optimization instead of a NumPyro guide.

• In the filter strand, we optimize only rho, using the Kalman filter to marginalize latent states.
• In the joint strand, we optimize both rho and an explicit state_path_params array with dsx.log_prob.

For this linear-Gaussian example both objectives are differentiable, so plain gradient-based optimization works well.

In [4]:

import jax
import optax

from dynestyx import Filter
from dynestyx.inference.filters import KFConfig


def neg_loglik_filter(rho):
    with Filter(filter_config=KFConfig(filter_source="cuthbert")):
        result = dsx.condition(
            "f",
            make_dynamics(rho),
            obs_times=obs_times,
            obs_values=obs_values,
            ctrl_times=ctrl_times,
            ctrl_values=ctrl_values,
        )
    return -result.marginal_loglik


def neg_loglik_joint(params):
    return -dsx.log_prob(
        make_dynamics(params["rho"]),
        state_path_params=params["state_path_params"],
        state_path_param_times=obs_times,
        obs_times=obs_times,
        obs_values=obs_values,
        ctrl_times=ctrl_times,
        ctrl_values=ctrl_values,
    )


num_steps = 300
num_joint_steps = 1500
optimizer = optax.adam(learning_rate=1e-2)

filter_rho = jnp.array(0.0)
filter_opt_state = optimizer.init(filter_rho)
filter_loss_and_grad = jax.jit(jax.value_and_grad(neg_loglik_filter))

filter_loss_history, filter_rho_history = [], []
for _ in range(num_steps):
    loss, grad = filter_loss_and_grad(filter_rho)
    updates, filter_opt_state = optimizer.update(grad, filter_opt_state)
    filter_rho = optax.apply_updates(filter_rho, updates)
    filter_loss_history.append(loss)
    filter_rho_history.append(filter_rho)

filter_loss_history = jnp.stack(filter_loss_history)
filter_rho_history = jnp.stack(filter_rho_history)

joint_params = {
    "rho": jnp.array(0.0),
    "state_path_params": jnp.zeros((len(obs_times), state_dim)),
}
joint_opt_state = optimizer.init(joint_params)
joint_loss_and_grad = jax.jit(jax.value_and_grad(neg_loglik_joint))

joint_loss_history, joint_rho_history = [], []
for _ in range(num_joint_steps):
    loss, grad = joint_loss_and_grad(joint_params)
    updates, joint_opt_state = optimizer.update(grad, joint_opt_state)
    joint_params = optax.apply_updates(joint_params, updates)
    joint_loss_history.append(loss)
    joint_rho_history.append(joint_params["rho"])

joint_loss_history = jnp.stack(joint_loss_history)
joint_rho_history = jnp.stack(joint_rho_history)

print(f"Filter estimate: rho_hat = {filter_rho_history[-1]:.4f} (true value {rho_true})")
print(f"Joint estimate:  rho_hat = {joint_rho_history[-1]:.4f} (true value {rho_true})")

import jax import optax from dynestyx import Filter from dynestyx.inference.filters import KFConfig def neg_loglik_filter(rho): with Filter(filter_config=KFConfig(filter_source="cuthbert")): result = dsx.condition( "f", make_dynamics(rho), obs_times=obs_times, obs_values=obs_values, ctrl_times=ctrl_times, ctrl_values=ctrl_values, ) return -result.marginal_loglik def neg_loglik_joint(params): return -dsx.log_prob( make_dynamics(params["rho"]), state_path_params=params["state_path_params"], state_path_param_times=obs_times, obs_times=obs_times, obs_values=obs_values, ctrl_times=ctrl_times, ctrl_values=ctrl_values, ) num_steps = 300 num_joint_steps = 1500 optimizer = optax.adam(learning_rate=1e-2) filter_rho = jnp.array(0.0) filter_opt_state = optimizer.init(filter_rho) filter_loss_and_grad = jax.jit(jax.value_and_grad(neg_loglik_filter)) filter_loss_history, filter_rho_history = [], [] for _ in range(num_steps): loss, grad = filter_loss_and_grad(filter_rho) updates, filter_opt_state = optimizer.update(grad, filter_opt_state) filter_rho = optax.apply_updates(filter_rho, updates) filter_loss_history.append(loss) filter_rho_history.append(filter_rho) filter_loss_history = jnp.stack(filter_loss_history) filter_rho_history = jnp.stack(filter_rho_history) joint_params = { "rho": jnp.array(0.0), "state_path_params": jnp.zeros((len(obs_times), state_dim)), } joint_opt_state = optimizer.init(joint_params) joint_loss_and_grad = jax.jit(jax.value_and_grad(neg_loglik_joint)) joint_loss_history, joint_rho_history = [], [] for _ in range(num_joint_steps): loss, grad = joint_loss_and_grad(joint_params) updates, joint_opt_state = optimizer.update(grad, joint_opt_state) joint_params = optax.apply_updates(joint_params, updates) joint_loss_history.append(loss) joint_rho_history.append(joint_params["rho"]) joint_loss_history = jnp.stack(joint_loss_history) joint_rho_history = jnp.stack(joint_rho_history) print(f"Filter estimate: rho_hat = {filter_rho_history[-1]:.4f} (true value {rho_true})") print(f"Joint estimate: rho_hat = {joint_rho_history[-1]:.4f} (true value {rho_true})")

Filter estimate: rho_hat = 0.2466 (true value 0.3)
Joint estimate:  rho_hat = 0.2482 (true value 0.3)

How did the optimization do?¶

As in the original SVI notebook, we compare the two inference strands side by side. Here the comparison is between two optimization traces rather than two posterior distributions.

In [5]:

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(10, 3.5), constrained_layout=True)
axes[0].plot(filter_loss_history, color="C0", label="filter")
axes[0].plot(joint_loss_history, color="C1", label="joint")
axes[0].set_xlabel("optimization step")
axes[0].set_ylabel("loss")
axes[0].set_title("Optimization objective")
axes[0].legend()
axes[0].grid(True, alpha=0.3)
axes[1].plot(filter_rho_history, color="C0", label="filter")
axes[1].plot(joint_rho_history, color="C1", label="joint")
axes[1].axhline(rho_true, color="k", linestyle="--", label=r"$\rho_{\mathrm{true}}$")
axes[1].set_xlabel("optimization step")
axes[1].set_ylabel(r"$\rho$")
axes[1].set_title("Parameter estimate")
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.show()

import matplotlib.pyplot as plt fig, axes = plt.subplots(1, 2, figsize=(10, 3.5), constrained_layout=True) axes[0].plot(filter_loss_history, color="C0", label="filter") axes[0].plot(joint_loss_history, color="C1", label="joint") axes[0].set_xlabel("optimization step") axes[0].set_ylabel("loss") axes[0].set_title("Optimization objective") axes[0].legend() axes[0].grid(True, alpha=0.3) axes[1].plot(filter_rho_history, color="C0", label="filter") axes[1].plot(joint_rho_history, color="C1", label="joint") axes[1].axhline(rho_true, color="k", linestyle="--", label=r"$\rho_{\mathrm{true}}$") axes[1].set_xlabel("optimization step") axes[1].set_ylabel(r"$\rho$") axes[1].set_title("Parameter estimate") axes[1].legend() axes[1].grid(True, alpha=0.3) plt.show()

Latent-state recovery¶

As in the original notebook, we compare the latent-state outputs of the two inference strands.

• In the joint strand, the optimized state_path_params are themselves the fitted latent trajectory.
• In the filter strand, we rerun the filter at the optimized rho and read out the filtered state means and variances.

In [6]:

# Rerun the filter at the optimized rho and read the filtered state distributions.
with Filter(filter_config=KFConfig(filter_source="cuthbert")):
    result = dsx.condition(
        "f",
        make_dynamics(filter_rho_history[-1]),
        obs_times=obs_times,
        obs_values=obs_values,
        ctrl_times=ctrl_times,
        ctrl_values=ctrl_values,
    )

filtered_means = jnp.stack([d.mean for d in result.dists])  # (T, state_dim)
filtered_covs = jnp.stack([d.covariance_matrix for d in result.dists])  # (T, state_dim, state_dim)
filtered_stds = jnp.sqrt(jnp.diagonal(filtered_covs, axis1=1, axis2=2))  # (T, state_dim)

joint_state_path = joint_params["state_path_params"]

obs_times_np = jnp.asarray(obs_times)
obs_values_np = jnp.asarray(obs_values).squeeze()  # (T,)


def plot_latent_recovery(mean_states, lo, hi, title, obs_times, obs_values, state_dim=2):
    """mean/lo/hi: (T, state_dim). One figure, one subplot per state component."""
    n_comp = state_dim
    fig, axes = plt.subplots(
        n_comp, 1, figsize=(7, 2.5 * n_comp), sharex=True, constrained_layout=True
    )
    if n_comp == 1:
        axes = [axes]
    for i in range(n_comp):
        ax = axes[i]
        ax.fill_between(obs_times_np, lo[:, i], hi[:, i], alpha=0.3)
        ax.plot(obs_times_np, mean_states[:, i], label=f"$x_{i}$ (mean)")
        if i == 0:  # observed component: overlay data
            ax.scatter(
                obs_times_np,
                obs_values_np,
                s=8,
                alpha=0.7,
                color="k",
                label="observations",
                zorder=3,
            )
        ax.set_ylabel(f"$x_{i}$")
        ax.legend(loc="upper right", fontsize=8)
        ax.grid(True, alpha=0.3)
    axes[-1].set_xlabel("time")
    fig.suptitle(title)
    plt.show()


plot_latent_recovery(
    joint_state_path,
    joint_state_path,
    joint_state_path,
    "Latent state recovery — joint optimization",
    obs_times_np,
    obs_values_np,
    state_dim=state_dim,
)

plot_latent_recovery(
    filtered_means,
    filtered_means - 1.96 * filtered_stds,
    filtered_means + 1.96 * filtered_stds,
    "Latent state recovery — filter optimization",
    obs_times_np,
    obs_values_np,
    state_dim=state_dim,
)

# Rerun the filter at the optimized rho and read the filtered state distributions. with Filter(filter_config=KFConfig(filter_source="cuthbert")): result = dsx.condition( "f", make_dynamics(filter_rho_history[-1]), obs_times=obs_times, obs_values=obs_values, ctrl_times=ctrl_times, ctrl_values=ctrl_values, ) filtered_means = jnp.stack([d.mean for d in result.dists]) # (T, state_dim) filtered_covs = jnp.stack([d.covariance_matrix for d in result.dists]) # (T, state_dim, state_dim) filtered_stds = jnp.sqrt(jnp.diagonal(filtered_covs, axis1=1, axis2=2)) # (T, state_dim) joint_state_path = joint_params["state_path_params"] obs_times_np = jnp.asarray(obs_times) obs_values_np = jnp.asarray(obs_values).squeeze() # (T,) def plot_latent_recovery(mean_states, lo, hi, title, obs_times, obs_values, state_dim=2): """mean/lo/hi: (T, state_dim). One figure, one subplot per state component.""" n_comp = state_dim fig, axes = plt.subplots( n_comp, 1, figsize=(7, 2.5 * n_comp), sharex=True, constrained_layout=True ) if n_comp == 1: axes = [axes] for i in range(n_comp): ax = axes[i] ax.fill_between(obs_times_np, lo[:, i], hi[:, i], alpha=0.3) ax.plot(obs_times_np, mean_states[:, i], label=f"$x_{i}$ (mean)") if i == 0: # observed component: overlay data ax.scatter( obs_times_np, obs_values_np, s=8, alpha=0.7, color="k", label="observations", zorder=3, ) ax.set_ylabel(f"$x_{i}$") ax.legend(loc="upper right", fontsize=8) ax.grid(True, alpha=0.3) axes[-1].set_xlabel("time") fig.suptitle(title) plt.show() plot_latent_recovery( joint_state_path, joint_state_path, joint_state_path, "Latent state recovery — joint optimization", obs_times_np, obs_values_np, state_dim=state_dim, ) plot_latent_recovery( filtered_means, filtered_means - 1.96 * filtered_stds, filtered_means + 1.96 * filtered_stds, "Latent state recovery — filter optimization", obs_times_np, obs_values_np, state_dim=state_dim, )

This NumPyro-free workflow still mirrors the two-strand structure of the original SVI notebook: filter-based optimization versus explicit latent-path optimization. For full Bayesian posterior inference over these same two strands, see Part 4 (NUTS) and Part 5 (SVI).