Skip to content

CHS: freeze states after each person's last observed measurement #104

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

Open
hmgaudecker wants to merge 1 commit into
base: asymmetric-translog
Choose a base branch
from
Open

CHS: freeze states after each person's last observed measurement #104

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
11 changes: 9 additions & 2 deletions src/skillmodels/chs/kalman_filters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,12 +86,19 @@ def kalman_update(
)

_residuals = _safe_measurements.reshape(n_obs, 1) - _safe_expected_measurements
_f_stars = jnp.dot(upper_chols, loadings.reshape(n_states, 1))
# An identity covariance for missing measurements keeps the QR decomposition and
# its derivative well defined; the updated covariance is discarded below.
_safe_upper_chols = jnp.where(
not_missing.reshape(n_obs, 1, 1, 1),
upper_chols,
jnp.eye(n_states, dtype=upper_chols.dtype),
)
_f_stars = jnp.dot(_safe_upper_chols, loadings.reshape(n_states, 1))

_m = jnp.zeros((n_obs, n_mixtures, n_states + 1, n_states + 1))
_m = _m.at[..., 0, 0].set(meas_sd)
_m = _m.at[..., 1:, :1].set(_f_stars)
_m = _m.at[..., 1:, 1:].set(upper_chols)
_m = _m.at[..., 1:, 1:].set(_safe_upper_chols)

_r = array_qr_jax(_m)[1]

Expand Down
117 changes: 92 additions & 25 deletions src/skillmodels/chs/likelihood.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ def log_likelihood_obs(
"measurements": measurements,
"is_measurement_iteration": is_measurement_iteration,
"is_predict_iteration": is_predict_iteration,
"observed_later": get_observation_masks(measurements, iteration_to_period)[0],
}

