Skip to content

Update jigsaw_to_netcdf() for efficiency and robustness #632

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xylar merged 3 commits into from
Jun 3, 2025
Merged

Update jigsaw_to_netcdf() for efficiency and robustness #632

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
231ec68
Update circumcenter() to work with numpy arrays
xylar May 27, 2025
305ffbe
Update jigsaw_to_netcdf() for efficiency and robustness
xylar May 27, 2025
232146a
Add test for new circumcenter()
xylar May 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 95 additions & 107 deletions conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,12 @@
from __future__ import absolute_import, division, print_function, \
unicode_literals
import argparse

import numpy as np
import xarray as xr

from netCDF4 import Dataset as NetCDFFile
from mpas_tools.io import write_netcdf
from mpas_tools.mesh.creation.open_msh import readmsh
from mpas_tools.mesh.creation.util import circumcenter

import argparse


def jigsaw_to_netcdf(msh_filename, output_name, on_sphere, sphere_radius=None):
"""
Expand All @@ -18,150 +16,140 @@ def jigsaw_to_netcdf(msh_filename, output_name, on_sphere, sphere_radius=None):
----------
msh_filename : str
A JIGSAW mesh file name

output_name: str
The name of the output file

on_sphere : bool
Whether the mesh is spherical or planar

sphere_radius : float, optional
The radius of the sphere in meters. If ``on_sphere=True`` this argument
is required, otherwise it is ignored.
The radius of the sphere in meters. If ``on_sphere=True`` this
argument is required, otherwise it is ignored.
"""
# Authors: Phillip J. Wolfram, Matthew Hoffman and Xylar Asay-Davis

grid = NetCDFFile(output_name, 'w', format='NETCDF3_CLASSIC')
Copy link
Copy Markdown
Collaborator Author

@xylar xylar May 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was the line that mostly concerned time. We don't want to be restricted to JIGSAW meshes that are small enough for this old format.


# Get dimensions
# Get nCells
msh = readmsh(msh_filename)
nCells = msh['POINT'].shape[0]

# Get vertexDegree and nVertices
vertexDegree = 3 # always triangles with JIGSAW output
# always triangles with JIGSAW output
vertexDegree = 3
nVertices = msh['TRIA3'].shape[0]

if vertexDegree != 3:
ValueError("This script can only compute vertices with triangular "
"dual meshes currently.")

grid.createDimension('nCells', nCells)
grid.createDimension('nVertices', nVertices)
grid.createDimension('vertexDegree', vertexDegree)
raise ValueError(
'This script can only compute vertices with triangular '
'dual meshes currently.'
)

# Create cell variables and sphere_radius
xCell_full = msh['POINT'][:, 0]
yCell_full = msh['POINT'][:, 1]
zCell_full = []
if msh['NDIMS'] == 2:
zCell_full = np.zeros(yCell_full.shape)
elif msh['NDIMS'] == 3:
zCell_full = msh['POINT'][:, 2]
else:
raise ValueError("NDIMS must be 2 or 3; input mesh has NDIMS={}"
.format(msh['NDIMS']))
raise ValueError(
f'NDIMS must be 2 or 3; input mesh has NDIMS={msh["NDIMS"]}'
)
for cells in [xCell_full, yCell_full, zCell_full]:
assert cells.shape[0] == nCells, 'Number of anticipated nodes is' \
' not correct!'
assert cells.shape[0] == nCells, (
'Number of anticipated nodes is not correct!'
)

attrs = {}
if on_sphere:
grid.on_a_sphere = "YES"
grid.sphere_radius = sphere_radius
attrs['on_a_sphere'] = 'YES'
attrs['sphere_radius'] = sphere_radius
# convert from km to meters
xCell_full *= 1e3
yCell_full *= 1e3
zCell_full *= 1e3
xCell_full = xCell_full * 1e3
yCell_full = yCell_full * 1e3
zCell_full = zCell_full * 1e3
else:
grid.on_a_sphere = "NO"
grid.sphere_radius = 0.0
attrs['on_a_sphere'] = 'NO'
attrs['sphere_radius'] = 0.0

