Source code for e3sm_nex.api

from . import core
import numpy
import cdp
import cdms2
import MV2

[docs]def generateNEXGrid(lats, lons, elements_corners): """Generates NEx grid from latitudes, longitudes and elements corners :Example: .. doctest:: api_generateNEXGrid >>> grid = generateNEXGrid(lats, lons, elements_corners) :param lats: 1D pure numpy array containing latitudes of each column :type lats: `numpy.ndarray`_ :param lons: 1D pure numpy array containing longitudes of each column :type lons: `numpy.ndarray`_ :param elements_corners: 2D pure numpy array containing latitudes centers (grid_size, nvertices) :type lats_corners: `numpy.ndarray`_ :return: A cdms generic grid object representation of the nex grid :rtype: `cdms2.gengrid.TransientGenericGrid`_ """ ncols = len(lats) mesh = numpy.zeros((ncols,2,4)) params = [] for index in range(ncols): params.append(core.Input(index, elements_corners, lats, lons)) meshes = cdp.cdp_run.multiprocess(core.corners, params) for i in range(ncols): corners = meshes[i] index,l,L = corners mesh[index,0] = numpy.array(l) mesh[index,1] = numpy.array(L) lat_axis = cdms2.auxcoord.TransientAuxAxis1D(lats) lat_axis.id = "lat" lat_axis.setBounds(mesh[:,0]) lon_axis = cdms2.auxcoord.TransientAuxAxis1D(lons) lon_axis.id = "lon" lon_axis.setBounds(mesh[:,1]) grid = cdms2.gengrid.TransientGenericGrid(lat_axis,lon_axis) return grid
[docs]def rads2degrees(data, force=False): """ converting to degrees from radians if necessary""" if data.max()<7. or force: # likely radians data = data / numpy.pi * 180. return data
[docs]def generateMPASGrid(lats, lons, lats_corners, lons_corners, delta=180.): """Generates MPAS grid from latitudes, longitudes, latitudes corners, longitude_corners :Example: .. doctest:: api_generateMPASGrid >>> grid = generateMPASGrid(lats, lons, lats_corners, lons_corners) :param lats: 1D pure numpy array containing latitudes centers :type lats: `numpy.ndarray`_ :param lons: 1D pure numpy array containing longitudes centers :type lons: `numpy.ndarray`_ :param lats_corners: 2D pure numpy array containing latitudes centers (grid_size, nvertices) :type lats_corners: `numpy.ndarray`_ :param lons_corners: 2D pure numpy array containing longitudes centers (grid_size, nvertices) :type lons: `numpy.ndarray`_ :param delta: delta between longitude for which we assume `wrapping` needs to be done :type delta: `float`_ :return: A cdms generic grid object representation of the mpas grid :rtype: `cdms2.gengrid.TransientGenericGrid`_ """ lats = rads2degrees(lats) lons = rads2degrees(lons) lats_corners = rads2degrees(lats_corners) lons_corners = rads2degrees(lons_corners) # First clean up wrap effect for longitudes cont = True # work not finished? iterations = 0 while iterations<5 and cont: mx = lons_corners.max(axis=1) mn = lons_corners.min(axis=1) diff = mx - mn indices = numpy.argwhere(numpy.greater(diff,delta)) mns = mn[indices] lons_bad = lons_corners[indices] mns = numpy.resize(mns,(lons_bad.shape)) lons_corners[indices] = numpy.where(lons_corners[indices]>mns+delta,360.-lons_corners[indices],lons_corners[indices]) if len(indices) == 0: cont = False iterations+=1 # Now gen grid # Grid axes lat_axis = cdms2.auxcoord.TransientAuxAxis1D(lats) lat_axis.setBounds(lats_corners) lon_axis = cdms2.auxcoord.TransientAuxAxis1D(lons) lon_axis.setBounds(lons_corners) # Actual Grid grid = cdms2.gengrid.TransientGenericGrid(lat_axis,lon_axis) return grid
[docs]def applyGrid(data,grid): """Applies grid to data :Example: .. doctest:: api_applyGrid >>> grid = e3sm_nex.generateMPASGrid(lats, lons, lats_corners, lons_corners) >>> data = e3sm_nex.applyGrid(data,grid) :param data: cdms2.MV2 variable :type data: `cdms2.tvariable.TransientVariable`_ :param grid: grid to apply to data :type grid: `cdms2.gengrid.TransientGenericGrid`_ :return: Data with grid applied to it :rtype: `cdms2.tvariable.TransientVariable`_ """ if not isinstance(data,numpy.ndarray): raise RuntimeError("input data to e3smnex's applyGrid function must be numpy array") if not isinstance(data,cdms2.avariable.AbstractVariable): data = MV2.array(data) axes = data.getAxisList() gridAxis = grid.getAxis(0) for i, ax in enumerate(axes): if len(ax) == len(gridAxis): axes[i] = gridAxis data.setAxisList(axes) data.setGrid(grid) return data