From 93e62d6e47c06e5e35613dce7c15790d8b01ac89 Mon Sep 17 00:00:00 2001
From: Brian Henn <brianhenn@allenai.org>
Date: Mon, 24 Aug 2026 23:04:27 +0000
Subject: [PATCH 1/5] add options to AIMIP IC workflow

---
 scripts/aimip_forcing/Makefile                |  6 +-
 scripts/aimip_forcing/README.md               | 22 ++++++
 .../aimip_forcing/create_aimip_ic_datasets.py | 25 +++++-
 .../test_create_aimip_ic_datasets.py          | 76 +++++++++++++++++++
 4 files changed, 125 insertions(+), 4 deletions(-)
 create mode 100644 scripts/aimip_forcing/test_create_aimip_ic_datasets.py

diff --git a/scripts/aimip_forcing/Makefile b/scripts/aimip_forcing/Makefile
index 25e039d13d..429ad69d4c 100644
--- a/scripts/aimip_forcing/Makefile
+++ b/scripts/aimip_forcing/Makefile
@@ -11,6 +11,10 @@ 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
+# Extra flags for create_aimip_ic_datasets.py. Set to --include-near-surface to add
+# TMP2m/Q2m/UGRD10m/VGRD10m for models that carry them as prognostic inputs; the
+# source zarr given by ERA5_GCS_DATA must provide them.
+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
@@ -56,7 +60,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
diff --git a/scripts/aimip_forcing/README.md b/scripts/aimip_forcing/README.md
index e604825ae1..cd4eb33b9b 100644
--- a/scripts/aimip_forcing/README.md
+++ b/scripts/aimip_forcing/README.md
@@ -25,6 +25,28 @@ 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 initial-condition files an ACE model is
+initialized from for AIMIP inference. Each IC member is a single timestamp selected from
+the 1-degree 8-layer ERA5 zarr and restamped to a common target time, so the members are
+an ensemble over initial state rather than over start date.
+
+```make create_aimip_ics```
+
+Relevant variables:
+
+- `ERA5_GCS_DATA` — source zarr. Must be the same ERA5 version the model was trained on.
+- `GCS_IC_PATH` — where the resulting `.nc` files are uploaded.
+- `IC_EXTRA_ARGS` — extra flags. Use `--include-near-surface` for models that carry
+  `TMP2m`/`Q2m`/`UGRD10m`/`VGRD10m` as prognostic inputs; the source zarr must have them.
+  `--target-timestamp` and `--ic-timestamp` set the restamped time and the source times,
+  and must exist in the source zarr.
+
+Note the ICs supply only *prognostic* variables. Forcings (`DSWRFtoa`, `HGTsfc`,
+`land_fraction`, `ocean_fraction`, `sea_ice_fraction`, `surface_temperature`) come from the
+forcing zarr at inference time, so they are deliberately absent here.
+
 ### 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:
