Spaces:
Paused
Paused
| """ | |
| SHAP attribution for the Phase 5 hybrid model. | |
| The fusion engine is a *pipeline*, not a single estimator, so there is no | |
| one model to hand to SHAP. The final risk is | |
| risk = w · P_weather(weather features) + (1 - w) · P_satellite(scene) | |
| which means an input's effect on the decision runs through whichever | |
| branch consumes it. This module therefore attributes by the chain rule: | |
| SHAP on the branch that actually reads the input, scaled by that branch's | |
| exact fusion weight. | |
| Why this matters concretely: the shipped v1 fusion is `WeightedFusion` | |
| with `w = 1.00`. Perturbing humidity, pressure or cloud density changes | |
| its output by exactly 0.000000 — the blender reads only the two upstream | |
| probability vectors. Running SHAP against the blender would honestly | |
| return zero for every meteorological feature, and any non-zero number | |
| printed next to "Humidity" would be fabricated. Attribution therefore | |
| targets the Phase 3 estimator (which does read humidity) and is then | |
| scaled by `w`. | |
| Three strategies are supported, selected automatically: | |
| | Selected fusion | Attribution path | SHAP method | | |
| |---|---|---| | |
| | `WeightedFusion` | per-branch, scaled by the branch weight | exact TreeSHAP on the Phase 3 forest | | |
| | `FeatureFusion` | directly over the 22-feature fusion vector | exact TreeSHAP on the fusion tree ensemble | | |
| | `NeuralFusion` | directly over the 22-feature fusion vector | KernelSHAP (model-agnostic, sampled) | | |
| Inputs the selected model provably does not consume are reported in | |
| `unattributed_inputs` with the reason, never given an invented score. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Any | |
| import numpy as np | |
| import pandas as pd | |
| from explainability.shap_explainer.config import ShapConfig | |
| logger = logging.getLogger(__name__) | |
| class ExplainerNotAvailableError(RuntimeError): | |
| """Raised when the artifacts SHAP needs are missing.""" | |
| class InvalidExplanationInputError(ValueError): | |
| """Raised when the row handed to the explainer is malformed.""" | |
| class FeatureContribution: | |
| """One feature's signed effect on the fused risk score.""" | |
| feature: str | |
| label: str | |
| value: float | None | |
| #: Signed SHAP value, already expressed in fused-risk probability units. | |
| contribution: float | |
| #: Share of the total absolute attribution, signed, in [-1, 1]. | |
| share: float | |
| impact: str | |
| direction: str | |
| branch: str | |
| #: True when the model was scored on a training median because the | |
| #: record carried no observation. The attribution is real, but it | |
| #: describes the median, so the narrative layer skips these. | |
| imputed: bool = False | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "feature": self.feature, | |
| "label": self.label, | |
| "value": None if self.value is None else round(float(self.value), 4), | |
| "contribution": round(float(self.contribution), 6), | |
| "contribution_pct": round(float(self.share) * 100, 1), | |
| "impact": self.impact, | |
| "direction": self.direction, | |
| "branch": self.branch, | |
| "imputed": self.imputed, | |
| } | |
| class ShapExplanation: | |
| """Everything the narrative and report layers need from SHAP.""" | |
| contributions: list[FeatureContribution] | |
| base_value: float | |
| #: Risk reconstructed by the explained models: base + sum(attributions). | |
| predicted_risk: float | |
| method: str | |
| branch_weights: dict[str, float] | |
| branch_contributions: dict[str, float] | |
| #: What the stored fusion strategy returns for this row. Differs from | |
| #: `predicted_risk` on dataset rows, which carry out-of-fold upstream | |
| #: probabilities rather than the shipped artifacts' output. | |
| served_risk: float | None = None | |
| unattributed_inputs: list[dict[str, str]] = field(default_factory=list) | |
| notes: list[str] = field(default_factory=list) | |
| def top(self, count: int) -> list[FeatureContribution]: | |
| return sorted(self.contributions, key=lambda c: -abs(c.contribution))[:count] | |
| def to_dict(self, top_k: int) -> dict[str, Any]: | |
| return { | |
| "method": self.method, | |
| "base_value": round(float(self.base_value), 6), | |
| "predicted_risk": round(float(self.predicted_risk), 6), | |
| "served_risk": None if self.served_risk is None else round(float(self.served_risk), 6), | |
| "branch_weights": {k: round(float(v), 4) for k, v in self.branch_weights.items()}, | |
| "branch_contributions": { | |
| k: round(float(v), 6) for k, v in self.branch_contributions.items() | |
| }, | |
| "top_features": [c.to_dict() for c in self.top(top_k)], | |
| "all_features": [c.to_dict() for c in self.contributions], | |
| "unattributed_inputs": self.unattributed_inputs, | |
| "notes": self.notes, | |
| } | |
| class FusionShapExplainer: | |
| """Computes SHAP attributions for a trained Phase 5 fusion bundle.""" | |
| def __init__(self, bundle: dict[str, Any], config: ShapConfig | None = None) -> None: | |
| self.config = config or ShapConfig() | |
| self.bundle = bundle | |
| self.strategy = bundle["strategy"] | |
| self._weather_bundle: dict[str, Any] | None = None | |
| self._tree_explainer: Any = None | |
| self._kernel_explainer: Any = None | |
| self._region_coordinates_cache: dict[str, dict[str, float]] | None = None | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def explain(self, row: pd.DataFrame) -> ShapExplanation: | |
| """Attribute one region-day's fused risk to its inputs. | |
| Args: | |
| row: single-row frame carrying the fusion feature columns and | |
| the upstream probability columns. | |
| """ | |
| if not isinstance(row, pd.DataFrame) or len(row) != 1: | |
| raise InvalidExplanationInputError( | |
| f"Expected a single-row DataFrame, got {type(row).__name__} of length " | |
| f"{len(row) if hasattr(row, '__len__') else 'n/a'}" | |
| ) | |
| from ai_models.fusion_model.fusion import FeatureFusion, NeuralFusion, WeightedFusion | |
| if isinstance(self.strategy, WeightedFusion): | |
| return self._explain_weighted(row) | |
| if isinstance(self.strategy, FeatureFusion): | |
| return self._explain_tree_fusion(row) | |
| if isinstance(self.strategy, NeuralFusion): | |
| return self._explain_kernel_fusion(row) | |
| raise ExplainerNotAvailableError( | |
| f"No SHAP path implemented for fusion strategy {type(self.strategy).__name__}" | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Case A — weighted late fusion: attribute through the branches | |
| # ------------------------------------------------------------------ | |
| def _explain_weighted(self, row: pd.DataFrame) -> ShapExplanation: | |
| weight = float(self.strategy.weather_weight) | |
| weights = {"weather": weight, "satellite": 1.0 - weight} | |
| weather_shap, weather_base, weather_risk, scored, imputed = self._weather_branch_shap(row) | |
| contributions = [ | |
| self._contribution( | |
| feature=name, | |
| # The value the forest actually saw, median-filled or not. | |
| value=float(scored.iloc[0][name]), | |
| contribution=float(value) * weight, | |
| branch="weather", | |
| imputed=name in imputed, | |
| ) | |
| for name, value in weather_shap.items() | |
| ] | |
| # The satellite branch has no per-feature tabular attribution — it | |
| # reads pixels. Its effect on the fused score is exact and analytic: | |
| # its weight times how far its risk sits from the branch baseline. | |
| satellite_risk = self._risk_from(row, "satellite_prob") | |
| satellite_base = self._satellite_baseline() | |
| satellite_contribution = (1.0 - weight) * (satellite_risk - satellite_base) | |
| if abs(1.0 - weight) > 1e-9: | |
| contributions.append( | |
| self._contribution( | |
| feature="satellite_risk_score", | |
| value=satellite_risk, | |
| contribution=satellite_contribution, | |
| branch="satellite", | |
| ) | |
| ) | |
| contributions = self._rescore_shares(contributions) | |
| base_value = weight * weather_base + (1.0 - weight) * satellite_base | |
| # Computed from the branches that were actually explained, so that | |
| # base + sum(attributions) reconciles exactly. The stored column may | |
| # differ: dataset rows carry out-of-fold probabilities, while serving | |
| # (and this explanation) use the shipped Phase 3 artifact. | |
| predicted = weight * weather_risk + (1.0 - weight) * satellite_risk | |
| served = float(self.strategy.predict_proba(row)[0][1:].sum()) | |
| notes = [ | |
| "Attribution follows the fusion chain rule: SHAP on the branch that reads each " | |
| f"input, scaled by that branch's weight (weather {weight:.2f}, " | |
| f"satellite {1 - weight:.2f}).", | |
| "Weather attributions are exact TreeSHAP values on the Phase 3 forest, so the " | |
| "baseline plus the attributions reconstruct the explained risk exactly.", | |
| ] | |
| if abs(served - predicted) > 1e-6: | |
| notes.append( | |
| f"The stored fusion output for this row is {served:.4f} against the explained " | |
| f"{predicted:.4f}. Dataset rows carry out-of-fold upstream probabilities, while " | |
| "this explanation runs the shipped Phase 3 artifact — the one that serves live " | |
| "predictions. The gap is the out-of-fold correction, not an explanation error." | |
| ) | |
| unattributed: list[dict[str, str]] = [] | |
| if imputed: | |
| unattributed.append( | |
| { | |
| "inputs": ", ".join(imputed), | |
| "reason": "not present in the region-day record; the forest was scored on " | |
| "the training median, so any attribution shown reflects that " | |
| "median rather than an observed value", | |
| } | |
| ) | |
| if abs(1.0 - weight) < 1e-9: | |
| notes.append( | |
| "The validation weight sweep gave the satellite branch weight 0.00, so it " | |
| "contributes exactly nothing to this decision. The Grad-CAM heatmap therefore " | |
| "explains the satellite model's own same-day reading, not the fused forecast." | |
| ) | |
| unattributed.append( | |
| { | |
| "inputs": "satellite scene features and satellite risk score", | |
| "reason": "satellite branch weight is 0.00 in the selected fusion", | |
| } | |
| ) | |
| # Scene statistics only enter a feature-level fusion, never this one. | |
| scene_features = [ | |
| name | |
| for name in ("cloud_density", "brightness_mean", "spatial_dispersion", | |
| "cold_top_fraction", "cloud_growth_rate") | |
| if name in row.columns | |
| ] | |
| if scene_features: | |
| unattributed.append( | |
| { | |
| "inputs": ", ".join(scene_features), | |
| "reason": "the selected weighted fusion consumes only the two upstream " | |
| "probability vectors, so these columns cannot move its output", | |
| } | |
| ) | |
| return ShapExplanation( | |
| contributions=contributions, | |
| base_value=base_value, | |
| predicted_risk=predicted, | |
| served_risk=served, | |
| method="chain-rule TreeSHAP through the weather branch", | |
| branch_weights=weights, | |
| branch_contributions={ | |
| "weather": weight * (weather_risk - weather_base), | |
| "satellite": satellite_contribution, | |
| }, | |
| unattributed_inputs=unattributed, | |
| notes=notes, | |
| ) | |
| def _weather_branch_shap( | |
| self, row: pd.DataFrame | |
| ) -> tuple[dict[str, float], float, float, pd.DataFrame, list[str]]: | |
| """Exact TreeSHAP over the Phase 3 forest, reduced to risk units. | |
| Returns the attributions, the branch baseline and prediction, the | |
| exact frame the forest was scored on, and the names that had to be | |
| filled from training medians — the caller reports those rather than | |
| showing a blank value beside a real attribution. | |
| """ | |
| import shap | |
| bundle = self._load_weather_bundle() | |
| names = list(bundle["feature_names"]) | |
| medians = bundle.get("feature_medians", {}) | |
| # The fusion dataset aggregates to region-days and drops the point | |
| # coordinates the Phase 3 forest was trained on. Falling straight | |
| # back to the global median put ~20% of the attribution mass on a | |
| # constant, so resolve the region's real centroid first. | |
| coordinates = self._region_coordinates(_region_of(row)) | |
| values: dict[str, float] = {} | |
| imputed: list[str] = [] | |
| for name in names: | |
| supplied = self._value_of(row, name) | |
| if supplied is None: | |
| supplied = coordinates.get(name) | |
| if supplied is None: | |
| supplied = float(medians.get(name, 0.0)) | |
| imputed.append(name) | |
| values[name] = float(supplied) | |
| frame = pd.DataFrame([values]) | |
| if self._tree_explainer is None: | |
| self._tree_explainer = shap.TreeExplainer(bundle["estimator"]) | |
| values = np.asarray( | |
| self._tree_explainer.shap_values(frame, check_additivity=False), dtype=np.float64 | |
| ) | |
| # TreeSHAP returns (samples, features, classes) for multiclass forests. | |
| if values.ndim == 3: | |
| risk_values = values[0, :, 1:].sum(axis=1) | |
| else: # binary estimator — the positive column is already the risk | |
| risk_values = values[0] | |
| expected = np.atleast_1d(np.asarray(self._tree_explainer.expected_value, dtype=np.float64)) | |
| base = float(expected[1:].sum()) if expected.size >= 3 else float(expected[-1]) | |
| probabilities = bundle["estimator"].predict_proba(frame)[0] | |
| risk = float(np.asarray(probabilities)[1:].sum()) | |
| return dict(zip(names, risk_values)), base, risk, frame, imputed | |
| def _satellite_baseline(self) -> float: | |
| """Mean satellite risk over the training record — the branch's prior.""" | |
| frame = self._reference_frame() | |
| if frame is None or "satellite_risk_score" not in frame: | |
| return 0.0 | |
| return float(frame["satellite_risk_score"].mean()) | |
| # ------------------------------------------------------------------ | |
| # Case B — feature-level fusion: SHAP directly on the fusion estimator | |
| # ------------------------------------------------------------------ | |
| def _explain_tree_fusion(self, row: pd.DataFrame) -> ShapExplanation: | |
| import shap | |
| names = list(self.bundle["feature_names"]) | |
| frame = row[names].astype(float) | |
| if self._tree_explainer is None: | |
| self._tree_explainer = shap.TreeExplainer(self.strategy.estimator) | |
| values = np.asarray( | |
| self._tree_explainer.shap_values(frame, check_additivity=False), dtype=np.float64 | |
| ) | |
| risk_values = values[0, :, 1:].sum(axis=1) if values.ndim == 3 else values[0] | |
| expected = np.atleast_1d(np.asarray(self._tree_explainer.expected_value, dtype=np.float64)) | |
| base = float(expected[1:].sum()) if expected.size >= 3 else float(expected[-1]) | |
| contributions = self._rescore_shares( | |
| [ | |
| self._contribution( | |
| feature=name, | |
| value=self._value_of(row, name), | |
| contribution=float(value), | |
| branch=self._branch_of(name), | |
| ) | |
| for name, value in zip(names, risk_values) | |
| ] | |
| ) | |
| return ShapExplanation( | |
| contributions=contributions, | |
| base_value=base, | |
| predicted_risk=float(self.strategy.predict_proba(row)[0][1:].sum()), | |
| method="exact TreeSHAP on the feature-level fusion", | |
| branch_weights={"weather": float("nan"), "satellite": float("nan")}, | |
| branch_contributions=self._branch_totals(contributions), | |
| notes=[ | |
| "The feature-level fusion reads every input directly, so attributions are " | |
| "exact TreeSHAP values on the fusion estimator itself.", | |
| ], | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Case C — neural fusion: model-agnostic KernelSHAP | |
| # ------------------------------------------------------------------ | |
| def _explain_kernel_fusion(self, row: pd.DataFrame) -> ShapExplanation: | |
| import shap | |
| names = list(self.bundle["feature_names"]) | |
| background = self._background_sample(names) | |
| frame = row[names].astype(float) | |
| def predict(matrix: np.ndarray) -> np.ndarray: | |
| batch = pd.DataFrame(matrix, columns=names) | |
| return self.strategy.predict_proba(batch)[:, 1:].sum(axis=1) | |
| if self._kernel_explainer is None: | |
| self._kernel_explainer = shap.KernelExplainer(predict, background) | |
| values = np.asarray( | |
| self._kernel_explainer.shap_values(frame.to_numpy(), silent=True), dtype=np.float64 | |
| ) | |
| risk_values = values[0] if values.ndim > 1 else values | |
| contributions = self._rescore_shares( | |
| [ | |
| self._contribution( | |
| feature=name, | |
| value=self._value_of(row, name), | |
| contribution=float(value), | |
| branch=self._branch_of(name), | |
| ) | |
| for name, value in zip(names, risk_values) | |
| ] | |
| ) | |
| return ShapExplanation( | |
| contributions=contributions, | |
| base_value=float(np.asarray(self._kernel_explainer.expected_value).ravel()[0]), | |
| predicted_risk=float(self.strategy.predict_proba(row)[0][1:].sum()), | |
| method=f"KernelSHAP on the neural fusion ({len(background)} background rows)", | |
| branch_weights={"weather": float("nan"), "satellite": float("nan")}, | |
| branch_contributions=self._branch_totals(contributions), | |
| notes=[ | |
| "KernelSHAP is sampled, not exact: values are approximations whose precision " | |
| f"is bounded by the {len(background)}-row background set.", | |
| ], | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Shared helpers | |
| # ------------------------------------------------------------------ | |
| def _contribution( | |
| self, *, feature: str, value: float | None, contribution: float, branch: str, | |
| imputed: bool = False, | |
| ) -> FeatureContribution: | |
| inverted = feature in self.config.inverted_features | |
| rising = contribution > 0 | |
| if abs(contribution) < 1e-12: | |
| direction = "no effect" | |
| elif inverted: | |
| direction = "falling — raises risk" if rising else "rising — lowers risk" | |
| else: | |
| direction = "raises risk" if rising else "lowers risk" | |
| return FeatureContribution( | |
| feature=feature, | |
| label=self.config.label_for(feature), | |
| value=value, | |
| contribution=contribution, | |
| share=0.0, # filled by _rescore_shares once the total is known | |
| impact="Low", | |
| direction=direction, | |
| branch=branch, | |
| imputed=imputed, | |
| ) | |
| def _rescore_shares(self, contributions: list[FeatureContribution]) -> list[FeatureContribution]: | |
| """Express each attribution as a signed share of the total magnitude, | |
| then label its impact from that measured share.""" | |
| total = sum(abs(c.contribution) for c in contributions) | |
| for item in contributions: | |
| item.share = 0.0 if total <= 0 else item.contribution / total | |
| item.impact = self.config.impact_for(item.share) | |
| return contributions | |
| def _branch_totals(contributions: list[FeatureContribution]) -> dict[str, float]: | |
| totals: dict[str, float] = {} | |
| for item in contributions: | |
| totals[item.branch] = totals.get(item.branch, 0.0) + item.contribution | |
| return totals | |
| def _branch_of(feature: str) -> str: | |
| satellite = { | |
| "cloud_density", "brightness_mean", "brightness_std", "spatial_dispersion", | |
| "cold_top_fraction", "cloud_growth_rate", "valid_fraction", "satellite_risk_score", | |
| } | |
| return "satellite" if feature in satellite else "weather" | |
| def _value_of(row: pd.DataFrame, name: str, default: float | None = None) -> float | None: | |
| if name not in row.columns: | |
| return default | |
| value = row.iloc[0][name] | |
| if value is None or (isinstance(value, float) and np.isnan(value)): | |
| return default | |
| return float(value) | |
| def _risk_from(row: pd.DataFrame, prefix: str) -> float: | |
| columns = [c for c in row.columns if c.startswith(prefix)] | |
| if not columns: | |
| return 0.0 | |
| ordered = sorted(columns, key=lambda c: int(c.rsplit("_", 1)[-1])) | |
| return float(sum(float(row.iloc[0][c]) for c in ordered[1:])) | |
| def _load_weather_bundle(self) -> dict[str, Any]: | |
| if self._weather_bundle is None: | |
| from ai_models.baseline.config import BaselineConfig | |
| from ai_models.baseline.model import BUNDLE_FILENAME_TEMPLATE, load_bundle | |
| path = self.config.saved_models_dir / BUNDLE_FILENAME_TEMPLATE.format( | |
| version=BaselineConfig().model_version | |
| ) | |
| if not path.exists(): | |
| raise ExplainerNotAvailableError( | |
| f"Phase 3 model missing at {path}; SHAP cannot attribute the weather branch." | |
| ) | |
| self._weather_bundle = load_bundle(path) | |
| return self._weather_bundle | |
| def _region_coordinates(self, region: str | None) -> dict[str, float]: | |
| """Mean latitude/longitude of a region's observation points. | |
| The Phase 3 forest reads coordinates as features, so a region-day | |
| must be scored at the region's real location rather than at the | |
| dataset-wide median, which belongs to no region at all. | |
| """ | |
| if region is None: | |
| return {} | |
| if self._region_coordinates_cache is None: | |
| self._region_coordinates_cache = {} | |
| from ai_models.fusion_model.config import FusionConfig | |
| path = FusionConfig().weather_dataset_path | |
| if not path.exists(): | |
| logger.warning("Weather dataset absent at %s; coordinates unresolved", path) | |
| return {} | |
| points = pd.read_parquet(path, columns=["region", "latitude", "longitude"]) | |
| means = points.groupby("region")[["latitude", "longitude"]].mean() | |
| self._region_coordinates_cache = { | |
| str(name): {"latitude": float(values["latitude"]), | |
| "longitude": float(values["longitude"])} | |
| for name, values in means.iterrows() | |
| } | |
| return self._region_coordinates_cache.get(str(region), {}) | |
| def _reference_frame(self) -> pd.DataFrame | None: | |
| if not self.config.fusion_dataset_path.exists(): | |
| logger.warning("Fusion dataset absent at %s", self.config.fusion_dataset_path) | |
| return None | |
| return pd.read_parquet(self.config.fusion_dataset_path) | |
| def _background_sample(self, names: list[str]) -> np.ndarray: | |
| frame = self._reference_frame() | |
| if frame is None: | |
| raise ExplainerNotAvailableError( | |
| "KernelSHAP needs the fusion dataset as a background distribution." | |
| ) | |
| sample = frame[names].astype(float) | |
| if len(sample) > self.config.background_size: | |
| sample = sample.sample(self.config.background_size, random_state=0) | |
| return sample.to_numpy() | |
| def _region_of(row: pd.DataFrame) -> str | None: | |
| """The region name carried by a fusion row, if any.""" | |
| if "region" not in row.columns: | |
| return None | |
| value = row.iloc[0]["region"] | |
| return None if value is None else str(value) | |
| def load_fusion_bundle(config: ShapConfig | None = None) -> dict[str, Any]: | |
| """Load the trained Phase 5 bundle, or explain why it is unavailable.""" | |
| import joblib | |
| from ai_models.fusion_model.config import FUSION_BUNDLE_TEMPLATE, FusionConfig | |
| config = config or ShapConfig() | |
| fusion_config = FusionConfig() | |
| path: Path = config.saved_models_dir / FUSION_BUNDLE_TEMPLATE.format( | |
| version=fusion_config.model_version | |
| ) | |
| if not path.exists(): | |
| raise ExplainerNotAvailableError( | |
| f"No fusion model at {path}. Run `python -m ai_models.fusion_model.train` first." | |
| ) | |
| return joblib.load(path) | |