%load_ext autoreload
%autoreload 2

RWRAP-test

Test du modèle warp + residual sur HEALPix pour la prévision SST.

Ce notebook reprend la préparation de données du notebook U-NET existant, mais remplace le HealPixUNet par HealPixAdvectionResidualNet.

Idée du modèle:

• tête 1 : vitesse tangentielle est-ouest u

• tête 2 : vitesse tangentielle nord-sud v

• tête 3 : résidu local r

• incrément prédit : ΔT = -(u ∂T/∂east + v ∂T/∂north) + r

Ici on garde la même cible que dans ton notebook préparé :

Y[t] = SST_norm[t+1] - SST_norm[t]

Donc le réseau est instancié avec output_mode="delta".

import numpy as np
import matplotlib.pyplot as plt
import torch

from healpix_plot.mollview import mollview
import healpix_geo as hg

from healpix_analyse.rwrap_healpix import HealPixAdvectionResidualNet

SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")

sst_fields = np.load("../../Healpix_UNET/Notebooks/train_heal_32.npy")[:65*12, -1]
mollview(sst_fields[10], nest=True, cmap="coolwarm")

# Données océan uniquement
ocean_mask = np.isfinite(sst_fields[0]).astype(np.float32)
sst_ocean = np.where(np.isfinite(sst_fields), sst_fields, np.nan)

print("sst_ocean shape:", sst_ocean.shape)

# Standardisation mois par mois
mu = np.nanmean(sst_ocean.reshape(65, 12, sst_ocean.shape[1]), axis=0, keepdims=True)
std = np.nanstd(sst_ocean.reshape(65, 12, sst_ocean.shape[1]), axis=0, keepdims=True)

sst_norm = (sst_ocean.reshape(65, 12, sst_ocean.shape[1]) - mu) / std
sst_norm = sst_norm.reshape(65*12, sst_ocean.shape[1])
sst_norm[np.isnan(sst_norm)] = 0.0
sst_norm[:, ocean_mask == 0] = 0.0

print("norm shape:", sst_norm.shape)

# Input / target
# X[t] = [SST_norm(t), mask]
# Y[t] = SST_norm(t+1) - SST_norm(t)

X = np.stack([
    sst_norm,
    np.broadcast_to(ocean_mask, sst_norm.shape).copy(),
], axis=1).astype(np.float32)
X = X[:-1]

Y = (sst_norm[1:, None, :] - sst_norm[:-1, None, :]).astype(np.float32)

N_train = int(0.8 * Y.shape[0])
X_train, Y_train = X[:N_train], Y[:N_train]
X_val, Y_val = X[N_train:], Y[N_train:]

print(f"Train: {X_train.shape} / {Y_train.shape}")
print(f"Val  : {X_val.shape} / {Y_val.shape}")

NSIDE = 32
depth = 5
enso_idx, _, _ = hg.nested.cone_coverage((130, 0), 5.0, depth)

_MASK_T = torch.as_tensor(ocean_mask, dtype=torch.float32, device=DEVICE)

