Transformed Eulerian Mean #3

Activity

1. dcherian commented on Sep 17, 2024

   

   dcherian

   on Sep 17, 2024

   Author

   More actions

   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)