Update jigsaw_to_netcdf() for efficiency and robustness

This merge:
* switches to xarray instead of netcdf4
* it does not hard-code `NetCDF3-Classic` format but instead uses
  `write_netcdf()` with the default format.  This will be important
  for meshes with large numbers of vertices.
* it uses numpy arrays directly rather than looping over vertices

Some formatting has been cleaned up for better consistency with
other parts of `mpas_tools`.

Author: xylar

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,123 @@ 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), \
         '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)
     parser.add_argument(
-        "-m",
-        "--msh",
-        dest="msh",
-        required=True,
-        help="input .msh file generated by JIGSAW.",
-        metavar="FILE")
+        '-m', '--msh', dest='msh', required=True,
+        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",
-        type=float,
-        help="The radius of the sphere in meters")
+        '-r',
+        '--sphere_radius',
+        dest='sphere_radius', type=float,
+        help='The radius of the sphere in meters')
 
     args = parser.parse_args()
-
     jigsaw_to_netcdf(args.msh, args.output, args.spherical, args.sphere_radius)