kit/export_fme_latents.py

"""Record an fme (ACE / Samudra) run's latents as an xaig latent archive.

Runs ``fme.ace.inference`` (a YAML with ``initial_condition`` + ``forcing_loader``) or
``fme.ace.inference.evaluator`` (a YAML with ``loader``) with forward hooks on the
network inside the stepper.

- SFNO (atmosphere): layer 0 is the input to the first block (encoder output plus
  position embedding); layer L (1..num_layers) is the output of block L.
- Samudra U-Net (ocean): only the full-resolution levels, which xaig can map --
  layer 0 is the first encoder block's output, layer 1 the last decoder block's.

Optionally *writes* a layer as well (steering): add ``alpha`` times a direction to
layer L at every step (or a window of steps). The direction is one channel, or one
feature of an xaig basis file (a Dictionary's unit direction scaled by its
``input_scale``, so ``alpha`` is in activation units; a PCA component in layer units).

Network inputs are kept too, denormalized, exactly as the network saw them (so an
intervention on CO2 or a frozen flux shows up in ``reference.nc``).

Needs the model's environment (torch + fme) and xaig's archive writer, which is
pure numpy: point ``--xaig-src`` at a checkout's ``src`` when xaig is not installed.

    python export_fme_latents.py atm-ctrl.yaml --out archives/atm-E05S01-ctrl --steps 20

The archive's ``experiment`` block records the checkpoint, the source config, the
seed, the intervention (anything under ``--experiment``), the steering, and the
recorded steps.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import logging
import shutil
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

DEFAULT_FIELDS = [
    "PS", "TS", "Tat2m", "Qat2m", "Uat10m", "Vat10m",
    "surface_precipitation_rate", "FLUT", "FSUTOA", "FLUS", "FLDS", "FSNS",
    "LHFLX", "SHFLX", "TAUX", "TAUY", "lwp", "lcc",
    "U_1", "V_1", "T_1", "U_3", "V_3", "U_5", "V_5", "T_5", "STW_5", "U_7", "V_7", "STW_7",
]  # fmt: skip


DEFAULT_INPUTS = [
    "SOLIN", "global_mean_co2", "aerindexall", "colccn.3", "OCNFRAC", "ICEFRAC", "LANDFRAC",
    "PHIS", "TS", "FLUS", "FLDS", "FSNS", "LHFLX", "SHFLX", "TAUX", "TAUY", "sea_surface_fraction",
]  # fmt: skip
OCEAN_FIELDS = [
    "sst", "ssh", "ocean_sea_ice_fraction", "iceVolumeTotal", "temperatureCoarsened_0",
    "temperatureCoarsened_5", "temperatureCoarsened_10", "salinityCoarsened_0",
    "velocityZonalCoarsened_0", "velocityMeridionalCoarsened_0",
]  # fmt: skip


def find_step_with_packer(stepper):
    """The step object that packs named inputs into the network's channels."""
    candidates = [stepper, getattr(stepper, "_step_obj", None)]
    for obj in list(candidates):
        if obj is not None:
            candidates.extend(v for v in vars(obj).values() if hasattr(v, "in_packer"))
    for obj in candidates:
        if obj is not None and hasattr(obj, "in_packer") and hasattr(obj, "_normalizer"):
            return obj
    raise RuntimeError("could not find the step that packs the network inputs")


class SkipRescale:
    """A pre-hook on an SFNO's decoder that rescales one input in the big-skip half only.

    The decoder reads ``cat(block output, normalized inputs)``; this changes channel
    ``embed_dim + index(name)`` from ``(x - mean) / std`` to ``(factor * x - mean) / std``,
    so the skip sees the input scaled while every block saw it as it was."""

    def __init__(self, stepper, net, name, factor):
        if not getattr(net, "big_skip", False):
            raise SystemExit("--skip-rescale needs an SFNO with a big skip")
        step = find_step_with_packer(stepper)
        packed = list(step.in_packer.names)
        if name not in packed:
            raise SystemExit(f"{name} is not a network input; inputs are {packed}")
        self.channel = net.embed_dim + packed.index(name)
        self.mean = float(step._normalizer.means[name])
        self.std = float(step._normalizer.stds[name])
        self.factor = factor
        net.decoder.register_forward_pre_hook(self.hook)

    def hook(self, module, args):
        x = args[0].clone()
        physical = x[:, self.channel] * self.std + self.mean
        x[:, self.channel] = (self.factor * physical - self.mean) / self.std
        return (x, *args[1:])


