scripts/inference.py

#!/usr/bin/env python3
"""Run an official NeuralGCM pressure-level rollout."""
from __future__ import annotations
import argparse
import math
import pickle
import sys
import time
from pathlib import Path
import numpy as np
try:
    from common import PROJECT_ROOT, add_static_features, as_time_major_frames, era5_data_is_synthetic, era5_sample_to_xarray, load_config, load_era5_dataset, regrid_for_neuralgcm, resolve_path, validate_synthetic_era5_version
except ModuleNotFoundError:  # supports ``python -m scripts.inference``
    from scripts.common import PROJECT_ROOT, add_static_features, as_time_major_frames, era5_data_is_synthetic, era5_sample_to_xarray, load_config, load_era5_dataset, regrid_for_neuralgcm, resolve_path, validate_synthetic_era5_version
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
from model.NeuralGCM import checkpoint_mode, load_checkpoint, validate_checkpoint_mode

MODE_ALIASES = {"forecast": "weather_forecast", "weather_forecast": "weather_forecast", "climate": "climate_scale", "climate_scale": "climate_scale", "forecast_2_8_deg": "forecast_2_8_deg", "stochastic_1_4_deg": "stochastic_1_4_deg"}


def _validate_checkpoint_mode(candidate: Path, mode: str) -> None:
    """Reject a checkpoint whose declared/profile grid differs from --mode."""
    try:
        with candidate.open("rb") as handle:
            payload = pickle.load(handle)
    except Exception:
        return
    validate_checkpoint_mode(payload, mode, candidate)


