Add a tool to create transect cell and edge fields with edge sign and distance

ocean/transects/python/README.md

@@ -0,0 +1,57 @@
+# Python Transect Tools
+
+## compute_transects.py
+
+Computes transport through sections.
+
+Example call:
+```
+  ./compute_transects.py
+  -k transect_masks.nc
+  -m MPAS_mesh.nc
+  -t 'RUN_PATH/analysis_members/timeSeriesStatsMonthly.*.nc'
+  -n 'all'
+```
+To create the `transect_masks.nc` file, load e3sm-unified and:
+```
+   MpasMaskCreator.x MPAS_mesh.nc  transect_masks.nc -f transect_definitions.geojson
+```
+where the `transect_definitions.geojson` file includes a sequence of lat/lon
+points for each transect.
+
+On LANL IC, example file is at
+```
+/usr/projects/climate/mpeterse/analysis_input_files/geojson_files/SingleRegionAtlanticWTransportTransects.geojson
+```
+
+## create_transect_masks.py
+
+Requires a conda environment with:
+* `python`
+* `geometric_features`
+* `matplotlib`
+* `mpas_tools`
+* `netcdf4`
+* `numpy`
+* `scipy`
+* `shapely`
+* `xarray`
+
+The tools creates cell and edge masks, distance along the transect of cells
+and edges in the mask, and the edge sign on edges.  It also includes
+information (distance, cell and edge indices, interpolation weights, etc.)
+along the transect itself to aid plotting.
+
+The required inputs are an MPAS mesh file and a geojson file or the name of an
+ocean transect from `geometric_features`.  The required output is a filename
+with the masks and other information about the transect.
+
+## cut_closed_transect.py
+
+If a transect is a closed loop, the path of edges and edge signs don't work
+correctly (the shortest path between the beginning and end of the transect is
+trivial and involves a single edge). To avoid this, we provide a tool for
+cutting a square (in lat/lon space) out of the transect to sever the loop.
+The user provides a latitude and longitude (used to locate the closest point)
+on the transect and the size of the square to cut out.
+

