Add formula support to covars() KLD diagnostics (#376)

Summary:
- Closes #326

Pull Request resolved: #376

Differential Revision: D98613915

Pulled By: talgalili

fbshipit-source-id: 8a8c042f317cee6c711842acaa0b362caa1f35ef

Author: neuralsorcerer

CHANGELOG.md

@@ -1,4 +1,4 @@
-# 0.19.0 (2026-03-25)
+# 0.19.0 (Unreleased - TBD)
 
 ## Breaking Changes
 
@@ -16,6 +16,17 @@
 - **Removed `Sample.outcome_variance_ratio()`** — use `sample.outcomes().outcome_variance_ratio()` instead.
   Deprecated since 0.18.0.
 
+## New Features
+
+- **Added formula support to `Sample.covars()` for downstream diagnostics**
+  - `Sample.covars()` now accepts a `formula` argument and stores it on the
+    returned `BalanceDFCovars` object.
+  - `BalanceDFCovars.kld()` now honors formula-driven model matrices (including
+    interactions such as `"age_group * gender"`) when a formula is provided via
+    `covars(formula=...)`.
+  - Formula settings are now propagated to linked covariate views (`target`,
+    `unadjusted`) so comparative diagnostics run on consistent design matrices.
+
 # 0.18.0 (2026-03-24)
 
 ## New Features

balance/balancedf_class.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import logging
+import re
 from typing import Any, Callable, Dict, Literal, Tuple
 
 import numpy as np
@@ -285,14 +286,23 @@ def _BalanceDF_child_from_linked_samples(
             | "BalanceDFOutcomes"
             | None,
         ] = {"self": self}
+        linked_child_kwargs = self._linked_child_kwargs()
         d.update(
             {
-                k: getattr(v, BalanceDF_child_method)()
+                k: getattr(v, BalanceDF_child_method)(**linked_child_kwargs)
                 for k, v in self._sample._links.items()
             }
         )
         return d
 
+    def _linked_child_kwargs(self: "BalanceDF") -> dict[str, Any]:
+        """Keyword arguments used when creating linked BalanceDF children.
+
+        Subclasses can override this to preserve construction options across linked
+        samples (for example, formula settings for covariates).
+        """
+        return {}
+
     def _call_on_linked(
         self: "BalanceDF",
         method: str,
@@ -1276,8 +1286,18 @@ def _kld_BalanceDF(
     ) -> pd.Series:
         """Run KLD on two BalanceDF objects.
 
-        Prepares the BalanceDF objects by using their raw df (with NA indicators), and
-        then passes the df and weights from the two objects into :func:`weighted_comparisons_stats.kld`.
+        By default, this prepares the BalanceDF objects by using their raw df
+        (with NA indicators), and then passes the df and weights from the two
+        objects into :func:`weighted_comparisons_stats.kld`.
+
+        If either BalanceDF provides a formula for KLD comparisons (currently
+        :class:`BalanceDFCovars` with a stored formula), this method builds a
+        *shared* model matrix from the combined sample+target data using that
+        single effective formula and compares those aligned matrices instead of
+        raw covariates.
+
+        If both objects provide formulas and they differ, a ``ValueError`` is
+        raised to prevent comparing mismatched design matrices.
 
         Args:
             sample_BalanceDF (BalanceDF): Object
@@ -1287,6 +1307,57 @@ def _kld_BalanceDF(
         Returns:
             pd.Series: See :func:`weighted_comparisons_stats.kld`.
         """
+        BalanceDF._check_if_not_BalanceDF(sample_BalanceDF, "sample_BalanceDF")
+        BalanceDF._check_if_not_BalanceDF(target_BalanceDF, "target_BalanceDF")
+
+        sample_formula = sample_BalanceDF._kld_formula()
+        target_formula = target_BalanceDF._kld_formula()
+
+        if sample_formula is not None and target_formula is not None:
+            normalized_sample_formula = BalanceDF._normalize_formula_for_comparison(
+                sample_formula
+            )
+            normalized_target_formula = BalanceDF._normalize_formula_for_comparison(
+                target_formula
+            )
+            if normalized_sample_formula != normalized_target_formula:
+                raise ValueError(
+                    "KLD formula mismatch between sample and target. "
+                    f"Got sample formula {sample_formula!r} and target formula {target_formula!r}. "
+                    "Use a single shared formula for both."
+                )
+
+        effective_formula = (
+            sample_formula if sample_formula is not None else target_formula
+        )
+        use_model_matrix = effective_formula is not None
+
+        if use_model_matrix:
+            mm = balance_util.model_matrix(
+                sample_BalanceDF.df,
+                target_BalanceDF.df,
+                add_na=True,
+                return_type="two",
+                formula=effective_formula,
+            )
+            sample_weights = (
+                sample_BalanceDF._weights.values
+                if sample_BalanceDF._weights is not None
+                else None
+            )
+            target_weights = (
+                target_BalanceDF._weights.values
+                if target_BalanceDF._weights is not None
+                else None
+            )
+            return weighted_comparisons_stats.kld(
+                _assert_type(mm["sample"], pd.DataFrame),
+                _assert_type(mm["target"], pd.DataFrame),
+                sample_weights,
+                target_weights,
+                aggregate_by_main_covar=aggregate_by_main_covar,
+            )
+
         return BalanceDF._apply_comparison_stat_to_BalanceDF(
             weighted_comparisons_stats.kld,
             sample_BalanceDF,
@@ -1295,6 +1366,23 @@ def _kld_BalanceDF(
             use_model_matrix=False,
         )
 
+    def _kld_formula(self: "BalanceDF") -> str | list[str] | None:
+        """Formula to use for KLD comparison matrices, if applicable."""
+        return None
+
+    @staticmethod
+    def _normalize_formula_for_comparison(
+        formula: str | list[str],
+    ) -> tuple[str, ...]:
+        """Normalize formulas for robust equality checks.
+
+        Insignificant whitespace is removed so equivalent formulas such as
+        ``\"a*b\"`` and ``\"a * b\"`` compare equal.
+        """
+        if isinstance(formula, str):
+            formula = [formula]
+        return tuple(re.sub(r"\s+", "", f) for f in formula)
+
     @staticmethod
     def _emd_BalanceDF(
         sample_BalanceDF: "BalanceDF",
@@ -2570,7 +2658,11 @@ def outcome_variance_ratio(self: "BalanceDFOutcomes") -> pd.Series:
 
 
 class BalanceDFCovars(BalanceDF):
-    def __init__(self: "BalanceDFCovars", sample: Sample) -> None:
+    def __init__(
+        self: "BalanceDFCovars",
+        sample: Sample,
+        formula: str | list[str] | None = None,
+    ) -> None:
         """A factory function to create BalanceDFCovars
 
         This is used through :func:`Sample.covars`.
@@ -2580,23 +2672,47 @@ def __init__(self: "BalanceDFCovars", sample: Sample) -> None:
         Args:
             self (BalanceDFCovars): Object that is initiated.
             sample (Sample): Object
+            formula (str | list[str] | None, optional): Optional formula to use
+                as the default when constructing model matrices for this object.
         """
         super().__init__(sample._covar_columns(), sample, name="covars")
+        self._formula: str | list[str] | None = formula
 
+    def model_matrix(
+        self: "BalanceDFCovars", formula: str | list[str] | None = None
+    ) -> pd.DataFrame:
+        """Return a model matrix, defaulting to the formula provided at construction."""
+        effective_formula = self._formula if formula is None else formula
+        return super().model_matrix(formula=effective_formula)
+
+    def _linked_child_kwargs(self: "BalanceDFCovars") -> dict[str, Any]:
+        """Propagate formula choice to linked covariate views."""
+        if self._formula is None:
+            return {}
+        return {"formula": self._formula}
+
+    def _kld_formula(self: "BalanceDFCovars") -> str | list[str] | None:
+        """Formula to use for KLD when comparing covariates."""
+        return self._formula
+
+    @classmethod
     def from_frame(
-        self: "BalanceDFCovars",
+        cls: type["BalanceDFCovars"],
         df: pd.DataFrame,
         weights: pd.Series | None = None,
+        formula: str | list[str] | None = None,
     ) -> "BalanceDFCovars":
         """A factory function to create a BalanceDFCovars from a df.
 
         Although generally the main way the object is created is through the __init__ method.
         This method is useful when you need to create a BalanceDFCovars object directly from a DataFrame.
 
         Args:
-            self (BalanceDFCovars): Object
+            cls (type[BalanceDFCovars]): Class object.
             df (pd.DataFrame): A df.
             weights (Optional[pd.Series], optional): _description_. Defaults to None.
+            formula (str | list[str] | None, optional): Optional formula to set on
+                the returned ``BalanceDFCovars`` object.
 
         Returns:
             BalanceDFCovars: Object.
@@ -2618,7 +2734,7 @@ def from_frame(
         if weights is not None:
             concat_list.append(weights)
         df = pd.concat(concat_list, axis=1)
-        return Sample.from_frame(df, id_column="id").covars()
+        return Sample.from_frame(df, id_column="id").covars(formula=formula)
 
 
 class BalanceDFWeights(BalanceDF):

balance/sample_class.py

@@ -766,7 +766,7 @@ def weights(
         return BalanceDFWeights(self)
 
     def covars(
-        self: "Sample",
+        self: "Sample", formula: str | list[str] | None = None
     ) -> (
         Any
     ):  # -> "Optional[Type[BalanceDFCovars]]" (not imported due to circular dependency)
@@ -776,6 +776,12 @@ def covars(
 
         Args:
             self (Sample): Sample object.
+            formula (str | list[str] | None, optional): Optional formula string
+                (or list of formulas) used when creating model matrices from the
+                returned ``BalanceDFCovars`` object. If provided, methods that
+                rely on model matrices (or distribution-comparison methods that
+                opt into model-matrix behavior for formula-based covariates) will
+                use this formula. Defaults to None.
 
         Returns:
             BalanceDFCovars
@@ -805,7 +811,7 @@ def covars(
         # NOTE: must import here so to avoid circular dependency
         from balance.balancedf_class import BalanceDFCovars
 
-        return BalanceDFCovars(self)
+        return BalanceDFCovars(self, formula=formula)
 
     def ignored_columns(self: "Sample") -> pd.DataFrame | None:
         """Return columns marked as ignored on the sample.