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