Skip to content

Commit d2bbd6f

Browse files
xylarxylar
authored
Merge pull request #632 from xylar/clean-up-jigsaw-conversion
Update `jigsaw_to_netcdf()` for efficiency and robustness
2 parents 1557e3f + 232146a commit d2bbd6f

3 files changed

Lines changed: 295 additions & 138 deletions

File tree

‎conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py‎

Lines changed: 95 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,12 @@
1-
from __future__ import absolute_import, division, print_function, \
2-
unicode_literals
1+
import argparse
32

43
import numpy as np
4+
import xarray as xr
55

6-
from netCDF4 import Dataset as NetCDFFile
6+
from mpas_tools.io import write_netcdf
77
from mpas_tools.mesh.creation.open_msh import readmsh
88
from mpas_tools.mesh.creation.util import circumcenter
99

10-
import argparse
11-
1210

1311
def jigsaw_to_netcdf(msh_filename, output_name, on_sphere, sphere_radius=None):
1412
"""
@@ -18,150 +16,140 @@ def jigsaw_to_netcdf(msh_filename, output_name, on_sphere, sphere_radius=None):
1816
----------
1917
msh_filename : str
2018
A JIGSAW mesh file name
19+
2120
output_name: str
2221
The name of the output file
22+
2323
on_sphere : bool
2424
Whether the mesh is spherical or planar
25+
2526
sphere_radius : float, optional
26-
The radius of the sphere in meters. If ``on_sphere=True`` this argument
27-
is required, otherwise it is ignored.
27+
The radius of the sphere in meters. If ``on_sphere=True`` this
28+
argument is required, otherwise it is ignored.
2829
"""
29-
# Authors: Phillip J. Wolfram, Matthew Hoffman and Xylar Asay-Davis
30-
31-
grid = NetCDFFile(output_name, 'w', format='NETCDF3_CLASSIC')
3230

3331
# Get dimensions
3432
# Get nCells
3533
msh = readmsh(msh_filename)
3634
nCells = msh['POINT'].shape[0]
37-
38-
# Get vertexDegree and nVertices
39-
vertexDegree = 3 # always triangles with JIGSAW output
35+
# always triangles with JIGSAW output
36+
vertexDegree = 3
4037
nVertices = msh['TRIA3'].shape[0]
4138

4239
if vertexDegree != 3:
43-
ValueError("This script can only compute vertices with triangular "
44-
"dual meshes currently.")
45-
46-
grid.createDimension('nCells', nCells)
47-
grid.createDimension('nVertices', nVertices)
48-
grid.createDimension('vertexDegree', vertexDegree)
40+
raise ValueError(
41+
'This script can only compute vertices with triangular '
42+
'dual meshes currently.'
43+
)
4944

5045
# Create cell variables and sphere_radius
5146
xCell_full = msh['POINT'][:, 0]
5247
yCell_full = msh['POINT'][:, 1]
53-
zCell_full = []
5448
if msh['NDIMS'] == 2:
5549
zCell_full = np.zeros(yCell_full.shape)
5650
elif msh['NDIMS'] == 3:
5751
zCell_full = msh['POINT'][:, 2]
5852
else:
59-
raise ValueError("NDIMS must be 2 or 3; input mesh has NDIMS={}"
60-
.format(msh['NDIMS']))
53+
raise ValueError(
54+
f'NDIMS must be 2 or 3; input mesh has NDIMS={msh["NDIMS"]}'
55+
)
6156
for cells in [xCell_full, yCell_full, zCell_full]:
62-
assert cells.shape[0] == nCells, 'Number of anticipated nodes is' \
63-
' not correct!'
57+
assert cells.shape[0] == nCells, (
58+
'Number of anticipated nodes is not correct!'
59+
)
60+
61+
attrs = {}
6462
if on_sphere:
65-
grid.on_a_sphere = "YES"
66-
grid.sphere_radius = sphere_radius
63+
attrs['on_a_sphere'] = 'YES'
64+
attrs['sphere_radius'] = sphere_radius
6765
# convert from km to meters
68-
xCell_full *= 1e3
69-
yCell_full *= 1e3
70-
zCell_full *= 1e3
66+
xCell_full = xCell_full * 1e3
67+
yCell_full = yCell_full * 1e3
68+
zCell_full = zCell_full * 1e3
7169
else:
72-
grid.on_a_sphere = "NO"
73-
grid.sphere_radius = 0.0
70+
attrs['on_a_sphere'] = 'NO'
71+
attrs['sphere_radius'] = 0.0
7472

