def add_MITgcm_missing_metrics(dset, periodic=None):
"""
Infer missing metrics from MITgcm output files.
Parameters
----------
dset : xarray.Dataset
A dataset open from a file
Return
-------
dset : xarray.Dataset
Input dataset with appropriated metrics added
grid : xgcm.Grid
The grid with appropriated metrics
"""
coords = dset.coords
grid = Grid(dset, periodic=periodic)
BCx = 'periodic' if grid.axes['X']._periodic else 'fill'
BCy = 'periodic' if grid.axes['Y']._periodic else 'fill'
if 'drW' not in coords: # vertical cell size at u point
coords['drW'] = dset.hFacW * dset.drF
if 'drS' not in coords: # vertical cell size at v point
coords['drS'] = dset.hFacS * dset.drF
if 'drC' not in coords: # vertical cell size at tracer point
coords['drC'] = dset.hFacC * dset.drF
if 'dxF' not in coords:
coords['dxF'] = grid.interp(dset.dxC, 'X', boundary=BCx)
if 'dyF' not in coords:
coords['dyF'] = grid.interp(dset.dyC, 'Y', boundary=BCy)
if 'dxV' not in coords:
coords['dxV'] = grid.interp(dset.dxG, 'X', boundary=BCx)
if 'dyU' not in coords:
coords['dyU'] = grid.interp(dset.dyG, 'Y', boundary=BCy)
if 'hFacZ' not in coords:
coords['hFacZ'] = grid.interp(dset.hFacS, 'X', boundary=BCx)
if 'maskZ' not in coords:
coords['maskZ'] = coords['hFacZ']
# Calculate vertical distances located on the cellboundary
# ds.coords['dzC'] = grid.diff(ds.depth, 'Z', boundary='extrapolate')
# Calculate vertical distances located on the cellcenter
# ds.coords['dzT'] = grid.diff(ds.depth_left, 'Z', boundary='extrapolate')
metrics = {
('X', ): ['dxG', 'dxF', 'dxC', 'dxV'], # X distances
('Y', ): ['dyG', 'dyF', 'dyC', 'dyU'], # Y distances
('Z', ): ['drW', 'drS', 'drC', 'drF'], # Z distances
('X', 'Y'): ['rAw', 'rAs', 'rA', 'rAz']
} # Areas
grid._assign_metrics(metrics)
return dset, grid
def add_latlon_metrics(dset, dims=None):
"""
Infer 2D metrics (latitude/longitude) from gridded data file.
Parameters
----------
dset : xarray.Dataset
A dataset open from a file
dims : dict
Dimension pair in a dict, e.g., {'lat':'latitude', 'lon':'longitude'}
Return
-------
dset : xarray.Dataset
Input dataset with appropriated metrics added
grid : xgcm.Grid
The grid with appropriated metrics
"""
lon, lat = None, None
if dims is None:
for dim in dimXList:
if dim in dset or dim in dset.coords:
lon = dim
break
for dim in dimYList:
if dim in dset or dim in dset.coords:
lat = dim
break
if lon is None or lat is None:
raise Exception('unknown dimension names in dset, should be in ' +
str(dimXList + dimYList))
else:
lon, lat = dims['lon'], dims['lat']
ds = generate_grid_ds(dset, {'X': lon, 'Y': lat})
coords = ds.coords
if __is_periodic(coords[lon], 360.0):
periodic = 'X'
else:
periodic = []
grid = Grid(ds, periodic=periodic)
na = np.nan
if 'X' in periodic:
dlonG = grid.diff(ds[lon], 'X', boundary_discontinuity=360)
dlonC = grid.diff(ds[lon + '_left'], 'X', boundary_discontinuity=360)
else:
dlonG = grid.diff(ds[lon], 'X', boundary='fill', fill_value=na)
dlonC = grid.diff(ds[lon + '_left'],
'X',
boundary='fill',
fill_value=na)
dlatG = grid.diff(ds[lat], 'Y', boundary='fill', fill_value=na)
dlatC = grid.diff(ds[lat + '_left'], 'Y', boundary='fill', fill_value=na)
coords['dxG'], coords['dyG'] = __dll_dist(dlonG, dlatG, ds[lon], ds[lat])
coords['dxC'], coords['dyC'] = __dll_dist(dlonC, dlatC, ds[lon], ds[lat])
coords['rAc'] = ds['dyC'] * ds['dxC']
metrics = {
('X', ): ['dxG', 'dxC'], # X distances
('Y', ): ['dyG', 'dyC'], # Y distances
('X', 'Y'): ['rAc']
}
grid._assign_metrics(metrics)
return ds, grid
