Open dcherian opened 1 year ago
Rewritten to use the weatherbench2 dataset
import xarray as xr
ds = xr.open_zarr(
"gs://weatherbench2/datasets/era5/1959-2023_01_10-full_37-1h-0p25deg-chunk-1.zarr",
chunks={},
)
ds = ds[
["u_component_of_wind", "v_component_of_wind", "temperature", "vertical_velocity"]
].rename(
{
"u_component_of_wind": "U",
"v_component_of_wind": "V",
"temperature": "T",
"vertical_velocity": "W",
}
)
zonal_means = ds.mean("longitude")
anomaly = ds - zonal_means
anomaly["uv"] = anomaly.U * anomaly.V
anomaly["vt"] = anomaly.V * anomaly.T
anomaly["uw"] = anomaly.U * anomaly.W
temdiags = zonal_means.merge(anomaly[["uv", "vt", "uw"]].mean("longitude"))
# This is incredibly slow, takes a while for flox to construct the graph
# daily = temdiags.resample(time="D").mean()
# Option 2: rechunk to make it a blockwise problem
# we should do this automatically
from xarray.groupers import TimeResampler
daily = temdiags.chunk(time=TimeResampler("D")).resample(time="D").mean()
daily.to_zarr(SOMEWHERE)
This used to be such a pain, people were looping over a years worth of files for 60 years of data.
This is ERA5