def _official_checkpoint(config: dict, mode: str, explicit: str | None) -> Path:
    if explicit:
        candidate = resolve_path(explicit)
        if candidate.exists():
            _validate_checkpoint_mode(candidate, mode)
        return candidate
    configured = config["inference"].get("checkpoint")
    if configured:
        candidate = resolve_path(configured)
        if candidate.exists():
            try:
                with candidate.open("rb") as handle:
                    payload = pickle.load(handle)
                if isinstance(payload, dict) and {"model_config_str", "aux_ds_dict", "params"}.issubset(payload):
                    try:
                        _validate_checkpoint_mode(candidate, mode)
                    except ValueError as exc:
                        # The default path is a convenience pointer. If a
                        # previous run left a checkpoint for another profile,
                        # use the bundled official checkpoint for the requested
                        # mode; explicit --checkpoint remains strict.
                        print(f"Warning: {exc}; falling back to profile checkpoint")
                        return resolve_path(config["model"]["profiles"][mode]["official_reference"])
                    return candidate
            except Exception:
                pass
        if Path(configured).name != "model_bak.pkl":
            return candidate
    return resolve_path(config["model"]["profiles"][mode]["official_reference"])


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--mode")
    parser.add_argument("--checkpoint", help="official NeuralGCM checkpoint (.pkl)")
    parser.add_argument("--input-nc")
    parser.add_argument("--data-dir")
    parser.add_argument("--years", nargs="+", type=int)
    parser.add_argument("--sample-index", type=int, default=0)
    parser.add_argument("--steps", type=int)
    parser.add_argument(
        "--output-interval-hours",
        type=int,
        help="hours between saved forecasts; defaults to inference.output_interval_hours",
    )
    parser.add_argument("--output")
    parser.add_argument(
        "--seed",
        type=int,
        help="PRNG seed for stochastic profiles; overrides inference.seed",
    )
    parser.add_argument("--device", default="auto", help="jax platform: auto, cpu, or gpu")
    args = parser.parse_args()
    config = load_config(args.config)
    if args.data_dir:
        config["data"]["data_dir"] = args.data_dir
        paired_static = resolve_path(args.data_dir, args.config) / "static.nc"
        if paired_static.exists():
            config["data"]["static_file"] = str(paired_static)
    requested = args.mode or config["inference"].get("mode", "weather_forecast")
    mode = MODE_ALIASES.get(requested, requested)
    if mode not in config["model"].get("profiles", {}):
        raise ValueError(f"Unknown mode {requested!r}")
    if args.device != "auto":
        import jax
        jax.config.update("jax_platform_name", args.device)
    checkpoint = _official_checkpoint(config, mode, args.checkpoint)
    model = load_checkpoint(checkpoint)
    print(f"Official checkpoint: {checkpoint}")
    grid = (model.data_coords.horizontal.latitudes.size, model.data_coords.horizontal.longitudes.size)
    print(f"Model mode={mode}, timestep={model.timestep}, grid={grid}")
    steps = int(args.steps or config["inference"].get("prediction_steps", 8))
    if steps <= 0:
        raise ValueError("steps must be positive")
    output_interval_hours = int(
        args.output_interval_hours
        or config["inference"].get("output_interval_hours", 6)
    )
    if output_interval_hours <= 0:
        raise ValueError("output-interval-hours must be positive")
    output_interval = np.timedelta64(output_interval_hours, "h")
    if args.input_nc:
        import xarray as xr
        dataset = xr.load_dataset(resolve_path(args.input_nc))
        synthetic_input = False
    else:
        years = args.years or list(config["data"].get("test_years", [2002]))
        synthetic_input = era5_data_is_synthetic(config, years)
        if synthetic_input:
            validate_synthetic_era5_version(config, years)
        data_interval_hours = int(config["data"].get("time_step_hours", 6))
        forcing_steps = math.ceil(steps * output_interval_hours / data_interval_hours)
        # Keep the forcing trajectory available at every requested lead time.
        # ERA5Dataset returns (initial state, future frames, ...); using its
        # default output_steps=1 would silently hold SST/sea-ice forcing fixed.
        source = load_era5_dataset(
            config, years, input_steps=1, output_steps=forcing_steps
        )
        # OneScience exposes ``total_samples`` directly, but its ``__len__``
        # returns the raw (possibly negative) window count.  Calling len() on
        # an undersized file therefore raises Python's own ``ValueError``
        # before we can explain which forcing window is missing.
        dataset_size = int(getattr(source, "total_samples", 0))
        if dataset_size <= 0:
            raise ValueError(
                "ERA5Dataset has no complete inference window: "
                f"T={getattr(source, 'T', '?')}, input_steps=1, "
                f"output_steps={forcing_steps}. Provide at least "
                f"{forcing_steps + 1} consecutive frames."
            )
        if not 0 <= args.sample_index < dataset_size:
            raise IndexError(
                f"sample-index {args.sample_index} outside dataset of length "
                f"{dataset_size}"
            )
        sample = source[args.sample_index]
        target_frames = as_time_major_frames(sample[1], name="ERA5 target")
        frames = [sample[0]] + [
            target_frames[t] for t in range(min(forcing_steps, len(target_frames)))
        ]
        import xarray as xr
        sample_time = sample[4][0]
        if not isinstance(sample_time, str) or len(sample_time) != 10 or not sample_time.isdigit():
            raise ValueError(f"invalid ERA5Dataset time index {sample_time!r}")
        start_timestamp = np.datetime64(
            f"{sample_time[:4]}-{sample_time[4:6]}-{sample_time[6:8]}T{sample_time[8:10]}:00:00"
        )
        dataset = xr.concat(
            [
                era5_sample_to_xarray(
                    (frame, frame),
                    config,
                    timestamp=start_timestamp
                    + np.timedelta64(int(config["data"].get("time_step_hours", 6)) * t, "h"),
                )
                for t, frame in enumerate(frames)
            ],
            dim="time",
        )
    dataset = regrid_for_neuralgcm(dataset, model)
    dataset = add_static_features(
        dataset, config, mode=mode, prefer_profile=not synthetic_input
    )
    data_in, forcings_in = model.data_from_xarray(dataset.isel(time=0))
    forcing_trajectory = model.forcings_from_xarray(dataset)
    import jax
    seed = int(
        args.seed
        if args.seed is not None
        else config["inference"].get("seed", config["project"].get("seed", 0))
    )
    state = model.encode(data_in, forcings_in, rng_key=jax.random.key(seed))
    start = time.perf_counter()
    _, outputs = model.unroll(
        state,
        forcing_trajectory,
        steps=steps,
        timedelta=output_interval,
    )
    times = np.arange(
        dataset.time.values[0] + output_interval,
        dataset.time.values[0] + (steps + 1) * output_interval,
        output_interval,
    )
    result = model.data_to_xarray(outputs, times=times)
    nonfinite = [
        name
        for name, values in result.data_vars.items()
        if np.issubdtype(values.dtype, np.number)
        and not bool(np.isfinite(values.values).all())
    ]
    if nonfinite:
        raise FloatingPointError(
            "NeuralGCM rollout produced NaN/Inf in variables "
            f"{nonfinite}. The input passed shape checks but is not a stable "
            "physical initial condition for this checkpoint."
        )
    result.attrs.update({"neuralgcm_mode": mode, "checkpoint": str(checkpoint), "official_api": "PressureLevelModel", "random_seed": seed})
    output = resolve_path(args.output or config["inference"].get("output", "results/predictions.nc"), args.config)
    output.parent.mkdir(parents=True, exist_ok=True)
    result.to_netcdf(output)
    print(
        f"Official rollout steps={steps}, output_interval={output_interval_hours}h, "
        f"horizon={steps * output_interval_hours / 24:g} days, "
        f"elapsed={time.perf_counter() - start:.2f}s"
    )
    print(f"Saved {output} with variables={list(result.data_vars)}")


if __name__ == "__main__":
    main()