scripts/train_2026_photosynthesis.py

"""Experimental Stage-2 photosynthesis ML training on the 2026-schema parquet.

This is an exploratory smoke test on ~10 days of data — **not** a production
model. The goal is to (a) validate the training pipeline end-to-end, (b) see
whether the 2026-fleet sensor signals (NDVI, PRI, PSRI, soil moisture, leaf
temp) carry information beyond IMS weather, and (c) establish a baseline that
we can revisit once enough growing-season data accumulates (~Aug 2026).

Two feature sets are trained side by side:

  Model A — forecast-grade
      Features: IMS weather only (ghi, air_temp, rh, wind)
      Use case: day-ahead forecasting (no on-site data needed at inference)

  Model B — full
      Features: Model A + 2026 canopy state (leaf_temp, NDVI, PRI, PSRI,
      soil_moisture_shallow, soil_temp_shallow)
      Use case: nowcasting / control-loop input

Excluded from both:
  par_umol_derived, vpd_kpa_derived, cwsi_proxy — these feed the Farquhar
  equation that generates the label, so including them reduces ML to
  reverse-engineering. The whole point of Stage-2 ML is to predict A from
  signals the Farquhar inputs are derived from.

Honest expectation: with 964 daytime rows and a 2-day hold-out, both models
will likely look very good on this hold-out (the test conditions are close
to training) and the contrast between A and B may not be statistically
meaningful. Treat numbers as directional.

Usage::

    python -m scripts.train_2026_photosynthesis
    python -m scripts.train_2026_photosynthesis --holdout-days 2
"""

from __future__ import annotations

import argparse
import pickle
import sys
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

_PARQUET = _PROJECT_ROOT / "Data" / "2026" / "sensor_history.parquet"
_OUT_PKL = _PROJECT_ROOT / "Data" / "2026" / "photosynthesis_model.pkl"

_TARGET = "a_farquhar_umol"
_FEATURES_A = ["ghi_w_m2", "air_temperature_c", "rh_percent", "wind_speed_ms"]
_FEATURES_B = _FEATURES_A + [
    "leaf_temperature", "ndvi", "pri", "psri",
    "soil_moisture_shallow_pct", "soil_temp_shallow_c",
]


def _load_daytime() -> pd.DataFrame:
    if not _PARQUET.exists():
        raise FileNotFoundError(
            f"{_PARQUET} not found. Run scripts.collect_2026_training_data first."
        )
    df = pd.read_parquet(_PARQUET)
    # Daytime only — A is 0 at night by construction, includes those would
    # inflate R² without testing model skill.
    df = df[df["ghi_w_m2"] > 50].copy()
    df = df.dropna(subset=_FEATURES_B + [_TARGET])
    df = df.sort_index()
    return df


def _train_and_eval(name: str, X_tr, y_tr, X_te, y_te, model):
    from sklearn.metrics import mean_absolute_error, r2_score
    model.fit(X_tr, y_tr)
    pred_tr = model.predict(X_tr)
    pred_te = model.predict(X_te)
    mae_tr = mean_absolute_error(y_tr, pred_tr)
    mae_te = mean_absolute_error(y_te, pred_te)
    r2_tr = r2_score(y_tr, pred_tr)
    r2_te = r2_score(y_te, pred_te)
    print(f"  {name:<22}  MAE  train={mae_tr:.3f}  test={mae_te:.3f}   "
          f"R²  train={r2_tr:.3f}  test={r2_te:.3f}")
    return model, mae_te, r2_te