ocean/transects/python/create_transect_masks.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+
+import argparse
+
+import numpy as np
+import xarray as xr
+from geometric_features import read_feature_collection
+from geometric_features import GeometricFeatures
+from mpas_tools.cime.constants import constants
+from mpas_tools.logging import LoggingContext
+from mpas_tools.io import write_netcdf
+from mpas_tools.mesh.mask import compute_mpas_transect_masks
+from mpas_tools.parallel import create_pool
+
+from transect.vert import compute_transect
+
+
+def combine_transect_datasets(ds_mesh, fc_transect, out_filename, pool,
+                              logger):
+    """
+    Combine transects masks on cells and edges with a dataset for plotting
+    that describes how the transect slices through cell and edge geometry.
+    Add fields on edges and cells that define the (mean) distance along the
+    transect for each cell or edge in the transect
+    """
+
+    earth_radius = constants['SHR_CONST_REARTH']
+
+    ds_mask = compute_mpas_transect_masks(dsMesh=ds_mesh, fcMask=fc_transect,
+                                          earthRadius=earth_radius,
+                                          maskTypes=('cell', 'edge',),
+                                          logger=logger,
+                                          pool=pool, addEdgeSign=True)
+
+    feature = fc_transect.features[0]
+    geom_type = feature['geometry']['type']
+    if geom_type == 'LineString':
+        coordinates = [feature['geometry']['coordinates']]
+    elif geom_type == 'MultiLineString':
+        coordinates = feature['geometry']['coordinates']
+    else:
+        raise ValueError(
+            f'Unexpected geometry type for the transect {geom_type}')
+
+    lon = []
+    lat = []
+    for coords in coordinates:
+        lon_local, lat_local = zip(*coords)
+        lon.extend(lon_local)
+        lat.extend(lat_local)
+    lon = xr.DataArray(data=lon, dims='nTransectPoints')
+    lat = xr.DataArray(data=lat, dims='nTransectPoints')
+
+    layer_thickness = ds_mesh.layerThickness
+    bottom_depth = ds_mesh.bottomDepth
+    min_level_cell = ds_mesh.minLevelCell
+    max_level_cell = ds_mesh.maxLevelCell
+
+    ds_transect = compute_transect(lon, lat, ds_mesh, layer_thickness,
+                                   bottom_depth, min_level_cell,
+                                   max_level_cell, spherical=True)
+
+    ds = ds_mask
+    for var in ds_transect.data_vars:
+        ds[var] = ds_transect[var]
+
+    add_distance_field(ds, logger)
+
+    write_netcdf(ds, out_filename)
+
+
+def add_distance_field(ds, logger):
+    """
+    Add fields on edges and cells that define the (mean) distance along the
+    transect for each cell or edge in the transect
+    """
+
+    dist_cell = np.zeros(ds.sizes['nCells'])
+    count_cell = np.zeros(ds.sizes['nCells'], dtype=int)
+    dist_edge = np.zeros(ds.sizes['nEdges'])
+    count_edge = np.zeros(ds.sizes['nEdges'], dtype=int)
+
+    logger.info('Adding transect distance fields on cells and edges...')
+
+    for segment in range(ds.sizes['nSegments']):
+        icell = ds.horizCellIndices.isel(nSegments=segment).values
+        iedge = ds.horizEdgeIndices.isel(nSegments=segment).values
+        # the distance for the midpoint of the segment is the mean
+        # of the distances of the end points
+        dist = 0.5 * (ds.dNode.isel(nHorizNodes=segment) +
+                      ds.dNode.isel(nHorizNodes=segment + 1))
+        dist_cell[icell] += dist
+        count_cell[icell] += 1
+        dist_edge[iedge] += dist
+        count_edge[iedge] += 1
+
+    mask = count_cell > 0
+    dist_cell[mask] /= count_cell[mask]
+    dist_cell[np.logical_not(mask)] = np.nan
+
+    mask = count_edge > 0
+    dist_edge[mask] /= count_edge[mask]
+    dist_edge[np.logical_not(mask)] = np.nan
+
+    ds['transectDistanceCell'] = ('nCells', dist_cell)
+    ds['transectDistanceEdge'] = ('nEdges', dist_edge)
+    logger.info('Done.')
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description='''
+        creates transect edge and cell masks along with edge sign and distance
+        along the transect''')
+    parser.add_argument('-m', dest='mesh_filename',
+                        help='MPAS-Ocean horizontal and vertical filename',
+                        required=True)
+    parser.add_argument('-g', dest='geojson_filename',
+                        help='Geojson filename with transect', required=False)
+    parser.add_argument('-f', dest='feature_name',
+                        help='Name of an ocean transect from '
+                             'geometric_features',
+                        required=False)
+    parser.add_argument('-o', dest='out_filename',
+                        help='Edge transect filename', required=True)
+    parser.add_argument(
+        "--process_count", required=False, dest="process_count", type=int,
+        help="The number of processes to use to compute masks.  The "
+             "default is to use all available cores")
+    parser.add_argument(
+        "--multiprocessing_method", dest="multiprocessing_method",
+        default='forkserver',
+        help="The multiprocessing method use for python mask creation "
+             "('fork', 'spawn' or 'forkserver')")
+    args = parser.parse_args()
+
+    if args.geojson_filename is None and args.feature_name is None:
+        raise ValueError('Must supply either a geojson file or a transect '
+                         'name')
+
+    if args.geojson_filename is not None:
+        fc_transect = read_feature_collection(args.geojson_filename)
+    else:
+        gf = GeometricFeatures()
+        fc_transect = gf.read(componentName='ocean', objectType='transect',
+                              featureNames=[args.feature_name])
+
+    ds_mesh = xr.open_dataset(args.mesh_filename)
+    if 'Time' in ds_mesh.dims:
+        ds_mesh = ds_mesh.isel(Time=0)
+
+    pool = create_pool(process_count=args.process_count,
+                       method=args.multiprocessing_method)
+
+    with LoggingContext('create_transect_masks') as logger:
+
+        combine_transect_datasets(ds_mesh, fc_transect, args.out_filename,
+                                  pool, logger)
+
+
+if __name__ == '__main__':
+    main()

