diff --git a/conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py b/conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py index 062f28d78..75664e72d 100644 --- a/conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py +++ b/conda_package/mpas_tools/mesh/creation/jigsaw_to_netcdf.py @@ -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): """ @@ -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') # 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) diff --git a/conda_package/mpas_tools/mesh/creation/util.py b/conda_package/mpas_tools/mesh/creation/util.py index f8488ae4a..b81bec2d7 100644 --- a/conda_package/mpas_tools/mesh/creation/util.py +++ b/conda_package/mpas_tools/mesh/creation/util.py @@ -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): @@ -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 diff --git a/conda_package/tests/test_mesh_creation_util.py b/conda_package/tests/test_mesh_creation_util.py new file mode 100644 index 000000000..fa5e66108 --- /dev/null +++ b/conda_package/tests/test_mesh_creation_util.py @@ -0,0 +1,150 @@ +import collections + +import numpy as np + +from mpas_tools.mesh.creation.util import circumcenter + +point = collections.namedtuple('Point', ['x', 'y', 'z']) + + +def test_circumcenter_planar(): + # List of test cases: (cell0, cell1, cell2, expected) + test_cases = [ + # From a regular hex mesh + ( + (0.0, 17320.50807569, 0.0), + (10000.0, 17320.50807569, 0.0), + (5000.0, 25980.76211353, 0.0), + (5000.0, 20207.25942164, 0.0), + ), + # From an Antarctic mesh + ( + (2194652.05105505, 1144816.29348609, 0.0), + (2198726.62725077, 1145337.65538494, 0.0), + (2195474.22967839, 1148983.95766333, 0.0), + (2196492.12915069, 1146618.22089598, 0.0), + ), + # From an Antarctic mesh + ( + (2195474.22967839, 1148983.95766333, 0.0), + (2198726.62725077, 1145337.65538494, 0.0), + (2199279.33328257, 1149240.68069275, 0.0), + (2197485.28573298, 1147504.08822421, 0.0), + ), + # From an Antarctic mesh + ( + (2150629.6663189, 1441451.4899963, 0.0), + (2149349.39244275, 1445285.59533973, 0.0), + (2146912.55160507, 1442187.732174, 0.0), + (2149013.33940856, 1443042.57600619, 0.0), + ), + # From an Antarctic mesh + ( + (2570868.91779026, -97880.09977304, 0.0), + (2575332.71614832, -97980.66509542, 0.0), + (2573585.71564672, -95138.63592864, 0.0), + (2573113.05578246, -97387.13758518, 0.0), + ), + ] + + for i, (cell0, cell1, cell2, expected) in enumerate(test_cases): + x1, y1, z1 = cell0 + x2, y2, z2 = cell1 + x3, y3, z3 = cell2 + + xv, yv, zv = circumcenter(False, x1, y1, z1, x2, y2, z2, x3, y3, z3) + assert np.allclose([xv, yv, zv], expected, atol=1e-8), ( + f'Planar circumcenter failed for case {i}: got {[xv, yv, zv]}, ' + f'expected {expected}' + ) + + xv_old, yv_old, zv_old = _old_circumcenter( + False, x1, y1, z1, x2, y2, z2, x3, y3, z3 + ) + print(f'Planar circumcenter (case {i}):') + print(' circumcenter: ', (xv, yv, zv)) + print(' _old_circumcenter: ', (xv_old, yv_old, zv_old)) + print(' difference: ', (xv - xv_old, yv - yv_old, zv - zv_old)) + assert np.array_equal([xv, yv, zv], [xv_old, yv_old, zv_old]), ( + f'circumcenter and _old_circumcenter results differ in case {i}: ' + f'({xv}, {yv}, {zv}) != ({xv_old}, {yv_old}, {zv_old})' + ) + + +def test_circumcenter_sphere(): + # Use provided cell coordinates and vertex as expected output + x1, y1, z1 = -3590.95704026, -684.52527287, -6371218.95125671 + x2, y2, z2 = -9926.94769501, 18734.26792821, -6371184.72274307 + x3, y3, z3 = 12921.92507847, 11340.69768405, -6371196.8028643 + expected = (296.83807146, 11327.05862805, -6371209.92418473) + + xv, yv, zv = circumcenter(True, x1, y1, z1, x2, y2, z2, x3, y3, z3) + assert np.allclose([xv, yv, zv], expected, atol=1e-8), ( + f'Spherical circumcenter failed: got {[xv, yv, zv]}, expected ' + f'{expected}' + ) + + xv_old, yv_old, zv_old = _old_circumcenter( + True, x1, y1, z1, x2, y2, z2, x3, y3, z3 + ) + # Print both sets and their difference + print('Spherical circumcenter:') + print(' circumcenter: ', (xv, yv, zv)) + print(' _old_circumcenter: ', (xv_old, yv_old, zv_old)) + print(' difference: ', (xv - xv_old, yv - yv_old, zv - zv_old)) + # Check for exact equality + assert np.array_equal([xv, yv, zv], [xv_old, yv_old, zv_old]), ( + f'circumcenter and _old_circumcenter results differ: ' + f'({xv}, {yv}, {zv}) != ({xv_old}, {yv_old}, {zv_old})' + ) + + +def _old_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. + + Returns + ------- + center : point + The circumcenter of the triangle with x, y and z attributes + + """ + p1 = point(x1, y1, z1) + p2 = point(x2, y2, z2) + p3 = point(x3, y3, 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 + + 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) + else: + d = 2 * ( + p1.x * (p2.y - p3.y) + p2.x * (p3.y - p1.y) + p3.x * (p1.y - p2.y) + ) + + 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 + + # 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)