Support AIMIP IC and forcing generation for a different ERA5 build

scripts/aimip_forcing/Makefile

@@ -6,14 +6,18 @@ GCS_PATH_PUBLIC_AIMIP_FORCING := gs://vcm-ml-intermediate/2025-09-05-era5-0.25de
 ZENODO_FORCING_URL := https://zenodo.org/api/records/17065758/files/$(PUBLIC_FORCING_FILE_NAME)/content
 LOCAL_FORCING_FILE := $(LOCAL_DATA_DIR)/$(PUBLIC_FORCING_FILE_NAME)
 LOCAL_REGRIDDED_FILE := $(LOCAL_DATA_DIR)/ERA5-1deg-monthly-mean-forcing-1978-2024.nc
-OUTPUT_ZARR_NAME = 2025-09-09-aimip-era5-1deg-forcing-1978-2024.zarr
+OUTPUT_ZARR_NAME ?= 2025-09-09-aimip-era5-1deg-forcing-1978-2024.zarr
 LOCAL_OUTPUT_ZARR := $(LOCAL_DATA_DIR)/$(OUTPUT_ZARR_NAME)
 GCS_PATH_PROCESSED_FORCING ?= gs://vcm-ml-intermediate/$(OUTPUT_ZARR_NAME)
 ENVIRONMENT_NAME=regrid-aimip-forcing
 IC_OUTPUT_DIR := $(LOCAL_DATA_DIR)/aimip_ics
+# e.g. --include-near-surface to add TMP2m/Q2m/UGRD10m/VGRD10m
+IC_EXTRA_ARGS ?=
 GCS_IC_PATH ?= gs://vcm-ml-intermediate/2025-09-12-AIMIP-evaluation-ICs/
 ERA5_GCS_DATA ?= gs://vcm-ml-intermediate/2024-06-20-era5-1deg-8layer-1940-2022.zarr
-MODIFIED_FORCING_NAME = 2025-09-12-aimip-era5-1deg-forcing-1978-2024-repeat-first.zarr
+# e.g. --ace2-era5-gcs-data <zarr> to match the model's ERA5 build
+FORCING_EXTRA_ARGS ?=
+MODIFIED_FORCING_NAME ?= 2025-09-12-aimip-era5-1deg-forcing-1978-2024-repeat-first.zarr
 LOCAL_MODIFIED_FORCING := $(LOCAL_DATA_DIR)/$(MODIFIED_FORCING_NAME)
 GCS_PATH_MODIFIED_FORCING ?= gs://vcm-ml-intermediate/$(MODIFIED_FORCING_NAME)
 
@@ -45,7 +49,8 @@ $(LOCAL_REGRIDDED_FILE): $(LOCAL_FORCING_FILE)
 	python regrid_aimip_forcing.py $(LOCAL_FORCING_FILE) $(LOCAL_REGRIDDED_FILE)
 
 $(LOCAL_OUTPUT_ZARR): $(LOCAL_REGRIDDED_FILE)
-	python interpolate_aimip_forcing.py $(LOCAL_REGRIDDED_FILE) $(LOCAL_OUTPUT_ZARR)
+	python interpolate_aimip_forcing.py $(LOCAL_REGRIDDED_FILE) $(LOCAL_OUTPUT_ZARR) \
+	  $(FORCING_EXTRA_ARGS)
 
 upload_processed_aimip_forcing: $(LOCAL_OUTPUT_ZARR)
 	gsutil -m cp -r $(LOCAL_OUTPUT_ZARR) $(GCS_PATH_PROCESSED_FORCING)
@@ -56,7 +61,7 @@ process_aimip_forcing: upload_processed_aimip_forcing
 $(IC_OUTPUT_DIR)/.done: data_dir
 	mkdir -p $(IC_OUTPUT_DIR)
 	python create_aimip_ic_datasets.py $(IC_OUTPUT_DIR) \
-	  --era5-gcs-data $(ERA5_GCS_DATA)
+	  --era5-gcs-data $(ERA5_GCS_DATA) $(IC_EXTRA_ARGS)
 	touch $(IC_OUTPUT_DIR)/.done
 
 upload_aimip_ics: $(IC_OUTPUT_DIR)/.done

scripts/aimip_forcing/README.md

@@ -25,6 +25,36 @@ The GCS path for the resulting zarr dataset can be specified:
 
 Note that the workflow is memory-intensive and was run on a high-memory (128GB) GCP VM.
 
