LZ96_CONV.py

# borrowed @ francois-rozet

# import h5py
import json
import torch
import math
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle
import random
import seaborn
import time
import wandb

from dawgz import job, schedule
from pathlib import Path
from tqdm import trange
from typing import *

from zuko.distributions import DiagNormal
from zuko.flows import Unconditional
from dasbi.inference.models import ConvNPE as NPE
from dasbi.networks.embedding import EmbedObs
from dasbi.simulators.sim_lorenz96 import LZ96 as sim


SCRATCH = os.environ.get("SCRATCH", ".")
PATH = Path(SCRATCH) / "npe_conv/lz96"
PATH.mkdir(parents=True, exist_ok=True)

fact = 5
N_grid = [2**i for i in range(3,9)]
Y_grid = [int(np.ceil(x/4)) for x in N_grid]
lN = len(N_grid)
window = 1
max_epochs = 2048
nms_dict = {
    8: 2,
    16: 2,
    32: 2,
    64: 3,
    128: 3,
    256: 4,
    512: 4,
}

ms_mod = {
    8: 1,
    16: 1,
    32: 1,
    64: 2,
    128: 3,
    256: 4,
    512: 4,
}
CONFIG = {
    # Architecture
    "embedding": [3]*lN,
    "kernel_size": [2]*lN,
    "ms_modules": [ms_mod[k] for k in N_grid],
    "num_conv": [2]*lN,
    "N_ms": [nms_dict[k] for k in N_grid],
    # Training
    # "epochs": [512]*lN,
    "batch_size": [128]*lN,
    "step_per_batch": [512]*lN,
    "optimizer": ["AdamW"]*lN,
    "learning_rate": [1e-3]*lN,  # np.geomspace(1e-3, 1e-4).tolist(),
    "weight_decay": [1e-4]*lN,  # np.geomspace(1e-2, 1e-4).tolist(),
    "scheduler": ["linear"]*lN,  # , 'cosine', 'exponential'],
    # Data
    "points": N_grid,
    "noise": [0.5]*lN,
    "train_sim": [2**10]*lN,
    "val_sim": [2**8]*lN,
    "device": ['cuda']*lN,
    # Test with assimilation window
    "x_dim": [(1, 1, sp, 1) for sp in N_grid],
    "y_dim": [(1, window, spy, 1) for spy in Y_grid],
    "y_dim_emb": [(1, 11, sp, 1) for sp in N_grid],
    'obs_mask': [True]*lN, #+1 in y_dim
    'ar': [False]*lN, #+1 in y_dim_emb (for modargs not embnet)
    'roll':[True]*lN,
    "observer_fp": [f"experiments/observer{N}LZ.pickle" for N in N_grid],
}


def build(**config):
    mod_args = {
        "x_dim": torch.tensor(config["x_dim"]),
        "y_dim": torch.tensor(config["y_dim_emb"]),
        "n_modules": config["ms_modules"],
        "n_c": config["num_conv"],
        "k_sz": torch.tensor((config["kernel_size"], 1)),
        "type": "1D",
    }

    N = config["points"]
    base = Unconditional(
        DiagNormal,
        torch.zeros(N),
        torch.ones(N),
        buffer=True,
    )

    mask = None
    if config['obs_mask']:
        with open(config["observer_fp"], "rb") as handle:
            observer = pickle.load(handle)
        mask = observer.get_mask().to(config['device'])

    emb_out = torch.tensor(config["y_dim_emb"]) 
    if config['ar']:
        emb_out[1] -= 1

    emb_net = EmbedObs(
        torch.tensor(config["y_dim"]),
        emb_out,
        conv_lay=config["embedding"],
        observer_mask=mask
    )
    return NPE(config["N_ms"], base, emb_net, mod_args, roll = config["roll"], ar=config["ar"])


def process_sim(simulator):
    MUX = simulator.data.mean(dim=(0, 1))
    SIGMAX = simulator.data.std(dim=(0, 1))

    MUY = simulator.obs.mean(dim=(0, 1))
    SIGMAY = simulator.obs.std(dim=(0, 1))

    MUT = simulator.time.mean(dim=(0, 1))
    SIGMAT = simulator.time.std(dim=(0, 1))

    simulator.data = (simulator.data - MUX) / SIGMAX
    simulator.obs = (simulator.obs - MUY) / SIGMAY
    simulator.time = (simulator.time - MUT) / SIGMAT

    ret_ls = [MUX, SIGMAX, MUY, SIGMAY, MUT, SIGMAT]

    return ret_ls