7573
# Create cellsOnVertex
7674
cellsOnVertex_full = msh['TRIA3'][:, :3] + 1
77-
assert cellsOnVertex_full.shape == (nVertices, vertexDegree), \
75+
assert cellsOnVertex_full.shape == (nVertices, vertexDegree), (
7876
'cellsOnVertex_full is not the right shape!'
79-
80-
# Create vertex variables
81-
xVertex_full = np.zeros((nVertices,))
82-
yVertex_full = np.zeros((nVertices,))
83-
zVertex_full = np.zeros((nVertices,))
84-
85-
for iVertex in np.arange(0, nVertices):
86-
cell1 = cellsOnVertex_full[iVertex, 0]
87-
cell2 = cellsOnVertex_full[iVertex, 1]
88-
cell3 = cellsOnVertex_full[iVertex, 2]
89-
90-
x1 = xCell_full[cell1 - 1]
91-
y1 = yCell_full[cell1 - 1]
92-
z1 = zCell_full[cell1 - 1]
93-
x2 = xCell_full[cell2 - 1]
94-
y2 = yCell_full[cell2 - 1]
95-
z2 = zCell_full[cell2 - 1]
96-
x3 = xCell_full[cell3 - 1]
97-
y3 = yCell_full[cell3 - 1]
98-
z3 = zCell_full[cell3 - 1]
99-
100-
pv = circumcenter(on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3)
101-
xVertex_full[iVertex] = pv.x
102-
yVertex_full[iVertex] = pv.y
103-
zVertex_full[iVertex] = pv.z
104-
105-
meshDensity_full = grid.createVariable(
106-
'meshDensity', 'f8', ('nCells',))
107-
108-
for iCell in np.arange(0, nCells):
109-
meshDensity_full[iCell] = 1.0
110-
111-
del meshDensity_full
112-
113-
var = grid.createVariable('xCell', 'f8', ('nCells',))
114-
var[:] = xCell_full
115-
var = grid.createVariable('yCell', 'f8', ('nCells',))
116-
var[:] = yCell_full
117-
var = grid.createVariable('zCell', 'f8', ('nCells',))
118-
var[:] = zCell_full
119-
var = grid.createVariable('xVertex', 'f8', ('nVertices',))
120-
var[:] = xVertex_full
121-
var = grid.createVariable('yVertex', 'f8', ('nVertices',))
122-
var[:] = yVertex_full
123-
var = grid.createVariable('zVertex', 'f8', ('nVertices',))
124-
var[:] = zVertex_full
125-
var = grid.createVariable(
126-
'cellsOnVertex', 'i4', ('nVertices', 'vertexDegree',))
127-
var[:] = cellsOnVertex_full
128-
129-
grid.sync()
130-
grid.close()
77+
)
78+
79+
# Vectorized circumcenter calculation
80+
# shape (nVertices, vertexDegree)
81+
cell_indices = cellsOnVertex_full - 1
82+
x_cells = xCell_full[cell_indices]
83+
y_cells = yCell_full[cell_indices]
84+
z_cells = zCell_full[cell_indices]
85+
86+
# Vectorized circumcenter computation
87+
x1, x2, x3 = x_cells[:, 0], x_cells[:, 1], x_cells[:, 2]
88+
y1, y2, y3 = y_cells[:, 0], y_cells[:, 1], y_cells[:, 2]
89+
z1, z2, z3 = z_cells[:, 0], z_cells[:, 1], z_cells[:, 2]
90+
xVertex_full, yVertex_full, zVertex_full = circumcenter(
91+
on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3
92+
)
93+
94+
meshDensity_full = np.ones((nCells,), dtype=np.float64)
95+
96+
# Build xarray.Dataset
97+
ds = xr.Dataset(
98+
data_vars=dict(
99+
xCell=(('nCells',), xCell_full),
100+
yCell=(('nCells',), yCell_full),
101+
zCell=(('nCells',), zCell_full),
102+
xVertex=(('nVertices',), xVertex_full),
103+
yVertex=(('nVertices',), yVertex_full),
104+
zVertex=(('nVertices',), zVertex_full),
105+
cellsOnVertex=(('nVertices', 'vertexDegree'), cellsOnVertex_full),
106+
meshDensity=(('nCells',), meshDensity_full),
107+
),
108+
attrs=attrs,
109+
)
110+
111+
# Write to NetCDF using write_netcdf
112+
write_netcdf(ds, output_name)
131113