class InputCapture:
    """A pre-hook on the network that keeps chosen inputs, denormalized, at recorded steps."""

    def __init__(self, stepper, net_outer, names, recorder, n_field_times):
        step = find_step_with_packer(stepper)
        packed = list(step.in_packer.names)
        missing = [n for n in names if n not in packed]
        if missing:
            logging.warning("inputs not among the network's inputs, skipped: %s", missing)
        self.names = [n for n in names if n in packed]
        self.index = {n: packed.index(n) for n in self.names}
        self.mean = {n: float(step._normalizer.means[n]) for n in self.names}
        self.std = {n: float(step._normalizer.stds[n]) for n in self.names}
        self.rec = recorder
        self.n_field_times = n_field_times
        self.values = {}
        net_outer.register_forward_pre_hook(self.hook)  # after the recorder's counter

    def hook(self, module, args):
        k = self.rec.call
        x = args[0]
        if k >= self.n_field_times or x.shape[0] != 1:
            return
        for n in self.names:
            v = x[0, self.index[n]].detach().float() * self.std[n] + self.mean[n]
            v = v.reshape(-1).cpu().numpy()
            if n not in self.values:
                self.values[n] = np.full((self.n_field_times, v.size), np.nan, dtype=np.float32)
            self.values[n][k] = v


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("config", help="fme inference YAML (initial_condition + forcing_loader)")
    p.add_argument("--out", required=True, help="archive directory to write")
    p.add_argument("--steps", type=int, required=True, help="forward steps to run")
    p.add_argument("--every", type=int, default=1, help="record every Nth step")
    p.add_argument("--layers", default="all", help="'all' or comma list, e.g. 0,4,8")
    p.add_argument(
        "--fields",
        default=None,
        help="output variables kept in reference.nc (default: a set per component)",
    )
    p.add_argument(
        "--inputs",
        default=",".join(DEFAULT_INPUTS),
        help="network inputs kept in reference.nc, as the network saw them "
        "(interventions included); named <name>_input when an output shares the name",
    )
    p.add_argument("--dtype", default="float16")
    p.add_argument(
        "--seed",
        type=int,
        default=0,
        help="rollout seed; the noise-conditioned SFNO is stochastic, so runs meant "
        "to be compared node by node must share it",
    )
    p.add_argument("--steer", help="L:BASIS.npz:FEATURE:ALPHA  add alpha x a basis direction")
    p.add_argument("--steer-channel", help="L:CHANNEL:AMOUNT  add a constant to one channel")
    p.add_argument(
        "--steer-vector", help="L:VECTOR.npy:ALPHA  add alpha x a saved (channels,) vector"
    )
    p.add_argument("--steer-window", default=None, help="FIRST:LAST forward steps to steer")
    p.add_argument(
        "--skip-rescale",
        help="NAME:FACTOR  SFNO with a big skip only: scale one input, in physical units, in the "
        "copy of the inputs the skip hands the decoder, leaving what the blocks see alone",
    )
    p.add_argument("--model-name", default=None, help="manifest 'model' (default: run id)")
    p.add_argument("--experiment", default="{}", help="JSON merged into the manifest's experiment")
    p.add_argument("--xaig-src", default=None, help="path to an xaig checkout's src/")
    p.add_argument("--keep-run", action="store_true", help="keep the fme experiment_dir")
    return p.parse_args(argv)


def find_network(stepper):
    """(kind, outermost module, network, taps): taps are (module, 'pre'|'post', label)."""
    from fme.core.models.conditional_sfno.sfnonet import SphericalFourierNeuralOperatorNet

    modules = list(stepper.modules.modules())
    sfno = [m for m in modules if isinstance(m, SphericalFourierNeuralOperatorNet)]
    if len(sfno) == 1:
        net = sfno[0]
        from fme.ace.registry.stochastic_sfno import NoiseConditionedModel

        outer = [m for m in modules if isinstance(m, NoiseConditionedModel)]
        taps = [(net.blocks[0], "pre", "encoder output (input to block 1)")]
        taps += [(b, "post", f"SFNO block {i} output") for i, b in enumerate(net.blocks, 1)]
        return "sfno", (outer[0] if outer else net), net, taps
    from fme.ace.models.ocean.m2lines.samudra import Samudra

    samudra = [m for m in modules if isinstance(m, Samudra)]
    if len(samudra) == 1:
        net = samudra[0]
        taps = [
            (
                net.layers[0],
                "post",
                f"encoder level 1 output ({net.ch_width[0]} ch, full resolution)",
            ),
            (
                net.layers[-2],
                "post",
                f"last decoder block output ({net.ch_width[0]} ch, full resolution)",
            ),
        ]
        return "samudra", net, net, taps
    raise RuntimeError("found neither one SFNO nor one Samudra in the stepper")


