From 3516108caa4031e809351ff615d0a0acc019339a Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Wed, 4 Feb 2026 16:52:52 +0000
Subject: [PATCH 01/13] add time coarsen script

---
 fme/ace/testing/__init__.py               |  1 +
 scripts/time_coarsen/test_time_coarsen.py | 72 +++++++++++++++++++
 scripts/time_coarsen/time_coarsen.py      | 87 +++++++++++++++++++++++
 3 files changed, 160 insertions(+)
 create mode 100644 scripts/time_coarsen/test_time_coarsen.py
 create mode 100644 scripts/time_coarsen/time_coarsen.py

diff --git a/fme/ace/testing/__init__.py b/fme/ace/testing/__init__.py
index 034a41a7e..717162fc2 100644
--- a/fme/ace/testing/__init__.py
+++ b/fme/ace/testing/__init__.py
@@ -4,6 +4,7 @@
     FV3GFSData,
     MonthlyReferenceData,
     StatsData,
+    get_nd_dataset,
     save_nd_netcdf,
     save_scalar_netcdf,
 )
diff --git a/scripts/time_coarsen/test_time_coarsen.py b/scripts/time_coarsen/test_time_coarsen.py
new file mode 100644
index 000000000..a1450b7e6
--- /dev/null
+++ b/scripts/time_coarsen/test_time_coarsen.py
@@ -0,0 +1,72 @@
+import numpy as np
+import xarray as xr
+from time_coarsen import Config, PathPair, main
+
+from fme.ace.testing import DimSize, DimSizes, get_nd_dataset
+
+
+def test_time_coarsen() -> None:
+    n_input_times = 5
+    nz_interface = 3
+    # Create a small but non-trivial dataset (coords + attrs)
+    ds = get_nd_dataset(
+        dim_sizes=DimSizes(
+            n_time=n_input_times,
+            horizontal=[
+                DimSize(name="lat", size=4),
+                DimSize(name="lon", size=8),
+            ],
+            nz_interface=nz_interface,
+        ),
+        variable_names=["temp", "temp_tendency", "flag"],
+        timestep_days=1.0,
+        include_vertical_coordinate=True,
+    )
+    ds["temp"].attrs["units"] = "K"
+    ds["temp_tendency"].attrs["units"] = "K/day"
+    constant_names = []
+    for name in ds.data_vars:
+        if name.startswith("ak_") or name.startswith("bk_"):
+            constant_names.append(name)
+    assert len(constant_names) == 2 * nz_interface  # sanity check
+    input_path = "memory://test-time-coarsen/dataset.zarr"
+    output_path = "memory://test-time-coarsen/dataset_coarsened.zarr"
+
+    # Write to an in-memory filesystem (xarray recognizes the memory:// protocol)
+    ds.to_zarr(input_path)
+    ds = xr.open_zarr(input_path)  # for comparison later, use fresh read
+
+    config = Config(
+        paths=[PathPair(input=input_path, output=output_path)],
+        coarsen_factor=2,
+        snapshot_names=["temp"],
+        window_names=["temp_tendency"],
+    )
+    main(config)
+
+    # Read back the coarsened dataset
+    ds_coarsened = xr.open_zarr(output_path)
+    # Note on each timestep, we have the tendencies which led to the current
+    # snapshot alongside the snapshot itself. This means the first snapshot
+    # of the coarsened dataset is not the first snapshot of the input dataset.
+    assert ds_coarsened.dims["time"] == n_input_times // config.coarsen_factor
+    expected_snapshot_slice = slice(
+        config.coarsen_factor - 1, None, config.coarsen_factor
+    )
+    np.testing.assert_array_equal(
+        ds_coarsened["time"].values,
+        ds["time"].isel(time=expected_snapshot_slice).values,
+    )
+    np.testing.assert_array_equal(
+        ds_coarsened["temp"].values,
+        ds["temp"].isel(time=expected_snapshot_slice).values,
+    )
+    np.testing.assert_array_equal(
+        ds_coarsened["temp_tendency"].isel(time=0).values,
+        ds["temp_tendency"]
+        .isel(time=slice(0, config.coarsen_factor))
+        .mean("time")
+        .values,
+    )
+    for name in constant_names:
+        np.testing.assert_array_equal(ds_coarsened[name].values, ds[name].values)
diff --git a/scripts/time_coarsen/time_coarsen.py b/scripts/time_coarsen/time_coarsen.py
new file mode 100644
index 000000000..33ae50674
--- /dev/null
+++ b/scripts/time_coarsen/time_coarsen.py
@@ -0,0 +1,87 @@
+import argparse
+import dataclasses
+
+import dacite
+import xarray as xr
+import yaml
+
+
+@dataclasses.dataclass
+class PathPair:
+    input: str
+    output: str
+
+
+@dataclasses.dataclass
+class Config:
+    """
+    Configuration for time coarsening of a dataset.
+
+    Attributes:
+        input_path: Path to the input dataset.
+        output_path: Path to save the coarsened dataset as a zarr store.
+        coarsen_factor: Factor by which to coarsen the time dimension.
+        snapshot_names: List of snapshot variable names to coarsen. These will be
+            coarsened by skipping each coarsen_factor times.
+        window_names: List of window variable names to coarsen. These will be
+            coarsened by averaging over each coarsen_factor times.
+        constant_prefixes: List of prefixes for constant data variables to copy without
+            modification. Raises an exception if any of these have a "time" dimension.
+    """
+
+    paths: list[PathPair]
+    coarsen_factor: int
+    snapshot_names: list[str]
+    window_names: list[str]
+    constant_prefixes: list[str] = dataclasses.field(
+        default_factory=lambda: ["ak_", "bk_"]
+    )
+
+
+def main(config: Config):
+    for paths in config.paths:
+        process_path_pair(paths, config)
+
+
+def process_path_pair(pair: PathPair, config: Config):
+    ds = xr.open_dataset(pair.input)
+    constant_names = [
+        name
+        for name in ds.data_vars
+        if any(name.startswith(prefix) for prefix in config.constant_prefixes)
+    ]
+    assert set(constant_names).intersection(set(config.snapshot_names)) == set()
+    assert set(constant_names).intersection(set(config.window_names)) == set()
+    for name in constant_names:
+        if "time" in ds[name].dims:
+            raise ValueError(
+                f"Constant data variable {name} has a 'time' dimension, "
+                "which is not allowed."
+            )
+    ds_constants = ds[constant_names]
+    ds_snapshot = ds[config.snapshot_names].isel(
+        time=slice(config.coarsen_factor - 1, None, config.coarsen_factor)
+    )
+    ds_window = (
+        ds[config.window_names]
+        .coarsen(time=config.coarsen_factor, boundary="trim")
+        .mean()
+        .drop("time")
+    )  # use time of snapshots
+    # ds_coarsened = ds_snapshot#xr.merge([ds_snapshot, ds_window, ds_coords])
+    ds_coarsened = xr.merge([ds_snapshot, ds_window, ds_constants])
+    ds_coarsened.to_zarr(pair.output)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Time Coarsen Script")
+    parser.add_argument(
+        "config_yaml", type=str, required=True, help="Path to configuration yaml file"
+    )
+    args = parser.parse_args()
+    with open(args.config_yaml, "r") as f:
+        config_dict = yaml.safe_load(f)
+    config = dacite.from_dict(
+        data_class=Config, data=config_dict, config=dacite.Config(strict=True)
+    )
+    main(config)

