src/continual_demo.py

"""
Demonstration of continual learning on successive scientific domains
using the hierarchical Poincaré model + Replay + EWC.

Uses the exact best hyperparameters from the long Optuna study:
  lr=3.82e-4, curvature=0.455, hidden=96, batch_size=8,
  pred_steps=4, w_phys=9.6e-4, levels=2
"""
from __future__ import annotations
import os, sys, copy
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from torch.utils.data import DataLoader
import numpy as np
from tqdm import tqdm

from src.normalization import FieldNormalizer
from src.synthetic_fields import SyntheticWellLike
from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor
from src.physics_losses import combined_physics_loss
from src.continual import ReplayBuffer, DiagonalEWC, hyperbolic_distillation_loss, fit_normalizer_for_domain
from src.config import BEST_HPARAMS as BEST

def collate(batch):
    return torch.stack([b["fields"] for b in batch])

def make_domain(seed: int, n_channels: int = 2, n_samples=96, n_steps=14):
    """Slightly different synthetic regimes act as successive scientific domains."""
    torch.manual_seed(seed)
    return SyntheticWellLike(n_samples=n_samples, n_steps=n_steps, height=32, width=32,
                              n_channels=n_channels, noise=0.12 + 0.04*(seed%3))

def evaluate(model, norm, ds, device, pred_steps=None):
    if pred_steps is None:
        pred_steps = model.pred_steps
    model.eval()
    loader = DataLoader(ds, batch_size=8, collate_fn=collate)
    losses = []
    with torch.no_grad():
        for batch in loader:
            B,T,C,H,W = batch.shape
            batch = batch.to(device)
            flat = norm.transform(batch.view(B*T,C,H,W)).view(B,T,C,H,W)
            win = 4
            if T < win + pred_steps:
                continue
            x = flat[:, :win]
            tgt = torch.stack([model.encode(flat[:, win+s]) for s in range(pred_steps)], 1)
            pred = model(x)
            losses.append(model.hyperbolic_loss(pred, tgt).item())
    model.train()
    return float(np.mean(losses)) if losses else 1e6

def train_domain(model, norm, ds, opt, device, epochs=5, replay: ReplayBuffer=None,
                 ewc: DiagonalEWC=None, teacher=None, mix_replay=0.4):
    loader = DataLoader(ds, batch_size=BEST["batch_size"], shuffle=True, collate_fn=collate)
    w_phys = BEST["w_phys"]
    ps = model.pred_steps
    probe = ds[0]["fields"]
    domain_c = int(probe.shape[1])
    domain_hw = (int(probe.shape[2]), int(probe.shape[3]))
    for ep in range(epochs):
        for batch in loader:
            B,T,C,H,W = batch.shape
            batch = batch.to(device)
            if replay is not None and len(replay) > 0 and np.random.rand() < mix_replay:
                old = replay.sample(max(1, B//2), channels=domain_c, spatial=domain_hw)
                if old is not None:
                    old = old.to(device)
                    tmin = min(old.size(1), T)
                    batch = torch.cat([batch[:, :tmin], old[:, :tmin]], dim=0)
                    B = batch.size(0)
                    T = tmin
            flat = norm.transform(batch.view(B*T if batch.dim()==5 else B*batch.size(1), C, H, W))
            if batch.dim() == 5:
                flat = flat.view(B, T, C, H, W)
            else:
                flat = flat.view(B, -1, C, H, W)
                T = flat.size(1)
            win = 4
            if T < win + ps:
                continue
            x = flat[:, :win]
            with torch.no_grad():
                tgt = torch.stack([model.encode(flat[:, win+s]) for s in range(ps)], 1)
            pred = model(x)
            loss = model.hyperbolic_loss(pred, tgt)
            loss = loss + combined_physics_loss(flat[:, :win+ps], w_smooth=w_phys, w_temp=w_phys)
            if ewc is not None:
                loss = loss + ewc.ewc_loss(model)
            if teacher is not None:
                with torch.no_grad():
                    t_lat = teacher.encode(x[:, -1] if x.dim()==5 else x)
                s_lat = model.encode(x[:, -1] if x.dim()==5 else x)
                loss = loss + 0.1 * hyperbolic_distillation_loss(s_lat, t_lat, model.poincare)
            opt.zero_grad()
            loss.backward()
            if ewc is not None and np.random.rand() < 0.25:
                ewc.accumulate_fisher(model, loss)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        if replay is not None:
            for i in range(min(8, len(ds))):
                replay.add(ds[i]["fields"])

def main():
    device = "cpu"
    print("=" * 64)
    print("Continual learning demo – Replay + EWC on Poincaré hierarchical model")
    print("Using Optuna best HPs:", BEST)
    print("=" * 64)

    domain_specs = [
        {"seed": 11, "n_channels": 2},
        {"seed": 22, "n_channels": 5},
        {"seed": 33, "n_channels": 2},
    ]
    domains = [make_domain(seed=s["seed"], n_channels=s["n_channels"]) for s in domain_specs]

    enc = MultiScaleEncoder(hidden=BEST["hidden"], out_dim=8)
    model = HierarchicalHyperbolicPredictor(
        enc, c=BEST["curvature"], pred_steps=BEST["pred_steps"], levels=BEST["levels"]
    ).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=BEST["lr"])

    replay = ReplayBuffer(capacity=128)
    ewc = DiagonalEWC(model, lambda_ewc=500.0)
    teacher = None
    normalizers = []

    history = {f"domain_{i}": [] for i in range(3)}

    for d_idx, ds in enumerate(domains):
        c = domain_specs[d_idx]["n_channels"]
        print(f"\n--- Training domain {d_idx+1}/3 (C={c}) ---")
        norm = fit_normalizer_for_domain(ds, max_fit=40)
        normalizers.append(norm)

        train_domain(model, norm, ds, opt, device, epochs=4,
                     replay=replay,
                     ewc=ewc if d_idx > 0 else None,
                     teacher=teacher, mix_replay=0.0 if d_idx == 0 else 0.35)
        print(f"  replay buffer by channel count: {replay.counts_by_channels()}")
        ewc.finalize_domain(model)
        teacher = copy.deepcopy(model).eval()
        for p in teacher.parameters():
            p.requires_grad = False

        for j in range(d_idx + 1):
            loss_j = evaluate(model, normalizers[j], domains[j], device)  # derives pred_steps from model itself
            history[f"domain_{j}"].append(loss_j)
            print(f"  Eval domain {j+1} (C={domain_specs[j]['n_channels']}) loss: {loss_j:.4f}")

    print("\n" + "=" * 64)
    print("Retention summary (loss after each successive domain)")
    for k, v in history.items():
        print(f"  {k}: {[round(x,4) for x in v]}")

    os.makedirs("logs", exist_ok=True)
    path = "logs/poincare8d_continual.pt"
    torch.save({
        "model": model.state_dict(),
        "params": BEST,
        "normalizers": [n.state_dict() for n in normalizers],
        "domain_specs": domain_specs,
        "history": history,
    }, path)
    print(f"\nSaved {path}")
    print("Continual learning module integrated and demonstrated (cross-channel-count).")

if __name__ == "__main__":
    main()