class Recorder:
    """Forward hooks that record chosen layers at chosen steps and optionally steer one."""

    def __init__(
        self, outer, taps, layers, steps, every, staging: Path, dtype, steer=None, window=None
    ):
        self.outer = outer
        self.taps = taps
        self.n_layers = len(taps)
        self.layers = (
            list(range(self.n_layers)) if layers == "all" else [int(x) for x in layers.split(",")]
        )
        for layer in self.layers:
            if not 0 <= layer < self.n_layers:
                raise ValueError(f"layer {layer} not in 0..{self.n_layers - 1}")
        self.recorded_steps = list(range(0, steps, every))
        self.slot = {s: i for i, s in enumerate(self.recorded_steps)}
        self.staging = staging
        self.dtype = dtype
        self.steer = steer  # (layer, torch vector (C,))
        self.window = window
        self.arrays = {}
        self.call = -1
        self.shape = None

    def _array(self, layer, x):
        if layer not in self.arrays:
            n_ch, h, w = x.shape[1:]
            self.shape = (h, w)
            path = self.staging / f"raw_{layer:02d}.npy"
            self.arrays[layer] = np.lib.format.open_memmap(
                path, mode="w+", dtype=self.dtype, shape=(len(self.recorded_steps), h * w, n_ch)
            )
        return self.arrays[layer]

    def _steer(self, layer, x):
        if self.steer is None or self.steer[0] != layer:
            return x
        if self.window is not None and not (self.window[0] <= self.call <= self.window[1]):
            return x
        vec = self.steer[1].to(device=x.device, dtype=x.dtype)
        if vec.numel() < x.shape[1]:  # a one-channel push, given only up to its channel
            vec = torch.nn.functional.pad(vec, (0, x.shape[1] - vec.numel()))
        if vec.numel() != x.shape[1]:
            raise RuntimeError(
                f"steering direction has {vec.numel()} channels, layer {layer} has {x.shape[1]}"
            )
        return x + vec.view(1, -1, 1, 1)

    def _record(self, layer, x):
        if layer in self.layers and self.call in self.slot:
            if x.shape[0] != 1:
                raise RuntimeError(f"recording needs batch 1, got {x.shape[0]}")
            flat = x[0].detach().float().reshape(x.shape[1], -1).T  # (nodes, channels), lat-major
            self._array(layer, x)[self.slot[self.call]] = flat.cpu().numpy().astype(self.dtype)

    def attach(self):
        def count(module, args):
            self.call += 1

        self.outer.register_forward_pre_hook(count, prepend=True)
        for layer, (module, kind, _) in enumerate(self.taps):
            if kind == "pre":

                def pre(module, args, layer=layer):
                    x = self._steer(layer, args[0])
                    self._record(layer, x)
                    return (x, *args[1:])

                module.register_forward_pre_hook(pre)
            else:

                def post(module, args, out, layer=layer):
                    out = self._steer(layer, out)
                    self._record(layer, out)
                    return out

                module.register_forward_hook(post)

    def labels(self):
        return [label for _, _, label in self.taps]


def steering_vector(args):
    import torch

    if sum(bool(x) for x in (args.steer, args.steer_channel, args.steer_vector)) > 1:
        raise SystemExit("give one of --steer, --steer-channel, --steer-vector")
    if args.steer_vector:
        layer, path, alpha = args.steer_vector.split(":")
        vec = float(alpha) * np.load(path).astype(np.float64).ravel()
        info = {"layer": int(layer), "vector": Path(path).name, "by": float(alpha)}
        return int(layer), torch.as_tensor(vec, dtype=torch.float32), info
    if args.steer_channel:
        layer, ch, amount = args.steer_channel.split(":")
        vec = torch.zeros(int(ch) + 1)  # padded to the layer's width at first use
        vec[int(ch)] = float(amount)
        return int(layer), vec, {"layer": int(layer), "channel": int(ch), "by": float(amount)}
    if args.steer:
        from xaig.daig.latent import load_basis

        layer, path, feature, alpha = args.steer.split(":")
        basis = load_basis(path)
        feature, alpha = int(feature), float(alpha)
        # a Dictionary's activation a adds a * input_scale of its unit direction to the
        # layer; a PCA component is a unit direction in the layer's own units
        scale = float(getattr(basis, "input_scale", 1.0))
        vec = alpha * scale * np.asarray(basis.directions()[feature], dtype=np.float64)
        info = {
            "layer": int(layer),
            "basis": Path(path).name,
            "kind": basis.kind,
            "feature": feature,
            "by": alpha,
            "input_scale": scale,
        }
        return int(layer), torch.as_tensor(vec, dtype=torch.float32), info
    return None, None, None