ocean/transects/python/cut_closed_transect.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+
+import argparse
+
+import numpy as np
+from geometric_features import (
+    GeometricFeatures,
+    read_feature_collection
+)
+from shapely.geometry import (
+    mapping,
+    Polygon,
+    shape
+)
+
+
+def cut_transect(fc_transect, lat, lon, size, out_filename):
+    """
+    Cut a square out of the given closed-loop transect to break the loop.
+    """
+
+    # find the closest point in the transect to the specificed lat/lon
+
+    feature = fc_transect.features[0]
+    coordinates = feature['geometry']['coordinates']
+    feature_type = feature['geometry']['type']
+    if feature_type == 'LineString':
+        coordinates = [coordinates]
+    elif feature_type != 'MultiLineString':
+        raise ValueError(
+            f'Unexpected geometry type for transect {feature_type}')
+
+    min_dist = None
+    center_lon = None
+    center_lan = None
+    for coords in coordinates:
+        lon_local, lat_local = zip(*coords)
+        dist = np.sqrt((np.array(lon_local) - lon)**2 +
+                       (np.array(lat_local) - lat)**2)
+        index_min = np.argmin(dist)
+        if min_dist is None or dist[index_min] < min_dist:
+            center_lon = lon_local[index_min]
+            center_lan = lat_local[index_min]
+            min_dist = dist[index_min]
+
+    square = Polygon([(center_lon - 0.5 * size, center_lan - 0.5 * size),
+                      (center_lon - 0.5 * size, center_lan + 0.5 * size),
+                      (center_lon + 0.5 * size, center_lan + 0.5 * size),
+                      (center_lon + 0.5 * size, center_lan - 0.5 * size),
+                      (center_lon - 0.5 * size, center_lan - 0.5 * size)])
+
+    feature = fc_transect.features[0]
+    transect_shape = shape(feature['geometry'])
+    transect_shape = transect_shape.difference(square)
+
+    # now sort the coordinates so the start and end of the transect are at the
+    # dividing point
+
+    feature['geometry'] = mapping(transect_shape)
+
+    feature_type = feature['geometry']['type']
+    if feature_type == 'MultiLineString':
+        coordinates = feature['geometry']['coordinates']
+
+        # reorder the LineStrings so the first one starts right after the cut
+
+        closest = None
+        min_dist = None
+        for index, coords in enumerate(coordinates):
+            lon_first, lat_first = coords[0]
+            dist = np.sqrt((lon_first - lon)**2 + (lat_first - lat)**2)
+            if min_dist is None or dist < min_dist:
+                closest = index
+                min_dist = dist
+        new_coords = list(coordinates[closest:])
+        new_coords.extend(list(coordinates[:closest]))
+        feature['geometry']['coordinates'] = tuple(new_coords)
+
+    fc_transect.to_geojson(out_filename)
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description='''
+        cut the given transect loop as close as possible to the given
+        latitude and longitude''')
+    parser.add_argument('-g', dest='geojson_filename',
+                        help='Geojson filename with transect', required=False)
+    parser.add_argument('-f', dest='feature_name',
+                        help='Name of an ocean transect from '
+                             'geometric_features',
+                        required=False)
+    parser.add_argument('--lat', dest='lat', type=float,
+                        help='The approx. latitude at which to cut the loop',
+                        required=True)
+    parser.add_argument('--lon', dest='lon', type=float,
+                        help='The approx. longitude at which to cut the loop',
+                        required=True)
+    parser.add_argument('--size', dest='size', type=float,
+                        help='The size in degrees of the square used to cut '
+                             'the loop',
+                        required=True)
+    parser.add_argument('-o', dest='out_filename',
+                        help='The geojson file with the cut transect to write '
+                             'out',
+                        required=True)
+    args = parser.parse_args()
+
+    if args.geojson_filename is None and args.feature_name is None:
+        raise ValueError('Must supply either a geojson file or a transect '
+                         'name')
+
+    if args.geojson_filename is not None:
+        fc_transect = read_feature_collection(args.geojson_filename)
+    else:
+        gf = GeometricFeatures()
+        fc_transect = gf.read(componentName='ocean', objectType='transect',
+                              featureNames=[args.feature_name])
+
+    cut_transect(fc_transect, args.lat, args.lon, args.size, args.out_filename)
+
+
+if __name__ == '__main__':
+    main()

ocean/transects/python/transect/__init__.py

@@ -0,0 +1,3 @@
+"""
+Copied from Polaris
+"""