#!/usr/bin/env python
from __future__ import absolute_import, division, print_function, \
unicode_literals
import numpy
import xarray
import argparse
from mpas_tools.io import write_netcdf
[docs]def make_planar_hex_mesh(nx, ny, dc, nonperiodic_x,
nonperiodic_y, outFileName=None,
compareWithFileName=None,
format=None, engine=None):
"""
Builds an MPAS periodic, planar hexagonal mesh with the requested
dimensions, optionally saving it to a file, and returns it as an
``xarray.Dataset``.
Parameters
----------
nx : int
The number of cells in the x direction
ny : even int
The number of cells in the y direction (must be an even number for
periodicity to work out)
dc : float
The distance in meters between adjacent cell centers.
nonperiodic_x, nonperiodic_y : bool
is the mesh non-periodic in x and y directions?
outFileName : str, optional
The name of a file to save the mesh to. The mesh is not saved to a
file if no file name is supplied.
compareWithFileName : str, optional
The name of a grid file to compare with to see if they are identical,
used for testing purposes
format : {'NETCDF4', 'NETCDF4_CLASSIC', 'NETCDF3_64BIT', '
NETCDF3_CLASSIC'}, optional
The NetCDF format to use for output
engine : {'netcdf4', 'scipy', 'h5netcdf'}, optional
The library to use for NetCDF output
Returns
-------
mesh : xarray.Dataset
The mesh data set, available for further manipulation such as culling
cells or removing periodicity.
"""
mesh = initial_setup(nx, ny, dc, nonperiodic_x, nonperiodic_y)
compute_indices_on_cell(mesh)
if nonperiodic_x:
mark_cull_cell_nonperiodic_x(mesh)
if nonperiodic_y:
mark_cull_cell_nonperiodic_y(mesh)
compute_indices_on_edge(mesh)
compute_indices_on_vertex(mesh)
compute_weights_on_edge(mesh)
compute_coordinates(mesh)
add_one_to_indices(mesh)
# drop some arrays that aren't standard for MPAS but were used to compute
# the hex mesh
mesh = mesh.drop_vars(['cellIdx', 'cellRow', 'cellCol'])
if outFileName is not None:
write_netcdf(mesh, outFileName, format=format, engine=engine)
if compareWithFileName is not None:
# used to make sure results are exactly identical to periodic_hex
make_diff(mesh, compareWithFileName, 'diff.nc')
return mesh
def initial_setup(nx, ny, dc, nonperiodic_x, nonperiodic_y):
"""Setup the dimensions and add placeholders for some index variables"""
if ny % 2 != 0:
raise ValueError('ny must be divisible by 2 for the grid\'s '
'periodicity to work properly.')
mesh = xarray.Dataset()
if nonperiodic_x and nonperiodic_y:
mesh.attrs['is_periodic'] = 'NO'
else:
mesh.attrs['is_periodic'] = 'YES'
if nonperiodic_x:
mesh.attrs['x_period'] = 0.
else:
mesh.attrs['x_period'] = nx * dc
if nonperiodic_y:
mesh.attrs['y_period'] = 0.
else:
mesh.attrs['y_period'] = ny * dc * numpy.sqrt(3.) / 2.
mesh.attrs['dc'] = dc
mesh.attrs['nx'] = nx
mesh.attrs['ny'] = ny
mesh.attrs['on_a_sphere'] = 'NO'
mesh.attrs['sphere_radius'] = 0.
if nonperiodic_x:
nx = nx + 2
if nonperiodic_y:
ny = ny + 2
nCells = nx * ny
nEdges = 3 * nCells
nVertices = 2 * nCells
vertexDegree = 3
maxEdges = 6
# add some basic arrays to get all the dimensions in place
indexToCellID = numpy.arange(nCells, dtype='i4')
indexToEdgeID = numpy.arange(nEdges, dtype='i4')
indexToVertexID = numpy.arange(nVertices, dtype='i4')
cellIdx = indexToCellID.reshape(ny, nx)
cellCol, cellRow = numpy.meshgrid(numpy.arange(nx, dtype='i4'),
numpy.arange(ny, dtype='i4'))
mesh['cellIdx'] = (('ny', 'nx'), cellIdx)
mesh['cellRow'] = (('nCells',), cellRow.ravel())
mesh['cellCol'] = (('nCells',), cellCol.ravel())
mesh['indexToCellID'] = (('nCells',), indexToCellID)
mesh['indexToEdgeID'] = (('nEdges',), indexToEdgeID)
mesh['indexToVertexID'] = (('nVertices',), indexToVertexID)
mesh['cullCell'] = (('nCells',), numpy.zeros(nCells, 'i4'))
mesh['nEdgesOnCell'] = (('nCells',), 6 * numpy.ones((nCells,), 'i4'))
mesh['cellsOnCell'] = (('nCells', 'maxEdges'),
numpy.zeros((nCells, maxEdges), 'i4'))
mesh['edgesOnCell'] = (('nCells', 'maxEdges'),
numpy.zeros((nCells, maxEdges), 'i4'))
mesh['verticesOnCell'] = (('nCells', 'maxEdges'),
numpy.zeros((nCells, maxEdges), 'i4'))
mesh['nEdgesOnEdge'] = (('nEdges',), 10 * numpy.ones((nEdges,), 'i4'))
mesh['cellsOnEdge'] = (('nEdges', 'TWO'),
numpy.zeros((nEdges, 2), 'i4'))
mesh['edgesOnEdge'] = (('nEdges', 'maxEdges2'),
-1 * numpy.ones((nEdges, 2 * maxEdges), 'i4'))
mesh['verticesOnEdge'] = (('nEdges', 'TWO'),
numpy.zeros((nEdges, 2), 'i4'))
mesh['cellsOnVertex'] = (('nVertices', 'vertexDegree'),
numpy.zeros((nVertices, vertexDegree), 'i4'))
mesh['edgesOnVertex'] = (('nVertices', 'vertexDegree'),
numpy.zeros((nVertices, vertexDegree), 'i4'))
return mesh
def mark_cull_cell_nonperiodic_y(mesh):
cullCell = mesh.cullCell
nCells = mesh.sizes['nCells']
nx = mesh.sizes['nx']
cullCell[0:nx] = 1
cullCell[nCells - nx:nCells + 1] = 1
def mark_cull_cell_nonperiodic_x(mesh):
cullCell = mesh.cullCell
nCells = mesh.sizes['nCells']
nx = mesh.sizes['nx']
cullCell[::nx] = 1
cullCell[nx - 1:nCells + 1:nx] = 1
def compute_indices_on_cell(mesh):
cellIdx = mesh.cellIdx
cellRow = mesh.cellRow
cellCol = mesh.cellCol
indexToCellID = mesh.indexToCellID
nx = mesh.sizes['nx']
ny = mesh.sizes['ny']
mx = numpy.mod(cellCol - 1, nx)
my = numpy.mod(cellRow - 1, ny)
px = numpy.mod(cellCol + 1, nx)
py = numpy.mod(cellRow + 1, ny)
mask = numpy.mod(cellRow, 2) == 0
cellsOnCell = mesh.cellsOnCell
cellsOnCell[:, 0] = cellIdx[cellRow, mx]
cellsOnCell[:, 1] = cellIdx[my, mx].where(mask, cellIdx[my, cellCol])
cellsOnCell[:, 2] = cellIdx[my, cellCol].where(mask, cellIdx[my, px])
cellsOnCell[:, 3] = cellIdx[cellRow, px]
cellsOnCell[:, 4] = cellIdx[py, cellCol].where(mask, cellIdx[py, px])
cellsOnCell[:, 5] = cellIdx[py, mx].where(mask, cellIdx[py, cellCol])
edgesOnCell = mesh.edgesOnCell
edgesOnCell[:, 0] = 3 * indexToCellID
edgesOnCell[:, 1] = 3 * indexToCellID + 1
edgesOnCell[:, 2] = 3 * indexToCellID + 2
edgesOnCell[:, 3] = 3 * cellsOnCell[:, 3]
edgesOnCell[:, 4] = 3 * cellsOnCell[:, 4] + 1
edgesOnCell[:, 5] = 3 * cellsOnCell[:, 5] + 2
verticesOnCell = mesh.verticesOnCell
verticesOnCell[:, 0] = 2 * indexToCellID
verticesOnCell[:, 1] = 2 * indexToCellID + 1
verticesOnCell[:, 2] = 2 * cellsOnCell[:, 2]
verticesOnCell[:, 3] = 2 * cellsOnCell[:, 3] + 1
verticesOnCell[:, 4] = 2 * cellsOnCell[:, 3]
verticesOnCell[:, 5] = 2 * cellsOnCell[:, 4] + 1
def compute_indices_on_edge(mesh):
edgesOnCell = mesh.edgesOnCell
verticesOnCell = mesh.verticesOnCell
indexToCellID = mesh.indexToCellID
cellsOnEdge = mesh.cellsOnEdge
for j in range(3):
cellsOnEdge[edgesOnCell[:, j], 1] = indexToCellID
for j in range(3, 6):
cellsOnEdge[edgesOnCell[:, j], 0] = indexToCellID
verticesOnEdge = mesh.verticesOnEdge
verticesOnEdge[edgesOnCell[:, 0], 0] = verticesOnCell[:, 1]
verticesOnEdge[edgesOnCell[:, 0], 1] = verticesOnCell[:, 0]
verticesOnEdge[edgesOnCell[:, 1], 0] = verticesOnCell[:, 2]
verticesOnEdge[edgesOnCell[:, 1], 1] = verticesOnCell[:, 1]
verticesOnEdge[edgesOnCell[:, 2], 0] = verticesOnCell[:, 3]
verticesOnEdge[edgesOnCell[:, 2], 1] = verticesOnCell[:, 2]
edgesOnEdge = mesh.edgesOnEdge
edgesOnEdge[edgesOnCell[:, 3], 0] = edgesOnCell[:, 4]
edgesOnEdge[edgesOnCell[:, 3], 1] = edgesOnCell[:, 5]
edgesOnEdge[edgesOnCell[:, 3], 2] = edgesOnCell[:, 0]
edgesOnEdge[edgesOnCell[:, 3], 3] = edgesOnCell[:, 1]
edgesOnEdge[edgesOnCell[:, 3], 4] = edgesOnCell[:, 2]
edgesOnEdge[edgesOnCell[:, 4], 0] = edgesOnCell[:, 5]
edgesOnEdge[edgesOnCell[:, 4], 1] = edgesOnCell[:, 0]
edgesOnEdge[edgesOnCell[:, 4], 2] = edgesOnCell[:, 1]
edgesOnEdge[edgesOnCell[:, 4], 3] = edgesOnCell[:, 2]
edgesOnEdge[edgesOnCell[:, 4], 4] = edgesOnCell[:, 3]
edgesOnEdge[edgesOnCell[:, 5], 0] = edgesOnCell[:, 0]
edgesOnEdge[edgesOnCell[:, 5], 1] = edgesOnCell[:, 1]
edgesOnEdge[edgesOnCell[:, 5], 2] = edgesOnCell[:, 2]
edgesOnEdge[edgesOnCell[:, 5], 3] = edgesOnCell[:, 3]
edgesOnEdge[edgesOnCell[:, 5], 4] = edgesOnCell[:, 4]
edgesOnEdge[edgesOnCell[:, 0], 5] = edgesOnCell[:, 1]
edgesOnEdge[edgesOnCell[:, 0], 6] = edgesOnCell[:, 2]
edgesOnEdge[edgesOnCell[:, 0], 7] = edgesOnCell[:, 3]
edgesOnEdge[edgesOnCell[:, 0], 8] = edgesOnCell[:, 4]
edgesOnEdge[edgesOnCell[:, 0], 9] = edgesOnCell[:, 5]
edgesOnEdge[edgesOnCell[:, 1], 5] = edgesOnCell[:, 2]
edgesOnEdge[edgesOnCell[:, 1], 6] = edgesOnCell[:, 3]
edgesOnEdge[edgesOnCell[:, 1], 7] = edgesOnCell[:, 4]
edgesOnEdge[edgesOnCell[:, 1], 8] = edgesOnCell[:, 5]
edgesOnEdge[edgesOnCell[:, 1], 9] = edgesOnCell[:, 0]
edgesOnEdge[edgesOnCell[:, 2], 5] = edgesOnCell[:, 3]
edgesOnEdge[edgesOnCell[:, 2], 6] = edgesOnCell[:, 4]
edgesOnEdge[edgesOnCell[:, 2], 7] = edgesOnCell[:, 5]
edgesOnEdge[edgesOnCell[:, 2], 8] = edgesOnCell[:, 0]
edgesOnEdge[edgesOnCell[:, 2], 9] = edgesOnCell[:, 1]
def compute_indices_on_vertex(mesh):
edgesOnCell = mesh.edgesOnCell
verticesOnCell = mesh.verticesOnCell
indexToCellID = mesh.indexToCellID
cellsOnVertex = mesh.cellsOnVertex
cellsOnVertex[verticesOnCell[:, 1], 2] = indexToCellID
cellsOnVertex[verticesOnCell[:, 3], 0] = indexToCellID
cellsOnVertex[verticesOnCell[:, 5], 1] = indexToCellID
cellsOnVertex[verticesOnCell[:, 0], 0] = indexToCellID
cellsOnVertex[verticesOnCell[:, 2], 1] = indexToCellID
cellsOnVertex[verticesOnCell[:, 4], 2] = indexToCellID
edgesOnVertex = mesh.edgesOnVertex
edgesOnVertex[verticesOnCell[:, 0], 0] = edgesOnCell[:, 0]
edgesOnVertex[verticesOnCell[:, 1], 0] = edgesOnCell[:, 0]
edgesOnVertex[verticesOnCell[:, 2], 2] = edgesOnCell[:, 1]
edgesOnVertex[verticesOnCell[:, 1], 2] = edgesOnCell[:, 1]
edgesOnVertex[verticesOnCell[:, 2], 1] = edgesOnCell[:, 2]
edgesOnVertex[verticesOnCell[:, 3], 1] = edgesOnCell[:, 2]
def compute_weights_on_edge(mesh):
edgesOnCell = mesh.edgesOnCell
nEdges = mesh.sizes['nEdges']
maxEdges2 = mesh.sizes['maxEdges2']
mesh['weightsOnEdge'] = (('nEdges', 'maxEdges2'),
numpy.zeros((nEdges, maxEdges2), 'f8'))
weightsOnEdge = mesh.weightsOnEdge
weights = (1. / numpy.sqrt(3.)) * numpy.array(
[[1. / 3., 1. / 6., 0., 1. / 6., 1. / 3.],
[1. / 3., -1. / 6., 0., 1. / 6., -1. / 3.],
[-1. / 3., -1. / 6., 0., -1. / 6., -1. / 3.]])
for i in range(3):
for j in range(5):
weightsOnEdge[edgesOnCell[:, i + 3], j] = weights[i, j]
for i in range(3):
for j in range(5):
weightsOnEdge[edgesOnCell[:, i], j + 5] = weights[i, j]
def compute_coordinates(mesh):
dc = mesh.attrs['dc']
edgesOnCell = mesh.edgesOnCell
verticesOnCell = mesh.verticesOnCell
nCells = mesh.sizes['nCells']
nEdges = mesh.sizes['nEdges']
nVertices = mesh.sizes['nVertices']
vertexDegree = mesh.sizes['vertexDegree']
mesh['latCell'] = (('nCells',), numpy.zeros((nCells,), 'f8'))
mesh['lonCell'] = (('nCells',), numpy.zeros((nCells,), 'f8'))
mesh['latEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh['lonEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh['latVertex'] = (('nVertices',), numpy.zeros((nVertices,), 'f8'))
mesh['lonVertex'] = (('nVertices',), numpy.zeros((nVertices,), 'f8'))
cellRow = mesh.cellRow
cellCol = mesh.cellCol
mask = numpy.mod(cellRow, 2) == 0
mesh['xCell'] = (dc * (cellCol + 0.5)).where(mask, dc * (cellCol + 1))
mesh['yCell'] = dc * (cellRow + 1) * numpy.sqrt(3.) / 2.
mesh['zCell'] = (('nCells',), numpy.zeros((nCells,), 'f8'))
mesh['xEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh['yEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh['zEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh.xEdge[edgesOnCell[:, 0]] = mesh.xCell - 0.5 * dc
mesh.yEdge[edgesOnCell[:, 0]] = mesh.yCell
mesh.xEdge[edgesOnCell[:, 1]] = mesh.xCell - \
0.5 * dc * numpy.cos(numpy.pi / 3.)
mesh.yEdge[edgesOnCell[:, 1]] = mesh.yCell - \
0.5 * dc * numpy.sin(numpy.pi / 3.)
mesh.xEdge[edgesOnCell[:, 2]] = mesh.xCell + \
0.5 * dc * numpy.cos(numpy.pi / 3.)
mesh.yEdge[edgesOnCell[:, 2]] = mesh.yCell - \
0.5 * dc * numpy.sin(numpy.pi / 3.)
mesh['xVertex'] = (('nVertices',), numpy.zeros((nVertices,), 'f8'))
mesh['yVertex'] = (('nVertices',), numpy.zeros((nVertices,), 'f8'))
mesh['zVertex'] = (('nVertices',), numpy.zeros((nVertices,), 'f8'))
mesh.xVertex[verticesOnCell[:, 0]] = mesh.xCell - 0.5 * dc
mesh.yVertex[verticesOnCell[:, 0]] = mesh.yCell + dc * numpy.sqrt(3.) / 6.
mesh.xVertex[verticesOnCell[:, 1]] = mesh.xCell - 0.5 * dc
mesh.yVertex[verticesOnCell[:, 1]] = mesh.yCell - dc * numpy.sqrt(3.) / 6.
mesh['angleEdge'] = (('nEdges',), numpy.zeros((nEdges,), 'f8'))
mesh.angleEdge[edgesOnCell[:, 1]] = numpy.pi / 3.
mesh.angleEdge[edgesOnCell[:, 2]] = 2. * numpy.pi / 3.
mesh['dcEdge'] = (('nEdges',), dc * numpy.ones((nEdges,), 'f8'))
mesh['dvEdge'] = mesh.dcEdge * numpy.sqrt(3.) / 3.
mesh['areaCell'] = \
(('nCells',), dc**2 * numpy.sqrt(3.) / 2. * numpy.ones((nCells,), 'f8'))
mesh['areaTriangle'] = \
(('nVertices',), dc**2 * numpy.sqrt(3.) /
4. * numpy.ones((nVertices,), 'f8'))
mesh['kiteAreasOnVertex'] = \
(('nVertices', 'vertexDegree'),
dc**2 * numpy.sqrt(3.) / 12. * numpy.ones((nVertices, vertexDegree),
'f8'))
mesh['meshDensity'] = (('nCells',), numpy.ones((nCells,), 'f8'))
def add_one_to_indices(mesh):
"""Needed to adhere to Fortran indexing"""
indexVars = ['indexToCellID', 'indexToEdgeID', 'indexToVertexID',
'cellsOnCell', 'edgesOnCell', 'verticesOnCell',
'cellsOnEdge', 'edgesOnEdge', 'verticesOnEdge',
'cellsOnVertex', 'edgesOnVertex']
for var in indexVars:
mesh[var] = mesh[var] + 1
def make_diff(mesh, refMeshFileName, diffFileName):
refMesh = xarray.open_dataset(refMeshFileName)
diff = xarray.Dataset()
for variable in mesh.data_vars:
if variable in refMesh:
diff[variable] = mesh[variable] - refMesh[variable]
print(diff[variable].name, float(numpy.abs(diff[variable]).max()))
else:
print('mesh has extra variable {}'.format(mesh[variable].name))
for variable in refMesh.data_vars:
if variable not in mesh:
print('mesh mising variable {}'.format(refMesh[variable].name))
for attr in refMesh.attrs:
if attr not in mesh.attrs:
print('mesh mising attribute {}'.format(attr))
for attr in mesh.attrs:
if attr not in refMesh.attrs:
print('mesh has extra attribute {}'.format(attr))
write_netcdf(diff, diffFileName)
def main():
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('--nx', dest='nx', type=int, required=True,
help='Cells in x direction')
parser.add_argument('--ny', dest='ny', type=int, required=True,
help='Cells in y direction')
parser.add_argument('--dc', dest='dc', type=float, required=True,
help='Distance between cell centers in meters')
parser.add_argument('--npx', '--nonperiodic_x', dest='nonperiodic_x',
action="store_true",
help='non-periodic in x direction')
parser.add_argument('--npy', '--nonperiodic_y', dest='nonperiodic_y',
action="store_true",
help='non-periodic in y direction')
parser.add_argument('-o', '--outFileName', dest='outFileName', type=str,
required=False, default='grid.nc',
help='The name of the output file')
args = parser.parse_args()
make_planar_hex_mesh(args.nx, args.ny, args.dc,
args.nonperiodic_x, args.nonperiodic_y,
args.outFileName)
if __name__ == '__main__':
main()