Remove aug_periods from public-facing functions.

Author: hmgaudecker

CLAUDE.md

@@ -174,8 +174,7 @@ The codebase uses:
 **Simulation:**
 
 - `simulate_dataset(model_spec, params, n_obs=None, data=None, policies=None, seed=None)`
-  — returns dict with `"unanchored_states"`, `"anchored_states"`,
-  `"aug_unanchored_states"`, `"aug_measurements"`
+  — returns dict with `"unanchored_states"`, `"anchored_states"`
 - `simulate_policy_effect(model_spec, params, data, policies, seed=None)` — returns
   DataFrame of factor mean differences between policy and baseline
 
@@ -184,7 +183,7 @@ The codebase uses:
 - `plot_likelihood_contributions(model_spec, data, params, period=None)`
 - `plot_residual_boxplots(model_spec, data, params, period=None)`
 - `decompose_measurement_variance(model_spec, params, data)` — returns DataFrame indexed
-  by `(aug_period, measurement, factor)` with signal/noise columns
+  by `(period, measurement, factor)` with signal/noise columns
 - `summarize_measurement_reliability(variance_decomposition)`
 - `create_state_ranges(filtered_states, factors, quantile_cutoff=None)`
 
@@ -239,18 +238,18 @@ These are not in `__all__` but are imported directly by application projects:
 ### Period vs Aug_period
 
 Models with endogenous factors split each calendar period into multiple **augmented
-periods** (`aug_period`). The public API should use `period` (user-facing); `aug_period`
-is an internal concept. Current status:
+periods** (`aug_period`). The public API uses `period` (user-facing); `aug_period` is
+strictly internal. All public functions now return `period`:
 
 - `ModelSpec` — clean, no `aug_period` exposure.
 - `get_transition_plots()` — clean, accepts `period`/`periods`.
-- `get_filtered_states()` — **leaks `aug_period`**: returned states DataFrames have an
-  `aug_period` column, not `period`.
-- `simulate_dataset()` — mixed: `"anchored_states"` / `"unanchored_states"` use
-  `period`, but `"aug_unanchored_states"` / `"aug_measurements"` use `aug_period`.
-- `plot_residual_boxplots()` / `plot_likelihood_contributions()` — accept `period` but
-  return figures keyed by `aug_period`.
-- `decompose_measurement_variance()` — returns DataFrame indexed by `aug_period`.
+- `get_filtered_states()` — clean, returns `period` column.
+- `simulate_dataset()` — clean, returns `period` in states DataFrames.
+- `plot_residual_boxplots()` / `plot_likelihood_contributions()` — clean, accept and
+  return `period`.
+- `decompose_measurement_variance()` — clean, indexed by
+  `(period, measurement, factor)`.
+- `simulate_policy_effect()` / `simulate_dataset()` policies — accept `"period"` key.
 - `ProcessedModel.labels` — exposes `aug_periods_to_periods` mapping (acceptable for
   internal/advanced use).
 

src/skillmodels/diagnostic_plots.py

@@ -61,19 +61,25 @@ def plot_residual_boxplots(
         how="left",
     )
 
+    # Map aug_period → period for the public API
+    ap_to_p = processed_model.labels.aug_periods_to_periods
+
     available_periods = sorted(residuals_df[period_col].unique())
 
     if period is not None:
