Source code for pypago.grid


''' Class that handles grid scale factors for a model configuration '''

import numpy as np
import pypago.disp
import pypago.coords
import pylab as plt

[docs]class Grid(object): """ Class that handles grid extraction on a specific domain :param int jmin: Index of the southernmost point of the subdomain :param int jmax: Index of the northernmost point of the subdomain :param int imin: Index of the westernmost point of the subdomain :param int imax: Index of the easternnmost point of the subdomain :param :py:class:`pypago.coords.Coords` coord: Coordinate object associated with the grid file (longitude, latitude, mask of the entire domain, used in the file extraction) """ def __init__(self, coord, jmin, jmax, imin, imax): ''' Initialisation of the grid class ''' # here are the coords and scale factors # extracted on the domain grid. self.latt = None self.lont = None self.mask = None self.dyw = None self.bathy = None self.dxt = None self.dyt = None self.dyte = None self.dxn = None self.dzt = None self.dzw = None self.areaw = None self.arean = None self.volume = None self.surface = None self.jmin = None self.jmax = None self.imin = None self.imax = None self.filename = None self.modelname = None self.nlon = None if coord is not None: for arg in [jmin, jmax, imin, imax]: if arg is None: message = 'Either jmin, jmax, imin or imax ' message += 'is None.' raise PypagoErrors(message) nlat, nlon = coord.lont.shape # Checking that the jmin/jmax/imin/imin are within the bounds # if not, correct their values if (imin<0) | (imin>nlon-1): print('The imin argument must be ' + \ 'between %d and %d. Currently, %d' %(0, nlon-1, imin)) print('imin has been set to %d' %(0)) imin = 0 if (imax<0) | (imax>nlon-1): print('The imax argument must be ' + \ 'between %d and %d. Currently, %d' %(0, nlon-1, imax)) print('imax has been set to %d' %(nlon-1)) imax = nlon-1 if (jmin<1) | (jmin>nlat-2): print('The jmin argument must be ' + \ 'between %d and %d. Currently, %d' %(1, nlat-2, jmin)) print('jmin has been set to %d' %(1)) jmin = 1 if (jmax<1) | (jmax>nlat-2): print('The jmax argument must be ' + \ 'between %d and %d. Currently, %d' %(1, nlat-2, jmax)) print('jmax has been set to %d' %(nlat-2)) jmax = nlat-2 if (jmax<jmin): print("Currently, jmax<jmin. " + \ "The variables have been switched") jmax, jmin = jmin, jmax self.filename = coord.filename self.modelname = coord.modelname self.nlon = nlon self.jmin = jmin self.jmax = jmax self.imin = imin self.imax = imax self.extract_all_var(coord)
[docs] def plot_dom(self, ax=None): """ Draws the domain defined by the jmax, jmin, imin, and imax attributes. :param matplotlib.axes.Axes ax: Axis on which to draw the figure. If None, draws on current axis. """ if ax is None: ax = plt.gca() if self.imin < self.imax: lontp = [self.imin, self.imax, self.imax, self.imin, self.imin] lattp = [self.jmin, self.jmin, self.jmax, self.jmax, self.jmin] ax.plot(lontp, lattp, color='r') else: lontp = [self.nlon, self.imin, self.imin, self.nlon] lattp = [self.jmin, self.jmin, self.jmax, self.jmax] ax.plot(lontp, lattp, color='r') lontp = [0, self.imax, self.imax, 0] lattp = [self.jmin, self.jmin, self.jmax, self.jmax] ax.plot(lontp, lattp, color='r')
def __str__(self): output = 'Model grid for the %s model:\n' %self.modelname output += ' - mesh file: %s\n' %self.filename output += ' - jmin: %d\n' %self.jmin output += ' - jmax: %d\n' %self.jmax output += ' - imin: %d\n' %self.imin output += ' - imax: %d\n' %self.imax output += ' - nlat: %d\n' %self.mask.shape[0] output += ' - nlon: %d\n' %self.mask.shape[1] if self.volume is not None: output += ' - nz: %d\n' %self.volume.shape[0] return output
[docs] def compute_areas(self): """ Computes the cell area at the northern and western faces, and the volume and surface of the T grid cells. """ nz = self.dzt.shape[0] self.areaw = self.dzw * np.tile(self.dyw, (nz, 1, 1)) self.arean = self.dzn * np.tile(self.dxn, (nz, 1, 1)) self.areaw[self.areaw==0] = np.nan self.areaw[np.ma.getmaskarray(self.areaw)==1] = np.nan self.arean[self.arean==0] = np.nan self.arean[np.ma.getmaskarray(self.arean)==1] = np.nan self.surface = self.dxt * self.dyt self.volume = self.dzt * np.tile(self.surface, (nz, 1, 1))
[docs] def extract_2d_var(self, varin): """ Extraction of a 2D variable on the subdomain defined by the four domain limits. It provides the possibility to read a variable "discontinuously" (if `imin` > `imax`). :param numpy.array varin: Input 2D variable :return: The input array extracted on the specified subdomain :rtype: numpy.array """ if self.imin > self.imax: out1 = varin[self.jmin:self.jmax+1, self.imin:] out2 = varin[self.jmin:self.jmax+1, :self.imax+1] varout = np.concatenate((out1, out2), axis=-1) else: varout = varin[self.jmin:self.jmax+1, self.imin:self.imax+1] return varout
[docs] def extract_3d_var(self, varin): """ Extraction of a 2D variable on the subdomain defined by the four domain limits. It provides the possibility to read a variable "discontinuously" (if `imin` > `imax`). :param numpy.array varin: Input 3D variable :return: The input array extracted on the specified subdomain :rtype: numpy.array """ if self.imin > self.imax: out1 = varin[:, self.jmin:self.jmax+1, self.imin:] out2 = varin[:, self.jmin:self.jmax+1, :self.imax+1] varout = np.concatenate((out1, out2), axis=-1) else: varout = varin[:, self.jmin:self.jmax+1, self.imin:self.imax+1] return varout
[docs] def create_3d_vert_scalefact(self): """ If the vertical scale factor dzt is 1D, it reconstructs a 3D variable by repeating the 1D variable along the lon and lat dimensions """ ny, nx = self.dxt.shape self.dzt = np.transpose(np.tile(self.dzt, [ny, nx, 1]), [2, 0, 1]) self.dzw = self.dzt.copy() self.dzn = self.dzt.copy()
[docs] def extract_all_var(self, coord): """ Extracts all the variables (coordinates, scale factors, etc) on the sub-domain defined by the jmin, jmax, imin and imax variables. """ print("Extraction of lont on the domain") self.lont = self.extract_2d_var(coord.lont) print("Extraction of latt on the domain") self.latt = self.extract_2d_var(coord.latt) print("Extraction of bathy on the domain") self.bathy = self.extract_2d_var(coord.bathy) print("Extraction of mask on the domain") self.mask = self.extract_2d_var(coord.mask) print("Extraction of dxt on the domain") self.dxt = self.extract_2d_var(coord.dxt) print("Extraction of dyt on the domain") self.dyt = self.extract_2d_var(coord.dyt) print("Extraction of dxn on the domain") self.dxn = self.extract_2d_var(coord.dxn) if coord.dyw is not None: # If the model contains the dyw variable (HYCO, OFAM, MICO models for instance) print("Extraction of dyw on the domain") self.dyw = self.extract_2d_var(coord.dyw) else: # If the model does not contain the dyw variable, we reconstruct it from the # dye variable print("Reconstruction of the dyw variable from dye") if self.imin < self.imax: if self.imin > 0: self.dyw = coord.dye[self.jmin:self.jmax+1, self.imin-1:self.imax] else: self.dyw = np.concatenate((coord.dye[self.jmin:self.jmax+1, -1:], coord.dye[self.jmin:self.jmax+1, self.imin:self.imax]), axis=-1) else: # self.imax < self.imin if self.imax > 0: self.dyw = np.concatenate((coord.dye[self.jmin:self.jmax+1, self.imin-1:], coord.dye[self.jmin:self.jmax+1, :self.imax]), axis=-1) else: self.dyw = coord.dye[self.jmin:self.jmax, self.imin-1:]
[docs] def plot_mask(self, ax=None): ''' Contours the mask attribute ''' if ax is None: ax = plt.gca() cs = ax.contour(self.mask, levels=[1 - np.spacing(1), 1], colors='k') ax.set_xlim(0, self.mask.shape[1]-1) ax.set_ylim(0, self.mask.shape[0]-1) return cs
[docs]class CommonGrid(Grid): ''' Grid class associated with all models but GFDL ''' def __init__(self, coord, jmin, jmax, imin, imax): ''' Initialisation of the common grid class ''' # initialisation of the attributes to None # using the mother initialisation super(CommonGrid, self).__init__(coord, jmin, jmax, imin, imax) if coord.dzt.ndim > 2: self.dzt = self.extract_3d_var(coord.dzt) self.dzw = self.extract_3d_var(coord.dzw) self.dzn = self.extract_3d_var(coord.dzn) else: # if dzt is 1D, we create 3D scale factors # from 1D scale factors self.create_3d_vert_scalefact() self.compute_areas()
[docs]class GfdlGrid(Grid): ''' GFDL grid class ''' def __init__(self, coord, jmin, jmax, imin, imax): ''' Initialisation of the GFDL grid class ''' # initialisation of the attributes to None # using the mother initialisation super(GfdlGrid, self).__init__(coord, jmin, jmax, imin, imax) self.dzc = None self.extract_3d_var_gfdl(coord) self.compute_areas()
[docs] def extract_3d_var_gfdl(self, coord): """ Creates the vertical scale factors for the GFDL model. """ if coord.dzt.ndims > 2: if self.imax < self.imin: dzc = np.concatenate((coord.dzc[:, self.jmin-1:self.jmax+1, self.imin:], coord.dzc[:, self.jmin-1:self.jmax+1, :self.imax+2]), axis = -1) temp = np.array([dzc[:, 1:, :-1], dzc[:, :-1, :-1]]) temp = np.ma.masked_where(temp != temp, temp) self.dzw = np.mean(temp, axis=0) temp = np.array([dzc[:, 1:, 1:], dzc[:, 1:, :-1]]) temp = np.ma.masked_where(temp != temp, temp) self.dzn = np.mean(temp, axis=0) self.dzt = self.extract_3d_var(coord.dzt) else: # imin < imax if self.imax < self.mask.shape[-1]: dzc = coord.dzc[:, self.jmin-1:self.jmax+1, self.imin:self.imax+2] else: dzc = np.concatenate((coord.dzc[:, self.jmin-1:self.jmax+1, self.imin:self.imax+1], coord.dzc[:, self.jmin-1:self.jmax+1, 0:1]), axis=-1) temp = np.array([dzc[:, 1:, :-1], dzc[:, :-1, :-1]]) temp = np.ma.masked_where(temp != temp, temp) self.dzw = np.mean(temp, axis=0) temp = np.array([dzc[:, 1:, 1:], dzc[:, 1:, :-1]]) temp = np.ma.masked_where(temp != temp, temp) self.dzn = np.mean(temp, axis=0) self.dzt = self.extract_3d_var(coord.dzt) else: # if dzt is 1D, we create 3D scale factors # from 1D scale factors self.create_3d_vert_scalefact()
[docs]def create_grid(coord, jmin=None, jmax=None, imin=None, imax=None): """ Returns a Grid instance associated with the input coord argument. If it is a `pypago.coord.GfdlCoords` object, then it returns a :py:class:`pypago.grid.GfdlGrid` object. Else, it returns a :py:class:`pypago.grid.CommonGrid` object. :param int jmin: Index of the southernmost point of the subdomain :param int jmax: Index of the northernmost point of the subdomain :param int imin: Index of the westernmost point of the subdomain :param int imax: Index of the easternnmost point of the subdomain :param :py:class:`pypago.coords.Coords` coord: Coordinate object associated with the grid file (longitude, latitude, mask of the entire domain, used in the file extraction) """ if imin is None: print('The imin argument is None. Set to 0') imin = 0 if imax is None: print('The imax argument is None. Set to %d' %(coord.lont.shape[1]-1)) imax = coord.lont.shape[1] - 1 if jmin is None: jmin = 1 print('The jmin argument is None. Set to 1') if jmax is None: print('The jmax argument is None. Set to %d' %(coord.lont.shape[0]-2)) jmax = coord.lont.shape[0] - 2 if isinstance(coord, pypago.coords.GfdlCoords): return GfdlGrid(coord, jmin, jmax, imin, imax) else: return CommonGrid(coord, jmin, jmax, imin, imax)