# Create cellsOnVertex
cellsOnVertex_full = msh['TRIA3'][:, :3] + 1
assert cellsOnVertex_full.shape == (nVertices, vertexDegree), \
assert cellsOnVertex_full.shape == (nVertices, vertexDegree), (
'cellsOnVertex_full is not the right shape!'

# Create vertex variables
xVertex_full = np.zeros((nVertices,))
yVertex_full = np.zeros((nVertices,))
zVertex_full = np.zeros((nVertices,))

for iVertex in np.arange(0, nVertices):
cell1 = cellsOnVertex_full[iVertex, 0]
cell2 = cellsOnVertex_full[iVertex, 1]
cell3 = cellsOnVertex_full[iVertex, 2]

x1 = xCell_full[cell1 - 1]
y1 = yCell_full[cell1 - 1]
z1 = zCell_full[cell1 - 1]
x2 = xCell_full[cell2 - 1]
y2 = yCell_full[cell2 - 1]
z2 = zCell_full[cell2 - 1]
x3 = xCell_full[cell3 - 1]
y3 = yCell_full[cell3 - 1]
z3 = zCell_full[cell3 - 1]

pv = circumcenter(on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3)
xVertex_full[iVertex] = pv.x
yVertex_full[iVertex] = pv.y
zVertex_full[iVertex] = pv.z

meshDensity_full = grid.createVariable(
'meshDensity', 'f8', ('nCells',))

for iCell in np.arange(0, nCells):
meshDensity_full[iCell] = 1.0

del meshDensity_full

var = grid.createVariable('xCell', 'f8', ('nCells',))
var[:] = xCell_full
var = grid.createVariable('yCell', 'f8', ('nCells',))
var[:] = yCell_full
var = grid.createVariable('zCell', 'f8', ('nCells',))
var[:] = zCell_full
var = grid.createVariable('xVertex', 'f8', ('nVertices',))
var[:] = xVertex_full
var = grid.createVariable('yVertex', 'f8', ('nVertices',))
var[:] = yVertex_full
var = grid.createVariable('zVertex', 'f8', ('nVertices',))
var[:] = zVertex_full
var = grid.createVariable(
'cellsOnVertex', 'i4', ('nVertices', 'vertexDegree',))
var[:] = cellsOnVertex_full

grid.sync()
grid.close()
)

# Vectorized circumcenter calculation
# shape (nVertices, vertexDegree)
cell_indices = cellsOnVertex_full - 1
x_cells = xCell_full[cell_indices]
y_cells = yCell_full[cell_indices]
z_cells = zCell_full[cell_indices]

# Vectorized circumcenter computation
x1, x2, x3 = x_cells[:, 0], x_cells[:, 1], x_cells[:, 2]
y1, y2, y3 = y_cells[:, 0], y_cells[:, 1], y_cells[:, 2]
z1, z2, z3 = z_cells[:, 0], z_cells[:, 1], z_cells[:, 2]
xVertex_full, yVertex_full, zVertex_full = circumcenter(
on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3
)

meshDensity_full = np.ones((nCells,), dtype=np.float64)

# Build xarray.Dataset
ds = xr.Dataset(
data_vars=dict(
xCell=(('nCells',), xCell_full),
yCell=(('nCells',), yCell_full),
zCell=(('nCells',), zCell_full),
xVertex=(('nVertices',), xVertex_full),
yVertex=(('nVertices',), yVertex_full),
zVertex=(('nVertices',), zVertex_full),
cellsOnVertex=(('nVertices', 'vertexDegree'), cellsOnVertex_full),
meshDensity=(('nCells',), meshDensity_full),
),
attrs=attrs,
)

# Write to NetCDF using write_netcdf
write_netcdf(ds, output_name)


def main():
"""
Convert a JIGSAW mesh file to NetCDF format for MPAS-Tools.
"""
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawTextHelpFormatter)
description=__doc__, formatter_class=argparse.RawTextHelpFormatter
)
parser.add_argument(
"-m",
"--msh",
dest="msh",
'-m',
'--msh',
dest='msh',
required=True,
help="input .msh file generated by JIGSAW.",
metavar="FILE")
help='input .msh file generated by JIGSAW.',
metavar='FILE',
)
parser.add_argument(
"-o",
"--output",
dest="output",
default="grid.nc",
help="output file name.",
metavar="FILE")
'-o',
'--output',
dest='output',
default='grid.nc',
help='output file name.',
metavar='FILE',
)
parser.add_argument(
"-s",
"--spherical",
dest="spherical",
action="store_true",
'-s',
'--spherical',
dest='spherical',
action='store_true',
default=False,
help="Determines if the input/output should be spherical or not.")
help='Determines if the input/output should be spherical or not.',
)
parser.add_argument(
"-r",
"--sphere_radius",
dest="sphere_radius",
'-r',
'--sphere_radius',
dest='sphere_radius',
type=float,
help="The radius of the sphere in meters")
help='The radius of the sphere in meters',
)

args = parser.parse_args()

jigsaw_to_netcdf(args.msh, args.output, args.spherical, args.sphere_radius)
81 changes: 50 additions & 31 deletions conda_package/mpas_tools/mesh/creation/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,51 +1,70 @@
import collections
import numpy as np

