checkpoints, kit and exporter

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+kit/forcing/v3.LR.historical_0101.aigo.eam.h0.2015-01.nc filter=lfs diff=lfs merge=lfs -text
+kit/forcing/v3.LR.historical_0101.aigo.eam.h0.2015-02.nc filter=lfs diff=lfs merge=lfs -text
+kit/forcing-aerofrz/aero.2015-01.nc filter=lfs diff=lfs merge=lfs -text
+kit/forcing-aerofrz/aero.2015-02.nc filter=lfs diff=lfs merge=lfs -text
+kit/ic/E05-FT-S01.restart.2015-01-03T12.nc filter=lfs diff=lfs merge=lfs -text

E05-FT-S01/best_inference_ckpt.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:086629bcecc824a6a4c21bb08051a972f3d34b9ed0fcc002587f86bfbadcb3d2
+size 1824999519

E11-FT-S01/best_inference_ckpt.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f50c886354f30076309b53fbb41e6aef33202a0aa32a774de45e08ae36cf9a3e
+size 341473131

README.md

@@ -1,3 +1,61 @@
 ---
 license: bsd-3-clause
+tags:
+  - climate
+  - emulator
+  - e3sm
+  - ace
+  - samudra
 ---
+
+# E3SMv3 emulator checkpoints for the E3SM AI group hackathon, 25 September 2026
+
+> **Not fit for forcing or scenario experiments.** Both checkpoints reproduce the E3SMv3 historical run well
+> when the other half of the climate system is handed to them, but neither responds to a change in forcing
+> physically: the atmosphere uses CO₂ as a clock, and the ocean reads its answer off its upward-longwave input.
+> Use them to study how an emulator represents the climate, not what the climate does. Results from them show
+> how the model *represents* forcing.
+
+| Folder | Checkpoint | Fit for | Not fit for |
+| --- | --- | --- | --- |
+| `E05-FT-S01/` | `E05-FT.aug26.atm.A3_B16_C1_L0_O5_W0_X0.S01`, atmosphere, noise-conditioned SFNO (8 blocks × 384 channels), 1°, 6-hourly, SST prescribed | Historical runs with prescribed SST; looking inside the network | Changed CO₂ or SST, any forcing experiment; precipitation extremes (about 2× too frequent) |
+| `E11-FT-S01/` | `E11-FT.aug26.ocn.A0_B16_C0_L0_O5_W0_X0.S01`, ocean, Samudra U-Net, 1°, 5-daily, surface fluxes prescribed | Replaying the historical run from reference fluxes; looking inside the network | Any change to the fluxes; ENSO studies; subsurface heat uptake |
+
+These are the checkpoints the latents in
+[E3SM-Project/aigs-hack-sep26-latents](https://huggingface.co/datasets/E3SM-Project/aigs-hack-sep26-latents)
+were recorded from. E05-FT S01 is not the shipped atmosphere (S03): it is the one the probes behind the
+"not fit for" column used.
+
+## Run an experiment yourself
+
+The `kit/` folder holds what a 2015 atmosphere run needs besides the checkpoint: the initial condition
+(2015-01-03 12:00), two months of forcing (the inputs only), and configs for the control and each intervention.
+One run uses about 4 GB of GPU memory.
+
+```console
+$ pip install "fme @ git+https://github.com/E3SM-Project/ace@e3sm/exps/hist-v2026.8.0"
+$ hf download E3SM-Project/aigs-hack-sep26-models --local-dir aigs-models && cd aigs-models
+$ python -m fme.ace.inference kit/configs/atm-co2x1.25.yaml
+```
+
+| Config | Intervention |
+| --- | --- |
+| `atm-ctrl.yaml` | none |
+| `atm-co2x1.25.yaml`, `atm-co2x0.8.yaml` | `overwrite.multiply_scalar.global_mean_co2` |
+| `atm-p4K.yaml` | `overwrite.add_scalar.TS: 4.0` |
+| `atm-aerofrz.yaml` | aerosol diagnostics swapped for their 1965 values (`kit/forcing-aerofrz/`) |
+
+Every config sets `seed: 0`. Keep it: the model is stochastic, and a shared seed is what makes two runs
+comparable node for node. Change the intervention by editing the `overwrite` block; make it longer with
+`n_forward_steps` (the forcing covers January and February 2015).
+
+To record the latents as well, and to steer a layer along a feature, use the exporter in the kit
+(`kit/export_fme_latents.py`; it needs [xaig](https://pypi.org/project/xaig/) 0.2 or later in the same
+environment):
+
+```console
+$ python kit/export_fme_latents.py kit/configs/atm-ctrl.yaml --out my-run --steps 20 \
+    --steer 0:sae_L00.npz:142:3      # push SAE feature 142 of layer 0 by +3 at every node
+```
+
+It reproduces the published latents bit for bit.

kit/configs/atm-aerofrz.yaml

@@ -0,0 +1,39 @@
+experiment_dir: output/atm-aerofrz
+n_forward_steps: 20
+forward_steps_in_memory: 20
+checkpoint_path: E05-FT-S01/best_inference_ckpt.tar
+logging:
+  log_to_screen: true
+  log_to_wandb: false
+  log_to_file: true
+initial_condition:
+  path: kit/ic/E05-FT-S01.restart.2015-01-03T12.nc
+  engine: netcdf4
+forcing_loader:
+  dataset:
+    merge:
+    - data_path: kit/forcing
+      file_pattern: v3.LR.historical_0101.aigo.eam.h0.*.nc
+      rename:
+        co2vmr: global_mean_co2
+        aerindexall: aerindexall_orig
+        colccn.3: colccn3_orig
+      reference_pressure_name: P0
+    - data_path: kit/forcing-aerofrz
+      file_pattern: aero.*.nc
+      rename:
+        aerindexall_F: aerindexall
+        colccn.3_F: colccn.3
+  num_data_workers: 2
+data_writer:
+  save_prediction_files: true
+  save_monthly_files: false
+  names:
+  - PS
+  - TS
+  - Tat2m
+  - surface_precipitation_rate
+  - FLUT
+  - FSUTOA
+  - lwp
+seed: 0

kit/configs/atm-co2x0.8.yaml

@@ -0,0 +1,34 @@
+experiment_dir: output/atm-co2x0.8
+n_forward_steps: 20
+forward_steps_in_memory: 20
+checkpoint_path: E05-FT-S01/best_inference_ckpt.tar
+logging:
+  log_to_screen: true
+  log_to_wandb: false
+  log_to_file: true
+initial_condition:
+  path: kit/ic/E05-FT-S01.restart.2015-01-03T12.nc
+  engine: netcdf4
+forcing_loader:
+  dataset:
+    data_path: kit/forcing
+    file_pattern: v3.LR.historical_0101.aigo.eam.h0.*.nc
+    rename:
+      co2vmr: global_mean_co2
+    reference_pressure_name: P0
+    overwrite:
+      multiply_scalar:
+        global_mean_co2: 0.8
+  num_data_workers: 2
+data_writer:
+  save_prediction_files: true
+  save_monthly_files: false
+  names:
+  - PS
+  - TS
+  - Tat2m
+  - surface_precipitation_rate
+  - FLUT
+  - FSUTOA
+  - lwp
+seed: 0

kit/configs/atm-co2x1.25.yaml

@@ -0,0 +1,34 @@
+experiment_dir: output/atm-co2x1.25
+n_forward_steps: 20
+forward_steps_in_memory: 20
+checkpoint_path: E05-FT-S01/best_inference_ckpt.tar
+logging:
+  log_to_screen: true
+  log_to_wandb: false
+  log_to_file: true
+initial_condition:
+  path: kit/ic/E05-FT-S01.restart.2015-01-03T12.nc
+  engine: netcdf4
+forcing_loader:
+  dataset:
+    data_path: kit/forcing
+    file_pattern: v3.LR.historical_0101.aigo.eam.h0.*.nc
+    rename:
+      co2vmr: global_mean_co2
+    reference_pressure_name: P0
+    overwrite:
+      multiply_scalar:
+        global_mean_co2: 1.25
+  num_data_workers: 2
+data_writer:
+  save_prediction_files: true
+  save_monthly_files: false
+  names:
+  - PS
+  - TS
+  - Tat2m
+  - surface_precipitation_rate
+  - FLUT
+  - FSUTOA
+  - lwp
+seed: 0

kit/configs/atm-ctrl.yaml

@@ -0,0 +1,31 @@
+experiment_dir: output/atm-ctrl
+n_forward_steps: 20
+forward_steps_in_memory: 20
+checkpoint_path: E05-FT-S01/best_inference_ckpt.tar
+logging:
+  log_to_screen: true
+  log_to_wandb: false
+  log_to_file: true
+initial_condition:
+  path: kit/ic/E05-FT-S01.restart.2015-01-03T12.nc
+  engine: netcdf4
+forcing_loader:
+  dataset:
+    data_path: kit/forcing
+    file_pattern: v3.LR.historical_0101.aigo.eam.h0.*.nc
+    rename:
+      co2vmr: global_mean_co2
+    reference_pressure_name: P0
+  num_data_workers: 2
+data_writer:
+  save_prediction_files: true
+  save_monthly_files: false
+  names:
+  - PS
+  - TS
+  - Tat2m
+  - surface_precipitation_rate
+  - FLUT
+  - FSUTOA
+  - lwp
+seed: 0

kit/configs/atm-p4K.yaml

@@ -0,0 +1,34 @@
+experiment_dir: output/atm-p4K
+n_forward_steps: 20
+forward_steps_in_memory: 20
+checkpoint_path: E05-FT-S01/best_inference_ckpt.tar
+logging:
+  log_to_screen: true
+  log_to_wandb: false
+  log_to_file: true
+initial_condition:
+  path: kit/ic/E05-FT-S01.restart.2015-01-03T12.nc
+  engine: netcdf4
+forcing_loader:
+  dataset:
+    data_path: kit/forcing
+    file_pattern: v3.LR.historical_0101.aigo.eam.h0.*.nc
+    rename:
+      co2vmr: global_mean_co2
+    reference_pressure_name: P0
+    overwrite:
+      add_scalar:
+        TS: 4.0
+  num_data_workers: 2
+data_writer:
+  save_prediction_files: true
+  save_monthly_files: false
+  names:
+  - PS
+  - TS
+  - Tat2m
+  - surface_precipitation_rate
+  - FLUT
+  - FSUTOA
+  - lwp
+seed: 0

kit/export_fme_latents.py

@@ -0,0 +1,533 @@
+"""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)]
+    if sorted(rec.arrays) != list(range(len(layers))):
+        # the reader indexes layers by position; keep the network's numbering in the label
+        layers = [(f"{labels[i]} [network layer {i}]", rec.arrays[i]) for i in 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,
+        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()

kit/forcing-aerofrz/aero.2015-01.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4852019df49bfe240e0c2c0234a2ab234aaaa425bff081dbd57713528a83719
+size 64297664

kit/forcing-aerofrz/aero.2015-02.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5b3201db3d9654b2e18d421aafa298ee0b91001b1082f888f4840131ce143de
+size 58076768

kit/forcing/v3.LR.historical_0101.aigo.eam.h0.2015-01.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93c82a6f0211c593b7b26f5c503b2c8f16b8f4ba2b7252f97f26e98906526378
+size 134716164

kit/forcing/v3.LR.historical_0101.aigo.eam.h0.2015-02.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:198700789f49c2ef215ca3bdc61531d7ad03ced68efc916a626312d474833881
+size 121586871

kit/ic/E05-FT-S01.restart.2015-01-03T12.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e28123a6d0a4cfba3c701e35995ad744fe19cf6ad88caf8dd3a40c49024eebc
+size 9896940