Source code for gcpy.grid

#!/usr/bin/env python3
"""
Module containing variables and functions that define and
manipulate GEOS-Chem horizontal and vertical grids
"""
from itertools import product
import xarray as xr
import numpy as np
from gcpy.util import get_shape_of_data, verify_variable_type
from gcpy.grid_stretching_transforms import scs_transform
from gcpy.constants import \
    GLOBAL_LL_EXTENT, NO_STRETCH_SG_PARAMS, R_EARTH_m
from gcpy.vgrid_defs import \
    _GCAP2_102L_AP, _GCAP2_102L_BP, _GCAP2_74L_AP, _GCAP2_74L_BP, \
    _GCAP2_40L_AP, _GCAP2_40L_BP, _GEOS_72L_AP, _GEOS_72L_BP, \
    _GEOS_47L_AP, _GEOS_47L_BP, _CAM_26L_AP, _CAM_26L_BP
from gcpy.cstools import find_index, is_cubed_sphere


[docs] def get_troposphere_mask(ds): """ Returns a mask array for picking out the tropospheric grid boxes. Parameters ---------- ds : xarray.Dataset Dataset containing certain met field variables (i.e. Met_TropLev, Met_BXHEIGHT). Returns ------- tropmask : numpy.ndarray Tropospheric mask. False denotes grid boxes that are in the troposphere and True in the stratosphere (as per Python masking logic). """ # ================================================================== # Initialization # ================================================================== # Make sure ds is an xarray Dataset object if not isinstance(ds, xr.Dataset): raise TypeError("The ds argument must be an xarray Dataset!") # Make sure certain variables are found if "Met_BXHEIGHT" not in ds.data_vars.keys(): raise ValueError("Met_BXHEIGHT could not be found!") if "Met_TropLev" not in ds.data_vars.keys(): raise ValueError("Met_TropLev could not be found!") # Mask of tropospheric grid boxes in the Ref dataset shape = get_shape_of_data(np.squeeze(ds["Met_BXHEIGHT"])) # Determine if this is GCHP data is_gchp = "nf" in ds["Met_BXHEIGHT"].dims # ================================================================== # Create the mask arrays for the troposphere # # Convert the Met_TropLev DataArray objects to numpy ndarrays of # integer. Also subtract 1 to convert from Fortran to Python # array index notation. # ================================================================== multi_time_slices = (is_gchp and len(shape) == 5) or \ (not is_gchp and len(shape) == 4) if multi_time_slices: # -------------------------------------------------------------- # GCC: There are multiple time slices # -------------------------------------------------------------- # Create the tropmask array with dims # (time, lev, nf*lat*lon) for GCHP, or # (time, lev, lat*lon ) for GCC tropmask = np.ones((shape[0], shape[1], np.prod(np.array(shape[2:]))), bool) # Loop over each time for t in range(tropmask.shape[0]): # Pick the tropopause level and make a 1-D array values = ds["Met_TropLev"].isel(time=t).values lev = np.int_(np.squeeze(values) - 1) lev_1d = lev.flatten() # Create the tropospheric mask array for x in range(tropmask.shape[2]): tropmask[t, 0: lev_1d[x], x] = False else: # -------------------------------------------------------------- # There is only one time slice # -------------------------------------------------------------- # Create the tropmask array with dims (lev, lat*lon) tropmask = np.ones((shape[0], np.prod(np.array(shape[1:]))), bool) # Pick the tropopause level and make a 1-D array values = ds["Met_TropLev"].values lev = np.int_(np.squeeze(values) - 1) lev_1d = lev.flatten() # Create the tropospheric mask array for x in range(tropmask.shape[1]): tropmask[0: lev_1d[x], x] = False # Reshape into the same shape as Met_BxHeight return tropmask.reshape(shape)
[docs] def get_input_res(data): """ Returns resolution of dataset passed to compare_single_level or compare_zonal_means. Parameters ---------- data : xarray.Dataset Input GEOS-Chem dataset. Returns ------- res : str or int Lat/lon res of the form 'latresxlonres' or cubed-sphere resolution. gridtype : str 'll' for lat/lon or 'cs' for cubed-sphere. """ vdims = data.dims if "lat" in vdims and "lon" in vdims: lat = data["lat"].values lon = data["lon"].values if lat.size / 6 == lon.size: return lon.size, "cs" # Make writable copies before sorting. # This will avoid an error in numpy 2.x lat_sort = lat.copy() lon_sort = lon.copy() lat_sort.sort() lon_sort.sort() # use increment of second and third coordinates # to avoid polar mischief lat_res = np.abs(lat_sort[2] - lat_sort[1]) lon_res = np.abs(lon_sort[2] - lon_sort[1]) return str(lat_res) + "x" + str(lon_res), "ll" #print("grid is cs: ", vdims) # GCHP data using MAPL v1.0.0+ has dims time, lev, nf, Ydim, and Xdim if isinstance(data.dims, tuple): return len(data["Xdim"].values), "cs" return data.dims["Xdim"], "cs"
[docs] def call_make_grid( res, gridtype, in_extent=None, out_extent=None, sg_params=None ): """ Create a mask with NaN values removed from an input array Parameters ---------- res : str or int Resolution of grid (format 'latxlon' or csres). gridtype : str 'll' for lat/lon or 'cs' for cubed-sphere. in_extent : list of float, optional Describes minimum and maximum latitude and longitude of input data in the format [minlon, maxlon, minlat, maxlat]. Default value: GLOBAL_LL_EXTENT, i.e. [-180, 180, -90, 90] out_extent : list of float, optional Desired minimum and maximum latitude and longitude of output grid in the format [minlon, maxlon, minlat, maxlat]. Default value: GLOBAL_LL_EXTENT, i.e. [-180, 180, -90, 90] sg_params : list, optional Desired stretched-grid parameters in the format [stretch_factor, target_longitude, target_latitude]. Will trigger stretched-grid creation if not default values. Default value: NO_STRETCH_SG_PARAMS (no stretching, [1, 170, -90]) Returns ------- grid : dict The created grid. grid_list : list of dict or None List of grids if gridtype is 'cs', otherwise None. """ verify_variable_type(res, (str, int)) verify_variable_type(gridtype, str) # Set defaults for keyword args if in_extent is None: in_extent = GLOBAL_LL_EXTENT if out_extent is None: out_extent = GLOBAL_LL_EXTENT if sg_params is None: sg_params = NO_STRETCH_SG_PARAMS # Lat-lon grid if gridtype == "ll": return [make_grid_ll(res, in_extent, out_extent), None] # Standard cubed-sphere if sg_params == NO_STRETCH_SG_PARAMS: return make_grid_cs(res) # Stretched-grid return make_grid_sg(res, *sg_params)
[docs] def get_grid_extents(data, edges=True): """ Get min and max lat and lon from an input GEOS-Chem xarray dataset or grid dict. Parameters ---------- data : xarray.Dataset or dict A GEOS-Chem dataset or a grid dict. edges : bool, optional Whether grid extents should use cell edges instead of centers. Default value: True Returns ------- minlon : float Minimum longitude of data grid. maxlon : float Maximum longitude of data grid. minlat : float Minimum latitude of data grid. maxlat : float Maximum latitude of data grid. """ if isinstance(data, dict): if "lon_b" in data and edges: return \ np.min(data["lon_b"]), np.max(data["lon_b"]), \ np.min(data["lat_b"]), np.max(data["lat_b"]) if not edges: return \ np.min(data["lon"]), np.max(data["lon"]), \ np.min(data["lat"]), np.max(data["lat"]) return \ GLOBAL_LL_EXTENT[0], GLOBAL_LL_EXTENT[1], \ GLOBAL_LL_EXTENT[2], GLOBAL_LL_EXTENT[3] if "lat" in data.dims and "lon" in data.dims: lat = data["lat"].values lon = data["lon"].values if lat.size / 6 == lon.size: # No extents for CS plots right now return -180, 180, -90, 90 lat = np.sort(lat) minlat = np.min(lat) if abs(abs(lat[1]) - abs(lat[0])) != abs(abs(lat[2]) - abs(lat[1])): #pole is cutoff minlat = minlat - 1 maxlat = np.max(lat) if abs(abs(lat[-1]) - abs(lat[-2])) != abs(abs(lat[-2]) - abs(lat[-3])): maxlat = maxlat + 1 # add longitude res to max longitude lon = np.sort(lon) minlon = np.min(lon) maxlon = np.max(lon) + abs(abs(lon[-1]) - abs(lon[-2])) return minlon, maxlon, minlat, maxlat # GCHP data using MAPL v1.0.0+ has dims time, lev, nf, Ydim, and Xdim return \ GLOBAL_LL_EXTENT[0], GLOBAL_LL_EXTENT[1], \ GLOBAL_LL_EXTENT[2], GLOBAL_LL_EXTENT[3]
[docs] def get_vert_grid( dataset, AP=None, BP=None, p_sfc=1013.25): """ Determine vertical grid of input dataset. Parameters ---------- dataset : xr.Dataset A GEOS-Chem output dataset. AP : list-like, optional Hybrid grid parameter A (hPa). BP : list-like, optional Hybrid grid parameter B (unitless). p_sfc : float, optional Surface pressure in hPa. Default value: 1013.25 Returns ------- pedge : numpy.ndarray Edge pressure values for vertical grid. p_mid : numpy.ndarray Midpoint pressure values for vertical grid. nlev : int Number of levels in vertical grid. """ # 72L GEOS grid if dataset.sizes["lev"] in (72, 73): grid = VertGrid(_GEOS_72L_AP, _GEOS_72L_BP, p_sfc) return grid.p_edge(), grid.p_mid(), 72 # 47L GEOS grid if dataset.sizes["lev"] in (47, 48): grid = VertGrid(_GEOS_47L_AP, _GEOS_47L_BP, p_sfc) return grid.p_edge(), grid.p_mid(), 47 # Grid without specified AP, BP if AP is None or BP is None: if dataset.sizes["lev"] == 1: AP = [1, 1] BP = [1] grid = VertGrid(AP, BP, p_sfc) return grid.p_edge(), grid.p_mid(), np.size(AP) raise ValueError( "Only 72/73 or 47/48 level vertical grids are automatically\n" + "determined from input dataset by get_vert_grid().\n" + "please pass grid parameters AP and BP as keyword arguments" ) # Grid with specified AP, BP grid = VertGrid(AP, BP, p_sfc) return grid.p_edge(), grid.p_mid(), np.size(AP)
[docs] def get_ilev_coord( n_lev=72, AP_edge=None, BP_edge=None, top_down=False, gchp_indices=False ): """ Returns the eta values (defined as (A/P0) + B) at vertical level edges. These are used to define the "ilev" netCDF coordinate variable. Parameters ---------- n_lev : int, optional Number of levels in the grid. Default value: 72 AP_edge : list-like, optional Hybrid grid parameter A (hPa), with values placed on level edges. If not specified, values from the _GEOS_72L_AP array in this module will be used. BP_edge : list-like, optional Hybrid grid parameter B (unitless), with values placed on level edges. If not specified, values from the _GEOS_72L_BP array in this module will be used. top_down : bool, optional Set this to True if the eta coordinate will be arranged from top-of-atm downward (True) or from the surface upward (False). Default value: False gchp_indices : bool, optional Set this to True to return an array of indices (as is used in GCHP files). Default value: False Returns ------- ilev : numpy.ndarray List of eta values at vertical grid edges. """ if n_lev is None: n_lev = 72 # Return GCHP-style indices for the level dimension if gchp_indices: return np.linspace(1, n_lev+1, n_lev+1, dtype=np.float64) # Get eta values at vertical level edges if AP_edge is None and n_lev == 72: AP_edge = _GEOS_72L_AP if AP_edge is None and n_lev == 47: AP_edge = _GEOS_47L_AP if BP_edge is None and n_lev == 72: BP_edge = _GEOS_72L_BP if BP_edge is None and n_lev == 47: BP_edge = _GEOS_47L_BP ilev = np.array((AP_edge/1000.0) + BP_edge, dtype=np.float64) if top_down: ilev = ilev[::-1] return ilev
[docs] def get_lev_coord( n_lev=72, AP_edge=None, BP_edge=None, top_down=False, gchp_indices=False ): """ Returns the eta values (defined as (A/P0) + B) at vertical level midpoints. These are used to define the "lev" netCDF coordinate variable. Parameters ---------- n_lev : int, optional Number of levels in the grid. Default value: 72 AP_edge : list-like, optional Hybrid grid parameter A (hPa), with values placed on level edges. If not specified, values from the _GEOS_72L_AP array in this module will be used. BP_edge : list-like, optional Hybrid grid parameter B (unitless), with values placed on level edges. If not specified, values from the _GEOS_72L_BP array in this module will be used. top_down : bool, optional Set this to true if the eta coordinate will be arranged from top-of-atm downward (True) or from the surface upward (False). Default value: False gchp_indices : bool, optional Set this to True to return an array of indices (as is used in GCHP files). Default value: False Returns ------- lev : numpy.ndarray List of eta values at vertical grid midpoints. """ if n_lev is None: n_lev = 72 # Return GCHP-style indices for the level dimension if gchp_indices: return np.linspace(1, n_lev, n_lev, dtype=np.float64) # Compute AP, BP at midpoints. # Convert inputs to numpy.ndarray for fast computation if AP_edge is None and n_lev == 72: AP_edge = _GEOS_72L_AP if AP_edge is None and n_lev == 47: AP_edge = _GEOS_47L_AP if BP_edge is None and n_lev == 72: BP_edge = _GEOS_72L_BP if BP_edge is None and n_lev == 47: BP_edge = _GEOS_47L_BP AP_edge = np.array(AP_edge) BP_edge = np.array(BP_edge) AP_mid = (AP_edge[0:n_lev:1] + AP_edge[1:n_lev+1:1]) * 0.5 BP_mid = (BP_edge[0:n_lev:1] + BP_edge[1:n_lev+1:1]) * 0.5 lev = np.array((AP_mid / 1000.0) + BP_mid, dtype=np.float64) if top_down: lev = lev[::-1] return lev
[docs] def get_pressure_indices(pedge, pres_range): """ Get indices where edge pressure values are within a given pressure range. Parameters ---------- pedge : numpy.ndarray Edge pressure values for the vertical grid. pres_range : list of float Contains minimum and maximum pressure. Returns ------- indices : numpy.ndarray Indices where edge pressure values are within the given pressure range. """ return np.where( (pedge <= np.max(pres_range)) & ( pedge >= np.min(pres_range)))[0]
[docs] def pad_pressure_edges(pedge_ind, max_ind, pmid_len): """ Add outer indices to edge pressure index list. Parameters ---------- pedge_ind : list List of edge pressure indices. max_ind : int Maximum index. pmid_len : int Length of pmid which should not be exceeded by indices. Returns ------- pedge_ind : list List of edge pressure indices, possibly with new minimum and maximum indices. """ if max_ind > pmid_len: # don't overstep array bounds for full array max_ind = max_ind - 1 if min(pedge_ind) != 0: pedge_ind = np.append(min(pedge_ind) - 1, pedge_ind) if max(pedge_ind) != max_ind: pedge_ind = np.append(pedge_ind, max(pedge_ind) + 1) return pedge_ind
[docs] def get_ind_of_pres(dataset, pres): """ Get index of pressure level that contains the requested pressure value. Parameters ---------- dataset : xarray.Dataset GEOS-Chem dataset. pres : int or float Desired pressure value. Returns ------- index : int Index of level in dataset that corresponds to requested pressure. """ pedge, pmid, _ = get_vert_grid(dataset) converted_dataset = convert_lev_to_pres(dataset, pmid, pedge) return np.argmin(np.abs(converted_dataset['lev'] - pres).values)
[docs] def convert_lev_to_pres(dataset, pmid, pedge, lev_type='pmid'): """ Convert lev dimension to pressure in a GEOS-Chem dataset. Parameters ---------- dataset : xarray.Dataset GEOS-Chem dataset. pmid : numpy.ndarray Midpoint pressure values. pedge : numpy.ndarray Edge pressure values. lev_type : str, optional Denote whether lev is 'pedge' or 'pmid' if grid is not 72/73 or 47/48 levels. Default value: 'pmid' Returns ------- dataset : xarray.Dataset Input dataset with "lev" dimension values replaced with pressure values. """ if dataset.sizes["lev"] in (72, 47): dataset["lev"] = pmid elif dataset.sizes["lev"] in (73, 48): dataset["lev"] = pedge elif lev_type == 'pmid': print('Warning: Assuming levels correspond with midpoint pressures') dataset["lev"] = pmid else: dataset["lev"] = pedge dataset["lev"].attrs["unit"] = "hPa" dataset["lev"].attrs["long_name"] = "level pressure" return dataset
[docs] class VertGrid: """ Class that defines a vertical grid given the Ap and Bp grid parameters and surface pressure. Parameters ---------- AP : list-like Hybrid-grid A parameter (hPa). BP : list-like Hybrid-grid B parameter (unitless). p_sfc : float, optional Surface pressure in hPa. Default value: 1013.25 """ def __init__(self, AP=None, BP=None, p_sfc=1013.25): if (len(AP) != len(BP)) or (AP is None): # Throw error? print('Inconsistent vertical grid specification') self.AP = np.array(AP) self.BP = np.array(BP) self.p_sfc = p_sfc
[docs] def p_edge(self): """ Compute pressure at grid cell edges. Returns ------- p_edge : numpy.ndarray Edge pressure values (hPa) computed as AP + BP * p_sfc. """ # Calculate pressure edges using eta coordinate return self.AP + self.BP * self.p_sfc
[docs] def p_mid(self): """ Compute pressure at grid cell midpoints. Returns ------- p_mid : numpy.ndarray Midpoint pressure values (hPa) computed as the average of adjacent edge pressure values. """ p_edge = self.p_edge() return (p_edge[1:] + p_edge[:-1]) / 2.0
# Define commonly-used vertical grids GEOS_72L_grid = VertGrid(_GEOS_72L_AP, _GEOS_72L_BP) GEOS_47L_grid = VertGrid(_GEOS_47L_AP, _GEOS_47L_BP) GCAP2_102L_grid = VertGrid(_GCAP2_102L_AP, _GCAP2_102L_BP) GCAP2_74L_grid = VertGrid(_GCAP2_74L_AP, _GCAP2_74L_BP) GCAP2_40L_grid = VertGrid(_GCAP2_40L_AP, _GCAP2_40L_BP) CAM_26L_grid = VertGrid(_CAM_26L_AP, _CAM_26L_BP)
[docs] def make_grid_ll(llres, in_extent=None, out_extent=None): """ Creates a lat/lon grid description. Parameters ---------- llres : str Lat/lon resolution in 'latxlon' format (e.g. '4x5'). in_extent : list of float, optional Describes minimum and maximum latitude and longitude of initial grid in the format [minlon, maxlon, minlat, maxlat]. Default value: [-180, 180, -90, 90] out_extent : list of float, optional Describes minimum and maximum latitude and longitude of target grid in the format [minlon, maxlon, minlat, maxlat]. Needed to trim extent of input data. Default value: [] (assumes value of in_extent) Returns ------- llgrid : dict Dict grid description of format {'lat' : lat midpoints, 'lon' : lon midpoints, 'lat_b' : lat edges, 'lon_b' : lon edges}. """ verify_variable_type(llres, str) # Default values for keyword args if in_extent is None: in_extent = GLOBAL_LL_EXTENT if out_extent is None: out_extent = in_extent # get initial bounds of grid [minlon, maxlon, minlat, maxlat] = in_extent [dlat, dlon] = list(map(float, llres.split('x'))) lon_b = np.linspace(minlon - dlon / 2, maxlon - dlon / 2, int((maxlon - minlon) / dlon) + 1) lat_b = np.linspace(minlat - dlat / 2, maxlat + dlat / 2, int((maxlat - minlat) / dlat) + 2) if minlat <= -90: lat_b = lat_b.clip(-90, None) if maxlat >= 90: lat_b = lat_b.clip(None, 90) lat = (lat_b[1:] + lat_b[:-1]) / 2 lon = (lon_b[1:] + lon_b[:-1]) / 2 # trim grid bounds when your desired extent is not the same as your # initial grid extent if out_extent != in_extent: [minlon, maxlon, minlat, maxlat] = out_extent minlon_ind = np.nonzero(lon >= minlon) maxlon_ind = np.nonzero(lon <= maxlon) lon_inds = np.intersect1d(minlon_ind, maxlon_ind) lon = lon[lon_inds] # make sure to get edges of grid correctly lon_inds = np.append(lon_inds, np.max(lon_inds) + 1) lon_b = lon_b[lon_inds] minlat_ind = np.nonzero(lat >= minlat) maxlat_ind = np.nonzero(lat <= maxlat) lat_inds = np.intersect1d(minlat_ind, maxlat_ind) lat = lat[lat_inds] # make sure to get edges of grid correctly lat_inds = np.append(lat_inds, np.max(lat_inds) + 1) lat_b = lat_b[lat_inds] llgrid = {'lat': lat, 'lon': lon, 'lat_b': lat_b, 'lon_b': lon_b} return llgrid
[docs] def make_grid_cs(csres): """ Creates a cubed-sphere grid description. Parameters ---------- csres : int Cubed-sphere resolution of target grid. Returns ------- csgrid : dict Dict of format {'lat': lat midpoints, 'lon': lon midpoints, 'lat_b': lat edges, 'lon_b': lon edges} where each value has an extra face dimension of length 6. csgrid_list : list of dict List of dicts separated by face index. """ csgrid = csgrid_gmao(csres) csgrid_list = [None] * 6 for i in range(6): csgrid_list[i] = {'lat': csgrid['lat'][i], 'lon': csgrid['lon'][i], 'lat_b': csgrid['lat_b'][i], 'lon_b': csgrid['lon_b'][i]} return [csgrid, csgrid_list]
[docs] def make_grid_sg(csres, stretch_factor, target_lon, target_lat): """ Creates a stretched-grid grid description. Parameters ---------- csres : int Cubed-sphere resolution of target grid. stretch_factor : float Stretch factor of target grid. target_lon : float Target stretching longitude of target grid. target_lat : float Target stretching latitude of target grid. Returns ------- csgrid : dict Dict of format {'lat': lat midpoints, 'lon': lon midpoints, 'lat_b': lat edges, 'lon_b': lon edges} where each value has an extra face dimension of length 6. csgrid_list : list of dict List of dicts separated by face index. """ csgrid = csgrid_gmao(csres, offset=0) csgrid_list = [None] * 6 for i in range(6): lat = csgrid['lat'][i].flatten() lon = csgrid['lon'][i].flatten() lon, lat = scs_transform( lon, lat, stretch_factor, target_lon, target_lat) lat = lat.reshape((csres, csres)) lon = lon.reshape((csres, csres)) lat_b = csgrid['lat_b'][i].flatten() lon_b = csgrid['lon_b'][i].flatten() lon_b, lat_b = scs_transform( lon_b, lat_b, stretch_factor, target_lon, target_lat) lat_b = lat_b.reshape((csres + 1, csres + 1)) lon_b = lon_b.reshape((csres + 1, csres + 1)) csgrid_list[i] = {'lat': lat, 'lon': lon, 'lat_b': lat_b, 'lon_b': lon_b} for i in range(6): csgrid['lat'][i] = csgrid_list[i]['lat'] csgrid['lon'][i] = csgrid_list[i]['lon'] csgrid['lat_b'][i] = csgrid_list[i]['lat_b'] csgrid['lon_b'][i] = csgrid_list[i]['lon_b'] return [csgrid, csgrid_list]
[docs] def calc_rectilinear_lon_edge(lon_stride, center_at_180): """ Compute longitude edge vector for a rectilinear grid. Parameters ---------- lon_stride : float Stride length in degrees. For example, for a standard GEOS-Chem Classic 4x5 grid, lon_stride would be 5. center_at_180 : bool Whether or not the grid should have a cell center at 180 degrees (i.e. on the date line). If true, the first grid cell is centered on the date line; if false, the first grid edge is on the date line. Returns ------- lon_edge : numpy.ndarray Longitudes of cell edges in degrees East. Notes ----- All values are forced to be between [-180,180]. For a grid with N cells in each band, N+1 edges will be returned, with the first and last value being duplicates. """ n_lon_edge = int(np.round(360.0 / lon_stride)) + 1 lon_edge = np.linspace(-180.0, 180.0, num=n_lon_edge) if center_at_180: lon_edge = lon_edge - (lon_stride / 2.0) lon_edge[lon_edge < -180.0] = lon_edge[lon_edge < -180] + 360.0 lon_edge[lon_edge > 180.0] = lon_edge[lon_edge > 180.0] - 360.0 return lon_edge
[docs] def calc_rectilinear_lat_edge(lat_stride, half_polar_grid): """ Compute latitude edge vector for a rectilinear grid. Parameters ---------- lat_stride : float Stride length in degrees. For example, for a standard GEOS-Chem Classic 4x5 grid, lat_stride would be 4. half_polar_grid : bool Whether or not the grid should be "half-polar" (i.e. bands at poles are half the size). In either case the grid will start and end at -/+ 90, but when half_polar_grid is True, the first and last bands will have a width of 1/2 the normal lat_stride. Returns ------- lat_edge : numpy.ndarray Latitudes of cell edges in degrees North. Notes ----- All values are forced to be between [-90,90]. For a grid with N cells in each band, N+1 edges will be returned, with the first and last value being duplicates. """ if half_polar_grid: start_pt = 90.0 + (lat_stride / 2.0) else: start_pt = 90.0 n_lat_edge = int(np.round(2.0 * start_pt / lat_stride)) + 1 lat_edge = np.linspace(-1.0 * start_pt, start_pt, num=n_lat_edge) # Force back onto +/- 90 lat_edge[lat_edge > 90.0] = 90.0 lat_edge[lat_edge < -90.0] = -90.0 return lat_edge
[docs] def calc_rectilinear_grid_area(lon_edge, lat_edge): """ Compute grid cell areas (in m2) for a rectilinear grid. Parameters ---------- lon_edge : numpy.ndarray Grid box longitude edges (in degrees east). lat_edge : numpy.ndarray Grid box latitude edges (in degrees north). Returns ------- grid_area : numpy.ndarray Array of grid box areas in m2. """ # Convert from km to m _radius_earth_m = R_EARTH_m lon_edge = np.asarray(lon_edge, dtype=float) lat_edge = np.asarray(lat_edge, dtype=float) n_lon = (lon_edge.size) - 1 n_lat = (lat_edge.size) - 1 grid_area = np.zeros((n_lat, n_lon)) sfc_area_const = 2.0 * np.pi * _radius_earth_m * _radius_earth_m # Longitudes loop, so need to be careful lon_delta = calc_delta_lon(lon_edge) # Convert into weights relative to the total circle lon_delta = lon_delta / 360.0 # Precalculate this sin_lat_edge = np.sin(np.deg2rad(lat_edge)) for i_lat in range(0, n_lat): sin_diff = sin_lat_edge[i_lat + 1] - sin_lat_edge[i_lat] grid_area[i_lat, :] = sin_diff * sfc_area_const * lon_delta return grid_area
[docs] def calc_delta_lon(lon_edge): """ Compute grid cell longitude widths from an edge vector. Parameters ---------- lon_edge : numpy.ndarray Vector of longitude edges, in degrees East. Returns ------- lon_delta : numpy.ndarray Width of each cell in degrees East. Notes ----- Accounts for looping over the domain. """ n_lon = (lon_edge.size) - 1 lon_edge = np.asarray(lon_edge) # Set up output array lon_delta = np.zeros((n_lon)) offset = 0.0 next_lon = lon_edge[0] for i_lon in range(0, n_lon): last_lon = next_lon next_lon = lon_edge[i_lon + 1] + offset while next_lon < last_lon: offset = offset + 360.0 next_lon = next_lon + 360.0 lon_delta[i_lon] = next_lon - last_lon return lon_delta
[docs] def csgrid_gmao(res, offset=-10): """ Return cubed-sphere coordinates with GMAO face orientation. Parameters ---------- res : int Cubed-sphere resolution. offset : float, optional Degrees to offset grid in the longitudinal direction. Default value: -10 Returns ------- grid : dict Dictionary with keys 'lon', 'lat', 'lon_b', 'lat_b' containing the cubed-sphere grid coordinates. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ cs = CSGrid(res, offset=offset) lon = cs.lon_center.transpose(2, 0, 1) lon_b = cs.lon_edge.transpose(2, 0, 1) lat = cs.lat_center.transpose(2, 0, 1) lat_b = cs.lat_edge.transpose(2, 0, 1) lon[lon < 0] += 360 lon_b[lon_b < 0] += 360 for a in [lon, lon_b, lat, lat_b]: for tile in [0, 1, 3, 4]: a[tile] = a[tile].T for tile in [3, 4]: a[tile] = np.flip(a[tile], 1) for tile in [3, 4, 2, 5]: a[tile] = np.flip(a[tile], 0) a[2], a[5] = a[5].copy(), a[2].copy() # swap north&south pole return {'lon': lon, 'lat': lat, 'lon_b': lon_b, 'lat_b': lat_b}
_INV_SQRT_3 = 1.0 / np.sqrt(3.0) _ASIN_INV_SQRT_3 = np.arcsin(_INV_SQRT_3)
[docs] class CSGrid(object): """ Generator for cubed-sphere grid geometries. CSGrid computes the latitutde and longitudes of cell centers and edges on a cubed-sphere grid, providing a way to retrieve these geometries on-the-fly if your model output data does not include them. Attributes ---------- lon_center, lat_center : numpy.ndarray Lat/lon coordinates for each cell center along the cubed-sphere mesh. lon_edge, lat_edge : numpy.ndarray Lat/lon coordinates for the midpoint of the edges separating each element on the cubed-sphere mesh. xyz_center, xyz_edge : numpy.ndarray As above, except coordinates are projected into a 3D cartesian space with common origin to the original lat/lon coordinate system, assuming a unit sphere. Notes ----- This class was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """
[docs] def __init__(self, c, offset=None): """ Parameters ---------- c : int Number of edges along each cubed-sphere edge. ======= ==================== C Lat/Lon Resolution ------- -------------------- 24 4 deg x 5 deg 48,45 2 deg x 2.5 deg 96,90 1 deg x 1.25 deg 192,180 0.5 deg x 0.625 deg 384,360 0.25 deg x 0.3125 deg 720 0.125 deg x 0.15625 deg ======= ==================== offset : float, optional Degrees to offset the first faces' edge in the latitudinal direction. If not passed, then the western edge of the first face will align with the prime meridian. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ self.c = c self.delta_y = 2. * _ASIN_INV_SQRT_3 / c self.nx = self.ny = c + 1 self.offset = offset self._initialize()
def _initialize(self): c = self.c nx, ny = self.nx, self.ny lambda_rad = np.zeros((nx, ny)) lambda_rad[0, :] = 3. * np.pi / 4. # West edge lambda_rad[-1, :] = 5. * np.pi / 4. # East edge theta_rad = np.zeros((nx, ny)) theta_rad[0, :] = -_ASIN_INV_SQRT_3 + \ (self.delta_y * np.arange(c + 1)) # West edge theta_rad[-1, :] = theta_rad[0, :] # East edge # Cache the reflection points - our upper-left and lower-right corners lon_mir1, lon_mir2 = lambda_rad[0, 0], lambda_rad[-1, -1] lat_mir1, lat_mir2 = theta_rad[0, 0], theta_rad[-1, -1] xyz_mir1 = latlon_to_cartesian(lon_mir1, lat_mir1) xyz_mir2 = latlon_to_cartesian(lon_mir2, lat_mir2) xyz_cross = np.cross(xyz_mir1, xyz_mir2) norm = np.sqrt(np.sum(xyz_cross**2)) xyz_cross /= norm for i in range(1, c): lon_ref, lat_ref = lambda_rad[0, i], theta_rad[0, i] xyz_ref = np.asarray(latlon_to_cartesian(lon_ref, lat_ref)) xyz_dot = np.sum(xyz_cross * xyz_ref) xyz_img = xyz_ref - (2. * xyz_dot * xyz_cross) xs_img, ys_img, zs_img = xyz_img lon_img, lat_img = cartesian_to_latlon(xs_img, ys_img, zs_img) lambda_rad[i, 0] = lon_img lambda_rad[i, -1] = lon_img theta_rad[i, 0] = lat_img theta_rad[i, -1] = -lat_img pp = np.zeros([3, c + 1, c + 1]) # Set the four corners # print("CORNERS") for i, j in product([0, -1], [0, -1]): # print(i, j) pp[:, i, j] = latlon_to_cartesian( lambda_rad[i, j], theta_rad[i, j]) # Map the edges on the sphere back to the cube. #Note that all intersections are at x = -rsq3 # print("EDGES") for ij in range(1, c + 1): # print(ij) pp[:, 0, ij] = latlon_to_cartesian( lambda_rad[0, ij], theta_rad[0, ij]) pp[1, 0, ij] = -pp[1, 0, ij] * _INV_SQRT_3 / pp[0, 0, ij] pp[2, 0, ij] = -pp[2, 0, ij] * _INV_SQRT_3 / pp[0, 0, ij] pp[:, ij, 0] = latlon_to_cartesian( lambda_rad[ij, 0], theta_rad[ij, 0]) pp[1, ij, 0] = -pp[1, ij, 0] * _INV_SQRT_3 / pp[0, ij, 0] pp[2, ij, 0] = -pp[2, ij, 0] * _INV_SQRT_3 / pp[0, ij, 0] # # Map interiors pp[0, :, :] = -_INV_SQRT_3 # print("INTERIOR") for i in range(1, c + 1): for j in range(1, c + 1): # Copy y-z face of the cube along j=1 pp[1, i, j] = pp[1, i, 0] # Copy along i=1 pp[2, i, j] = pp[2, 0, j] _pp = pp.copy() llr, ttr = vec_cartesian_to_latlon(_pp[0], _pp[1], _pp[2]) lambda_rad, theta_rad = llr.copy(), ttr.copy() # Make grid symmetrical to i = im/2 + 1 for j in range(1, c + 1): for i in range(1, c + 1): # print("({}, {}) -> ({}, {})".format(i, 0, i, j)) lambda_rad[i, j] = lambda_rad[i, 0] for j in range(c + 1): for i in range(c // 2): isymm = c - i # print(isymm) avg_pt = 0.5 * (lambda_rad[i, j] - lambda_rad[isymm, j]) # print(lambda_rad[i, j], lambda_rad[isymm, j], avg_pt) lambda_rad[i, j] = avg_pt + np.pi lambda_rad[isymm, j] = np.pi - avg_pt avg_pt = 0.5 * (theta_rad[i, j] + theta_rad[isymm, j]) theta_rad[i, j] = avg_pt theta_rad[isymm, j] = avg_pt # Make grid symmetrical to j = im/2 + 1 for j in range(c // 2): jsymm = c - j for i in range(1, c + 1): avg_pt = 0.5 * (lambda_rad[i, j] + lambda_rad[i, jsymm]) lambda_rad[i, j] = avg_pt lambda_rad[i, jsymm] = avg_pt avg_pt = 0.5 * (theta_rad[i, j] - theta_rad[i, jsymm]) theta_rad[i, j] = avg_pt theta_rad[i, jsymm] = -avg_pt # Final correction lambda_rad -= np.pi llr, ttr = lambda_rad.copy(), theta_rad.copy() ####################################################################### # MIRROR GRIDS ####################################################################### new_xgrid = np.zeros((c + 1, c + 1, 6)) new_ygrid = np.zeros((c + 1, c + 1, 6)) xgrid = llr.copy() ygrid = ttr.copy() new_xgrid[..., 0] = xgrid.copy() new_ygrid[..., 0] = ygrid.copy() # radius = 6370.0e3 radius = 1. for face in range(1, 6): for j in range(c + 1): for i in range(c + 1): x = xgrid[i, j] y = ygrid[i, j] z = radius if face == 1: # Rotate about z only new_xyz = rotate_sphere_3d(x, y, z, -np.pi / 2., 'z') elif face == 2: # Rotate about z, then x temp_xyz = rotate_sphere_3d(x, y, z, -np.pi / 2., 'z') x, y, z = temp_xyz[:] new_xyz = rotate_sphere_3d(x, y, z, np.pi / 2., 'x') elif face == 3: temp_xyz = rotate_sphere_3d(x, y, z, np.pi, 'z') x, y, z = temp_xyz[:] new_xyz = rotate_sphere_3d(x, y, z, np.pi / 2., 'x') if ((c % 2) != 0) and (j == c // 2 - 1): print(i, j, face) new_xyz = (np.pi, *new_xyz) elif face == 4: temp_xyz = rotate_sphere_3d(x, y, z, np.pi / 2., 'z') x, y, z = temp_xyz[:] new_xyz = rotate_sphere_3d(x, y, z, np.pi / 2., 'y') elif face == 5: temp_xyz = rotate_sphere_3d(x, y, z, np.pi / 2., 'y') x, y, z = temp_xyz[:] new_xyz = rotate_sphere_3d(x, y, z, 0., 'z') # print((x, y, z), "\n", new_xyz, "\n" + "--"*40) new_x, new_y, _ = new_xyz new_xgrid[i, j, face] = new_x new_ygrid[i, j, face] = new_y lon_edge, lat_edge = new_xgrid.copy(), new_ygrid.copy() ####################################################################### # CLEANUP GRID ####################################################################### for i, j, f in product(range(c + 1), range(c + 1), range(6)): new_lon = lon_edge[i, j, f] if new_lon < 0: new_lon += 2 * np.pi if np.abs(new_lon) < 1e-10: new_lon = 0. lon_edge[i, j, f] = new_lon if np.abs(lat_edge[i, j, f]) < 1e-10: lat_edge[i, j, f] = 0. lon_edge_deg = np.rad2deg(lon_edge) lat_edge_deg = np.rad2deg(lat_edge) ####################################################################### # COMPUTE CELL CENTROIDS ####################################################################### lon_ctr = np.zeros((c, c, 6)) lat_ctr = np.zeros((c, c, 6)) xyz_ctr = np.zeros((3, c, c, 6)) xyz_edge = np.zeros((3, c + 1, c + 1, 6)) for f in range(6): for i in range(c): last_x = i == (c - 1) for j in range(c): last_y = j == (c - 1) # Get the four corners lat_corner = [ lat_edge[i, j, f], lat_edge[i + 1, j, f], lat_edge[i + 1, j + 1, f], lat_edge[i, j + 1, f]] lon_corner = [ lon_edge[i, j, f], lon_edge[i + 1, j, f], lon_edge[i + 1, j + 1, f], lon_edge[i, j + 1, f]] # Convert from lat-lon back to cartesian xyz_corner = np.asarray( vec_latlon_to_cartesian( lon_corner, lat_corner)) # Store the edge information xyz_edge[:, i, j, f] = xyz_corner[:, 0] if last_x: xyz_edge[:, i + 1, j, f] = xyz_corner[:, 1] if last_x or last_y: xyz_edge[:, i + 1, j + 1, f] = xyz_corner[:, 2] if last_y: xyz_edge[:, i, j + 1, f] = xyz_corner[:, 3] e_mid = np.sum(xyz_corner, axis=1) e_abs = np.sqrt(np.sum(e_mid * e_mid)) if e_abs > 0: e_mid = e_mid / e_abs xyz_ctr[:, i, j, f] = e_mid _lon, _lat = cartesian_to_latlon(*e_mid) lon_ctr[i, j, f] = _lon lat_ctr[i, j, f] = _lat lon_ctr_deg = np.rad2deg(lon_ctr) lat_ctr_deg = np.rad2deg(lat_ctr) if self.offset is not None: lon_edge_deg += self.offset lon_ctr_deg += self.offset ####################################################################### # CACHE ####################################################################### self.lon_center = lon_ctr_deg self.lat_center = lat_ctr_deg self.lon_edge = lon_edge_deg self.lat_edge = lat_edge_deg self.xyz_center = xyz_ctr self.xyz_edge = xyz_edge
[docs] def latlon_to_cartesian(lon, lat): """ Convert latitude/longitude coordinates along the unit sphere to cartesian coordinates defined by a vector pointing from the sphere's center to its surface. Parameters ---------- lon : float Longitude coordinate in radians. lat : float Latitude coordinate in radians. Returns ------- x, y, z : float Cartesian coordinates on the unit sphere. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ x = np.cos(lat) * np.cos(lon) y = np.cos(lat) * np.sin(lon) z = np.sin(lat) return x, y, z
vec_latlon_to_cartesian = np.vectorize(latlon_to_cartesian)
[docs] def cartesian_to_latlon(x, y, z, ret_xyz=False): """ Convert a cartesian coordinate to latitude/longitude coordinates. Optionally return the original cartesian coordinate as a tuple. Parameters ---------- x : float Cartesian x-coordinate. y : float Cartesian y-coordinate. z : float Cartesian z-coordinate. ret_xyz : bool, optional If True, also return the normalized cartesian coordinate vector. Default value: False Returns ------- lon : float Longitude in radians. lat : float Latitude in radians. xyz : numpy.ndarray, optional Normalized cartesian coordinate vector. Only returned if ret_xyz is True. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ xyz = np.array([x, y, z]) vector_length = np.sqrt(np.sum(xyz * xyz, axis=0)) xyz /= vector_length x, y, z = xyz if (np.abs(x) + np.abs(y)) < 1e-20: lon = 0.0 else: lon = np.arctan2(y, x) if lon < 0.0: lon += 2 * np.pi lat = np.arcsin(z) # If not normalizing vector, take lat = np.arcsin(z/vector_length) if ret_xyz: return lon, lat, xyz return lon, lat
vec_cartesian_to_latlon = np.vectorize(cartesian_to_latlon)
[docs] def spherical_to_cartesian(theta, phi, r=1): """ Convert spherical coordinates in the form (theta, phi[, r]) to cartesian, with the origin at the center of the original spherical coordinate system. Parameters ---------- theta : float Azimuthal angle in radians. phi : float Polar angle in radians. r : float, optional Radius. Default value: 1 Returns ------- x, y, z : float Cartesian coordinates. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ x = r * np.cos(phi) * np.cos(theta) y = r * np.cos(phi) * np.sin(theta) z = r * np.sin(phi) return x, y, z
vec_spherical_to_cartesian = np.vectorize(spherical_to_cartesian)
[docs] def cartesian_to_spherical(x, y, z): """ Convert cartesian coordinates to spherical in the form (theta, phi[, r]) with the origin remaining at the center of the original spherical coordinate system. Parameters ---------- x : float Cartesian x-coordinate. y : float Cartesian y-coordinate. z : float Cartesian z-coordinate. Returns ------- theta : float Azimuthal angle in radians. phi : float Polar angle in radians. r : float Radius. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ r = np.sqrt(x**2 + y**2 + z**2) #theta = np.arccos(z / r) theta = np.arctan2(y, x) phi = np.arctan2(z, np.sqrt(x**2 + y**2)) # if np.abs(x) < 1e-16: # phi = np.pi # else: # phi = np.arctan(y / x) return theta, phi, r
vec_cartesian_to_spherical = np.vectorize(cartesian_to_spherical)
[docs] def rotate_sphere_3d(theta, phi, r, rot_ang, rot_axis='x'): """ Rotate a spherical coordinate in the form (theta, phi[, r]) about the indicated axis, rot_axis. This method accomplishes the rotation by projecting to a cartesian coordinate system and performing a solid body rotation around the requested axis. Parameters ---------- theta : float Azimuthal angle in radians. phi : float Polar angle in radians. r : float Radius. rot_ang : float Rotation angle in radians. rot_axis : str, optional Axis about which to rotate ('x', 'y', or 'z'). Default value: 'x' Returns ------- theta_new : float New azimuthal angle in radians after rotation. phi_new : float New polar angle in radians after rotation. r_new : float New radius after rotation. Notes ----- This function was originally written by Jiawei Zhuange and included in package cubedsphere: https://github.com/JiaweiZhuang/cubedsphere """ cos_ang = np.cos(rot_ang) sin_ang = np.sin(rot_ang) x, y, z = spherical_to_cartesian(theta, phi, r) if rot_axis == 'x': x_new = x y_new = cos_ang * y + sin_ang * z z_new = -sin_ang * y + cos_ang * z elif rot_axis == 'y': x_new = cos_ang * x - sin_ang * z y_new = y z_new = sin_ang * x + cos_ang * z elif rot_axis == 'z': x_new = cos_ang * x + sin_ang * y y_new = -sin_ang * x + cos_ang * y z_new = z else: raise ValueError(f"Invalid value for rot_axis: {rot_axis}") theta_new, phi_new, r_new = cartesian_to_spherical(x_new, y_new, z_new) return theta_new, phi_new, r_new
[docs] def get_nearest_model_data_cs( gc_data, gc_cs_grid, lon_value, lat_value, varlist=None ): """ Returns GEOS-Chem model data (on a cubed-sphere grid) at the grid box closest to a given (lat, lon) location. Parameters ---------- gc_data : xarray.DataArray or xarray.Dataset GEOS-Chem model data for a single variable. gc_cs_grid : xarray.Dataset Coordinate arrays defining the cubed-sphere grid. lon_value : float Longitude at the location of interest. lat_value : float Latitude at the location of interest. varlist : list of str, optional List of data variables to include in the output. Returns ------- dataframe : pandas.DataFrame Model data closest to the observation site. """ verify_variable_type(gc_data, (xr.DataArray, xr.Dataset)) verify_variable_type(gc_cs_grid, xr.Dataset) # Prevent the latitude from getting too close to the N or S poles lat_value = max(min(lat_value, 89.75), -89.75) # Indices (nf, yInd, xInd) of box nearest to observation site cs_indices = find_index( lat_value, lon_value, gc_cs_grid ) if varlist is not None: return gc_data[varlist].isel( nf=cs_indices[0, 0], Ydim=cs_indices[1, 0], Xdim=cs_indices[2, 0], ).to_dataframe() return gc_data.isel( nf=cs_indices[0, 0], Ydim=cs_indices[1, 0], Xdim=cs_indices[2, 0], ).to_dataframe()
[docs] def get_nearest_model_data_ll( gc_data, lon_value, lat_value, varlist=None ): """ Returns GEOS-Chem model data (on a lat/lon grid) at the grid box closest to a given (lat, lon) location. Parameters ---------- gc_data : xarray.DataArray or xarray.Dataset GEOS-Chem model data. lon_value : float Longitude at the location of interest. lat_value : float Latitude at the location of interest. varlist : list of str, optional List of data variables to include in the output. Returns ------- dataframe : pandas.DataFrame Model data closest to the observation site. """ verify_variable_type(gc_data, (xr.DataArray, xr.Dataset)) x_idx=( np.abs( gc_data.lon.values - float(lon_value) ) ).argmin() y_idx=( np.abs( gc_data.lat.values - float(lat_value) ) ).argmin() if varlist is not None: return gc_data[varlist].isel( lon=x_idx, lat=y_idx, ).to_dataframe() return gc_data.isel( lon=x_idx, lat=y_idx, ).to_dataframe()
[docs] def get_nearest_model_data( gc_data, lon_value, lat_value, gc_cs_grid=None, varlist=None ): """ Returns GEOS-Chem model data at the grid box closest to a given (lat, lon) location, dispatching to the appropriate function depending on the grid type. Parameters ---------- gc_data : xarray.DataArray or xarray.Dataset GEOS-Chem data for a single variable. lon_value : float Longitude at the location of interest. lat_value : float Latitude at the location of interest. gc_cs_grid : xarray.Dataset, optional Coordinate arrays defining the cubed-sphere grid. This can be obtained as the output of function gcpy.util.extract_data(). varlist : list of str, optional List of data variables to include in the output. Returns ------- dataframe : pandas.DataFrame Model data closest to the observation site. """ verify_variable_type(gc_data, (xr.DataArray, xr.Dataset)) verify_variable_type(gc_cs_grid, (xr.Dataset, type(None))) if is_cubed_sphere(gc_data): return get_nearest_model_data_cs( gc_data, gc_cs_grid, lon_value, lat_value, varlist=varlist ) return get_nearest_model_data_ll( gc_data, lon_value, lat_value, varlist=varlist, )