point = collections.namedtuple('Point', ['x', 'y', 'z'])


def circumcenter(on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3):
"""
Compute the circumcenter of the triangle (possibly on a sphere)
with the three given vertices in Cartesian coordinates.

Parameters
----------
on_sphere : bool
If True, the circumcenter is computed on a sphere.
If False, the circumcenter is computed in Cartesian coordinates.

x1, y1, z1 : float or numpy.ndarray
Cartesian coordinates of the first vertex

x2, y2, z2 : float or numpy.ndarray
Cartesian coordinates of the second vertex

x3, y3, z3 : float or numpy.ndarray
Cartesian coordinates of the third vertex

Returns
-------
center : point
The circumcenter of the triangle with x, y and z attributes

xv, yv, zv : float or numpy.ndarray
The circumcenter(s) of the triangle(s)
"""
p1 = point(x1, y1, z1)
p2 = point(x2, y2, z2)
p3 = point(x3, y3, z3)
x1 = np.asarray(x1)
y1 = np.asarray(y1)
z1 = np.asarray(z1)
x2 = np.asarray(x2)
y2 = np.asarray(y2)
z2 = np.asarray(z2)
x3 = np.asarray(x3)
y3 = np.asarray(y3)
z3 = np.asarray(z3)

if on_sphere:
a = (p2.x - p3.x)**2 + (p2.y - p3.y)**2 + (p2.z - p3.z)**2
b = (p3.x - p1.x)**2 + (p3.y - p1.y)**2 + (p3.z - p1.z)**2
c = (p1.x - p2.x)**2 + (p1.y - p2.y)**2 + (p1.z - p2.z)**2
a = (x2 - x3) ** 2 + (y2 - y3) ** 2 + (z2 - z3) ** 2
b = (x3 - x1) ** 2 + (y3 - y1) ** 2 + (z3 - z1) ** 2
c = (x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2

pbc = a * (-a + b + c)
apc = b * (a - b + c)
abp = c * (a + b - c)

xv = (pbc * p1.x + apc * p2.x + abp * p3.x) / (pbc + apc + abp)
yv = (pbc * p1.y + apc * p2.y + abp * p3.y) / (pbc + apc + abp)
zv = (pbc * p1.z + apc * p2.z + abp * p3.z) / (pbc + apc + abp)
denom = pbc + apc + abp
xv = (pbc * x1 + apc * x2 + abp * x3) / denom
yv = (pbc * y1 + apc * y2 + abp * y3) / denom
zv = (pbc * z1 + apc * z2 + abp * z3) / denom
else:
d = 2 * (p1.x * (p2.y - p3.y) + p2.x *
(p3.y - p1.y) + p3.x * (p1.y - p2.y))
d = 2 * (x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2))

xv = ((p1.x**2 + p1.y**2) * (p2.y - p3.y) + (p2.x**2 + p2.y**2)
* (p3.y - p1.y) + (p3.x**2 + p3.y**2) * (p1.y - p2.y)) / d
yv = ((p1.x**2 + p1.y**2) * (p3.x - p2.x) + (p2.x**2 + p2.y**2)
* (p1.x - p3.x) + (p3.x**2 + p3.y**2) * (p2.x - p1.x)) / d
zv = 0.0
xv = (
(x1**2 + y1**2) * (y2 - y3)
+ (x2**2 + y2**2) * (y3 - y1)
+ (x3**2 + y3**2) * (y1 - y2)
) / d
yv = (
(x1**2 + y1**2) * (x3 - x2)
+ (x2**2 + y2**2) * (x1 - x3)
+ (x3**2 + y3**2) * (x2 - x1)
) / d
zv = np.zeros_like(xv)

# Optional method to use barycenter instead.
# xv = p1.x + p2.x + p3.x
# xv = xv / 3.0
# yv = p1.y + p2.y + p3.y
# yv = yv / 3.0
return point(xv, yv, zv)
return xv, yv, zv


def lonlat2xyz(lon, lat, R=6378206.4):
Expand All @@ -69,8 +88,8 @@ def lonlat2xyz(lon, lat, R=6378206.4):

lon = np.deg2rad(lon)
lat = np.deg2rad(lat)
x = R*np.multiply(np.cos(lon), np.cos(lat))
y = R*np.multiply(np.sin(lon), np.cos(lat))
z = R*np.sin(lat)
x = R * np.multiply(np.cos(lon), np.cos(lat))
y = R * np.multiply(np.sin(lon), np.cos(lat))
z = R * np.sin(lat)

return x, y, z
Loading
Loading