diff --git a/docs/getting_started/tutorial.ipynb b/docs/getting_started/tutorial.ipynb index 46082e77..9a9ef30c 100644 --- a/docs/getting_started/tutorial.ipynb +++ b/docs/getting_started/tutorial.ipynb @@ -534,7 +534,7 @@ ], "source": [ "chs_filtered = get_individual_states_from_params(\n", - " model_spec=model, data=data, params=chs_params\n", + " model_spec=model, data=data, params=chs_params, chs_options=chs_options\n", ")\n", "chs_states = chs_filtered[\"unanchored_states\"][\"states\"]\n", "chs_decomp = decompose_measurement_variance(\n", diff --git a/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb b/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb index 56fc9a95..e950d71e 100644 --- a/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb +++ b/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb @@ -33,7 +33,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": "model = MODEL2\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "model = MODEL2\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "markdown", @@ -254,7 +254,7 @@ "id": "20", "metadata": {}, "outputs": [], - "source": "params = get_maximization_inputs(\n model_spec=model, data=data, chs_options=MODEL2_CHS_OPTIONS\n)[\"params_template\"]\nparams[\"value\"] = 0.1\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "params = get_maximization_inputs(\n model_spec=model, data=data, chs_options=MODEL2_CHS_OPTIONS\n)[\"params_template\"]\nparams[\"value\"] = 0.1\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "code", diff --git a/docs/how_to_guides/how_to_visualize_transition_equations.ipynb b/docs/how_to_guides/how_to_visualize_transition_equations.ipynb index 1225f49c..96ee44c7 100644 --- a/docs/how_to_guides/how_to_visualize_transition_equations.ipynb +++ b/docs/how_to_guides/how_to_visualize_transition_equations.ipynb @@ -6,7 +6,7 @@ "id": "0", "metadata": {}, "outputs": [], - "source": "import pandas as pd\n\nfrom skillmodels.common.config import REGRESSION_VAULT, TEST_DATA_DIR\nfrom skillmodels.common.individual_states import get_individual_states_from_params\nfrom skillmodels.common.visualize_transition_equations import (\n combine_transition_plots,\n get_transition_plots,\n)\nfrom skillmodels.test_data.model2 import MODEL2" + "source": "import pandas as pd\n\nfrom skillmodels.common.config import REGRESSION_VAULT, TEST_DATA_DIR\nfrom skillmodels.common.individual_states import get_individual_states_from_params\nfrom skillmodels.common.visualize_transition_equations import (\n combine_transition_plots,\n get_transition_plots,\n)\nfrom skillmodels.test_data.model2 import MODEL2, MODEL2_CHS_OPTIONS" }, { "cell_type": "markdown", @@ -37,7 +37,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": "model = MODEL2\n\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "model = MODEL2\n\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "markdown", diff --git a/docs/reference_guides/estimator_prerequisites.md b/docs/reference_guides/estimator_prerequisites.md index 7a91d36f..9e22066b 100644 --- a/docs/reference_guides/estimator_prerequisites.md +++ b/docs/reference_guides/estimator_prerequisites.md @@ -13,20 +13,20 @@ three estimators. ## Prerequisites matrix -| Dimension | CHS | AF | AMN | -| --------------------- | ------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Estimator | Joint maximum likelihood with square-root Gaussian-component filtering. | Sequential period-by-period likelihood with Halton integration. | Three stages: Gaussian-mixture EM → minimum distance → simulate-and-regress. | -| Latent distribution | Finite mixture of Gaussian initial states via `ModelSpec.n_mixtures`; Gaussian component filtering thereafter. | Initial finite mixture plus sequentially carried distributions; Halton nodes approximate the integrals. | Stage-1 Gaussian mixture over the augmented measurement vector; `n_mixtures` sets the component count. | -| Measurement families | **Gaussian only** (standard Kalman update). Probit/Tobit measurements are not consumed by the CHS path. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | -| Missing data | Gaussian measurement updates skip individually missing measurements. | Measurement masks skip missing rows in the per-step likelihood contributions. | `mixture_em_method="complete_case"` (default) or `"missing_data"` (marginalises over missing entries under MAR); never-observed columns require `allow_never_observed_measurements=True`. | -| Endogenous investment | Reconstructed endogenous factors plus a full `CorrectionSpec` control-function basis. | Reconstructed endogenous investment with independent shocks (source/destination calendar adapter); rejects `CorrectionSpec` / nonzero `kappa`. | Endogenous investment in Stage 3; correction support is a linearised control-function term, narrower than CHS. | -| Corrections | Full `CorrectionSpec` / `kappa` polynomial basis in the processed transition DAG. | Not implemented — the validator raises if a `CorrectionSpec` is attached. | Linear `cf` term only; a higher-order `kappa_terms` request raises. | -| Custom transitions | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables via the Stage-3 generic NLS path, but `log_ces` / `log_ces_with_constant` reject fixed parameters there. | -| `fixed_params` | Honoured through the shared parameter index. | Honoured (with public `constraints=` limited to `select_by_loc` equality groups). | Honoured only for the categories each stage estimates (Stage-2 loadings/intercepts/SDs, Stage-3 `transition`); other categories raise. `start_params` and `constraints` raise. | -| Normalization | Shared `Normalizations` / `fixed_params` / equality constraints. The checker is a precheck, not an identification proof. | Same public `ModelSpec`, with the AF-specific source/destination calendar and `af_state_role` metadata. | Stage-2 imposes its own structural moment restrictions and a mean-zero mixture convention. | -| Anchoring | Supported through `get_maximization_inputs` / the CHS path. | Not part of the AF likelihood; downstream visualization only. | Not wired through the AMN stages (the result reports unanchored scales). | -| Cost / scaling | Potentially expensive joint ML; the JAX square-root filter helps numerical stability. | Sequential but quadrature-heavy; cost grows with node count and state dimension. | Fast when the stages are well behaved; the cluster bootstrap is expensive because it re-estimates every stage. | -| Use it for | Likelihood benchmark, anchoring, correction-heavy models, Gaussian measurement systems. | Sequential AF-style models, initial-period limited measurements, period-by-period diagnostics. | Mixture-heavy Gaussian-measurement models, fast start values for CHS, Stage-1/2 structural diagnostics. | +| Dimension | CHS | AF | AMN | +| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Estimator | Joint maximum likelihood with square-root Gaussian-component filtering. | Sequential period-by-period likelihood with Halton integration. | Three stages: Gaussian-mixture EM → minimum distance → simulate-and-regress. | +| Latent distribution | Finite mixture of Gaussian initial states via `ModelSpec.n_mixtures`; Gaussian component filtering thereafter. | Initial finite mixture plus sequentially carried distributions; Halton nodes approximate the integrals. | Stage-1 Gaussian mixture over the augmented measurement vector; `n_mixtures` sets the component count. | +| Measurement families | **Gaussian only** (standard Kalman update): `get_maximization_inputs` / `estimate_chs` raise `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | +| Missing data | Gaussian measurement updates skip individually missing measurements. | Measurement masks skip missing rows in the per-step likelihood contributions. | `mixture_em_method="complete_case"` (default) or `"missing_data"` (marginalises over missing entries under MAR); never-observed columns require `allow_never_observed_measurements=True`. | +| Endogenous investment | Reconstructed endogenous factors plus a full `CorrectionSpec` control-function basis. | Reconstructed endogenous investment with independent shocks (source/destination calendar adapter); rejects `CorrectionSpec` / nonzero `kappa`. | Endogenous investment in Stage 3; correction support is a linearised control-function term, narrower than CHS. | +| Corrections | Full `CorrectionSpec` / `kappa` polynomial basis in the processed transition DAG. | Not implemented — the validator raises if a `CorrectionSpec` is attached. | Linear `cf` term only; a higher-order `kappa_terms` request raises. | +| Custom transitions | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables via the Stage-3 generic NLS path, but `log_ces` / `log_ces_with_constant` reject fixed parameters there. | +| `fixed_params` | Honoured through the shared parameter index. | Honoured (with public `constraints=` limited to `select_by_loc` equality groups). | Honoured only for the categories each stage estimates (Stage-2 loadings/intercepts/SDs, Stage-3 `transition`); other categories raise. `start_params` and `constraints` raise. | +| Normalization | Shared `Normalizations` / `fixed_params` / equality constraints. The checker is a precheck, not an identification proof. | Same public `ModelSpec`, with the AF-specific source/destination calendar and `af_state_role` metadata. | Stage-2 imposes its own structural moment restrictions and a mean-zero mixture convention. | +| Anchoring | Supported through `get_maximization_inputs` / the CHS path. | Not part of the AF likelihood; downstream visualization only. | Not wired through the AMN stages (the result reports unanchored scales). | +| Cost / scaling | Potentially expensive joint ML; the JAX square-root filter helps numerical stability. | Sequential but quadrature-heavy; cost grows with node count and state dimension. | Fast when the stages are well behaved; the cluster bootstrap is expensive because it re-estimates every stage. | +| Use it for | Likelihood benchmark, anchoring, correction-heavy models, Gaussian measurement systems. | Sequential AF-style models, initial-period limited measurements, period-by-period diagnostics. | Mixture-heavy Gaussian-measurement models, fast start values for CHS, Stage-1/2 structural diagnostics. | ## Reading the matrix diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index a7f22b12..d9fc21e0 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -11,6 +11,7 @@ hatch for callers who want to drive the optimiser themselves. """ +import numpy as np import optimagic as om import pandas as pd from beartype import beartype @@ -22,6 +23,7 @@ from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.constraints import ( enforce_fixed_constraints, + project_start_to_linear_equalities, reconcile_start_to_equality, ) from skillmodels.common.identification import ( @@ -40,6 +42,8 @@ def estimate_chs( start_params: pd.DataFrame | None = None, fixed_params: pd.DataFrame | None = None, constraints: list[om.constraints.Constraint] | None = None, + lower_bounds: pd.Series | None = None, + upper_bounds: pd.Series | None = None, *, require_initial_anchors: bool = True, ) -> CHSEstimationResult: @@ -64,6 +68,15 @@ def estimate_chs( fixed. constraints: Optional extra `optimagic` constraints, appended to the model-implied constraints from `get_maximization_inputs`. + lower_bounds: Optional lower bounds, indexed like the `params_template`, + that replace the template's bounds for those parameters: a floor on + a measurement SD that would otherwise collapse, or `-inf` to drop a + bound (optimagic counts each bound on a parameter in a + `LinearConstraint` as one more linear constraint). Start values + below them are raised onto them. + upper_bounds: Optional upper bounds, the counterpart of `lower_bounds` + (e.g. `0` to restrict a loading's sign). Start values above them + are lowered onto them. require_initial_anchors: When True (the default), run the estimator-agnostic period-0 anchor check (`fail_if_initial_state_unanchored`) and raise `ValueError` if the @@ -101,6 +114,22 @@ def estimate_chs( overlay = start_params["value"].reindex(start.index) start.loc[overlay.notna(), "value"] = overlay[overlay.notna()] + if lower_bounds is not None: + replacement = lower_bounds.reindex(start.index) + locs = replacement.dropna().index + start.loc[locs, "lower_bound"] = replacement[locs] + start.loc[locs, "value"] = np.maximum( + start.loc[locs, "value"], start.loc[locs, "lower_bound"] + ) + + if upper_bounds is not None: + replacement = upper_bounds.reindex(start.index) + locs = replacement.dropna().index + start.loc[locs, "upper_bound"] = replacement[locs] + start.loc[locs, "value"] = np.minimum( + start.loc[locs, "value"], start.loc[locs, "upper_bound"] + ) + all_constraints = [*max_inputs["constraints"], *(constraints or [])] # Write every `FixedConstraintWithValue`'s target into the start vector. @@ -118,11 +147,13 @@ def estimate_chs( # group's seeded value onto the constraint surface -- honouring any fixed # member's enforced value so the line above is not averaged away. start = reconcile_start_to_equality(start, all_constraints) + start = project_start_to_linear_equalities(start, all_constraints) optimize_options = { "algorithm": options.optimizer_algorithm, "algo_options": to_plain_dict(options.optimizer_options) or None, "fun_and_jac": max_inputs["loglike_and_gradient"], + "error_handling": options.error_handling, } # Default to OPG/jacobian-based inference (`hessian=False`): the numerical diff --git a/src/skillmodels/chs/kalman_filters.py b/src/skillmodels/chs/kalman_filters.py index 72360cd9..ad47fd6a 100644 --- a/src/skillmodels/chs/kalman_filters.py +++ b/src/skillmodels/chs/kalman_filters.py @@ -131,6 +131,13 @@ def kalman_update( log_mixture_weights, ) + if n_mixtures >= 2: + new_states, new_upper_chols = reset_negligible_mixture_components( + states=new_states, + upper_chols=new_upper_chols, + log_mixture_weights=new_log_mixture_weights, + ) + return ( new_states, new_upper_chols, @@ -139,6 +146,43 @@ def kalman_update( ) +def reset_negligible_mixture_components( + states: Array, + upper_chols: Array, + log_mixture_weights: Array, +) -> tuple[Array, Array]: + """Give components whose weight underflows the dominant component's state. + + A component whose weight is below the smallest normal float contributes exactly + nothing to the likelihood, but its state is still propagated. Under a nonlinear + (e.g. translog) transition it can then explode to inf over later periods, and + the gradient through its zero weight becomes 0 * inf = NaN for every parameter. + Overwriting its state and covariance with the dominant component's keeps all + values finite; the likelihood is unchanged and the gradient is exact up to the + underflowed weight. The weights themselves are left untouched. + + Args: + states: Array of shape (n_obs, n_mixtures, n_states). + upper_chols: Array of shape (n_obs, n_mixtures, n_states, n_states). + log_mixture_weights: Array of shape (n_obs, n_mixtures). + + Returns: + states and upper_chols with negligible components replaced. + + """ + log_weights = jax.lax.stop_gradient(log_mixture_weights) + negligible = log_weights < jnp.log(jnp.finfo(log_weights.dtype).tiny) + dominant = jnp.argmax(log_weights, axis=1) + dominant_states = jnp.take_along_axis(states, dominant[:, None, None], axis=1) + dominant_chols = jnp.take_along_axis( + upper_chols, dominant[:, None, None, None], axis=1 + ) + return ( + jnp.where(negligible[..., None], dominant_states, states), + jnp.where(negligible[..., None, None], dominant_chols, upper_chols), + ) + + def calculate_sigma_scaling_factor_and_weights( n_states: int, kappa: float = 2, @@ -175,9 +219,19 @@ def kalman_predict( anchoring_scaling_factors: Array, anchoring_constants: Array, observed_factors: Array, + observed_factor_vars: Array | None = None, + *, + missing_observed_indices: tuple[int, ...] = (), ) -> tuple[Array, Array]: """Make a unscented Kalman predict. + Observed factors at `missing_observed_indices` may be missing. They augment + the sigma-point input as extra Gaussian dimensions, independent of the latent + states, with mean `observed_factors` (the fill mean where missing) and + variance `observed_factor_vars` (0 where observed, so observed values pass + through unchanged). The number of sigma points, and hence `sigma_weights`, + is based on `n_latent + len(missing_observed_indices)` for every observation. + Args: transition_func: The transition function. states: Array of shape (n_obs, n_mixtures, n_states) with @@ -201,17 +255,31 @@ def kalman_predict( period, the second to the output period (i.e. input period + 1). observed_factors: Array of shape (n_obs, n_observed_factors) with data on the observed factors in period t. + observed_factor_vars: Array of shape + (n_obs, len(missing_observed_indices)) with the variance of each + possibly-missing observed factor (0 where it is observed). + missing_observed_indices: Positions (among the observed factors) of the + observed factors that may be missing. Returns: jax.numpy.array: Predicted states, same shape as states. jax.numpy.array: Predicted upper_chols, same shape as upper_chols. """ + if missing_observed_indices: + states, upper_chols = _augment_by_missing_observed_factors( + states=states, + upper_chols=upper_chols, + observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, + missing_observed_indices=missing_observed_indices, + ) sigma_points = _calculate_sigma_points( states=states, upper_chols=upper_chols, scaling_factor=sigma_scaling_factor, observed_factors=observed_factors, + missing_observed_indices=missing_observed_indices, ) transformed = transform_sigma_points( sigma_points=sigma_points, @@ -248,10 +316,12 @@ def linear_kalman_predict( anchoring_scaling_factors: Array, anchoring_constants: Array, observed_factors: Array, + observed_factor_vars: Array | None = None, *, latent_factors: tuple[str, ...], constant_factor_indices: frozenset[int], n_all_factors: int, + missing_observed_indices: tuple[int, ...] = (), ) -> tuple[Array, Array]: """Make a linear Kalman predict (square-root form). @@ -263,6 +333,13 @@ def linear_kalman_predict( `sigma_weights` are accepted for signature compatibility with `kalman_predict` but are ignored. + Observed factors at `missing_observed_indices` may be missing, in which case + they are `N(m, v)` independent of the latent states, with `m` given in + `observed_factors` and `v` in `observed_factor_vars`. The prediction is exact: + the mean uses `m`, and the covariance adds `B diag(v) B'` with `B` the + unanchored coefficients of those factors, so the cross-covariance induced by + the same factor entering several equations is kept. + Args: transition_func: Ignored (kept for signature compatibility). states: Array of shape (n_obs, n_mixtures, n_states). @@ -274,9 +351,14 @@ def linear_kalman_predict( anchoring_scaling_factors: Array of shape (2, n_states). anchoring_constants: Array of shape (2, n_states). observed_factors: Array of shape (n_obs, n_observed_factors). + observed_factor_vars: Array of shape + (n_obs, len(missing_observed_indices)) with the variance of each + possibly-missing observed factor (0 where it is observed). latent_factors: Tuple of latent factor names. constant_factor_indices: Indices of factors with `constant` transition. n_all_factors: Total number of factors (latent + observed). + missing_observed_indices: Positions (among the observed factors) of the + observed factors that may be missing. Returns: Predicted states, same shape as states. @@ -316,9 +398,23 @@ def linear_kalman_predict( chol_g = upper_chols @ g_mat.T # (n_obs, n_mix, n_latent, n_latent) shock_diag = jnp.diag(shock_sds / s_out) # (n_latent, n_latent) - stack = jnp.concatenate( - [chol_g, jnp.broadcast_to(shock_diag, chol_g.shape)], axis=-2 - ) # (n_obs, n_mix, 2*n_latent, n_latent) + blocks = [chol_g, jnp.broadcast_to(shock_diag, chol_g.shape)] + if missing_observed_indices: + if observed_factor_vars is None: + msg = "observed_factor_vars is required with missing_observed_indices." + raise ValueError(msg) + # Row j is sqrt(v_j) * B[:, j], so the rows add B diag(v) B' to the cov. + b_missing = f_mat[:, n_latent + jnp.array(missing_observed_indices)] + b_missing = b_missing / s_out[:, jnp.newaxis] # (n_latent, n_missing) + # Promote before the root: loader moments are float32, while the + # covariance factor can be float64. A later cast cannot undo root rounding. + missing_sds = jnp.sqrt(observed_factor_vars.astype(chol_g.dtype)) + rows = missing_sds[..., jnp.newaxis] * b_missing.T + blocks.append( + jnp.broadcast_to(rows[:, jnp.newaxis], (n_obs, n_mix, *b_missing.T.shape)) + ) + # Stack shape is (n_obs, n_mix, 2 * n_latent + n_missing, n_latent). + stack = jnp.concatenate(blocks, axis=-2) predicted_covs = array_qr_jax(stack)[1][:, :, :n_latent] @@ -377,6 +473,7 @@ def _calculate_sigma_points( upper_chols: Array, scaling_factor: float | Array, observed_factors: Array, + missing_observed_indices: tuple[int, ...] = (), ) -> Array: """Calculate the array of sigma_points for the unscented transform. @@ -391,6 +488,9 @@ def _calculate_sigma_points( the sigma_point algorithm chosen. observed_factors: Array of shape (n_obs, n_observed_factors) with data on the observed factors in period t. + missing_observed_indices: Positions (among the observed factors) of the + observed factors held in the trailing `len(missing_observed_indices)` + entries of `states`. Their sigma-point values replace the data. Returns: jax.numpy.array: Array of shape n_obs, n_mixtures, n_sigma, n_fac (where n_sigma @@ -400,6 +500,7 @@ def _calculate_sigma_points( n_obs, n_mixtures, n_fac = states.shape n_sigma = 2 * n_fac + 1 n_observed = observed_factors.shape[1] + n_missing = len(missing_observed_indices) scaled_upper_chols = upper_chols * scaling_factor sigma_points = jnp.repeat(states, n_sigma, axis=1).reshape( @@ -415,10 +516,62 @@ def _calculate_sigma_points( observed_factors[:, jnp.newaxis, jnp.newaxis, :], (n_obs, n_mixtures, n_sigma, n_observed), ) + if n_missing: + # Promote first: observed factors are stored as float32, and writing the + # float64 sigma-point values into them would round those values. + observed_part = observed_part.astype(sigma_points.dtype) + observed_part = observed_part.at[..., jnp.array(missing_observed_indices)].set( + sigma_points[..., n_fac - n_missing :] + ) + sigma_points = sigma_points[..., : n_fac - n_missing] return jnp.concatenate([sigma_points, observed_part], axis=-1) +def _augment_by_missing_observed_factors( + states: Array, + upper_chols: Array, + observed_factors: Array, + observed_factor_vars: Array | None, + missing_observed_indices: tuple[int, ...], +) -> tuple[Array, Array]: + """Append possibly-missing observed factors as independent Gaussian states. + + Args: + states: Array of shape (n_obs, n_mixtures, n_latent). + upper_chols: Array of shape (n_obs, n_mixtures, n_latent, n_latent). + observed_factors: Array of shape (n_obs, n_observed_factors) with the fill + mean where a factor is missing. + observed_factor_vars: Array of shape (n_obs, n_missing) with the fill + variance where a factor is missing and 0 where it is observed. + missing_observed_indices: Positions of the possibly-missing factors among + the observed factors. + + Returns: + Augmented states of shape (n_obs, n_mixtures, n_latent + n_missing) and + block-diagonal upper Cholesky factors of matching shape. + + """ + if observed_factor_vars is None: + msg = "observed_factor_vars is required with missing_observed_indices." + raise ValueError(msg) + n_obs, n_mixtures, n_latent = states.shape + n_missing = len(missing_observed_indices) + n_aug = n_latent + n_missing + means = observed_factors[:, jnp.array(missing_observed_indices)] + aug_states = jnp.concatenate( + [states, jnp.broadcast_to(means[:, None], (n_obs, n_mixtures, n_missing))], + axis=-1, + ) + # Preserve the covariance-factor precision during the root, not just scatter. + sds = jnp.sqrt(observed_factor_vars.astype(upper_chols.dtype)) + aug_chols = jnp.zeros((n_obs, n_mixtures, n_aug, n_aug), dtype=upper_chols.dtype) + aug_chols = aug_chols.at[..., :n_latent, :n_latent].set(upper_chols) + diag = jnp.arange(n_latent, n_aug) + aug_chols = aug_chols.at[..., diag, diag].set(sds[:, None, :]) + return aug_states, aug_chols + + def transform_sigma_points( sigma_points: Array, transition_func: Callable, diff --git a/src/skillmodels/chs/kalman_filters_debug.py b/src/skillmodels/chs/kalman_filters_debug.py index c7dd20a8..27d49665 100644 --- a/src/skillmodels/chs/kalman_filters_debug.py +++ b/src/skillmodels/chs/kalman_filters_debug.py @@ -6,6 +6,8 @@ import jax.numpy as jnp from jax import Array +from skillmodels.chs.kalman_filters import reset_negligible_mixture_components + array_qr_jax = jax.vmap(jax.vmap(jnp.linalg.qr)) @@ -118,6 +120,13 @@ def kalman_update( log_mixture_weights, ) + if n_mixtures >= 2: + new_states, new_upper_chols = reset_negligible_mixture_components( + states=new_states, + upper_chols=new_upper_chols, + log_mixture_weights=new_log_mixture_weights, + ) + debug_info = {} residuals = jnp.where(not_missing.reshape(n_obs, 1), _residuals, jnp.nan) debug_info["residuals"] = residuals diff --git a/src/skillmodels/chs/likelihood.py b/src/skillmodels/chs/likelihood.py index 2f9a5225..463bf7b9 100644 --- a/src/skillmodels/chs/likelihood.py +++ b/src/skillmodels/chs/likelihood.py @@ -36,6 +36,7 @@ def log_likelihood( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> Array: """Aggregated log likelihood of a skill formation model. @@ -66,6 +67,9 @@ def log_likelihood( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: Scalar aggregated log likelihood. @@ -86,6 +90,7 @@ def log_likelihood( is_predict_iteration=is_predict_iteration, iteration_to_period=iteration_to_period, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ).sum() @@ -104,6 +109,7 @@ def log_likelihood_obs( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> Array: """Log likelihood of a skill formation model. @@ -143,6 +149,9 @@ def log_likelihood_obs( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: jnp.array: 1d array of length N, the aggregated log likelihood. @@ -181,6 +190,7 @@ def log_likelihood_obs( sigma_weights=sigma_weights, predict_func=predict_func, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) _body = jax.checkpoint(_body, prevent_cse=False) static_out = jax.lax.scan(_body, carry, loop_args, unroll=False)[1] @@ -205,6 +215,7 @@ def _scan_body( sigma_weights: Array, predict_func: Callable, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> tuple[dict[str, Array], dict[str, Array]]: # ================================================================================== # create arguments needed for update @@ -250,6 +261,7 @@ def _scan_body( ], "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])], "observed_factors": observed_factors[t], + "observed_factor_vars": observed_factor_vars[t], } fixed_kwargs = {"predict_func": predict_func} diff --git a/src/skillmodels/chs/likelihood_debug.py b/src/skillmodels/chs/likelihood_debug.py index 6a77bc27..88a80481 100644 --- a/src/skillmodels/chs/likelihood_debug.py +++ b/src/skillmodels/chs/likelihood_debug.py @@ -36,6 +36,7 @@ def log_likelihood( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> dict[str, Any]: """Log likelihood of a skill formation model, returning debug data on top. @@ -65,6 +66,9 @@ def log_likelihood( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: All data relevant for debugging, e.g. the log likelihood contribution of @@ -104,6 +108,7 @@ def log_likelihood( sigma_weights=sigma_weights, predict_func=predict_func, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) static_out = jax.lax.scan(_body, carry, loop_args)[1] @@ -156,6 +161,7 @@ def _scan_body( sigma_weights: Array, predict_func: Callable[..., tuple[Array, Array]], observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> tuple[dict[str, Array], dict[str, Any]]: # ================================================================================== # create arguments needed for update @@ -201,6 +207,7 @@ def _scan_body( ], "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])], "observed_factors": observed_factors[t], + "observed_factor_vars": observed_factor_vars[t], } fixed_kwargs = {"predict_func": predict_func} diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 566bb052..e51cc2ff 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -40,6 +40,7 @@ get_constraints, project_to_probability_constraints, ) +from skillmodels.common.measurement_models import GaussianMeasurement from skillmodels.common.model_spec import ModelSpec from skillmodels.common.params_index import get_params_index from skillmodels.common.parse_params import create_parsing_info @@ -54,9 +55,8 @@ def get_maximization_inputs( # noqa: C901, PLR0915 model_spec: ModelSpec, data: pd.DataFrame, - split_dataset: int = 1, *, - chs_options: CHSEstimationOptions | None = None, + chs_options: CHSEstimationOptions, fixed_params: pd.DataFrame | None = None, ) -> dict[str, Any]: """Create inputs for optimagic's maximize function. @@ -64,10 +64,11 @@ def get_maximization_inputs( # noqa: C901, PLR0915 Args: model_spec: The model specification. See: :ref:`model_specs` data: Dataset in long format. - split_dataset: Controls into how many slices to split the dataset - during the gradient computation. - chs_options: CHS-specific tuning parameters. Defaults to - ``CHSEstimationOptions()`` when not provided. + chs_options: CHS-specific tuning parameters. Its `start_params_strategy` + decides how the `params_template` values are seeded; pass + `CHSEstimationOptions(start_params_strategy="none")` when only the + likelihood functions or the bare template are needed, since the + `"amn"` strategy runs a full AMN estimation. fixed_params: Optional DataFrame with a ``"value"`` column pinning specified parameters to fixed values. Uses the same 4-level MultiIndex as the returned ``params_template``. Each matching @@ -101,7 +102,7 @@ def get_maximization_inputs( # noqa: C901, PLR0915 endogenous factors, we double up the number of periods in order to add """ - chs_options = chs_options or CHSEstimationOptions() + _fail_if_non_gaussian_measurements(model_spec) processed_model = process_model(model_spec) p_index = get_params_index( update_info=processed_model.update_info, @@ -127,8 +128,26 @@ def get_maximization_inputs( # noqa: C901, PLR0915 purpose="estimation", ) + # A physical observed input is repeated in both endogenous half-steps. + # Marginalizing it separately in each half drops its posterior/cross-covariance. + # Until it is carried jointly, reject uncertain inputs in prediction periods. + # The final two augmented rows have no outgoing transition. + if processed_model.endogenous_factors_info.has_endogenous_factors and np.any( + np.asarray(processed_data["observed_factor_vars"][:-2]) > 0 + ): + msg = ( + "CHS endogenous half-steps require a " + "persistent joint missing-observed-factor state for " + "positive-variance missing inputs. " + "This combination is not yet supported; zero-variance fills and " + "complete observed factors remain supported." + ) + raise NotImplementedError(msg) + + # Possibly-missing observed factors augment the sigma-point input. sigma_scaling_factor, sigma_weights = calculate_sigma_scaling_factor_and_weights( - n_states=processed_model.dimensions.n_latent_factors, + n_states=processed_model.dimensions.n_latent_factors + + len(processed_data["missing_observed_factors"]), kappa=chs_options.sigma_points_scale, ) @@ -149,6 +168,8 @@ def get_maximization_inputs( # noqa: C901, PLR0915 measurements=processed_data["measurements"], controls=processed_data["controls"], observed_factors=processed_data["observed_factors"], + observed_factor_vars=processed_data["observed_factor_vars"], + missing_observed_factors=processed_data["missing_observed_factors"], model=processed_model, sigma_weights=sigma_weights, sigma_scaling_factor=sigma_scaling_factor, @@ -175,6 +196,7 @@ def loglike_and_gradient( n_obs = processed_data["measurements"].shape[1] _grad = jnp.zeros_like(params_vec) start = 0 + split_dataset = chs_options.split_dataset stop = int(n_obs / split_dataset) step = int(n_obs / split_dataset) for i in range(split_dataset): @@ -189,10 +211,21 @@ def loglike_and_gradient( measurements=measurements_slice, controls=controls_slice, observed_factors=observed_factors_slice, + observed_factor_vars=processed_data["observed_factor_vars"][ + :, start:stop, : + ], ) start += step stop += step grad = _to_numpy(_grad) + if not (np.isfinite(crit) and np.isfinite(grad).all()): + # Raise rather than return, so `error_handling="continue"` can replace + # the point by a penalty that rejects the step. + msg = ( + f"The log-likelihood ({crit}) or its gradient " + f"({int((~np.isfinite(grad)).sum())} entries) is non-finite." + ) + raise FloatingPointError(msg) return crit, grad def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: @@ -262,6 +295,27 @@ def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: } +def _fail_if_non_gaussian_measurements(model_spec: ModelSpec) -> None: + """Reject probit/Tobit measures: the CHS update is a Gaussian Kalman update. + + Such a measure would otherwise be silently estimated as a continuous Gaussian one. + """ + non_gaussian = sorted( + name + for name, model in model_spec.measurement_models.items() + if not isinstance(model, GaussianMeasurement) + ) + if non_gaussian: + msg = ( + f"CHS supports only Gaussian measurements, but {non_gaussian} are " + "non-Gaussian (probit/Tobit). The CHS likelihood runs a Gaussian Kalman " + "update for every measurement. Declare them Gaussian (a linear " + "probability / linear working model) or use estimate_af, which honours " + "probit/Tobit families in the initial period." + ) + raise NotImplementedError(msg) + + def _estimate_amn_for_chs_seeding( model_spec: ModelSpec, data: pd.DataFrame, @@ -378,6 +432,8 @@ def _partial_some_log_likelihood( measurements: Array, controls: Array, observed_factors: Array, + observed_factor_vars: Array, + missing_observed_factors: tuple[str, ...], model: ProcessedModel, sigma_weights: Array, sigma_scaling_factor: Array, @@ -415,6 +471,9 @@ def _partial_some_log_likelihood( # states (cf is a non-linear function of the latent states), so the linear # fast-path would silently bypass the cf DAG node. Force the unscented path. has_control_function = model.endogenous_factors_info.control_function is not None + missing_observed_indices = tuple( + model.labels.observed_factors.index(f) for f in missing_observed_factors + ) if is_all_linear(model.transition_info.function_names) and not has_control_function: constant_factor_indices = frozenset( i @@ -427,9 +486,14 @@ def _partial_some_log_likelihood( latent_factors=model.labels.latent_factors, constant_factor_indices=constant_factor_indices, n_all_factors=model.dimensions.n_all_factors, + missing_observed_indices=missing_observed_indices, ) else: - predict_func = functools.partial(kalman_predict, model.transition_info.func) + predict_func = functools.partial( + kalman_predict, + model.transition_info.func, + missing_observed_indices=missing_observed_indices, + ) return functools.partial( fun, @@ -446,6 +510,7 @@ def _partial_some_log_likelihood( is_predict_iteration=is_predict_iteration, iteration_to_period=iteration_to_period, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) diff --git a/src/skillmodels/chs/options.py b/src/skillmodels/chs/options.py index 48ad358b..f8058df1 100644 --- a/src/skillmodels/chs/options.py +++ b/src/skillmodels/chs/options.py @@ -20,6 +20,10 @@ class CHSEstimationOptions: """Distance for bounds. Zeroed out if `robust_bounds` is False.""" sigma_points_scale: float = 2 """Scaling factor for sigma points in unscented transform.""" + split_dataset: int = 1 + """Number of slices of the individuals over which the gradient is accumulated. + + More slices lower the peak memory of the gradient; the result is the same.""" clipping_lower_bound: float = -1e30 """Lower bound for soft clipping.""" clipping_upper_bound: float | None = None @@ -59,6 +63,14 @@ class CHSEstimationOptions: Forwarded by `estimate_chs` to `estimate_ml` (e.g. convergence tolerances, trust-region settings for `fides`).""" + error_handling: Literal["raise", "continue"] = "raise" + """What `optimagic` does when the criterion raises (`optimize_options`). + + The criterion raises whenever the log-likelihood or its gradient is + non-finite. `"raise"` aborts the estimation. `"continue"` replaces that + point by a penalty worse than the start value, so a trust-region step into + a numerically explosive region is rejected and the radius shrinks.""" + estimate_ml_options: Mapping[str, Any] = field( default_factory=lambda: MappingProxyType({}) ) diff --git a/src/skillmodels/chs/qr.py b/src/skillmodels/chs/qr.py index d42862f7..d9ac11ee 100644 --- a/src/skillmodels/chs/qr.py +++ b/src/skillmodels/chs/qr.py @@ -53,6 +53,28 @@ def _tril(m: Array, k: int = 0) -> Array: return jax.lax.select(jax.lax.broadcast(mask, m.shape[:-2]), m, jnp.zeros_like(m)) +# Floor for the magnitude of R's diagonal in the tangent solve below. The QR +# derivative solves against R, so a zero on R's diagonal -- a rank-deficient +# input, which the square-root UKF produces when a predicted covariance collapses +# as the optimiser inflates a nonlinear transition -- would send the tangent to +# +/-inf and poison the entire gradient with NaN, aborting the optimisation. +_QR_JVP_DIAGONAL_FLOOR = 1e-6 + + +def _floor_diagonal_magnitude(r: Array, floor: float) -> Array: + """Return `r` with each diagonal entry pushed to magnitude at least `floor`. + + Off-diagonal entries and the sign of each diagonal entry are preserved; a + diagonal entry already larger than `floor` in magnitude is untouched, so a + well-conditioned R (and its gradient) is unchanged. + """ + n = r.shape[-1] + diag = jnp.diagonal(r, axis1=-2, axis2=-1) + sign = jnp.where(diag < 0.0, -1.0, 1.0) + safe = jnp.where(jnp.abs(diag) < floor, sign * floor, diag) + return r + (safe - diag)[..., :, None] * jnp.eye(n, dtype=r.dtype) + + @qr_gpu.defjvp def qr_jvp_rule( primals: tuple[Array], @@ -63,7 +85,13 @@ def qr_jvp_rule( (x,) = primals (dx,) = tangents q, r = qr_gpu(x) - dx_rinv = jax.lax.linalg.triangular_solve(r, dx) # Right side solve by default + # Solve against a diagonal-floored R so a (near-)singular decomposition yields + # a large but finite tangent instead of NaN. The primal `q, r` and every other + # term below use the unmodified `r`. + r_solve = _floor_diagonal_magnitude(r, _QR_JVP_DIAGONAL_FLOOR) + dx_rinv = jax.lax.linalg.triangular_solve( + r_solve, dx + ) # Right side solve by default qt_dx_rinv = _h(q) @ dx_rinv qt_dx_rinv_lower = _tril(m=qt_dx_rinv, k=-1) do = qt_dx_rinv_lower - _h(qt_dx_rinv_lower) # This is skew-symmetric diff --git a/src/skillmodels/common/constraints.py b/src/skillmodels/common/constraints.py index 285aac29..6a78e390 100644 --- a/src/skillmodels/common/constraints.py +++ b/src/skillmodels/common/constraints.py @@ -196,6 +196,43 @@ def reconcile_start_to_equality( return out +def project_start_to_linear_equalities( + params: pd.DataFrame, + constraints: Iterable[om.constraints.Constraint], +) -> pd.DataFrame: + """Move the start point onto every linear equality constraint. + + `om.minimize` raises `InvalidParamsError` when a constraint is violated at the + starting point. All `om.LinearConstraint`s with a `value` are stacked into one + system `A x = b` over the parameters they involve, and those parameters' `value` + is replaced by its orthogonal projection onto the solution set (e.g. the + least-squares line for second-difference constraints). Other parameters are + untouched. Returns a copy; `params` is not modified. + """ + rows = [] + for c in constraints: + if isinstance(c, om.LinearConstraint) and c.value is not None: + selected = pd.Series(c.selector(params)).index + weights = np.broadcast_to( + np.asarray(c.weights, dtype=float), (len(selected),) + ) + rows.append((selected, weights, float(c.value))) + if not rows: + return params + involved = pd.Index( + sorted({tup for idx, _, _ in rows for tup in idx}), tupleize_cols=False + ) + a = np.zeros((len(rows), len(involved))) + b = np.array([value for _, _, value in rows]) + for i, (idx, weights, _) in enumerate(rows): + a[i, involved.get_indexer(idx)] = weights + x = params.loc[list(involved), "value"].to_numpy(dtype=float) + projected = x - np.linalg.pinv(a) @ (a @ x - b) + out = params.copy() + out.loc[list(involved), "value"] = projected + return out + + def collect_fixed_locs( constraints: Iterable[om.constraints.Constraint], ) -> set[tuple[Any, ...]]: diff --git a/src/skillmodels/common/correlation_heatmap.py b/src/skillmodels/common/correlation_heatmap.py index 9225e4b0..595127b8 100644 --- a/src/skillmodels/common/correlation_heatmap.py +++ b/src/skillmodels/common/correlation_heatmap.py @@ -411,6 +411,8 @@ def _get_measurement_data_for_single_period( )["variable"].to_list() for fac in observed_factors: measurements.append(fac) + # A measurement that loads on several factors is one column. + measurements = list(dict.fromkeys(measurements)) return data.query(f"{update_info_by_period.index.names[0]}=={period}")[measurements] diff --git a/src/skillmodels/common/individual_states.py b/src/skillmodels/common/individual_states.py index 9d073c04..ee884fff 100644 --- a/src/skillmodels/common/individual_states.py +++ b/src/skillmodels/common/individual_states.py @@ -10,6 +10,7 @@ result machinery, so there is no raw-params path for them). """ +from dataclasses import replace from typing import Any import pandas as pd @@ -23,6 +24,7 @@ from skillmodels.af.types import AFEstimationResult from skillmodels.amn.types import AMNEstimationResult from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.anchoring import anchor_states_df from skillmodels.common.model_spec import ModelSpec @@ -34,6 +36,8 @@ def get_individual_states( data: pd.DataFrame, result: CHSEstimationResult | AFEstimationResult | AMNEstimationResult, + *, + chs_options: CHSEstimationOptions | None = None, ) -> dict[str, dict[str, Any]]: """Compute per-individual latent state estimates from a fitted result. @@ -47,6 +51,9 @@ def get_individual_states( data: Dataset in long format with MultiIndex (id, period). result: A `CHSEstimationResult`, `AFEstimationResult`, or `AMNEstimationResult`. + chs_options: The options the CHS result was estimated with, so the + filter uses the same likelihood settings. Required for a + `CHSEstimationResult`; unused for AF and AMN. Return: Nested dict. The return shape is **asymmetric** across estimators: @@ -78,10 +85,17 @@ def get_individual_states( data=data, ) + if chs_options is None: + msg = ( + "get_individual_states needs `chs_options` for a CHSEstimationResult: " + "pass the CHSEstimationOptions used for the estimation." + ) + raise TypeError(msg) return get_individual_states_from_params( model_spec=result.model_spec, data=data, params=result.params, + chs_options=chs_options, ) @@ -90,6 +104,8 @@ def get_individual_states_from_params( model_spec: ModelSpec, data: pd.DataFrame, params: pd.DataFrame, + *, + chs_options: CHSEstimationOptions, ) -> dict[str, dict[str, Any]]: """Compute CHS filtered states directly from raw parameters. @@ -102,13 +118,20 @@ def get_individual_states_from_params( model_spec: Model specification. data: Dataset in long format with MultiIndex (id, period). params: Estimated parameter DataFrame (standard 4-level MultiIndex). + chs_options: The options used for the estimation, so the filter uses the + same likelihood settings (e.g. the sigma-point scale). Its + `start_params_strategy` is ignored: no start values are seeded. Return: Dict with both `"anchored_states"` and `"unanchored_states"`, each `{"states": DataFrame, "state_ranges": ...}`. """ - max_inputs = get_maximization_inputs(model_spec=model_spec, data=data) + max_inputs = get_maximization_inputs( + model_spec=model_spec, + data=data, + chs_options=replace(chs_options, start_params_strategy="none"), + ) params = params.loc[max_inputs["params_template"].index] debug_loglike = max_inputs["debug_loglike"] debug_data = debug_loglike(params) diff --git a/src/skillmodels/common/measurement_models.py b/src/skillmodels/common/measurement_models.py index 9a500c95..b955d900 100644 --- a/src/skillmodels/common/measurement_models.py +++ b/src/skillmodels/common/measurement_models.py @@ -2,9 +2,10 @@ One JAX-compatible kernel evaluates the log density/probability contribution of a single measurement given its linear predictor ``eta = c + x'beta + lambda'theta``, -its scale, and a `MeasurementFamily` code. AF estimation, CHS, simulation and -posterior-state reweighting all route their measurement contributions through this -kernel, so the measurement system is defined in exactly one place. +its scale, and a `MeasurementFamily` code. AF estimation, simulation and +posterior-state reweighting route their measurement contributions through this +kernel. CHS (a Gaussian Kalman update) and AMN do not, and reject non-Gaussian +measurements. The three families: diff --git a/src/skillmodels/common/model_spec.py b/src/skillmodels/common/model_spec.py index b9f57d9a..0c886cb2 100644 --- a/src/skillmodels/common/model_spec.py +++ b/src/skillmodels/common/model_spec.py @@ -22,6 +22,20 @@ Normalizations, ensure_containers_are_immutable, ) +from skillmodels.exceptions import ModelSpecInitializationError + +_ALLOWED_MODEL_SPEC_KEYS = frozenset( + { + "factors", + "observed_factors", + "controls", + "stagemap", + "anchoring", + "n_mixtures", + "measurement_models", + "missing_observed_factor_moments", + } +) def _parse_measurement_models( @@ -211,6 +225,24 @@ class ModelSpec: ) """Per-variable observation model (probit / Tobit / Gaussian). A measurement omitted from this mapping is Gaussian, so existing specs are unchanged.""" + missing_observed_factor_moments: MappingProxyType[str, tuple[str, str]] = field( + default_factory=lambda: MappingProxyType({}) + ) + """Map an observed factor that may be missing to its `(mean, variance)` columns. + + Where the factor is missing for a person-period, the CHS filter treats it as + `N(mean, variance)`, independent of the latent states, with both moments read + from that person-period's row of the data, and integrates it out in the + predict step. Where the factor is observed, the moment columns are ignored. + An observed factor not listed here must never be missing. Only the CHS + likelihood (and the filtered states built on it) integrates missing values + out; simulation and transition plots still require complete observed + factors. In CHS models with endogenous factors, positive-variance missing + inputs in a physical period with an outgoing transition are not yet + supported: both half-steps would need to carry the same missing input's + joint posterior. Such inputs raise `NotImplementedError`; zero-variance + fills, complete inputs, and terminal-only missing inputs remain supported. + """ def __init__( self, @@ -221,6 +253,7 @@ def __init__( anchoring: AnchoringSpec | None = None, n_mixtures: int = 1, measurement_models: Mapping[str, MeasurementModel] | None = None, + missing_observed_factor_moments: Mapping[str, tuple[str, str]] | None = None, ) -> None: """Create ModelSpec, wrapping factors dict in MappingProxyType.""" object.__setattr__(self, "_factors", ensure_containers_are_immutable(factors)) @@ -230,6 +263,17 @@ def __init__( object.__setattr__(self, "anchoring", anchoring) object.__setattr__(self, "n_mixtures", n_mixtures) self._set_measurement_models(measurement_models) + moments = dict(missing_observed_factor_moments or {}) + unknown = sorted(name for name in moments if name not in observed_factors) + if unknown: + msg = ( + f"missing_observed_factor_moments references {unknown}, which is not " + f"an observed factor. Observed factors: {list(observed_factors)}." + ) + raise ValueError(msg) + object.__setattr__( + self, "missing_observed_factor_moments", MappingProxyType(moments) + ) def _set_measurement_models( self, measurement_models: Mapping[str, MeasurementModel] | None @@ -272,7 +316,21 @@ def from_dict(cls, d: dict[str, Any]) -> Self: Returns: A ModelSpec instance. + Raises: + ModelSpecInitializationError: If `d` contains an unrecognized top-level + key. Silently dropping unknown keys hides typos and stale config + blocks (e.g. an `estimation_options` mapping whose settings never + take effect). + """ + unknown = set(d) - _ALLOWED_MODEL_SPEC_KEYS + if unknown: + msg = ( + f"Unrecognized ModelSpec keys: {sorted(unknown)}. " + f"Allowed keys: {sorted(_ALLOWED_MODEL_SPEC_KEYS)}." + ) + raise ModelSpecInitializationError(msg) + factors = {} auto_instruments: list[str] = [] for name, spec in d["factors"].items(): @@ -324,6 +382,12 @@ def from_dict(cls, d: dict[str, Any]) -> Self: observed += tuple(i for i in auto_instruments if i not in observed) measurement_models = _parse_measurement_models(d.get("measurement_models")) + moments = { + name: (mean_col, var_col) + for name, (mean_col, var_col) in d.get( + "missing_observed_factor_moments", {} + ).items() + } return cls( factors=factors, @@ -333,6 +397,7 @@ def from_dict(cls, d: dict[str, Any]) -> Self: anchoring=anchoring, n_mixtures=d.get("n_mixtures", 1), measurement_models=measurement_models, + missing_observed_factor_moments=moments, ) @property @@ -352,6 +417,10 @@ def _replace(self, **changes: Any) -> Self: # noqa: ANN401 measurement_models=changes.get( "measurement_models", self.measurement_models ), + missing_observed_factor_moments=changes.get( + "missing_observed_factor_moments", + self.missing_observed_factor_moments, + ), ) def with_transition_functions( diff --git a/src/skillmodels/common/process_data.py b/src/skillmodels/common/process_data.py index 8e3f48de..73e955da 100644 --- a/src/skillmodels/common/process_data.py +++ b/src/skillmodels/common/process_data.py @@ -42,7 +42,13 @@ def process_data( observed control variables for the measurement equations. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. - Only returned if estimation==True + Where a factor listed in `labels.missing_observed_factor_moments` is + missing, it holds the fill mean. + missing_observed_factors: Names of the observed factors with fill moments + that are missing somewhere in the data (estimation only; else empty). + observed_factor_vars: Array of shape + (n_periods, n_obs, len(missing_observed_factors)) with the fill + variance where the factor is missing and 0 where it is observed. """ df = pre_process_data(df=df, periods=labels.periods) @@ -63,6 +69,20 @@ def process_data( df=df, controls=labels.controls, update_info=update_info ) out["controls"] = _generate_controls_array(df=df, labels=labels, n_obs=n_obs) + missing = tuple( + f + for f in labels.observed_factors + if purpose == "estimation" + and f in labels.missing_observed_factor_moments + and df[f].isna().any() + ) + out["missing_observed_factors"] = missing + out["observed_factor_vars"] = _generate_observed_factor_vars_array( + df=df, labels=labels, n_obs=n_obs, factors=missing + ) + for factor in missing: + mean_col, _ = labels.missing_observed_factor_moments[factor] + df[factor] = df[factor].fillna(df[mean_col]) out["observed_factors"] = _generate_observed_factor_array( df=df, labels=labels, n_obs=n_obs ) @@ -119,18 +139,27 @@ def _get_period_data_for_endogenous_factors( meas = _get_period_measurements(update_info=update_info, aug_period=aug_period) controls = labels.controls observed = labels.observed_factors - - out = df.query(f"period == {period}")[ - [ - "id", - *meas, - *controls, - *observed, - "period", - "__old_id__", - "__old_period__", - ] + moment_cols = [ + col for cols in labels.missing_observed_factor_moments.values() for col in cols ] + + # Several factor mappings may share a physical moment column, or reuse a + # column already selected above. Preserve its first position, not duplicates. + columns = list( + dict.fromkeys( + [ + "id", + *meas, + *controls, + *observed, + *moment_cols, + "period", + "__old_id__", + "__old_period__", + ] + ) + ) + out = df.query(f"period == {period}")[columns] out["aug_period"] = aug_period return out @@ -209,12 +238,14 @@ def _check_data( # noqa: C901 ) for factor in labels.observed_factors: - if factor not in period_data.columns: - var_report.loc[(aug_period, factor), "problem"] = "Variable is missing" - elif period_data[factor].isna().any(): - var_report.loc[(aug_period, factor), "problem"] = ( - "Variable has missings" - ) + moments = ( + labels.missing_observed_factor_moments.get(factor) + if purpose == "estimation" + else None + ) + problems = _observed_factor_problems(period_data, factor, moments) + for var, problem in problems.items(): + var_report.loc[(aug_period, var), "problem"] = problem var_report = var_report.to_string() if len(var_report) > 0 else "" @@ -222,6 +253,32 @@ def _check_data( # noqa: C901 raise ValueError(var_report) +def _observed_factor_problems( + period_data: pd.DataFrame, factor: str, moments: tuple[str, str] | None +) -> dict[str, str]: + """Map variables to problems with an observed factor and its fill moments. + + Without `moments`, any missing value is a problem. With `moments`, missing + values need a non-missing fill mean and a non-negative fill variance. + """ + if factor not in period_data.columns: + return {factor: "Variable is missing"} + missing = period_data[factor].isna() + if moments is None: + return {factor: "Variable has missings"} if missing.any() else {} + absent = [col for col in moments if col not in period_data.columns] + if absent: + return dict.fromkeys(absent, "Variable is missing") + mean, var = (period_data.loc[missing, col] for col in moments) + if mean.isna().any() or var.isna().any(): + return {factor: f"Variable has missings without fill moments {moments}"} + if (var < 0).any(): + return { + factor: f"Variable has missings with negative fill variance in {moments[1]}" + } + return {} + + def _handle_controls_with_missings( df: pd.DataFrame, controls: tuple[str, ...], @@ -282,6 +339,23 @@ def _generate_controls_array( return jnp.array(arr, dtype="float32") +def _generate_observed_factor_vars_array( + df: pd.DataFrame, + labels: Labels, + n_obs: int, + factors: tuple[str, ...], +) -> Array: + arr = np.zeros((len(labels.aug_periods), n_obs, len(factors))) + for aug_period in labels.aug_periods: + period_data = df.query(f"aug_period == {aug_period}") + for k, factor in enumerate(factors): + var_col = labels.missing_observed_factor_moments[factor][1] + arr[aug_period, :, k] = np.where( + period_data[factor].isna(), period_data[var_col], 0.0 + ) + return jnp.array(arr, dtype="float32") + + def _generate_observed_factor_array( df: pd.DataFrame, labels: Labels, diff --git a/src/skillmodels/common/process_model.py b/src/skillmodels/common/process_model.py index 71c504c2..31f48438 100644 --- a/src/skillmodels/common/process_model.py +++ b/src/skillmodels/common/process_model.py @@ -237,6 +237,7 @@ def _get_labels( aug_stagemap=tuple(aug_stagemap), aug_stages=tuple(sorted(int(v) for v in np.unique(aug_stagemap))), aug_stages_to_stages=MappingProxyType(aug_stages_to_stages), + missing_observed_factor_moments=model_spec.missing_observed_factor_moments, ) diff --git a/src/skillmodels/common/types.py b/src/skillmodels/common/types.py index 2cf2de35..ded3f768 100644 --- a/src/skillmodels/common/types.py +++ b/src/skillmodels/common/types.py @@ -153,6 +153,11 @@ class Labels: """Map each augmented stage to its original stage.""" transition_names: tuple[str, ...] = () """Names of the transition functions per factor.""" + missing_observed_factor_moments: MappingProxyType[str, tuple[str, str]] = field( + default_factory=lambda: MappingProxyType({}) + ) + """Map an observed factor that may be missing to its `(mean, variance)` data + columns, see `ModelSpec.missing_observed_factor_moments`.""" @property def all_factors(self) -> tuple[str, ...]: diff --git a/src/skillmodels/common/visualize_factor_distributions.py b/src/skillmodels/common/visualize_factor_distributions.py index 8f77f9e5..bef4becb 100644 --- a/src/skillmodels/common/visualize_factor_distributions.py +++ b/src/skillmodels/common/visualize_factor_distributions.py @@ -365,10 +365,11 @@ def bivariate_density_contours( factors=pair, n_points=n_points, ) + # `z[i, j]` is the density at `(x[i], y[j])`; plotly reads rows as y. contour = go.Contour( x=x[:, 0], y=y[0, :], - z=z, + z=z.T, line={"color": getattr(px.colors.qualitative, lines_colorscale)[i]}, ) fig.add_trace(contour) diff --git a/src/skillmodels/common/visualize_transition_equations.py b/src/skillmodels/common/visualize_transition_equations.py index da1d4dbd..6abd3b20 100644 --- a/src/skillmodels/common/visualize_transition_equations.py +++ b/src/skillmodels/common/visualize_transition_equations.py @@ -193,7 +193,8 @@ def get_transition_plots( aggregation_method: How to handle factors not being varied in a plot. - "quantiles": Fix other factors at specified quantiles (default). - "median": Fix other factors at their median (0.5 quantile). - - "average": Average over random draws of other factors. + - "average": Average over random draws of other factors, from the + persons with all other factors observed. n_points: Number of grid points per input. Default 50. n_draws: Number of randomly drawn values of the factors that are averaged out. Only relevant if aggregation_method is "average". Default 50. @@ -546,14 +547,19 @@ def _get_state_ranges( all_factors: tuple[str, ...], quantile_cutoff: float | None = None, ) -> dict[str, pd.DataFrame]: - """Create state ranges if none is given.""" - if state_ranges is None: - state_ranges = create_state_ranges( + """Complete the given state ranges with those of the other factors in the data.""" + given = state_ranges or {} + missing = [factor for factor in all_factors if factor not in given] + if not missing: + return given + return { + **given, + **create_state_ranges( filtered_states=states_data, - factors=list(all_factors), + factors=missing, quantile_cutoff=quantile_cutoff, - ) - return state_ranges + ), + } def _get_parsed_params( @@ -760,7 +766,14 @@ def _prepare_data_for_one_plot_average_2d( period_data = states_data.query(f"aug_period == {aug_period}") sampled_factors = [factor for factor in all_factors if factor != input_factor] - draws = period_data[sampled_factors].sample(n=n_draws) + # Persons with a missing factor (e.g. an observed input) would make every mean + # NaN; draw from the complete rows, with a fixed seed for reproducible figures. + complete = period_data[sampled_factors].dropna() + draws = complete.sample( + n=n_draws, + replace=len(complete) < n_draws, + random_state=np.random.default_rng(0), + ) input_min = state_ranges[input_factor].loc[aug_period]["minimum"] input_max = state_ranges[input_factor].loc[aug_period]["maximum"] diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index 1f2946be..4061eaf3 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -1,8 +1,10 @@ """End-to-end smoke test for `skillmodels.chs.estimate.estimate_chs`.""" +import functools from pathlib import Path import numpy as np +import optimagic as om import pandas as pd import pytest @@ -11,6 +13,7 @@ from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.config import TEST_DATA_DIR from skillmodels.common.fixed_constraint import FixedConstraintWithValue +from skillmodels.common.selector import select_by_loc from skillmodels.test_data.model2 import MODEL2 REGRESSION_VAULT = Path(__file__).parent / "regression_vault" @@ -125,3 +128,147 @@ def test_estimate_chs_enforces_user_fixed_constraint_value( ) assert result.params.loc[loc, "value"] == pytest.approx(target) + + +def test_estimate_chs_forwards_error_handling( + model2_data, anchoring_start_params, monkeypatch +): + """`error_handling` reaches optimagic so failed evaluations can be penalised.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs["optimize_options"]) + raise StopIteration + + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions( + start_params_strategy="none", error_handling="continue" + ), + start_params=anchoring_start_params, + require_initial_anchors=False, + ) + assert captured["error_handling"] == "continue" + + +def test_estimate_chs_applies_lower_bounds_and_clips_start( + model2_data, anchoring_start_params, monkeypatch +): + """`lower_bounds` tightens the template bounds; the start moves onto them.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = anchoring_start_params.query("category == 'meas_sds'").index[:1] + start = anchoring_start_params.copy() + start.loc[loc, "value"] = 0.05 + lower_bounds = pd.Series(0.1, index=loc) + + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + lower_bounds=lower_bounds, + require_initial_anchors=False, + ) + assert captured["bounds"].lower.loc[loc].tolist() == [0.1] + assert captured["params"].loc[loc, "value"].tolist() == [0.1] + other = captured["bounds"].lower.drop(loc) + assert other[np.isfinite(other)].max() < 0.1 + + +def test_estimate_chs_lower_bounds_can_remove_a_bound( + model2_data, anchoring_start_params, monkeypatch +): + """A `-inf` entry in `lower_bounds` replaces the template's bound.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = anchoring_start_params.query("category == 'meas_sds'").index[:1] + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=anchoring_start_params, + lower_bounds=pd.Series(-np.inf, index=loc), + require_initial_anchors=False, + ) + assert captured["bounds"].lower.loc[loc].tolist() == [-np.inf] + + +def test_estimate_chs_applies_upper_bounds_and_clips_start( + model2_data, anchoring_start_params, monkeypatch +): + """`upper_bounds` replaces the template bounds; the start moves onto them.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = pd.MultiIndex.from_tuples( + [("controls", 0, "y1", "x1")], names=anchoring_start_params.index.names + ) + start = anchoring_start_params.copy() + start.loc[loc, "value"] = 0.5 + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + upper_bounds=pd.Series(0.0, index=loc), + require_initial_anchors=False, + ) + assert captured["bounds"].upper.loc[loc].tolist() == [0.0] + assert captured["params"].loc[loc, "value"].tolist() == [0.0] + + +def test_estimate_chs_moves_start_onto_linear_equalities( + model2_data, anchoring_start_params, monkeypatch +): + """Start values satisfy user `LinearConstraint` equalities (a linear trend).""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + sds = anchoring_start_params.query("category == 'meas_sds' & name1 == 'y1'").index + start = anchoring_start_params.copy() + start.loc[sds, "value"] = np.linspace(0.5, 1.5, len(sds)) ** 2 + constraints: list[om.constraints.Constraint] = [ + om.LinearConstraint( + selector=functools.partial(select_by_loc, loc=list(sds[t : t + 3])), + weights=np.array([1.0, -2.0, 1.0]), + value=0.0, + ) + for t in range(len(sds) - 2) + ] + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + constraints=constraints, + lower_bounds=pd.Series(-np.inf, index=sds), + require_initial_anchors=False, + ) + values = captured["params"].loc[sds, "value"].to_numpy() + np.testing.assert_allclose(np.diff(values, n=2), 0.0, atol=1e-12) diff --git a/tests/test_chs_limited_measurement_rejection.py b/tests/test_chs_limited_measurement_rejection.py new file mode 100644 index 00000000..37ca86de --- /dev/null +++ b/tests/test_chs_limited_measurement_rejection.py @@ -0,0 +1,71 @@ +"""CHS must reject non-Gaussian measurements rather than treat them as continuous. + +The CHS likelihood runs a Gaussian Kalman update for every measurement, so a probit +or Tobit measure would be silently estimated as a continuous Gaussian one. +""" + +import numpy as np +import pandas as pd +import pytest + +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.common.measurement_models import ( + ProbitMeasurement, + TobitMeasurement, +) +from skillmodels.common.model_spec import ( + FactorSpec, + ModelSpec, + Normalizations, +) + + +def _model(measurement_models) -> ModelSpec: + return ModelSpec( + factors={ + "skills": FactorSpec( + measurements=(("y1", "y2", "y3"),) * 2, + normalizations=Normalizations( + loadings=({"y1": 1},) * 2, + intercepts=({"y1": 0},) * 2, + ), + transition_function="linear", + ), + }, + measurement_models=measurement_models, + ) + + +def _tiny_data() -> pd.DataFrame: + rng = np.random.default_rng(0) + rows = [ + { + "caseid": caseid, + "period": period, + **dict(zip(("y1", "y2", "y3"), rng.normal(size=3), strict=True)), + } + for caseid in range(20) + for period in (0, 1) + ] + return pd.DataFrame(rows).set_index(["caseid", "period"]) + + +@pytest.mark.parametrize("strategy", ["amn", "none"]) +@pytest.mark.parametrize( + "models", + [ + {"y2": ProbitMeasurement()}, + {"y3": TobitMeasurement(lower=0.0)}, + ], +) +def test_get_maximization_inputs_rejects_non_gaussian_measurements( + models, strategy +) -> None: + """Rejected whatever the start-value strategy, not only via AMN seeding.""" + with pytest.raises(NotImplementedError, match="CHS supports only Gaussian"): + get_maximization_inputs( + model_spec=_model(models), + data=_tiny_data(), + chs_options=CHSEstimationOptions(start_params_strategy=strategy), + ) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index 5fa52218..1076ca4a 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -22,6 +22,7 @@ _get_transition_constraints, add_bounds, get_constraints, + project_start_to_linear_equalities, reconcile_start_to_equality, ) from skillmodels.common.process_model import process_model @@ -787,3 +788,26 @@ def test_get_constraints_for_augmented_periods(simplest_augmented_model) -> None {"loc": ("shock_sds", 1, "fac2", "-"), "type": "fixed", "value": 0.00000001}, ] assert_list_equal_except_for_order(as_dicts, expected) + + +def test_project_start_to_linear_equalities_fits_a_line() -> None: + """Second-difference equalities make the start the least-squares line.""" + index = pd.MultiIndex.from_tuples( + [("meas_sds", t, "alive", "-") for t in range(5)] + [("loadings", 0, "y", "f")], + names=["category", "aug_period", "name1", "name2"], + ) + values = np.array([0.2, 0.25, 0.2, 0.35, 0.5, 7.0]) + params = pd.DataFrame({"value": values}, index=index) + sds = list(index[:5]) + constraints = [ + om.LinearConstraint( + selector=functools.partial(select_by_loc, loc=sds[t : t + 3]), + weights=np.array([1.0, -2.0, 1.0]), + value=0.0, + ) + for t in range(3) + ] + slope, intercept = np.polyfit(np.arange(5), values[:5], deg=1) + expected = np.append(intercept + slope * np.arange(5), 7.0) + projected = project_start_to_linear_equalities(params, constraints) + np.testing.assert_allclose(projected["value"].to_numpy(), expected, atol=1e-12) diff --git a/tests/test_correlation_heatmap.py b/tests/test_correlation_heatmap.py index 8251b5f4..100e0461 100644 --- a/tests/test_correlation_heatmap.py +++ b/tests/test_correlation_heatmap.py @@ -464,3 +464,60 @@ def test_get_scores_corr_single_period( periods=0, ) assert isinstance(result, pd.DataFrame) + + +def test_get_measurement_data_with_single_period_lists_a_shared_measurement_once() -> ( + None +): + """A measurement that loads on two factors appears as one column.""" + update_info = pd.DataFrame( + { + "f1": [True, True, False], + "f2": [False, True, True], + "purpose": "measurement", + }, + index=pd.MultiIndex.from_product( + [[0], ["y0", "y1", "y2"]], names=["stage", "variable"] + ), + ) + data = pd.DataFrame({"stage": [0, 0], "y0": [1, 2], "y1": [3, 4], "y2": [5, 6]}) + result = _get_measurement_data_for_single_period( + data, update_info, 0, latent_factors=["f1", "f2"], observed_factors=[] + ) + assert list(result.columns) == ["y0", "y1", "y2"] + + +def test_get_measurement_data_with_multiple_periods_lists_shared_measurement_once() -> ( + None +): + """A measurement that loads on two factors appears once per period.""" + update_info = pd.DataFrame( + { + "f1": [True, True, False] * 2, + "f2": [False, True, True] * 2, + "purpose": "measurement", + }, + index=pd.MultiIndex.from_product( + [[0, 1], ["y0", "y1", "y2"]], names=["stage", "variable"] + ), + ) + data = pd.DataFrame( + { + "stage": [0, 0, 1, 1], + "y0": [1, 2, 3, 4], + "y1": [3, 4, 5, 6], + "y2": [5, 6, 7, 8], + }, + index=[0, 1, 0, 1], + ) + result = _get_measurement_data_for_multiple_periods( + data, update_info, [0, 1], latent_factors=["f1", "f2"], observed_factors=[] + ) + assert list(result.columns) == [ + "y0, 0", + "y1, 0", + "y2, 0", + "y0, 1", + "y1, 1", + "y2, 1", + ] diff --git a/tests/test_filtered_states.py b/tests/test_filtered_states.py index f8024f44..58369711 100644 --- a/tests/test_filtered_states.py +++ b/tests/test_filtered_states.py @@ -6,9 +6,15 @@ import pandas as pd import pytest +import skillmodels.chs.maximization_inputs as mi from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.config import TEST_DATA_DIR -from skillmodels.common.individual_states import get_individual_states_from_params +from skillmodels.common.individual_states import ( + get_individual_states, + get_individual_states_from_params, +) from skillmodels.test_data.model2 import MODEL2, MODEL2_CHS_OPTIONS REGRESSION_VAULT = Path(__file__).parent / "regression_vault" @@ -35,7 +41,10 @@ def test_get_individual_states_from_params(model2, model2_data) -> None: params = params.loc[max_inputs["params_template"].index] calculated = get_individual_states_from_params( - model_spec=model2, data=model2_data, params=params + model_spec=model2, + data=model2_data, + params=params, + chs_options=MODEL2_CHS_OPTIONS, ) factors = ["fac1", "fac2", "fac3"] @@ -45,3 +54,66 @@ def test_get_individual_states_from_params(model2, model2_data) -> None: unanch_ranges = calculated["unanchored_states"]["state_ranges"][factor] ratio = (anch_ranges / unanch_ranges).to_numpy() assert np.allclose(ratio, expected_ratio) + + +@pytest.fixture +def model2_params(model2, model2_data): + params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv") + params = params.set_index(["category", "period", "name1", "name2"]) + template = get_maximization_inputs( + model2, + model2_data, + chs_options=CHSEstimationOptions(start_params_strategy="none"), + )["params_template"] + return params.loc[template.index] + + +def _fail(*args, **kwargs): + raise AssertionError("start values must not be seeded") + + +def test_individual_states_from_params_never_seed_start_values( + model2, model2_data, model2_params, monkeypatch +) -> None: + """Filtering at given params skips AMN / Spearman start-value seeding.""" + monkeypatch.setattr(mi, "_estimate_amn_for_chs_seeding", _fail) + monkeypatch.setattr(mi, "get_spearman_start_params", _fail) + states = get_individual_states_from_params( + model_spec=model2, + data=model2_data, + params=model2_params, + chs_options=CHSEstimationOptions(start_params_strategy="amn"), + )["unanchored_states"]["states"] + assert len(states) == len(model2_data) + + +def test_individual_states_from_params_use_the_sigma_point_scale( + model2, model2_data, model2_params +) -> None: + """The filter runs with the likelihood settings of the given options.""" + + def _states(scale): + return get_individual_states_from_params( + model_spec=model2, + data=model2_data, + params=model2_params, + chs_options=CHSEstimationOptions( + start_params_strategy="none", sigma_points_scale=scale + ), + )["unanchored_states"]["states"]["fac1"].to_numpy() + + assert not np.allclose(_states(2.0), _states(0.5)) + + +def test_get_individual_states_requires_chs_options_for_chs_result( + model2, model2_data, model2_params +) -> None: + result = CHSEstimationResult( + model_spec=model2, + params=model2_params, + success=True, + loglikelihood=0.0, + optimize_result=None, + ) + with pytest.raises(TypeError, match="chs_options"): + get_individual_states(data=model2_data, result=result) diff --git a/tests/test_kalman_filters.py b/tests/test_kalman_filters.py index 73b04040..2774b8e4 100644 --- a/tests/test_kalman_filters.py +++ b/tests/test_kalman_filters.py @@ -570,3 +570,49 @@ def test_sigma_points_multiple_mixtures() -> None: (n_sigma, n_observed), ) aaae(observed_slice, expected) + + +@pytest.mark.parametrize("update_func", UPDATE_FUNCS) +def test_kalman_update_resets_negligible_mixture_component(update_func) -> None: + """A component with underflowing weight takes the dominant component's state. + + Otherwise a nonlinear transition can blow its (irrelevant) state up to inf over + later periods, and the gradient through its zero weight becomes 0 * inf = NaN. + """ + n_states = 2 + states = jnp.array([[[1e100, -1e100], [0.3, -0.2]]]) + chols = jnp.broadcast_to(jnp.eye(n_states), (1, 2, n_states, n_states)) + log_weights = jnp.array([[-1e6, 0.0]]) + kwargs = { + "loadings": jnp.array([1.0, 0.5]), + "control_params": jnp.ones(1), + "meas_sd": 1.0, + "controls": jnp.ones((1, 1)), + } + + def two_updates(s: jax.Array) -> jax.Array: + s1, c1, w1, ll1 = update_func( + states=s, + upper_chols=chols, + measurements=jnp.array([0.4]), + log_mixture_weights=log_weights, + **kwargs, + )[:4] + _, _, _, ll2 = update_func( + states=s1**4, + upper_chols=c1, + measurements=jnp.array([0.1]), + log_mixture_weights=w1, + **kwargs, + )[:4] + return (ll1 + ll2).sum() + + new_states = update_func( + states=states, + upper_chols=chols, + measurements=jnp.array([0.4]), + log_mixture_weights=log_weights, + **kwargs, + )[0] + aaae(new_states[0, 0], new_states[0, 1]) + assert jnp.isfinite(jax.grad(two_updates)(states)).all() diff --git a/tests/test_likelihood_regression.py b/tests/test_likelihood_regression.py index cc540be4..5a2da600 100644 --- a/tests/test_likelihood_regression.py +++ b/tests/test_likelihood_regression.py @@ -103,7 +103,9 @@ def test_splitting_does_not_change_gradient(model2, model2_data) -> None: model2, model2_data, chs_options=MODEL2_CHS_OPTIONS ) inputs_split = get_maximization_inputs( - model2, model2_data, 13, chs_options=MODEL2_CHS_OPTIONS + model2, + model2_data, + chs_options=replace(MODEL2_CHS_OPTIONS, split_dataset=13), ) params = inputs["params_template"] @@ -115,6 +117,18 @@ def test_splitting_does_not_change_gradient(model2, model2_data) -> None: aaae(gradient, gradient_split) +def test_loglike_and_gradient_raises_on_non_finite_value(model2, model2_data) -> None: + """A non-finite value must raise, so optimagic can penalise and reject the step.""" + inputs = get_maximization_inputs( + model2, model2_data, chs_options=MODEL2_CHS_OPTIONS + ) + params = inputs["params_template"] + params["value"] = 0.1 + params.loc[params.index[0], "value"] = np.nan + with pytest.raises(FloatingPointError, match="non-finite"): + inputs["loglike_and_gradient"](params) + + @pytest.mark.parametrize( ("model_name", "fun_key"), product(MODEL_NAMES, ["loglikeobs"]) ) diff --git a/tests/test_missing_observed_factors.py b/tests/test_missing_observed_factors.py new file mode 100644 index 00000000..2751823a --- /dev/null +++ b/tests/test_missing_observed_factors.py @@ -0,0 +1,529 @@ +"""Tests for integrating out missing observed factors in the CHS Kalman filter. + +Where an observed factor `x` is missing for a person-period, it is treated as +`x ~ N(m, v)`, independent of the latent states, with `(m, v)` taken from data +columns named in `ModelSpec.missing_observed_factor_moments`. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pandas as pd +import pytest +from numpy.testing import assert_allclose + +from skillmodels.chs.kalman_filters import ( + calculate_sigma_scaling_factor_and_weights, + kalman_predict, + linear_kalman_predict, +) +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.common.model_spec import FactorSpec, ModelSpec + +jax.config.update("jax_enable_x64", True) + +LATENT = ("fac1", "fac2") +OBSERVED = ("ob1", "ob2") +N_PERIODS = 3 +N_PERSONS = 6 +MOMENTS = {"ob1": ("ob1_mean", "ob1_var"), "ob2": ("ob2_mean", "ob2_var")} +OPTIONS = CHSEstimationOptions(start_params_strategy="none") + +# (person, period) cells where each observed factor is missing. Person 3 misses +# both in period 1 (cross-equation covariance), person 5 misses ob1 only in the +# last period (never used by a transition), person 4 misses nothing. +MISSING_CELLS = { + "ob1": ((0, 0), (1, 0), (1, 1), (3, 1), (5, 2)), + "ob2": ((1, 0), (2, 1), (3, 1)), +} + +# Log-likelihood of `_model(n_mixtures=2, transition=...)` on `_data()` without +# missings at `_params(..., nonlinear=0.05)`: reference values pinning the +# complete-data filter, which supports observed factors without fill moments. +BASELINE_LOGLIKE = { + "linear": -198.6803696097199, + "translog": -206.7069287270014, +} + + +def _model( + n_mixtures: int = 1, + transition: str = "linear", + moments: dict | None = None, +) -> ModelSpec: + factors = { + f: FactorSpec( + measurements=((f"{f}_m1", f"{f}_m2"),) * N_PERIODS, + transition_function=transition, + ) + for f in LATENT + } + return ModelSpec( + factors=factors, + observed_factors=OBSERVED, + n_mixtures=n_mixtures, + missing_observed_factor_moments=moments, + ) + + +def _round(arr: np.ndarray) -> np.ndarray: + """Round to a float32-representable grid so dtype casts are exact.""" + return np.round(arr * 64) / 64 + + +def _data(*, with_missings: bool = True) -> pd.DataFrame: + rng = np.random.default_rng(1234) + idx = pd.MultiIndex.from_product( + [range(N_PERSONS), range(N_PERIODS)], names=["id", "period"] + ) + n_rows = len(idx) + meas = [f"{f}_m{j}" for f in LATENT for j in (1, 2)] + df = pd.DataFrame(_round(rng.normal(size=(n_rows, 4))), index=idx, columns=meas) + df.loc[(2, 1), "fac1_m1"] = np.nan + for ob, (mean_col, var_col) in MOMENTS.items(): + df[ob] = _round(rng.normal(size=n_rows)) + df[mean_col] = _round(rng.normal(scale=0.5, size=n_rows)) + df[var_col] = _round(rng.uniform(0.3, 1.2, size=n_rows)) + if with_missings: + for ob, cells in MISSING_CELLS.items(): + for cell in cells: + df.loc[cell, ob] = np.nan + return df + + +def _params(template: pd.DataFrame, nonlinear: float = 0.0) -> pd.DataFrame: + """Fill a CHS params template with fixed, well-conditioned values.""" + params = template.copy() + lin = { + "fac1": {"fac1": 0.7, "fac2": 0.2, "ob1": 0.9, "ob2": -0.6, "constant": 0.1}, + "fac2": {"fac1": -0.3, "fac2": 0.6, "ob1": 0.5, "ob2": 0.8, "constant": -0.2}, + } + values = [] + for category, period, name1, name2 in params.index: + if category == "controls": + val = 0.1 + elif category == "loadings": + val = 1.0 if name1.endswith("m1") else 0.8 + elif category == "meas_sds": + val = 0.5 + elif category == "shock_sds": + val = 0.4 + 0.1 * period + elif category == "initial_states": + val = 0.0 if name1 == "mixture_0" else 1.0 + val += 0.3 if name2 == "fac2" else 0.0 + elif category == "mixture_weights": + val = 0.4 if name1 == "mixture_0" else 0.6 + elif category == "initial_cholcovs": + a, b = name2.split("-") + val = (1.0 if name1 == "mixture_0" else 0.7) if a == b else 0.3 + elif category == "transition": + val = nonlinear if "*" in name2 else lin[name1][name2] * (1 + 0.1 * period) + else: + raise ValueError(category) + values.append(val) + params["value"] = values + return params + + +def _loglikeobs(model: ModelSpec, data: pd.DataFrame, nonlinear: float = 0.0): + inputs = get_maximization_inputs(model, data, chs_options=OPTIONS) + params = _params(inputs["params_template"], nonlinear=nonlinear) + return inputs["loglikeobs"](params) + + +def _monte_carlo_likelihood( + n_mixtures: int, n_draws: int, seed: int +) -> tuple[np.ndarray, np.ndarray]: + """Average the per-person likelihood over draws of the missing cells. + + Each draw fills every missing observed-factor cell with `N(m, v)` and runs the + ordinary filter (all observed factors present). The draws are stacked as + replicas of the panel so one likelihood evaluation covers all of them. + + Returns: + Per-person Monte-Carlo likelihood and its standard error. + + """ + rng = np.random.default_rng(seed) + base = _data().reset_index() + rep = pd.concat([base] * n_draws, ignore_index=True) + replica = np.repeat(np.arange(n_draws), len(base)) + rep["id"] = np.tile(base["id"].to_numpy(), n_draws) + N_PERSONS * replica + for ob, (mean_col, var_col) in MOMENTS.items(): + values = rep[ob].to_numpy(copy=True) + miss = np.isnan(values) + mean = rep[mean_col].to_numpy()[miss] + sd = np.sqrt(rep[var_col].to_numpy()[miss]) + values[miss] = rng.normal(mean, sd) + rep[ob] = values + rep = rep.set_index(["id", "period"]) + loglikes = _loglikeobs(_model(n_mixtures), rep).reshape(n_draws, N_PERSONS) + likes = np.exp(loglikes) + return likes.mean(axis=0), likes.std(axis=0, ddof=1) / np.sqrt(n_draws) + + +@pytest.fixture(scope="module", params=[1, 2], ids=["1mix", "2mix"]) +def mc_reference(request): + n_mixtures = request.param + mc_like, mc_se = _monte_carlo_likelihood(n_mixtures, n_draws=40_000, seed=7) + return n_mixtures, mc_like, mc_se + + +def test_linear_filter_likelihood_matches_monte_carlo_integral(mc_reference) -> None: + """The exact marginal likelihood equals the MC average of per-draw likelihoods.""" + n_mixtures, mc_like, mc_se = mc_reference + model = _model(n_mixtures, moments=MOMENTS) + exact = np.exp(_loglikeobs(model, _data())) + assert np.all(np.abs(exact - mc_like) <= 4 * mc_se + 1e-12 * mc_like) + + +def test_monte_carlo_reference_rejects_mean_imputation(mc_reference) -> None: + """Plugging in `m` without its variance is far outside the MC error band. + + This shows the MC comparison has power to detect a missing variance term. + """ + n_mixtures, mc_like, mc_se = mc_reference + data = _data() + for ob, (mean_col, _) in MOMENTS.items(): + data[ob] = data[ob].fillna(data[mean_col]) + naive = np.exp(_loglikeobs(_model(n_mixtures), data)) + affected = mc_se > 0 + assert np.all(np.abs(naive - mc_like)[affected] > 10 * mc_se[affected]) + + +@pytest.mark.parametrize("n_mixtures", [1, 2]) +def test_unscented_path_matches_linear_path_for_linear_translog(n_mixtures) -> None: + """A translog with all squares/interactions at 0 reproduces the linear filter.""" + data = _data() + linear = _loglikeobs(_model(n_mixtures, "linear", MOMENTS), data) + unscented = _loglikeobs(_model(n_mixtures, "translog", MOMENTS), data) + assert_allclose(unscented, linear, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("transition", ["linear", "translog"]) +def test_declared_moments_without_missings_leave_loglike_unchanged( + transition, +) -> None: + """Declaring moments for fully observed factors does not change the filter.""" + data = _data(with_missings=False) + plain = _loglikeobs(_model(2, transition), data, nonlinear=0.05) + declared = _loglikeobs(_model(2, transition, MOMENTS), data, nonlinear=0.05) + assert np.array_equal(plain, declared) + + +@pytest.mark.parametrize("transition", ["linear", "translog"]) +def test_loglike_without_missings_matches_baseline(transition) -> None: + """Models with complete observed factors keep their pre-feature likelihood.""" + data = _data(with_missings=False) + loglike = _loglikeobs(_model(2, transition), data, nonlinear=0.05).sum() + assert_allclose(loglike, BASELINE_LOGLIKE[transition], rtol=1e-12) + + +def _endogenous_model(moments: dict | None) -> ModelSpec: + """Endogenous investment `inv` and state `s` share the observed inputs.""" + return ModelSpec( + factors={ + "s": FactorSpec( + measurements=(("y",), ("y",)), transition_function="linear" + ), + "inv": FactorSpec( + measurements=(("i",), ()), + is_endogenous=True, + transition_function="linear", + ), + }, + observed_factors=("x", "z"), + missing_observed_factor_moments=moments, + ) + + +def _endogenous_data(variance: float | None = None) -> pd.DataFrame: + idx = pd.MultiIndex.from_product([[0, 1], [0, 1]], names=["id", "period"]) + data = pd.DataFrame( + {"y": [-1.0, 0.0, 1.0, 2.0], "i": np.nan, "x": 0.0, "z": 0.0}, index=idx + ) + data[["m", "v"]] = 0.0 + if variance is not None: + data.loc[(0, 0), ["x", "v"]] = [np.nan, variance] + return data + + +def _endogenous_loglike(model: ModelSpec, data: pd.DataFrame) -> float: + inputs = get_maximization_inputs(model, data, chs_options=OPTIONS) + params = inputs["params_template"].copy() + params["value"] = np.where( + params.index.get_level_values("category").isin( + ["meas_sds", "shock_sds", "loadings"] + ), + 0.5, + 0.1, + ) + for idx in params.query("category == 'initial_cholcovs'").index: + row, col = idx[3].split("-") + params.loc[idx, "value"] = float(row == col) + return float(inputs["loglike"](params)) + + +@pytest.mark.parametrize( + "moments", + [ + {"x": ("m", "v"), "z": ("m", "v")}, + {"x": ("z", "v"), "z": ("m", "v")}, + ], + ids=["shared_pair", "moment_is_observed_factor"], +) +def test_moment_columns_shared_or_aliased_leave_loglike_unchanged(moments) -> None: + """Moment columns may be shared across factors or alias a selected column.""" + data = _endogenous_data() + plain = _endogenous_loglike(_endogenous_model(None), data) + declared = _endogenous_loglike(_endogenous_model(moments), data) + assert_allclose(declared, plain, rtol=0, atol=1e-12) + + +def test_endogenous_model_rejects_uncertain_missing_input() -> None: + """Both half-steps would need the input's joint posterior, which is not carried.""" + model = _endogenous_model({"x": ("m", "v")}) + with pytest.raises(NotImplementedError, match="joint missing-observed-factor"): + get_maximization_inputs(model, _endogenous_data(1.0), chs_options=OPTIONS) + + +def test_endogenous_model_plugs_in_zero_variance_fill() -> None: + """A zero-variance fill acts like observing the fill mean.""" + model = _endogenous_model({"x": ("m", "v")}) + filled = _endogenous_loglike(model, _endogenous_data(0.0)) + plugged_in = _endogenous_loglike(model, _endogenous_data()) + assert_allclose(filled, plugged_in, rtol=0, atol=1e-12) + + +def test_moments_where_factor_is_observed_are_ignored() -> None: + """Fill moments only matter in cells where the observed factor is missing.""" + data = _data() + garbled = data.copy() + for ob, (mean_col, var_col) in MOMENTS.items(): + observed = garbled[ob].notna() + garbled.loc[observed, mean_col] = np.nan + garbled.loc[observed, var_col] = -1.0 + model = _model(1, moments=MOMENTS) + assert np.array_equal(_loglikeobs(model, data), _loglikeobs(model, garbled)) + + +@pytest.mark.parametrize("bad_col", ["ob1_mean", "ob1_var"]) +def test_missing_factor_without_moments_raises(bad_col) -> None: + data = _data() + data.loc[(1, 1), bad_col] = np.nan + with pytest.raises(ValueError, match=r"ob1[\s\S]*fill moments") as excinfo: + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + assert "1" in str(excinfo.value) + + +def test_negative_fill_variance_raises() -> None: + data = _data() + data.loc[(1, 1), "ob1_var"] = -0.5 + with pytest.raises(ValueError, match=r"ob1[\s\S]*negative fill variance"): + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + + +def test_missing_moment_column_raises() -> None: + data = _data().drop(columns=["ob2_var"]) + with pytest.raises(ValueError, match=r"ob2_var[\s\S]*Variable is missing"): + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + + +def test_missing_factor_without_declared_moments_raises() -> None: + moments = {"ob1": MOMENTS["ob1"]} + with pytest.raises(ValueError, match=r"ob2[\s\S]*Variable has missings"): + get_maximization_inputs( + _model(1, moments=moments), _data(), chs_options=OPTIONS + ) + + +def test_model_spec_rejects_moments_for_unknown_observed_factor() -> None: + with pytest.raises(ValueError, match="not_an_observed_factor"): + _model(1, moments={"not_an_observed_factor": ("a", "b")}) + + +def test_model_spec_from_dict_reads_moments() -> None: + spec = ModelSpec.from_dict( + { + "factors": {"fac1": {"measurements": [["y1"], ["y1"]]}}, + "observed_factors": ["ob1"], + "missing_observed_factor_moments": {"ob1": ["ob1_mean", "ob1_var"]}, + } + ) + assert dict(spec.missing_observed_factor_moments) == { + "ob1": ("ob1_mean", "ob1_var") + } + + +def _predict_inputs(): + """Inputs for one predict step: 3 persons, 2 mixtures, 2 latent, 2 observed. + + Person 0 misses ob1, person 1 misses both, person 2 misses none. + """ + rng = np.random.default_rng(0) + n_obs, n_mix, n_lat = 3, 2, 2 + states = rng.normal(size=(n_obs, n_mix, n_lat)) + covs = np.empty((n_obs, n_mix, n_lat, n_lat)) + chols = np.empty_like(covs) + for i in range(n_obs): + for k in range(n_mix): + a = rng.normal(size=(n_lat, n_lat)) + covs[i, k] = a @ a.T + np.eye(n_lat) + chols[i, k] = np.linalg.cholesky(covs[i, k]).T + # Means are already substituted where the factor is missing. + observed = rng.normal(size=(n_obs, 2)) + # float32-representable variances with irrational roots. + obs_vars = np.array([[0.75, 0.0], [0.375, 1.25], [0.0, 0.0]]) + f_obs = np.array([[0.9, -0.6], [0.5, 0.8]]) + f_lat = np.array([[0.7, 0.2], [-0.3, 0.6]]) + const = np.array([0.1, -0.2]) + shock_sds = np.array([0.4, 0.3]) + scaling = np.array([[1.5, 0.8, 1.0, 1.0], [1.2, 0.9, 1.0, 1.0]]) + constants = np.array([[0.3, -0.1, 0.0, 0.0], [0.2, 0.4, 0.0, 0.0]]) + return { + "states": states, + "covs": covs, + "chols": chols, + "observed": observed, + "obs_vars": obs_vars, + "f_lat": f_lat, + "f_obs": f_obs, + "const": const, + "shock_sds": shock_sds, + "scaling": scaling, + "constants": constants, + } + + +def _analytic_prediction(inp): + """Predicted mean and covariance with `x ~ N(m, diag(v))` integrated out.""" + s_in, s_out = inp["scaling"][0, :2], inp["scaling"][1, :2] + c_in, c_out = inp["constants"][0, :2], inp["constants"][1, :2] + n_obs, n_mix, _ = inp["states"].shape + means = np.empty_like(inp["states"]) + covs = np.empty_like(inp["covs"]) + for i in range(n_obs): + for k in range(n_mix): + anchored = inp["states"][i, k] * s_in + c_in + pred = inp["f_lat"] @ anchored + inp["f_obs"] @ inp["observed"][i] + means[i, k] = (pred + inp["const"] - c_out) / s_out + g_mat = np.diag(1 / s_out) @ inp["f_lat"] @ np.diag(s_in) + b_mat = np.diag(1 / s_out) @ inp["f_obs"] + covs[i, k] = ( + g_mat @ inp["covs"][i, k] @ g_mat.T + + np.diag((inp["shock_sds"] / s_out) ** 2) + + b_mat @ np.diag(inp["obs_vars"][i]) @ b_mat.T + ) + return means, covs + + +def _linear_trans_coeffs(inp): + return { + f: jnp.array([*inp["f_lat"][j], *inp["f_obs"][j], inp["const"][j]]) + for j, f in enumerate(LATENT) + } + + +@pytest.mark.parametrize("variance_dtype", [jnp.float64, jnp.float32]) +def test_linear_predict_integrates_out_missing_observed_factors( + variance_dtype, +) -> None: + inp = _predict_inputs() + exp_means, exp_covs = _analytic_prediction(inp) + means, chols = linear_kalman_predict( + None, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + 1.0, + jnp.ones(1), + _linear_trans_coeffs(inp), + jnp.array(inp["shock_sds"]), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + observed_factor_vars=jnp.array(inp["obs_vars"], dtype=variance_dtype), + latent_factors=LATENT, + constant_factor_indices=frozenset(), + n_all_factors=4, + missing_observed_indices=(0, 1), + ) + covs = np.einsum("nkja,nkjb->nkab", chols, chols) + assert_allclose(means, exp_means, rtol=1e-12) + assert_allclose(covs, exp_covs, rtol=1e-12) + + +def _linear_transition(params, states): + return jnp.column_stack([states @ params[f][:-1] + params[f][-1] for f in LATENT]) + + +@pytest.mark.parametrize("variance_dtype", [jnp.float64, jnp.float32]) +def test_unscented_predict_integrates_out_missing_observed_factors( + variance_dtype, +) -> None: + inp = _predict_inputs() + exp_means, exp_covs = _analytic_prediction(inp) + scaling_factor, weights = calculate_sigma_scaling_factor_and_weights(4, 2) + # The unscented predict adds shocks in the unanchored output scale. + shock_sds = inp["shock_sds"] / inp["scaling"][1, :2] + means, chols = kalman_predict( + _linear_transition, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + scaling_factor, + weights, + _linear_trans_coeffs(inp), + jnp.array(shock_sds), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + observed_factor_vars=jnp.array(inp["obs_vars"], dtype=variance_dtype), + missing_observed_indices=(0, 1), + ) + covs = np.einsum("nkja,nkjb->nkab", chols, chols) + assert_allclose(means, exp_means, rtol=1e-10) + assert_allclose(covs, exp_covs, rtol=1e-10) + + +def _nonlinear_transition(params, states): + lat = states[:, :2] + out = _linear_transition(params, states) + return out + 0.2 * lat**2 - 0.1 * lat[:, ::-1] * states[:, 2:3] + + +def test_unscented_predict_augmented_by_observed_cells_equals_larger_kappa() -> None: + """With zero fill variance, the augmented sigma points act like `kappa + n_aug`. + + The latent marginal of the augmented sigma-point set is the plain set built + with `kappa` raised by the number of augmented dimensions. + """ + inp = _predict_inputs() + common = ( + _nonlinear_transition, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + ) + rest = ( + _linear_trans_coeffs(inp), + jnp.array(inp["shock_sds"]), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + ) + aug_scaling, aug_weights = calculate_sigma_scaling_factor_and_weights(4, 2) + aug_means, aug_chols = kalman_predict( + *common, + aug_scaling, + aug_weights, + *rest, + observed_factor_vars=jnp.zeros((3, 2)), + missing_observed_indices=(0, 1), + ) + plain_scaling, plain_weights = calculate_sigma_scaling_factor_and_weights(2, 4) + plain_means, plain_chols = kalman_predict( + *common, plain_scaling, plain_weights, *rest + ) + aug_covs = np.einsum("nkja,nkjb->nkab", aug_chols, aug_chols) + plain_covs = np.einsum("nkja,nkjb->nkab", plain_chols, plain_chols) + assert_allclose(aug_means, plain_means, rtol=1e-12) + assert_allclose(aug_covs, plain_covs, rtol=1e-12) diff --git a/tests/test_model_spec.py b/tests/test_model_spec.py index 97d34aa1..f0a258da 100644 --- a/tests/test_model_spec.py +++ b/tests/test_model_spec.py @@ -9,6 +9,7 @@ ModelSpec, Normalizations, ) +from skillmodels.exceptions import ModelSpecInitializationError def _minimal_dict(): @@ -81,6 +82,28 @@ def test_from_dict_with_stagemap() -> None: assert spec.stagemap == (0,) +def test_from_dict_rejects_unknown_keys() -> None: + """An unrecognized top-level key raises instead of being silently dropped.""" + d = _minimal_dict() + d["estimation_options"] = {"n_mixtures": 2, "bounds_distance": 0.001} + with pytest.raises(ModelSpecInitializationError, match="estimation_options"): + ModelSpec.from_dict(d) + + +def test_from_dict_accepts_every_key_it_reads() -> None: + d = _minimal_dict() + d |= { + "observed_factors": ["x"], + "controls": ["c"], + "stagemap": [0], + "n_mixtures": 1, + "measurement_models": {}, + "missing_observed_factor_moments": {"x": ("x_mean", "x_var")}, + } + spec = ModelSpec.from_dict(d) + assert spec.observed_factors == ("x",) + + def test_correction_spec_defaults() -> None: cf = CorrectionSpec(instruments=("z1",)) assert cf.instruments == ("z1",) diff --git a/tests/test_visualize_factor_distributions.py b/tests/test_visualize_factor_distributions.py index e021dff2..54356067 100644 --- a/tests/test_visualize_factor_distributions.py +++ b/tests/test_visualize_factor_distributions.py @@ -2,7 +2,9 @@ from pathlib import Path +import numpy as np import pandas as pd +from scipy.stats import gaussian_kde from skillmodels.chs.maximization_inputs import get_maximization_inputs from skillmodels.common.config import TEST_DATA_DIR @@ -31,7 +33,7 @@ def _load_model2_filtered() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: params = params.loc[max_inputs["params_template"].index] states = get_individual_states_from_params( - model_spec=MODEL2, data=data, params=params + model_spec=MODEL2, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS )["anchored_states"]["states"] return data, params, states @@ -115,3 +117,22 @@ def test_visualize_factor_distributions_with_period_indexed_states() -> None: contour_plots=contours, surface_plots=None, ) + + +def test_bivariate_density_contours_put_first_factor_on_x_axis() -> None: + """Contour value at (x[col], y[row]) is the KDE of (first, second) factor there.""" + data, _params, states = _load_model2_filtered() + contours = bivariate_density_contours( + data=data, + model_spec=MODEL2, + period=1, + filtered_states=states, + n_points=7, + ) + (fac_x, fac_y), fig = next(iter(contours.items())) + trace = fig.data[0] + period_states = states.query("period == 1") + kernel = gaussian_kde(np.vstack([period_states[fac_x], period_states[fac_y]])) + xx, yy = np.meshgrid(trace.x, trace.y) + expected = kernel(np.vstack([xx.ravel(), yy.ravel()])).reshape(xx.shape) + np.testing.assert_allclose(np.asarray(trace.z), expected, rtol=1e-10) diff --git a/tests/test_visualize_transition_equations.py b/tests/test_visualize_transition_equations.py index b0dfe37f..13a4969b 100644 --- a/tests/test_visualize_transition_equations.py +++ b/tests/test_visualize_transition_equations.py @@ -11,6 +11,8 @@ from skillmodels.common.config import TEST_DATA_DIR from skillmodels.common.individual_states import get_individual_states_from_params from skillmodels.common.visualize_transition_equations import ( + _get_state_ranges, + _prepare_data_for_one_plot_average_2d, combine_transition_plots, get_transition_plots, ) @@ -35,7 +37,7 @@ def test_visualize_transition_equations_runs() -> None: params["value"] = params["value"].fillna(0) states = get_individual_states_from_params( - model_spec=model, data=data, params=params + model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS )["anchored_states"]["states"] subplots = get_transition_plots( @@ -143,7 +145,10 @@ def test_get_transition_plots_runs_for_correction_model() -> None: assert not template["value"].isna().any() states = get_individual_states_from_params( - model_spec=model, data=panel, params=template + model_spec=model, + data=panel, + params=template, + chs_options=CHSEstimationOptions(), )["anchored_states"]["states"] plots = get_transition_plots( @@ -156,3 +161,41 @@ def test_get_transition_plots_runs_for_correction_model() -> None: ) # A target factor's plot exists -> its cf-corrected transition evaluated. assert ("inv", "fac1") in plots + + +def test_get_state_ranges_adds_factors_missing_from_given_ranges() -> None: + # Callers pass ranges of the latent factors; observed factors are inputs too. + states = pd.DataFrame( + {"aug_period": [0, 0, 1, 1], "fac1": [0.0, 1.0, 2.0, 3.0], "obs": [5, 7, 6, 9]} + ) + given = {"fac1": pd.DataFrame({"minimum": [-1.0], "maximum": [1.0]})} + ranges = _get_state_ranges( + state_ranges=given, states_data=states, all_factors=("fac1", "obs") + ) + assert ranges["fac1"] is given["fac1"] + assert ranges["obs"]["maximum"].tolist() == [7, 9] + + +def test_average_over_draws_skips_persons_with_a_missing_factor() -> None: + # Observed factors can be missing; a draw with a NaN would make every mean NaN. + states = pd.DataFrame( + { + "aug_period": [0] * 4, + "fac1": [0.0, 1.0, 2.0, 3.0], + "obs": [1.0, np.nan, 1.0, 1.0], + } + ) + ranges = {"fac1": pd.DataFrame({"minimum": [0.0], "maximum": [1.0]})} + out = _prepare_data_for_one_plot_average_2d( + states_data=states, + state_ranges=ranges, + aug_period=0, + input_factor="fac1", + output_factor="fac1", + n_points=3, + n_draws=10, + transition_function=lambda _params, x: x.sum(axis=1), + transition_params={}, + all_factors=("fac1", "obs"), + ) + assert out["output_fac1"].tolist() == [1.0, 1.5, 2.0]