diff --git a/scripts/aimip_forcing/create_aimip_ic_datasets.py b/scripts/aimip_forcing/create_aimip_ic_datasets.py
index 035154326f..6ae6dffbcc 100644
--- a/scripts/aimip_forcing/create_aimip_ic_datasets.py
+++ b/scripts/aimip_forcing/create_aimip_ic_datasets.py
@@ -26,6 +26,9 @@
     + [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(
@@ -33,8 +36,11 @@ def create_ic(
     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]
diff --git a/scripts/aimip_forcing/test_create_aimip_ic_datasets.py b/scripts/aimip_forcing/test_create_aimip_ic_datasets.py
new file mode 100644
index 0000000000..34ee9d3c25
--- /dev/null
+++ b/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

From 732f5d79ab929b8ef85016d17b6722b572a91dd3 Mon Sep 17 00:00:00 2001
From: Brian Henn <brianhenn@allenai.org>
Date: Tue, 25 Aug 2026 21:16:26 +0000
Subject: [PATCH 2/5] Support regenerating AIMIP forcing from a later ERA5
 build

HGTsfc, DSWRFtoa and land_fraction are read from an ACE ERA5 zarr and differ
between builds, so the forcing should be generated from the build the model was
trained on. Makes the source selectable via FORCING_EXTRA_ARGS.

The insolation extension is now opt-in via --extension-start rather than always
appending 2023 onward, since a source that already spans the window needs none.
Without it, a source that falls short errors instead of silently truncating.

Also makes OUTPUT_ZARR_NAME and MODIFIED_FORCING_NAME overridable so a regeneration
does not overwrite the existing datasets.
---
 scripts/aimip_forcing/Makefile                | 13 +--
 scripts/aimip_forcing/README.md               | 36 +++++----
 .../interpolate_aimip_forcing.py              | 79 ++++++++++++++-----
 .../test_interpolate_aimip_forcing.py         | 79 +++++++++++++++++++
 4 files changed, 168 insertions(+), 39 deletions(-)
 create mode 100644 scripts/aimip_forcing/test_interpolate_aimip_forcing.py

diff --git a/scripts/aimip_forcing/Makefile b/scripts/aimip_forcing/Makefile
index 429ad69d4c..9fdd33cbbb 100644
--- a/scripts/aimip_forcing/Makefile
+++ b/scripts/aimip_forcing/Makefile
@@ -6,18 +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
-# Extra flags for create_aimip_ic_datasets.py. Set to --include-near-surface to add
-# TMP2m/Q2m/UGRD10m/VGRD10m for models that carry them as prognostic inputs; the
-# source zarr given by ERA5_GCS_DATA must provide them.
+# 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)
 
@@ -49,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)
diff --git a/scripts/aimip_forcing/README.md b/scripts/aimip_forcing/README.md
index cd4eb33b9b..1a43d637c0 100644
--- a/scripts/aimip_forcing/README.md
+++ b/scripts/aimip_forcing/README.md
@@ -27,25 +27,33 @@ Note that the workflow is memory-intensive and was run on a high-memory (128GB)
 
 ### Generating AIMIP evaluation initial conditions
 
-`create_aimip_ic_datasets.py` builds the initial-condition files an ACE model is
-initialized from for AIMIP inference. Each IC member is a single timestamp selected from
-the 1-degree 8-layer ERA5 zarr and restamped to a common target time, so the members are
-an ensemble over initial state rather than over start date.
+`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```
 
-Relevant variables:
+- `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.
 
-- `ERA5_GCS_DATA` — source zarr. Must be the same ERA5 version the model was trained on.
-- `GCS_IC_PATH` — where the resulting `.nc` files are uploaded.
-- `IC_EXTRA_ARGS` — extra flags. Use `--include-near-surface` for models that carry
-  `TMP2m`/`Q2m`/`UGRD10m`/`VGRD10m` as prognostic inputs; the source zarr must have them.
-  `--target-timestamp` and `--ic-timestamp` set the restamped time and the source times,
-  and must exist in the source zarr.
+ICs carry only prognostic variables; forcings come from the forcing zarr at inference time.
 
-Note the ICs supply only *prognostic* variables. Forcings (`DSWRFtoa`, `HGTsfc`,
-`land_fraction`, `ocean_fraction`, `sea_ice_fraction`, `surface_temperature`) come from the
-forcing zarr at inference time, so they are deliberately absent here.
+### 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
 
diff --git a/scripts/aimip_forcing/interpolate_aimip_forcing.py b/scripts/aimip_forcing/interpolate_aimip_forcing.py
index 2b0e0877c8..fe2f77722c 100644
--- a/scripts/aimip_forcing/interpolate_aimip_forcing.py
+++ b/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(
diff --git a/scripts/aimip_forcing/test_interpolate_aimip_forcing.py b/scripts/aimip_forcing/test_interpolate_aimip_forcing.py
new file mode 100644
index 0000000000..347b45e509
--- /dev/null
+++ b/scripts/aimip_forcing/test_interpolate_aimip_forcing.py
@@ -0,0 +1,79 @@
+import interpolate_aimip_forcing as mod
+import numpy as np
+import xarray as xr
+from click.testing import CliRunner
+from interpolate_aimip_forcing import (
+    DEFAULT_EXTENSION_START,
+    SURFACE_TEMPERATURE_NAME,
+    main,
+)
+
+
+def _times(start, end, freq):
+    return xr.DataArray(
+        xr.date_range(start=start, end=end, freq=freq, use_cftime=False).values,
+        dims=["time"],
+        name="time",
+    )
+
+
+def _fake_era5(start, end, freq="6h"):
+    t = xr.date_range(start=start, end=end, freq=freq, use_cftime=False).values
+    shape = (t.size, 2, 2)
+    return xr.Dataset(
+        {
+            "DSWRFtoa": (("time", "latitude", "longitude"), np.zeros(shape)),
+            "HGTsfc": (("time", "latitude", "longitude"), np.zeros(shape)),
+            "land_fraction": (("time", "latitude", "longitude"), np.zeros(shape)),
+        },
+        coords={"time": t, "latitude": [0.0, 1.0], "longitude": [0.0, 1.0]},
+    )
+
+
+def test_disabled_extension_errors_when_source_falls_short(tmp_path, monkeypatch):
+    """Without an extension a short source would silently truncate the output."""
+    monthly = xr.Dataset(
+        {
+            SURFACE_TEMPERATURE_NAME: (
+                ("time", "latitude", "longitude"),
+                np.zeros((2, 2, 2)),
+            )
+        },
+        coords={
+            "time": xr.date_range(
+                "1979-01-01", periods=2, freq="MS", use_cftime=False
+            ).values,
+            "latitude": [0.0, 1.0],
+            "longitude": [0.0, 1.0],
+        },
+    )
+    monkeypatch.setattr(mod, "open_aimip_forcing_data", lambda *a, **k: monthly)
+    monkeypatch.setattr(
+        mod,
+        "get_existing_era5_forcing",
+        lambda *a, **k: _fake_era5("1979-01-01", "1979-01-02"),
+    )
+    src = tmp_path / "in.nc"
+    src.write_bytes(b"")
+
+    result = CliRunner().invoke(
+        main,
+        [
+            str(src),
+            str(tmp_path / "out.zarr"),
+            "--start-time",
+            "1979-01-01T00:00:00",
+            "--end-time",
+            "1979-06-01T00:00:00",
+            "--extension-start",
+            "",
+        ],
+    )
+    assert result.exit_code != 0
+    assert isinstance(result.exception, ValueError)
+    assert "before --end-time" in str(result.exception)
+
+
+def test_extension_is_off_by_default():
+    """Extension is opt-in; sources are expected to span the full window."""
+    assert DEFAULT_EXTENSION_START == ""

From 13983d3c5a32a3a1a012579db5a29dc4f012df1d Mon Sep 17 00:00:00 2001
From: Brian Henn <brianhenn@allenai.org>
Date: Tue, 25 Aug 2026 21:26:46 +0000
Subject: [PATCH 3/5] Preserve static and scalar variables when prepending a
 forcing timestep

xr.concat's defaults broadcast variables that have no time dimension across every
timestep, inflating the store and slowing the data loader, which handles static
variables natively.

Skip 0-d variables when setting chunks and shards; zarr rejects empty tuples,
which the preserved scalars now hit.
---
 scripts/aimip_forcing/encoding.py             |  4 ++
 .../prepend_first_timestep_forcing.py         | 10 ++++-
 scripts/aimip_forcing/test_encoding.py        | 40 +++++++++++++++++
 .../test_prepend_first_timestep_forcing.py    | 45 +++++++++++++++++++
 4 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 scripts/aimip_forcing/test_encoding.py
 create mode 100644 scripts/aimip_forcing/test_prepend_first_timestep_forcing.py

diff --git a/scripts/aimip_forcing/encoding.py b/scripts/aimip_forcing/encoding.py
index 2a25a8952b..234d229673 100644
--- a/scripts/aimip_forcing/encoding.py
+++ b/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 = {}
diff --git a/scripts/aimip_forcing/prepend_first_timestep_forcing.py b/scripts/aimip_forcing/prepend_first_timestep_forcing.py
index c009b88217..dab3cdf3de 100644
--- a/scripts/aimip_forcing/prepend_first_timestep_forcing.py
+++ b/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)
diff --git a/scripts/aimip_forcing/test_encoding.py b/scripts/aimip_forcing/test_encoding.py
new file mode 100644
index 0000000000..1f86857e43
--- /dev/null
+++ b/scripts/aimip_forcing/test_encoding.py
@@ -0,0 +1,40 @@
+import numpy as np
+import xarray as xr
+from encoding import clear_encoding, set_shards_chunks
+
+
+def _ds(n_time=4):
+    return xr.Dataset(
+        {
+            "tvar": (("time", "latitude", "longitude"), np.random.rand(n_time, 2, 2)),
+            "static2d": (("latitude", "longitude"), np.ones((2, 2))),
+            "scalar": ((), np.float64(3.0)),
+        },
+        coords={
+            "time": xr.date_range(
+                "1978-10-01", periods=n_time, freq="6h", use_cftime=False
+            ).values,
+            "latitude": [0.0, 1.0],
+            "longitude": [0.0, 1.0],
+        },
+    )
+
+
+def test_scalars_get_no_chunk_or_shard_encoding():
+    """zarr rejects empty chunk/shard tuples, so 0-d variables must be left alone."""
+    out = set_shards_chunks(clear_encoding(_ds()))
+    assert out["scalar"].encoding.get("chunks") is None
+    assert out["scalar"].encoding.get("shards") is None
+
+
+def test_every_variable_is_writable(tmp_path):
+    out = set_shards_chunks(clear_encoding(_ds()))
+    for name in out.data_vars:
+        out[[name]].to_zarr(tmp_path / f"{name}.zarr", mode="w")
+
+
+def test_dask_chunks_match_shards():
+    """xarray needs dask chunks equal to the zarr shards to write sharded stores."""
+    out = set_shards_chunks(clear_encoding(_ds(n_time=40)))
+    assert out["tvar"].encoding["shards"] == tuple(c[0] for c in out["tvar"].chunks)
+    assert out["tvar"].encoding["chunks"] == (1, 2, 2)
diff --git a/scripts/aimip_forcing/test_prepend_first_timestep_forcing.py b/scripts/aimip_forcing/test_prepend_first_timestep_forcing.py
new file mode 100644
index 0000000000..cf160fd5f5
--- /dev/null
+++ b/scripts/aimip_forcing/test_prepend_first_timestep_forcing.py
@@ -0,0 +1,45 @@
+import numpy as np
+import xarray as xr
+from click.testing import CliRunner
+from prepend_first_timestep_forcing import main
+
+
+def _forcing():
+    t = xr.date_range("1978-10-01", periods=4, freq="6h", use_cftime=False).values
+    return xr.Dataset(
+        {
+            "surface_temperature": (
+                ("time", "latitude", "longitude"),
+                np.random.rand(4, 2, 2),
+            ),
+            "HGTsfc": (("latitude", "longitude"), np.ones((2, 2))),
+            "ak_0": ((), np.float64(3.0)),
+        },
+        coords={"time": t, "latitude": [0.0, 1.0], "longitude": [0.0, 1.0]},
+    )
+
+
+def test_statics_are_not_broadcast_along_time(tmp_path, monkeypatch):
+    """Broadcasting statics stores one constant per timestep and slows the loader."""
+    monkeypatch.setattr(xr, "open_zarr", lambda *a, **k: _forcing())
+    out = tmp_path / "out.zarr"
+    result = CliRunner().invoke(
+        main,
+        [
+            str(out),
+            "--input-forcing-path",
+            "unused",
+            "--input-timestamp",
+            "1978-10-01T00:00:00",
+            "--output-timestamp",
+            "1978-09-30T18:00:00",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    ds = xr.open_dataset(out, engine="zarr")
+    assert ds["HGTsfc"].dims == ("latitude", "longitude")
+    assert ds["ak_0"].dims == ()
+    assert ds["surface_temperature"].dims == ("time", "latitude", "longitude")
+    assert ds.sizes["time"] == 5
+    assert ds.time.values[0] == np.datetime64("1978-09-30T18:00:00")

From 97491f5c7d454c120c2cf70de000349af286de60 Mon Sep 17 00:00:00 2001
From: Brian Henn <brianhenn@allenai.org>
Date: Tue, 25 Aug 2026 22:09:27 +0000
Subject: [PATCH 4/5] only import dask progress bar if installed so tests can
 run in CI

---
 scripts/aimip_forcing/create_aimip_ic_datasets.py       | 2 +-
 scripts/aimip_forcing/create_public_aimip_forcing.py    | 2 +-
 scripts/aimip_forcing/interpolate_aimip_forcing.py      | 2 +-
 scripts/aimip_forcing/prepend_first_timestep_forcing.py | 2 +-
 scripts/aimip_forcing/progress.py                       | 8 ++++++++
 scripts/aimip_forcing/regrid_aimip_forcing.py           | 2 +-
 6 files changed, 13 insertions(+), 5 deletions(-)
 create mode 100644 scripts/aimip_forcing/progress.py

diff --git a/scripts/aimip_forcing/create_aimip_ic_datasets.py b/scripts/aimip_forcing/create_aimip_ic_datasets.py
index 6ae6dffbcc..f4aab765b4 100644
--- a/scripts/aimip_forcing/create_aimip_ic_datasets.py
+++ b/scripts/aimip_forcing/create_aimip_ic_datasets.py
@@ -5,7 +5,7 @@
 import click
 import numpy as np
 import xarray as xr
-from dask.diagnostics import ProgressBar
+from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/2025-10-16-make-AIMIP-evaluation-IC-datasets-v3.ipynb
diff --git a/scripts/aimip_forcing/create_public_aimip_forcing.py b/scripts/aimip_forcing/create_public_aimip_forcing.py
index 5e9b4f20d1..4e1f0c071b 100644
--- a/scripts/aimip_forcing/create_public_aimip_forcing.py
+++ b/scripts/aimip_forcing/create_public_aimip_forcing.py
@@ -4,8 +4,8 @@
 import cftime
 import click
 import xarray as xr
-from dask.diagnostics import ProgressBar
 from encoding import clear_encoding
+from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/troya/2025-08-06-AIMP-ERA5/2025-08-05-ARCO-ERA5-monthly-average-forcing-AIMIP.ipynb
diff --git a/scripts/aimip_forcing/interpolate_aimip_forcing.py b/scripts/aimip_forcing/interpolate_aimip_forcing.py
index fe2f77722c..1bec4f886a 100644
--- a/scripts/aimip_forcing/interpolate_aimip_forcing.py
+++ b/scripts/aimip_forcing/interpolate_aimip_forcing.py
@@ -4,8 +4,8 @@
 import click
 import numpy as np
 import xarray as xr
-from dask.diagnostics import ProgressBar
 from encoding import clear_encoding, set_shards_chunks
+from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/regrid/2025-08-26-ACE2-ERA5-forcing-AIMIP.ipynb
diff --git a/scripts/aimip_forcing/prepend_first_timestep_forcing.py b/scripts/aimip_forcing/prepend_first_timestep_forcing.py
index dab3cdf3de..3dce07eb7f 100644
--- a/scripts/aimip_forcing/prepend_first_timestep_forcing.py
+++ b/scripts/aimip_forcing/prepend_first_timestep_forcing.py
@@ -3,8 +3,8 @@
 import click
 import numpy as np
 import xarray as xr
-from dask.diagnostics import ProgressBar
 from encoding import clear_encoding, set_shards_chunks
+from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/2025-11-24-repeat-first-timestep-AIMIP-forcing.ipynb
diff --git a/scripts/aimip_forcing/progress.py b/scripts/aimip_forcing/progress.py
new file mode 100644
index 0000000000..b6c54c9c8c
--- /dev/null
+++ b/scripts/aimip_forcing/progress.py
@@ -0,0 +1,8 @@
+from contextlib import nullcontext
+
+try:
+    from dask.diagnostics import ProgressBar
+except ImportError:
+    # dask only drives progress output here, and is absent from the repo-wide test
+    # environment, so degrade to a no-op rather than making it a hard import.
+    ProgressBar = nullcontext
diff --git a/scripts/aimip_forcing/regrid_aimip_forcing.py b/scripts/aimip_forcing/regrid_aimip_forcing.py
index 1b6143fd64..d4a121b9a8 100644
--- a/scripts/aimip_forcing/regrid_aimip_forcing.py
+++ b/scripts/aimip_forcing/regrid_aimip_forcing.py
@@ -4,7 +4,7 @@
 import numpy as np
 import xarray as xr
 import xesmf as xe
-from dask.diagnostics import ProgressBar
+from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/regrid/2025-08-20-regrid-AIMIP-forcing.ipynb

From f70da9b24ec96564fb03c112f378cbef4627c58d Mon Sep 17 00:00:00 2001
From: Brian Henn <brianhenn@allenai.org>
Date: Tue, 25 Aug 2026 22:53:32 +0000
Subject: [PATCH 5/5] don't run tests on CI due to env differences

---
 scripts/aimip_forcing/conftest.py                       | 7 +++++++
 scripts/aimip_forcing/create_aimip_ic_datasets.py       | 2 +-
 scripts/aimip_forcing/create_public_aimip_forcing.py    | 2 +-
 scripts/aimip_forcing/interpolate_aimip_forcing.py      | 2 +-
 scripts/aimip_forcing/prepend_first_timestep_forcing.py | 2 +-
 scripts/aimip_forcing/progress.py                       | 8 --------
 scripts/aimip_forcing/regrid_aimip_forcing.py           | 2 +-
 7 files changed, 12 insertions(+), 13 deletions(-)
 create mode 100644 scripts/aimip_forcing/conftest.py
 delete mode 100644 scripts/aimip_forcing/progress.py

diff --git a/scripts/aimip_forcing/conftest.py b/scripts/aimip_forcing/conftest.py
new file mode 100644
index 0000000000..6d3260747a
--- /dev/null
+++ b/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"]
diff --git a/scripts/aimip_forcing/create_aimip_ic_datasets.py b/scripts/aimip_forcing/create_aimip_ic_datasets.py
index f4aab765b4..6ae6dffbcc 100644
--- a/scripts/aimip_forcing/create_aimip_ic_datasets.py
+++ b/scripts/aimip_forcing/create_aimip_ic_datasets.py
@@ -5,7 +5,7 @@
 import click
 import numpy as np
 import xarray as xr
-from progress import ProgressBar
+from dask.diagnostics import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/2025-10-16-make-AIMIP-evaluation-IC-datasets-v3.ipynb
diff --git a/scripts/aimip_forcing/create_public_aimip_forcing.py b/scripts/aimip_forcing/create_public_aimip_forcing.py
index 4e1f0c071b..5e9b4f20d1 100644
--- a/scripts/aimip_forcing/create_public_aimip_forcing.py
+++ b/scripts/aimip_forcing/create_public_aimip_forcing.py
@@ -4,8 +4,8 @@
 import cftime
 import click
 import xarray as xr
+from dask.diagnostics import ProgressBar
 from encoding import clear_encoding
-from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/troya/2025-08-06-AIMP-ERA5/2025-08-05-ARCO-ERA5-monthly-average-forcing-AIMIP.ipynb
diff --git a/scripts/aimip_forcing/interpolate_aimip_forcing.py b/scripts/aimip_forcing/interpolate_aimip_forcing.py
index 1bec4f886a..fe2f77722c 100644
--- a/scripts/aimip_forcing/interpolate_aimip_forcing.py
+++ b/scripts/aimip_forcing/interpolate_aimip_forcing.py
@@ -4,8 +4,8 @@
 import click
 import numpy as np
 import xarray as xr
+from dask.diagnostics import ProgressBar
 from encoding import clear_encoding, set_shards_chunks
-from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/regrid/2025-08-26-ACE2-ERA5-forcing-AIMIP.ipynb
diff --git a/scripts/aimip_forcing/prepend_first_timestep_forcing.py b/scripts/aimip_forcing/prepend_first_timestep_forcing.py
index 3dce07eb7f..dab3cdf3de 100644
--- a/scripts/aimip_forcing/prepend_first_timestep_forcing.py
+++ b/scripts/aimip_forcing/prepend_first_timestep_forcing.py
@@ -3,8 +3,8 @@
 import click
 import numpy as np
 import xarray as xr
+from dask.diagnostics import ProgressBar
 from encoding import clear_encoding, set_shards_chunks
-from progress import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/2025-11-24-repeat-first-timestep-AIMIP-forcing.ipynb
diff --git a/scripts/aimip_forcing/progress.py b/scripts/aimip_forcing/progress.py
deleted file mode 100644
index b6c54c9c8c..0000000000
--- a/scripts/aimip_forcing/progress.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from contextlib import nullcontext
-
-try:
-    from dask.diagnostics import ProgressBar
-except ImportError:
-    # dask only drives progress output here, and is absent from the repo-wide test
-    # environment, so degrade to a no-op rather than making it a hard import.
-    ProgressBar = nullcontext
diff --git a/scripts/aimip_forcing/regrid_aimip_forcing.py b/scripts/aimip_forcing/regrid_aimip_forcing.py
index d4a121b9a8..1b6143fd64 100644
--- a/scripts/aimip_forcing/regrid_aimip_forcing.py
+++ b/scripts/aimip_forcing/regrid_aimip_forcing.py
@@ -4,7 +4,7 @@
 import numpy as np
 import xarray as xr
 import xesmf as xe
-from progress import ProgressBar
+from dask.diagnostics import ProgressBar
 
 # this script is based on the notebook at
 # https://github.com/ai2cm/explore2/blob/main/brianh/2025-06-18-ERA5-SHiELD-AMIP-forcing/regrid/2025-08-20-regrid-AIMIP-forcing.ipynb