generate_pop_grid.py

def preprocess(ds):
    """Drop all unnecessary variables and coordinates"""
    vars_to_drop = [vname for vname in ds.data_vars if vname not in variables]
    coord_vars = [vname for vname in ds.data_vars if 'time' not in ds[vname].dims or 'bound' in vname]
    ds_fixed = ds.set_coords(coord_vars)
    data_vars_dims = []
    for data_var in ds_fixed.data_vars:
        data_vars_dims.extend(list(ds_fixed[data_var].dims))
    coords_to_drop = [coord for coord in ds_fixed.coords if coord not in data_vars_dims]
    grid_vars = list(set(vars_to_drop + coords_to_drop) - set(['time', 'time_bound']))
    # ds_fixed = ds_fixed.drop(grid_vars)
    if 'history' in ds_fixed.attrs:
        del ds_fixed.attrs['history']
    return ds_fixed, grid_vars
  
ds = dsets['ocn_CTRL_pop.h_SST'].copy()
x, y = preprocess(ds)
d = x[y]

data_vars = list(d.data_vars)
grid = d.isel(time=0).drop('time').set_coords(data_vars).load()
grid.attrs = {}
grid.to_zarr("/glade/scratch/abanihi/lens-aws/ocean/grid.zarr", consolidated=True)