def create_source_grid(self,
filename,
from_global,
coord_names,
x_coords=None,
y_coords=None,
autocrop=True):
''' create ESMF grid object for source grid '''
# new way to create source grid
# TO DO : move into separate function, has to be called before drown
# so that we know the periodicity
# Allow to provide lon/lat from existing array
if x_coords is not None and y_coords is not None:
lon_src = x_coords
lat_src = y_coords
else:
lons = _ncdf.read_field(filename, coord_names[0])
lats = _ncdf.read_field(filename, coord_names[1])
if len(lons.shape) == 1:
lon_src, lat_src = _np.meshgrid(lons, lats)
else:
lon_src = lons
lat_src = lats
# autocrop
if autocrop:
imin_src, imax_src, jmin_src, jmax_src = \
_lc.find_subset(self.grid_target,lon_src,lat_src)
lon_src = lon_src[jmin_src:jmax_src, imin_src:imax_src]
lat_src = lat_src[jmin_src:jmax_src, imin_src:imax_src]
ny_src, nx_src = lon_src.shape
if not autocrop:
imin_src = 0
imax_src = nx_src
jmin_src = 0
jmax_src = ny_src
if from_global and not autocrop:
gridsrc = _ESMF.Grid(_np.array([nx_src, ny_src]), num_peri_dims=1)
self.gtype = 1 # 1 = periodic for drown NCL
self.kew = 0 # 0 = periodic for drown sosie
else:
gridsrc = _ESMF.Grid(_np.array([nx_src, ny_src]))
self.gtype = 0 # 1 = non periodic for drown NCL
self.kew = -1 # -1 = non periodic for drown sosie
gridsrc.add_coords(staggerloc=[_ESMF.StaggerLoc.CENTER])
gridsrc.coords[_ESMF.StaggerLoc.CENTER][0][:] = lon_src.T
gridsrc.coords[_ESMF.StaggerLoc.CENTER][1][:] = lat_src.T
# original from RD
#self.gridsrc = _ESMF.Grid(filename=filename,filetype=_ESMF.FileFormat.GRIDSPEC,\
#is_sphere=from_global,coord_names=coord_names)
return gridsrc, imin_src, imax_src, jmin_src, jmax_src
def perform_extrapolation(self, datasrc, maskfile, maskvar, missing_value,
drown):
# 2.1 read mask or compute it
if maskvar is not None:
mask = _ncdf.read_field(maskfile, maskvar)
# to do, needs imin/imax_src,...
else:
mask = self.compute_mask_from_missing_value(
datasrc, missing_value=missing_value)
# 2.2 mask the source data
if _np.ma.is_masked(datasrc):
datasrc = datasrc.data
datasrc[_np.where(mask == 0)] = self.xmsg
datamin = datasrc[_np.where(mask == 1)].min()
datamax = datasrc[_np.where(mask == 1)].max()
if self.debug:
datasrc_plt = _np.ma.masked_values(datasrc, self.xmsg)
_plt.figure()
_plt.contourf(datasrc_plt[0, :, :], 40)
_plt.title('original')
_plt.colorbar()
# 2.3 perform land extrapolation on reduced variable
datanorm = self.normalize(datasrc, datamin, datamax, mask)
if self.debug:
print(datanorm.min(), datanorm.max(), datamin, datamax)
datanormextrap = self.drown_field(datanorm, mask, drown)
dataextrap = self.unnormalize(datanormextrap, datamin, datamax)
return dataextrap
def interpolate_from(self,
filename,
variable_u,
variable_v,
frame=None,
drown='sosie',
maskfile=None,
maskvar=None,
missing_value=None,
from_global=True,
depthname='z',
timename='time',
coord_names_u=['lon', 'lat'],
coord_names_v=['lon', 'lat'],
x_coords_u=None,
y_coords_u=None,
x_coords_v=None,
y_coords_v=None,
method='bilinear',
interpolator_u=None,
interpolator_v=None,
autocrop=True):
''' interpolate_from performs a serie of operation :
* read input data
* perform extrapolation over land if desired
* read or create mask if extrapolation
* call ESMF for regridding
Optional arguments (=default) :
* frame=None : time record from input data (e.g. 1,2,..,12) when input
file contains more than one record.
* drown=True : perform extrapolation of ocean values onto land
* maskfile=None : to read mask from a file (else uses missing
value of variable)
* maskvar=None : if maskfile is defined, we need to provide name of
mask array in maskfile
* missing_value=None : when missing value attribute not defined
in input file, this allows to pass it
* use_locstream=False : interpolate from ESMF grid to ESMF locstream
instead of ESMF grid, a bit faster.
use only to interpolate to boundary.
* from_global=True : if input file is global leave to true. If input
is regional, set to False.
interpolating from a regional extraction can
significantly speed up processing.
'''
# 1. Create ESMF source grids
if maskfile is not None:
self.gridsrc_u, imin_src_u, imax_src_u, jmin_src_u, jmax_src_u = \
self.create_source_grid(maskfile, from_global, coord_names_u,
x_coords=x_coords_u,
y_coords=y_coords_u,
autocrop=autocrop)
self.gridsrc_v, imin_src_v, imax_src_v, jmin_src_v, jmax_src_v = \
self.create_source_grid(maskfile, from_global, coord_names_v,
x_coords=x_coords_v,
y_coords=y_coords_v,
autocrop=autocrop)
else:
self.gridsrc_u, imin_src_u, imax_src_u, jmin_src_u, jmax_src_u = \
self.create_source_grid(filename, from_global, coord_names_u,
x_coords=x_coords_u,
y_coords=y_coords_u,
autocrop=autocrop)
self.gridsrc_v, imin_src_v, imax_src_v, jmin_src_v, jmax_src_v = \
self.create_source_grid(filename, from_global, coord_names_v,
x_coords=x_coords_v,
y_coords=y_coords_v,
autocrop=autocrop)
# 2. read the original field
datasrc_u = _ncdf.read_field(filename, variable_u, frame=frame)
datasrc_v = _ncdf.read_field(filename, variable_v, frame=frame)
if self.geometry == 'surface':
datasrc_u = datasrc_u[:, jmin_src_u:jmax_src_u,
imin_src_u:imax_src_u]
datasrc_v = datasrc_v[:, jmin_src_v:jmax_src_v,
imin_src_v:imax_src_v]
self.depth, self.nz, self.dz = _ncdf.read_vert_coord(
filename, depthname, self.nx, self.ny)
else:
datasrc_u = datasrc_u[jmin_src_u:jmax_src_u, imin_src_u:imax_src_u]
datasrc_v = datasrc_v[jmin_src_v:jmax_src_v, imin_src_v:imax_src_v]
self.depth = 0.
self.nz = 1
self.dz = 0.
# read time
try:
self.timesrc = _ncdf.read_time(filename, timename, frame=frame)
except:
print('input data time variable not read')
# TODO !! make rotation to east,north from source grid.
# important : if the grid is regular, we don't need to colocate u,v and
# the interpolation will be more accurate.
# Run colocation only if grid is non-regular.
# angle_src_u = _lc.compute_angle_from_lon_lat(self.gridsrc_u.coords[0][0].T,\
# self.gridsrc_u.coords[0][1].T)
# angle_src_v = _lc.compute_angle_from_lon_lat(self.gridsrc_v.coords[0][0].T,\
# self.gridsrc_v.coords[0][1].T)
# 3. perform extrapolation over land
print('drown')
start = ptime.time()
if drown in self.drown_methods:
dataextrap_u = self.perform_extrapolation(datasrc_u, maskfile,
maskvar, missing_value,
drown)
dataextrap_v = self.perform_extrapolation(datasrc_v, maskfile,
maskvar, missing_value,
drown)
else:
dataextrap_u = datasrc_u.copy()
dataextrap_v = datasrc_v.copy()
end = ptime.time()
print('end drown', end - start)
# 4. ESMF interpolation
# Create a field on the centers of the grid
field_src_u = _ESMF.Field(self.gridsrc_u,
staggerloc=_ESMF.StaggerLoc.CENTER)
field_src_v = _ESMF.Field(self.gridsrc_v,
staggerloc=_ESMF.StaggerLoc.CENTER)
# Set up a regridding object between source and destination
if interpolator_u is None:
print('create regridding for u')
if method == 'bilinear':
regridme_u = _ESMF.Regrid(
field_src_u,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.BILINEAR)
elif method == 'patch':
regridme_u = _ESMF.Regrid(
field_src_u,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.PATCH)
else:
regridme_u = interpolator_u
if interpolator_v is None:
print('create regridding for v')
if method == 'bilinear':
regridme_v = _ESMF.Regrid(
field_src_v,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.BILINEAR)
elif method == 'patch':
regridme_v = _ESMF.Regrid(
field_src_v,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.PATCH)
else:
regridme_v = interpolator_v
print('regridding u')
self.data_u = self.perform_interpolation(dataextrap_u, regridme_u,
field_src_u,
self.field_target,
self.use_locstream)
print('regridding v')
self.data_v = self.perform_interpolation(dataextrap_v, regridme_v,
field_src_v,
self.field_target,
self.use_locstream)
# vector rotation to output grid
self.data_u_out = self.data_u * _np.cos(self.angle_dx[self.jmin:self.jmax+1,self.imin:self.imax+1]) + \
self.data_v * _np.sin(self.angle_dx[self.jmin:self.jmax+1,self.imin:self.imax+1])
self.data_v_out = self.data_v * _np.cos(self.angle_dx[self.jmin:self.jmax+1,self.imin:self.imax+1]) - \
self.data_u * _np.sin(self.angle_dx[self.jmin:self.jmax+1,self.imin:self.imax+1])
# free memory (ESMPy has memory leak)
self.gridsrc_u.destroy()
self.gridsrc_v.destroy()
field_src_u.destroy()
field_src_v.destroy()
return regridme_u, regridme_v
def interpolate_from(self,
filename,
variable,
frame=None,
drown='sosie',
maskfile=None,
maskvar=None,
missing_value=None,
from_global=True,
depthname='z',
timename='time',
coord_names=['lon', 'lat'],
x_coords=None,
y_coords=None,
method='bilinear',
interpolator=None,
autocrop=True,
use_gridspec=False):
''' interpolate_from performs a serie of operation :
* read input data
* perform extrapolation over land if desired
* read or create mask if extrapolation
* call ESMF for regridding
Optional arguments (=default) :
* frame=None : time record from input data (e.g. 1,2,..,12) when input
file contains more than one record.
* drown=True : perform extrapolation of ocean values onto land
* maskfile=None : to read mask from a file (else uses missing value
of variable)
* maskvar=None : if maskfile is defined, we need to provide name of
mask array in maskfile
* missing_value=None : when missing value attribute not defined in
input file, this allows to pass it
* use_locstream=False : interpolate from ESMF grid to ESMF locstream
instead of ESMF grid, a bit faster.
use only to interpolate to boundary.
* from_global=True : if input file is global leave to true. If input
is regional, set to False.
interpolating from a regional extraction can
significantly speed up processing.
'''
# 1. Create ESMF source grid
if maskfile is not None:
self.gridsrc = self.create_source_grid(maskfile,
from_global,
coord_names,
x_coords=x_coords,
y_coords=y_coords,
autocrop=autocrop,
use_gridspec=use_gridspec)
else:
self.gridsrc = self.create_source_grid(filename,
from_global,
coord_names,
x_coords=x_coords,
y_coords=y_coords,
autocrop=autocrop,
use_gridspec=use_gridspec)
# 2. read the original field
datasrc = _ncdf.read_field(filename, variable, frame=frame)
if self.geometry == 'surface':
datasrc = datasrc[:, self.jmin_src:self.jmax_src,
self.imin_src:self.imax_src]
self.depth, self.nz, self.dz = _ncdf.read_vert_coord(
filename, depthname, self.nx, self.ny)
else:
datasrc = datasrc[self.jmin_src:self.jmax_src,
self.imin_src:self.imax_src]
self.depth = 0.
self.nz = 1
self.dz = 0.
# read time
try:
self.timesrc = _ncdf.read_time(filename, timename, frame=frame)
except:
print('input data time variable not read')
# 3. perform extrapolation over land
print('drown')
start = ptime.time()
if drown in self.drown_methods:
dataextrap = self.perform_extrapolation(datasrc, maskfile, maskvar,
missing_value, drown)
else:
dataextrap = datasrc.copy()
end = ptime.time()
print('end drown', end - start)
# 4. ESMF interpolation
# Create a field on the centers of the grid
field_src = _ESMF.Field(self.gridsrc,
staggerloc=_ESMF.StaggerLoc.CENTER)
# Set up a regridding object between source and destination
if interpolator is None:
if method == 'bilinear':
regridme = _ESMF.Regrid(
field_src,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.BILINEAR)
elif method == 'patch':
regridme = _ESMF.Regrid(
field_src,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.PATCH)
elif method == 'conserve':
regridme = _ESMF.Regrid(
field_src,
self.field_target,
unmapped_action=_ESMF.UnmappedAction.IGNORE,
regrid_method=_ESMF.RegridMethod.CONSERVE)
else:
regridme = interpolator
self.data = self.perform_interpolation(dataextrap, regridme, field_src,
self.field_target,
self.use_locstream)
# free memory (ESMPy has memory leak)
self.gridsrc.destroy()
field_src.destroy()
return regridme
def __init__(self, segment_name, target_grid_file, target_model='MOM6',
**kwargs):
''' read target grid and create associated ESMF grid and
locstream objects.
Args:
segment_name (string): name of the segment
target_grid_file (netcdf file): grid file of the ocean model
target_model: name of the ocean model (default MOM6)
Examples:
south = obc_segment('segment_001','./ocean_hgrid.nc',
imin=0,imax=360,jmin=0,jmax=0)
north = obc_segment('segment_002','./ocean_hgrid.nc',
imin=0,imax=360,jmin=960,jmax=960)
south = obc_segment('south','./roms_grd.nc',target_model='ROMS',
imin=0,imax=180,jmin=0,jmax=0)
north = obc_segment('north','./roms_grd.nc',target_model='ROMS',
imin=0,imax=180,jmin=480,jmax=480)
'''
# * istart : along x axis, where the segment begins
#
# * iend : along x axis, where the segment ends
#
# * jstart : along y axis, where the segment begins
#
# * jend : along y axis, where the segment end
#
# '''
# read args
self.segment_name = segment_name
self.target_grid_file = target_grid_file
self.items = []
self.items.append('segment_name')
self.items.append('target_grid')
self.debug = False
# iterate over all kwargs and store them as attributes for the object
if kwargs is not None:
self.__dict__.update(kwargs)
for key, value in kwargs.items():
self.items.append(key)
if self.istart > self.iend:
self.imin = self.iend
self.imax = self.istart
self.orientation = 2 # E-W
elif self.iend > self.istart:
self.imin = self.istart
self.imax = self.iend
self.orientation = 0 # W-E
else:
self.imin = self.istart
self.imax = self.istart
if self.jstart > self.jend:
self.jmin = self.jend
self.jmax = self.jstart
self.orientation = 3 # N-S
elif self.jend > self.jstart:
self.jmin = self.jstart
self.jmax = self.jend
self.orientation = 1 # S-N
else:
self.jmin = self.jstart
self.jmax = self.jstart
# compute dimensions
self.nx = self.imax - self.imin + 1
self.ny = self.jmax - self.jmin + 1
self.ilist = _np.empty((self.ny, self.nx))
self.jlist = _np.empty((self.ny, self.nx))
for kx in _np.arange(self.nx):
self.jlist[:, kx] = _np.arange(self.jmin, self.jmax+1) / 2.
for ky in _np.arange(self.ny):
self.ilist[ky, :] = _np.arange(self.imin, self.imax+1) / 2.
# coordinate names depend on ocean model
# MOM6 has all T,U,V points in one big grid
# ROMS has in 3 separate ones.
if target_model == 'MOM6':
coord_names = ["x", "y"]
self.angle_dx = _ncdf.read_field(target_grid_file, 'angle_dx')
lon_target = _ncdf.read_field(target_grid_file, 'x')
lat_target = _ncdf.read_field(target_grid_file, 'y')
ny_target, nx_target = lat_target.shape
self.grid_target = _ESMF.Grid(_np.array([nx_target, ny_target]))
self.grid_target.add_coords(staggerloc=[_ESMF.StaggerLoc.CENTER])
self.grid_target.coords[_ESMF.StaggerLoc.CENTER]
self.grid_target.coords[_ESMF.StaggerLoc.CENTER][0][:] = \
lon_target.T
self.grid_target.coords[_ESMF.StaggerLoc.CENTER][1][:] = \
lat_target.T
elif target_model == 'ROMS':
coord_names = ["lon_rho", "lat_rho"]
self.angle_dx = _ncdf.read_field(target_grid_file, 'angle')
# import target grid into ESMF grid object
self.grid_target = _ESMF.Grid(filename=target_grid_file,
filetype=_ESMF.FileFormat.GRIDSPEC,
coord_names=coord_names)
elif target_model == 'regular':
coord_names = ["lon", "lat"]
self.angle_dx = 0. # for now
# import target grid into ESMF grid object
self.grid_target = _ESMF.Grid(filename=target_grid_file,
filetype=_ESMF.FileFormat.GRIDSPEC,
coord_names=coord_names,
add_corner_stagger=True)
# import same target grid into ESMF locstream object
self.locstream_target = \
_ESMF.LocStream(self.nx * self.ny,
coord_sys=_ESMF.CoordSys.SPH_DEG)
self.locstream_target["ESMF:Lon"] = \
self.grid_target.coords[0][0][self.imin:self.imax+1,
self.jmin:self.jmax+1].flatten()
self.locstream_target["ESMF:Lat"] = \
self.grid_target.coords[0][1][self.imin:self.imax+1,
self.jmin:self.jmax+1].flatten()
# save lon/lat on this segment
self.lon = self.grid_target.coords[0][0][self.imin:self.imax+1,
self.jmin:self.jmax+1]
self.lon = self.lon.transpose().squeeze()
self.lat = self.grid_target.coords[0][1][self.imin:self.imax+1,
self.jmin:self.jmax+1]
self.lat = self.lat.transpose().squeeze()
# nc dimensions for horizontal coords
self.hdimensions_name = ('ny_' + self.segment_name,
'nx_' + self.segment_name,)
return None