From 1d46c9f2e867bef28039361c50ea2e3164ea4391 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Wed, 4 Feb 2026 16:54:15 +0000
Subject: [PATCH 02/13] add time coarsen submission script and yaml

---
 scripts/time_coarsen/c96-shield.yaml | 93 ++++++++++++++++++++++++++++
 scripts/time_coarsen/run.sh          | 45 ++++++++++++++
 2 files changed, 138 insertions(+)
 create mode 100644 scripts/time_coarsen/c96-shield.yaml
 create mode 100644 scripts/time_coarsen/run.sh

diff --git a/scripts/time_coarsen/c96-shield.yaml b/scripts/time_coarsen/c96-shield.yaml
new file mode 100644
index 000000000..2ab3c4f66
--- /dev/null
+++ b/scripts/time_coarsen/c96-shield.yaml
@@ -0,0 +1,93 @@
+paths:
+  - input: /climate-default/2026-01-28-vertically-resolved-c96-1deg-shield-amip-ensemble-dataset/ic_0001.zarr
+    output: /climate-default/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset/ic_0001.zarr
+  - input: /climate-default/2026-01-28-vertically-resolved-c96-1deg-shield-amip-ensemble-dataset/ic_0002.zarr
+    output: /climate-default/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset/ic_0002.zarr
+coarsen_factor: 4
+snapshot_names:
+  - PRESsfc
+  - Q2m
+  - RH200
+  - RH500
+  - RH800
+  - TMP200
+  - TMP2m
+  - TMP500
+  - TMP850
+  - UGRD1000
+  - UGRD10m
+  - UGRD200
+  - UGRD500
+  - UGRD850
+  - VGRD1000
+  - VGRD10m
+  - VGRD200
+  - VGRD500
+  - VGRD850
+  - air_temperature_0
+  - air_temperature_1
+  - air_temperature_2
+  - air_temperature_3
+  - air_temperature_4
+  - air_temperature_5
+  - air_temperature_6
+  - air_temperature_7
+  - eastward_wind_0
+  - eastward_wind_1
+  - eastward_wind_2
+  - eastward_wind_3
+  - eastward_wind_4
+  - eastward_wind_5
+  - eastward_wind_6
+  - eastward_wind_7
+  - global_mean_co2
+  - h1000
+  - h50
+  - h500
+  - h850
+  - land_fraction
+  - northward_wind_0
+  - northward_wind_1
+  - northward_wind_2
+  - northward_wind_3
+  - northward_wind_4
+  - northward_wind_5
+  - northward_wind_6
+  - northward_wind_7
+  - ocean_fraction
+  - sea_ice_fraction
+  - snow_cover_fraction
+  - soil_moisture_0
+  - soil_moisture_1
+  - soil_moisture_2
+  - soil_moisture_3
+  - specific_total_water_0
+  - specific_total_water_1
+  - specific_total_water_2
+  - specific_total_water_3
+  - specific_total_water_4
+  - specific_total_water_5
+  - specific_total_water_6
+  - specific_total_water_7
+  - surface_temperature
+  - total_water_path
+window_names:
+  - DLWRFsfc
+  - DSWRFsfc
+  - DSWRFtoa
+  - LHTFLsfc
+  - PRATEsfc
+  - SHTFLsfc
+  - ULWRFsfc
+  - ULWRFtoa
+  - USWRFsfc
+  - USWRFtoa
+  - tendency_of_total_water_path
+  - tendency_of_total_water_path_due_to_advection
+  - total_frozen_precipitation_rate
+constant_prefixes:
+  - ak_
+  - bk_
+  - HGTsfc
+  - grid_xt
+  - grid_yt
diff --git a/scripts/time_coarsen/run.sh b/scripts/time_coarsen/run.sh
new file mode 100644
index 000000000..bd8690a97
--- /dev/null
+++ b/scripts/time_coarsen/run.sh
@@ -0,0 +1,45 @@
+
+#!/bin/bash
+
+set -e
+
+SCRIPT_PATH=$(git rev-parse --show-prefix)  # relative to the root of the repository
+REPO_ROOT=$(git rev-parse --show-toplevel)
+
+cd "$REPO_ROOT"
+
+run_coarsen() {
+  local config_filename="$1"
+  local job_name="$2"
+  local CONFIG_PATH="$SCRIPT_PATH/$config_filename"
+
+  # Extract additional args from config header
+  local extra_args=()
+  while IFS= read -r line; do
+    [[ "$line" =~ ^#\ arg:\ (.*) ]] && extra_args+=(${BASH_REMATCH[1]})
+  done < "$CONFIG_PATH"
+
+  gantry run \
+    --name "$job_name" \
+    --description 'Run ACE training' \
+    --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \
+    --workspace ai2/ace \
+    --priority high \
+    --non-preemptible \
+    --cluster ai2/jupiter \
+    --cluster ai2/ceres \
+    --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \
+    --dataset-secret google-credentials:/tmp/google_application_credentials.json \
+    --gpus 0 \
+    --shared-memory 400GiB \
+    --weka climate-default:/climate-default \
+    --budget ai2/climate \
+    --system-python \
+    --install "pip install --no-deps ." \
+    "${extra_args[@]}" \
+    -- python $SCRIPT_PATH/time_coarsen.py "$CONFIG_PATH"
+}
+
+base_name="time-coarsen"
+
+run_coarsen "c96-shield.yaml" "$base_name-c96-shield"

From f21a46e83bd01df23d44b84f05cebcbc19ac0545 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 17:05:05 +0000
Subject: [PATCH 03/13] change workspace and priority

---
 scripts/time_coarsen/run.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/time_coarsen/run.sh b/scripts/time_coarsen/run.sh
index bd8690a97..8a32a2e36 100644
--- a/scripts/time_coarsen/run.sh
+++ b/scripts/time_coarsen/run.sh
@@ -23,8 +23,8 @@ run_coarsen() {
     --name "$job_name" \
     --description 'Run ACE training' \
     --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \
-    --workspace ai2/ace \
-    --priority high \
+    --workspace ai2/climate-titan \
+    --priority urgent \
     --non-preemptible \
     --cluster ai2/jupiter \
     --cluster ai2/ceres \

From 29d39836c2583113c236a8a43a782b7dd26338da Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 17:05:23 +0000
Subject: [PATCH 04/13] make run.sh runnable

---
 scripts/time_coarsen/run.sh | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 mode change 100644 => 100755 scripts/time_coarsen/run.sh

diff --git a/scripts/time_coarsen/run.sh b/scripts/time_coarsen/run.sh
old mode 100644
new mode 100755

From d0691359095604993f3f0923993b5b7d1631b795 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 17:05:46 +0000
Subject: [PATCH 05/13] fix not-preemptible flag

---
 scripts/time_coarsen/run.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/time_coarsen/run.sh b/scripts/time_coarsen/run.sh
index 8a32a2e36..aae3b11d8 100755
--- a/scripts/time_coarsen/run.sh
+++ b/scripts/time_coarsen/run.sh
@@ -25,7 +25,7 @@ run_coarsen() {
     --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \
     --workspace ai2/climate-titan \
     --priority urgent \
-    --non-preemptible \
+    --not-preemptible \
     --cluster ai2/jupiter \
     --cluster ai2/ceres \
     --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \

From e68a5c5dc7ff0f22511265725e01467e550179d0 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 17:09:51 +0000
Subject: [PATCH 06/13] remove invalid positional argument

---
 scripts/time_coarsen/time_coarsen.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/scripts/time_coarsen/time_coarsen.py b/scripts/time_coarsen/time_coarsen.py
index 33ae50674..4d468e7d0 100644
--- a/scripts/time_coarsen/time_coarsen.py
+++ b/scripts/time_coarsen/time_coarsen.py
@@ -75,9 +75,7 @@ def process_path_pair(pair: PathPair, config: Config):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Time Coarsen Script")
-    parser.add_argument(
-        "config_yaml", type=str, required=True, help="Path to configuration yaml file"
-    )
+    parser.add_argument("config_yaml", type=str, help="Path to configuration yaml file")
     args = parser.parse_args()
     with open(args.config_yaml, "r") as f:
         config_dict = yaml.safe_load(f)

From a16c12435e1e7b3f3bcc5306580a72b7f9b8a4d3 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 17:20:42 +0000
Subject: [PATCH 07/13] fix snapshot name

---
 scripts/time_coarsen/c96-shield.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/time_coarsen/c96-shield.yaml b/scripts/time_coarsen/c96-shield.yaml
index 2ab3c4f66..94843e7dd 100644
--- a/scripts/time_coarsen/c96-shield.yaml
+++ b/scripts/time_coarsen/c96-shield.yaml
@@ -9,7 +9,7 @@ snapshot_names:
   - Q2m
   - RH200
   - RH500
-  - RH800
+  - RH850
   - TMP200
   - TMP2m
   - TMP500

From 23842e62e733f00f65762ff6f500be71a85f70b1 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:16:43 +0000
Subject: [PATCH 08/13] add script to generate stats

---
 scripts/time_coarsen/c96-shield-stats.yaml | 12 ++++++
 scripts/time_coarsen/run_stats.sh          | 46 ++++++++++++++++++++++
 scripts/time_coarsen/time_coarsen.py       |  6 +++
 3 files changed, 64 insertions(+)
 create mode 100644 scripts/time_coarsen/c96-shield-stats.yaml
 create mode 100644 scripts/time_coarsen/run_stats.sh

diff --git a/scripts/time_coarsen/c96-shield-stats.yaml b/scripts/time_coarsen/c96-shield-stats.yaml
new file mode 100644
index 000000000..c8fdbaf88
--- /dev/null
+++ b/scripts/time_coarsen/c96-shield-stats.yaml
@@ -0,0 +1,12 @@
+runs:
+  ic_0001: gs://vcm-ml-raw-flexible-retention/2024-06-29-C96-SHiELD-AMIP/regridded-zarrs/gaussian_grid_180_by_360/ic_0001
+  ic_0002: gs://vcm-ml-raw-flexible-retention/2024-06-29-C96-SHiELD-AMIP/regridded-zarrs/gaussian_grid_180_by_360/ic_0002
+data_output_directory: gs://vcm-ml-intermediate/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset
+stats:
+  output_directory: gs://vcm-ml-intermediate/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
+  beaker_dataset: 2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
+  start_date: "1940-01-01"
+  end_date: "2021-12-31"
+  data_type: FV3GFS
+  exclude_runs:
+    - "ic_0002"
diff --git a/scripts/time_coarsen/run_stats.sh b/scripts/time_coarsen/run_stats.sh
new file mode 100644
index 000000000..54a08a7f4
--- /dev/null
+++ b/scripts/time_coarsen/run_stats.sh
@@ -0,0 +1,46 @@
+
+#!/bin/bash
+
+set -e
+
+SCRIPT_PATH=$(git rev-parse --show-prefix)  # relative to the root of the repository
+REPO_ROOT=$(git rev-parse --show-toplevel)
+
+cd "$REPO_ROOT"
+
+run_stats() {
+  local config_filename="$1"
+  local job_name="$2"
+  local CONFIG_PATH="$SCRIPT_PATH/$config_filename"
+
+  # Extract additional args from config header
+  local extra_args=()
+  while IFS= read -r line; do
+    [[ "$line" =~ ^#\ arg:\ (.*) ]] && extra_args+=(${BASH_REMATCH[1]})
+  done < "$CONFIG_PATH"
+
+  gantry run \
+    --name "$job_name" \
+    --description 'Run ACE training' \
+    --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \
+    --workspace ai2/climate-titan \
+    --priority urgent \
+    --not-preemptible \
+    --cluster ai2/jupiter \
+    --cluster ai2/ceres \
+    --env GOOGLE_APPLICATION_CREDENTIALS=/tmp/google_application_credentials.json \
+    --dataset-secret google-credentials:/tmp/google_application_credentials.json \
+    --gpus 0 \
+    --shared-memory 400GiB \
+    --weka climate-default:/climate-default \
+    --budget ai2/climate \
+    --system-python \
+    --install "pip install --no-deps ." \
+    "${extra_args[@]}" \
+    -- $SCRIPT_PATH/../data_process/compute_stats.sh --config $CONFIG_PATH
+}
+
+base_name="time-coarsen-stats"
+
+run_stats "c96-shield-stats.yaml" "$base_name-c96-shield"
+
diff --git a/scripts/time_coarsen/time_coarsen.py b/scripts/time_coarsen/time_coarsen.py
index 4d468e7d0..b1a617f0b 100644
--- a/scripts/time_coarsen/time_coarsen.py
+++ b/scripts/time_coarsen/time_coarsen.py
@@ -1,5 +1,7 @@
 import argparse
 import dataclasses
+import logging
+import os
 
 import dacite
 import xarray as xr
@@ -44,6 +46,10 @@ def main(config: Config):
 
 
 def process_path_pair(pair: PathPair, config: Config):
+    logging.info(f"Processing input: {pair.input} to output: {pair.output}")
+    if os.path.exists(pair.output):
+        logging.warning(f"Output path {pair.output} already exists. Skipping.")
+        return
     ds = xr.open_dataset(pair.input)
     constant_names = [
         name

From c9aae9b4ef6d9be2550b6936441f02a58ca4daca Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:17:02 +0000
Subject: [PATCH 09/13] make script executable

---
 scripts/time_coarsen/run_stats.sh | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 mode change 100644 => 100755 scripts/time_coarsen/run_stats.sh

diff --git a/scripts/time_coarsen/run_stats.sh b/scripts/time_coarsen/run_stats.sh
old mode 100644
new mode 100755

From 4f8e810b4b3fc857704cf794ebf1e8a310e13ac6 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:22:38 +0000
Subject: [PATCH 10/13] run get_stats directly

---
 scripts/time_coarsen/run_stats.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/time_coarsen/run_stats.sh b/scripts/time_coarsen/run_stats.sh
index 54a08a7f4..aa6853e47 100755
--- a/scripts/time_coarsen/run_stats.sh
+++ b/scripts/time_coarsen/run_stats.sh
@@ -21,7 +21,7 @@ run_stats() {
 
   gantry run \
     --name "$job_name" \
-    --description 'Run ACE training' \
+    --description 'Run ACE stats computation' \
     --beaker-image "$(cat $REPO_ROOT/latest_deps_only_image.txt)" \
     --workspace ai2/climate-titan \
     --priority urgent \
@@ -37,7 +37,7 @@ run_stats() {
     --system-python \
     --install "pip install --no-deps ." \
     "${extra_args[@]}" \
-    -- $SCRIPT_PATH/../data_process/compute_stats.sh --config $CONFIG_PATH
+    -- python $SCRIPT_PATH/../data_process/get_stats.py "$CONFIG_PATH" 0
 }
 
 base_name="time-coarsen-stats"

From bcbbf5a219787a6369cd7114111812bbfc61a787 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:30:25 +0000
Subject: [PATCH 11/13] make dask optional

---
 scripts/data_process/get_stats.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/scripts/data_process/get_stats.py b/scripts/data_process/get_stats.py
index aabd2deb4..202a5d3b7 100644
--- a/scripts/data_process/get_stats.py
+++ b/scripts/data_process/get_stats.py
@@ -88,11 +88,18 @@ def get_stats(
     debug: bool,
 ):
     # Import dask-related things here to enable testing in environments without dask.
-    import dask
-    import distributed
+    try:
+        import dask
+        import distributed
+
+        client = distributed.Client(n_workers=16)
+    except ImportError as e:
+        # warn and continue
+        logging.warning(f"Could not import dask ({e}), chunking is disabled.")
+        client = None
+        dask = None
 
     initial_time = time.time()
-    client = distributed.Client(n_workers=16)
 
     xr.set_options(keep_attrs=True, display_max_rows=100)
     logging.info(f"Reading data from {input_zarr}")
@@ -100,8 +107,11 @@ def get_stats(
     # Open data with roughly 128 MiB chunks via dask's automatic chunking. This
     # is useful when opening sharded zarr stores with an inner chunk size of 1,
     # which is otherwise inefficient for the type of computation done here.
-    with dask.config.set({"array.chunk-size": "128MiB"}):
-        ds = xr.open_zarr(input_zarr, chunks={"time": "auto"})
+    if dask is not None:
+        with dask.config.set({"array.chunk-size": "128MiB"}):
+            ds = xr.open_zarr(input_zarr, chunks={"time": "auto"})
+    else:
+        ds = xr.open_zarr(input_zarr)
 
     ds = ds.drop_vars(DROP_VARIABLES, errors="ignore")
     ds = ds.sel(time=slice(config.start_date, config.end_date))
@@ -186,7 +196,8 @@ def get_stats(
     total_time = time.time() - initial_time
     logging.info(f"Total time for computing stats: {total_time:0.2f} seconds.")
 
-    client.close()
+    if client is not None:
+        client.close()
     client = None
 
 

From de585b44658f3b15ba380beeb28f94084bdb6979 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:32:24 +0000
Subject: [PATCH 12/13] fix stats config maybe

---
 scripts/time_coarsen/c96-shield-stats.yaml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/scripts/time_coarsen/c96-shield-stats.yaml b/scripts/time_coarsen/c96-shield-stats.yaml
index c8fdbaf88..da8b91bc8 100644
--- a/scripts/time_coarsen/c96-shield-stats.yaml
+++ b/scripts/time_coarsen/c96-shield-stats.yaml
@@ -1,7 +1,7 @@
 runs:
-  ic_0001: gs://vcm-ml-raw-flexible-retention/2024-06-29-C96-SHiELD-AMIP/regridded-zarrs/gaussian_grid_180_by_360/ic_0001
-  ic_0002: gs://vcm-ml-raw-flexible-retention/2024-06-29-C96-SHiELD-AMIP/regridded-zarrs/gaussian_grid_180_by_360/ic_0002
-data_output_directory: gs://vcm-ml-intermediate/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset
+  ic_0001: ""
+  ic_0002: ""
+data_output_directory: /climate-default/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset
 stats:
   output_directory: gs://vcm-ml-intermediate/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
   beaker_dataset: 2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats

From bf9e693690c4e5259c3f80c5ae2b0ae682e630a2 Mon Sep 17 00:00:00 2001
From: Jeremy McGibbon <jeremy.mcgibbon@gmail.com>
Date: Thu, 5 Feb 2026 21:33:58 +0000
Subject: [PATCH 13/13] write stats to weka

---
 scripts/time_coarsen/c96-shield-stats.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/time_coarsen/c96-shield-stats.yaml b/scripts/time_coarsen/c96-shield-stats.yaml
index da8b91bc8..b46588c7c 100644
--- a/scripts/time_coarsen/c96-shield-stats.yaml
+++ b/scripts/time_coarsen/c96-shield-stats.yaml
@@ -3,7 +3,7 @@ runs:
   ic_0002: ""
 data_output_directory: /climate-default/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset
 stats:
-  output_directory: gs://vcm-ml-intermediate/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
+  output_directory: /climate-default/2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
   beaker_dataset: 2026-01-28-vertically-resolved-c96-1deg-daily-shield-amip-ensemble-dataset-stats
   start_date: "1940-01-01"
   end_date: "2021-12-31"