Part 11b: Missing Observations with LatentPathBuilder + MCMC¶

Part 11 handled missing observations with Filter and Smoother, where the latent states are integrated out analytically or approximately.

In this sequel we keep the same discrete-time AR(1)-style latent system from Part 11, but switch to MCMC + LatentPathBuilder. That means we now sample the latent trajectory and the structural parameter jointly.

The most important conceptual point is how missingness is interpreted at the level of discrete-time model indexing.

1. Two missing-data cases in a generic discrete-time system¶

Consider a generic discrete-time state-space model

$$ x_0 \,\sim\, p(x_0 \mid \theta), \qquad x_k \,\sim\, p(x_k \mid x_{k-1}, \theta), \qquad y_k \,\sim\, p(y_k \mid x_k, \theta), \qquad k = 1, \dots, T. $$

Here obs_values[k] always refers to the observation attached to the latent index $k$. The array obs_times[k] is just the physical time label associated with that same index.

So if obs_values[k] contains NaN, it means that there was a missing measurement at index $k$, which corresponds to time obs_times[k].

That leaves two cases.

Case A: the whole observation row is missing¶

If every coordinate of $y_k$ is missing, then there is simply no observation factor at index $k$. If $O_k$ is the set of observed coordinates and $O_k = \varnothing$, then

$$ p(y_{k, O_k} \mid x_k, \theta) = 1. $$

So the latent-path formulation still keeps $x_k$ in the model, but it adds zero observation log-probability at that index.

This is why all observation models can handle full-row missingness under LatentPathBuilder: there is no need to marginalize anything inside the observation distribution at that step.

Case B: only part of the observation row is missing¶

If only some coordinates are observed, then LatentPathBuilder needs the marginal likelihood of the observed subset,

$$ p(y_{k, O_k} \mid x_k, \theta). $$

That is trickier. It requires some exploitable structure in the observation model. For example:

• if the observation is multivariate Gaussian, we can restrict to the observed subvector and covariance submatrix;
• if the observation factorizes coordinate-wise, we can sum the log-probabilities of the observed coordinates only;
• if the observation is exact identity, the observed coordinates are fixed directly and only the missing coordinates remain free latent variables;
• if none of those direct routes is available, we can introduce explicit latent variables for the missing observation coordinates and score the completed observation rows instead.

Without one of those structures, there is no single generic masked-likelihood formula.

What LatentPathBuilder can currently do¶

For latent-path inference in this repo:

• full-row missingness works for all observation models;
• partial missingness works when the observation model has supported structure, notably:
  • MultivariateNormal observations,
  • factorizable Independent(..., 1) observations,
  • DiracIdentityObservation, where state_path_params compress down to just the missing coordinates,
  • explicit missing-observation augmentation for continuous observation families, where LatentPathBuilder creates extra latents f_missing_obs_values and scores the completed observations.

We will illustrate those scopes below with four observation families:

1. multivariate Gaussian,
2. independent asymmetric Laplace observations,
3. exact identity observations with per-coordinate compression,
4. a correlated multivariate Student $t$ using explicit missing-observation augmentation.

In [1]:

import arviz as az
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive
from numpyro.infer.initialization import init_to_value

import dynestyx as dsx
from dynestyx import DiscreteTimeSimulator, Smoother
from dynestyx.inference.configs.smoother import KFSmootherConfig
from dynestyx.models import (
    DynamicalModel,
    LinearGaussianObservation,
    LinearGaussianStateEvolution,
)

plt.style.use("seaborn-v0_8-whitegrid")
np.set_printoptions(precision=3, suppress=True)

import arviz as az import jax.numpy as jnp import jax.random as jr import matplotlib.pyplot as plt import numpy as np import numpyro import numpyro.distributions as dist from numpyro.infer import MCMC, NUTS, Predictive from numpyro.infer.initialization import init_to_value import dynestyx as dsx from dynestyx import DiscreteTimeSimulator, Smoother from dynestyx.inference.configs.smoother import KFSmootherConfig from dynestyx.models import ( DynamicalModel, LinearGaussianObservation, LinearGaussianStateEvolution, ) plt.style.use("seaborn-v0_8-whitegrid") np.set_printoptions(precision=3, suppress=True)

In [2]:

obs_times = jnp.arange(0.0, 100.0, 1.0)
true_alpha = 0.4
state_dim = 2

block_start, block_end = 35, 60
partial_idx_dim0 = np.array([8, 18, 72, 90])
partial_idx_dim1 = np.array([25, 33, 84, 95])

transition_cov = jnp.array([[0.1, 0.01], [0.01, 0.15]])
transition_cov_non_gaussian = jnp.array([[0.1, 0.01], [0.01, 0.15]])
mvn_obs_cov = jnp.diag(jnp.array([0.25, 0.25]))
asym_laplace_scale = 0.25
asym_laplace_k = 1.7
student_df = 5.0
student_scale_tril = jnp.linalg.cholesky(jnp.array([[0.22, 0.08], [0.08, 0.18]]))


def full_only_mask(T, D):
    mask = np.ones((T, D), dtype=bool)
    mask[block_start:block_end, :] = False
    return mask


def full_plus_partial_mask(T, D):
    mask = full_only_mask(T, D)
    mask[partial_idx_dim0, 0] = False
    mask[partial_idx_dim1, 1] = False
    return mask


mask_full_only = full_only_mask(len(obs_times), state_dim)
mask_full_plus_partial = full_plus_partial_mask(len(obs_times), state_dim)


def apply_nan_mask(obs_values, mask):
    return jnp.where(jnp.asarray(mask), obs_values, jnp.nan)


def summarize_1d(samples):
    arr = np.asarray(samples)
    return {
        "mean": float(arr.mean()),
        "p05": float(np.percentile(arr, 5)),
        "p95": float(np.percentile(arr, 95)),
    }


def summarize_state_draws(state_draws):
    return {
        "mean": state_draws.mean(axis=0),
        "lo": np.percentile(state_draws, 5, axis=0),
        "hi": np.percentile(state_draws, 95, axis=0),
    }


def rmse(a, b):
    return float(np.sqrt(np.mean((np.asarray(a) - np.asarray(b)) ** 2)))


def summarize_gaussian_mixture(mean_draws, cov_diag_draws, rng_key, n_draws=400):
    mean_draws = np.asarray(mean_draws)
    cov_diag_draws = np.asarray(cov_diag_draws)
    S, T, D = mean_draws.shape
    flat_mean = mean_draws.reshape(S * T * D)
    flat_scale = np.sqrt(np.maximum(cov_diag_draws.reshape(S * T * D), 1e-8))

    keys = jr.split(rng_key, flat_mean.size)
    samples = jnp.stack(
        [
            dist.Normal(loc=m, scale=s).sample(k, sample_shape=(n_draws,))
            for m, s, k in zip(flat_mean, flat_scale, keys, strict=False)
        ],
        axis=1,
    )
    samples = np.asarray(samples).reshape(n_draws, S, T, D)
    samples = samples.reshape(n_draws * S, T, D)
    return {
        "mean": samples.mean(axis=0),
        "lo": np.percentile(samples, 5, axis=0),
        "hi": np.percentile(samples, 95, axis=0),
    }


def _shade_gap(ax, t, start, end):
    ax.axvspan(t[start], t[end - 1], color="orange", alpha=0.15, zorder=0, label="missing block")


def _draw_partial_flags(ax, t, indices, y_value):
    ax.scatter(
        t[indices],
        np.full(len(indices), y_value),
        s=55,
        marker="|",
        color="purple",
        linewidths=2,
        zorder=4,
        label="partial missing",
    )


def plot_missing_data(t, states, obs_values, *, partial_indices=None, title=""):
    fig, axes = plt.subplots(1, 2, figsize=(12, 3), sharey=False)
    for d, ax in enumerate(axes):
        _shade_gap(ax, t, block_start, block_end)
        ax.plot(t, states[:, d], color="C0", lw=1.0, alpha=0.6, label="true state")
        obs_d = np.asarray(obs_values[:, d])
        obs_mask = ~np.isnan(obs_d)
        ax.plot(t[obs_mask], obs_d[obs_mask], ".", ms=3, color="C1", alpha=0.8, label="observed")
        if partial_indices is not None and partial_indices[d].size > 0:
            ylo = np.nanmin(states[:, d]) - 0.5
            _draw_partial_flags(ax, t, partial_indices[d], ylo)
        ax.set_xlabel("time")
        ax.set_title(f"Dimension {d}")
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), fontsize=8)
    fig.suptitle(title, fontsize=12)
    plt.tight_layout()
    plt.show()


