import argparse
import numpy as np
import xarray as xr
from mpas_tools.io import write_netcdf
from mpas_tools.mesh.creation.open_msh import readmsh
from mpas_tools.mesh.creation.util import circumcenter
[docs]
def jigsaw_to_netcdf(msh_filename, output_name, on_sphere, sphere_radius=None):
    """
    Converts mesh data defined in triangle format to NetCDF
    Parameters
    ----------
    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.
    """
    # Get dimensions
    # Get nCells
    msh = readmsh(msh_filename)
    nCells = msh['POINT'].shape[0]
    # always triangles with JIGSAW output
    vertexDegree = 3
    nVertices = msh['TRIA3'].shape[0]
    if vertexDegree != 3:
        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]
    if msh['NDIMS'] == 2:
        zCell_full = np.zeros(yCell_full.shape)
    elif msh['NDIMS'] == 3:
        zCell_full = msh['POINT'][:, 2]
    else:
        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!'
        )
    attrs = {}
    if on_sphere:
        attrs['on_a_sphere'] = 'YES'
        attrs['sphere_radius'] = sphere_radius
        # convert from km to meters
        xCell_full = xCell_full * 1e3
        yCell_full = yCell_full * 1e3
        zCell_full = zCell_full * 1e3
    else:
        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!'
    )
    # 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',
    )
    parser.add_argument(
        '-o',
        '--output',
        dest='output',
        default='grid.nc',
        help='output file name.',
        metavar='FILE',
    )
    parser.add_argument(
        '-s',
        '--spherical',
        dest='spherical',
        action='store_true',
        default=False,
        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',
    )
    args = parser.parse_args()
    jigsaw_to_netcdf(args.msh, args.output, args.spherical, args.sphere_radius)