_body = functools.partial(
Expand All @@ -206,6 +207,90 @@ def log_likelihood_obs(
).sum(axis=0)


def get_observation_masks(
measurements: Array | np.ndarray,
iteration_to_period: Array | np.ndarray,
) -> tuple[Array, Array]:
"""Return masks of persons with an observed measurement after each update.

Args:
measurements: Array of shape (n_updates, n_obs). NaN if not observed.
iteration_to_period: Array of length n_updates mapping each update to its
period. Updates of one period are contiguous.

Returns:
- observed_later: Array of shape (n_updates, n_obs). True if the person has
an observed measurement in a later update. Only these persons are
predicted after the update.
- observed_from_period: Array of shape (n_updates, n_obs). True if the
person has an observed measurement in the period of the update or in a
later period. Filtered states are reported only for these persons.

"""
observed = jnp.isfinite(measurements)
observed_from = jnp.flip(jnp.cumsum(jnp.flip(observed, axis=0), axis=0) > 0, axis=0)
observed_later = jnp.concatenate(
[observed_from[1:], jnp.zeros_like(observed_from[:1])], axis=0
)
periods = jnp.asarray(iteration_to_period)
is_period_start = jnp.concatenate([jnp.array([True]), periods[1:] != periods[:-1]])
period_start = jax.lax.cummax(
jnp.where(is_period_start, jnp.arange(len(periods)), 0), axis=0
)
return observed_later, observed_from[period_start]


def predict_if_observed_later(
*,
predict_kwargs: dict[str, Any],
observed_later: Array | np.ndarray,
is_predict_iteration: Array | np.ndarray,
predict_func: Callable,
) -> tuple[Array, Array]:
"""Predict the persons with a later observation; keep all others unchanged.

After the last observed measurement of a person, states and covariances stay at
their last filtered values. The predict step runs on placeholder inputs (zero
states, identity covariance) for these persons and its result is discarded.
The placeholders keep the predict step and its derivatives finite, which a
covariance that is predicted over many unobserved periods does not do.

Args:
predict_kwargs: Keyword arguments of `predict_func`, with `states` of shape
(n_obs, n_mixtures, n_states) and `upper_chols` of shape
(n_obs, n_mixtures, n_states, n_states).
observed_later: Boolean array of length n_obs.
is_predict_iteration: Whether a predict step follows this update.
predict_func: The predict function.

Returns:
new_states: Same shape as `predict_kwargs["states"]`.
new_upper_chols: Same shape as `predict_kwargs["upper_chols"]`.

"""
states = predict_kwargs["states"]
upper_chols = predict_kwargs["upper_chols"]
keep_states = observed_later.reshape(-1, 1, 1)
keep_chols = observed_later.reshape(-1, 1, 1, 1)
safe_kwargs = {
**predict_kwargs,
"states": jnp.where(keep_states, states, 0.0),
"upper_chols": jnp.where(
keep_chols, upper_chols, jnp.eye(states.shape[-1], dtype=upper_chols.dtype)
),
}
new_states, new_upper_chols = jax.lax.cond(
is_predict_iteration,
lambda kwargs: predict_func(**kwargs),
lambda kwargs: (kwargs["states"], kwargs["upper_chols"]),
safe_kwargs,
)
return (
jnp.where(keep_states, new_states, states),
jnp.where(keep_chols, new_upper_chols, upper_chols),
)


def _scan_body(
carry: dict[str, Array],
loop_args: dict[str, Array],
Expand Down Expand Up @@ -264,16 +349,15 @@ def _scan_body(
"observed_factor_vars": observed_factor_vars[t],
}

fixed_kwargs = {"predict_func": predict_func}

# ==================================================================================
# Do a predict step or a do-nothing fake predict step
# Do a predict step for persons observed later; freeze all others
# ==================================================================================
states, upper_chols, filtered_states = jax.lax.cond(
loop_args["is_predict_iteration"],
functools.partial(_one_arg_predict, **fixed_kwargs),
functools.partial(_one_arg_no_predict, **fixed_kwargs),
predict_kwargs,
filtered_states = states
states, upper_chols = predict_if_observed_later(
predict_kwargs=predict_kwargs,
observed_later=loop_args["observed_later"],
is_predict_iteration=loop_args["is_predict_iteration"],
predict_func=predict_func,
)

new_state = {
Expand Down Expand Up @@ -302,20 +386,3 @@ def _one_arg_anchoring_update(
new_log_mixture_weights,
new_loglikes,
)


def _one_arg_no_predict(
kwargs: dict[str, Any],
predict_func: Callable, # noqa: ARG001
) -> tuple[Array, Array, Array]:
"""Just return the states cond chols without any changes."""
return kwargs["states"], kwargs["upper_chols"], kwargs["states"]


def _one_arg_predict(
kwargs: dict[str, Any],
predict_func: Callable,
) -> tuple[Array, Array, Array]:
"""Do a predict step but also return the input states as filtered states."""
new_states, new_upper_chols = predict_func(**kwargs)
return new_states, new_upper_chols, kwargs["states"]
41 changes: 15 additions & 26 deletions src/skillmodels/chs/likelihood_debug.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

from skillmodels.chs.clipping import soft_clipping
from skillmodels.chs.kalman_filters_debug import kalman_update
from skillmodels.chs.likelihood import get_observation_masks, predict_if_observed_later
from skillmodels.chs.options import CHSEstimationOptions
from skillmodels.common.parse_params import parse_params
from skillmodels.common.types import (
Expand Down Expand Up @@ -90,6 +91,9 @@ def log_likelihood(
"log_mixture_weights": log_mixture_weights,
}

observed_later, observed_from_period = get_observation_masks(
measurements, iteration_to_period
)
loop_args = {
"period": iteration_to_period,
"loadings": parsed_params.loadings,
Expand All @@ -98,6 +102,7 @@ def log_likelihood(
"measurements": measurements,
"is_measurement_iteration": is_measurement_iteration,
"is_predict_iteration": is_predict_iteration,
"observed_later": observed_later,
}

_body = functools.partial(
Expand Down Expand Up @@ -146,7 +151,9 @@ def log_likelihood(
out["initial_states"] = initial_states
out["initial_log_mixture_weights"] = initial_log_mixture_weights

out["filtered_states"] = static_out["states"]
out["filtered_states"] = jnp.where(
observed_from_period[..., None, None], static_out["states"], jnp.nan
)
out["log_mixture_weights"] = static_out["log_mixture_weights"]

return out
Expand Down Expand Up @@ -210,16 +217,15 @@ def _scan_body(
"observed_factor_vars": observed_factor_vars[t],
}

fixed_kwargs = {"predict_func": predict_func}

# ==================================================================================
# Do a predict step or a do-nothing fake predict step
# Do a predict step for persons observed later; freeze all others
# ==================================================================================
states, upper_chols, filtered_states = jax.lax.cond(
loop_args["is_predict_iteration"],
functools.partial(_one_arg_predict, **fixed_kwargs),
functools.partial(_one_arg_no_predict, **fixed_kwargs),
predict_kwargs,
filtered_states = states
states, upper_chols = predict_if_observed_later(
predict_kwargs=predict_kwargs,
observed_later=loop_args["observed_later"],
is_predict_iteration=loop_args["is_predict_iteration"],
predict_func=predict_func,
)

new_state = {
Expand Down Expand Up @@ -249,20 +255,3 @@ def _one_arg_anchoring_update(
new_loglikes,
debug_info,
)


def _one_arg_no_predict(
kwargs: dict[str, Any],
predict_func: Callable[..., tuple[Array, Array]], # noqa: ARG001
) -> tuple[Array, Array, Array]:
"""Just return the states cond chols without any changes."""
return kwargs["states"], kwargs["upper_chols"], kwargs["states"]


def _one_arg_predict(
kwargs: dict[str, Any],
predict_func: Callable[..., tuple[Array, Array]],
) -> tuple[Array, Array, Array]:
"""Do a predict step but also return the input states as filtered states."""
new_states, new_upper_chols = predict_func(**kwargs)
return new_states, new_upper_chols, kwargs["states"]
Loading
Loading