{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3a2a7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b70b1452",
   "metadata": {},
   "source": [
    "# RWRAP-test\n",
    "\n",
    "Test du modèle **warp + residual** sur HEALPix pour la prévision SST.\n",
    "\n",
    "Ce notebook reprend la préparation de données du notebook U-NET existant, mais remplace le `HealPixUNet` par `HealPixAdvectionResidualNet`.\n",
    "\n",
    "Idée du modèle:\n",
    "- tête 1 : vitesse tangentielle est-ouest `u`\n",
    "- tête 2 : vitesse tangentielle nord-sud `v`\n",
    "- tête 3 : résidu local `r`\n",
    "- incrément prédit : `ΔT = -(u ∂T/∂east + v ∂T/∂north) + r`\n",
    "\n",
    "Ici on garde la même cible que dans ton notebook préparé :\n",
    "\n",
    "\n",
    "`Y[t] = SST_norm[t+1] - SST_norm[t]`\n",
    "\n",
    "Donc le réseau est instancié avec `output_mode=\"delta\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d2db56a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "from healpix_plot.mollview import mollview\n",
    "import healpix_geo as hg\n",
    "\n",
    "from healpix_analyse.rwrap_healpix import HealPixAdvectionResidualNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad2697a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "torch.manual_seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Device: {DEVICE}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c5705f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst_fields = np.load(\"../../Healpix_UNET/Notebooks/train_heal_32.npy\")[:65*12, -1]\n",
    "mollview(sst_fields[10], nest=True, cmap=\"coolwarm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33552443",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Données océan uniquement\n",
    "ocean_mask = np.isfinite(sst_fields[0]).astype(np.float32)\n",
    "sst_ocean = np.where(np.isfinite(sst_fields), sst_fields, np.nan)\n",
    "\n",
    "print(\"sst_ocean shape:\", sst_ocean.shape)\n",
    "\n",
    "# Standardisation mois par mois\n",
    "mu = np.nanmean(sst_ocean.reshape(65, 12, sst_ocean.shape[1]), axis=0, keepdims=True)\n",
    "std = np.nanstd(sst_ocean.reshape(65, 12, sst_ocean.shape[1]), axis=0, keepdims=True)\n",
    "\n",
    "sst_norm = (sst_ocean.reshape(65, 12, sst_ocean.shape[1]) - mu) / std\n",
    "sst_norm = sst_norm.reshape(65*12, sst_ocean.shape[1])\n",
    "sst_norm[np.isnan(sst_norm)] = 0.0\n",
    "sst_norm[:, ocean_mask == 0] = 0.0\n",
    "\n",
    "print(\"norm shape:\", sst_norm.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c4a9a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input / target\n",
    "# X[t] = [SST_norm(t), mask]\n",
    "# Y[t] = SST_norm(t+1) - SST_norm(t)\n",
    "\n",
    "X = np.stack([\n",
    "    sst_norm,\n",
    "    np.broadcast_to(ocean_mask, sst_norm.shape).copy(),\n",
    "], axis=1).astype(np.float32)\n",
    "X = X[:-1]\n",
    "\n",
    "Y = (sst_norm[1:, None, :] - sst_norm[:-1, None, :]).astype(np.float32)\n",
    "\n",
    "N_train = int(0.8 * Y.shape[0])\n",
    "X_train, Y_train = X[:N_train], Y[:N_train]\n",
    "X_val, Y_val = X[N_train:], Y[N_train:]\n",
    "\n",
    "print(f\"Train: {X_train.shape} / {Y_train.shape}\")\n",
    "print(f\"Val  : {X_val.shape} / {Y_val.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb79df1",
   "metadata": {},
   "outputs": [],
   "source": [
    "NSIDE = 32\n",
    "depth = 5\n",
    "enso_idx, _, _ = hg.nested.cone_coverage((130, 0), 5.0, depth)\n",
    "\n",
    "_MASK_T = torch.as_tensor(ocean_mask, dtype=torch.float32, device=DEVICE)\n",
    "\n",
    "def loss_ocean_delta(y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Loss sur l'incrément SST, océan uniquement, avec régularisation\n",
    "    flow smoothness + petit résidu implicites dans model.total_loss.\n",
    "    \"\"\"\n",
    "    mask = _MASK_T.to(y_pred.device)\n",
    "    err = (y_pred - y_true) ** 2\n",
    "    data_term = (err * mask).sum() / (mask.sum() * y_pred.shape[0] * y_pred.shape[1])\n",
    "\n",
    "    # On ajoute les régularisations internes du modèle à l'extérieur\n",
    "    # pour garder une logique explicite dans le notebook.\n",
    "    return data_term + 0.05 * model.flow_smoothness_loss() + 0.01 * model.residual_l1_loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e151b6ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = HealPixAdvectionResidualNet(\n",
    "    nside=NSIDE,\n",
    "    in_channels=2,\n",
    "    out_channels=1,\n",
    "    feature_channels=[16, 8, 4, 2],\n",
    "    vars_per_t=1,\n",
    "    time_steps=1,\n",
    "    down_mode=\"maxpool\",\n",
    "    weight_norm=\"l1\",\n",
    "    residual_clip=4.0,\n",
    "    gate_scale=1.0,\n",
    "    tau=1.0,\n",
    "    kernel_sz=5,\n",
    "    n_gauges=1,\n",
    "    gauge_type=\"two_ref\",\n",
    "    singularity_lonlat=[\n",
    "        (-55.0, -10.0),\n",
    "        ( 90.0,  45.0),\n",
    "    ],\n",
    "    ellipsoid=\"WGS84\",\n",
    "    grad_ring=1,\n",
    "    smoothness_lambda=0.05,\n",
    "    residual_lambda=0.01,\n",
    "    output_mode=\"delta\",\n",
    "    device=DEVICE,\n",
    ")\n",
    "\n",
    "print(model)\n",
    "print(f\"\\nModel parameters: {model.count_parameters():,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff30752a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n--- Training -------------------------------------------------------\")\n",
    "history = model.fit(\n",
    "    X_train[-250:], Y_train[-250:],\n",
    "    x_val=X_val, y_val=Y_val,\n",
    "    n_epoch=300,\n",
    "    batch_size=16,\n",
    "    lr=3e-3,\n",
    "    weight_decay=1e-5,\n",
    "    optimizer=\"adam\",\n",
    "    view_epoch=1,\n",
    "    loss_fn=loss_ocean_delta,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8598d82",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation sur incréments\n",
    "Y_ptrain = model.predict(X_train, batch_size=16).numpy()[:, 0, :]\n",
    "Y_ptrue  = Y_train[:, 0, :]\n",
    "\n",
    "Y_pred = model.predict(X_val, batch_size=16).numpy()[:, 0, :]\n",
    "Y_true = Y_val[:, 0, :]\n",
    "\n",
    "mask_np = ocean_mask.astype(bool)\n",
    "mse_val = np.mean((Y_pred[:, mask_np] - Y_true[:, mask_np]) ** 2)\n",
    "rmse_val = np.sqrt(mse_val)\n",
    "corr_val = float(np.mean([\n",
    "    np.corrcoef(Y_pred[i, mask_np], Y_true[i, mask_np])[0, 1]\n",
    "    for i in range(Y_pred.shape[0])\n",
    "]))\n",
    "\n",
    "print(f\"Validation RMSE: {rmse_val:.4f}\")\n",
    "print(f\"Validation Corr: {corr_val:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0715fb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualisation incrément prédit vs vrai\n",
    "a = 2\n",
    "plt.figure(figsize=(16, 4))\n",
    "tmp = Y_ptrain[300] / np.where(ocean_mask > 0, ocean_mask, np.nan)\n",
    "tmp[enso_idx] = a * 2\n",
    "mollview(tmp, nest=True, hold=False, sub=(1, 2, 1), cmap=\"coolwarm\", vmin=-a, vmax=a)\n",
    "\n",
    "tmp = Y_ptrue[300] / np.where(ocean_mask > 0, ocean_mask, np.nan)\n",
    "tmp[enso_idx] = a * 2\n",
    "mollview(tmp, nest=True, hold=False, sub=(1, 2, 2), cmap=\"coolwarm\", vmin=-a, vmax=a)\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "tmp = Y_pred[70] / np.where(ocean_mask > 0, ocean_mask, np.nan)\n",
    "tmp[enso_idx] = -a * 2\n",
    "mollview(tmp, nest=True, hold=False, sub=(1, 2, 1), cmap=\"coolwarm\", vmin=-a, vmax=a)\n",
    "\n",
    "tmp = Y_true[70] / np.where(ocean_mask > 0, ocean_mask, np.nan)\n",
    "tmp[enso_idx] = -a * 2\n",
    "mollview(tmp, nest=True, hold=False, sub=(1, 2, 2), cmap=\"coolwarm\", vmin=-a, vmax=a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b4af5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Diagnostic des composantes du modèle\n",
    "components = model.predict_components(X_val[:8], batch_size=8)\n",
    "flow_e = components[\"flow_east\"].numpy()[:, 0, :]\n",
    "flow_n = components[\"flow_north\"].numpy()[:, 0, :]\n",
    "resid  = components[\"residual\"].numpy()[:, 0, :]\n",
    "\n",
    "idx = 0\n",
    "A = 2\n",
    "plt.figure(figsize=(16, 9))\n",
    "mollview(flow_e[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 1), cmap=\"coolwarm\")\n",
    "plt.title(\"flow east\")\n",
    "mollview(flow_n[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 2), cmap=\"coolwarm\")\n",
    "plt.title(\"flow north\")\n",
    "mollview(resid[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 3), cmap=\"coolwarm\", vmin=-A, vmax=A)\n",
    "plt.title(\"residual\")\n",
    "mollview(Y_pred[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 4), cmap=\"coolwarm\", vmin=-A, vmax=A)\n",
    "plt.title(\"predicted delta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20d16f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconstruction de SST(t+1) absolue à partir de ΔT prédit\n",
    "last_sst_val = X_val[:, 0, :]\n",
    "SST_next_pred = last_sst_val + Y_pred\n",
    "SST_next_true = last_sst_val + Y_true\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "mollview(SST_next_pred[10] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(1, 2, 1), cmap=\"coolwarm\")\n",
    "mollview(SST_next_true[10] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(1, 2, 2), cmap=\"coolwarm\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}