+### Generating AIMIP evaluation initial conditions
+
+`create_aimip_ic_datasets.py` builds the ICs a model is initialized from for AIMIP inference.
+Each member is one timestamp from the ERA5 zarr, restamped to a common target time.
+
+```make create_aimip_ics```
+
+- `ERA5_GCS_DATA` — source zarr; use the build the model was trained on.
+- `GCS_IC_PATH` — upload destination.
+- `IC_EXTRA_ARGS` — e.g. `--include-near-surface` for models carrying
+  `TMP2m`/`Q2m`/`UGRD10m`/`VGRD10m`. `--target-timestamp` and `--ic-timestamp` set the
+  restamped and source times.
+
+ICs carry only prognostic variables; forcings come from the forcing zarr at inference time.
+
+### Regenerating the forcing for a different ERA5 build
+
+`HGTsfc`, `DSWRFtoa` and `land_fraction` are read from an ACE ERA5 zarr, and differ between
+builds (the March 2026 pipeline rewrite moves `HGTsfc` by tens of metres in mountains), so
+that zarr should match the model's training build:
+
+```
+FORCING_EXTRA_ARGS="--ace2-era5-gcs-data gs://vcm-ml-intermediate/<era5>.zarr" \
+  make process_aimip_forcing
+```
+
+The source must span the whole window or the run errors; `--extension-start` opts into
+synthesizing the remainder, needed only for short sources like the 2022-ending
+`2024-06-20-…` store.
+
 ### Generating the public AIMIP forcing dataset
 
 Additionally, the public forcing dataset at 0.25° resolution [available on Zenodo as version 2](https://zenodo.org/records/17065758) can also be generated and uploaded to GCS here. To do so, run:

scripts/aimip_forcing/conftest.py

@@ -0,0 +1,7 @@
+# These scripts run in their own conda environment (see README) and require dask,
+# both for progress output and because set_shards_chunks calls DataArray.chunk.
+# The repo-wide test environment has no dask, so skip rather than fail there.
+try:
+    import dask.diagnostics  # noqa: F401
+except ImportError:
+    collect_ignore_glob = ["test_*.py"]

scripts/aimip_forcing/create_aimip_ic_datasets.py

@@ -26,15 +26,21 @@
     + [f"eastward_wind_{i}" for i in range(8)]
     + [f"northward_wind_{i}" for i in range(8)]
 )
+# Near-surface prognostics carried by some ACE models but not others. Included only
+# on request, since the source zarr must carry them and older models do not use them.
+NEAR_SURFACE_VARIABLES = ["TMP2m", "Q2m", "UGRD10m", "VGRD10m"]
 
 
 def create_ic(
     era5: xr.Dataset,
     ic_timestamp: str,
     target_timestamp: np.datetime64,
 ) -> xr.Dataset:
-    ic = era5.sel(time=ic_timestamp)
-    return ic.assign_coords(time=target_timestamp)
+    # Select with a single-element list so `time` stays a length-1 dimension rather
+    # than collapsing to a scalar coordinate. `get_initial_condition` requires the
+    # prognostic variables to have shape (n_samples, [spatial dims]).
+    ic = era5.sel(time=[ic_timestamp])
+    return ic.assign_coords(time=[target_timestamp])
 
 
 @click.command()
@@ -62,17 +68,30 @@ def create_ic(
         "Output files are named {target_date}_IC{i}.nc."
     ),
 )
+@click.option(
+    "--include-near-surface/--no-include-near-surface",
+    default=False,
+    help=(
+        "Include the near-surface prognostics (TMP2m, Q2m, UGRD10m, VGRD10m) required "
+        "by models that carry them as inputs. The source zarr must provide them."
+    ),
+)
 def main(
     local_output_dir: str,
     era5_gcs_data: str,
     target_timestamp: str,
     ic_timestamps: Tuple[str, ...],
+    include_near_surface: bool,
 ):
     logging.basicConfig(level=logging.INFO)
     os.makedirs(local_output_dir, exist_ok=True)
 
+    prognostic_variables = PROGNOSTIC_VARIABLES + (
+        NEAR_SURFACE_VARIABLES if include_near_surface else []
+    )
+
     logging.info(f"Opening ERA5 data from {era5_gcs_data}")
-    era5 = xr.open_zarr(era5_gcs_data)[PROGNOSTIC_VARIABLES]
+    era5 = xr.open_zarr(era5_gcs_data)[prognostic_variables]
 
     target_dt = np.datetime64(target_timestamp)
     target_date = target_timestamp.split("T")[0]

