Skip to content

Add future forecasting functionality for LatentPathBuilder #270

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DanWaxman merged 1 commit into from
Jul 13, 2026
Merged

Add future forecasting functionality for LatentPathBuilder #270

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
201 changes: 137 additions & 64 deletions dynestyx/inference/latent/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@
from __future__ import annotations

import dataclasses
from typing import Any
from typing import Any, cast

import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from effectful.ops.semantics import fwd
from effectful.ops.syntax import ObjectInterpretation, implements
from jaxtyping import Array, Real

Expand All @@ -35,6 +36,10 @@
_prepare_latent_path_request,
_PreparedLatentPathRequest,
)
from dynestyx.inference.posterior_rollout import (
_final_times_for_rollout,
_validate_future_only_predict_times,
)
from dynestyx.inference.state_paths.reconstruct import AssembledStatePath
from dynestyx.inference.state_paths.score import (
TrajectoryLogProbTerms,
Expand All @@ -50,6 +55,7 @@ def _build_latent_state_result(
prepared: _PreparedLatentPathRequest,
assembled_state_path: AssembledStatePath | None,
log_prob_terms: TrajectoryLogProbTerms | None,
state_dists: list[dist.Distribution] | None,
) -> LatentStateResult:
"""Package one latent-path evaluation into the public result dataclass.

Expand Down Expand Up @@ -103,10 +109,22 @@ def _build_latent_state_result(
completed_obs_values=(
None if log_prob_terms is None else log_prob_terms.completed_obs_values
),
state_dists=None,
state_dists=state_dists,
)


def _build_state_path_distributions(
dynamics,
state_path: Array,
) -> list[dist.Distribution]:
"""Wrap each inferred state in a Delta distribution for rollout forwarding."""
event_dim = len(dynamics.initial_condition.event_shape)
time_major_state_path = jnp.moveaxis(jnp.asarray(state_path), -(event_dim + 1), 0)
return [
dist.Delta(state_t, event_dim=event_dim) for state_t in time_major_state_path
]


@dataclasses.dataclass(init=False)
class LatentPathBuilder(ObjectInterpretation, HandlesSelf):
"""Build latent path parameters and score ``log p(x, y | ...)``.
Expand Down Expand Up @@ -312,6 +330,9 @@ def _sample_single(
prepared=prepared,
assembled_state_path=assembled_state_path,
log_prob_terms=log_prob_terms,
state_dists=_build_state_path_distributions(
dynamics, assembled_state_path.state_path
),
)

@implements(_condition_intp)
Expand All @@ -332,6 +353,7 @@ def _sample_ds(
ctrl_values: Real[Array, "*ctrl_value_plate ctrl_time control_dim"]
| Real[Array, "*ctrl_value_plate ctrl_time"]
| None = None,
predict_times: Real[Array, "*predict_time_plate predict_time"] | None = None,
state_path_params: Array | None = None,
missing_obs_values: Array | None = None,
_dsx_sample_mode: bool = False,
Expand All @@ -345,7 +367,7 @@ def _sample_ds(
)

if not plate_shapes:
return self._sample_single(
result = self._sample_single(
name,
dynamics,
obs_times=obs_times,
Expand All @@ -358,73 +380,124 @@ def _sample_ds(
state_path_params=state_path_params,
missing_obs_values=missing_obs_values,
)
else:
member_specs = _plate_member_specs(
name=name,
dynamics=dynamics,
plate_shapes=plate_shapes,
obs_times=obs_times,
obs_values=obs_values,
obs_values_filled=_obs_values_filled,
obs_mask=_obs_mask,
ctrl_times=ctrl_times,
ctrl_values=ctrl_values,
state_path_params=state_path_params,
missing_obs_values=missing_obs_values,
)

member_specs = _plate_member_specs(
name=name,
dynamics=dynamics,
member_results = []
for member in member_specs:
with _suspend_numpyro_plate_frames():
member_results.append(
self._sample_single(
member.name,
member.dynamics,
obs_times=member.obs_times,
obs_values=member.obs_values,
_obs_values_filled=member.obs_values_filled,
_obs_mask=member.obs_mask,
_obs_has_missing=_obs_has_missing,
ctrl_times=member.ctrl_times,
ctrl_values=member.ctrl_values,
state_path_params=member.state_path_params,
missing_obs_values=member.missing_obs_values,
)
)

state_path = _stack_member_attr(member_results, "state_path", plate_shapes)
result = LatentStateResult(
joint_log_prob=_stack_member_attr(
member_results, "joint_log_prob", plate_shapes
),
state_path_params=_stack_member_attr(
member_results, "state_path_params", plate_shapes
),
state_path_param_times=_stack_member_attr(
member_results, "state_path_param_times", plate_shapes
),
state_path_param_coordinate_indices=_stack_member_attr(
member_results,
"state_path_param_coordinate_indices",
plate_shapes,
),
state_path=state_path,
state_path_times=_stack_member_attr(
member_results, "state_path_times", plate_shapes
),
missing_obs_values=_stack_member_attr(
member_results, "missing_obs_values", plate_shapes
),
missing_obs_times=_stack_member_attr(
member_results, "missing_obs_times", plate_shapes
),
missing_obs_coordinate_indices=_stack_member_attr(
member_results, "missing_obs_coordinate_indices", plate_shapes
),
completed_obs_values=_stack_member_attr(
member_results, "completed_obs_values", plate_shapes
),
state_dists=(
None
if state_path is None
else _build_state_path_distributions(dynamics, state_path)
),
)

predict_times = _validate_future_only_predict_times(
predict_times,
cast(Array | None, result.state_path_times),
error_message=(
"LatentPathBuilder rollout only supports predict_times >= "
"max(state_path_times); in-window latent-path predictions are "
"not implemented yet."
),
)
filtered_times = None
filtered_dists = None
posterior_rollout_final_only = False
smoothed_times = result.state_path_times
smoothed_dists = result.state_dists
if predict_times is not None and smoothed_dists:
assert result.state_path_times is not None
filtered_times = _final_times_for_rollout(
cast(Array, result.state_path_times)
)
filtered_dists = [smoothed_dists[-1]]
posterior_rollout_final_only = True
smoothed_times = None
smoothed_dists = None

forwarded_result = fwd(
name,
dynamics,
plate_shapes=plate_shapes,
obs_times=obs_times,
obs_values=obs_values,
obs_values_filled=_obs_values_filled,
obs_mask=_obs_mask,
obs_times=None,
obs_values=None,
ctrl_times=ctrl_times,
ctrl_values=ctrl_values,
state_path_params=state_path_params,
missing_obs_values=missing_obs_values,
predict_times=predict_times,
filtered_times=filtered_times,
filtered_dists=filtered_dists,
smoothed_times=smoothed_times,
smoothed_dists=smoothed_dists,
_posterior_rollout_final_only=posterior_rollout_final_only,
**kwargs,
)
downstream_register = getattr(forwarded_result, "_register_numpyro_sites", None)
if callable(downstream_register):
downstream_register(name)

member_results = []
for member in member_specs:
with _suspend_numpyro_plate_frames():
member_results.append(
self._sample_single(
member.name,
member.dynamics,
obs_times=member.obs_times,
obs_values=member.obs_values,
_obs_values_filled=member.obs_values_filled,
_obs_mask=member.obs_mask,
_obs_has_missing=_obs_has_missing,
ctrl_times=member.ctrl_times,
ctrl_values=member.ctrl_values,
state_path_params=member.state_path_params,
missing_obs_values=member.missing_obs_values,
)
)

return LatentStateResult(
joint_log_prob=_stack_member_attr(
member_results, "joint_log_prob", plate_shapes
),
state_path_params=_stack_member_attr(
member_results, "state_path_params", plate_shapes
),
state_path_param_times=_stack_member_attr(
member_results, "state_path_param_times", plate_shapes
),
state_path_param_coordinate_indices=_stack_member_attr(
member_results,
"state_path_param_coordinate_indices",
plate_shapes,
),
state_path=_stack_member_attr(member_results, "state_path", plate_shapes),
state_path_times=_stack_member_attr(
member_results, "state_path_times", plate_shapes
),
missing_obs_values=_stack_member_attr(
member_results, "missing_obs_values", plate_shapes
),
missing_obs_times=_stack_member_attr(
member_results, "missing_obs_times", plate_shapes
),
missing_obs_coordinate_indices=_stack_member_attr(
member_results, "missing_obs_coordinate_indices", plate_shapes
),
completed_obs_values=_stack_member_attr(
member_results, "completed_obs_values", plate_shapes
),
state_dists=None,
)
return result


__all__ = ["LatentPathBuilder"]
38 changes: 38 additions & 0 deletions dynestyx/inference/posterior_rollout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
"""Shared helpers for posterior-rollout time handling."""

from __future__ import annotations

import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
from jaxtyping import Array, Real


def _validate_future_only_predict_times(
predict_times: Real[Array, "*predict_time_plate predict_time"] | None,
anchor_times: Real[Array, "*anchor_time_plate anchor_time"] | None,
*,
error_message: str,
) -> Real[Array, "*predict_time_plate predict_time"] | None:
"""Validate the current future-only posterior-rollout contract."""
if predict_times is None or anchor_times is None:
return predict_times
anchor_end = anchor_times[..., -1:]
_ = eqx.error_if(
predict_times,
jnp.any(predict_times < anchor_end),
error_message,
)
return predict_times


def _final_times_for_rollout(
times: Real[Array, "*time_plate time"],
) -> Real[Array, "*time_plate one"]:
"""Return the final anchor time while keeping simulator segmentation host-safe."""
try:
times_host = np.asarray(jax.device_get(times))
return jnp.asarray(times_host[..., -1:], dtype=times.dtype)
except Exception:
return times[..., -1:]
44 changes: 14 additions & 30 deletions dynestyx/inference/smoothers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from effectful.ops.semantics import fwd
from effectful.ops.syntax import ObjectInterpretation, implements
Expand Down Expand Up @@ -46,6 +45,10 @@
from dynestyx.inference.integrations.cuthbert.discrete_smoother import (
run_discrete_smoother as run_cuthbert_discrete_smoother,
)
from dynestyx.inference.posterior_rollout import (
_final_times_for_rollout,
_validate_future_only_predict_times,
)
from dynestyx.inference.utils.distribution_utils import (
_cholesky_state_sequence_to_dists,
_posterior_sequence_to_dists,
Expand Down Expand Up @@ -86,33 +89,6 @@ def _valid_smoother_config_names(*, continuous_time: bool) -> list[str]:
return [c.__name__ for c in DiscreteTimeSmootherConfigs]


def _validate_future_only_predict_times(
predict_times: Real[Array, "*predict_time_plate predict_time"] | None,
obs_times: Real[Array, "*obs_time_plate obs_time"] | None,
) -> Real[Array, "*predict_time_plate predict_time"] | None:
"""Validate the current smoother prediction contract."""
if predict_times is None or obs_times is None:
return predict_times
obs_end = obs_times[..., -1:]
_ = eqx.error_if(
predict_times,
jnp.any(predict_times < obs_end),
"Smoother prediction only supports predict_times >= max(obs_times); in-window smoothing predictions are not implemented yet. Please use `Filter` for in-window predictions for now.",
)
return predict_times


def _final_obs_times_for_rollout(
obs_times: Real[Array, "*obs_time_plate obs_time"],
) -> Real[Array, "*obs_time_plate one"]:
"""Return the final observation time while keeping simulator segmentation host-safe."""
try:
obs_times_host = np.asarray(jax.device_get(obs_times))
return jnp.asarray(obs_times_host[..., -1:], dtype=obs_times.dtype)
except Exception:
return obs_times[..., -1:]


class BaseSmootherLogFactorAdder(ObjectInterpretation, HandlesSelf, ABC):
"""Base class for smoother handlers."""

Expand Down Expand Up @@ -147,14 +123,22 @@ def _sample_ds(
**kwargs,
)

predict_times = _validate_future_only_predict_times(predict_times, obs_times)
predict_times = _validate_future_only_predict_times(
predict_times,
obs_times,
error_message=(
"Smoother prediction only supports predict_times >= max(obs_times); "
"in-window smoothing predictions are not implemented yet. "
"Please use `Filter` for in-window predictions for now."
),
)
filtered_times = None
filtered_dists = None
posterior_rollout_final_only = False
smoothed_times = obs_times
if predict_times is not None and smoothed_dists:
assert obs_times is not None
filtered_times = _final_obs_times_for_rollout(obs_times)
filtered_times = _final_times_for_rollout(obs_times)
filtered_dists = [smoothed_dists[-1]]
posterior_rollout_final_only = True
smoothed_times = None
Expand Down
Loading
Loading