def plot_state_recovery(t, states_true, obs_values, summary_a, *, label_a, color_a, summary_b=None, label_b=None, color_b=None, partial_indices=None, title=""):
    fig, axes = plt.subplots(1, 2, figsize=(12, 3.5), sharey=False)
    for d, ax in enumerate(axes):
        _shade_gap(ax, t, block_start, block_end)
        ax.plot(t, states_true[:, d], color="black", lw=0.8, ls="--", alpha=0.5, label="true state")
        obs_d = np.asarray(obs_values[:, d])
        obs_mask = ~np.isnan(obs_d)
        ax.plot(t[obs_mask], obs_d[obs_mask], ".", ms=3, color="C1", alpha=0.8, label="observed")
        ax.fill_between(t, summary_a["lo"][:, d], summary_a["hi"][:, d], color=color_a, alpha=0.18, label=f"{label_a} 90% CI")
        ax.plot(t, summary_a["mean"][:, d], color=color_a, lw=1.2, label=f"{label_a} mean")
        if summary_b is not None:
            ax.fill_between(t, summary_b["lo"][:, d], summary_b["hi"][:, d], color=color_b, alpha=0.18, label=f"{label_b} 90% CI")
            ax.plot(t, summary_b["mean"][:, d], color=color_b, lw=1.2, ls="--", label=f"{label_b} mean")
        if partial_indices is not None and partial_indices[d].size > 0:
            ylo = np.nanmin(states_true[:, d]) - 0.5
            _draw_partial_flags(ax, t, partial_indices[d], ylo)
        ax.set_xlabel("time")
        ax.set_title(f"Dimension {d}")
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), fontsize=7)
    fig.suptitle(title, fontsize=12)
    plt.tight_layout()
    plt.show()

obs_times = jnp.arange(0.0, 100.0, 1.0) true_alpha = 0.4 state_dim = 2 block_start, block_end = 35, 60 partial_idx_dim0 = np.array([8, 18, 72, 90]) partial_idx_dim1 = np.array([25, 33, 84, 95]) transition_cov = jnp.array([[0.1, 0.01], [0.01, 0.15]]) transition_cov_non_gaussian = jnp.array([[0.1, 0.01], [0.01, 0.15]]) mvn_obs_cov = jnp.diag(jnp.array([0.25, 0.25])) asym_laplace_scale = 0.25 asym_laplace_k = 1.7 student_df = 5.0 student_scale_tril = jnp.linalg.cholesky(jnp.array([[0.22, 0.08], [0.08, 0.18]])) def full_only_mask(T, D): mask = np.ones((T, D), dtype=bool) mask[block_start:block_end, :] = False return mask def full_plus_partial_mask(T, D): mask = full_only_mask(T, D) mask[partial_idx_dim0, 0] = False mask[partial_idx_dim1, 1] = False return mask mask_full_only = full_only_mask(len(obs_times), state_dim) mask_full_plus_partial = full_plus_partial_mask(len(obs_times), state_dim) def apply_nan_mask(obs_values, mask): return jnp.where(jnp.asarray(mask), obs_values, jnp.nan) def summarize_1d(samples): arr = np.asarray(samples) return { "mean": float(arr.mean()), "p05": float(np.percentile(arr, 5)), "p95": float(np.percentile(arr, 95)), } def summarize_state_draws(state_draws): return { "mean": state_draws.mean(axis=0), "lo": np.percentile(state_draws, 5, axis=0), "hi": np.percentile(state_draws, 95, axis=0), } def rmse(a, b): return float(np.sqrt(np.mean((np.asarray(a) - np.asarray(b)) ** 2))) def summarize_gaussian_mixture(mean_draws, cov_diag_draws, rng_key, n_draws=400): mean_draws = np.asarray(mean_draws) cov_diag_draws = np.asarray(cov_diag_draws) S, T, D = mean_draws.shape flat_mean = mean_draws.reshape(S * T * D) flat_scale = np.sqrt(np.maximum(cov_diag_draws.reshape(S * T * D), 1e-8)) keys = jr.split(rng_key, flat_mean.size) samples = jnp.stack( [ dist.Normal(loc=m, scale=s).sample(k, sample_shape=(n_draws,)) for m, s, k in zip(flat_mean, flat_scale, keys, strict=False) ], axis=1, ) samples = np.asarray(samples).reshape(n_draws, S, T, D) samples = samples.reshape(n_draws * S, T, D) return { "mean": samples.mean(axis=0), "lo": np.percentile(samples, 5, axis=0), "hi": np.percentile(samples, 95, axis=0), } def _shade_gap(ax, t, start, end): ax.axvspan(t[start], t[end - 1], color="orange", alpha=0.15, zorder=0, label="missing block") def _draw_partial_flags(ax, t, indices, y_value): ax.scatter( t[indices], np.full(len(indices), y_value), s=55, marker="|", color="purple", linewidths=2, zorder=4, label="partial missing", ) def plot_missing_data(t, states, obs_values, *, partial_indices=None, title=""): fig, axes = plt.subplots(1, 2, figsize=(12, 3), sharey=False) for d, ax in enumerate(axes): _shade_gap(ax, t, block_start, block_end) ax.plot(t, states[:, d], color="C0", lw=1.0, alpha=0.6, label="true state") obs_d = np.asarray(obs_values[:, d]) obs_mask = ~np.isnan(obs_d) ax.plot(t[obs_mask], obs_d[obs_mask], ".", ms=3, color="C1", alpha=0.8, label="observed") if partial_indices is not None and partial_indices[d].size > 0: ylo = np.nanmin(states[:, d]) - 0.5 _draw_partial_flags(ax, t, partial_indices[d], ylo) ax.set_xlabel("time") ax.set_title(f"Dimension {d}") handles, labels = ax.get_legend_handles_labels() by_label = dict(zip(labels, handles)) ax.legend(by_label.values(), by_label.keys(), fontsize=8) fig.suptitle(title, fontsize=12) plt.tight_layout() plt.show() def plot_state_recovery(t, states_true, obs_values, summary_a, *, label_a, color_a, summary_b=None, label_b=None, color_b=None, partial_indices=None, title=""): fig, axes = plt.subplots(1, 2, figsize=(12, 3.5), sharey=False) for d, ax in enumerate(axes): _shade_gap(ax, t, block_start, block_end) ax.plot(t, states_true[:, d], color="black", lw=0.8, ls="--", alpha=0.5, label="true state") obs_d = np.asarray(obs_values[:, d]) obs_mask = ~np.isnan(obs_d) ax.plot(t[obs_mask], obs_d[obs_mask], ".", ms=3, color="C1", alpha=0.8, label="observed") ax.fill_between(t, summary_a["lo"][:, d], summary_a["hi"][:, d], color=color_a, alpha=0.18, label=f"{label_a} 90% CI") ax.plot(t, summary_a["mean"][:, d], color=color_a, lw=1.2, label=f"{label_a} mean") if summary_b is not None: ax.fill_between(t, summary_b["lo"][:, d], summary_b["hi"][:, d], color=color_b, alpha=0.18, label=f"{label_b} 90% CI") ax.plot(t, summary_b["mean"][:, d], color=color_b, lw=1.2, ls="--", label=f"{label_b} mean") if partial_indices is not None and partial_indices[d].size > 0: ylo = np.nanmin(states_true[:, d]) - 0.5 _draw_partial_flags(ax, t, partial_indices[d], ylo) ax.set_xlabel("time") ax.set_title(f"Dimension {d}") handles, labels = ax.get_legend_handles_labels() by_label = dict(zip(labels, handles)) ax.legend(by_label.values(), by_label.keys(), fontsize=7) fig.suptitle(title, fontsize=12) plt.tight_layout() plt.show()

2. Example 1: multivariate Gaussian observations¶

This first model is identical to Part 11.

The latent dynamics are the same 2D linear-Gaussian AR(1)-style system, and the observation model is the same linear Gaussian sensor model,

$$ y_k \mid x_k, \alpha \sim \mathcal N(x_k, R). $$

