convert ROMS output to zarr

roms_to_zarr.py

#!/usr/bin/env python

# useful ref:
# http://xarray.pydata.org/en/stable/examples/ROMS_ocean_model.html

import xarray as xr
import glob


# list of files to convert
filelist = glob.glob('../monthly/CCS2-RD.HCcobalt40R_avg_*.nc')
# path to grid file
gridfile = '/Volumes/P1/ROMS-Inputs/CCS2/grid/CCS2_smoother_0-360.nc'
# output store name
store_name = 'CCS2-RD.HCcobalt40R_avg_monthly'

# ------------------------------------------------------------------------------------------------
# functions

def remove_all_junk_var(ds):
    for v in ds.variables:
        if v not in ['Cs_r', 'Cs_w', 'hc', 'Vtransform', 'ocean_time']:
            ds = ds.drop_vars(v) if len(ds[v].dims) < 3 else ds
    return ds


def select_interior(ds):
    ds = ds.isel(xi_rho=slice(1,-1), eta_rho=slice(1,-1))
    if 'xi_v' in ds.dims:
        ds = ds.isel(xi_v=slice(1,-1))
    if 'eta_u' in ds.dims:
        ds = ds.isel(eta_u=slice(1,-1))
    return ds


def rename_dims(ds):
    ds = ds.rename({'xi_rho': 'xh', 'xi_v': 'xh', 'xi_u': 'xq', 'xi_psi': 'xq',
                    'eta_rho': 'yh', 'eta_v': 'yq', 'eta_u': 'yh', 'eta_psi': 'yq',
                    'ocean_time': 'time'
                    })
    return ds


def add_coords(ds):
    ds = ds.assign_coords({'lon_rho': ds['lon_rho'],
                           'lon_v': ds['lon_v'],
                           'lon_u': ds['lon_u'],
                           'lon_psi': ds['lon_psi'],
                           'lat_rho': ds['lat_rho'],
                           'lat_v': ds['lat_v'],
                           'lat_u': ds['lat_u'],
                           'lat_psi': ds['lat_psi'],
                           'time': ds['time'],
                           })
    ds = ds.set_coords(['Cs_r', 'Cs_w', 'hc', 'h', 'Vtransform'])
    return ds


# ------------------------------------------------------------------------------------------------
# main

# load the grid
ds_grid = xr.open_dataset(gridfile)
ds_grid = ds_grid.drop_vars(['hraw', 'spherical', 'lon_vert', 'lat_vert', 'x_vert', 'y_vert',
                             'Cs_r', 'Cs_w', 'hc'])
ds_grid = select_interior(ds_grid)

first=True
# iterate on files
for ncfile in filelist:
    ds = xr.open_dataset(ncfile)
    ds = remove_all_junk_var(ds)
    ds = select_interior(ds)
    ds = xr.merge([ds, ds_grid])
    ds = rename_dims(ds)
    ds = add_coords(ds)
    ds = ds.chunk({'time':1})
    print(f"working on {ncfile}")
    if first:
        ds.to_zarr(store_name, consolidated=True, mode='w')
    else:
        ds.to_zarr(store_name, consolidated=True, mode='a', append_dim='time')
    first=False