+        # Find aug_period(s) matching the requested period
+        aug_periods_for_period = [ap for ap, p in ap_to_p.items() if p == period]
+        aug_period = aug_periods_for_period[0] if aug_periods_for_period else period
         return _create_residual_boxplot_for_period(
             residuals_df=residuals_df,
-            period=period,
+            period=aug_period,
             period_col=period_col,
             show_reference_line=show_reference_line,
             layout_kwargs=layout_kwargs,
         )
 
     return {
-        p: _create_residual_boxplot_for_period(
+        ap_to_p.get(p, p): _create_residual_boxplot_for_period(
             residuals_df=residuals_df,
             period=p,
             period_col=period_col,
@@ -173,18 +179,24 @@ def plot_likelihood_contributions(
         how="left",
     )
 
+    # Map aug_period → period for the public API
+    ap_to_p = processed_model.labels.aug_periods_to_periods
+
     available_periods = sorted(contributions_df[period_col].unique())
 
     if period is not None:
+        # Find aug_period(s) matching the requested period
+        aug_periods_for_period = [ap for ap, p in ap_to_p.items() if p == period]
+        aug_period = aug_periods_for_period[0] if aug_periods_for_period else period
         return _create_likelihood_boxplot_for_period(
             contributions_df=contributions_df,
-            period=period,
+            period=aug_period,
             period_col=period_col,
             layout_kwargs=layout_kwargs,
         )
 
     return {
-        p: _create_likelihood_boxplot_for_period(
+        ap_to_p.get(p, p): _create_likelihood_boxplot_for_period(
             contributions_df=contributions_df,
             period=p,
             period_col=period_col,

src/skillmodels/filtered_states.py

@@ -35,10 +35,20 @@ def get_filtered_states(
         use_aug_period=True,
     )
 
+    # Map aug_period → period for the public API
+    ap_to_p = processed_model.labels.aug_periods_to_periods
+    for df in (anchored_states_df, unanchored_states_df):
+        df["period"] = df["aug_period"].map(ap_to_p)
+        df.drop(columns="aug_period", inplace=True)  # noqa: PD002
+
     anchored_ranges = create_state_ranges(
         filtered_states=anchored_states_df,
         factors=processed_model.labels.latent_factors,
     )
+    unanchored_ranges = create_state_ranges(
+        filtered_states=unanchored_states_df,
+        factors=processed_model.labels.latent_factors,
+    )
 
     return {
         "anchored_states": {

src/skillmodels/process_debug_data.py

@@ -182,8 +182,8 @@ def create_state_ranges(
     """Compute minimum and maximum state values for each factor by period.
 
     Args:
-        filtered_states: DataFrame with filtered states. Must have a "period" or
-            "aug_period" column.
+        filtered_states: DataFrame with filtered states. Must have a "period"
+            column.
         factors: List of factor names to compute ranges for.
         quantile_cutoff: If provided, use quantiles instead of min/max. The cutoff
             is applied symmetrically: the minimum is the `quantile_cutoff` quantile

src/skillmodels/simulate_data.py

@@ -44,8 +44,8 @@ def simulate_dataset(
         data: Dataset in the same format as for estimation, containing
             information about observed factors and control variables.
         policies: Each dictionary specifies a stochastic shock to a latent factor
-            AT THE END of "period" for "factor" with mean "effect_size" and
-            "standard deviation".
+            AT THE END of ``"period"`` for ``"factor"`` with mean
+            ``"effect_size"`` and ``"standard_deviation"``.
         seed: Random seed for reproducibility. If None, uses numpy's default random
             state.
 
@@ -125,7 +125,14 @@ def simulate_dataset(
         n_obs=n_obs,
     )
 
-    aug_measurements, aug_latent_data = _simulate_dataset(
+    # Convert "period" keys in policies to "aug_period" for internal use
+    if policies is not None:
+        policies = _convert_policy_periods(
+            policies=policies,
+            endogenous_factors_info=processed_model.endogenous_factors_info,
+        )
+
+    _aug_measurements, aug_latent_data = _simulate_dataset(
         latent_states=states,
         covs=covs,
         log_weights=log_weights,
@@ -173,14 +180,6 @@ def simulate_dataset(
                 factors=processed_model.labels.latent_factors,
             ),
         },
-        "aug_unanchored_states": {
-            "states": aug_latent_data,
-            "state_ranges": create_state_ranges(
-                filtered_states=aug_latent_data,
-                factors=processed_model.labels.latent_factors,
-            ),
-        },
-        "aug_measurements": aug_measurements,
     }
 
 
@@ -396,6 +395,31 @@ def _collapse_aug_periods_to_periods(
     )
 
 
+def _convert_policy_periods(
+    policies: list[dict],
+    endogenous_factors_info: EndogenousFactorsInfo,
+) -> list[dict]:
+    """Convert ``"period"`` keys in policy dicts to ``"aug_period"``.
+
+    Policies may specify either ``"period"`` (public API) or ``"aug_period"``
+    (legacy/internal). This normalises to ``"aug_period"`` for the simulation loop.
+    """
+    converted = []
+    for policy in policies:
+        if "aug_period" in policy:
+            converted.append(policy)
+        elif "period" in policy:
+            p = dict(policy)
+            period = p.pop("period")
+            aug_periods = endogenous_factors_info.aug_periods_from_period(period)
+            # Use the first aug_period for the given period
+            p["aug_period"] = aug_periods[0]
+            converted.append(p)
+        else:
+            raise ValueError("Each policy dict must contain a 'period' key.")
+    return converted
+
+
 def _get_shock(
     rng: np.random.Generator,
     mean: float,
@@ -514,7 +538,7 @@ def simulate_policy_effect(
         data: Dataset with observed factors and control variables.
         policies: List of policy dictionaries. Each dictionary specifies a
             stochastic shock to a latent factor with keys:
-            - "period" or "aug_period": When to apply the shock
+            - "period": When to apply the shock
             - "factor": Which factor to shock
             - "effect_size": Mean of the shock
             - "standard_deviation": Standard deviation of the shock (use 0 for
@@ -564,8 +588,6 @@ def simulate_policy_effect(
     policy_means = policy_states.groupby("period").mean()
 
     # Drop non-factor columns
-    factor_cols = [
-        c for c in baseline_means.columns if c not in ("id", "aug_period", "period")
-    ]
+    factor_cols = [c for c in baseline_means.columns if c not in ("id", "period")]
 
     return policy_means[factor_cols] - baseline_means[factor_cols]

src/skillmodels/variance_decomposition.py

@@ -5,10 +5,13 @@
 Section 4.2.2.
 """
 
+from collections.abc import Mapping
+
 import pandas as pd
 
 from skillmodels.filtered_states import get_filtered_states
 from skillmodels.model_spec import ModelSpec
+from skillmodels.process_model import process_model
 
 
 def decompose_measurement_variance(
@@ -34,7 +37,7 @@ def decompose_measurement_variance(
         data: Empirical dataset used to estimate the model.
 
     Returns:
-        DataFrame indexed by (aug_period, measurement, factor) with columns:
+        DataFrame indexed by (period, measurement, factor) with columns:
         - loading: The factor loading (L)
         - factor_variance: Var(F) for that period
         - meas_sd: The measurement error standard deviation
@@ -54,32 +57,50 @@ def decompose_measurement_variance(
     )
     filtered_states = filtered_result["anchored_states"]["states"]
 
+    processed_model = process_model(model_spec)
     return _compute_variance_decomposition(
         filtered_states=filtered_states,
         params=params,
+        aug_periods_to_periods=processed_model.labels.aug_periods_to_periods,
     )
 
 
 def _compute_variance_decomposition(
     filtered_states: pd.DataFrame,
     params: pd.DataFrame,
+    aug_periods_to_periods: Mapping[int, int],
 ) -> pd.DataFrame:
     """Compute variance decomposition from filtered states and parameters.
 
     Args:
         filtered_states: DataFrame with filtered states, must have columns for
-            each factor plus "aug_period" and "id".
+            each factor plus "period" and "id".
         params: DataFrame with model parameters indexed by
             (category, aug_period, name1, name2).
+        aug_periods_to_periods: Mapping from aug_period to period.
 
     Returns:
         DataFrame with variance decomposition results.
 
     """
+    # Build reverse mapping: period → aug_period (pick first aug_period per period)
+    periods_to_aug_periods = {}
+    for ap, p in aug_periods_to_periods.items():
+        if p not in periods_to_aug_periods:
+            periods_to_aug_periods[p] = ap
+
+    # Add aug_period column for internal merges with params
+    filtered_states = filtered_states.copy()
+    filtered_states["aug_period"] = filtered_states["period"].map(
+        periods_to_aug_periods
+    )
+
     # Compute factor variances by period
     periods = filtered_states["aug_period"].unique()
     factor_cols = [
-        c for c in filtered_states.columns if c not in ("aug_period", "id", "mixture")
+        c
+        for c in filtered_states.columns
+        if c not in ("aug_period", "period", "id", "mixture")
     ]
 
     factor_variances = {}
@@ -133,8 +154,11 @@ def _compute_variance_decomposition(
     merged["fraction_noise"] = noise_var / total_var
     merged["signal_to_noise_ratio"] = signal_var / noise_var
 
+    # Map aug_period → period for the public API
+    merged["period"] = merged["aug_period"].map(aug_periods_to_periods)
+
     # Set index and select columns
-    return merged.set_index(["aug_period", "measurement", "factor"])[
+    return merged.set_index(["period", "measurement", "factor"])[
         [
             "loading",
             "factor_variance",

src/skillmodels/visualize_transition_equations.py

@@ -196,7 +196,7 @@ def get_transition_plots(  # noqa: C901, PLR0912
         layout_kwargs: Dictionary of key word arguments used to
             update layout of plotly image object. If None, the default kwargs
             defined in the function will be used.
-        states: Pre-computed filtered states DataFrame (with an `aug_period`
+        states: Pre-computed filtered states DataFrame (with a `period`
             column). If provided, skip the internal `get_filtered_states` call.
         include_correction_factors: Whether to include correction factors in the
             plots. Default False.

tests/test_variance_decomposition.py

@@ -22,7 +22,7 @@ def setup_variance_decomposition():
             "fac1": [0.1, 0.1, 0.1, 0.2],
             "fac2": [0.1, 0.1, 0.1, 0.1],
             "fac3": [0.2, 0.2, 0.2, 0.4],
-            "aug_period": [0, 0, 0, 0],
+            "period": [0, 0, 0, 0],
             "id": [0, 1, 2, 3],
         }
     )
@@ -60,7 +60,11 @@ def setup_variance_decomposition():
 
     params = pd.concat([loadings_df, meas_sds_df], keys=["loadings", "meas_sds"])
 
-    return {"filtered_states": filtered_states, "params": params}
+    return {
+        "filtered_states": filtered_states,
+        "params": params,
+        "aug_periods_to_periods": {0: 0},
+    }
 
 
 @pytest.fixture
@@ -74,7 +78,7 @@ def expected_variance_decomposition():
     """
     index = pd.MultiIndex.from_tuples(
         [(0, "y1", "fac1"), (0, "y2", "fac2"), (0, "y3", "fac3")],
-        names=["aug_period", "measurement", "factor"],
+        names=["period", "measurement", "factor"],
     )
     return pd.DataFrame(
         {