Because the observation distribution is multivariate Gaussian, partial missingness can be handled exactly by restricting to the observed subvector and covariance submatrix:

$$ y_{k, O_k} \mid x_k, \alpha \sim \mathcal N\big(x_{k, O_k}, R_{O_k O_k}\big). $$

So the simulator scores

$$ \log p(y_{k, O_k} \mid x_k, \alpha) = \log \mathcal N\big(y_{k, O_k}; x_{k, O_k}, R_{O_k O_k}\big). $$

Here we use a mask that combines:

• a contiguous fully missing block, and
• a few partially missing coordinates outside that block.

In [3]:

def ar1_model(obs_times=None, obs_values=None, predict_times=None):
    alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7))

    dynamics = DynamicalModel(
        initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)),
        state_evolution=LinearGaussianStateEvolution(
            A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]),
            cov=transition_cov,
        ),
        observation_model=LinearGaussianObservation(
            H=jnp.eye(state_dim),
            R=mvn_obs_cov,
        ),
        control_dim=0,
    )
    return dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
        predict_times=predict_times,
    )


data_key, mcmc_key_sim, mcmc_key_smooth, pp_key = jr.split(jr.PRNGKey(0), 4)

n_mcmc_warmup = 150
n_mcmc_samples = 150

with DiscreteTimeSimulator():
    synthetic = Predictive(
        ar1_model,
        params={"alpha": jnp.array(true_alpha)},
        num_samples=1,
        exclude_deterministic=False,
    )(data_key, predict_times=obs_times)

states_clean = np.asarray(synthetic["f_states"].squeeze((0, 1)))
obs_clean = np.asarray(synthetic["f_observations"].squeeze((0, 1)))
t = np.asarray(obs_times)

fig, axes = plt.subplots(1, 2, figsize=(12, 3), sharey=False)
for d, ax in enumerate(axes):
    ax.plot(t, states_clean[:, d], color="C0", lw=1.2, label=f"latent $x^{{({d})}}_t$")
    ax.plot(t, obs_clean[:, d], ".", ms=3, color="C1", alpha=0.6, label=f"obs $y^{{({d})}}_t$")
    ax.set_xlabel("time")
    ax.set_title(f"Dimension {d}")
    ax.legend(fontsize=8)
plt.tight_layout()
plt.show()

obs_mvn_missing = np.asarray(apply_nan_mask(jnp.asarray(obs_clean), mask_full_plus_partial))
plot_missing_data(
    t,
    states_clean,
    obs_mvn_missing,
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Gaussian observations: full block + partial missing coordinates",
)

def ar1_model(obs_times=None, obs_values=None, predict_times=None): alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7)) dynamics = DynamicalModel( initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)), state_evolution=LinearGaussianStateEvolution( A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]), cov=transition_cov, ), observation_model=LinearGaussianObservation( H=jnp.eye(state_dim), R=mvn_obs_cov, ), control_dim=0, ) return dsx.sample( "f", dynamics, obs_times=obs_times, obs_values=obs_values, predict_times=predict_times, ) data_key, mcmc_key_sim, mcmc_key_smooth, pp_key = jr.split(jr.PRNGKey(0), 4) n_mcmc_warmup = 150 n_mcmc_samples = 150 with DiscreteTimeSimulator(): synthetic = Predictive( ar1_model, params={"alpha": jnp.array(true_alpha)}, num_samples=1, exclude_deterministic=False, )(data_key, predict_times=obs_times) states_clean = np.asarray(synthetic["f_states"].squeeze((0, 1))) obs_clean = np.asarray(synthetic["f_observations"].squeeze((0, 1))) t = np.asarray(obs_times) fig, axes = plt.subplots(1, 2, figsize=(12, 3), sharey=False) for d, ax in enumerate(axes): ax.plot(t, states_clean[:, d], color="C0", lw=1.2, label=f"latent $x^{{({d})}}_t$") ax.plot(t, obs_clean[:, d], ".", ms=3, color="C1", alpha=0.6, label=f"obs $y^{{({d})}}_t$") ax.set_xlabel("time") ax.set_title(f"Dimension {d}") ax.legend(fontsize=8) plt.tight_layout() plt.show() obs_mvn_missing = np.asarray(apply_nan_mask(jnp.asarray(obs_clean), mask_full_plus_partial)) plot_missing_data( t, states_clean, obs_mvn_missing, partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Gaussian observations: full block + partial missing coordinates", )

In [4]:

def conditioned_latent_path(obs_times=None, obs_values=None):
    with dsx.LatentPathBuilder():
        ar1_model(obs_times=obs_times, obs_values=obs_values)


def conditioned_smoother(obs_times=None, obs_values=None):
    with Smoother(
        smoother_config=KFSmootherConfig(
            filter_source="cuthbert",
            record_smoothed_states_mean=True,
            record_smoothed_states_cov_diag=True,
        )
    ):
        ar1_model(obs_times=obs_times, obs_values=obs_values)


mcmc_latent_path = MCMC(
    NUTS(conditioned_latent_path),
    num_warmup=n_mcmc_warmup,
    num_samples=n_mcmc_samples,
    progress_bar=False,
)
mcmc_latent_path.run(mcmc_key_sim, obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing))

mcmc_smoother = MCMC(
    NUTS(conditioned_smoother),
    num_warmup=n_mcmc_warmup,
    num_samples=n_mcmc_samples,
    progress_bar=False,
)
mcmc_smoother.run(mcmc_key_smooth, obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing))

posterior_alpha_latent_path = np.asarray(mcmc_latent_path.get_samples()["alpha"])
posterior_alpha_smoother = np.asarray(mcmc_smoother.get_samples()["alpha"])

latent_path_state_samples = np.asarray(mcmc_latent_path.get_samples()["f_state_path"])
latent_path_summary = summarize_state_draws(latent_path_state_samples)

with Smoother(
    smoother_config=KFSmootherConfig(
        filter_source="cuthbert",
        record_smoothed_states_mean=True,
        record_smoothed_states_cov_diag=True,
    )
):
    pp_smoother = Predictive(
        ar1_model,
        posterior_samples=mcmc_smoother.get_samples(),
        exclude_deterministic=False,
    )(jr.fold_in(pp_key, 1), obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing))

smoothed_states = np.asarray(pp_smoother["f_smoothed_states_mean"])
smoothed_cov_diag = np.asarray(pp_smoother["f_smoothed_states_cov_diag"])
smoother_summary = summarize_gaussian_mixture(
    smoothed_states,
    smoothed_cov_diag,
    rng_key=jr.fold_in(pp_key, 11),
)

def conditioned_latent_path(obs_times=None, obs_values=None): with dsx.LatentPathBuilder(): ar1_model(obs_times=obs_times, obs_values=obs_values) def conditioned_smoother(obs_times=None, obs_values=None): with Smoother( smoother_config=KFSmootherConfig( filter_source="cuthbert", record_smoothed_states_mean=True, record_smoothed_states_cov_diag=True, ) ): ar1_model(obs_times=obs_times, obs_values=obs_values) mcmc_latent_path = MCMC( NUTS(conditioned_latent_path), num_warmup=n_mcmc_warmup, num_samples=n_mcmc_samples, progress_bar=False, ) mcmc_latent_path.run(mcmc_key_sim, obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing)) mcmc_smoother = MCMC( NUTS(conditioned_smoother), num_warmup=n_mcmc_warmup, num_samples=n_mcmc_samples, progress_bar=False, ) mcmc_smoother.run(mcmc_key_smooth, obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing)) posterior_alpha_latent_path = np.asarray(mcmc_latent_path.get_samples()["alpha"]) posterior_alpha_smoother = np.asarray(mcmc_smoother.get_samples()["alpha"]) latent_path_state_samples = np.asarray(mcmc_latent_path.get_samples()["f_state_path"]) latent_path_summary = summarize_state_draws(latent_path_state_samples) with Smoother( smoother_config=KFSmootherConfig( filter_source="cuthbert", record_smoothed_states_mean=True, record_smoothed_states_cov_diag=True, ) ): pp_smoother = Predictive( ar1_model, posterior_samples=mcmc_smoother.get_samples(), exclude_deterministic=False, )(jr.fold_in(pp_key, 1), obs_times=obs_times, obs_values=jnp.asarray(obs_mvn_missing)) smoothed_states = np.asarray(pp_smoother["f_smoothed_states_mean"]) smoothed_cov_diag = np.asarray(pp_smoother["f_smoothed_states_cov_diag"]) smoother_summary = summarize_gaussian_mixture( smoothed_states, smoothed_cov_diag, rng_key=jr.fold_in(pp_key, 11), )

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(10, 3), sharey=True)
az.plot_posterior(posterior_alpha_latent_path, hdi_prob=0.95, ref_val=true_alpha, ax=axes[0])
axes[0].set_title(r"Latent-path posterior of $\alpha$")
az.plot_posterior(posterior_alpha_smoother, hdi_prob=0.95, ref_val=true_alpha, ax=axes[1])
axes[1].set_title(r"KF smoother posterior of $\alpha$")
plt.tight_layout()
plt.show()

