Add rake and poststratify diagnostics (#530)

Summary: Pull Request resolved: #530

Differential Revision: D112668817

Pulled By: talgalili

fbshipit-source-id: 03fd8da0a8bcd1a230de831be1bb439a569c2e9f

Author: neuralsorcerer

CHANGELOG.md

@@ -1,3 +1,17 @@
+# 0.23.0 (unreleased)
+
+## New Features
+
+- Diagnostics now include compact model-glance rows for rake and poststratification adjustments, covering rake convergence metadata and persisted poststratification matching-cell metadata when available.
+
+## Documentation
+
+- Added docstring and statistical-method and notebook tutorial examples showing the new rake and poststratification `model_glance` diagnostics output.
+
+## Tests
+
+- Added summary-helper regression coverage for rake and poststratification model diagnostics, including examples that mirror the new docstring outputs.
+
 # 0.22.0 (2026-07-15)
 
 ## New Features

balance/summary_utils.py

@@ -18,6 +18,7 @@ class hierarchy.
 from __future__ import annotations
 
 import logging
+from collections.abc import Sized
 from typing import Any
 
 import numpy as np
@@ -161,6 +162,151 @@ def _concat_metric_val_var(
     return pd.concat((diagnostics, rows), ignore_index=True)
 
 
+def _safe_len(value: Any) -> int | float:
+    """Return ``len(value)`` when available, otherwise ``np.nan``.
+
+    Args:
+        value: Candidate object to measure.
+
+    Returns:
+        The object's length, or ``np.nan`` when the object is scalar or its
+        length implementation raises.
+    """
+    if value is None or isinstance(value, (str, bytes)):
+        return np.nan
+    if not isinstance(value, Sized):
+        return np.nan
+
+    try:
+        return len(value)
+    except Exception:
+        return np.nan
+
+
+def _append_rake_model_diagnostics(
+    diagnostics: pd.DataFrame, model: dict[str, Any]
+) -> pd.DataFrame:
+    """Append compact diagnostics for a fitted rake model.
+
+    Args:
+        diagnostics: Existing diagnostics table to append to.
+        model: Rake model metadata dictionary, usually the ``model`` object
+            returned by ``balance.weighting_methods.rake.rake``.
+
+    Returns:
+        A diagnostics table with any available rake ``model_glance`` rows
+        appended. Missing optional metadata is skipped except for the rake
+        convergence flag, which is emitted as ``NaN`` when absent.
+
+    Examples:
+        A rake model with two iterations emits compact convergence rows::
+
+            >>> import pandas as pd
+            >>> diagnostics = pd.DataFrame(columns=["metric", "val", "var"])
+            >>> model = {
+            ...     "method": "rake",
+            ...     "converged": 1,
+            ...     "iterations": pd.DataFrame({"conv": [0.5, 0.01]}),
+            ...     "variables": ["gender", "age_group"],
+            ... }
+            >>> _append_rake_model_diagnostics(diagnostics, model).to_dict("records")
+            [{'metric': 'model_glance', 'val': 1, 'var': 'converged'}, {'metric': 'model_glance', 'val': 2, 'var': 'iterations'}, {'metric': 'model_glance', 'val': 0.01, 'var': 'final_conv'}, {'metric': 'model_glance', 'val': 2, 'var': 'n_variables'}]
+    """
+    diagnostics = _concat_metric_val_var(
+        diagnostics,
+        "model_glance",
+        [_coerce_scalar(model.get("converged"))],
+        ["converged"],
+    )
+
+    iterations = model.get("iterations")
+    if isinstance(iterations, pd.DataFrame):
+        diagnostics = _concat_metric_val_var(
+            diagnostics,
+            "model_glance",
+            [len(iterations)],
+            ["iterations"],
+        )
+        if "conv" in iterations.columns and len(iterations) > 0:
+            diagnostics = _concat_metric_val_var(
+                diagnostics,
+                "model_glance",
+                [_coerce_scalar(iterations["conv"].iloc[-1])],
+                ["final_conv"],
+            )
+
+    variables = model.get("variables")
+    if variables is not None:
+        diagnostics = _concat_metric_val_var(
+            diagnostics,
+            "model_glance",
+            [_safe_len(variables)],
+            ["n_variables"],
+        )
+
+    return diagnostics
+
+
+def _append_poststratify_model_diagnostics(
+    diagnostics: pd.DataFrame, model: dict[str, Any]
+) -> pd.DataFrame:
+    """Append compact diagnostics for a fitted poststratification model.
+
+    Args:
+        diagnostics: Existing diagnostics table to append to.
+        model: Poststratification model metadata dictionary, usually the
+            ``model`` object returned by
+            ``balance.weighting_methods.poststratify.poststratify``.
+
+    Returns:
+        A diagnostics table with any available poststratification
+        ``model_glance`` rows appended. Rows that require persisted fit
+        metadata are omitted when that metadata is unavailable.
+
+    Examples:
+        Persisted poststratification metadata emits variable, matching, and
+        cell-count rows::
+
+            >>> import pandas as pd
+            >>> diagnostics = pd.DataFrame(columns=["metric", "val", "var"])
+            >>> model = {
+            ...     "method": "poststratify",
+            ...     "variables": ["gender", "age_group"],
+            ...     "strict_matching": True,
+            ...     "cell_weight_ratio": pd.Series([0.5, 2.0]),
+            ... }
+            >>> _append_poststratify_model_diagnostics(diagnostics, model).to_dict("records")
+            [{'metric': 'model_glance', 'val': 2, 'var': 'n_variables'}, {'metric': 'model_glance', 'val': 1, 'var': 'strict_matching'}, {'metric': 'model_glance', 'val': 2, 'var': 'n_cells'}]
+    """
+    variables = model.get("variables")
+    if variables is not None:
+        diagnostics = _concat_metric_val_var(
+            diagnostics,
+            "model_glance",
+            [_safe_len(variables)],
+            ["n_variables"],
+        )
+
+    if "strict_matching" in model:
+        diagnostics = _concat_metric_val_var(
+            diagnostics,
+            "model_glance",
+            [int(bool(model["strict_matching"]))],
+            ["strict_matching"],
+        )
+
+    cell_weight_ratio = model.get("cell_weight_ratio")
+    if cell_weight_ratio is not None:
+        diagnostics = _concat_metric_val_var(
+            diagnostics,
+            "model_glance",
+            [_safe_len(cell_weight_ratio)],
+            ["n_cells"],
+        )
+
+    return diagnostics
+
+
 def _build_summary(
     *,
     is_adjusted: bool,
@@ -553,7 +699,11 @@ def _build_diagnostics(
             optimizations = pd.DataFrame({"metric": metric, "var": var, "val": val})
             diagnostics = pd.concat((diagnostics, optimizations))
 
-        # TODO: add model diagnostics for other models
+        elif model["method"] == "rake":
+            diagnostics = _append_rake_model_diagnostics(diagnostics, model)
+
+        elif model["method"] == "poststratify":
+            diagnostics = _append_poststratify_model_diagnostics(diagnostics, model)
 
     # ----------------------------------------------------
     # Diagnostics on the covariates correction

tests/test_sample_diagnostics_helper.py

@@ -13,6 +13,8 @@
 import pandas as pd
 from balance.sample_class import Sample
 from balance.summary_utils import (
+    _append_poststratify_model_diagnostics,
+    _append_rake_model_diagnostics,
     _build_diagnostics,
     _build_summary,
     _concat_metric_val_var,
@@ -756,3 +758,158 @@ def get_params(self, deep=False):
         covars_asmd_main=covars_asmd,
     )
     assert not out.empty
+
+
+class _BrokenLen:
+    def __len__(self) -> int:
+        raise RuntimeError("length unavailable")
+
+
+def _minimal_diagnostics_inputs() -> dict[str, Any]:
+    covars_df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]})
+    covars_asmd = pd.DataFrame(
+        {"self": [0.1], "unadjusted": [0.2], "unadjusted - self": [0.1]},
+        index=pd.Index(["a"]),
+    )
+
+    return {
+        "covars_df": covars_df,
+        "target_covars_df": covars_df.copy(),
+        "weights_summary": pd.DataFrame({"var": ["design_effect"], "val": [1.0]}),
+        "covars_asmd": covars_asmd,
+        "covars_asmd_main": covars_asmd,
+    }
+
+
+def test_build_diagnostics_includes_rake_model_glance() -> None:
+    iterations = pd.DataFrame({"conv": [0.5, 0.01]}, index=pd.Index([0, 1]))
+    model = {
+        "method": "rake",
+        "converged": 1,
+        "iterations": iterations,
+        "variables": ["a", "b"],
+    }
+
+    out = _build_diagnostics(
+        **_minimal_diagnostics_inputs(),
+        model_dict=model,
+    )
+
+    glance = out[out["metric"] == "model_glance"].set_index("var")["val"]
+    assert glance["converged"] == 1
+    assert glance["iterations"] == 2
+    assert glance["final_conv"] == 0.01
+    assert glance["n_variables"] == 2
+
+
+def test_build_diagnostics_handles_sparse_rake_metadata() -> None:
+    model = {
+        "method": "rake",
+        "iterations": pd.DataFrame({"other": [1.0]}),
+        "variables": _BrokenLen(),
+    }
+
+    out = _build_diagnostics(
+        **_minimal_diagnostics_inputs(),
+        model_dict=model,
+    )
+
+    glance = out[out["metric"] == "model_glance"].set_index("var")["val"]
+    assert np.isnan(float(glance["converged"]))
+    assert glance["iterations"] == 1
+    assert "final_conv" not in glance.index
+    assert np.isnan(float(glance["n_variables"]))
+
+
+def test_build_diagnostics_treats_scalar_strings_as_missing_lengths() -> None:
+    for method in ("rake", "poststratify"):
+        model = {
+            "method": method,
+            "variables": "ab",
+            "cell_weight_ratio": b"xy",
+        }
+
+        out = _build_diagnostics(
+            **_minimal_diagnostics_inputs(),
+            model_dict=model,
+        )
+
+        glance = out[out["metric"] == "model_glance"].set_index("var")["val"]
+        assert np.isnan(float(glance["n_variables"]))
+        if method == "poststratify":
+            assert np.isnan(float(glance["n_cells"]))
+
+
+def test_build_diagnostics_includes_poststratify_model_glance() -> None:
+    model = {
+        "method": "poststratify",
+        "variables": ["a"],
+        "strict_matching": True,
+        "cell_weight_ratio": pd.Series([0.5, 2.0], index=["x", "y"]),
+    }
+
+    out = _build_diagnostics(
+        **_minimal_diagnostics_inputs(),
+        model_dict=model,
+    )
+
+    glance = out[out["metric"] == "model_glance"].set_index("var")["val"]
+    assert glance["n_variables"] == 1
+    assert glance["strict_matching"] == 1
+    assert glance["n_cells"] == 2
+
+
+def test_build_diagnostics_handles_sparse_poststratify_metadata() -> None:
+    model = {
+        "method": "poststratify",
+        "variables": _BrokenLen(),
+        "strict_matching": False,
+        "cell_weight_ratio": _BrokenLen(),
+    }
+
+    out = _build_diagnostics(
+        **_minimal_diagnostics_inputs(),
+        model_dict=model,
+    )
+
+    glance = out[out["metric"] == "model_glance"].set_index("var")["val"]
+    assert np.isnan(float(glance["n_variables"]))
+    assert glance["strict_matching"] == 0
+    assert np.isnan(float(glance["n_cells"]))
+
+
+def test_rake_model_diagnostics_docstring_example_output() -> None:
+    diagnostics = pd.DataFrame(columns=["metric", "val", "var"])
+    model = {
+        "method": "rake",
+        "converged": 1,
+        "iterations": pd.DataFrame({"conv": [0.5, 0.01]}),
+        "variables": ["gender", "age_group"],
+    }
+
+    out = _append_rake_model_diagnostics(diagnostics, model)
+
+    assert out.to_dict("records") == [
+        {"metric": "model_glance", "val": 1, "var": "converged"},
+        {"metric": "model_glance", "val": 2, "var": "iterations"},
+        {"metric": "model_glance", "val": 0.01, "var": "final_conv"},
+        {"metric": "model_glance", "val": 2, "var": "n_variables"},
+    ]
+
+
+def test_poststratify_model_diagnostics_docstring_example_output() -> None:
+    diagnostics = pd.DataFrame(columns=["metric", "val", "var"])
+    model = {
+        "method": "poststratify",
+        "variables": ["gender", "age_group"],
+        "strict_matching": True,
+        "cell_weight_ratio": pd.Series([0.5, 2.0]),
+    }
+
+    out = _append_poststratify_model_diagnostics(diagnostics, model)
+
+    assert out.to_dict("records") == [
+        {"metric": "model_glance", "val": 2, "var": "n_variables"},
+        {"metric": "model_glance", "val": 1, "var": "strict_matching"},
+        {"metric": "model_glance", "val": 2, "var": "n_cells"},
+    ]