def time_labels(ds):
    """Valid-time labels (ISO, 4-digit year) of the prediction file, initial condition first."""
    t = ds["valid_time"].values if "valid_time" in ds else ds["time"].values
    t = np.asarray(t).reshape(-1)
    return [x.strftime("%Y-%m-%dT%H:%M:%S") if hasattr(x, "strftime") else str(x)[:19] for x in t]


def main(argv=None):
    args = parse_args(argv)
    if args.xaig_src:
        sys.path.insert(0, args.xaig_src)
    import fme.ace.inference.evaluator as evaluator
    import fme.ace.inference.inference as inference
    import xarray as xr

    from xaig.adapters.latent_archive import write_archive
    from xaig.daig.grid import Grid

    out = Path(args.out).resolve()
    staging = out.parent / f".{out.name}.work"
    if staging.exists():
        shutil.rmtree(staging)
    staging.mkdir(parents=True)
    run_dir = staging / "fme"

    cfg = yaml.safe_load(Path(args.config).read_text())
    source_experiment_dir = cfg.get("experiment_dir")
    cfg["experiment_dir"] = str(run_dir)
    cfg["n_forward_steps"] = args.steps
    cfg["seed"] = args.seed
    cfg["forward_steps_in_memory"] = min(cfg.get("forward_steps_in_memory", 20), args.steps)
    evaluating = "loader" in cfg
    fields = None if args.fields is None else [f for f in args.fields.split(",") if f]
    if fields is None:
        import torch as _torch

        ckpt_cfg = _torch.load(cfg["checkpoint_path"], map_location="cpu", weights_only=False)
        out_names = ckpt_cfg["stepper"]["config"]["step"]["config"].get("out_names", [])
        del ckpt_cfg
        fields = OCEAN_FIELDS if "sst" in out_names else DEFAULT_FIELDS
    cfg["data_writer"] = {
        "save_prediction_files": bool(fields),
        "save_monthly_files": False,
        **({"names": fields} if fields else {}),
    }
    cfg.setdefault("logging", {})["log_to_wandb"] = False
    run_yaml = staging / "run.yaml"
    run_yaml.write_text(yaml.safe_dump(cfg, sort_keys=False))

    state = {}
    config_class = evaluator.InferenceEvaluatorConfig if evaluating else inference.InferenceConfig
    original_load = config_class.load_stepper

    def load_stepper_with_hooks(self):
        stepper = original_load(self)
        kind, outer, net, taps = find_network(stepper)
        s_layer, s_vec, s_info = steering_vector(args)
        window = tuple(int(v) for v in args.steer_window.split(":")) if args.steer_window else None
        rec = Recorder(
            outer, taps, args.layers, args.steps, args.every, staging, args.dtype,
            steer=(s_layer, s_vec) if s_vec is not None else None, window=window,
        )  # fmt: skip
        rec.attach()
        inputs = [n for n in args.inputs.split(",") if n]
        if inputs:
            state["inputs"] = InputCapture(stepper, outer, inputs, rec, args.steps + 1)
        if args.skip_rescale:
            if kind != "sfno":
                raise SystemExit("--skip-rescale applies to an SFNO only")
            name, factor = args.skip_rescale.rsplit(":", 1)
            SkipRescale(stepper, net, name, float(factor))
            state["skip"] = {"input": name, "factor": float(factor)}
        state.update(rec=rec, steer=s_info, window=window, kind=kind)
        return stepper

    config_class.load_stepper = load_stepper_with_hooks
    try:
        (evaluator if evaluating else inference).main(str(run_yaml))
    finally:
        config_class.load_stepper = original_load

    rec: Recorder = state["rec"]
    if rec.call + 1 != args.steps:
        logging.warning("model was called %d times for %d steps", rec.call + 1, args.steps)

    pred = (
        xr.open_dataset(run_dir / "autoregressive_predictions.nc", decode_timedelta=False)
        if fields
        else None
    )
    ic = xr.open_dataset(run_dir / "initial_condition.nc") if fields else None
    if pred is None:
        raise SystemExit(
            "--fields '' is not supported yet: the grid and times come from the prediction file"
        )
    lat1 = pred["lat"].values
    lon1 = pred["lon"].values
    lon, lat = np.meshgrid(lon1, lat1)
    ocean = state.get("kind") == "samudra"
    mask = None
    if ocean and "sst" in pred:  # land holds no ocean: the first predicted SST says where
        mask = np.isfinite(pred["sst"].values.reshape(-1, lat1.size * lon1.size)[0])
    grid = Grid(lat=lat.ravel(), lon=lon.ravel(), shape=(lat1.size, lon1.size), mask=mask)
    if rec.shape != (lat1.size, lon1.size):
        raise SystemExit(
            f"latent grid {rec.shape} differs from the output grid {(lat1.size, lon1.size)}"
        )

    all_times = time_labels(pred)  # prediction valid times, one per step
    ic_time = time_labels(ic)[0] if "valid_time" in ic or "time" in ic else None
    if len(all_times) == args.steps and ic_time is not None:
        field_times = [ic_time, *all_times]
    else:
        field_times = all_times
    step_input_times = field_times[: args.steps]  # the state each forward call started from
    times = [step_input_times[s] for s in rec.recorded_steps]

    field_arrays = {}
    for name in fields:
        if name not in pred:
            logging.warning("field %s not in predictions; skipped", name)
            continue
        values = pred[name].values.reshape(-1, lat1.size * lon1.size)  # (steps, nodes)
        if len(field_times) == values.shape[0] + 1:
            first = (
                ic[name].values.reshape(1, -1)
                if name in ic
                else np.full((1, values.shape[1]), np.nan)
            )
            values = np.concatenate([first, values])
        field_arrays[name] = values.astype(np.float32)
    captured = state.get("inputs")
    if captured is not None:
        if len(field_times) != args.steps + 1:
            raise SystemExit(
                "input capture needs the initial condition's time in the prediction files"
            )
        for name, values in captured.values.items():
            key = f"{name}_input" if name in field_arrays else name
            field_arrays[key] = values

    ckpt = cfg["checkpoint_path"]
    run_id = Path(ckpt).parent.parent.name
    experiment = {
        "run_id": run_id,
        "source_config": str(Path(args.config).resolve()),
        "source_experiment_dir": source_experiment_dir,
        "initial_condition": cfg.get("initial_condition", {}).get("path"),
        "seed": cfg.get("seed"),
        "recorded_steps": {"n": args.steps, "every": args.every},
        "exported": dt.datetime.now().isoformat(timespec="seconds"),
        **json.loads(args.experiment),
    }
    if state.get("skip"):
        experiment["skip_rescale"] = state["skip"]
    if state.get("steer"):
        experiment["steer"] = {
            **state["steer"],
            **({"window": list(state["window"])} if state["window"] else {}),
        }

    labels = rec.labels()
    layers = [(labels[i], rec.arrays[i]) for i in sorted(rec.arrays)]
    network_layers = None
    if sorted(rec.arrays) != list(range(len(layers))):
        # the reader indexes layers by position; record where each sits in the network, so
        # that a basis is matched by that place (and keep it in the label, for people)
        layers = [(f"{labels[i]} [network layer {i}]", rec.arrays[i]) for i in sorted(rec.arrays)]
        network_layers = sorted(rec.arrays)
    step_seconds = (
        int(
            round(
                (np.datetime64(field_times[1]) - np.datetime64(field_times[0]))
                / np.timedelta64(1, "s")
            )
        )
        if len(field_times) > 1
        else None
    )
    write_archive(
        out,
        grid=grid,
        times=times,
        layers=layers,
        network_layers=network_layers,
        fields=field_arrays or None,
        field_times=field_times if field_arrays else None,
        model=args.model_name or run_id,
        component="ocean" if ocean else "atmosphere",
        checkpoint=ckpt,
        calendar=str(pred["time"].attrs.get("calendar", "noleap")) if "time" in pred else "noleap",
        timestep_seconds=step_seconds,
        experiment=experiment,
        dtype=args.dtype,
        overwrite=True,
    )
    if args.keep_run:
        shutil.move(str(run_dir), str(out / "_fme_run"))
    shutil.rmtree(staging, ignore_errors=True)
    print(
        f"wrote {out}: {len(times)} times x {len(layers)} layers; fields: {', '.join(field_arrays)}"
    )


if __name__ == "__main__":
    main()