missing_rows_mvn = ~mask_full_plus_partial.all(axis=1)
print("LatentPathBuilder state RMSE:", rmse(latent_path_summary["mean"], states_clean))
print("KF smoother state RMSE:", rmse(smoother_summary["mean"], states_clean))
print("LatentPathBuilder missing-row RMSE:", rmse(latent_path_summary["mean"][missing_rows_mvn], states_clean[missing_rows_mvn]))
print("KF smoother missing-row RMSE:", rmse(smoother_summary["mean"][missing_rows_mvn], states_clean[missing_rows_mvn]))

plot_state_recovery(
    t,
    states_clean,
    obs_mvn_missing,
    latent_path_summary,
    label_a="latent path",
    color_a="C2",
    summary_b=smoother_summary,
    label_b="smoother",
    color_b="C3",
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Gaussian case: latent path vs. smoother state recovery",
)

fig, axes = plt.subplots(1, 2, figsize=(10, 3), sharey=True) az.plot_posterior(posterior_alpha_latent_path, hdi_prob=0.95, ref_val=true_alpha, ax=axes[0]) axes[0].set_title(r"Latent-path posterior of $\alpha$") az.plot_posterior(posterior_alpha_smoother, hdi_prob=0.95, ref_val=true_alpha, ax=axes[1]) axes[1].set_title(r"KF smoother posterior of $\alpha$") plt.tight_layout() plt.show() missing_rows_mvn = ~mask_full_plus_partial.all(axis=1) print("LatentPathBuilder state RMSE:", rmse(latent_path_summary["mean"], states_clean)) print("KF smoother state RMSE:", rmse(smoother_summary["mean"], states_clean)) print("LatentPathBuilder missing-row RMSE:", rmse(latent_path_summary["mean"][missing_rows_mvn], states_clean[missing_rows_mvn])) print("KF smoother missing-row RMSE:", rmse(smoother_summary["mean"][missing_rows_mvn], states_clean[missing_rows_mvn])) plot_state_recovery( t, states_clean, obs_mvn_missing, latent_path_summary, label_a="latent path", color_a="C2", summary_b=smoother_summary, label_b="smoother", color_b="C3", partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Gaussian case: latent path vs. smoother state recovery", )

LatentPathBuilder state RMSE: 0.3901934027671814
KF smoother state RMSE: 0.37186411023139954
LatentPathBuilder missing-row RMSE: 0.5262763500213623
KF smoother missing-row RMSE: 0.4925110340118408

3. Example 2: independent asymmetric Laplace observations¶

Now keep the same latent AR(1)-style discrete-time setup, but switch the observation family to a coordinate-wise asymmetric Laplace model.

If we want to perform Filtering/Smoothing as before, we would need to use a variant that can handle these non-Gaussian observations (e.g. PF/PS or EKF/EKS); however, as noted in the previous tutorial, these methods are not yet set up to automatically support missing data. The EnKF, while configured to automatically support missing data, is not set up to support non-Gaussian observations.

Nevertheless, the MCMC + LatentPathBuilder strategy can handle all forms of missingness (partial and full) under non-Gaussian observation models when the observation model is independent per observation dimension, which is the case here. Later we will see how explicit missing-observation augmentation extends this to correlated continuous observation families as well.

In [6]:

def asymmetric_laplace_model(obs_times=None, obs_values=None, predict_times=None):
    alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7))

    def observation_model(x, u, t):
        mean_shift = asym_laplace_scale * (1.0 / asym_laplace_k - asym_laplace_k)
        loc = x - mean_shift
        return dist.Independent(
            dist.AsymmetricLaplace(
                loc=loc,
                scale=asym_laplace_scale,
                asymmetry=asym_laplace_k,
            ),
            1,
        )

    dynamics = DynamicalModel(
        initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)),
        state_evolution=LinearGaussianStateEvolution(
            A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]),
            cov=transition_cov_non_gaussian,
        ),
        observation_model=observation_model,
        control_dim=0,
    )
    return dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
        predict_times=predict_times,
    )


with DiscreteTimeSimulator():
    asym_laplace_synth = Predictive(
        asymmetric_laplace_model,
        params={"alpha": jnp.array(true_alpha)},
        num_samples=1,
        exclude_deterministic=False,
    )(jr.PRNGKey(10), predict_times=obs_times)

asym_laplace_states_true = np.asarray(asym_laplace_synth["f_states"].squeeze((0, 1)))
asym_laplace_obs_clean = np.asarray(asym_laplace_synth["f_observations"].squeeze((0, 1))).astype(float)
asym_laplace_obs_missing = np.asarray(apply_nan_mask(jnp.asarray(asym_laplace_obs_clean), mask_full_plus_partial))

plot_missing_data(
    t,
    asym_laplace_states_true,
    asym_laplace_obs_missing,
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Independent asymmetric-Laplace observations: full block + partial missing coordinates",
)