scripts/aimip_forcing/encoding.py

@@ -24,6 +24,10 @@ def set_shards_chunks(ds, shards=OUTPUT_SHARDING, chunks=OUTPUT_CHUNKING):
     """
     out_ds = xr.Dataset()
     for name, da in ds.data_vars.items():
+        if not da.dims:
+            # zarr rejects empty chunk/shard tuples, and a scalar needs neither.
+            out_ds[name] = da
+            continue
         da_chunks = []
         da_shards = []
         chunking_dict = {}

scripts/aimip_forcing/interpolate_aimip_forcing.py

@@ -24,6 +24,11 @@
 ]
 START_TIME = "1978-10-01T00:00:00"
 END_TIME = "2024-12-31T18:00:00"
+# Empty extension start = no extension; the repeat-source defaults below pair with
+# the 2022-ending 2024-06-20 store.
+DEFAULT_EXTENSION_START = ""
+DEFAULT_REPEAT_SOURCE_START = "2020-12-31T00:00:00"
+DEFAULT_REPEAT_SOURCE_END = "2022-12-31T18:00:00"
 
 
 def open_aimip_forcing_data(
@@ -197,12 +202,37 @@ def write_output_zarr(ds: xr.Dataset, output_data_file: str):
     default=ACE2_ERA5_DATA,
     help="Path to ACE2 ERA5 data in GCS.",
 )
+@click.option(
+    "--extension-start",
+    type=str,
+    default=DEFAULT_EXTENSION_START,
+    help=(
+        "Start of a synthetic period beyond the source's coverage, over which "
+        "insolation is repeated. Empty (default) requires the source to span the "
+        "full window."
+    ),
+)
+@click.option(
+    "--repeat-source-start",
+    type=str,
+    default=DEFAULT_REPEAT_SOURCE_START,
+    help="Start of the insolation window repeated over the extension period.",
+)
+@click.option(
+    "--repeat-source-end",
+    type=str,
+    default=DEFAULT_REPEAT_SOURCE_END,
+    help="End of the insolation window repeated over the extension period.",
+)
 def main(
     input_data_file: str,
     output_data_file: str,
     ace2_era5_gcs_data: str,
     start_time: str,
     end_time: str,
+    extension_start: str,
+    repeat_source_start: str,
+    repeat_source_end: str,
 ):
     logging.basicConfig(level=logging.INFO)
     monthly_aimip_forcing = open_aimip_forcing_data(input_data_file)
@@ -218,11 +248,20 @@ def main(
         end_time,
     )
 
-    time_coord = get_time_coordinate(
-        existing_era5_forcing.time.drop_vars("time"),
-        extension_start="2023-01-01T00:00:00",
-        extension_end=end_time,
-    )
+    if extension_start:
+        time_coord = get_time_coordinate(
+            existing_era5_forcing.time.drop_vars("time"),
+            extension_start=extension_start,
+            extension_end=end_time,
+        )
+    else:
+        era5_end = existing_era5_forcing.time.values[-1]
+        if era5_end < np.datetime64(end_time):
+            raise ValueError(
+                f"Forcing source ends at {era5_end}, before --end-time {end_time}. "
+                "Set --extension-start, or an --end-time the source covers."
+            )
+        time_coord = existing_era5_forcing.time.drop_vars("time")
 
     logging.info("Interpolating AIMIP forcing data to ACE2-ERA5 time coordinate.")
     interpolated_aimip_forcing = monthly_aimip_forcing.interp(time=time_coord)
@@ -232,21 +271,23 @@ def main(
     ].where(sst_mask)
 
     logging.info("Merging interpolated AIMIP forcing with existing ERA5 forcing.")
-    repeated_era5_forcing_DSWRFtoa = get_repeated_insolation(
-        existing_era5_forcing.DSWRFtoa,
-        start_repeat="2023-01-01T00:00:00",
-        end_repeat=end_time,
-        source_start="2020-12-31T00:00:00",
-        source_end="2022-12-31T18:00:00",
-    )
-
-    era5_forcing_DSWRFtoa = xr.concat(
-        [
+    if extension_start:
+        repeated_era5_forcing_DSWRFtoa = get_repeated_insolation(
             existing_era5_forcing.DSWRFtoa,
-            repeated_era5_forcing_DSWRFtoa,
-        ],
-        dim="time",
-    )
+            start_repeat=extension_start,
+            end_repeat=end_time,
+            source_start=repeat_source_start,
+            source_end=repeat_source_end,
+        )
+        era5_forcing_DSWRFtoa = xr.concat(
+            [
+                existing_era5_forcing.DSWRFtoa,
+                repeated_era5_forcing_DSWRFtoa,
+            ],
+            dim="time",
+        )
+    else:
+        era5_forcing_DSWRFtoa = existing_era5_forcing.DSWRFtoa
 
     logging.info("Finalizing interpolated AIMIP forcing data.")
     interpolated_forcing = xr.merge(

scripts/aimip_forcing/prepend_first_timestep_forcing.py

@@ -69,7 +69,15 @@ def main(
     first_step = ds.sel(time=[input_timestamp]).assign_coords(
         time=[np.datetime64(output_timestamp)]
     )
-    ds_with_prepended = xr.concat([first_step, ds], dim="time")
+    # Concat defaults would broadcast variables without a time dimension along it,
+    # storing one constant per timestep.
+    ds_with_prepended = xr.concat(
+        [first_step, ds],
+        dim="time",
+        data_vars="minimal",
+        coords="minimal",
+        compat="override",
+    )
 
     logging.info("Setting chunking and sharding for output.")
     ds_with_prepended = clear_encoding(ds_with_prepended)

scripts/aimip_forcing/test_create_aimip_ic_datasets.py

@@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+import xarray as xr
+from click.testing import CliRunner
+from create_aimip_ic_datasets import (
+    NEAR_SURFACE_VARIABLES,
+    PROGNOSTIC_VARIABLES,
+    create_ic,
+    main,
+)
+
+TARGET = "1978-09-30T18:00:00"
+
+
+def _era5(times, names):
+    n_lat, n_lon = 4, 8
+    return xr.Dataset(
+        {
+            name: (
+                ("time", "latitude", "longitude"),
+                np.random.rand(len(times), n_lat, n_lon),
+            )
+            for name in names
+        },
+        coords={
+            "time": np.array(times, dtype="datetime64[ns]"),
+            "latitude": np.arange(n_lat, dtype=float),
+            "longitude": np.arange(n_lon, dtype=float),
+        },
+    )
+
+
+def test_create_ic_keeps_time_as_length_one_dimension():
+    """get_initial_condition requires shape (n_samples, [spatial dims]), so `time`
+    must survive selection as a dimension rather than collapsing to a scalar."""
+    era5 = _era5(["1978-09-29T00", "1978-09-30T00"], PROGNOSTIC_VARIABLES)
+    ic = create_ic(era5, "1978-09-29T00", np.datetime64(TARGET))
+    assert ic.sizes["time"] == 1
+    for name in PROGNOSTIC_VARIABLES:
+        assert ic[name].dims == ("time", "latitude", "longitude")
+
+
+def test_create_ic_restamps_time_to_target():
+    era5 = _era5(["1978-09-29T00", "1978-09-30T00"], PROGNOSTIC_VARIABLES)
+    ic = create_ic(era5, "1978-09-30T00", np.datetime64(TARGET))
+    assert ic.time.values[0] == np.datetime64(TARGET)
+
+
+def test_create_ic_selects_the_requested_timestamp():
+    era5 = _era5(["1978-09-29T00", "1978-09-30T00"], PROGNOSTIC_VARIABLES)
+    ic = create_ic(era5, "1978-09-30T00", np.datetime64(TARGET))
+    expected = era5[PROGNOSTIC_VARIABLES[0]].sel(time="1978-09-30T00").values
+    np.testing.assert_array_equal(ic[PROGNOSTIC_VARIABLES[0]].values[0], expected)
+
+
+@pytest.mark.parametrize(
+    "flag, expect_near_surface",
+    [([], False), (["--include-near-surface"], True)],
+)
+def test_near_surface_variables_included_only_on_request(
+    tmp_path, monkeypatch, flag, expect_near_surface
+):
+    era5 = _era5(["1978-09-29T00"], PROGNOSTIC_VARIABLES + NEAR_SURFACE_VARIABLES)
+    monkeypatch.setattr(xr, "open_zarr", lambda *a, **k: era5)
+
+    result = CliRunner().invoke(
+        main,
+        [str(tmp_path), "--ic-timestamp", "1978-09-29T00", *flag],
+    )
+    assert result.exit_code == 0, result.output
+
+    written = xr.load_dataset(tmp_path / "1978-09-30_IC0.nc")
+    for name in PROGNOSTIC_VARIABLES:
+        assert name in written
+    for name in NEAR_SURFACE_VARIABLES:
+        assert (name in written) is expect_near_surface