# This software is open source software available under the BSD-3 license.
#
# Copyright (c) 2019 Triad National Security, LLC. All rights reserved.
# Copyright (c) 2019 Lawrence Livermore National Security, LLC. All rights
# reserved.
# Copyright (c) 2019 UT-Battelle, LLC. All rights reserved.
#
# Additional copyright and license information can be found in the LICENSE file
# distributed with this code, or at
# https://raw.githubusercontent.com/MPAS-Dev/MPAS-Analysis/master/LICENSE
'''
Classes for describing meshes and grids , including creating SCRIP files,
used to create mapping files
Classes
-------
MpasMeshDescriptor - describes an MPAS mesh
LatLonGridDescriptor - describes a lat-lon grid
ProjectionGridDescriptor - describes a logically rectangular grid on a pyproj
projection
PointCollectionDescriptor - describes a collection of unconnected points (only
valid as a destination grid for remapping)
'''
# Authors
# -------
# Xylar Asay-Davis
import netCDF4
import numpy
import sys
import pyproj
import xarray
[docs]def get_lat_lon_descriptor(dLon, dLat, lonMin=-180., lonMax=180., latMin=-90.,
latMax=90.):
"""
Get a descriptor of a lat/lon grid, used for remapping
Parameters
----------
dLon, dLat : float`
Resolution of the lon-lat grid in degrees
Returns
-------
descriptor : ``LatLonGridDescriptor`` object
A descriptor of the lat/lon grid
"""
# Authors
# -------
# Xylar Asay-Davis
nLat = int((latMax - latMin) / dLat) + 1
nLon = int((lonMax - lonMin) / dLon) + 1
lat = numpy.linspace(latMin, latMax, nLat)
lon = numpy.linspace(lonMin, lonMax, nLon)
descriptor = LatLonGridDescriptor.create(lat, lon, units='degrees')
return descriptor
class MeshDescriptor(object): # {{{
'''
A class for describing a mesh
'''
# Authors
# -------
# Xylar Asay-Davis
def __init__(self): # {{{
'''
Constructor creates a common ``meshName`` member variable, ``None`` by
default. Each Subclass should define or use input arguments to set
``meshName`` to a short description of the mesh or grid.
'''
# Authors
# -------
# Xylar Asay-Davis
self.meshName = None # }}}
def to_scrip(self, scripFileName): # {{{
'''
Subclasses should overload this method to write a SCRIP file based on
the mesh.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
'''
# Authors
# ------
# Xylar Asay-Davis
return # }}}
def to_esmf(self, esmfFileName): # {{{
"""
Subclasses should overload this method to write an ESMF mesh file based
on the mesh.
Parameters
----------
esmfFileName : str
The path to which the SCRIP file should be written
"""
# Authors
# ------
# Xylar Asay-Davis
return # }}}
# }}}
[docs]class MpasMeshDescriptor(MeshDescriptor): # {{{
"""
A class for describing an MPAS mesh
"""
# Authors
# -------
# Xylar Asay-Davis
[docs] def __init__(self, fileName, meshName=None, vertices=False): # {{{
"""
Constructor stores the file name
Parameters
----------
fileName : str
The path of the file containing the MPAS mesh
meshName : str, optional
The name of the MPAS mesh (e.g. ``'oEC60to30'`` or
``'oRRS18to6'``). If not provided, the data set in ``fileName``
must have a global attribute ``meshName`` that will be used
instead.
vertices : bool, optional
Whether the mapping is to or from vertices instead of corners
(for non-conservative remapping)
"""
# Authors
# -------
# Xylar Asay-Davis
super().__init__()
self.vertices = vertices
with xarray.open_dataset(fileName) as ds:
if meshName is None:
if 'meshName' not in ds.attrs:
raise ValueError('No meshName provided or found in file.')
self.meshName = ds.attrs['meshName']
else:
self.meshName = meshName
self.fileName = fileName
self.regional = True
if vertices:
# build coords
self.coords = {'latVertex': {'dims': 'nVertices',
'data': ds.latVertex.values,
'attrs': {'units': 'radians'}},
'lonVertex': {'dims': 'nVertices',
'data': ds.lonVertex.values,
'attrs': {'units': 'radians'}}}
self.dims = ['nVertices']
else:
# build coords
self.coords = {'latCell': {'dims': 'nCells',
'data': ds.latCell.values,
'attrs': {'units': 'radians'}},
'lonCell': {'dims': 'nCells',
'data': ds.lonCell.values,
'attrs': {'units': 'radians'}}}
self.dims = ['nCells']
self.dimSize = [ds.dims[dim] for dim in self.dims]
# }}}
def to_scrip(self, scripFileName): # {{{
"""
Given an MPAS mesh file, create a SCRIP file based on the mesh.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
"""
# Authors
# -------
# Xylar Asay-Davis
if self.vertices:
raise ValueError('A SCRIP file won\'t work for remapping vertices')
inFile = netCDF4.Dataset(self.fileName, 'r')
outFile = netCDF4.Dataset(scripFileName, 'w')
# Get info from input file
latCell = inFile.variables['latCell'][:]
lonCell = inFile.variables['lonCell'][:]
latVertex = inFile.variables['latVertex'][:]
lonVertex = inFile.variables['lonVertex'][:]
verticesOnCell = inFile.variables['verticesOnCell'][:]
nEdgesOnCell = inFile.variables['nEdgesOnCell'][:]
nCells = len(inFile.dimensions['nCells'])
maxVertices = len(inFile.dimensions['maxEdges'])
areaCell = inFile.variables['areaCell'][:]
sphereRadius = float(inFile.sphere_radius)
_create_scrip(outFile, grid_size=nCells, grid_corners=maxVertices,
grid_rank=1, units='radians', meshName=self.meshName)
grid_area = outFile.createVariable('grid_area', 'f8', ('grid_size',))
grid_area.units = 'radian^2'
# SCRIP uses square radians
grid_area[:] = areaCell[:] / (sphereRadius**2)
outFile.variables['grid_center_lat'][:] = latCell[:]
outFile.variables['grid_center_lon'][:] = lonCell[:]
outFile.variables['grid_dims'][:] = nCells
outFile.variables['grid_imask'][:] = 1
# grid corners:
grid_corner_lon = numpy.zeros((nCells, maxVertices))
grid_corner_lat = numpy.zeros((nCells, maxVertices))
for iVertex in range(maxVertices):
cellIndices = numpy.arange(nCells)
# repeat the last vertex wherever iVertex > nEdgesOnCell
localVertexIndices = numpy.minimum(nEdgesOnCell - 1, iVertex)
vertexIndices = verticesOnCell[cellIndices, localVertexIndices] - 1
grid_corner_lat[cellIndices, iVertex] = latVertex[vertexIndices]
grid_corner_lon[cellIndices, iVertex] = lonVertex[vertexIndices]
outFile.variables['grid_corner_lat'][:] = grid_corner_lat[:]
outFile.variables['grid_corner_lon'][:] = grid_corner_lon[:]
# Update history attribute of netCDF file
if hasattr(inFile, 'history'):
newhist = '\n'.join([getattr(inFile, 'history'),
' '.join(sys.argv[:])])
else:
newhist = sys.argv[:]
setattr(outFile, 'history', newhist)
inFile.close()
outFile.close() # }}}
def to_esmf(self, esmfFileName): # {{{
"""
Create an ESMF mesh file for the mesh
Parameters
----------
esmfFileName : str
The path to which the ESMF mesh file should be written
"""
# Authors
# -------
# Xylar Asay-Davis
with xarray.open_dataset(self.fileName) as ds:
nodeCount = ds.sizes['nVertices']
elementCount = ds.sizes['nCells']
coordDim = 2
nodeCoords = numpy.zeros((nodeCount, coordDim))
nodeCoords[:, 0] = ds.lonVertex.values
nodeCoords[:, 1] = ds.latVertex.values
centerCoords = numpy.zeros((elementCount, coordDim))
centerCoords[:, 0] = ds.lonCell.values
centerCoords[:, 1] = ds.latCell.values
elementConn = ds.verticesOnCell.values
elementConn[elementConn == nodeCount + 1] = -1
ds_out = xarray.Dataset()
ds_out['nodeCoords'] = (('nodeCount', 'coordDim'), nodeCoords)
ds_out.nodeCoords.attrs['units'] = 'radians'
ds_out['centerCoords'] = \
(('elementCount', 'coordDim'), centerCoords)
ds_out.centerCoords.attrs['units'] = 'radians'
ds_out['elementConn'] = \
(('elementCount', 'maxNodePElement'), elementConn)
ds_out.elementConn.attrs['start_index'] = 1
ds_out.elementConn.attrs['_FillValue'] = -1
ds_out['numElementConn'] = \
(('elementCount',), ds.nEdgesOnCell.values)
ds_out.attrs['gridType'] = 'unstructured mesh'
ds_out.attrs['version'] = '0.9'
ds_out.to_netcdf(esmfFileName) # }}}
# }}}
[docs]class LatLonGridDescriptor(MeshDescriptor): # {{{
'''
A class for describing a lat-lon grid
'''
# Authors
# -------
# Xylar Asay-Davis
[docs] def __init__(self): # {{{
'''
Constructor stores the file name
Parameters
----------
fileName : str
The path of the file containing the MPAS mesh
'''
# Authors
# -------
# Xylar Asay-Davis
self.regional = False
self.meshName = None # }}}
@classmethod
def read(cls, fileName=None, ds=None, latVarName='lat',
lonVarName='lon'): # {{{
'''
Read the lat-lon grid from a file with the given lat/lon var names.
Parameters
----------
fileName : str, optional
The path of the file containing the lat-lon grid (if ``ds`` is not
supplied directly)
ds : ``xarray.Dataset`` object, optional
The path of the file containing the lat-lon grid (if supplied,
``fileName`` will be ignored)
latVarName, lonVarName : str, optional
The name of the latitude and longitude variables in the grid file
'''
# Authors
# -------
# Xylar Asay-Davis
if ds is None:
ds = xarray.open_dataset(fileName)
descriptor = cls()
if descriptor.meshName is None and 'meshName' in ds.attrs:
descriptor.meshName = ds.attrs['meshName']
# Get info from input file
descriptor.lat = numpy.array(ds[latVarName].values, float)
descriptor.lon = numpy.array(ds[lonVarName].values, float)
if 'degree' in ds[latVarName].units:
descriptor.units = 'degrees'
else:
descriptor.units = 'radians'
# interp/extrap corners
descriptor.lonCorner = interp_extrap_corner(descriptor.lon)
descriptor.latCorner = interp_extrap_corner(descriptor.lat)
descriptor._set_coords(latVarName, lonVarName, ds[latVarName].dims[0],
ds[lonVarName].dims[0])
if 'history' in ds.attrs:
descriptor.history = '\n'.join([ds.attrs['history'],
' '.join(sys.argv[:])])
else:
descriptor.history = sys.argv[:]
return descriptor # }}}
@classmethod
def create(cls, latCorner, lonCorner, units='degrees'): # {{{
'''
Create the lat-lon grid with the given arrays and units.
Parameters
----------
latCorner, lonCorner : 1D numpy.arrays
One dimensional arrays defining the latitude and longitude
coordinates of grid corners.
units : {'degrees', 'radians'}, optional
The units of `latCorner` and `lonCorner`
'''
# Authors
# -------
# Xylar Asay-Davis
descriptor = cls()
descriptor.latCorner = latCorner
descriptor.lonCorner = lonCorner
descriptor.lon = 0.5 * (lonCorner[0:-1] + lonCorner[1:])
descriptor.lat = 0.5 * (latCorner[0:-1] + latCorner[1:])
descriptor.units = units
descriptor.history = sys.argv[:]
descriptor._set_coords('lat', 'lon', 'lat', 'lon')
return descriptor # }}}
def to_scrip(self, scripFileName): # {{{
'''
Given a lat-lon grid file, create a SCRIP file based on the grid.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
'''
# Authors
# -------
# Xylar Asay-Davis
outFile = netCDF4.Dataset(scripFileName, 'w')
nLat = len(self.lat)
nLon = len(self.lon)
grid_size = nLat * nLon
_create_scrip(outFile, grid_size=grid_size, grid_corners=4,
grid_rank=2, units=self.units, meshName=self.meshName)
(Lon, Lat) = numpy.meshgrid(self.lon, self.lat)
(LonCorner, LatCorner) = numpy.meshgrid(self.lonCorner, self.latCorner)
outFile.variables['grid_center_lat'][:] = Lat.flat
outFile.variables['grid_center_lon'][:] = Lon.flat
outFile.variables['grid_dims'][:] = [nLon, nLat]
outFile.variables['grid_imask'][:] = 1
outFile.variables['grid_corner_lat'][:] = _unwrap_corners(LatCorner)
outFile.variables['grid_corner_lon'][:] = _unwrap_corners(LonCorner)
setattr(outFile, 'history', self.history)
outFile.close() # }}}
def to_esmf(self, esmfFileName): # {{{
"""
Create an ESMF mesh file for the mesh
Parameters
----------
esmfFileName : str
The path to which the ESMF mesh file should be written
"""
# Authors
# -------
# Xylar Asay-Davis
nLat = len(self.lat)
nLon = len(self.lon)
(Lon, Lat) = numpy.meshgrid(self.lon, self.lat)
(LonCorner, LatCorner) = numpy.meshgrid(self.lonCorner, self.latCorner)
(XIndices, YIndices) = numpy.meshgrid(numpy.arange(nLon + 1),
numpy.arange(nLat + 1))
elementCount = nLon * nLat
nodeCount = (nLon + 1) * (nLat + 1)
coordDim = 2
maxNodePElement = 4
nodeCoords = numpy.zeros((nodeCount, coordDim))
nodeCoords[:, 0] = LonCorner.flat
nodeCoords[:, 1] = LatCorner.flat
centerCoords = numpy.zeros((elementCount, coordDim))
centerCoords[:, 0] = Lon.flat
centerCoords[:, 1] = Lat.flat
elementConn = numpy.zeros((elementCount, maxNodePElement), dtype=int)
elementConn[:, 0] = (XIndices[0:-1, 0:-1].ravel() +
(nLon + 1) * YIndices[0:-1, 0:-1].ravel())
elementConn[:, 1] = (XIndices[0:-1, 1:].ravel() +
(nLon + 1) * YIndices[0:-1, 1:].ravel())
elementConn[:, 2] = (XIndices[1:, 1:].ravel() +
(nLon + 1) * YIndices[1:, 1:].ravel())
elementConn[:, 3] = (XIndices[1:, 0:-1].ravel() +
(nLon + 1) * YIndices[1:, 0:-1].ravel())
ds_out = xarray.Dataset()
ds_out['nodeCoords'] = (('nodeCount', 'coordDim'), nodeCoords)
ds_out.nodeCoords.attrs['units'] = self.units
ds_out['centerCoords'] = \
(('elementCount', 'coordDim'), centerCoords)
ds_out.centerCoords.attrs['units'] = self.units
ds_out['elementConn'] = \
(('elementCount', 'maxNodePElement'), elementConn + 1)
ds_out.elementConn.attrs['start_index'] = 1
ds_out.elementConn.attrs['_FillValue'] = -1
ds_out['numElementConn'] = \
(('elementCount',), maxNodePElement * numpy.ones(elementCount,
dtype=int))
ds_out['origGridDims'] = (('origGridRank',), [nLon, nLat])
ds_out.attrs['gridType'] = 'unstructured mesh'
ds_out.attrs['version'] = '0.9'
ds_out.to_netcdf(esmfFileName) # }}}
def _set_coords(self, latVarName, lonVarName, latDimName,
lonDimName): # {{{
'''
Set up a coords dict with lat and lon
'''
self.latVarName = latVarName
self.lonVarName = lonVarName
self.coords = {latVarName: {'dims': latDimName,
'data': self.lat,
'attrs': {'units': self.units}},
lonVarName: {'dims': lonDimName,
'data': self.lon,
'attrs': {'units': self.units}}}
self.dims = [latDimName, lonDimName]
self.dimSize = [len(self.lat), len(self.lon)]
# set the name of the grid
dLat = self.lat[1] - self.lat[0]
dLon = self.lon[1] - self.lon[0]
lonRange = self.lonCorner[-1] - self.lonCorner[0]
latRange = self.latCorner[-1] - self.latCorner[0]
if 'degree' in self.units:
units = 'degree'
if numpy.abs(lonRange - 360.) > 1e-10:
self.regional = True
if numpy.abs(latRange - 180.) > 1e-10:
self.regional = True
elif 'rad' in self.units:
if numpy.abs(lonRange - 2. * numpy.pi) > 1e-10:
self.regional = True
if numpy.abs(latRange - numpy.pi) > 1e-10:
self.regional = True
units = 'radian'
else:
raise ValueError('Could not figure out units {}'.format(
self.units))
if self.meshName is None:
self.meshName = '{}x{}{}'.format(_round_res(abs(dLat)),
_round_res(abs(dLon)), units)
# }}}
# }}}
[docs]class LatLon2DGridDescriptor(MeshDescriptor): # {{{
'''
A class for describing a lat-lon grid that may not be a tensor grid
(lat/lon are 2D arrays). The grid is assumed to be regional, since this
is difficult to determine just from the lat/lon values. The calling code
should set ``regional = False`` for global grids with 2D lat/lon
'''
# Authors
# -------
# Xylar Asay-Davis
[docs] def __init__(self): # {{{
'''
Constructor stores the file name
Parameters
----------
fileName : str
The path of the file containing the MPAS mesh
'''
# Authors
# -------
# Xylar Asay-Davis
self.regional = True
self.meshName = None # }}}
@classmethod
def read(cls, fileName=None, ds=None, latVarName='lat',
lonVarName='lon'): # {{{
'''
Read the lat-lon grid from a file with the given lat/lon var names.
Parameters
----------
fileName : str, optional
The path of the file containing the lat-lon grid (if ``ds`` is not
supplied directly)
ds : ``xarray.Dataset`` object, optional
The path of the file containing the lat-lon grid (if supplied,
``fileName`` will be ignored)
latVarName, lonVarName : str, optional
The name of the latitude and longitude variables in the grid file
'''
# Authors
# -------
# Xylar Asay-Davis
if ds is None:
ds = xarray.open_dataset(fileName)
descriptor = cls()
if descriptor.meshName is None and 'meshName' in ds.attrs:
descriptor.meshName = ds.attrs['meshName']
# Get info from input file
descriptor.lat = numpy.array(ds[latVarName].values, float)
descriptor.lon = numpy.array(ds[lonVarName].values, float)
if 'degree' in ds[latVarName].units:
descriptor.units = 'degrees'
else:
descriptor.units = 'radians'
# interp/extrap corners
descriptor.lonCorner = interp_extrap_corners_2d(descriptor.lon)
descriptor.latCorner = interp_extrap_corners_2d(descriptor.lat)
descriptor._set_coords(latVarName, lonVarName, ds[latVarName].dims[0],
ds[latVarName].dims[1])
if 'history' in ds.attrs:
descriptor.history = '\n'.join([ds.attrs['history'],
' '.join(sys.argv[:])])
else:
descriptor.history = sys.argv[:]
return descriptor # }}}
def to_scrip(self, scripFileName): # {{{
'''
Given a lat-lon grid file, create a SCRIP file based on the grid.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
'''
# Authors
# -------
# Xylar Asay-Davis
outFile = netCDF4.Dataset(scripFileName, 'w')
nLat, nLon = self.lat.shape
grid_size = nLat * nLon
_create_scrip(outFile, grid_size=grid_size, grid_corners=4,
grid_rank=2, units=self.units, meshName=self.meshName)
outFile.variables['grid_center_lat'][:] = self.lat.flat
outFile.variables['grid_center_lon'][:] = self.lon.flat
outFile.variables['grid_dims'][:] = [nLon, nLat]
outFile.variables['grid_imask'][:] = 1
outFile.variables['grid_corner_lat'][:] = \
_unwrap_corners(self.lonCorner)
outFile.variables['grid_corner_lon'][:] = \
_unwrap_corners(self.latCorner)
setattr(outFile, 'history', self.history)
outFile.close() # }}}
def _set_coords(self, latVarName, lonVarName, latDimName,
lonDimName): # {{{
'''
Set up a coords dict with lat and lon
'''
self.latVarName = latVarName
self.lonVarName = lonVarName
self.coords = {latVarName: {'dims': (latDimName, lonDimName),
'data': self.lat,
'attrs': {'units': self.units}},
lonVarName: {'dims': (latDimName, lonDimName),
'data': self.lon,
'attrs': {'units': self.units}}}
self.dims = [latDimName, lonDimName]
self.dimSize = self.lat.shape
# set the name of the grid
dLat = self.lat[1, 0] - self.lat[0, 0]
dLon = self.lon[0, 1] - self.lon[0, 0]
if 'degree' in self.units:
units = 'degree'
elif 'rad' in self.units:
units = 'radian'
else:
raise ValueError('Could not figure out units {}'.format(
self.units))
if self.meshName is None:
self.meshName = '{}x{}{}'.format(_round_res(abs(dLat)),
_round_res(abs(dLon)), units)
# }}}
# }}}
[docs]class ProjectionGridDescriptor(MeshDescriptor): # {{{
'''
A class for describing a general logically rectangular grid that can be
defined by a `pyproj` projection.
'''
# Authors
# -------
# Xylar Asay-Davis
[docs] def __init__(self, projection): # {{{
'''
Constructor stores the projection
Parameters
----------
projection : ``pyproj.Proj`` object
The projection used to map from grid x-y space to latitude and
longitude
'''
# Authors
# -------
# Xylar Asay-Davis
super().__init__()
self.projection = projection
self.latLonProjection = pyproj.Proj(proj='latlong', datum='WGS84')
self.regional = True
self.x = None
self.y = None
self.xCorner = None
self.yCorner = None
self.history = None
@classmethod
def read(cls, projection, fileName, meshName=None, xVarName='x',
yVarName='y'): # {{{
'''
Given a grid file with x and y coordinates defining the axes of the
logically rectangular grid, read in the x and y coordinates and
interpolate/extrapolate to locate corners.
Parameters
----------
projection : pyproj.Proj object
The projection used to map from grid x-y space to latitude and
longitude
fileName : str
The path of the file containing the grid data
meshName : str, optional
The name of the grid (e.g. ``'10km_Antarctic_stereo'``). If not
provided, the data set in ``fileName`` must have a global
attribute ``meshName`` that will be used instead.
xVarName, yVarName : str, optional
The name of the x and y (in meters) variables in the grid file
'''
# Authors
# -------
# Xylar Asay-Davis
descriptor = cls(projection)
ds = xarray.open_dataset(fileName)
if meshName is None:
if 'meshName' not in ds.attrs:
raise ValueError('No meshName provided or found in file.')
descriptor.meshName = ds.attrs['meshName']
else:
descriptor.meshName = meshName
# Get info from input file
descriptor.x = numpy.array(ds[xVarName].values, float)
descriptor.y = numpy.array(ds[yVarName].values, float)
descriptor._set_coords(xVarName, yVarName, ds[xVarName].dims[0],
ds[yVarName].dims[0])
# interp/extrap corners
descriptor.xCorner = interp_extrap_corner(descriptor.x)
descriptor.yCorner = interp_extrap_corner(descriptor.y)
# Update history attribute of netCDF file
if 'history' in ds.attrs:
descriptor.history = '\n'.join([ds.attrs['history'],
' '.join(sys.argv[:])])
else:
descriptor.history = sys.argv[:]
return descriptor # }}}
@classmethod
def create(cls, projection, x, y, meshName): # {{{
'''
Given x and y coordinates defining the axes of the logically
rectangular grid, save the coordinates interpolate/extrapolate to
locate corners.
Parameters
----------
x, y : 1D numpy.arrays
One dimensional arrays defining the x and y coordinates of grid
cell centers.
meshName : str
The name of the grid (e.g. ``'10km_Antarctic_stereo'``)
'''
# Authors
# -------
# Xylar Asay-Davis
descriptor = cls(projection)
descriptor.meshName = meshName
descriptor.x = x
descriptor.y = y
descriptor._set_coords('x', 'y', 'x', 'y')
# interp/extrap corners
descriptor.xCorner = interp_extrap_corner(descriptor.x)
descriptor.yCorner = interp_extrap_corner(descriptor.y)
descriptor.history = sys.argv[:]
return descriptor # }}}
def to_scrip(self, scripFileName): # {{{
'''
Create a SCRIP file based on the grid and projection.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
'''
# Authors
# -------
# Xylar Asay-Davis
outFile = netCDF4.Dataset(scripFileName, 'w')
nx = len(self.x)
ny = len(self.y)
grid_size = nx * ny
_create_scrip(outFile, grid_size=grid_size, grid_corners=4,
grid_rank=2, units='degrees', meshName=self.meshName)
(X, Y) = numpy.meshgrid(self.x, self.y)
(XCorner, YCorner) = numpy.meshgrid(self.xCorner, self.yCorner)
(Lat, Lon) = self.project_to_lat_lon(X, Y)
(LatCorner, LonCorner) = self.project_to_lat_lon(XCorner, YCorner)
outFile.variables['grid_center_lat'][:] = Lat.flat
outFile.variables['grid_center_lon'][:] = Lon.flat
outFile.variables['grid_dims'][:] = [nx, ny]
outFile.variables['grid_imask'][:] = 1
outFile.variables['grid_corner_lat'][:] = _unwrap_corners(LatCorner)
outFile.variables['grid_corner_lon'][:] = _unwrap_corners(LonCorner)
setattr(outFile, 'history', self.history)
outFile.close() # }}}
def to_esmf(self, esmfFileName): # {{{
"""
Create an ESMF mesh file for the mesh
Parameters
----------
esmfFileName : str
The path to which the ESMF mesh file should be written
"""
# Authors
# -------
# Xylar Asay-Davis
nx = len(self.x)
ny = len(self.y)
(X, Y) = numpy.meshgrid(self.x, self.y)
(XCorner, YCorner) = numpy.meshgrid(self.xCorner, self.yCorner)
(XIndices, YIndices) = numpy.meshgrid(numpy.arange(nx + 1),
numpy.arange(ny + 1))
(Lat, Lon) = self.project_to_lat_lon(X, Y)
(LatCorner, LonCorner) = self.project_to_lat_lon(XCorner, YCorner)
elementCount = nx * ny
nodeCount = (nx + 1) * (ny + 1)
coordDim = 2
maxNodePElement = 4
nodeCoords = numpy.zeros((nodeCount, coordDim))
nodeCoords[:, 0] = LonCorner.flat
nodeCoords[:, 1] = LatCorner.flat
centerCoords = numpy.zeros((elementCount, coordDim))
centerCoords[:, 0] = Lon.flat
centerCoords[:, 1] = Lat.flat
elementConn = numpy.zeros((elementCount, maxNodePElement), dtype=int)
elementConn[:, 0] = (XIndices[0:-1, 0:-1].ravel() +
(nx + 1) * YIndices[0:-1, 0:-1].ravel())
elementConn[:, 1] = (XIndices[0:-1, 1:].ravel() +
(nx + 1) * YIndices[0:-1, 1:].ravel())
elementConn[:, 2] = (XIndices[1:, 1:].ravel() +
(nx + 1) * YIndices[1:, 1:].ravel())
elementConn[:, 3] = (XIndices[1:, 0:-1].ravel() +
(nx + 1) * YIndices[1:, 0:-1].ravel())
ds_out = xarray.Dataset()
ds_out['nodeCoords'] = (('nodeCount', 'coordDim'), nodeCoords)
ds_out.nodeCoords.attrs['units'] = 'degrees'
ds_out['centerCoords'] = \
(('elementCount', 'coordDim'), centerCoords)
ds_out.centerCoords.attrs['units'] = 'degrees'
ds_out['elementConn'] = \
(('elementCount', 'maxNodePElement'), elementConn + 1)
ds_out.elementConn.attrs['start_index'] = 1
ds_out.elementConn.attrs['_FillValue'] = -1
ds_out['numElementConn'] = \
(('elementCount',), maxNodePElement * numpy.ones(elementCount,
dtype=int))
ds_out['origGridDims'] = (('origGridRank',), [ny, nx])
ds_out.attrs['gridType'] = 'unstructured mesh'
ds_out.attrs['version'] = '0.9'
ds_out.to_netcdf(esmfFileName) # }}}
def project_to_lat_lon(self, X, Y): # {{{
'''
Given X and Y locations of points in a projection, returns the
corresponding latitude and longitude of each point.
Parameters
----------
outFile : file pointer
A SCRIP file opened in write mode
X, Y : 1D or 2D numpy.array
X and y arrays of points in in the projection
Returns
-------
Lat, Lon : numpy.array with same shape as X and Y
the latitude and longitude in degrees of the points
'''
# Authors
# -------
# Xylar Asay-Davis
transformer = pyproj.Transformer.from_proj(self.projection,
self.latLonProjection)
Lon, Lat = transformer.transform(X, Y)
return (Lat, Lon) # }}}
def _set_coords(self, xVarName, yVarName, xDimName, yDimName): # {{{
'''
Set up a coords dict with x, y, lat and lon
'''
self.xVarName = xVarName
self.yVarName = yVarName
(X, Y) = numpy.meshgrid(self.x, self.y)
(Lat, Lon) = self.project_to_lat_lon(X, Y)
self.coords = {xVarName: {'dims': xDimName,
'data': self.x,
'attrs': {'units': 'meters'}},
yVarName: {'dims': yDimName,
'data': self.y,
'attrs': {'units': 'meters'}},
'lat': {'dims': (yDimName, xDimName),
'data': Lat,
'attrs': {'units': 'degrees'}},
'lon': {'dims': (yDimName, xDimName),
'data': Lon,
'attrs': {'units': 'degrees'}}}
self.dims = [yDimName, xDimName]
self.dimSize = [len(self.y), len(self.x)]
# }}}
# }}}
[docs]class PointCollectionDescriptor(MeshDescriptor): # {{{
'''
A class for describing a collection of points
Author
------
Xylar Asay-Davis
Last Modified
-------------
04/19/2017
'''
[docs] def __init__(self, lats, lons, collectionName,
units='degrees', outDimension='nPoints'): # {{{
'''
Constructor stores
Parameters
----------
lats, lons : 1D numpy arrays
The latitude and longitude of each point
collectionName : str
A unique name for the collection of transects, used in the names
of files containing data mapped to these points.
units : {'degrees', 'radians'}, optional
The units of ``lats`` and ``lons``
outDimension : str, optional
The name of the dimension corresponding to the points (i.e. the
"horizontal" dimension of the point collection)
Author
------
Xylar Asay-Davis
Last Modified
-------------
04/19/2017
'''
self.meshName = collectionName
self.regional = True
self.lat = lats
self.lon = lons
self.units = units
# build coords
self.coords = {'lat': {'dims': outDimension,
'data': self.lat,
'attrs': {'units': units}},
'lon': {'dims': outDimension,
'data': self.lon,
'attrs': {'units': units}}}
self.dims = [outDimension]
self.dimSize = [len(self.lat)]
# }}}
def to_scrip(self, scripFileName): # {{{
'''
Given an MPAS mesh file, create a SCRIP file based on the mesh.
Parameters
----------
scripFileName : str
The path to which the SCRIP file should be written
'''
# Authors
# ------
# Xylar Asay-Davis
outFile = netCDF4.Dataset(scripFileName, 'w')
nPoints = len(self.lat)
_create_scrip(outFile, grid_size=nPoints,
grid_corners=4,
grid_rank=1, units=self.units, meshName=self.meshName)
grid_area = outFile.createVariable('grid_area', 'f8', ('grid_size',))
grid_area.units = 'radian^2'
# SCRIP uses square radians
grid_area[:] = numpy.zeros(nPoints)
outFile.variables['grid_center_lat'][:] = self.lat
outFile.variables['grid_center_lon'][:] = self.lon
outFile.variables['grid_dims'][:] = nPoints
outFile.variables['grid_imask'][:] = 1
# grid corners:
grid_corner_lon = numpy.zeros((nPoints, 4))
grid_corner_lat = numpy.zeros((nPoints, 4))
# just repeat the center lat and lon
for iVertex in range(4):
grid_corner_lat[:, iVertex] = self.lat
grid_corner_lon[:, iVertex] = self.lon
outFile.variables['grid_corner_lat'][:] = grid_corner_lat[:]
outFile.variables['grid_corner_lon'][:] = grid_corner_lon[:]
# Update history attribute of netCDF file
setattr(outFile, 'history', ' '.join(sys.argv[:]))
outFile.close() # }}}
def to_esmf(self, esmfFileName): # {{{
"""
Create an ESMF mesh file for the mesh
Parameters
----------
esmfFileName : str
The path to which the ESMF mesh file should be written
"""
# Authors
# -------
# Xylar Asay-Davis
nPoints = len(self.lat)
(Lon, Lat) = numpy.meshgrid(self.lon, self.lat)
elementCount = nPoints
nodeCount = 3*nPoints
coordDim = 2
maxNodePElement = 3
nodeCoords = numpy.zeros((nodeCount, coordDim))
# just repeat the center lat and lon
for iVertex in range(maxNodePElement):
nodeCoords[iVertex::maxNodePElement, 0] = self.lon
nodeCoords[iVertex::maxNodePElement, 1] = self.lat
centerCoords = numpy.zeros((elementCount, coordDim))
centerCoords[:, 0] = self.lon
centerCoords[:, 1] = self.lat
elementConn = numpy.zeros((elementCount, maxNodePElement), dtype=int)
elementConn[:, 0] = maxNodePElement * numpy.arange(nPoints)
elementConn[:, 1] = maxNodePElement * numpy.arange(nPoints) + 1
elementConn[:, 2] = maxNodePElement * numpy.arange(nPoints) + 2
ds_out = xarray.Dataset()
ds_out['nodeCoords'] = (('nodeCount', 'coordDim'), nodeCoords)
ds_out.nodeCoords.attrs['units'] = self.units
ds_out['centerCoords'] = \
(('elementCount', 'coordDim'), centerCoords)
ds_out.centerCoords.attrs['units'] = self.units
ds_out['elementConn'] = \
(('elementCount', 'maxNodePElement'), elementConn + 1)
ds_out.elementConn.attrs['start_index'] = 1
ds_out.elementConn.attrs['_FillValue'] = -1
ds_out['numElementConn'] = \
(('elementCount',), maxNodePElement * numpy.ones(elementCount,
dtype=int))
ds_out.attrs['gridType'] = 'unstructured mesh'
ds_out.attrs['version'] = '0.9'
ds_out.to_netcdf(esmfFileName) # }}}
# }}}
def interp_extrap_corner(inField): # {{{
'''Interpolate/extrapolate a 1D field from grid centers to grid corners'''
outField = numpy.zeros(len(inField) + 1)
outField[1:-1] = 0.5 * (inField[0:-1] + inField[1:])
# extrapolate the ends
outField[0] = 1.5 * inField[0] - 0.5 * inField[1]
outField[-1] = 1.5 * inField[-1] - 0.5 * inField[-2]
return outField # }}}
def interp_extrap_corners_2d(inField): # {{{
'''Interpolate/extrapolate a 1D field from grid centers to grid corners'''
temp = numpy.zeros((inField.shape[0], inField.shape[1] + 1))
temp[:, 1:-1] = 0.5 * (inField[:, 0:-1] + inField[:, 1:])
# extrapolate the ends
temp[:, 0] = 1.5 * inField[:, 0] - 0.5 * inField[:, 1]
temp[:, -1] = 1.5 * inField[:, -1] - 0.5 * inField[:, -2]
outField = numpy.zeros((inField.shape[0] + 1, inField.shape[1] + 1))
outField[1:-1, :] = 0.5 * (temp[0:-1, :] + temp[1:, :])
# extrapolate the ends
outField[0, :] = 1.5 * temp[0, :] - 0.5 * temp[1, :]
outField[-1, :] = 1.5 * temp[-1, :] - 0.5 * temp[-2, :]
return outField # }}}
def _create_scrip(outFile, grid_size, grid_corners, grid_rank, units,
meshName): # {{{
'''
Given a SCRIP files, creates common variables and writes common values used
in various types of SCRIP files.
Parameters
----------
outFile : file pointer
A SCRIP file opened in write mode
grid_size : int
The number of elements in the grid or mesh
grid_corners : int
The number of corners in the grid or mesh
grid_rank : int
The dimensionality of the grid (1 for mesh, 2 for grid)
units : {'degrees', 'radians'}
The units for latitude and longitude
meshName : str
The name of the mesh
'''
# Authors
# -------
# Xylar Asay-Davis
# Write to output file
# Dimensions
outFile.createDimension("grid_size", grid_size)
outFile.createDimension("grid_corners", grid_corners)
outFile.createDimension("grid_rank", grid_rank)
# Variables
grid_center_lat = outFile.createVariable('grid_center_lat', 'f8',
('grid_size',))
grid_center_lat.units = units
grid_center_lon = outFile.createVariable('grid_center_lon', 'f8',
('grid_size',))
grid_center_lon.units = units
grid_corner_lat = outFile.createVariable('grid_corner_lat', 'f8',
('grid_size', 'grid_corners'))
grid_corner_lat.units = units
grid_corner_lon = outFile.createVariable('grid_corner_lon', 'f8',
('grid_size', 'grid_corners'))
grid_corner_lon.units = units
grid_imask = outFile.createVariable('grid_imask', 'i4', ('grid_size',))
grid_imask.units = 'unitless'
outFile.createVariable('grid_dims', 'i4', ('grid_rank',))
outFile.meshName = meshName
# }}}
def _unwrap_corners(inField):
'''Turn a 2D array of corners into an array of rectangular mesh elements'''
outField = numpy.zeros(((inField.shape[0] - 1) * (inField.shape[1] - 1),
4))
# corners are counterclockwise
outField[:, 0] = inField[0:-1, 0:-1].flat
outField[:, 1] = inField[0:-1, 1:].flat
outField[:, 2] = inField[1:, 1:].flat
outField[:, 3] = inField[1:, 0:-1].flat
return outField
def _round_res(res):
'''Round the resoltuion to a reasonable number for grid names'''
rounded = numpy.round(res*1000.)/1000.
return '{}'.format(rounded)
# vim: ai ts=4 sts=4 et sw=4 ft=python