def asymmetric_laplace_model(obs_times=None, obs_values=None, predict_times=None): alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7)) def observation_model(x, u, t): mean_shift = asym_laplace_scale * (1.0 / asym_laplace_k - asym_laplace_k) loc = x - mean_shift return dist.Independent( dist.AsymmetricLaplace( loc=loc, scale=asym_laplace_scale, asymmetry=asym_laplace_k, ), 1, ) dynamics = DynamicalModel( initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)), state_evolution=LinearGaussianStateEvolution( A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]), cov=transition_cov_non_gaussian, ), observation_model=observation_model, control_dim=0, ) return dsx.sample( "f", dynamics, obs_times=obs_times, obs_values=obs_values, predict_times=predict_times, ) with DiscreteTimeSimulator(): asym_laplace_synth = Predictive( asymmetric_laplace_model, params={"alpha": jnp.array(true_alpha)}, num_samples=1, exclude_deterministic=False, )(jr.PRNGKey(10), predict_times=obs_times) asym_laplace_states_true = np.asarray(asym_laplace_synth["f_states"].squeeze((0, 1))) asym_laplace_obs_clean = np.asarray(asym_laplace_synth["f_observations"].squeeze((0, 1))).astype(float) asym_laplace_obs_missing = np.asarray(apply_nan_mask(jnp.asarray(asym_laplace_obs_clean), mask_full_plus_partial)) plot_missing_data( t, asym_laplace_states_true, asym_laplace_obs_missing, partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Independent asymmetric-Laplace observations: full block + partial missing coordinates", )

In [7]:

def conditioned_asymmetric_laplace_latent_path(obs_times=None, obs_values=None):
    with dsx.LatentPathBuilder():
        asymmetric_laplace_model(obs_times=obs_times, obs_values=obs_values)


mcmc_asym_laplace = MCMC(
    NUTS(conditioned_asymmetric_laplace_latent_path),
    num_warmup=n_mcmc_warmup*10,
    num_samples=n_mcmc_samples*10,
    progress_bar=False,
)
mcmc_asym_laplace.run(jr.PRNGKey(11), obs_times=obs_times, obs_values=jnp.asarray(asym_laplace_obs_missing))

posterior_alpha_asym_laplace = np.asarray(mcmc_asym_laplace.get_samples()["alpha"])
asym_laplace_state_samples = np.asarray(mcmc_asym_laplace.get_samples()["f_state_path"])
asym_laplace_summary = summarize_state_draws(asym_laplace_state_samples)

def conditioned_asymmetric_laplace_latent_path(obs_times=None, obs_values=None): with dsx.LatentPathBuilder(): asymmetric_laplace_model(obs_times=obs_times, obs_values=obs_values) mcmc_asym_laplace = MCMC( NUTS(conditioned_asymmetric_laplace_latent_path), num_warmup=n_mcmc_warmup*10, num_samples=n_mcmc_samples*10, progress_bar=False, ) mcmc_asym_laplace.run(jr.PRNGKey(11), obs_times=obs_times, obs_values=jnp.asarray(asym_laplace_obs_missing)) posterior_alpha_asym_laplace = np.asarray(mcmc_asym_laplace.get_samples()["alpha"]) asym_laplace_state_samples = np.asarray(mcmc_asym_laplace.get_samples()["f_state_path"]) asym_laplace_summary = summarize_state_draws(asym_laplace_state_samples)

In [8]:

fig, ax = plt.subplots(1, 1, figsize=(5, 3))
az.plot_posterior(posterior_alpha_asym_laplace, hdi_prob=0.95, ref_val=true_alpha, ax=ax)
ax.set_title(r"Latent-path posterior of $\alpha$  (asymmetric-Laplace obs)")
plt.tight_layout()
plt.show()

print("Asymmetric-Laplace latent-path alpha:", summarize_1d(posterior_alpha_asym_laplace))
print("True alpha:", true_alpha)
print("Asymmetric-Laplace state RMSE:", rmse(asym_laplace_summary["mean"], asym_laplace_states_true))
print(
    "Asymmetric-Laplace missing-row RMSE:",
    rmse(
        asym_laplace_summary["mean"][~mask_full_plus_partial.all(axis=1)],
        asym_laplace_states_true[~mask_full_plus_partial.all(axis=1)],
    ),
)

plot_state_recovery(
    t,
    asym_laplace_states_true,
    asym_laplace_obs_missing,
    asym_laplace_summary,
    label_a="latent path",
    color_a="C2",
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Independent asymmetric-Laplace case: latent-path state recovery",
)

fig, ax = plt.subplots(1, 1, figsize=(5, 3)) az.plot_posterior(posterior_alpha_asym_laplace, hdi_prob=0.95, ref_val=true_alpha, ax=ax) ax.set_title(r"Latent-path posterior of $\alpha$ (asymmetric-Laplace obs)") plt.tight_layout() plt.show() print("Asymmetric-Laplace latent-path alpha:", summarize_1d(posterior_alpha_asym_laplace)) print("True alpha:", true_alpha) print("Asymmetric-Laplace state RMSE:", rmse(asym_laplace_summary["mean"], asym_laplace_states_true)) print( "Asymmetric-Laplace missing-row RMSE:", rmse( asym_laplace_summary["mean"][~mask_full_plus_partial.all(axis=1)], asym_laplace_states_true[~mask_full_plus_partial.all(axis=1)], ), ) plot_state_recovery( t, asym_laplace_states_true, asym_laplace_obs_missing, asym_laplace_summary, label_a="latent path", color_a="C2", partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Independent asymmetric-Laplace case: latent-path state recovery", )

Asymmetric-Laplace latent-path alpha: {'mean': 0.41981014609336853, 'p05': 0.13752658665180206, 'p95': 0.641951858997345}
True alpha: 0.4
Asymmetric-Laplace state RMSE: 0.28957778215408325
Asymmetric-Laplace missing-row RMSE: 0.3523925542831421

4. Example 3: exact identity observations with per-coordinate state latents¶

Now suppose the observation model is exact identity:

$$ y_k = x_k. $$

With partial missingness, the observed coordinates determine those state entries directly. The only free latent variables are the missing coordinates,

$$ z = \{x_{k,d} : (k,d) \text{ is missing}\}. $$

So LatentPathBuilder can represent state_path_params using exactly those missing coordinates, while f_state_path remains the full reconstructed trajectory.

This layout is now prepared automatically from obs_times and obs_values; there is no extra metadata object to pass into the builder.

In [9]:

def make_dirac_dynamics(alpha):
    return DynamicalModel(
        initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)),
        state_evolution=LinearGaussianStateEvolution(
            A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]),
            cov=transition_cov,
        ),
        observation_model=dsx.DiracIdentityObservation(),
        control_dim=0,
    )


small_obs_times = jnp.arange(10.0)
dirac_missing_pairs = np.array([[1, 0], [2, 1], [4, 0], [7, 0], [7, 1]])
dirac_mask = np.ones((len(small_obs_times), state_dim), dtype=bool)
dirac_mask[dirac_missing_pairs[:, 0], dirac_missing_pairs[:, 1]] = False
dirac_partial_indices = {
    d: dirac_missing_pairs[dirac_missing_pairs[:, 1] == d, 0] for d in range(state_dim)
}

dirac_sim = dsx.simulate(
    make_dirac_dynamics(true_alpha),
    rng_key=jr.PRNGKey(30),
    predict_times=small_obs_times,
)
dirac_states_true = np.asarray(dirac_sim.states[0])
dirac_obs_missing = np.asarray(
    apply_nan_mask(jnp.asarray(dirac_states_true), dirac_mask)
)
dirac_t = np.asarray(small_obs_times)


def dirac_model(obs_times=None, obs_values=None):
    alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7))
    return dsx.sample(
        "f",
        make_dirac_dynamics(alpha),
        obs_times=obs_times,
        obs_values=obs_values,
    )


def conditioned_dirac_latent_path(obs_times=None, obs_values=None):
    with dsx.LatentPathBuilder():
        dirac_model(obs_times=obs_times, obs_values=obs_values)


mcmc_dirac = MCMC(
    NUTS(conditioned_dirac_latent_path),
    num_warmup=n_mcmc_warmup,
    num_samples=n_mcmc_samples*10,
    progress_bar=False,
)
mcmc_dirac.run(
    jr.PRNGKey(31),
    obs_times=small_obs_times,
    obs_values=jnp.asarray(dirac_obs_missing),
)
dirac_samples = mcmc_dirac.get_samples()
posterior_alpha_dirac = np.asarray(dirac_samples["alpha"])

pp_dirac = Predictive(
    conditioned_dirac_latent_path,
    posterior_samples=dirac_samples,
    return_sites=[
        "f_state_path",
        "f_state_path_param_times",
        "f_state_path_param_coordinate_indices",
    ],
    exclude_deterministic=False,
)(
    jr.PRNGKey(32),
    obs_times=small_obs_times,
    obs_values=jnp.asarray(dirac_obs_missing),
)
dirac_state_samples = np.asarray(pp_dirac["f_state_path"])
dirac_summary = summarize_state_draws(dirac_state_samples)

print("Dirac latent-path alpha:", summarize_1d(posterior_alpha_dirac))
print("True alpha:", true_alpha)
print("Free state_path_params per draw:", dirac_samples["f_state_path_params"].shape[-1])
print("Dense state coordinates:", dirac_states_true.size)
print("state_path_param_times:", np.asarray(pp_dirac["f_state_path_param_times"][0]))
print(
    "state_path_param_coordinate_indices:",
    np.asarray(pp_dirac["f_state_path_param_coordinate_indices"][0]),
)

missing_rows = dirac_missing_pairs[:, 0]
missing_cols = dirac_missing_pairs[:, 1]
print(
    np.column_stack(
        [
            missing_rows,
            missing_cols,
            dirac_states_true[missing_rows, missing_cols],
            dirac_summary["mean"][missing_rows, missing_cols],
        ]
    )
)

fig, axes = plt.subplots(1, 2, figsize=(10, 3.5), sharey=False)
for d, ax in enumerate(axes):
    ax.plot(dirac_t, dirac_states_true[:, d], color="black", lw=0.9, ls="--", alpha=0.6, label="true state")
    obs_d = np.asarray(dirac_obs_missing[:, d])
    obs_mask_d = ~np.isnan(obs_d)
    ax.plot(dirac_t[obs_mask_d], obs_d[obs_mask_d], ".", ms=5, color="C1", alpha=0.85, label="observed")
    ax.fill_between(
        dirac_t,
        dirac_summary["lo"][:, d],
        dirac_summary["hi"][:, d],
        color="C2",
        alpha=0.18,
        label="latent path 90% CI",
    )
    ax.plot(dirac_t, dirac_summary["mean"][:, d], color="C2", lw=1.3, label="latent path mean")
    if dirac_partial_indices[d].size > 0:
        ax.scatter(
            dirac_t[dirac_partial_indices[d]],
            np.full(dirac_partial_indices[d].shape[0], np.nanmin(dirac_states_true[:, d]) - 0.25),
            s=60,
            marker="|",
            color="purple",
            linewidths=2,
            zorder=4,
            label="missing coordinate",
        )
    ax.set_xlabel("time")
    ax.set_title(f"Dimension {d}")
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys(), fontsize=7)