132114

133115
def main():
116+
"""
117+
Convert a JIGSAW mesh file to NetCDF format for MPAS-Tools.
118+
"""
134119
parser = argparse.ArgumentParser(
135-
description=__doc__,
136-
formatter_class=argparse.RawTextHelpFormatter)
120+
description=__doc__, formatter_class=argparse.RawTextHelpFormatter
121+
)
137122
parser.add_argument(
138-
"-m",
139-
"--msh",
140-
dest="msh",
123+
'-m',
124+
'--msh',
125+
dest='msh',
141126
required=True,
142-
help="input .msh file generated by JIGSAW.",
143-
metavar="FILE")
127+
help='input .msh file generated by JIGSAW.',
128+
metavar='FILE',
129+
)
144130
parser.add_argument(
145-
"-o",
146-
"--output",
147-
dest="output",
148-
default="grid.nc",
149-
help="output file name.",
150-
metavar="FILE")
131+
'-o',
132+
'--output',
133+
dest='output',
134+
default='grid.nc',
135+
help='output file name.',
136+
metavar='FILE',
137+
)
151138
parser.add_argument(
152-
"-s",
153-
"--spherical",
154-
dest="spherical",
155-
action="store_true",
139+
'-s',
140+
'--spherical',
141+
dest='spherical',
142+
action='store_true',
156143
default=False,
157-
help="Determines if the input/output should be spherical or not.")
144+
help='Determines if the input/output should be spherical or not.',
145+
)
158146
parser.add_argument(
159-
"-r",
160-
"--sphere_radius",
161-
dest="sphere_radius",
147+
'-r',
148+
'--sphere_radius',
149+
dest='sphere_radius',
162150
type=float,
163-
help="The radius of the sphere in meters")
151+
help='The radius of the sphere in meters',
152+
)
164153

165154
args = parser.parse_args()
166-
167155
jigsaw_to_netcdf(args.msh, args.output, args.spherical, args.sphere_radius)

‎conda_package/mpas_tools/mesh/creation/util.py‎

Lines changed: 50 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,70 @@
1-
import collections
21
import numpy as np
32

4-
point = collections.namedtuple('Point', ['x', 'y', 'z'])
5-
63

