features.py

352 lines · 14,273 bytes · view raw

"""Engineer features used across the analysis layers.

Adds census_region, urban_rural, denominator-specific charge ratios,
drg_mdc / apc_family, and produces provider_combined.parquet (one row per
provider joined across inpatient + outpatient).
"""
import pandas as pd
import numpy as np
from .clean import clean_inpatient, clean_outpatient
from .config import (
    DATA_PROC, FIPS_TO_REGION, ruca_to_urban_rural,
    COMPATIBILITY_RATIO_ALIAS, MAPPING_METADATA, MAPPING_STATUS,
    MAPPING_VALIDATION_EXAMPLES, drg_mapping_quarantine_reason,
    drg_to_mdc, apc_to_family,
)


# ---------------------------------------------------------------------------
# Per-row feature engineering
# ---------------------------------------------------------------------------
def _safe_ratio(numerator: pd.Series, denominator: pd.Series) -> pd.Series:
    return (numerator / denominator).replace([np.inf, -np.inf], np.nan)


def add_features_inpatient(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["census_region"] = df["state_fips"].map(FIPS_TO_REGION).fillna("Unknown")
    df["urban_rural"] = df["ruca_code"].map(ruca_to_urban_rural)
    df["drg_mdc"] = df.apply(lambda r: drg_to_mdc(r["drg_code"], r["drg_desc"]), axis=1)
    df["mapping_quarantined"] = df.apply(
        lambda r: drg_mapping_quarantine_reason(r["drg_code"], r["drg_desc"]) is not None,
        axis=1,
    )
    df["is_maryland_exception"] = df["state"].eq("MD")
    df["charge_to_total_payment_ratio"] = _safe_ratio(
        df["avg_submitted_charge"], df["avg_total_payment"]
    )
    # Compatibility alias required by the unchanged analysis modules. The
    # canonical field above names the actual inpatient denominator.
    df[COMPATIBILITY_RATIO_ALIAS] = df["charge_to_total_payment_ratio"]
    return df


def add_features_outpatient(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["census_region"] = df["state_fips"].map(FIPS_TO_REGION).fillna("Unknown")
    df["urban_rural"] = df["ruca_code"].map(ruca_to_urban_rural)
    df["apc_family"] = df["apc_desc"].map(apc_to_family)
    df["mapping_quarantined"] = False
    df["is_maryland_exception"] = df["state"].eq("MD")
    # The outpatient denominator is the allowed amount, not Medicare-only
    # payment. Suppressed allowed amounts therefore produce missing ratios.
    df["charge_to_allowed_amount_ratio"] = _safe_ratio(
        df["avg_submitted_charge"], df["avg_allowed_amount"]
    )
    # Compatibility alias required by the unchanged analysis modules. The
    # canonical field above names the actual outpatient denominator.
    df[COMPATIBILITY_RATIO_ALIAS] = df["charge_to_allowed_amount_ratio"]
    return df


# ---------------------------------------------------------------------------
# Provider-level aggregation + cross-dataset join
# ---------------------------------------------------------------------------
def provider_summary_inpatient(df: pd.DataFrame) -> pd.DataFrame:
    """One row per provider summarizing inpatient activity."""
    g = df.groupby("provider_ccn", as_index=False).agg(
        provider_name=("provider_name", "first"),
        state=("state", "first"),
        state_fips=("state_fips", "first"),
        census_region=("census_region", "first"),
        urban_rural=("urban_rural", "first"),
        is_maryland_exception=("is_maryland_exception", "first"),
        inp_drgs_seen=("drg_code", "nunique"),
        inp_total_discharges=("total_discharges", "sum"),
        inp_avg_payment=("avg_total_payment", "mean"),
        inp_median_payment=("avg_total_payment", "median"),
        inp_avg_charge_to_total_payment_ratio=("charge_to_total_payment_ratio", "mean"),
    )
    return g


def provider_summary_outpatient(df: pd.DataFrame) -> pd.DataFrame:
    """One row per provider summarizing outpatient activity (cost rows only)."""
    g = df.groupby("provider_ccn", as_index=False).agg(
        provider_name=("provider_name", "first"),
        state=("state", "first"),
        state_fips=("state_fips", "first"),
        census_region=("census_region", "first"),
        urban_rural=("urban_rural", "first"),
        is_maryland_exception=("is_maryland_exception", "first"),
        out_apcs_seen=("apc_code", "nunique"),
        out_total_services=("apc_services", "sum"),
        out_avg_allowed=("avg_allowed_amount", "mean"),
        out_median_allowed=("avg_allowed_amount", "median"),
        out_avg_charge_to_allowed_amount_ratio=("charge_to_allowed_amount_ratio", "mean"),
    )
    return g


def _outpatient_presence_summary(df: pd.DataFrame) -> pd.DataFrame:
    """Summarize full-frame outpatient presence independently of suppression."""
    aggregations = {
        "provider_name": ("provider_name", "first"),
        "state": ("state", "first"),
        "state_fips": ("state_fips", "first"),
    }
    for column in ("census_region", "urban_rural"):
        if column in df.columns:
            aggregations[column] = (column, "first")
    result = df.groupby("provider_ccn", as_index=False).agg(**aggregations)
    result["out_full_rows"] = df.groupby("provider_ccn").size().reindex(
        result["provider_ccn"]
    ).to_numpy()
    result["out_full_apcs_seen"] = df.groupby("provider_ccn")["apc_code"].nunique().reindex(
        result["provider_ccn"]
    ).to_numpy()
    return result


def build_provider_combined(
    inp_feat: pd.DataFrame,
    out_feat: pd.DataFrame,
    out_full_feat: pd.DataFrame | None = None,
) -> pd.DataFrame:
    """Join inpatient metrics to cost-observed metrics and full outpatient presence."""
    inp_s = provider_summary_inpatient(inp_feat)
    out_s = provider_summary_outpatient(out_feat)

    merged = inp_s.merge(out_s, on="provider_ccn", how="outer", suffixes=("_inp", "_out"))
    if out_full_feat is not None:
        merged = merged.merge(
            _outpatient_presence_summary(out_full_feat),
            on="provider_ccn",
            how="outer",
        )

    # Reconcile geography explicitly: inpatient wins, then full outpatient
    # presence, then the cost-observed outpatient summary.
    for column in ("provider_name", "state", "state_fips", "census_region", "urban_rural"):
        candidates = [f"{column}_inp", column, f"{column}_out"]
        available = [merged[candidate] for candidate in candidates if candidate in merged]
        if available:
            value = available[0]
            for candidate in available[1:]:
                value = value.fillna(candidate)
            merged[column] = value
    maryland_candidates = [
        merged[column] for column in (
            "is_maryland_exception_inp", "is_maryland_exception_out"
        ) if column in merged
    ]
    merged["is_maryland_exception"] = (
        maryland_candidates[0].fillna(maryland_candidates[1])
        if len(maryland_candidates) == 2
        else maryland_candidates[0]
        if maryland_candidates
        else False
    ).fillna(False).astype(bool)
    merged = merged.drop(columns=[
        c for c in merged.columns
        if c.endswith(("_inp", "_out"))
        and c not in (
            "provider_name", "state", "state_fips", "census_region",
            "urban_rural", "is_maryland_exception",
        )
    ])
    # Boolean presence flags — useful for outlier profiling
    merged["has_inpatient"] = merged["inp_total_discharges"].notna()
    merged["has_outpatient_cost_observed"] = merged["out_total_services"].notna()
    merged["has_outpatient_full"] = merged.get("out_full_rows", pd.Series(index=merged.index)).notna()
    merged["has_outpatient"] = merged["has_outpatient_full"]
    merged["outpatient_presence_basis"] = "full outpatient frame"
    return merged


# ---------------------------------------------------------------------------
# Mapping validation
# ---------------------------------------------------------------------------
_UNMAPPED_LABELS = {"Unknown", "Other", "Other / Unclassified"}


def _normalise_mapping_code(value, width: int) -> str | None:
    if pd.isna(value):
        return None
    text = str(value).strip()
    if text.endswith(".0") and text[:-2].isdigit():
        text = text[:-2]
    return text.zfill(width) if text.isdigit() else text


def _validate_mapping(
    df: pd.DataFrame,
    code_col: str,
    desc_col: str,
    mapper,
    examples: list[dict],
    code_width: int,
    feature_field: str,
    quarantine_fn=None,
) -> dict:
    observed = df[[code_col, desc_col]].copy()
    observed["_mapping_code"] = observed[code_col].map(
        lambda value: _normalise_mapping_code(value, code_width)
    )
    observed = observed[observed["_mapping_code"].notna()]
    observed_rows = len(observed)
    observed_pairs = observed.drop_duplicates(["_mapping_code", desc_col])

    labels_by_code = {}
    quarantine_by_code = {}
    for code, group in observed_pairs.groupby("_mapping_code", sort=True):
        descriptions = group[desc_col].where(group[desc_col].notna(), "")
        labels_by_code[code] = sorted({mapper(code, desc) for desc in descriptions})
        if quarantine_fn:
            reasons = {
                quarantine_fn(code, desc)
                for desc in descriptions
                if quarantine_fn(code, desc)
            }
            if reasons:
                quarantine_by_code[code] = sorted(reasons)

    unmapped_codes = sorted(
        code for code, labels in labels_by_code.items()
        if not labels or all(label in _UNMAPPED_LABELS for label in labels)
    )
    ambiguous_codes = sorted(
        code for code, labels in labels_by_code.items() if len(labels) > 1
    )

    validation_examples = []
    for example in examples:
        code = _normalise_mapping_code(example["code"], code_width)
        matches = observed_pairs[observed_pairs["_mapping_code"] == code]
        description = ""
        if not matches.empty:
            value = matches.iloc[0][desc_col]
            description = "" if pd.isna(value) else str(value)
        actual = mapper(code, description)
        marker = example["description_contains"]
        description_check = marker.upper() in description.upper()
        validation_examples.append({
            "code": example["code"],
            "observed": not matches.empty,
            "observed_description": description,
            "expected_mapping": example["expected_mapping"],
            "actual_mapping": actual,
            "description_contains": marker,
            "description_check_passed": description_check,
            "passed": bool(
                not matches.empty
                and description_check
                and actual == example["expected_mapping"]
            ),
        })

    observed_code_count = len(labels_by_code)
    return {
        "mapping_status": MAPPING_STATUS,
        "feature_field": feature_field,
        "code_column": code_col,
        "description_column": desc_col,
        "observed_rows": int(observed_rows),
        "observed_code_count": int(observed_code_count),
        "observed_codes": sorted(labels_by_code),
        "mapped_code_count": int(observed_code_count - len(unmapped_codes)),
        "unmapped_code_count": int(len(unmapped_codes)),
        "unmapped_count": int(len(unmapped_codes)),
        "unmapped_codes": unmapped_codes,
        "ambiguity_rule": "A code is ambiguous when its observed descriptions produce more than one distinct exploratory label.",
        "ambiguous_code_count": int(len(ambiguous_codes)),
        "ambiguous_count": int(len(ambiguous_codes)),
        "ambiguous_codes": ambiguous_codes,
        "validation_examples": validation_examples,
        "quarantined_code_count": int(len(quarantine_by_code)),
        "quarantined_codes": [
            {"code": code, "reasons": reasons}
            for code, reasons in sorted(quarantine_by_code.items())
        ],
        "quarantined_row_count": int(
            observed["_mapping_code"].isin(quarantine_by_code).sum()
        ),
        "modeling_policy": (
            "Rows carrying a quarantined DRG mapping are excluded from predictor modeling."
            if quarantine_fn else "No APC mapping quarantine rules are defined."
        ),
    }


def validate_observed_mappings(
    inpatient_df: pd.DataFrame, outpatient_df: pd.DataFrame
) -> dict:
    """Validate observed code coverage without implying an official crosswalk."""
    result = dict(MAPPING_METADATA)
    result["mapping_status"] = MAPPING_STATUS
    result["drg"] = _validate_mapping(
        inpatient_df,
        "drg_code",
        "drg_desc",
        drg_to_mdc,
        MAPPING_VALIDATION_EXAMPLES["drg"],
        code_width=3,
        feature_field="drg_mdc",
        quarantine_fn=drg_mapping_quarantine_reason,
    )
    result["apc"] = _validate_mapping(
        outpatient_df,
        "apc_code",
        "apc_desc",
        lambda _code, desc: apc_to_family(desc),
        MAPPING_VALIDATION_EXAMPLES["apc"],
        code_width=4,
        feature_field="apc_family",
    )
    return result


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------
def run() -> dict:
    """Run the full feature engineering step. Returns dict of artifacts."""
    inp_clean = clean_inpatient()
    out_cost = clean_outpatient(drop_suppressed=True)
    out_full = clean_outpatient(drop_suppressed=False)

    inp_feat = add_features_inpatient(inp_clean)
    out_feat = add_features_outpatient(out_cost)
    out_full_feat = add_features_outpatient(out_full)

    provider_combined = build_provider_combined(inp_feat, out_feat, out_full_feat)

    inp_feat.to_parquet(DATA_PROC / "inpatient_features.parquet", index=False)
    out_feat.to_parquet(DATA_PROC / "outpatient_features.parquet", index=False)
    out_full_feat.to_parquet(DATA_PROC / "outpatient_full_features.parquet", index=False)
    provider_combined.to_parquet(DATA_PROC / "provider_combined.parquet", index=False)

    return {
        "inpatient_features": inp_feat,
        "outpatient_features": out_feat,
        "provider_combined": provider_combined,
    }


if __name__ == "__main__":  # pragma: no cover
    res = run()
    for k, df in res.items():
        print(f"{k}: {df.shape}")
        print(df.head(2).to_string())
        print()