fig.suptitle("Dirac identity observations: only missing coordinates become free latents", fontsize=12)
plt.tight_layout()
plt.show()

def make_dirac_dynamics(alpha): return DynamicalModel( initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)), state_evolution=LinearGaussianStateEvolution( A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]), cov=transition_cov, ), observation_model=dsx.DiracIdentityObservation(), control_dim=0, ) small_obs_times = jnp.arange(10.0) dirac_missing_pairs = np.array([[1, 0], [2, 1], [4, 0], [7, 0], [7, 1]]) dirac_mask = np.ones((len(small_obs_times), state_dim), dtype=bool) dirac_mask[dirac_missing_pairs[:, 0], dirac_missing_pairs[:, 1]] = False dirac_partial_indices = { d: dirac_missing_pairs[dirac_missing_pairs[:, 1] == d, 0] for d in range(state_dim) } dirac_sim = dsx.simulate( make_dirac_dynamics(true_alpha), rng_key=jr.PRNGKey(30), predict_times=small_obs_times, ) dirac_states_true = np.asarray(dirac_sim.states[0]) dirac_obs_missing = np.asarray( apply_nan_mask(jnp.asarray(dirac_states_true), dirac_mask) ) dirac_t = np.asarray(small_obs_times) def dirac_model(obs_times=None, obs_values=None): alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7)) return dsx.sample( "f", make_dirac_dynamics(alpha), obs_times=obs_times, obs_values=obs_values, ) def conditioned_dirac_latent_path(obs_times=None, obs_values=None): with dsx.LatentPathBuilder(): dirac_model(obs_times=obs_times, obs_values=obs_values) mcmc_dirac = MCMC( NUTS(conditioned_dirac_latent_path), num_warmup=n_mcmc_warmup, num_samples=n_mcmc_samples*10, progress_bar=False, ) mcmc_dirac.run( jr.PRNGKey(31), obs_times=small_obs_times, obs_values=jnp.asarray(dirac_obs_missing), ) dirac_samples = mcmc_dirac.get_samples() posterior_alpha_dirac = np.asarray(dirac_samples["alpha"]) pp_dirac = Predictive( conditioned_dirac_latent_path, posterior_samples=dirac_samples, return_sites=[ "f_state_path", "f_state_path_param_times", "f_state_path_param_coordinate_indices", ], exclude_deterministic=False, )( jr.PRNGKey(32), obs_times=small_obs_times, obs_values=jnp.asarray(dirac_obs_missing), ) dirac_state_samples = np.asarray(pp_dirac["f_state_path"]) dirac_summary = summarize_state_draws(dirac_state_samples) print("Dirac latent-path alpha:", summarize_1d(posterior_alpha_dirac)) print("True alpha:", true_alpha) print("Free state_path_params per draw:", dirac_samples["f_state_path_params"].shape[-1]) print("Dense state coordinates:", dirac_states_true.size) print("state_path_param_times:", np.asarray(pp_dirac["f_state_path_param_times"][0])) print( "state_path_param_coordinate_indices:", np.asarray(pp_dirac["f_state_path_param_coordinate_indices"][0]), ) missing_rows = dirac_missing_pairs[:, 0] missing_cols = dirac_missing_pairs[:, 1] print( np.column_stack( [ missing_rows, missing_cols, dirac_states_true[missing_rows, missing_cols], dirac_summary["mean"][missing_rows, missing_cols], ] ) ) fig, axes = plt.subplots(1, 2, figsize=(10, 3.5), sharey=False) for d, ax in enumerate(axes): ax.plot(dirac_t, dirac_states_true[:, d], color="black", lw=0.9, ls="--", alpha=0.6, label="true state") obs_d = np.asarray(dirac_obs_missing[:, d]) obs_mask_d = ~np.isnan(obs_d) ax.plot(dirac_t[obs_mask_d], obs_d[obs_mask_d], ".", ms=5, color="C1", alpha=0.85, label="observed") ax.fill_between( dirac_t, dirac_summary["lo"][:, d], dirac_summary["hi"][:, d], color="C2", alpha=0.18, label="latent path 90% CI", ) ax.plot(dirac_t, dirac_summary["mean"][:, d], color="C2", lw=1.3, label="latent path mean") if dirac_partial_indices[d].size > 0: ax.scatter( dirac_t[dirac_partial_indices[d]], np.full(dirac_partial_indices[d].shape[0], np.nanmin(dirac_states_true[:, d]) - 0.25), s=60, marker="|", color="purple", linewidths=2, zorder=4, label="missing coordinate", ) ax.set_xlabel("time") ax.set_title(f"Dimension {d}") handles, labels = ax.get_legend_handles_labels() by_label = dict(zip(labels, handles)) ax.legend(by_label.values(), by_label.keys(), fontsize=7) fig.suptitle("Dirac identity observations: only missing coordinates become free latents", fontsize=12) plt.tight_layout() plt.show()

Dirac latent-path alpha: {'mean': 0.016201939433813095, 'p05': -0.387282133102417, 'p95': 0.4468417465686798}
True alpha: 0.4
Free state_path_params per draw: 5
Dense state coordinates: 20
state_path_param_times: [1. 2. 4. 7. 7.]
state_path_param_coordinate_indices: [0 1 0 0 1]
[[ 1.     0.     0.297 -0.192]
 [ 2.     1.    -0.636 -0.862]
 [ 4.     0.    -0.495 -0.117]
 [ 7.     0.     0.395  0.074]
 [ 7.     1.     0.775  0.449]]

5. Example 4: a correlated multivariate Student $t$¶

Finally, consider a correlated non-Gaussian observation family: a multivariate Student $t$.

For this model there is no simple generic masked-likelihood formula for

$$ p(y_{k, O_k} \mid x_k, \alpha), $$

because the observed coordinates remain coupled to the missing ones.

The fallback is explicit missing-observation augmentation. Instead of marginalizing the missing coordinates analytically, we introduce latent variables

$$ y_{k, M_k} $$

for the missing coordinates, fill those values back into the dense observation array, and then score the complete-data observation density

$$ p(y_k \mid x_k, \alpha). $$

In dynestyx, LatentPathBuilder(missing_observation_strategy="auto") will take this route automatically once it sees the concrete missingness pattern in obs_values.

In [10]:

def make_student_t_dynamics(alpha):
    return DynamicalModel(
        initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)),
        state_evolution=LinearGaussianStateEvolution(
            A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]),
            cov=transition_cov_non_gaussian,
        ),
        observation_model=lambda x, u, t: dist.MultivariateStudentT(
            df=student_df,
            loc=x,
            scale_tril=student_scale_tril,
        ),
        control_dim=0,
    )


def student_t_model(obs_times=None, obs_values=None, predict_times=None):
    alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7))
    return dsx.sample(
        "f",
        make_student_t_dynamics(alpha),
        obs_times=obs_times,
        obs_values=obs_values,
        predict_times=predict_times,
    )


with DiscreteTimeSimulator():
    student_synth = Predictive(
        student_t_model,
        params={"alpha": jnp.array(true_alpha)},
        num_samples=1,
        exclude_deterministic=False,
    )(jr.PRNGKey(20), predict_times=obs_times)

student_states_true = np.asarray(student_synth["f_states"].squeeze((0, 1)))
student_obs_clean = np.asarray(student_synth["f_observations"].squeeze((0, 1)))
student_obs_missing = np.asarray(
    apply_nan_mask(jnp.asarray(student_obs_clean), mask_full_plus_partial)
)
student_missing_pairs = np.argwhere(~mask_full_plus_partial)

plot_missing_data(
    t,
    student_states_true,
    student_obs_missing,
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Multivariate Student t observations: partial missingness handled by explicit augmentation",
)