@job(array=fact*lN, cpus=2, gpus=1, ram="32GB", time="3-12:00:00")
def CONV_train(i: int):
    # config = {key: random.choice(values) for key, values in CONFIG.items()}
    config = {key : values[i%lN] for key,values in CONFIG.items()}

    with open(config["observer_fp"], "rb") as handle:
        observer = pickle.load(handle)

    gr = 'step' if window == 1 else 'assim'
    run = wandb.init(project="dasbi", config=config, group=f"LZ96_scaling_{gr}")
    runpath = PATH / f"runs/{run.name}_{run.id}"
    runpath.mkdir(parents=True, exist_ok=True)

    # with open(runpath / "config.json", "w") as f:
    #     json.dump(config, f)

    # Data
    tmax = 50
    traj_len = 1024 
    times = torch.linspace(0, tmax, traj_len)

    simt = sim(N=config["points"], noise=config["noise"])
    simt.init_observer(observer)
    simt.generate_steps(torch.randn((config["train_sim"], config["points"])), times)
    process_sim(simt)

    simv = sim(N=config["points"], noise=config["noise"])
    simv.init_observer(observer)
    simv.generate_steps(torch.randn((config["val_sim"], config["points"])), times)
    process_sim(simv)

    # Network
    conv_npe = build(**config).cuda()
    size = sum(param.numel() for param in conv_npe.parameters())
    run.config.num_param = size

    # Training
    # epochs = config["epochs"]
    batch_size = config["batch_size"]
    step_per_batch = config["step_per_batch"]
    best = 1000
    prev_loss = best
    time_buff = 256
    count = 0
    ## Optimizer
    if config["optimizer"] == "AdamW":
        optimizer = torch.optim.AdamW(
            conv_npe.parameters(),
            lr=config["learning_rate"],
            weight_decay=config["weight_decay"],
        )
    else:
        raise ValueError()

    if config["scheduler"] == "linear":
        lr = lambda t: 1 - (t / (1.5*max_epochs))
    # elif config["scheduler"] == "cosine":
    #     lr = lambda t: (1 + math.cos(math.pi * t / epochs)) / 2
    # elif config["scheduler"] == "exponential":
    #     lr = lambda t: math.exp(-7 * (t / epochs) ** 2)
    else:
        raise ValueError()

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr)

    ## Loop
    # for epoch in trange(epochs, ncols=88):
    epoch = 0
    while True:
        losses_train = []
        losses_val = []

        ### Train
        i = np.random.choice(
            len(simt.data),
            size=batch_size,
            replace=False
        )

        start = time.time()

        for xb, yb, tb in zip(
            simt.data[i].cuda(), simt.obs[i].cuda(), simt.time[i].cuda()
        ):
            subset_data = np.random.choice(
                np.arange(window - 1, traj_len),#because window of 10
                size=step_per_batch,
                replace=False,
            )

            x, y, t = (
                xb[subset_data],
                torch.cat([yb[i - window + 1 : i + 1].unsqueeze(0) for i in subset_data], dim=0),
                tb[subset_data],
            )
            x = x[:, None, ..., None]
            y = y[..., None]
            
            optimizer.zero_grad()
            l = conv_npe.loss(x, y, t)
            l.backward()
            norm = torch.nn.utils.clip_grad_norm_(conv_npe.embed.parameters(), 1)
            if torch.isfinite(norm):
                optimizer.step()
            
            losses_train.append(l.detach())

        end = time.time()

        ### Valid
        i = np.random.choice(
            len(simv.data),
            size=batch_size//2,
            replace=False,
        )

        with torch.no_grad():
            for xb, yb, tb in zip(
                simv.data[i].cuda(), simv.obs[i].cuda(), simv.time[i].cuda()
            ):
                subset_data = np.random.choice(
                    np.arange(window - 1, traj_len),
                    size=step_per_batch,
                    replace=False,
                )

                x, y, t = (
                    xb[subset_data],
                    torch.cat([yb[i - window + 1 : i + 1].unsqueeze(0) for i in subset_data], dim = 0),
                    tb[subset_data],
                )
                x = x[:, None, ..., None]
                y = y[..., None]
                
                losses_val.append(conv_npe.loss(x, y, t))

        ### Logs
        loss_train = torch.stack(losses_train).mean().item()
        loss_val = torch.stack(losses_val).mean().item()
        
        run.log(
            {
                "loss": loss_train,
                "loss_val": loss_val,
                "time_epoch": (end - start),
                "lr": optimizer.param_groups[0]["lr"],
                "plateau_buffer": count,
                "epoch": epoch
            }
        )

        ### Checkpoint
        if (prev_loss - loss_val) > 5e-2:
            prev_loss = loss_val
            torch.save(
                conv_npe.state_dict(),
                runpath / f"checkpoint.pth",
            )
            count = 0
        else:
            count += 1

        epoch += 1

        if count == time_buff or epoch == max_epochs:
            break

        scheduler.step()

    run.finish()


if __name__ == "__main__":
    schedule(
        CONV_train,
        name="Scaling LZ",
        backend="slurm",
        settings={"export": "ALL"},
        env=[
            "conda activate DASBI",
            "export WANDB_SILENT=true",
        ],
    )