def main() -> None:
    p = argparse.ArgumentParser(description="Train an exploratory 2026 PS model.")
    p.add_argument("--holdout-days", type=int, default=2)
    args = p.parse_args()

    from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
    from sklearn.linear_model import LinearRegression

    df = _load_daytime()
    print(f"Loaded {len(df):,} daytime rows  "
          f"{df.index.min()} → {df.index.max()}")

    # Temporal split — last N days as hold-out
    split = df.index.max() - pd.Timedelta(days=args.holdout_days)
    train = df[df.index < split]
    test = df[df.index >= split]
    print(f"Train: {len(train):,} rows   Hold-out: {len(test):,} rows  "
          f"(split at {split})")
    if len(test) < 50:
        raise RuntimeError("Hold-out window too small — collect more days first.")

    y_tr = train[_TARGET].to_numpy()
    y_te = test[_TARGET].to_numpy()

    naive_mean = float(y_tr.mean())
    naive_mae = float(np.abs(y_te - naive_mean).mean())
    print(f"\nNaive 'predict training-mean' baseline: MAE = {naive_mae:.3f}  "
          f"(model has to beat this to be useful)")

    results: dict[str, dict] = {}

    for set_name, feats in [("A_ims_only", _FEATURES_A), ("B_ims+canopy", _FEATURES_B)]:
        print(f"\n=== Feature set {set_name}  ({len(feats)} features) ===")
        print(f"  features: {feats}")
        X_tr = train[feats].to_numpy()
        X_te = test[feats].to_numpy()

        lr, lr_mae, lr_r2 = _train_and_eval("LinearRegression", X_tr, y_tr, X_te, y_te,
                                            LinearRegression())
        rf, rf_mae, rf_r2 = _train_and_eval("RandomForest", X_tr, y_tr, X_te, y_te,
                                            RandomForestRegressor(
                                                n_estimators=300,
                                                max_depth=8,
                                                min_samples_leaf=4,
                                                random_state=42,
                                                n_jobs=-1,
                                            ))
        gbr, gbr_mae, gbr_r2 = _train_and_eval("GradientBoosting", X_tr, y_tr, X_te, y_te,
                                               GradientBoostingRegressor(
                                                   n_estimators=300,
                                                   max_depth=4,
                                                   learning_rate=0.05,
                                                   random_state=42,
                                               ))

        # Feature importances from RF (most interpretable)
        imp = pd.Series(rf.feature_importances_, index=feats).sort_values(ascending=False)
        print("  RF feature importances:")
        for k, v in imp.items():
            print(f"    {k:<28}  {v:.3f}")

        results[set_name] = {
            "features": feats,
            "lr":  {"model": lr,  "mae": lr_mae,  "r2": lr_r2},
            "rf":  {"model": rf,  "mae": rf_mae,  "r2": rf_r2},
            "gbr": {"model": gbr, "mae": gbr_mae, "r2": gbr_r2},
            "feature_importances": imp.to_dict(),
        }

    # ---------- Summary ----------
    print("\n=== Summary (hold-out MAE / R²) ===")
    print(f"  Naive baseline                 MAE = {naive_mae:.3f}")
    for set_name, r in results.items():
        for mname in ("lr", "rf", "gbr"):
            d = r[mname]
            print(f"  {set_name:<14}  {mname:<4}     MAE = {d['mae']:.3f}    R² = {d['r2']:.3f}")

    # Save the best Model B variant for downstream consumption.
    best_b = min(("lr", "rf", "gbr"), key=lambda m: results["B_ims+canopy"][m]["mae"])
    bundle = {
        "model_a_features": _FEATURES_A,
        "model_a_rf":       results["A_ims_only"]["rf"]["model"],
        "model_b_features": _FEATURES_B,
        "model_b_best":     results["B_ims+canopy"][best_b]["model"],
        "model_b_best_name": best_b,
        "test_mae_b":       results["B_ims+canopy"][best_b]["mae"],
        "test_r2_b":        results["B_ims+canopy"][best_b]["r2"],
        "test_mae_a":       results["A_ims_only"]["rf"]["mae"],
        "naive_mae":        naive_mae,
        "train_rows":       int(len(train)),
        "test_rows":        int(len(test)),
        "train_range":      (str(train.index.min()), str(train.index.max())),
        "test_range":       (str(test.index.min()), str(test.index.max())),
        "trained_at":       datetime.now(tz=timezone.utc).isoformat(),
        "schema":           "2026",
        "source":           str(_PARQUET.relative_to(_PROJECT_ROOT)),
        "feature_importances_b": results["B_ims+canopy"]["rf"]["feature_importances"],
    }
    _OUT_PKL.parent.mkdir(parents=True, exist_ok=True)
    with open(_OUT_PKL, "wb") as f:
        pickle.dump(bundle, f)
    print(f"\nSaved bundle → {_OUT_PKL}")


if __name__ == "__main__":
    main()