def make_student_t_dynamics(alpha): return DynamicalModel( initial_condition=dist.MultivariateNormal(jnp.zeros(state_dim), jnp.eye(state_dim)), state_evolution=LinearGaussianStateEvolution( A=jnp.array([[alpha, 0.2], [-0.1, 0.8]]), cov=transition_cov_non_gaussian, ), observation_model=lambda x, u, t: dist.MultivariateStudentT( df=student_df, loc=x, scale_tril=student_scale_tril, ), control_dim=0, ) def student_t_model(obs_times=None, obs_values=None, predict_times=None): alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7)) return dsx.sample( "f", make_student_t_dynamics(alpha), obs_times=obs_times, obs_values=obs_values, predict_times=predict_times, ) with DiscreteTimeSimulator(): student_synth = Predictive( student_t_model, params={"alpha": jnp.array(true_alpha)}, num_samples=1, exclude_deterministic=False, )(jr.PRNGKey(20), predict_times=obs_times) student_states_true = np.asarray(student_synth["f_states"].squeeze((0, 1))) student_obs_clean = np.asarray(student_synth["f_observations"].squeeze((0, 1))) student_obs_missing = np.asarray( apply_nan_mask(jnp.asarray(student_obs_clean), mask_full_plus_partial) ) student_missing_pairs = np.argwhere(~mask_full_plus_partial) plot_missing_data( t, student_states_true, student_obs_missing, partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Multivariate Student t observations: partial missingness handled by explicit augmentation", )

In [11]:

student_t_builder = dsx.LatentPathBuilder(
    missing_observation_strategy="auto",
)


def conditioned_student_t_latent_path(obs_times=None, obs_values=None):
    with student_t_builder:
        student_t_model(obs_times=obs_times, obs_values=obs_values)


mcmc_student = MCMC(
    NUTS(conditioned_student_t_latent_path),
    num_warmup=n_mcmc_warmup,
    num_samples=n_mcmc_samples,
    progress_bar=False,
)
mcmc_student.run(jr.PRNGKey(21), obs_times=obs_times, obs_values=jnp.asarray(student_obs_missing))

student_samples = mcmc_student.get_samples()
posterior_alpha_student = np.asarray(student_samples["alpha"])
student_state_samples = np.asarray(student_samples["f_state_path"])
student_missing_obs_samples = np.asarray(student_samples["f_missing_obs_values"])
student_summary = summarize_state_draws(student_state_samples)
missing_rows_student = ~mask_full_plus_partial.all(axis=1)

pp_student = Predictive(
    conditioned_student_t_latent_path,
    posterior_samples=student_samples,
    return_sites=[
        "f_completed_obs_values",
        "f_missing_obs_times",
        "f_missing_obs_coordinate_indices",
    ],
    exclude_deterministic=False,
)(jr.PRNGKey(22), obs_times=obs_times, obs_values=jnp.asarray(student_obs_missing))

student_completed_obs = np.asarray(pp_student["f_completed_obs_values"])
student_completed_obs_mean = student_completed_obs.mean(axis=0)

student_t_builder = dsx.LatentPathBuilder( missing_observation_strategy="auto", ) def conditioned_student_t_latent_path(obs_times=None, obs_values=None): with student_t_builder: student_t_model(obs_times=obs_times, obs_values=obs_values) mcmc_student = MCMC( NUTS(conditioned_student_t_latent_path), num_warmup=n_mcmc_warmup, num_samples=n_mcmc_samples, progress_bar=False, ) mcmc_student.run(jr.PRNGKey(21), obs_times=obs_times, obs_values=jnp.asarray(student_obs_missing)) student_samples = mcmc_student.get_samples() posterior_alpha_student = np.asarray(student_samples["alpha"]) student_state_samples = np.asarray(student_samples["f_state_path"]) student_missing_obs_samples = np.asarray(student_samples["f_missing_obs_values"]) student_summary = summarize_state_draws(student_state_samples) missing_rows_student = ~mask_full_plus_partial.all(axis=1) pp_student = Predictive( conditioned_student_t_latent_path, posterior_samples=student_samples, return_sites=[ "f_completed_obs_values", "f_missing_obs_times", "f_missing_obs_coordinate_indices", ], exclude_deterministic=False, )(jr.PRNGKey(22), obs_times=obs_times, obs_values=jnp.asarray(student_obs_missing)) student_completed_obs = np.asarray(pp_student["f_completed_obs_values"]) student_completed_obs_mean = student_completed_obs.mean(axis=0)

In [12]:

fig, ax = plt.subplots(1, 1, figsize=(5, 3))
az.plot_posterior(posterior_alpha_student, hdi_prob=0.95, ref_val=true_alpha, ax=ax)
ax.set_title(r"Latent-path posterior of $\alpha$  (Student t obs, augmented missing y)")
plt.tight_layout()
plt.show()

observation_dim = student_obs_missing.shape[1]
n_missing_student = student_missing_pairs.shape[0]
fig, axes = plt.subplots(1, n_missing_student, figsize=(3.5 * n_missing_student, 3), squeeze=False)
for j, ax in enumerate(axes[0]):
    t_idx, d_idx = student_missing_pairs[j]
    az.plot_posterior(
        student_completed_obs[:, t_idx, d_idx],
        hdi_prob=0.95,
        ref_val=student_obs_clean[t_idx, d_idx],
        ax=ax,
    )
    ax.set_title(f"reconstructed y[{t_idx}, {d_idx}]")
fig.suptitle("Posterior over reconstructed missing observations")
plt.tight_layout()
plt.show()

student_completed_obs_lo = np.quantile(student_completed_obs, 0.05, axis=0)
student_completed_obs_hi = np.quantile(student_completed_obs, 0.95, axis=0)

fig, axes = plt.subplots(1, observation_dim, figsize=(10, 3.5), sharex=True, sharey=False)
if observation_dim == 1:
    axes = [axes]
for d, ax in enumerate(axes):
    ax.plot(
        t,
        student_obs_clean[:, d],
        color="black",
        lw=0.9,
        ls="--",
        alpha=0.6,
        label="true observation",
    )
    obs_d = np.asarray(student_obs_missing[:, d])
    observed_mask_d = ~np.isnan(obs_d)
    missing_mask_d = ~observed_mask_d
    ax.plot(
        t[observed_mask_d],
        obs_d[observed_mask_d],
        ".",
        ms=5,
        color="C1",
        alpha=0.85,
        label="observed y",
    )
    if np.any(missing_mask_d):
        y_mean = student_completed_obs_mean[missing_mask_d, d]
        y_lo = student_completed_obs_lo[missing_mask_d, d]
        y_hi = student_completed_obs_hi[missing_mask_d, d]
        ax.errorbar(
            t[missing_mask_d],
            y_mean,
            yerr=np.vstack([y_mean - y_lo, y_hi - y_mean]),
            fmt="o",
            color="C3",
            ecolor="C3",
            elinewidth=1.2,
            capsize=3,
            label="reconstructed missing y",
        )
    ax.set_xlabel("time")
    ax.set_title(f"Observation dim {d}")
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys(), fontsize=7)
fig.suptitle("Reconstructed missing observations on the observation scale")
plt.tight_layout()
plt.show()

print("Student t latent-path alpha:", summarize_1d(posterior_alpha_student))
print("True alpha:", true_alpha)
print("Student t state RMSE:", rmse(student_summary["mean"], student_states_true))
print("Student t missing-row RMSE:", rmse(student_summary["mean"][missing_rows_student], student_states_true[missing_rows_student]))
print("Explicit missing-observation latents per draw:", student_missing_obs_samples.shape[-1])
print("missing_obs_times:", np.asarray(pp_student["f_missing_obs_times"][0]))
print(
    "missing_obs_coordinate_indices:",
    np.asarray(pp_student["f_missing_obs_coordinate_indices"][0]),
)
print(
    np.column_stack(
        [
            student_missing_pairs[:, 0],
            student_missing_pairs[:, 1],
            student_obs_clean[student_missing_pairs[:, 0], student_missing_pairs[:, 1]],
            student_completed_obs_mean[student_missing_pairs[:, 0], student_missing_pairs[:, 1]],
        ]
    )
)

plot_state_recovery(
    t,
    student_states_true,
    student_obs_missing,
    student_summary,
    label_a="latent path",
    color_a="C2",
    partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1},
    title="Multivariate Student t case: latent-path state recovery via explicit missing-observation augmentation",
)

