diff --git a/scripts/data_process/Makefile b/scripts/data_process/Makefile index f075f9fb55..c5313a9cfd 100644 --- a/scripts/data_process/Makefile +++ b/scripts/data_process/Makefile @@ -356,6 +356,14 @@ cm4_piControl_ocean_stats: cm4_1pctCO2_ocean_stats: ./compute_stats.sh --config configs/CM4-1pctCO2-ocean-$(RESOLUTION)-140yr.yaml +.PHONY: cm4_abrupt_4xCO2_ensemble_ocean_stats +cm4_abrupt_4xCO2_ensemble_ocean_stats: + ./compute_stats.sh --config configs/CM4-abrupt-4xCO2-ensemble-ocean-$(RESOLUTION).yaml + +.PHONY: cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats +cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats: + ./compute_stats.sh --config configs/CM4-abrupt-4xCO2-ensemble-ocean-$(RESOLUTION)-1daily.yaml + .PHONY: cm4_piControl_ocean_1daily_stats cm4_piControl_ocean_1daily_stats: ./compute_stats.sh --config configs/CM4-piControl-ocean-$(RESOLUTION)-1daily-200yr.yaml @@ -377,6 +385,8 @@ CONFIG_CM4_1PCTCO2_COUPLED_1DAILY = configs/CM4-1pctCO2-coupled-$(RESOLUTION)-1d CONFIG_E3SMV3_PICONTROL_COUPLED = configs/E3SMv3-piControl-100yr-coupled.yaml CONFIG_CM4_LIKE_AM4_RANDCO2_COUPLED = configs/CM4-like-AM4-random-CO2-ensemble-coupled.yaml CONFIG_CM4_LIKE_AM4_RANDCO2_COUPLED_4DEG = configs/CM4-like-AM4-random-CO2-ensemble-coupled-4deg.yaml +CONFIG_CM4_ABRUPT_4XCO2_COUPLED = configs/CM4-abrupt-4xCO2-ensemble-coupled-$(RESOLUTION).yaml +CONFIG_CM4_ABRUPT_4XCO2_COUPLED_1DAILY = configs/CM4-abrupt-4xCO2-ensemble-coupled-$(RESOLUTION)-1daily.yaml .PHONY: cm4_piControl_coupled cm4_piControl_coupled: @@ -428,6 +438,10 @@ e3smv3_piControl_100yr_coupled: cm4_like_am4_random_co2_ensemble_atmosphere: ./compute_dataset.sh --dataset --stats $(if $(filter 4deg,$(RESOLUTION)),--time-coarsen) --config configs/CM4-like-AM4-random-CO2-ensemble-atmosphere-$(RESOLUTION)-8layer.yaml +.PHONY: cm4_abrupt_4xCO2_ensemble_atmosphere +cm4_abrupt_4xCO2_ensemble_atmosphere: + ./compute_dataset.sh --dataset --stats $(if $(filter 4deg,$(RESOLUTION)),--time-coarsen) --config configs/CM4-abrupt-4xCO2-ensemble-atmosphere-$(RESOLUTION)-8layer.yaml + .PHONY: cm4_like_am4_random_co2_ensemble_coupled cm4_like_am4_random_co2_ensemble_coupled: conda run --no-capture-output -n create_coupled_datasets python ./create_coupled_datasets.py --yaml $(CONFIG_CM4_LIKE_AM4_RANDCO2_COUPLED) @@ -440,6 +454,22 @@ cm4_like_am4_random_co2_ensemble_coupled_argo: cm4_like_am4_random_co2_ensemble_coupled_4deg_argo: ./create_coupled_datasets.sh --config $(CONFIG_CM4_LIKE_AM4_RANDCO2_COUPLED_4DEG) --image $(COUPLED_PROCESSING_IMAGE):$(COUPLED_PROCESSING_IMAGE_VERSION) +.PHONY: cm4_abrupt_4xCO2_ensemble_coupled +cm4_abrupt_4xCO2_ensemble_coupled: + conda run --no-capture-output -n create_coupled_datasets python ./create_coupled_datasets.py --yaml $(CONFIG_CM4_ABRUPT_4XCO2_COUPLED) + +.PHONY: cm4_abrupt_4xCO2_ensemble_coupled_argo +cm4_abrupt_4xCO2_ensemble_coupled_argo: + ./create_coupled_datasets.sh --config $(CONFIG_CM4_ABRUPT_4XCO2_COUPLED) --image $(COUPLED_PROCESSING_IMAGE):$(COUPLED_PROCESSING_IMAGE_VERSION) + +.PHONY: cm4_abrupt_4xCO2_ensemble_coupled_1daily +cm4_abrupt_4xCO2_ensemble_coupled_1daily: + conda run --no-capture-output -n create_coupled_datasets python ./create_coupled_datasets.py --yaml $(CONFIG_CM4_ABRUPT_4XCO2_COUPLED_1DAILY) + +.PHONY: cm4_abrupt_4xCO2_ensemble_coupled_1daily_argo +cm4_abrupt_4xCO2_ensemble_coupled_1daily_argo: + ./create_coupled_datasets.sh --config $(CONFIG_CM4_ABRUPT_4XCO2_COUPLED_1DAILY) --image $(COUPLED_PROCESSING_IMAGE):$(COUPLED_PROCESSING_IMAGE_VERSION) + .PHONY: shield_amip_constant_co2_c96_dataset shield_amip_constant_co2_c96_dataset: ./compute_dataset.sh --dataset --stats --time-coarsen --config configs/shield-amip-constant-co2-c96-$(RESOLUTION)-$(LAYERS).yaml diff --git a/scripts/data_process/combine_stats.py b/scripts/data_process/combine_stats.py index 80dbc8dca4..662671ef85 100644 --- a/scripts/data_process/combine_stats.py +++ b/scripts/data_process/combine_stats.py @@ -9,6 +9,7 @@ import fsspec import xarray as xr import yaml +from fs_utils import load_netcdf def copy(source: str, destination: str): @@ -59,12 +60,7 @@ def _make_history_string(config_filename: str, stats_output_dir: str): def open_datasets(roots: List[str], filename: str) -> List[xr.Dataset]: - datasets = [] - for root in roots: - with fsspec.open(root + filename) as file: - ds = xr.open_dataset(file, decode_timedelta=False).load() - datasets.append(ds) - return datasets + return [load_netcdf(root + filename) for root in roots] def _combined_stats_exist(output_directory: str, subdirectory: str) -> bool: diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-1deg-8layer.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-1deg-8layer.yaml new file mode 100644 index 0000000000..c2b5e50a42 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-1deg-8layer.yaml @@ -0,0 +1,146 @@ +runs: + ic_0001: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0001 + ic_0002: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0002 + ic_0003: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0003 + ic_0004: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0004 + ic_0005: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0005 + ic_0006: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0006 + ic_0007: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0007 + ic_0008: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0008 + ic_0009: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0009 + ic_0010: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0010 + ic_0011: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0011 + ic_0012: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0012 + ic_0013: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0013 + ic_0014: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0014 + ic_0015: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0015 + ic_0016: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0016 + ic_0017: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0017 + ic_0018: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0018 + ic_0019: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0019 + ic_0020: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0020 + ic_0021: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0021 + ic_0022: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0022 + ic_0023: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0023 + ic_0024: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0024 + ic_0025: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0025 + ic_0026: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0026 + ic_0027: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0027 + ic_0028: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0028 + ic_0029: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0029 + ic_0030: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0030 + ic_0031: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0031 + ic_0032: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0032 + ic_0033: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0033 + ic_0034: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0034 + ic_0035: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0035 + ic_0036: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_180_by_360/abrupt-4xCO2-ic_0036 +data_output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-stats + data_type: CM4 + beaker_dataset: 2026-10-16-CM4-abrupt-4xCO2-ensemble-stats +dataset_computation: + n_split: 10 + reference_vertical_coordinate_file: gs://vcm-ml-raw-flexible-retention/2024-08-10-pre-industrial-CM4-simulation/vertical-coordinate-file/fv_core.res.nc + reference_vertical_coordinate_file_land: gs://vcm-ml-raw-flexible-retention/2024-11-01-CM4-piControl/land-vertical-coordinate-file/land-vertical-coordinate.nc + vertical_coarsening_indices: + # computed here: https://github.com/ai2cm/explore/blob/master/jamesd/2024-08-13-pre-industiral-CM4-eda/2024-08-28-AM4-vertical-indices.ipynb + - [0, 7] + - [7, 10] + - [10, 13] + - [13, 16] + - [16, 18] + - [18, 22] + - [22, 25] + - [25, 33] + vertical_coarsening_indices_land: + # computed here: https://github.com/ai2cm/explore2/blob/main/troya/2025-02-28-LS-evaluation/soil-levels-lm4.ipynb + - [0, 2] + - [2, 6] + - [6, 10] + - [10, 15] + validate_vertical_coarsening_indices_land: false # land vertical coarsening indices intentionally do not span full depth + mask_soil_moisture: true + renaming: + specific_humidity_at_two_meters: Q2m + air_temperature_at_two_meters: TMP2m + eastward_wind_at_ten_meters: UGRD10m + northward_wind_at_ten_meters: VGRD10m + interval_average_carbon_dioxide: carbon_dioxide + variable_sources: + fluxes_2d.zarr: + - PRATEsfc + - SHTFLsfc + - DLWRFsfc + - DSWRFsfc + - DSWRFtoa + - ULWRFsfc + - ULWRFtoa + - USWRFsfc + - USWRFtoa + - eastward_surface_wind_stress + - northward_surface_wind_stress + - surface_evaporation_rate + - total_energy + - total_frozen_precipitation_rate + full_state.zarr: + # 2D vars + - HGTsfc # static + - PRESsfc + - PRMSL + - surface_temperature + - air_temperature_at_two_meters + - specific_humidity_at_two_meters + - eastward_wind_at_ten_meters + - northward_wind_at_ten_meters + - h500 + - TMP850 + # 3D vars: + - air_temperature + - specific_humidity # water species + - cloud_water_mixing_ratio # water species + - cloud_ice_mixing_ratio # water species + - eastward_wind + - northward_wind + full_state_land.zarr: + - column_soil_moisture + - GHTFLsfc + - surface_snow_amount + - surface_snow_area_fraction + - total_land_column_energy_without_rivers + - total_precipitation_heat_flux + - total_vapor_sensible_heat_flux + - downward_longwave_radiative_flux + - net_surface_longwave_radiative_flux + full_state_soil.zarr: + # below are vertically resolved soil layer variables -- need to be coarsened + - total_moisture_content_of_soil_layer + - temperature_of_soil_layer + land_static.zarr: + - land_fraction + ice_6hourly_snap.zarr: + - sea_ice_fraction + # Constructed an interval-average CO2 concentration from 6-hour snapshots + # prior to uploading since AM4 has a subtle off-by-one bug when computing + # interval averages of the scalar CO2 concentration online. + scalar_6hourly_snap.zarr: + - interval_average_carbon_dioxide + land_surface_type_fractions.zarr: + - glac_fraction + - soil_fraction + - lake_fraction + standard_names: + longitude_dim: lon + latitude_dim: lat + graupel_mixing_ratio: None + rain_mixing_ratio: None + snow_mixing_ratio: None + precipitable_water_path: None + surface_snow_rate: None + surface_ice_rate: None + surface_graupel_rate: None + land_names_to_vertically_coarsen_by_height_weighting: + - temperature_of_soil_layer + land_names_to_vertically_coarsen_by_sum: + - total_moisture_content_of_soil_layer diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-4deg-8layer.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-4deg-8layer.yaml new file mode 100644 index 0000000000..2a399c040f --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-atmosphere-4deg-8layer.yaml @@ -0,0 +1,241 @@ +runs: + ic_0001: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0001 + ic_0002: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0002 + ic_0003: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0003 + ic_0004: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0004 + ic_0005: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0005 + ic_0006: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0006 + ic_0007: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0007 + ic_0008: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0008 + ic_0009: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0009 + ic_0010: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0010 + ic_0011: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0011 + ic_0012: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0012 + ic_0013: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0013 + ic_0014: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0014 + ic_0015: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0015 + ic_0016: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0016 + ic_0017: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0017 + ic_0018: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0018 + ic_0019: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0019 + ic_0020: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0020 + ic_0021: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0021 + ic_0022: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0022 + ic_0023: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0023 + ic_0024: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0024 + ic_0025: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0025 + ic_0026: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0026 + ic_0027: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0027 + ic_0028: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0028 + ic_0029: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0029 + ic_0030: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0030 + ic_0031: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0031 + ic_0032: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0032 + ic_0033: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0033 + ic_0034: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0034 + ic_0035: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0035 + ic_0036: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-abrupt-4xCO2-ensemble/regridded-zarrs/gaussian_grid_45_by_90/abrupt-4xCO2-ic_0036 +data_output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-stats + data_type: CM4 + beaker_dataset: 2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-stats +dataset_computation: + sharding: + time_dim: 3600 + n_split: 3 + reference_vertical_coordinate_file: gs://vcm-ml-raw-flexible-retention/2024-08-10-pre-industrial-CM4-simulation/vertical-coordinate-file/fv_core.res.nc + reference_vertical_coordinate_file_land: gs://vcm-ml-raw-flexible-retention/2024-11-01-CM4-piControl/land-vertical-coordinate-file/land-vertical-coordinate.nc + vertical_coarsening_indices: + # computed here: https://github.com/ai2cm/explore/blob/master/jamesd/2024-08-13-pre-industiral-CM4-eda/2024-08-28-AM4-vertical-indices.ipynb + - [0, 7] + - [7, 10] + - [10, 13] + - [13, 16] + - [16, 18] + - [18, 22] + - [22, 25] + - [25, 33] + vertical_coarsening_indices_land: + # computed here: https://github.com/ai2cm/explore2/blob/main/troya/2025-02-28-LS-evaluation/soil-levels-lm4.ipynb + - [0, 2] + - [2, 6] + - [6, 10] + - [10, 15] + validate_vertical_coarsening_indices_land: false # land vertical coarsening indices intentionally do not span full depth + mask_soil_moisture: true + renaming: + specific_humidity_at_two_meters: Q2m + air_temperature_at_two_meters: TMP2m + eastward_wind_at_ten_meters: UGRD10m + northward_wind_at_ten_meters: VGRD10m + interval_average_carbon_dioxide: carbon_dioxide + variable_sources: + fluxes_2d.zarr: + - PRATEsfc + - SHTFLsfc + - DLWRFsfc + - DSWRFsfc + - DSWRFtoa + - ULWRFsfc + - ULWRFtoa + - USWRFsfc + - USWRFtoa + - eastward_surface_wind_stress + - northward_surface_wind_stress + - surface_evaporation_rate + - total_energy + - total_frozen_precipitation_rate + full_state.zarr: + # 2D vars + - HGTsfc # static + - PRESsfc + - PRMSL + - surface_temperature + - air_temperature_at_two_meters + - specific_humidity_at_two_meters + - eastward_wind_at_ten_meters + - northward_wind_at_ten_meters + - h500 + - TMP850 + # 3D vars: + - air_temperature + - specific_humidity # water species + - cloud_water_mixing_ratio # water species + - cloud_ice_mixing_ratio # water species + - eastward_wind + - northward_wind + full_state_land.zarr: + - column_soil_moisture + - GHTFLsfc + - surface_snow_amount + - surface_snow_area_fraction + - total_land_column_energy_without_rivers + - total_precipitation_heat_flux + - total_vapor_sensible_heat_flux + - downward_longwave_radiative_flux + - net_surface_longwave_radiative_flux + full_state_soil.zarr: + # below are vertically resolved soil layer variables -- need to be coarsened + - total_moisture_content_of_soil_layer + - temperature_of_soil_layer + land_static.zarr: + - land_fraction + ice_6hourly_snap.zarr: + - sea_ice_fraction + # Constructed an interval-average CO2 concentration from 6-hour snapshots + # prior to uploading since AM4 has a subtle off-by-one bug when computing + # interval averages of the scalar CO2 concentration online. + scalar_6hourly_snap.zarr: + - interval_average_carbon_dioxide + land_surface_type_fractions.zarr: + - glac_fraction + - soil_fraction + - lake_fraction + standard_names: + longitude_dim: lon + latitude_dim: lat + graupel_mixing_ratio: None + rain_mixing_ratio: None + snow_mixing_ratio: None + precipitable_water_path: None + surface_snow_rate: None + surface_ice_rate: None + surface_graupel_rate: None + land_names_to_vertically_coarsen_by_height_weighting: + - temperature_of_soil_layer + land_names_to_vertically_coarsen_by_sum: + - total_moisture_content_of_soil_layer +time_coarsen: + factor: 4 + data_output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-daily + stats_output_directory: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-daily-stats + snapshot_names: + - PRESsfc + - PRMSL + - Q2m + - TMP2m + - TMP850 + - UGRD10m + - VGRD10m + - air_temperature_0 + - air_temperature_1 + - air_temperature_2 + - air_temperature_3 + - air_temperature_4 + - air_temperature_5 + - air_temperature_6 + - air_temperature_7 + - column_soil_moisture + - eastward_wind_0 + - eastward_wind_1 + - eastward_wind_2 + - eastward_wind_3 + - eastward_wind_4 + - eastward_wind_5 + - eastward_wind_6 + - eastward_wind_7 + - h500 + - northward_wind_0 + - northward_wind_1 + - northward_wind_2 + - northward_wind_3 + - northward_wind_4 + - northward_wind_5 + - northward_wind_6 + - northward_wind_7 + - sea_ice_fraction + - 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_snow_amount + - surface_snow_area_fraction + - surface_temperature + - temperature_of_soil_layer_0 + - temperature_of_soil_layer_1 + - temperature_of_soil_layer_2 + - temperature_of_soil_layer_3 + - total_land_column_energy_without_rivers + - total_moisture_content_of_soil_layer_0 + - total_moisture_content_of_soil_layer_1 + - total_moisture_content_of_soil_layer_2 + - total_moisture_content_of_soil_layer_3 + - total_water_path + - total_energy + window_names: + - DLWRFsfc + - DSWRFsfc + - DSWRFtoa + - GHTFLsfc + - LHTFLsfc + - PRATEsfc + - SHTFLsfc + - ULWRFsfc + - ULWRFtoa + - USWRFsfc + - USWRFtoa + - carbon_dioxide + - downward_longwave_radiative_flux + - eastward_surface_wind_stress + - net_surface_longwave_radiative_flux + - northward_surface_wind_stress + - tendency_of_total_water_path + - tendency_of_total_water_path_due_to_advection + - total_frozen_precipitation_rate + - total_precipitation_heat_flux + - total_vapor_sensible_heat_flux + constant_prefixes: + - HGTsfc + - ak_ + - bk_ + - glac_fraction + - lake_fraction + - land_fraction + - lat + - lon + - soil_fraction diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg-1daily.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg-1daily.yaml new file mode 100644 index 0000000000..2d225a4e4e --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg-1daily.yaml @@ -0,0 +1,314 @@ +# make cm4_abrupt_4xCO2_ensemble_coupled_1daily RESOLUTION=1deg +# Requires the ocean stats from make cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats RESOLUTION=1deg. +version: "2026-10-20" +family_name: cm4-abrupt-4xCO2-ensemble-1deg-coupled-1daily +output_directory: gs://vcm-ml-intermediate + +coupled_datasets: + coupled_ts: + how: interpolate_sst + coupled_sea_surface: + precomputed_sea_ice_mask: + zarr_path: gs://vcm-ml-intermediate/2026-09-02-cm4-piControl-1deg-coupled-1daily-ocean.zarr + ocean_extra_masked_names: ["HI", "sea_ice_volume", "UI", "VI", "sfdsi"] + ocean_extra_fill_values: + UI: 0.0 + VI: 0.0 + # the ocean store is 1-daily; the default (120h) is for 5-daily stores + timedelta: 24h + surface_flux_window_avg: + window_timedelta: 24h + # No first_timestamp or last_timestamp: the members start on different + # dates, so each member's flux windows are right-labeled at its own + # ocean timesteps, i.e. each window (t - 24h, t] ends at an ocean + # timestep t, starting from the ocean store's first timestamp. + subset_names: + - DLWRFsfc + - DSWRFsfc + - ULWRFsfc + - USWRFsfc + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - eastward_surface_wind_stress + - northward_surface_wind_stress + - total_frozen_precipitation_rate + coupled_sea_ice: + # No window_avg and no sea_ice input dataset: the sea ice fraction is + # sourced from the ocean store's native 1-daily fields and forward-filled + # onto the atmosphere's 6-hourly time index. + include_ts: true + use_atmosphere_sea_ice_fraction_fallback: false + output_writer: + n_dask_workers: 32 + # Members span at most 1460 6-hourly steps, i.e. a few 360-step outer + # chunks, so the default 50 partitions would be mostly empty. + n_split: 1 + +stats: + beaker_dataset: 2026-10-20-cm4-abrupt-4xCO2-ensemble-1deg-coupled-1daily-stats + +input_datasets: + climate_data_type: "CM4" + stats: + atmosphere_dir: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-stats/combined + ocean_dir: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily-stats/combined + # Every member shares ic_0001's atmosphere and ocean settings (anchored + # below) and overrides only its zarr_path. + runs: + ic_0001: + atmosphere: &atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0001.zarr + time_chunk_size: 1460 + extra_fields: + names_and_prefixes: ["ak_", "bk_"] + ocean: &ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0001.zarr + time_chunk_size: 365 + extra_fields: + names_and_prefixes: ["idepth_", "mask_", "deptho"] + ic_0002: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0002.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0002.zarr + ic_0003: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0003.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0003.zarr + ic_0004: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0004.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0004.zarr + ic_0005: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0005.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0005.zarr + ic_0006: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0006.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0006.zarr + ic_0007: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0007.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0007.zarr + ic_0008: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0008.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0008.zarr + ic_0009: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0009.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0009.zarr + ic_0010: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0010.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0010.zarr + ic_0011: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0011.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0011.zarr + ic_0012: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0012.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0012.zarr + ic_0013: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0013.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0013.zarr + ic_0014: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0014.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0014.zarr + ic_0015: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0015.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0015.zarr + ic_0016: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0016.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0016.zarr + ic_0017: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0017.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0017.zarr + ic_0018: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0018.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0018.zarr + ic_0019: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0019.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0019.zarr + ic_0020: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0020.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0020.zarr + ic_0021: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0021.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0021.zarr + ic_0022: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0022.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0022.zarr + ic_0023: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0023.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0023.zarr + ic_0024: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0024.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0024.zarr + ic_0025: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0025.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0025.zarr + ic_0026: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0026.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0026.zarr + ic_0027: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0027.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0027.zarr + ic_0028: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0028.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0028.zarr + ic_0029: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0029.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0029.zarr + ic_0030: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0030.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0030.zarr + ic_0031: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0031.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0031.zarr + ic_0032: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0032.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0032.zarr + ic_0033: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0033.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0033.zarr + ic_0034: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0034.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0034.zarr + ic_0035: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0035.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0035.zarr + ic_0036: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0036.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/ic_0036.zarr diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg.yaml new file mode 100644 index 0000000000..401da02420 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-1deg.yaml @@ -0,0 +1,312 @@ +# make cm4_abrupt_4xCO2_ensemble_coupled +# Requires the ocean stats from make cm4_abrupt_4xCO2_ensemble_ocean_stats. +version: "2026-10-20" +family_name: cm4-abrupt-4xCO2-ensemble-1deg-coupled +output_directory: gs://vcm-ml-intermediate + +coupled_datasets: + coupled_ts: + how: interpolate_sst + coupled_sea_surface: + precomputed_sea_ice_mask: + zarr_path: gs://vcm-ml-intermediate/2026-07-15-om4-picontrol-1deg-coupled-ocean.zarr + ocean_extra_masked_names: ["HI", "sea_ice_volume", "UI", "VI", "sfdsi"] + ocean_extra_fill_values: + UI: 0.0 + VI: 0.0 + surface_flux_window_avg: + window_timedelta: 120h + # No first_timestamp or last_timestamp: the members start on different + # dates, so each member's flux windows are right-labeled at its own + # ocean timesteps, i.e. each window (t - 120h, t] ends at an ocean + # timestep t, starting from the ocean store's first timestamp. + subset_names: + - DLWRFsfc + - DSWRFsfc + - ULWRFsfc + - USWRFsfc + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - eastward_surface_wind_stress + - northward_surface_wind_stress + - total_frozen_precipitation_rate + coupled_sea_ice: + # No window_avg and no sea_ice input dataset: the sea ice fraction is + # sourced from the ocean store's native 5-daily fields and forward-filled + # onto the atmosphere's 6-hourly time index. + include_ts: true + use_atmosphere_sea_ice_fraction_fallback: false + output_writer: + n_dask_workers: 32 + # Members span at most 1460 6-hourly steps, i.e. a few 360-step outer + # chunks, so the default 50 partitions would be mostly empty. + n_split: 1 + +stats: + beaker_dataset: 2026-10-20-cm4-abrupt-4xCO2-ensemble-1deg-coupled-stats + +input_datasets: + climate_data_type: "CM4" + stats: + atmosphere_dir: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-stats/combined + ocean_dir: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily-stats/combined + # Every member shares ic_0001's atmosphere and ocean settings (anchored + # below) and overrides only its zarr_path. + runs: + ic_0001: + atmosphere: &atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0001.zarr + time_chunk_size: 1460 + extra_fields: + names_and_prefixes: ["ak_", "bk_"] + ocean: &ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0001.zarr + time_chunk_size: 73 + extra_fields: + names_and_prefixes: ["idepth_", "mask_", "deptho"] + ic_0002: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0002.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0002.zarr + ic_0003: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0003.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0003.zarr + ic_0004: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0004.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0004.zarr + ic_0005: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0005.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0005.zarr + ic_0006: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0006.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0006.zarr + ic_0007: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0007.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0007.zarr + ic_0008: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0008.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0008.zarr + ic_0009: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0009.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0009.zarr + ic_0010: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0010.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0010.zarr + ic_0011: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0011.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0011.zarr + ic_0012: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0012.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0012.zarr + ic_0013: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0013.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0013.zarr + ic_0014: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0014.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0014.zarr + ic_0015: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0015.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0015.zarr + ic_0016: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0016.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0016.zarr + ic_0017: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0017.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0017.zarr + ic_0018: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0018.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0018.zarr + ic_0019: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0019.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0019.zarr + ic_0020: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0020.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0020.zarr + ic_0021: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0021.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0021.zarr + ic_0022: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0022.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0022.zarr + ic_0023: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0023.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0023.zarr + ic_0024: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0024.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0024.zarr + ic_0025: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0025.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0025.zarr + ic_0026: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0026.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0026.zarr + ic_0027: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0027.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0027.zarr + ic_0028: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0028.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0028.zarr + ic_0029: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0029.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0029.zarr + ic_0030: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0030.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0030.zarr + ic_0031: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0031.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0031.zarr + ic_0032: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0032.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0032.zarr + ic_0033: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0033.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0033.zarr + ic_0034: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0034.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0034.zarr + ic_0035: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0035.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0035.zarr + ic_0036: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble/ic_0036.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/ic_0036.zarr diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg-1daily.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg-1daily.yaml new file mode 100644 index 0000000000..2fd6feccd4 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg-1daily.yaml @@ -0,0 +1,314 @@ +# make cm4_abrupt_4xCO2_ensemble_coupled_1daily RESOLUTION=4deg +# Requires the ocean stats from make cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats RESOLUTION=4deg. +version: "2026-10-20" +family_name: cm4-abrupt-4xCO2-ensemble-4deg-coupled-1daily +output_directory: gs://vcm-ml-intermediate + +coupled_datasets: + coupled_ts: + how: interpolate_sst + coupled_sea_surface: + precomputed_sea_ice_mask: + zarr_path: gs://vcm-ml-intermediate/2026-09-02-cm4-piControl-4deg-coupled-1daily-ocean.zarr + ocean_extra_masked_names: ["HI", "sea_ice_volume", "UI", "VI", "sfdsi"] + ocean_extra_fill_values: + UI: 0.0 + VI: 0.0 + # the ocean store is 1-daily; the default (120h) is for 5-daily stores + timedelta: 24h + surface_flux_window_avg: + window_timedelta: 24h + # No first_timestamp or last_timestamp: the members start on different + # dates, so each member's flux windows are right-labeled at its own + # ocean timesteps, i.e. each window (t - 24h, t] ends at an ocean + # timestep t, starting from the ocean store's first timestamp. + subset_names: + - DLWRFsfc + - DSWRFsfc + - ULWRFsfc + - USWRFsfc + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - eastward_surface_wind_stress + - northward_surface_wind_stress + - total_frozen_precipitation_rate + coupled_sea_ice: + # No window_avg and no sea_ice input dataset: the sea ice fraction is + # sourced from the ocean store's native 1-daily fields and forward-filled + # onto the atmosphere's 6-hourly time index. + include_ts: true + use_atmosphere_sea_ice_fraction_fallback: false + output_writer: + n_dask_workers: 32 + # Members span at most 1460 6-hourly steps, i.e. a few 360-step outer + # chunks, so the default 50 partitions would be mostly empty. + n_split: 1 + +stats: + beaker_dataset: 2026-10-20-cm4-abrupt-4xCO2-ensemble-4deg-coupled-1daily-stats + +input_datasets: + climate_data_type: "CM4" + stats: + atmosphere_dir: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-stats/combined + ocean_dir: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily-stats/combined + # Every member shares ic_0001's atmosphere and ocean settings (anchored + # below) and overrides only its zarr_path. + runs: + ic_0001: + atmosphere: &atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0001.zarr + time_chunk_size: 1460 + extra_fields: + names_and_prefixes: ["ak_", "bk_"] + ocean: &ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0001.zarr + time_chunk_size: 365 + extra_fields: + names_and_prefixes: ["idepth_", "mask_", "deptho"] + ic_0002: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0002.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0002.zarr + ic_0003: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0003.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0003.zarr + ic_0004: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0004.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0004.zarr + ic_0005: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0005.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0005.zarr + ic_0006: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0006.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0006.zarr + ic_0007: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0007.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0007.zarr + ic_0008: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0008.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0008.zarr + ic_0009: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0009.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0009.zarr + ic_0010: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0010.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0010.zarr + ic_0011: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0011.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0011.zarr + ic_0012: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0012.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0012.zarr + ic_0013: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0013.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0013.zarr + ic_0014: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0014.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0014.zarr + ic_0015: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0015.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0015.zarr + ic_0016: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0016.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0016.zarr + ic_0017: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0017.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0017.zarr + ic_0018: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0018.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0018.zarr + ic_0019: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0019.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0019.zarr + ic_0020: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0020.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0020.zarr + ic_0021: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0021.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0021.zarr + ic_0022: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0022.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0022.zarr + ic_0023: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0023.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0023.zarr + ic_0024: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0024.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0024.zarr + ic_0025: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0025.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0025.zarr + ic_0026: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0026.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0026.zarr + ic_0027: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0027.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0027.zarr + ic_0028: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0028.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0028.zarr + ic_0029: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0029.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0029.zarr + ic_0030: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0030.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0030.zarr + ic_0031: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0031.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0031.zarr + ic_0032: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0032.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0032.zarr + ic_0033: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0033.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0033.zarr + ic_0034: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0034.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0034.zarr + ic_0035: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0035.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0035.zarr + ic_0036: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0036.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/ic_0036.zarr diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg.yaml new file mode 100644 index 0000000000..4e5c0dad50 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-coupled-4deg.yaml @@ -0,0 +1,312 @@ +# make cm4_abrupt_4xCO2_ensemble_coupled RESOLUTION=4deg +# Requires the ocean stats from make cm4_abrupt_4xCO2_ensemble_ocean_stats RESOLUTION=4deg. +version: "2026-10-20" +family_name: cm4-abrupt-4xCO2-ensemble-4deg-coupled +output_directory: gs://vcm-ml-intermediate + +coupled_datasets: + coupled_ts: + how: interpolate_sst + coupled_sea_surface: + precomputed_sea_ice_mask: + zarr_path: gs://vcm-ml-intermediate/2026-07-22-cm4-picontrol-4deg-coupled-ocean.zarr + ocean_extra_masked_names: ["HI", "sea_ice_volume", "UI", "VI", "sfdsi"] + ocean_extra_fill_values: + UI: 0.0 + VI: 0.0 + surface_flux_window_avg: + window_timedelta: 120h + # No first_timestamp or last_timestamp: the members start on different + # dates, so each member's flux windows are right-labeled at its own + # ocean timesteps, i.e. each window (t - 120h, t] ends at an ocean + # timestep t, starting from the ocean store's first timestamp. + subset_names: + - DLWRFsfc + - DSWRFsfc + - ULWRFsfc + - USWRFsfc + - LHTFLsfc + - SHTFLsfc + - PRATEsfc + - eastward_surface_wind_stress + - northward_surface_wind_stress + - total_frozen_precipitation_rate + coupled_sea_ice: + # No window_avg and no sea_ice input dataset: the sea ice fraction is + # sourced from the ocean store's native 5-daily fields and forward-filled + # onto the atmosphere's 6-hourly time index. + include_ts: true + use_atmosphere_sea_ice_fraction_fallback: false + output_writer: + n_dask_workers: 32 + # Members span at most 1460 6-hourly steps, i.e. a few 360-step outer + # chunks, so the default 50 partitions would be mostly empty. + n_split: 1 + +stats: + beaker_dataset: 2026-10-20-cm4-abrupt-4xCO2-ensemble-4deg-coupled-stats + +input_datasets: + climate_data_type: "CM4" + stats: + atmosphere_dir: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg-stats/combined + ocean_dir: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily-stats/combined + # Every member shares ic_0001's atmosphere and ocean settings (anchored + # below) and overrides only its zarr_path. + runs: + ic_0001: + atmosphere: &atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0001.zarr + time_chunk_size: 1460 + extra_fields: + names_and_prefixes: ["ak_", "bk_"] + ocean: &ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0001.zarr + time_chunk_size: 73 + extra_fields: + names_and_prefixes: ["idepth_", "mask_", "deptho"] + ic_0002: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0002.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0002.zarr + ic_0003: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0003.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0003.zarr + ic_0004: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0004.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0004.zarr + ic_0005: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0005.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0005.zarr + ic_0006: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0006.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0006.zarr + ic_0007: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0007.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0007.zarr + ic_0008: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0008.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0008.zarr + ic_0009: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0009.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0009.zarr + ic_0010: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0010.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0010.zarr + ic_0011: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0011.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0011.zarr + ic_0012: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0012.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0012.zarr + ic_0013: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0013.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0013.zarr + ic_0014: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0014.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0014.zarr + ic_0015: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0015.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0015.zarr + ic_0016: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0016.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0016.zarr + ic_0017: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0017.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0017.zarr + ic_0018: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0018.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0018.zarr + ic_0019: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0019.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0019.zarr + ic_0020: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0020.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0020.zarr + ic_0021: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0021.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0021.zarr + ic_0022: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0022.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0022.zarr + ic_0023: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0023.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0023.zarr + ic_0024: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0024.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0024.zarr + ic_0025: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0025.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0025.zarr + ic_0026: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0026.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0026.zarr + ic_0027: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0027.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0027.zarr + ic_0028: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0028.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0028.zarr + ic_0029: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0029.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0029.zarr + ic_0030: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0030.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0030.zarr + ic_0031: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0031.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0031.zarr + ic_0032: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0032.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0032.zarr + ic_0033: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0033.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0033.zarr + ic_0034: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0034.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0034.zarr + ic_0035: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0035.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0035.zarr + ic_0036: + atmosphere: + <<: *atmosphere + zarr_path: gs://vcm-ml-intermediate/2026-10-16-CM4-abrupt-4xCO2-ensemble-4deg/ic_0036.zarr + ocean: + <<: *ocean + zarr_path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/ic_0036.zarr diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml new file mode 100644 index 0000000000..b46975bf35 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml @@ -0,0 +1,46 @@ +# make cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats RESOLUTION=1deg +# Stats of the per-member 1-daily ocean and sea-ice stores produced by the +# scripts/gfdl_om4 xrbeam Dataflow pipeline +# (gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml). +runs: + ic_0001: "" + ic_0002: "" + ic_0003: "" + ic_0004: "" + ic_0005: "" + ic_0006: "" + ic_0007: "" + ic_0008: "" + ic_0009: "" + ic_0010: "" + ic_0011: "" + ic_0012: "" + ic_0013: "" + ic_0014: "" + ic_0015: "" + ic_0016: "" + ic_0017: "" + ic_0018: "" + ic_0019: "" + ic_0020: "" + ic_0021: "" + ic_0022: "" + ic_0023: "" + ic_0024: "" + ic_0025: "" + ic_0026: "" + ic_0027: "" + ic_0028: "" + ic_0029: "" + ic_0030: "" + ic_0031: "" + ic_0032: "" + ic_0033: "" + ic_0034: "" + ic_0035: "" + ic_0036: "" +data_output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily-stats + data_type: CM4 + beaker_dataset: 2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily-stats diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg.yaml new file mode 100644 index 0000000000..ee6146f083 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-1deg.yaml @@ -0,0 +1,46 @@ +# make cm4_abrupt_4xCO2_ensemble_ocean_stats +# Stats of the per-member 5-daily ocean and sea-ice stores produced by the +# scripts/gfdl_om4 xrbeam Dataflow pipeline +# (gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily.yaml). +runs: + ic_0001: "" + ic_0002: "" + ic_0003: "" + ic_0004: "" + ic_0005: "" + ic_0006: "" + ic_0007: "" + ic_0008: "" + ic_0009: "" + ic_0010: "" + ic_0011: "" + ic_0012: "" + ic_0013: "" + ic_0014: "" + ic_0015: "" + ic_0016: "" + ic_0017: "" + ic_0018: "" + ic_0019: "" + ic_0020: "" + ic_0021: "" + ic_0022: "" + ic_0023: "" + ic_0024: "" + ic_0025: "" + ic_0026: "" + ic_0027: "" + ic_0028: "" + ic_0029: "" + ic_0030: "" + ic_0031: "" + ic_0032: "" + ic_0033: "" + ic_0034: "" + ic_0035: "" + ic_0036: "" +data_output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily-stats + data_type: CM4 + beaker_dataset: 2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily-stats diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml new file mode 100644 index 0000000000..4ed2909dd3 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml @@ -0,0 +1,46 @@ +# make cm4_abrupt_4xCO2_ensemble_ocean_1daily_stats RESOLUTION=4deg +# Stats of the per-member 1-daily ocean and sea-ice stores produced by the +# scripts/gfdl_om4 xrbeam Dataflow pipeline +# (gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml). +runs: + ic_0001: "" + ic_0002: "" + ic_0003: "" + ic_0004: "" + ic_0005: "" + ic_0006: "" + ic_0007: "" + ic_0008: "" + ic_0009: "" + ic_0010: "" + ic_0011: "" + ic_0012: "" + ic_0013: "" + ic_0014: "" + ic_0015: "" + ic_0016: "" + ic_0017: "" + ic_0018: "" + ic_0019: "" + ic_0020: "" + ic_0021: "" + ic_0022: "" + ic_0023: "" + ic_0024: "" + ic_0025: "" + ic_0026: "" + ic_0027: "" + ic_0028: "" + ic_0029: "" + ic_0030: "" + ic_0031: "" + ic_0032: "" + ic_0033: "" + ic_0034: "" + ic_0035: "" + ic_0036: "" +data_output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily-stats + data_type: CM4 + beaker_dataset: 2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily-stats diff --git a/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg.yaml b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg.yaml new file mode 100644 index 0000000000..68c6d92078 --- /dev/null +++ b/scripts/data_process/configs/CM4-abrupt-4xCO2-ensemble-ocean-4deg.yaml @@ -0,0 +1,46 @@ +# make cm4_abrupt_4xCO2_ensemble_ocean_stats RESOLUTION=4deg +# Stats of the per-member 5-daily ocean and sea-ice stores produced by the +# scripts/gfdl_om4 xrbeam Dataflow pipeline +# (gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily.yaml). +runs: + ic_0001: "" + ic_0002: "" + ic_0003: "" + ic_0004: "" + ic_0005: "" + ic_0006: "" + ic_0007: "" + ic_0008: "" + ic_0009: "" + ic_0010: "" + ic_0011: "" + ic_0012: "" + ic_0013: "" + ic_0014: "" + ic_0015: "" + ic_0016: "" + ic_0017: "" + ic_0018: "" + ic_0019: "" + ic_0020: "" + ic_0021: "" + ic_0022: "" + ic_0023: "" + ic_0024: "" + ic_0025: "" + ic_0026: "" + ic_0027: "" + ic_0028: "" + ic_0029: "" + ic_0030: "" + ic_0031: "" + ic_0032: "" + ic_0033: "" + ic_0034: "" + ic_0035: "" + ic_0036: "" +data_output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily +stats: + output_directory: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily-stats + data_type: CM4 + beaker_dataset: 2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily-stats diff --git a/scripts/data_process/coupled_dataset_utils.py b/scripts/data_process/coupled_dataset_utils.py index 49361a3a46..e707e58c76 100644 --- a/scripts/data_process/coupled_dataset_utils.py +++ b/scripts/data_process/coupled_dataset_utils.py @@ -239,7 +239,10 @@ class CoupledSeaSurfaceConfig: Parameters: surface_flux_window_avg: Configuration for computing windowed time-averaged - surface flux variables from the atmosphere dataset. + surface flux variables from the atmosphere dataset. If its + first_timestamp is None, the first window ends at the ocean + dataset's first timestamp, so that an ensemble whose runs start on + different dates can share this configuration. sst_threshold: Upper threshold for the SST time mean in degrees Kelvin. The resulting sea ice mask is 1 where the time mean SST < the threshold and 0 otherwise. @@ -322,9 +325,13 @@ def apply_sea_ice_window_avg(self, ds: xr.Dataset) -> xr.Dataset: return self.sea_ice_window_avg.get_window_avg(ds) return ds - def apply_surface_flux_window_avg(self, ds: xr.Dataset) -> xr.Dataset: + def apply_surface_flux_window_avg( + self, ds: xr.Dataset, ocean_time: xr.DataArray + ) -> xr.Dataset: logging.info("Computing window-averaged surface fluxes") - return self.surface_flux_window_avg.get_window_avg(ds) + return self.surface_flux_window_avg.get_window_avg( + ds, first_window_end=ocean_time[0].item() + ) @dataclasses.dataclass @@ -801,7 +808,7 @@ def compute_coupled_ocean( ds = config.apply_sea_ice_window_avg(ds) ds = ds.sel({tdim: ocean[tdim]}) - ds = xr.merge([ds, config.apply_surface_flux_window_avg(atmos)]) + ds = xr.merge([ds, config.apply_surface_flux_window_avg(atmos, ocean[tdim])]) ds[tdim] = _make_serializable_time_coord( ds=ocean, tdim=tdim, timedelta=config.timedelta ) diff --git a/scripts/data_process/create_window_avg_dataset.py b/scripts/data_process/create_window_avg_dataset.py index f0bf9bda7a..a7850cab30 100644 --- a/scripts/data_process/create_window_avg_dataset.py +++ b/scripts/data_process/create_window_avg_dataset.py @@ -3,6 +3,7 @@ import pdb import sys +import cftime import click import dacite import pandas as pd @@ -36,6 +37,8 @@ class WindowAvgDatasetConfig: first_timestamp: Optional start timestamp for time slicing, e.g. '2016-01-01T06:00:00'. The 'origin' used in resampling is one step (of size window_timedelta) earlier than first_timestamp. + If None, the first_window_end passed to get_window_avg is used + instead, and one of the two must be given. last_timestamp: Optional final timestamp in case of extra unneeded windows. Used after window averaging but before shifting timestamps. shift_timestamps_to_avg_interval_midpoint: If True, shift time axis labels @@ -51,7 +54,9 @@ class WindowAvgDatasetConfig: time_dim: str = "time" subset_names: list[str] | None = None - def _compute_window_avg(self, ds: xr.Dataset) -> xr.Dataset: + def _compute_window_avg( + self, ds: xr.Dataset, first_window_end: cftime.datetime | None + ) -> xr.Dataset: if self.subset_names is not None: ds = ds[self.subset_names] @@ -68,7 +73,16 @@ def _compute_window_avg(self, ds: xr.Dataset) -> xr.Dataset: ds_time_varying = ds[time_varying_vars] dt = pd.Timedelta(self.window_timedelta).to_pytimedelta() - origin = ds_time_varying.time.sel(time=self.first_timestamp).item() - dt + if self.first_timestamp is not None: + first_window_end = ds_time_varying.time.sel( + time=self.first_timestamp + ).item() + elif first_window_end is None: + raise ValueError( + "Either first_timestamp must be configured or first_window_end " + "must be passed to get_window_avg." + ) + origin = first_window_end - dt ds_time_varying_avg = ( ds_time_varying.resample( time=self.window_timedelta, closed="right", label="right", origin=origin @@ -112,9 +126,18 @@ def _shift_timestamps(self, ds: xr.Dataset): ) return ds - def get_window_avg(self, ds: xr.Dataset) -> xr.Dataset: + def get_window_avg( + self, ds: xr.Dataset, first_window_end: cftime.datetime | None = None + ) -> xr.Dataset: + """Compute the window average of ds. + + Args: + ds: Input dataset. + first_window_end: Right edge (and label) of the first window, used + only when first_timestamp is not configured. + """ logging.info(f"Input time coord:\n {str(ds[self.time_dim])}") - ds = self._compute_window_avg(ds) + ds = self._compute_window_avg(ds, first_window_end) ds = self._shift_timestamps(ds) return ds diff --git a/scripts/data_process/fs_utils.py b/scripts/data_process/fs_utils.py index ec7225381c..81209a62c3 100644 --- a/scripts/data_process/fs_utils.py +++ b/scripts/data_process/fs_utils.py @@ -4,7 +4,11 @@ use a single code path for local paths and remote URLs (e.g. ``gs://``). """ +import os +import tempfile + import fsspec +import xarray as xr from fsspec.implementations.local import LocalFileSystem @@ -54,3 +58,33 @@ def is_local(path: str) -> bool: """ fs, _ = _fs_and_path(path) return isinstance(fs, LocalFileSystem) + + +def load_netcdf(path: str) -> xr.Dataset: + """Load a small netCDF file, e.g. a stats file, fully into memory. + + A remote file is first downloaded to a temporary local file in a single + read. Parsing a netCDF4 (HDF5) file directly from a remote file object + instead issues a range request for each of its many small, scattered + metadata reads, which for a file with hundreds of variables on GCS is + tens of times slower. Reading from a local file rather than an in-memory + buffer also lets the default netCDF4 engine be used, so h5netcdf is not + required. + + The returned dataset's encoding["source"] is set to path. + """ + fs, fs_path = _fs_and_path(path) + if isinstance(fs, LocalFileSystem): + ds = _load_local_netcdf(fs_path) + else: + with tempfile.TemporaryDirectory() as tmpdir: + local_path = os.path.join(tmpdir, os.path.basename(fs_path)) + fs.get_file(fs_path, local_path) + ds = _load_local_netcdf(local_path) + ds.encoding["source"] = path + return ds + + +def _load_local_netcdf(path: str) -> xr.Dataset: + with xr.open_dataset(path, decode_timedelta=False) as ds: + return ds.load() diff --git a/scripts/data_process/merge_stats.py b/scripts/data_process/merge_stats.py index e7d639ab95..7370951377 100644 --- a/scripts/data_process/merge_stats.py +++ b/scripts/data_process/merge_stats.py @@ -5,10 +5,9 @@ import tempfile import dacite -import fsspec import xarray as xr import yaml -from fs_utils import makedirs +from fs_utils import load_netcdf, makedirs from get_stats import copy STATS_NC_FILE_NAMES = [ @@ -50,12 +49,9 @@ class MergeStatsConfig: exclude_names: list[str] = dataclasses.field(default_factory=list) def open_input_datasets(self, fname: str): - datasets = [] - for path in self.input_directories: - with fsspec.open(os.path.join(path, fname)) as file: - ds = xr.open_dataset(file, decode_timedelta=False).load() - datasets.append(ds) - return datasets + return [ + load_netcdf(os.path.join(path, fname)) for path in self.input_directories + ] def merge_stats(config: MergeStatsConfig): diff --git a/scripts/data_process/test_create_coupled_datasets.py b/scripts/data_process/test_create_coupled_datasets.py index a09c2f680c..f29defd343 100644 --- a/scripts/data_process/test_create_coupled_datasets.py +++ b/scripts/data_process/test_create_coupled_datasets.py @@ -56,10 +56,10 @@ def _field(times, offset=0.0, seed=0): ) -def _write_input_zarrs(input_dir): - atmos_times = _times("2000-01-01T06:00:00", N_ATMOS_TIMES, "6h") +def _write_input_zarrs(input_dir, year=2000): + atmos_times = _times(f"{year:04d}-01-01T06:00:00", N_ATMOS_TIMES, "6h") # 5-daily snapshots at the ends of the two 120h windows - ocean_times = _times(FIRST_WINDOW_END, 2, "120h") + ocean_times = _times(f"{year:04d}-01-06T00:00:00", 2, "120h") atmos = xr.Dataset( { @@ -79,8 +79,8 @@ def _write_input_zarrs(input_dir): ) ocean["sea_surface_fraction"] = _field(ocean_times, seed=8).isel(time=0, drop=True) - atmos_path = str(input_dir / "atmosphere.zarr") - ocean_path = str(input_dir / "ocean.zarr") + atmos_path = str(input_dir / f"atmosphere-{year:04d}.zarr") + ocean_path = str(input_dir / f"ocean-{year:04d}.zarr") atmos.to_zarr(atmos_path) ocean.to_zarr(ocean_path) return atmos_path, ocean_path @@ -411,3 +411,61 @@ def test_exclude_runs_on_single_dataset_raises(tmp_path): config = _make_config(tmp_path) with pytest.raises(ValueError, match="only supported for ensemble"): dataclasses.replace(config, stats=CoupledStatsConfig(exclude_runs=["run-a"])) + + +def test_ensemble_flux_windows_follow_each_runs_ocean_times(tmp_path): + """Runs starting on different dates share one surface_flux_window_avg; with + no first_timestamp configured, each run's flux windows end on its own + ocean timesteps.""" + single = _make_config(tmp_path) + years = {"run-a": 2001, "run-b": 2003} + runs = {} + input_paths = {} + for name, year in years.items(): + atmos_path, ocean_path = _write_input_zarrs(tmp_path / "inputs", year=year) + input_paths[name] = (atmos_path, ocean_path) + runs[name] = EnsembleRunConfig( + atmosphere=CoupledInputDatasetConfig( + zarr_path=atmos_path, time_chunk_size=N_ATMOS_TIMES + ), + ocean=CoupledInputDatasetConfig(zarr_path=ocean_path, time_chunk_size=2), + ) + coupled_datasets = dataclasses.replace( + single.coupled_datasets, + coupled_sea_ice=CoupledSeaIceConfig( + use_atmosphere_sea_ice_fraction_fallback=True + ), + coupled_sea_surface=dataclasses.replace( + single.coupled_datasets.coupled_sea_surface, + surface_flux_window_avg=WindowAvgDatasetConfig( + window_timedelta="120h", subset_names=["latent_heat_flux"] + ), + ), + ) + config = CreateCoupledDatasetsConfig( + version=single.version, + family_name=single.family_name, + output_directory=single.output_directory, + coupled_datasets=coupled_datasets, + input_datasets=InputEnsembleConfig( + runs=runs, + stats=single.input_datasets.stats, + climate_data_type=single.input_datasets.climate_data_type, + ), + ) + config.write_coupled_datasets(debug=False, subsample=False) + + ensemble_dir = tmp_path / "outputs" / "v1-synthetic" + for name, (atmos_path, ocean_path) in input_paths.items(): + atmos = xr.open_zarr(atmos_path) + ocean_in = xr.open_zarr(ocean_path) + ocean_out = xr.open_zarr(str(ensemble_dir / f"{name}-ocean.zarr")) + xr.testing.assert_equal(ocean_out.time, ocean_in.time) + # each 120h window holds 20 consecutive 6-hourly atmosphere samples + # and is labeled at its right edge + expected = ( + atmos["latent_heat_flux"] + .coarsen(time=N_ATMOS_TIMES // 2, coord_func="max") + .mean() + ) + xr.testing.assert_allclose(ocean_out["latent_heat_flux"], expected) diff --git a/scripts/data_process/test_fs_utils.py b/scripts/data_process/test_fs_utils.py index 97488bf705..745b9184f2 100644 --- a/scripts/data_process/test_fs_utils.py +++ b/scripts/data_process/test_fs_utils.py @@ -1,7 +1,9 @@ import os import fsspec -from fs_utils import is_dir, is_local, makedirs, path_exists +import numpy as np +import xarray as xr +from fs_utils import is_dir, is_local, load_netcdf, makedirs, path_exists from fsspec.implementations.memory import MemoryFileSystem @@ -63,3 +65,20 @@ def test_path_exists_sees_write_made_after_an_earlier_miss(): f.write(b"data") assert path_exists(store) + + +def test_load_netcdf_from_local_and_remote_paths(tmp_path): + ds = xr.Dataset( + {f"var_{i}": ("lat", np.arange(3.0) + i, {"units": "K"}) for i in range(5)}, + coords={"lat": [-45.0, 0.0, 45.0]}, + attrs={"input_samples": 1460}, + ) + local_path = str(tmp_path / "centering.nc") + ds.to_netcdf(local_path) + remote_path = "memory://stats/centering.nc" + fsspec.filesystem("memory").put_file(local_path, remote_path) + + for path in [local_path, remote_path]: + loaded = load_netcdf(path) + xr.testing.assert_identical(loaded, ds) + assert loaded.encoding["source"] == path diff --git a/scripts/gfdl_om4/README.md b/scripts/gfdl_om4/README.md index d821d97d03..252c71aaac 100644 --- a/scripts/gfdl_om4/README.md +++ b/scripts/gfdl_om4/README.md @@ -10,7 +10,7 @@ and writes one templated, sharded zarr v3 store. | Pipeline | Package | Configs | Make targets | Output | | --- | --- | --- | --- | --- | -| OM4 | `pipeline/om4/` | `configs/om4-*.yaml`, `configs/cm4-*-1daily.yaml` | `om4.mk` (unprefixed) | OM4/CM4 ocean stores, one per config | +| OM4 | `pipeline/om4/` | `configs/om4-*.yaml`, `configs/cm4-*-1daily.yaml`, `configs/cm4-abrupt-4xCO2-*.yaml` | `om4.mk` (unprefixed) | OM4/CM4 ocean stores, one per config (per member for an ensemble config) | | randco2 | `pipeline/randco2/` | `configs/cm4-like-am4-randco2-*.yaml` | `randco2.mk` (`randco2_*`) | CM4-like-AM4 random-CO2 sea-surface stores, one per member | Both run from `scripts/gfdl_om4/` through the one `Makefile`, which holds @@ -48,7 +48,11 @@ source chunk is fetched and decoded once rather than once per timestep it holds. OM4 reads 3D variables one level at a time. Before consolidating into output shards, read chunks are split to `gcd(read width, shard)` (`shard_aligned_chunk_size`), which puts every chunk boundary on a shard -boundary; where the read width divides the shard the split is a no-op. The +boundary; where the read width divides the shard the split is a no-op. An +OM4 stream with `time_block_mean_factor` is read at `lcm(source width, +factor)` source steps per chunk, so each read chunk holds whole source +chunks and whole averaging blocks; a beam step averages each chunk's blocks +and re-keys it in output timesteps before the per-chunk transforms. The output layout (time chunk and shard from the config's `output`) does not depend on the read width. @@ -107,7 +111,7 @@ Contents, under `pipeline/om4/`: - `pipeline/om4/config.py`: YAML→dataclass configuration (dacite). Stream options cover source-dim renaming (e.g. ice-model `xT/yT/xB/yB` onto the ocean `xh/yh/xq/yq` conventions), time subsampling to the shared snapshot - instants, full-cell (per-total-cell-area) regridding for selected + instants or block-mean time coarsening onto them, full-cell (per-total-cell-area) regridding for selected variables, and named postprocess transforms. - `pipeline/om4/postprocess.py`: named post-regrid transforms selectable per stream — Kelvin `sst`, `hfds_total_area`, and the sea-ice conventions @@ -121,8 +125,9 @@ Contents, under `pipeline/om4/`: See `run._rotate_pairs` for how the flagged faces and the wall-zero fill keep every velocity output on the tracer wetmask footprint (`mask_k`). -- `pipeline/om4/check_output.py`, `pipeline/om4/check_wetmask_equivalence.py`: - post-run and cross-config checks used by `om4.mk`. +- `pipeline/om4/check_output.py`, `pipeline/om4/check_wetmask_equivalence.py`, + `pipeline/om4/check_member_wetmask_equivalence.py`: post-run, + cross-config and cross-member checks used by `om4.mk`. ### Setup @@ -136,7 +141,8 @@ make generate_face_masks # per-simulation remap-born-zero face masks Face-mask artifacts are per-simulation (see `pipeline/om4/face_masks.py`): each `generate_face_masks_*` target scans that source's first year and checks the flagged surface-face counts against the independent census baked into the -target. +target. The abrupt-4xCO2 ensemble shares one artifact, scanned from its first +member (`make generate_face_masks_abrupt_4xco2`). ### Configs @@ -147,6 +153,17 @@ the permanent inputs prefix, sources are read from `gs://vcm-ml-intermediate/`; nothing consumed by a production run lives under `vcm-ml-scratch`. +The CM4 abrupt-4xCO2 ensemble has one 5-daily and one 1-daily config per +resolution, each carrying a `{member}` placeholder (as in randco2); pass +`--member` (or `MEMBER=` to the make targets) to resolve it, and `om4.mk` +holds the 36-member list. Members are named by their short `ic_NNNN` names +(the configs add the source's `abrupt-4xCO2-` prefix), so each output store +is `ic_NNNN.zarr`, matching the atmosphere stores from `scripts/data_process`. +Its sources are raw daily/6-hourly stores, so the 5-daily configs subsample +snapshots (`time_subsample_stride`) and average the daily-mean fluxes over +5-day blocks (`time_block_mean_factor`) in the pipeline; the 1-daily configs +only subsample the 6-hourly ice streams. + ### Running #### Smoke tests @@ -169,6 +186,11 @@ make smoke_test_picontrol_1deg # or any single config - `make smoke_test_repeat_fails` — a repeat run against an existing output store must refuse to initialize into it. +The abrupt-4xCO2 ensemble has its own set: `make smoke_tests_abrupt_4xco2` +smoke-tests one member at both resolutions, checks all 36 members derive the +same wetmask, and runs the repeat-run check; `make +smoke_tests_abrupt_4xco2_1daily` does the same for the 1-daily configs. + The pipeline can also be invoked directly, with any beam pipeline options after the script's own arguments: @@ -191,6 +213,9 @@ python -m pipeline.om4.run --config configs/om4-picontrol-1deg-5daily.yaml \ ``` make push_dataflow make dataflow_picontrol_1deg # one target per config + make dataflow_abrupt_4xco2_1deg_ic_0001 # one per ensemble member + make dataflow_abrupt_4xco2_1deg_all # or all 36 + make dataflow_abrupt_4xco2_1deg_1daily_ic_0001 # 1-daily, likewise ``` `make dataflow*` first checks that `$(IMAGE_NAME)` exists in Artifact diff --git a/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml new file mode 100644 index 0000000000..b9c6a9d006 --- /dev/null +++ b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml @@ -0,0 +1,119 @@ +# 1-daily 1-degree ocean and sea-ice (Samudra) store for one member of the +# CM4 abrupt-4xCO2 ensemble. `{member}` is filled from --member (see +# pipeline/config_io.py), e.g. ic_0001. +target_grid: F90 +weights_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/regridding-weights/om4-tripolar-0.25deg-to-F90 +expected_level_count: 19 +shift_timestamps_to_avg_interval_midpoint: false +output: + path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily/{member}.zarr + time_chunk_size: 1 + time_shard_size: 365 +wetmask: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variable: thetao +statics: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_static.zarr + variables: + - deptho + - hfgeou +streams: + # Daily ocean snapshots. + - name: snapshot_ocean + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variables: + - thetao + - so + - uo + - vo + - tos + - zos + - SSH + - mlotst + - rhoinsitu + rotated_pairs: + - [uo, vo] + # These sources carry coastal staggered velocity faces that are valid + # with value exactly 0.0 where the vertical remap unmasked land columns; + # the artifact flags them so they are not averaged into coastal centers. + # One artifact for the ensemble (make generate_face_masks_abrupt_4xco2). + face_mask_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/face-masks/cm4-abrupt-4xCO2-2026-10-16 + postprocess: + - name: kelvin_sst + sources: + celsius_sst: tos + # 1-day-mean surface fluxes/forcings, whose end-of-interval time labels + # coincide with the snapshot instants (enforced by the time-alignment + # assertion). + - name: flux + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily.zarr + variables: + - hfds + - hflso + - hfsso + - rlntds + - rsntds + - hfsifrazil + - hfevapds + - hfrainds + - hfrunoffds + - wfo + - prlq + - prsn + - evs + - sfdsi + - tauuo + - tauvo + - ustar + rotated_pairs: + - [tauuo, tauvo] + postprocess: + - hfds_total_area + # Ice-model true snapshots, subsampled from 6-hourly to the daily + # snapshot instants (every 4th step). The ice SSH is the sea surface as + # seen by the ice model, distinct from the ocean-store SSH and zos. + - name: ice_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly_snap.zarr + variables: + - sea_ice_fraction + - EXT + - HI + - HS + - UI + - VI + - TS + - TSN + - SST + - SSS + - simass + - sisnmass + - SSH + dim_renaming: + xT: xh + yT: yh + xB: xq + yB: yq + time_subsample_stride: 4 + rotated_pairs: + - [UI, VI] + renaming: + SSH: SSH_ice + sea_ice_fraction: ocean_sea_ice_fraction + full_cell_variables: + - sea_ice_fraction + postprocess: + - sea_ice + # Ice-model 6-hourly means at the snapshot instants (each mean covers the + # 6 hours ending at the instant), treated as snapshot estimates. + - name: ice_mean_as_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly.zarr + variables: + - T1 + - T2 + - T3 + - T4 + - T_bulkice + dim_renaming: + xT: xh + yT: yh + time_subsample_stride: 4 diff --git a/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily.yaml b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily.yaml new file mode 100644 index 0000000000..9f1e204a87 --- /dev/null +++ b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily.yaml @@ -0,0 +1,129 @@ +# 5-daily 1-degree ocean and sea-ice (Samudra) store for one member of the +# CM4 abrupt-4xCO2 ensemble. `{member}` is filled from --member (see +# pipeline/config_io.py), e.g. ic_0001. +# +# Unlike the 1pctCO2 and piControl 5-daily configs, whose sources were +# pre-subset / pre-coarsened to 5-daily, these sources are the raw daily and +# 6-hourly stores: snapshots are subsampled to every 5th day and the daily-mean +# fluxes are averaged over 5-day blocks in the pipeline. +target_grid: F90 +weights_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/regridding-weights/om4-tripolar-0.25deg-to-F90 +expected_level_count: 19 +shift_timestamps_to_avg_interval_midpoint: false +output: + path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily/{member}.zarr + time_chunk_size: 1 + time_shard_size: 365 +wetmask: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variable: thetao +statics: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_static.zarr + variables: + - deptho + - hfgeou +streams: + # Daily ocean snapshots, subsampled to the 5-daily instants (every 5th + # step). Every member's daily stores span a multiple of 5 days (90 or 365), + # so the block ends coincide with those of the 5-day flux means and the + # 20-step ice subsamples. + - name: snapshot_ocean + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variables: + - thetao + - so + - uo + - vo + - tos + - zos + - SSH + - mlotst + - rhoinsitu + time_subsample_stride: 5 + rotated_pairs: + - [uo, vo] + # These sources carry coastal staggered velocity faces that are valid + # with value exactly 0.0 where the vertical remap unmasked land columns; + # the artifact flags them so they are not averaged into coastal centers. + # One artifact for the ensemble (make generate_face_masks_abrupt_4xco2). + face_mask_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/face-masks/cm4-abrupt-4xCO2-2026-10-16 + postprocess: + - name: kelvin_sst + sources: + celsius_sst: tos + # Daily-mean surface fluxes/forcings averaged over 5-day blocks, labeled by + # the end of each block so they coincide with the snapshot instants + # (enforced by the time-alignment assertion). + - name: flux + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily.zarr + variables: + - hfds + - hflso + - hfsso + - rlntds + - rsntds + - hfsifrazil + - hfevapds + - hfrainds + - hfrunoffds + - wfo + - prlq + - prsn + - evs + - sfdsi + - tauuo + - tauvo + - ustar + time_block_mean_factor: 5 + rotated_pairs: + - [tauuo, tauvo] + postprocess: + - hfds_total_area + # Ice-model true snapshots, subsampled from 6-hourly to the 5-daily + # snapshot instants (every 20th step). The ice SSH is the sea surface as + # seen by the ice model, distinct from the ocean-store SSH and zos. + - name: ice_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly_snap.zarr + variables: + - sea_ice_fraction + - EXT + - HI + - HS + - UI + - VI + - TS + - TSN + - SST + - SSS + - simass + - sisnmass + - SSH + dim_renaming: + xT: xh + yT: yh + xB: xq + yB: yq + time_subsample_stride: 20 + rotated_pairs: + - [UI, VI] + renaming: + SSH: SSH_ice + sea_ice_fraction: ocean_sea_ice_fraction + full_cell_variables: + - sea_ice_fraction + postprocess: + - sea_ice + # Ice-model 6-hourly means at the snapshot instants (each mean covers the + # 6 hours ending at the instant), treated as snapshot estimates. + - name: ice_mean_as_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly.zarr + variables: + - T1 + - T2 + - T3 + - T4 + - T_bulkice + dim_renaming: + xT: xh + yT: yh + time_subsample_stride: 20 diff --git a/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml new file mode 100644 index 0000000000..6693a2a177 --- /dev/null +++ b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml @@ -0,0 +1,119 @@ +# 1-daily 4-degree ocean and sea-ice (Samudra) store for one member of the +# CM4 abrupt-4xCO2 ensemble. `{member}` is filled from --member (see +# pipeline/config_io.py), e.g. ic_0001. +target_grid: F22.5 +weights_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/regridding-weights/om4-tripolar-0.25deg-to-F22.5 +expected_level_count: 19 +shift_timestamps_to_avg_interval_midpoint: false +output: + path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily/{member}.zarr + time_chunk_size: 1 + time_shard_size: 365 +wetmask: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variable: thetao +statics: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_static.zarr + variables: + - deptho + - hfgeou +streams: + # Daily ocean snapshots. + - name: snapshot_ocean + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variables: + - thetao + - so + - uo + - vo + - tos + - zos + - SSH + - mlotst + - rhoinsitu + rotated_pairs: + - [uo, vo] + # These sources carry coastal staggered velocity faces that are valid + # with value exactly 0.0 where the vertical remap unmasked land columns; + # the artifact flags them so they are not averaged into coastal centers. + # One artifact for the ensemble (make generate_face_masks_abrupt_4xco2). + face_mask_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/face-masks/cm4-abrupt-4xCO2-2026-10-16 + postprocess: + - name: kelvin_sst + sources: + celsius_sst: tos + # 1-day-mean surface fluxes/forcings, whose end-of-interval time labels + # coincide with the snapshot instants (enforced by the time-alignment + # assertion). + - name: flux + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily.zarr + variables: + - hfds + - hflso + - hfsso + - rlntds + - rsntds + - hfsifrazil + - hfevapds + - hfrainds + - hfrunoffds + - wfo + - prlq + - prsn + - evs + - sfdsi + - tauuo + - tauvo + - ustar + rotated_pairs: + - [tauuo, tauvo] + postprocess: + - hfds_total_area + # Ice-model true snapshots, subsampled from 6-hourly to the daily + # snapshot instants (every 4th step). The ice SSH is the sea surface as + # seen by the ice model, distinct from the ocean-store SSH and zos. + - name: ice_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly_snap.zarr + variables: + - sea_ice_fraction + - EXT + - HI + - HS + - UI + - VI + - TS + - TSN + - SST + - SSS + - simass + - sisnmass + - SSH + dim_renaming: + xT: xh + yT: yh + xB: xq + yB: yq + time_subsample_stride: 4 + rotated_pairs: + - [UI, VI] + renaming: + SSH: SSH_ice + sea_ice_fraction: ocean_sea_ice_fraction + full_cell_variables: + - sea_ice_fraction + postprocess: + - sea_ice + # Ice-model 6-hourly means at the snapshot instants (each mean covers the + # 6 hours ending at the instant), treated as snapshot estimates. + - name: ice_mean_as_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly.zarr + variables: + - T1 + - T2 + - T3 + - T4 + - T_bulkice + dim_renaming: + xT: xh + yT: yh + time_subsample_stride: 4 diff --git a/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily.yaml b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily.yaml new file mode 100644 index 0000000000..8005a9d8f3 --- /dev/null +++ b/scripts/gfdl_om4/configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily.yaml @@ -0,0 +1,129 @@ +# 5-daily 4-degree ocean and sea-ice (Samudra) store for one member of the +# CM4 abrupt-4xCO2 ensemble. `{member}` is filled from --member (see +# pipeline/config_io.py), e.g. ic_0001. +# +# Unlike the 1pctCO2 and piControl 5-daily configs, whose sources were +# pre-subset / pre-coarsened to 5-daily, these sources are the raw daily and +# 6-hourly stores: snapshots are subsampled to every 5th day and the daily-mean +# fluxes are averaged over 5-day blocks in the pipeline. +target_grid: F22.5 +weights_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/regridding-weights/om4-tripolar-0.25deg-to-F22.5 +expected_level_count: 19 +shift_timestamps_to_avg_interval_midpoint: false +output: + path: gs://vcm-ml-intermediate/2026-10-19-cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily/{member}.zarr + time_chunk_size: 1 + time_shard_size: 365 +wetmask: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variable: thetao +statics: + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_static.zarr + variables: + - deptho + - hfgeou +streams: + # Daily ocean snapshots, subsampled to the 5-daily instants (every 5th + # step). Every member's daily stores span a multiple of 5 days (90 or 365), + # so the block ends coincide with those of the 5-day flux means and the + # 20-step ice subsamples. + - name: snapshot_ocean + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily_snap.zarr + variables: + - thetao + - so + - uo + - vo + - tos + - zos + - SSH + - mlotst + - rhoinsitu + time_subsample_stride: 5 + rotated_pairs: + - [uo, vo] + # These sources carry coastal staggered velocity faces that are valid + # with value exactly 0.0 where the vertical remap unmasked land columns; + # the artifact flags them so they are not averaged into coastal centers. + # One artifact for the ensemble (make generate_face_masks_abrupt_4xco2). + face_mask_url: gs://vcm-ml-intermediate/2026-07-15-gfdl-om4-pipeline-inputs/face-masks/cm4-abrupt-4xCO2-2026-10-16 + postprocess: + - name: kelvin_sst + sources: + celsius_sst: tos + # Daily-mean surface fluxes/forcings averaged over 5-day blocks, labeled by + # the end of each block so they coincide with the snapshot instants + # (enforced by the time-alignment assertion). + - name: flux + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ocean_1daily.zarr + variables: + - hfds + - hflso + - hfsso + - rlntds + - rsntds + - hfsifrazil + - hfevapds + - hfrainds + - hfrunoffds + - wfo + - prlq + - prsn + - evs + - sfdsi + - tauuo + - tauvo + - ustar + time_block_mean_factor: 5 + rotated_pairs: + - [tauuo, tauvo] + postprocess: + - hfds_total_area + # Ice-model true snapshots, subsampled from 6-hourly to the 5-daily + # snapshot instants (every 20th step). The ice SSH is the sea surface as + # seen by the ice model, distinct from the ocean-store SSH and zos. + - name: ice_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly_snap.zarr + variables: + - sea_ice_fraction + - EXT + - HI + - HS + - UI + - VI + - TS + - TSN + - SST + - SSS + - simass + - sisnmass + - SSH + dim_renaming: + xT: xh + yT: yh + xB: xq + yB: yq + time_subsample_stride: 20 + rotated_pairs: + - [UI, VI] + renaming: + SSH: SSH_ice + sea_ice_fraction: ocean_sea_ice_fraction + full_cell_variables: + - sea_ice_fraction + postprocess: + - sea_ice + # Ice-model 6-hourly means at the snapshot instants (each mean covers the + # 6 hours ending at the instant), treated as snapshot estimates. + - name: ice_mean_as_snapshot + store: gs://vcm-ml-raw-flexible-retention/2026-10-16-CM4-ocean-and-sea-ice-abrupt-4xCO2-ensemble/zarrs/abrupt-4xCO2-{member}/ice_6hourly.zarr + variables: + - T1 + - T2 + - T3 + - T4 + - T_bulkice + dim_renaming: + xT: xh + yT: yh + time_subsample_stride: 20 diff --git a/scripts/gfdl_om4/om4.mk b/scripts/gfdl_om4/om4.mk index a59eb54bee..1e757a3b3f 100644 --- a/scripts/gfdl_om4/om4.mk +++ b/scripts/gfdl_om4/om4.mk @@ -30,8 +30,13 @@ generate_face_masks_1pctco2: CONFIG ?= configs/om4-picontrol-1deg-5daily.yaml +# Ensemble configs carry a {member} placeholder in their URLs; set MEMBER to +# resolve it. Left empty for single-simulation configs. +MEMBER ?= +MEMBER_FLAG = $(if $(MEMBER),--member $(MEMBER)) + SMOKE_OUTPUT_ROOT ?= gs://vcm-ml-scratch/jamesd/gfdl-om4-test -SMOKE_OUTPUT_PATH = $(SMOKE_OUTPUT_ROOT)/smoke-$(notdir $(basename $(CONFIG))).zarr +SMOKE_OUTPUT_PATH = $(SMOKE_OUTPUT_ROOT)/smoke-$(notdir $(basename $(CONFIG)))$(if $(MEMBER),-$(MEMBER)).zarr # A few timesteps of $(CONFIG) through the DirectRunner into a throwaway # store: every chunk's footprint must equal the wetmask (the run fails @@ -40,13 +45,13 @@ SMOKE_OUTPUT_PATH = $(SMOKE_OUTPUT_ROOT)/smoke-$(notdir $(basename $(CONFIG))).z smoke_test: -gsutil -m -q rm -r $(SMOKE_OUTPUT_PATH) 2> /dev/null conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) python -m pipeline.om4.run \ - --config $(CONFIG) \ + --config $(CONFIG) $(MEMBER_FLAG) \ --num-timesteps $(SMOKE_NUM_TIMESTEPS) \ --output-path $(SMOKE_OUTPUT_PATH) \ --runner DirectRunner \ --job_server_timeout 3600 conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) python -m pipeline.om4.check_output \ - --config $(CONFIG) \ + --config $(CONFIG) $(MEMBER_FLAG) \ --output-path $(SMOKE_OUTPUT_PATH) \ --expected-timesteps $(SMOKE_NUM_TIMESTEPS) @@ -79,7 +84,7 @@ smoke_test_1pctco2_4deg_1daily: smoke_test_repeat_fails: gsutil -q ls $(SMOKE_OUTPUT_PATH)/zarr.json > /dev/null 2>&1 || $(MAKE) smoke_test ! conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) python -m pipeline.om4.run \ - --config $(CONFIG) \ + --config $(CONFIG) $(MEMBER_FLAG) \ --num-timesteps $(SMOKE_NUM_TIMESTEPS) \ --output-path $(SMOKE_OUTPUT_PATH) \ --runner DirectRunner \ @@ -111,7 +116,7 @@ smoke_tests_1daily: check_wetmask_equivalence_1daily \ dataflow: check_dataflow_image SDK_CONTAINER_IMAGE=$(IMAGE_NAME) \ conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) \ - ./run-dataflow.sh om4 DataflowRunner $(CONFIG) + ./run-dataflow.sh om4 DataflowRunner $(CONFIG) $(MEMBER_FLAG) dataflow_picontrol_1deg: $(MAKE) dataflow CONFIG=configs/om4-picontrol-1deg-5daily.yaml @@ -137,6 +142,88 @@ dataflow_1pctco2_1deg_1daily: dataflow_1pctco2_4deg_1daily: $(MAKE) dataflow CONFIG=configs/cm4-1pctCO2-4deg-1daily.yaml +# CM4 abrupt-4xCO2 ensemble: one 5-daily and one 1-daily store per member and +# resolution, from one config per resolution and frequency. Members are 90-day +# or 1-year runs started from different piControl states. +ABRUPT_4XCO2_MEMBERS = $(addprefix ic_,$(shell seq -f %04g 1 36)) +ABRUPT_4XCO2_1DEG_CONFIG = configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-5daily.yaml +ABRUPT_4XCO2_4DEG_CONFIG = configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-5daily.yaml +ABRUPT_4XCO2_1DEG_1DAILY_CONFIG = configs/cm4-abrupt-4xCO2-ensemble-ocean-1deg-1daily.yaml +ABRUPT_4XCO2_4DEG_1DAILY_CONFIG = configs/cm4-abrupt-4xCO2-ensemble-ocean-4deg-1daily.yaml +ABRUPT_4XCO2_SMOKE_MEMBER ?= $(firstword $(ABRUPT_4XCO2_MEMBERS)) + +# One artifact for the whole ensemble, scanned from the first (1-year) +# member. The ensemble-wide wetmask check below and the run-time staleness +# assertion guard its use for the other members. +generate_face_masks_abrupt_4xco2: + conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) python -m pipeline.om4.face_masks \ + --config $(ABRUPT_4XCO2_1DEG_CONFIG) \ + --member $(firstword $(ABRUPT_4XCO2_MEMBERS)) \ + --stream snapshot_ocean \ + --output-url $(FACE_MASKS_URL_ROOT)/cm4-abrupt-4xCO2-2026-10-16 \ + --start-time 0301-01-06 \ + --end-time 0302-01-01 \ + --expected-surface-count-u 16960 \ + --expected-surface-count-v 16694 + +check_wetmask_equivalence_abrupt_4xco2: + conda run --no-capture-output -n $(LOCAL_ENVIRONMENT) python -m pipeline.om4.check_member_wetmask_equivalence \ + --config $(ABRUPT_4XCO2_1DEG_CONFIG) \ + --members $(ABRUPT_4XCO2_MEMBERS) + +smoke_test_abrupt_4xco2_1deg: + $(MAKE) smoke_test CONFIG=$(ABRUPT_4XCO2_1DEG_CONFIG) MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +smoke_test_abrupt_4xco2_4deg: + $(MAKE) smoke_test CONFIG=$(ABRUPT_4XCO2_4DEG_CONFIG) MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +smoke_tests_abrupt_4xco2: check_wetmask_equivalence_abrupt_4xco2 \ + smoke_test_abrupt_4xco2_1deg smoke_test_abrupt_4xco2_4deg + $(MAKE) smoke_test_repeat_fails CONFIG=$(ABRUPT_4XCO2_1DEG_CONFIG) \ + MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +smoke_test_abrupt_4xco2_1deg_1daily: + $(MAKE) smoke_test CONFIG=$(ABRUPT_4XCO2_1DEG_1DAILY_CONFIG) MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +smoke_test_abrupt_4xco2_4deg_1daily: + $(MAKE) smoke_test CONFIG=$(ABRUPT_4XCO2_4DEG_1DAILY_CONFIG) MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +# The 1-daily configs read the same wetmask stores as the 5-daily ones, so +# the 5-daily ensemble wetmask check covers them. +smoke_tests_abrupt_4xco2_1daily: check_wetmask_equivalence_abrupt_4xco2 \ + smoke_test_abrupt_4xco2_1deg_1daily smoke_test_abrupt_4xco2_4deg_1daily + $(MAKE) smoke_test_repeat_fails CONFIG=$(ABRUPT_4XCO2_1DEG_1DAILY_CONFIG) \ + MEMBER=$(ABRUPT_4XCO2_SMOKE_MEMBER) + +# Per-member production launches, e.g. +# make dataflow_abrupt_4xco2_1deg_ic_0001 or +# make dataflow_abrupt_4xco2_1deg_1daily_ic_0001. Each is an independent +# Dataflow job writing its own store under the config's dated parent prefix. +# A 1-daily target also matches the 5-daily pattern of its resolution; make +# picks the pattern with the shortest stem, which is the 1-daily one. +dataflow_abrupt_4xco2_1deg_%: + $(MAKE) dataflow CONFIG=$(ABRUPT_4XCO2_1DEG_CONFIG) MEMBER=$* + +dataflow_abrupt_4xco2_4deg_%: + $(MAKE) dataflow CONFIG=$(ABRUPT_4XCO2_4DEG_CONFIG) MEMBER=$* + +dataflow_abrupt_4xco2_1deg_1daily_%: + $(MAKE) dataflow CONFIG=$(ABRUPT_4XCO2_1DEG_1DAILY_CONFIG) MEMBER=$* + +dataflow_abrupt_4xco2_4deg_1daily_%: + $(MAKE) dataflow CONFIG=$(ABRUPT_4XCO2_4DEG_1DAILY_CONFIG) MEMBER=$* + +dataflow_abrupt_4xco2_1deg_all: $(addprefix dataflow_abrupt_4xco2_1deg_,$(ABRUPT_4XCO2_MEMBERS)) + +dataflow_abrupt_4xco2_4deg_all: $(addprefix dataflow_abrupt_4xco2_4deg_,$(ABRUPT_4XCO2_MEMBERS)) + +dataflow_abrupt_4xco2_1deg_1daily_all: $(addprefix dataflow_abrupt_4xco2_1deg_1daily_,$(ABRUPT_4XCO2_MEMBERS)) + +dataflow_abrupt_4xco2_4deg_1daily_all: $(addprefix dataflow_abrupt_4xco2_4deg_1daily_,$(ABRUPT_4XCO2_MEMBERS)) + +# The per-member targets are deliberately absent: make skips implicit-rule +# search for a phony target, which would leave the pattern rules above with +# nothing to match. .PHONY: generate_face_masks \ generate_face_masks_picontrol generate_face_masks_1pctco2 smoke_test \ smoke_test_picontrol_1deg smoke_test_picontrol_4deg \ @@ -148,4 +235,10 @@ dataflow_1pctco2_4deg_1daily: dataflow dataflow_picontrol_1deg \ dataflow_picontrol_4deg dataflow_1pctco2_1deg dataflow_1pctco2_4deg \ dataflow_picontrol_1deg_1daily dataflow_picontrol_4deg_1daily \ - dataflow_1pctco2_1deg_1daily dataflow_1pctco2_4deg_1daily + dataflow_1pctco2_1deg_1daily dataflow_1pctco2_4deg_1daily \ + generate_face_masks_abrupt_4xco2 check_wetmask_equivalence_abrupt_4xco2 \ + smoke_test_abrupt_4xco2_1deg smoke_test_abrupt_4xco2_4deg \ + smoke_tests_abrupt_4xco2 dataflow_abrupt_4xco2_1deg_all \ + dataflow_abrupt_4xco2_4deg_all smoke_test_abrupt_4xco2_1deg_1daily \ + smoke_test_abrupt_4xco2_4deg_1daily smoke_tests_abrupt_4xco2_1daily \ + dataflow_abrupt_4xco2_1deg_1daily_all dataflow_abrupt_4xco2_4deg_1daily_all diff --git a/scripts/gfdl_om4/pipeline/dims.py b/scripts/gfdl_om4/pipeline/dims.py new file mode 100644 index 0000000000..3f8f501ce3 --- /dev/null +++ b/scripts/gfdl_om4/pipeline/dims.py @@ -0,0 +1,3 @@ +"""Dimension names shared across pipelines, kept free of heavy imports.""" + +TIME_DIM = "time" diff --git a/scripts/gfdl_om4/pipeline/om4/check_member_wetmask_equivalence.py b/scripts/gfdl_om4/pipeline/om4/check_member_wetmask_equivalence.py new file mode 100644 index 0000000000..3155131bb3 --- /dev/null +++ b/scripts/gfdl_om4/pipeline/om4/check_member_wetmask_equivalence.py @@ -0,0 +1,72 @@ +"""Check that every ensemble member derives the same ocean wetmask. + +Each member's wetmask comes from its own snapshot store's reference variable +(the NaN pattern of e.g. ``thetao`` at the first timestep). Members of one +ensemble run on the same grid and bathymetry, and downstream training and +analysis assume their output stores share one mask — so a difference is a +stop-and-report finding about the sources, not something to conform around. +""" + +import argparse +import logging + +from .config import load_config +from .run import LEVEL_DIM, load_wetmask + +logger = logging.getLogger(__name__) + + +def check_member_wetmask_equivalence(config_path: str, members: list[str]) -> None: + reference_member = members[0] + reference = None + for member in members: + config = load_config(config_path, member) + logger.info("loading wetmask for %s from %s", member, config.wetmask.store) + mask = load_wetmask(config) + if reference is None: + reference = mask + continue + if reference.sizes != mask.sizes: + raise AssertionError( + f"wetmask shape of {member} ({dict(mask.sizes)}) differs from " + f"{reference_member} ({dict(reference.sizes)})" + ) + differing = int((reference.values != mask.values).sum()) + if differing: + raise AssertionError( + f"wetmask of {member} differs from {reference_member} at " + f"{differing} cells; the sources do not share a mask — stop " + "and report, do not conform around this" + ) + assert reference is not None + logger.info( + "wetmasks identical across %d members at all %d levels " + "(%d ocean cells at the surface)", + len(members), + reference.sizes[LEVEL_DIM], + int(reference.isel({LEVEL_DIM: 0}).sum()), + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Assert every ensemble member derives the same wetmask." + ) + parser.add_argument( + "--config", required=True, help="Path to the pipeline YAML config" + ) + parser.add_argument( + "--members", nargs="+", required=True, help="Ensemble-member names to compare" + ) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + # apache_beam may have already configured the root logger on import, + # making basicConfig a no-op; raise the level explicitly. + logging.getLogger().setLevel(logging.INFO) + if len(args.members) < 2: + parser.error("at least two members are needed to compare wetmasks") + check_member_wetmask_equivalence(args.config, args.members) + + +if __name__ == "__main__": + main() diff --git a/scripts/gfdl_om4/pipeline/om4/check_output.py b/scripts/gfdl_om4/pipeline/om4/check_output.py index 7d335df55f..ce0a719344 100644 --- a/scripts/gfdl_om4/pipeline/om4/check_output.py +++ b/scripts/gfdl_om4/pipeline/om4/check_output.py @@ -21,10 +21,11 @@ def check_output( config_path: str, + member: str | None = None, output_path: str | None = None, expected_timesteps: int | None = None, ) -> None: - config = load_config(config_path) + config = load_config(config_path, member) path = output_path if output_path is not None else config.output.path stream_datasets = { stream.name: open_stream(stream, config) for stream in config.streams @@ -63,6 +64,9 @@ def main() -> None: parser.add_argument( "--config", required=True, help="Path to the pipeline YAML config" ) + parser.add_argument( + "--member", help="Ensemble-member name substituted into the config's URLs" + ) parser.add_argument( "--output-path", help="Store to check (defaults to the config's output path)", @@ -77,7 +81,7 @@ def main() -> None: # apache_beam may have already configured the root logger on import, # making basicConfig a no-op; raise the level explicitly. logging.getLogger().setLevel(logging.INFO) - check_output(args.config, args.output_path, args.expected_timesteps) + check_output(args.config, args.member, args.output_path, args.expected_timesteps) if __name__ == "__main__": diff --git a/scripts/gfdl_om4/pipeline/om4/config.py b/scripts/gfdl_om4/pipeline/om4/config.py index aeae802e6a..d1f75efe57 100644 --- a/scripts/gfdl_om4/pipeline/om4/config.py +++ b/scripts/gfdl_om4/pipeline/om4/config.py @@ -2,10 +2,13 @@ Each YAML file describes one output store: the source zarr stores and the streams of variables to read from them, the transforms to apply (vector -rotation, level splitting), the target grid and weight artifact, and the -output layout. Transforms are code; configs only select and parameterize -them, so a new simulation or a new output store is a config change, not a -code change. +rotation, level splitting, time subsampling or coarsening), the target grid +and weight artifact, and the output layout. An ensemble's members differ +only in a name inside URLs, so a config may carry the ``{member}`` +placeholder there and describe one store per member (see +pipeline/config_io.py). Transforms are code; configs only select and +parameterize them, so a new simulation or a new output store is a config +change, not a code change. """ import dataclasses @@ -49,6 +52,16 @@ class StreamConfig: (e.g. 20 subsamples 6-hourly data to 5-daily block ends). The cross-stream time-alignment assertion then guarantees the subsample lands exactly on the shared time coordinate. + time_block_mean_factor: if set, average every N consecutive timesteps, + in N-step blocks aligned to the source store's raw time grid like + ``time_subsample_stride`` (e.g. 5 coarsens daily means to 5-day + means), labeling each block mean with its block's final + timestamp. NaN anywhere in a block propagates to the mean, so a + time-varying footprint still fails the wetmask assertion. N must + evenly divide the source store's time axis (a trailing partial + block raises rather than being dropped). Mutually exclusive with + ``time_subsample_stride`` and with the pipeline's + ``shift_timestamps_to_avg_interval_midpoint``. full_cell_variables: variables additionally regridded with full-cell semantics — NaN filled with 0 over the whole grid (land included) and conservatively regridded without ocean-fraction @@ -74,6 +87,7 @@ class StreamConfig: renaming: dict[str, str] = dataclasses.field(default_factory=dict) dim_renaming: dict[str, str] = dataclasses.field(default_factory=dict) time_subsample_stride: int | None = None + time_block_mean_factor: int | None = None full_cell_variables: list[str] = dataclasses.field(default_factory=list) postprocess: list[str | PostprocessConfig] = dataclasses.field(default_factory=list) face_mask_url: str | None = None @@ -93,6 +107,17 @@ def __post_init__(self): f"time_subsample_stride must be >= 1; got " f"{self.time_subsample_stride}" ) + if self.time_block_mean_factor is not None: + if self.time_block_mean_factor < 1: + raise ValueError( + f"time_block_mean_factor must be >= 1; got " + f"{self.time_block_mean_factor}" + ) + if self.time_subsample_stride is not None: + raise ValueError( + f"stream {self.name!r} sets both time_subsample_stride and " + "time_block_mean_factor; choose one" + ) for name in self.full_cell_variables: if name not in self.variables: raise ValueError( @@ -160,7 +185,8 @@ class PipelineConfig: shift_timestamps_to_avg_interval_midpoint: if True, shift time labels backwards by half the timestep, the convention of legacy mean-state stores. Off for snapshot stores, whose raw timestamps - are kept verbatim. + are kept verbatim. Not allowed when any stream sets + ``time_block_mean_factor``. start_time: optional inclusive time-range start (e.g. "0151-01-06"), applied to all streams. end_time: optional inclusive time-range end. @@ -181,11 +207,25 @@ def __post_init__(self): names = [stream.name for stream in self.streams] if len(set(names)) != len(names): raise ValueError(f"stream names must be unique; got {names}") + if self.shift_timestamps_to_avg_interval_midpoint: + block_mean_streams = [ + stream.name + for stream in self.streams + if stream.time_block_mean_factor is not None + ] + if block_mean_streams: + raise ValueError( + "shift_timestamps_to_avg_interval_midpoint cannot be " + "combined with time_block_mean_factor; set on streams " + f"{block_mean_streams}" + ) -def load_config(path: str) -> PipelineConfig: +def load_config(path: str, member: str | None = None) -> PipelineConfig: + """Load a config, substituting ``member`` for the ``{member}`` placeholder + (see pipeline.config_io.load_yaml for the substitution rules).""" return dacite.from_dict( data_class=PipelineConfig, - data=load_yaml(path), + data=load_yaml(path, member), config=dacite.Config(strict=True), ) diff --git a/scripts/gfdl_om4/pipeline/om4/face_masks.py b/scripts/gfdl_om4/pipeline/om4/face_masks.py index a66dbbb07c..e1d5f189b1 100644 --- a/scripts/gfdl_om4/pipeline/om4/face_masks.py +++ b/scripts/gfdl_om4/pipeline/om4/face_masks.py @@ -184,6 +184,7 @@ def generate_face_masks( overwrite: bool = False, start_time: str | None = None, end_time: str | None = None, + member: str | None = None, ) -> None: """Scan a stream's C-grid velocity pair over the config's time window (optionally overridden by ``start_time``/``end_time``), compute the @@ -226,7 +227,7 @@ def generate_face_masks( "--overwrite if you really mean to replace it." ) - config = load_config(config_path) + config = load_config(config_path, member) if start_time is not None: config.start_time = start_time if end_time is not None: @@ -368,6 +369,9 @@ def main() -> None: parser.add_argument( "--config", required=True, help="Path to the pipeline YAML config" ) + parser.add_argument( + "--member", help="Ensemble-member name substituted into the config's URLs" + ) parser.add_argument( "--stream", required=True, @@ -410,6 +414,7 @@ def main() -> None: args.overwrite, args.start_time, args.end_time, + args.member, ) diff --git a/scripts/gfdl_om4/pipeline/om4/run.py b/scripts/gfdl_om4/pipeline/om4/run.py index 779d75bfce..13c76ec2b8 100644 --- a/scripts/gfdl_om4/pipeline/om4/run.py +++ b/scripts/gfdl_om4/pipeline/om4/run.py @@ -63,6 +63,7 @@ from ..postprocess import ( DERIVATION_ATTR, ChunkContext, + append_derivation, assert_postprocess_inputs, provenance_attrs, ) @@ -79,6 +80,12 @@ ) from .config import PipelineConfig, StreamConfig, load_config from .face_masks import count_masked_candidate_faces, get_face_masks +from .time_coarsening import ( + coarsen, + select_whole_blocks, + stream_output_time, + time_block_mean_factor, +) logger = logging.getLogger(__name__) @@ -108,7 +115,13 @@ def land_nan_exempt_names(level_count: int) -> list[str]: def open_stream(stream: StreamConfig, config: PipelineConfig) -> xr.Dataset: - """Open a stream's variables lazily, time-subset, with fail-fast checks.""" + """Open a stream's variables lazily, time-subset, with fail-fast checks. + + A coarsened stream is returned at its source time resolution, selected + in whole averaging blocks; the block means are taken per chunk in the beam + pipeline (see :func:`coarsen`) and its output time coordinate is + :func:`stream_output_time`. + """ ds = xr.open_zarr( make_zarr_store(stream.store), chunks=None, decode_timedelta=False ) @@ -138,7 +151,17 @@ def open_stream(stream: StreamConfig, config: PipelineConfig) -> xr.Dataset: # cross-stream time-alignment assertion enforces the coincidence. stride = stream.time_subsample_stride ds = ds.isel({TIME_DIM: slice(stride - 1, None, stride)}) - if config.start_time is not None or config.end_time is not None: + if stream.time_block_mean_factor is not None: + # Blocks are anchored to the raw time grid for the same reason as the + # subsample, and the time range selects whole blocks by their labels. + ds = select_whole_blocks( + ds, + stream.time_block_mean_factor, + config.start_time, + config.end_time, + f"stream {stream.name!r}", + ) + elif config.start_time is not None or config.end_time is not None: ds = ds.sel({TIME_DIM: slice(config.start_time, config.end_time)}) if ds.sizes[TIME_DIM] == 0: raise AssertionError( @@ -180,12 +203,12 @@ def load_wetmask(config: PipelineConfig) -> xr.DataArray: return wetmask -def _assert_time_alignment(datasets: dict[str, xr.Dataset]) -> xr.DataArray: - """Assert all streams share an identical time coordinate; return it.""" - names = list(datasets) - reference = datasets[names[0]][TIME_DIM] +def _assert_time_alignment(times: dict[str, xr.DataArray]) -> xr.DataArray: + """Assert all streams share an identical output time coordinate; return it.""" + names = list(times) + reference = times[names[0]] for name in names[1:]: - other = datasets[name][TIME_DIM] + other = times[name] if ( reference.sizes[TIME_DIM] != other.sizes[TIME_DIM] or not (reference.values == other.values).all() @@ -325,8 +348,8 @@ def _rotate_pairs( "valid zeros by the source) treated as invalid; " + rotation_note ) ds = ds.drop_vars([u_name, v_name]) - ds[u_name] = u.assign_attrs(u_attrs, **{DERIVATION_ATTR: rotation_note}) - ds[v_name] = v.assign_attrs(v_attrs, **{DERIVATION_ATTR: rotation_note}) + ds[u_name] = append_derivation(u.assign_attrs(u_attrs), rotation_note) + ds[v_name] = append_derivation(v.assign_attrs(v_attrs), rotation_note) return ds @@ -419,11 +442,18 @@ def process_chunk( weights_url: str | None = None, target_grid_name: str | None = None, ) -> tuple[xbeam.Key, xr.Dataset]: - """Beam entry point: process one (time, level) chunk of a stream.""" + """Beam entry point: process one (time, level) chunk of a stream. + + A chunk of a coarsened stream arrives at source time resolution and is + first averaged to output resolution (see :func:`coarsen`). + """ assert stream is not None assert wetmask is not None assert weights_url is not None assert target_grid_name is not None + time_offset = key.offsets[TIME_DIM] + if stream.time_block_mean_factor is not None: + time_offset, ds = coarsen(ds, time_offset, stream.time_block_mean_factor) level_index = key.offsets.get(LEVEL_DIM) if level_index is not None: wetmask = wetmask.isel({LEVEL_DIM: level_index}, drop=True) @@ -440,7 +470,7 @@ def process_chunk( level_index, ) new_key = xbeam.Key( - {TIME_DIM: key.offsets[TIME_DIM], "lat": 0, "lon": 0}, + {TIME_DIM: time_offset, "lat": 0, "lon": 0}, vars=frozenset(output.data_vars), ) return new_key, output @@ -595,23 +625,22 @@ def build_template( for name, ds in stream_datasets.items(): stream = streams_by_name[name] ds_3d, ds_2d = _split_stream(ds) - first = {TIME_DIM: slice(0, 1)} + # One output timestep: one averaging block of a coarsened stream. + first = {TIME_DIM: slice(0, time_block_mean_factor(stream))} + chunks = [] if ds_2d.data_vars: - key = xbeam.Key({TIME_DIM: 0}) - _, out = process_chunk( - key, - ds_2d.isel(first).load(), - stream, - wetmask, - config.weights_url, - config.target_grid, - ) - pieces.append(out) + chunks.append((xbeam.Key({TIME_DIM: 0}), ds_2d.isel(first))) for k in range(ds_3d.sizes.get(LEVEL_DIM, 0)): - key = xbeam.Key({TIME_DIM: 0, LEVEL_DIM: k}) + chunks.append( + ( + xbeam.Key({TIME_DIM: 0, LEVEL_DIM: k}), + ds_3d.isel({**first, LEVEL_DIM: slice(k, k + 1)}), + ) + ) + for key, chunk in chunks: _, out = process_chunk( key, - ds_3d.isel({**first, LEVEL_DIM: slice(k, k + 1)}).load(), + chunk.load(), stream, wetmask, config.weights_url, @@ -675,6 +704,10 @@ def _get_parser() -> argparse.ArgumentParser: "Unrecognized arguments are passed to beam as pipeline options." ) parser.add_argument("--config", required=True, help="Path to the YAML config") + parser.add_argument( + "--member", + help="Ensemble-member name substituted into the config's URLs", + ) parser.add_argument( "--start-time", help="Override the config's inclusive time-range start (e.g. 0151-01-06)", @@ -703,7 +736,7 @@ def main(): parser = _get_parser() args, pipeline_args = parser.parse_known_args() - config = load_config(args.config) + config = load_config(args.config, args.member) if args.start_time is not None: config.start_time = args.start_time if args.end_time is not None: @@ -736,17 +769,25 @@ def main(): int(wetmask.isel({LEVEL_DIM: -1}).sum()), ) + streams_by_name = {stream.name: stream for stream in config.streams} stream_datasets = { stream.name: open_stream(stream, config) for stream in config.streams } - output_time = _assert_time_alignment(stream_datasets) - if args.num_timesteps is not None: - output_time = output_time.isel({TIME_DIM: slice(0, args.num_timesteps)}) - stream_datasets = { - name: ds.isel({TIME_DIM: slice(0, args.num_timesteps)}) + output_time = _assert_time_alignment( + { + name: stream_output_time(streams_by_name[name], ds) for name, ds in stream_datasets.items() } + ) + if args.num_timesteps is not None: + output_time = output_time.isel({TIME_DIM: slice(0, args.num_timesteps)}) + for name, ds in stream_datasets.items(): + factor = time_block_mean_factor(streams_by_name[name]) + n_source = args.num_timesteps * factor + stream_datasets[name] = ds.isel({TIME_DIM: slice(0, n_source)}) if config.shift_timestamps_to_avg_interval_midpoint: + # Valid only for streams with one source step per output step, hence + # the config's exclusion of time_block_mean_factor. output_time = _shift_timestamps_to_midpoint(output_time) stream_datasets = { name: ds.assign_coords({TIME_DIM: output_time.values}) @@ -789,7 +830,6 @@ def main(): output_store = make_zarr_store(config.output.path, read_only=False) - streams_by_name = {stream.name: stream for stream in config.streams} logger.info("[pipeline] starting; writing to %s", config.output.path) with beam.Pipeline(options=PipelineOptions(pipeline_args)) as p: for name, ds in stream_datasets.items(): @@ -797,12 +837,15 @@ def main(): branches = [] ds_3d, ds_2d = _split_stream(ds) # Read at the source's own time-chunk width so each source chunk - # is fetched once, not once per timestep it holds. + # is fetched once, not once per timestep it holds; a coarsened + # stream widens that to a multiple of its factor so every read + # chunk holds whole averaging blocks. + factor = time_block_mean_factor(stream) if ds_3d.data_vars: - width = source_time_chunk_size(ds_3d) + width = math.lcm(source_time_chunk_size(ds_3d), factor) branches.append((f"{name}_3d", ds_3d, {TIME_DIM: width, LEVEL_DIM: 1})) if ds_2d.data_vars: - width = source_time_chunk_size(ds_2d) + width = math.lcm(source_time_chunk_size(ds_2d), factor) branches.append((f"{name}_2d", ds_2d, {TIME_DIM: width})) for label, branch_ds, chunks in branches: n_chunks = int( @@ -820,12 +863,14 @@ def main(): n_chunks, chunks[TIME_DIM], ) + # Output timesteps per chunk once any coarsening has run. + output_width = chunks[TIME_DIM] // factor split_chunks = { TIME_DIM: shard_aligned_chunk_size( - chunks[TIME_DIM], config.output.time_shard_size + output_width, config.output.time_shard_size ) } - if split_chunks[TIME_DIM] != chunks[TIME_DIM]: + if split_chunks[TIME_DIM] != output_width: logger.info( "[stream:%s] splitting to %d before consolidating into " "%d-wide shards; the %d-wide read does not align with " @@ -833,7 +878,7 @@ def main(): label, split_chunks[TIME_DIM], config.output.time_shard_size, - chunks[TIME_DIM], + output_width, ) ( p diff --git a/scripts/gfdl_om4/pipeline/om4/test_time_coarsening.py b/scripts/gfdl_om4/pipeline/om4/test_time_coarsening.py new file mode 100644 index 0000000000..226eb49537 --- /dev/null +++ b/scripts/gfdl_om4/pipeline/om4/test_time_coarsening.py @@ -0,0 +1,135 @@ +"""Tests of the time coarsening (``time_block_mean_factor``) of a stream.""" + +import numpy as np +import pytest +import xarray as xr +from packaging.version import Version + +from ..dims import TIME_DIM +from ..postprocess import DERIVATION_ATTR +from .config import StreamConfig +from .time_coarsening import coarsen, select_whole_blocks, stream_output_time + +VARIABLE = "tos" + + +# xarray 2026.9.0's Dataset coarsen mean rejects the skipna argument that +# coarsen passes. This is fixed upstream by +# https://github.com/pydata/xarray/pull/11686, which is not yet released. +xfail_with_xarray_2026_9_0 = pytest.mark.xfail( + Version(xr.__version__) == Version("2026.9.0"), + reason=( + "xarray 2026.9.0 only: Dataset coarsen mean rejects skipna " + "(pydata/xarray#11686)" + ), + raises=TypeError, + strict=True, +) + + +def _daily_source(n_days: int) -> xr.Dataset: + """A daily 2D source dataset on a noleap calendar starting 0151-01-01, + whose values are distinct at every point and timestep.""" + times = xr.date_range( + "0151-01-01", periods=n_days, calendar="noleap", use_cftime=True + ) + data = np.arange(n_days * 2 * 3, dtype=np.float32).reshape(n_days, 2, 3) + return xr.Dataset( + {VARIABLE: ((TIME_DIM, "yh", "xh"), data)}, + coords={TIME_DIM: times}, + ) + + +def _coarsened_stream(factor: int) -> StreamConfig: + return StreamConfig( + name="surface", + store="unused", + variables=[VARIABLE], + time_block_mean_factor=factor, + ) + + +def _expected_block_means(source: xr.Dataset, factor: int) -> xr.DataArray: + """The source variable's ``factor``-step block means, computed directly + with numpy (so NaN propagates), labeled with each block's final + timestamp.""" + da = source[VARIABLE] + data = da.values.reshape(-1, factor, *da.shape[1:]).mean(axis=1) + return da.isel({TIME_DIM: slice(factor - 1, None, factor)}).copy(data=data) + + +@xfail_with_xarray_2026_9_0 +def test_chunkwise_coarsening_matches_block_means(): + factor, read_width = 5, 10 + source = _daily_source(20) + + offsets, chunks = [], [] + for time_offset in range(0, source.sizes[TIME_DIM], read_width): + chunk = source.isel({TIME_DIM: slice(time_offset, time_offset + read_width)}) + output_offset, coarsened = coarsen(chunk, time_offset, factor) + offsets.append(output_offset) + chunks.append(coarsened) + result = xr.concat(chunks, dim=TIME_DIM) + + assert offsets == [0, 2] + expected = _expected_block_means(source, factor) + xr.testing.assert_allclose(result[VARIABLE], expected) + # Each block mean is labeled with its block's final timestamp, which is + # the stream's output time coordinate used for cross-stream alignment. + stream = _coarsened_stream(factor) + xr.testing.assert_equal(stream_output_time(stream, source), expected[TIME_DIM]) + assert "mean over 5-step blocks" in result[VARIABLE].attrs[DERIVATION_ATTR] + + +def test_time_range_selects_whole_blocks_by_their_labels(): + factor = 5 + source = _daily_source(20) + # Block labels are 01-05, 01-10, 01-15 and 01-20. The range selects the + # blocks whose labels fall within it, including the source days of the + # first block that precede the start time. + selected = select_whole_blocks(source, factor, "0151-01-03", "0151-01-17", "test") + + xr.testing.assert_identical(selected, source.isel({TIME_DIM: slice(0, 15)})) + labels = xr.date_range( + "0151-01-05", periods=3, freq="5D", calendar="noleap", use_cftime=True + ) + expected_time = xr.DataArray(labels, dims=TIME_DIM, coords={TIME_DIM: labels}) + stream = _coarsened_stream(factor) + xr.testing.assert_equal(stream_output_time(stream, selected), expected_time) + + +def test_time_range_without_whole_block_selects_nothing(): + selected = select_whole_blocks( + _daily_source(20), 5, "0151-01-06", "0151-01-09", "test" + ) + assert selected.sizes[TIME_DIM] == 0 + + +def test_factor_not_dividing_source_time_axis_raises(): + with pytest.raises(AssertionError, match="does not evenly divide"): + select_whole_blocks(_daily_source(21), 5, None, None, "test") + + +@xfail_with_xarray_2026_9_0 +def test_nan_in_block_propagates_to_its_mean_only(): + factor = 5 + source = _daily_source(10) + source[VARIABLE][6, 1, 2] = np.nan + + _, coarsened = coarsen(source, 0, factor) + + # The expected means are NaN only at that point of the second block, and + # xr.testing.assert_allclose requires NaN in exactly the same places. + xr.testing.assert_allclose( + coarsened[VARIABLE], _expected_block_means(source, factor) + ) + + +@pytest.mark.parametrize( + "time_offset, size", + [pytest.param(3, 10, id="offset"), pytest.param(0, 7, id="size")], +) +def test_chunk_without_whole_blocks_raises(time_offset, size): + chunk = _daily_source(size) + with pytest.raises(AssertionError, match="whole 5-step blocks"): + coarsen(chunk, time_offset, 5) diff --git a/scripts/gfdl_om4/pipeline/om4/time_coarsening.py b/scripts/gfdl_om4/pipeline/om4/time_coarsening.py new file mode 100644 index 0000000000..0023d79bb6 --- /dev/null +++ b/scripts/gfdl_om4/pipeline/om4/time_coarsening.py @@ -0,0 +1,91 @@ +"""Time coarsening of a stream by block means (``time_block_mean_factor``). + +A coarsened stream is selected at source time resolution in whole averaging +blocks (:func:`select_whole_blocks`), and each beam chunk is averaged to +output resolution (:func:`coarsen`). Kept free of beam so it can be tested +without the pipeline's runtime dependencies. +""" + +import xarray as xr + +from ..dims import TIME_DIM +from ..postprocess import append_derivation +from .config import StreamConfig + + +def time_block_mean_factor(stream: StreamConfig) -> int: + """Source timesteps per output timestep of a stream (1 unless coarsened).""" + return stream.time_block_mean_factor or 1 + + +def select_whole_blocks( + ds: xr.Dataset, + factor: int, + start_time: str | None, + end_time: str | None, + context: str, +) -> xr.Dataset: + """Select the ``factor``-step blocks of ``ds``, anchored to its first + timestep, whose final timestamps (their output labels) fall within the + inclusive time range. + + Raises if ``factor`` does not evenly divide the full time axis, rather + than silently dropping a trailing partial block. + """ + size = ds.sizes[TIME_DIM] + if size % factor: + raise AssertionError( + f"{context}: time_block_mean_factor {factor} does not evenly divide " + f"the {size} timesteps of the source store" + ) + n_blocks = size // factor + if start_time is None and end_time is None: + return ds + labels = ds.indexes[TIME_DIM][factor - 1 :: factor] + selected = labels.slice_indexer(start_time, end_time) + first, stop = selected.start or 0, selected.stop + stop = n_blocks if stop is None else stop + return ds.isel({TIME_DIM: slice(first * factor, max(stop, first) * factor)}) + + +def stream_output_time(stream: StreamConfig, ds: xr.Dataset) -> xr.DataArray: + """The output time coordinate of a stream opened by + :func:`pipeline.om4.run.open_stream`: its own time coordinate, or for a + coarsened stream the final timestamp of each block.""" + factor = time_block_mean_factor(stream) + return ds[TIME_DIM].isel({TIME_DIM: slice(factor - 1, None, factor)}) + + +def coarsen(ds: xr.Dataset, time_offset: int, factor: int) -> tuple[int, xr.Dataset]: + """Average a source-resolution chunk of a coarsened stream over its + ``factor``-step blocks. + + The chunk must hold whole blocks (the read width is a multiple of + ``factor``, and :func:`pipeline.om4.run.open_stream` selects the stream + in whole blocks). + Returns the chunk's time offset re-expressed in output timesteps, and the + block means labeled with each block's final timestamp. NaN anywhere in a + block propagates to its mean, so a time-varying footprint still fails the + wetmask assertion downstream. + """ + size = ds.sizes[TIME_DIM] + if time_offset % factor or size % factor: + raise AssertionError( + f"chunk at time offset {time_offset} with {size} timesteps does not " + f"hold whole {factor}-step blocks" + ) + start = time_offset // factor + labels = ds[TIME_DIM].values[factor - 1 :: factor] + coarsened = ( + ds.drop_vars(TIME_DIM) + .coarsen({TIME_DIM: factor}) + .mean(skipna=False, keep_attrs=True) + .assign_coords({TIME_DIM: labels}) + ) + note = ( + f"mean over {factor}-step blocks of the source time axis, labeled " + "with each block's final time" + ) + for name in coarsened.data_vars: + append_derivation(coarsened[name], note) + return start, coarsened diff --git a/scripts/gfdl_om4/pipeline/zarr_io.py b/scripts/gfdl_om4/pipeline/zarr_io.py index b563e2c022..9d0376af60 100644 --- a/scripts/gfdl_om4/pipeline/zarr_io.py +++ b/scripts/gfdl_om4/pipeline/zarr_io.py @@ -14,7 +14,8 @@ from obstore.store import from_url from zarr.storage import ObjectStore -TIME_DIM = "time" +from .dims import TIME_DIM + OUTPUT_DTYPE = np.float32