[FEATURE] Improve NA handeling inside modeling

[talgalili]

As seen here: https://import-balance.org/docs/tutorials/quickstart/
When we have NA in some variable, the default modeling does this:
[ipw/ipw (line 767)]: The formula used to build the model matrix: ['income + gender + age_group + _is_na_gender']

There are several ways to do better: (each of this should be it's own PR!)

1/ It's not clear what does the model do for na in gender? If it's imputed with the mean (or something similar), it's probably worth adding an INFO about it.

2/ Adding more information about missingness in variables (in the default outputs and others), is valuable. Since missingness can be a big deal. E.g., instead of just doing this:

print(adjusted)
        Adjusted balance Sample object with target set using ipw
        1000 observations x 3 variables: gender,age_group,income
        id_column: id, weight_column: weight,
        outcome_columns: happiness

To add missing info, such as:

print(adjusted)
        Adjusted balance Sample object with target set using ipw
        1000 observations x 3 variables: gender(8% NA),age_group,income
        id_column: id, weight_column: weight,
        outcome_columns: happiness

3/ Consider more carefully what should be in
adjusted.covars().summary()

4/
Include is_na as an interaction term, instead of just additive term.
Call this transformation default_with_na.
The new structure should be:
[ipw/ipw (line 767)]: The formula used to build the model matrix: ['income + gender + age_group + _is_na_gender + gender:_is_na_gender']

More ideas?

[talgalili]

@neuralsorcerer put up some ideas here. Consider if you want to explore any of them (or if you have more ideas) - since NA statistics can be a pain.

[neuralsorcerer]

5/ Add explicit NA policy modes, so behavior is always intentional and reproducible.

• Introduce na_policy with options like: drop, indicator_only, indicator_plus_interaction, impute_plus_indicator, and custom_formula.
• Save chosen policy in model metadata and print it in logs.
• Like:
  • Add na_policy="indicator_plus_interaction" to sample.adjust(...).
  • Log: NA policy: indicator_plus_interaction.
  • Store in adjusted.model()["na_policy"].

6/ Add model metadata focused on missingness.

• Store per-variable NA rates (sample + target), generated NA terms, and rows dropped due to NA.
• Make this accessible from adjusted.model() and diagnostics outputs.
• Return in model metadata:
  • na_rates_sample = {"gender": 0.08, "income": 0.01}
  • na_rates_target = {"gender": 0.03, "income": 0.00}
  • na_terms = ["_is_na_gender", "gender:_is_na_gender"]
  • rows_dropped = 17

7/ Add pre-fit NA warnings and post-fit NA diagnostics.

• Pre-fit warnings for very high NA, large NA mismatch between sample/target, and NA concentrated in high-weight units.
• Post-fit diagnostics: ASMD for NA indicators, ESS by NA subgroup, and weight distribution by NA status.
• Like:
• Pre-fit warning: gender has 42% NA in sample vs 9% in target (diff=33pp).
• Post-fit table:
  • _is_na_gender ASMD: 0.31 -> 0.07
  • ESS (gender NA=True): 120
  • ESS (gender NA=False): 740

8/ Add a dedicated missingness report API.

• New helper like covars().missingness_summary() with counts, weighted/unweighted NA rates, and sample-vs-target gaps.
• Sort by largest NA discrepancy to prioritize what needs attention.
• Like:
  • adjusted.covars().missingness_summary(sort_by="abs_gap", top_n=5)
  • Output columns: var, na_n_sample, na_pct_sample_w, na_pct_target_w, abs_gap_pp.

9/ Add per-variable controls for NA features.

• Allow na_interactions_for=[...] and exclude_na_indicators_for=[...].
• Useful when only a subset of variables needs richer NA treatment.
• Like:
  • sample.adjust(method="ipw", na_policy="indicator_plus_interaction", na_interactions_for=["gender"], exclude_na_indicators_for=["income"])
  • Result: interaction only for gender; no NA indicator term for income.