def loss_ocean_delta(y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
    """
    Loss sur l'incrément SST, océan uniquement, avec régularisation
    flow smoothness + petit résidu implicites dans model.total_loss.
    """
    mask = _MASK_T.to(y_pred.device)
    err = (y_pred - y_true) ** 2
    data_term = (err * mask).sum() / (mask.sum() * y_pred.shape[0] * y_pred.shape[1])

    # On ajoute les régularisations internes du modèle à l'extérieur
    # pour garder une logique explicite dans le notebook.
    return data_term + 0.05 * model.flow_smoothness_loss() + 0.01 * model.residual_l1_loss()

model = HealPixAdvectionResidualNet(
    nside=NSIDE,
    in_channels=2,
    out_channels=1,
    feature_channels=[16, 8, 4, 2],
    vars_per_t=1,
    time_steps=1,
    down_mode="maxpool",
    weight_norm="l1",
    residual_clip=4.0,
    gate_scale=1.0,
    tau=1.0,
    kernel_sz=5,
    n_gauges=1,
    gauge_type="two_ref",
    singularity_lonlat=[
        (-55.0, -10.0),
        ( 90.0,  45.0),
    ],
    ellipsoid="WGS84",
    grad_ring=1,
    smoothness_lambda=0.05,
    residual_lambda=0.01,
    output_mode="delta",
    device=DEVICE,
)

print(model)
print(f"\nModel parameters: {model.count_parameters():,}")

print("\n--- Training -------------------------------------------------------")
history = model.fit(
    X_train[-250:], Y_train[-250:],
    x_val=X_val, y_val=Y_val,
    n_epoch=300,
    batch_size=16,
    lr=3e-3,
    weight_decay=1e-5,
    optimizer="adam",
    view_epoch=1,
    loss_fn=loss_ocean_delta,
)

# Evaluation sur incréments
Y_ptrain = model.predict(X_train, batch_size=16).numpy()[:, 0, :]
Y_ptrue  = Y_train[:, 0, :]

Y_pred = model.predict(X_val, batch_size=16).numpy()[:, 0, :]
Y_true = Y_val[:, 0, :]

mask_np = ocean_mask.astype(bool)
mse_val = np.mean((Y_pred[:, mask_np] - Y_true[:, mask_np]) ** 2)
rmse_val = np.sqrt(mse_val)
corr_val = float(np.mean([
    np.corrcoef(Y_pred[i, mask_np], Y_true[i, mask_np])[0, 1]
    for i in range(Y_pred.shape[0])
]))

print(f"Validation RMSE: {rmse_val:.4f}")
print(f"Validation Corr: {corr_val:.4f}")

# Visualisation incrément prédit vs vrai
a = 2
plt.figure(figsize=(16, 4))
tmp = Y_ptrain[300] / np.where(ocean_mask > 0, ocean_mask, np.nan)
tmp[enso_idx] = a * 2
mollview(tmp, nest=True, hold=False, sub=(1, 2, 1), cmap="coolwarm", vmin=-a, vmax=a)

tmp = Y_ptrue[300] / np.where(ocean_mask > 0, ocean_mask, np.nan)
tmp[enso_idx] = a * 2
mollview(tmp, nest=True, hold=False, sub=(1, 2, 2), cmap="coolwarm", vmin=-a, vmax=a)

plt.figure(figsize=(16, 4))
tmp = Y_pred[70] / np.where(ocean_mask > 0, ocean_mask, np.nan)
tmp[enso_idx] = -a * 2
mollview(tmp, nest=True, hold=False, sub=(1, 2, 1), cmap="coolwarm", vmin=-a, vmax=a)

tmp = Y_true[70] / np.where(ocean_mask > 0, ocean_mask, np.nan)
tmp[enso_idx] = -a * 2
mollview(tmp, nest=True, hold=False, sub=(1, 2, 2), cmap="coolwarm", vmin=-a, vmax=a)

# Diagnostic des composantes du modèle
components = model.predict_components(X_val[:8], batch_size=8)
flow_e = components["flow_east"].numpy()[:, 0, :]
flow_n = components["flow_north"].numpy()[:, 0, :]
resid  = components["residual"].numpy()[:, 0, :]

idx = 0
A = 2
plt.figure(figsize=(16, 9))
mollview(flow_e[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 1), cmap="coolwarm")
plt.title("flow east")
mollview(flow_n[idx] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(2, 2, 2), cmap="coolwarm")
plt.title("flow north")
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)
plt.title("residual")
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)
plt.title("predicted delta")

# Reconstruction de SST(t+1) absolue à partir de ΔT prédit
last_sst_val = X_val[:, 0, :]
SST_next_pred = last_sst_val + Y_pred
SST_next_true = last_sst_val + Y_true

plt.figure(figsize=(16, 4))
mollview(SST_next_pred[10] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(1, 2, 1), cmap="coolwarm")
mollview(SST_next_true[10] / np.where(ocean_mask > 0, ocean_mask, np.nan), nest=True, hold=False, sub=(1, 2, 2), cmap="coolwarm")