test_helio_obs.py

#!/usr/bin/env python3
"""
test_helio_obs.py
=================
Diagnostic: prove SWPC helio reaches the RL observation vector, and that
quiet vs storm BasinContext produces a measurable basin_context delta.

Does NOT claim deterministic crop_risk_scorer benefit (scorer has no helio
term by design until a validated tropical crop mechanism is specified).

Run on Kaggle (ROOT on PYTHONPATH) or offline with synthetic BasinContext.
"""
from __future__ import annotations

import os
import sys
from datetime import datetime, timezone

import numpy as np

for p in (
    "/kaggle/working",
    "/kaggle/input/datasets/dhmmmreally/weather-modeller",
    os.path.dirname(os.path.abspath(__file__)),
):
    if p and os.path.isdir(p) and p not in sys.path:
        sys.path.insert(0, p)

from zone_observation import (
    BasinContext,
    DataSource,
    EpisodeContext,
    ForecastConfig,
    make_synthetic_forecast_result,
    make_synthetic_zone_obs,
)
from weather_forecast_env import WeatherForecastEnv, basin_context_vector


LABELS = (
    "enso_oni",
    "iod_dmi",
    "itcz_lat",
    "mslp_hpa",
    "solar_wind_kms",
    "kp_index",
    "goes_xray_log10",
    "helio_regime_ord",
)


def _print_vec(name: str, v: np.ndarray) -> None:
    print(f"\n{name}  shape={v.shape}")
    for i, lab in enumerate(LABELS):
        print(f"  [{i}] {lab:18s} {float(v[i]): .6g}")


def _episode_with_basin(bc: BasinContext) -> EpisodeContext:
    obs = make_synthetic_zone_obs("karawang_rice", seed=42)
    fc = make_synthetic_forecast_result(
        zone_id="karawang_rice",
        valid_time=obs.valid_time,
        horizon_days=7,
        seed=42,
    )
    return EpisodeContext(
        obs=obs,
        forecast=fc,
        zone_ids=["karawang_rice"],
        basin_context=bc,
    )


def main() -> int:
    print("=== helio → observation path diagnostic ===")

    quiet = BasinContext(
        valid_date=datetime.now(timezone.utc),
        enso_oni=0.0,
        iod_dmi=0.0,
        solar_wind_speed_kms=400.0,
        kp_index=2.0,
        goes_xray_flux=1e-7,
        helio_regime="quiet",
        source=DataSource.SYNTHETIC,
    )
    storm = BasinContext(
        valid_date=datetime.now(timezone.utc),
        enso_oni=0.0,
        iod_dmi=0.0,
        solar_wind_speed_kms=650.0,
        kp_index=6.5,
        goes_xray_flux=2e-5,
        helio_regime="storm",
        source=DataSource.SYNTHETIC,
    )
    v_q = basin_context_vector(quiet)
    v_s = basin_context_vector(storm)
    v_n = basin_context_vector(None)
    assert v_q.shape == (8,), v_q.shape
    assert v_s.shape == (8,), v_s.shape
    assert abs(float(v_q[7]) - 0.0) < 1e-6
    assert abs(float(v_s[7]) - 2.0) < 1e-6
    assert abs(float(v_s[5]) - 6.5) < 1e-6
    delta = float(np.linalg.norm(v_s - v_q))
    print(f"quiet vs storm L2 delta: {delta:.4f}  (must be > 0)")
    assert delta > 1.0, "storm/quiet vectors should differ substantially"
    _print_vec("quiet", v_q)
    _print_vec("storm", v_s)
    _print_vec("neutral (None)", v_n)

    env = WeatherForecastEnv(ForecastConfig(n_zones=1, horizon_days=7, max_steps=3))
    assert env.observation_space["basin_context"].shape == (8,), (
        env.observation_space["basin_context"].shape
    )
    obs_q, _ = env.reset(options={"context": _episode_with_basin(quiet)})
    obs_s, _ = env.reset(options={"context": _episode_with_basin(storm)})
    bq = obs_q["basin_context"]
    bs = obs_s["basin_context"]
    assert bq.shape == (8,) and bs.shape == (8,)
    env_delta = float(np.linalg.norm(bs - bq))
    print(f"\nenv obs quiet vs storm L2 delta: {env_delta:.4f}")
    assert env_delta > 1.0
    _print_vec("env quiet basin_context", bq)
    _print_vec("env storm basin_context", bs)
    print("\nENV OBS PATH OK — helio channels present and responsive")

    try:
        from era5_data_pipeline import _fetch_swpc_helio, fetch_basin_context

        now = datetime.now(timezone.utc)
        helio = _fetch_swpc_helio(now)
        print("\nSWPC live keys:", sorted(helio.keys()))
        print(
            "SWPC sample:",
            {
                k: helio.get(k)
                for k in (
                    "solar_wind_speed_kms",
                    "kp_index",
                    "goes_xray_flux",
                    "helio_regime",
                )
            },
        )
        cfg = ForecastConfig(
            include_basin_context=True, require_real_basin_context=False
        )
        bc = fetch_basin_context(now, cfg)
        v_live = basin_context_vector(bc)
        _print_vec("live BasinContext → vector", v_live)
        obs_live, _ = env.reset(options={"context": _episode_with_basin(bc)})
        _print_vec("env live basin_context", obs_live["basin_context"])
        print("LIVE SWPC → ENV PATH OK")
    except Exception as e:
        print(
            f"\nLIVE SWPC skipped or failed (offline-safe): "
            f"{type(e).__name__}: {e}"
        )

    print(
        "\nNOTE: crop_risk_scorer does not consume helio. Score deltas under "
        "quiet vs storm with identical ZoneObs/ForecastResult must be 0.0. "
        "Benefit path is RL observation → policy (requires retrain on 8-dim "
        "basin_context)."
    )
    print("\nAll helio observation diagnostics passed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())