74
def circumcenter(on_sphere, x1, y1, z1, x2, y2, z2, x3, y3, z3):
85
"""
96
Compute the circumcenter of the triangle (possibly on a sphere)
107
with the three given vertices in Cartesian coordinates.
118
9+
Parameters
10+
----------
11+
on_sphere : bool
12+
If True, the circumcenter is computed on a sphere.
13+
If False, the circumcenter is computed in Cartesian coordinates.
14+
15+
x1, y1, z1 : float or numpy.ndarray
16+
Cartesian coordinates of the first vertex
17+
18+
x2, y2, z2 : float or numpy.ndarray
19+
Cartesian coordinates of the second vertex
20+
21+
x3, y3, z3 : float or numpy.ndarray
22+
Cartesian coordinates of the third vertex
23+
1224
Returns
1325
-------
14-
center : point
15-
The circumcenter of the triangle with x, y and z attributes
16-
26+
xv, yv, zv : float or numpy.ndarray
27+
The circumcenter(s) of the triangle(s)
1728
"""
18-
p1 = point(x1, y1, z1)
19-
p2 = point(x2, y2, z2)
20-
p3 = point(x3, y3, z3)
29+
x1 = np.asarray(x1)
30+
y1 = np.asarray(y1)
31+
z1 = np.asarray(z1)
32+
x2 = np.asarray(x2)
33+
y2 = np.asarray(y2)
34+
z2 = np.asarray(z2)
35+
x3 = np.asarray(x3)
36+
y3 = np.asarray(y3)
37+
z3 = np.asarray(z3)
38+
2139
if on_sphere:
22-
a = (p2.x - p3.x)**2 + (p2.y - p3.y)**2 + (p2.z - p3.z)**2
23-
b = (p3.x - p1.x)**2 + (p3.y - p1.y)**2 + (p3.z - p1.z)**2
24-
c = (p1.x - p2.x)**2 + (p1.y - p2.y)**2 + (p1.z - p2.z)**2
40+
a = (x2 - x3) ** 2 + (y2 - y3) ** 2 + (z2 - z3) ** 2
41+
b = (x3 - x1) ** 2 + (y3 - y1) ** 2 + (z3 - z1) ** 2
42+
c = (x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2
2543

2644
pbc = a * (-a + b + c)
2745
apc = b * (a - b + c)
2846
abp = c * (a + b - c)
2947

30-
xv = (pbc * p1.x + apc * p2.x + abp * p3.x) / (pbc + apc + abp)
31-
yv = (pbc * p1.y + apc * p2.y + abp * p3.y) / (pbc + apc + abp)
32-
zv = (pbc * p1.z + apc * p2.z + abp * p3.z) / (pbc + apc + abp)
48+
denom = pbc + apc + abp
49+
xv = (pbc * x1 + apc * x2 + abp * x3) / denom
50+
yv = (pbc * y1 + apc * y2 + abp * y3) / denom
51+
zv = (pbc * z1 + apc * z2 + abp * z3) / denom
3352
else:
34-
d = 2 * (p1.x * (p2.y - p3.y) + p2.x *
35-
(p3.y - p1.y) + p3.x * (p1.y - p2.y))
53+
d = 2 * (x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2))
3654

37-
xv = ((p1.x**2 + p1.y**2) * (p2.y - p3.y) + (p2.x**2 + p2.y**2)
38-
* (p3.y - p1.y) + (p3.x**2 + p3.y**2) * (p1.y - p2.y)) / d
39-
yv = ((p1.x**2 + p1.y**2) * (p3.x - p2.x) + (p2.x**2 + p2.y**2)
40-
* (p1.x - p3.x) + (p3.x**2 + p3.y**2) * (p2.x - p1.x)) / d
41-
zv = 0.0
55+
xv = (
56+
(x1**2 + y1**2) * (y2 - y3)
57+
+ (x2**2 + y2**2) * (y3 - y1)
58+
+ (x3**2 + y3**2) * (y1 - y2)
59+
) / d
60+
yv = (
61+
(x1**2 + y1**2) * (x3 - x2)
62+
+ (x2**2 + y2**2) * (x1 - x3)
63+
+ (x3**2 + y3**2) * (x2 - x1)
64+
) / d
65+
zv = np.zeros_like(xv)
4266

43-
# Optional method to use barycenter instead.
44-
# xv = p1.x + p2.x + p3.x
45-
# xv = xv / 3.0
46-
# yv = p1.y + p2.y + p3.y
47-
# yv = yv / 3.0
48-
return point(xv, yv, zv)
67+
return xv, yv, zv
4968

5069

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

7089
lon = np.deg2rad(lon)
7190
lat = np.deg2rad(lat)
72-
x = R*np.multiply(np.cos(lon), np.cos(lat))
73-
y = R*np.multiply(np.sin(lon), np.cos(lat))
74-
z = R*np.sin(lat)
91+
x = R * np.multiply(np.cos(lon), np.cos(lat))
92+
y = R * np.multiply(np.sin(lon), np.cos(lat))
93+
z = R * np.sin(lat)
7594

7695
return x, y, z

0 commit comments

Comments
 (0)