fig, ax = plt.subplots(1, 1, figsize=(5, 3)) az.plot_posterior(posterior_alpha_student, hdi_prob=0.95, ref_val=true_alpha, ax=ax) ax.set_title(r"Latent-path posterior of $\alpha$ (Student t obs, augmented missing y)") plt.tight_layout() plt.show() observation_dim = student_obs_missing.shape[1] n_missing_student = student_missing_pairs.shape[0] fig, axes = plt.subplots(1, n_missing_student, figsize=(3.5 * n_missing_student, 3), squeeze=False) for j, ax in enumerate(axes[0]): t_idx, d_idx = student_missing_pairs[j] az.plot_posterior( student_completed_obs[:, t_idx, d_idx], hdi_prob=0.95, ref_val=student_obs_clean[t_idx, d_idx], ax=ax, ) ax.set_title(f"reconstructed y[{t_idx}, {d_idx}]") fig.suptitle("Posterior over reconstructed missing observations") plt.tight_layout() plt.show() student_completed_obs_lo = np.quantile(student_completed_obs, 0.05, axis=0) student_completed_obs_hi = np.quantile(student_completed_obs, 0.95, axis=0) fig, axes = plt.subplots(1, observation_dim, figsize=(10, 3.5), sharex=True, sharey=False) if observation_dim == 1: axes = [axes] for d, ax in enumerate(axes): ax.plot( t, student_obs_clean[:, d], color="black", lw=0.9, ls="--", alpha=0.6, label="true observation", ) obs_d = np.asarray(student_obs_missing[:, d]) observed_mask_d = ~np.isnan(obs_d) missing_mask_d = ~observed_mask_d ax.plot( t[observed_mask_d], obs_d[observed_mask_d], ".", ms=5, color="C1", alpha=0.85, label="observed y", ) if np.any(missing_mask_d): y_mean = student_completed_obs_mean[missing_mask_d, d] y_lo = student_completed_obs_lo[missing_mask_d, d] y_hi = student_completed_obs_hi[missing_mask_d, d] ax.errorbar( t[missing_mask_d], y_mean, yerr=np.vstack([y_mean - y_lo, y_hi - y_mean]), fmt="o", color="C3", ecolor="C3", elinewidth=1.2, capsize=3, label="reconstructed missing y", ) ax.set_xlabel("time") ax.set_title(f"Observation dim {d}") handles, labels = ax.get_legend_handles_labels() by_label = dict(zip(labels, handles)) ax.legend(by_label.values(), by_label.keys(), fontsize=7) fig.suptitle("Reconstructed missing observations on the observation scale") plt.tight_layout() plt.show() print("Student t latent-path alpha:", summarize_1d(posterior_alpha_student)) print("True alpha:", true_alpha) print("Student t state RMSE:", rmse(student_summary["mean"], student_states_true)) print("Student t missing-row RMSE:", rmse(student_summary["mean"][missing_rows_student], student_states_true[missing_rows_student])) print("Explicit missing-observation latents per draw:", student_missing_obs_samples.shape[-1]) print("missing_obs_times:", np.asarray(pp_student["f_missing_obs_times"][0])) print( "missing_obs_coordinate_indices:", np.asarray(pp_student["f_missing_obs_coordinate_indices"][0]), ) print( np.column_stack( [ student_missing_pairs[:, 0], student_missing_pairs[:, 1], student_obs_clean[student_missing_pairs[:, 0], student_missing_pairs[:, 1]], student_completed_obs_mean[student_missing_pairs[:, 0], student_missing_pairs[:, 1]], ] ) ) plot_state_recovery( t, student_states_true, student_obs_missing, student_summary, label_a="latent path", color_a="C2", partial_indices={0: partial_idx_dim0, 1: partial_idx_dim1}, title="Multivariate Student t case: latent-path state recovery via explicit missing-observation augmentation", )

Student t latent-path alpha: {'mean': 0.1667250245809555, 'p05': -0.38418588042259216, 'p95': 0.6336016654968262}
True alpha: 0.4
Student t state RMSE: 0.28437796235084534
Student t missing-row RMSE: 0.35274210572242737
Explicit missing-observation latents per draw: 58
missing_obs_times: [ 8. 18. 25. 33. 35. 35. 36. 36. 37. 37. 38. 38. 39. 39. 40. 40. 41. 41.
 42. 42. 43. 43. 44. 44. 45. 45. 46. 46. 47. 47. 48. 48. 49. 49. 50. 50.
 51. 51. 52. 52. 53. 53. 54. 54. 55. 55. 56. 56. 57. 57. 58. 58. 59. 59.
 72. 84. 90. 95.]
missing_obs_coordinate_indices: [0 0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0
 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1]
[[ 8.     0.     0.577  0.437]
 [18.     0.     0.358  0.161]
 [25.     1.     0.727  0.501]
 [33.     1.     0.578  0.261]
 [35.     0.    -0.246  0.022]
 [35.     1.    -0.345  0.194]
 [36.     0.    -0.448  0.021]
 [36.     1.     0.688  0.196]
 [37.     0.     0.026  0.011]
 [37.     1.     0.516  0.149]
 [38.     0.    -0.127  0.068]
 [38.     1.    -0.254  0.182]
 [39.     0.     0.283  0.057]
 [39.     1.    -0.228  0.161]
 [40.     0.     0.734  0.089]
 [40.     1.    -0.257  0.143]
 [41.     0.    -0.41   0.077]
 [41.     1.    -0.404  0.158]
 [42.     0.     0.644  0.056]
 [42.     1.    -1.251  0.031]
 [43.     0.    -0.269 -0.032]
 [43.     1.    -0.266  0.008]
 [44.     0.     0.794  0.154]
 [44.     1.    -0.812 -0.018]
 [45.     0.    -0.702  0.079]
 [45.     1.    -0.9    0.011]
 [46.     0.     0.69   0.   ]
 [46.     1.    -0.075 -0.08 ]
 [47.     0.     0.341  0.067]
 [47.     1.     0.303 -0.042]
 [48.     0.     0.151  0.035]
 [48.     1.    -0.107 -0.091]
 [49.     0.    -0.583  0.01 ]
 [49.     1.     1.044 -0.078]
 [50.     0.    -1.107  0.005]
 [50.     1.     0.329 -0.08 ]
 [51.     0.     2.026 -0.01 ]
 [51.     1.     0.895 -0.109]
 [52.     0.     1.03   0.028]
 [52.     1.    -0.852  0.052]
 [53.     0.    -0.733 -0.017]
 [53.     1.     0.162  0.088]
 [54.     0.     0.141  0.118]
 [54.     1.     0.302  0.077]
 [55.     0.     0.333  0.066]
 [55.     1.     1.698  0.045]
 [56.     0.     0.023  0.027]
 [56.     1.     0.148  0.061]
 [57.     0.    -0.638 -0.035]
 [57.     1.     0.359  0.132]
 [58.     0.    -0.276  0.008]
 [58.     1.     1.01   0.245]
 [59.     0.     0.122  0.087]
 [59.     1.    -0.092  0.333]
 [72.     0.     0.365  0.446]
 [84.     1.     0.889  0.316]
 [90.     0.    -1.136 -0.437]
 [95.     1.     0.148  0.29 ]]

Summary¶

The clean simulator-side interpretation is:

• obs_values[k] always belongs to latent index $k$;
• a fully missing row means: keep $x_k$, but add no observation factor at that index;
• a partially missing row means: keep $x_k$ and use whatever observation structure is available on the observed coordinates.

That immediately implies the key scope split:

• full-row missingness works for all observation models under LatentPathBuilder;
• partial missingness can be handled in several ways depending on model structure.

In this tutorial those routes were:

• multivariate Gaussian: restrict to the observed subvector and covariance submatrix;
• independent asymmetric-Laplace observations: sum only the observed-coordinate log-probabilities;
• exact identity observations: compress state_path_params down to only the missing coordinates, then reconstruct the full f_state_path deterministically;
• correlated continuous observations like multivariate Student $t$: introduce explicit missing-observation latents f_missing_obs_values, reconstruct f_completed_obs_values, and score the completed observations directly.

So LatentPathBuilder now has a practical fallback when direct marginalization is unavailable: explicit missing-observation augmentation. This is especially useful for continuous observation families that are incompatible with the filter/smoother missingness path or do not admit a simple masked likelihood formula.