Make MOC basin masks with southern transects

xylar · xylar · commit cd11e1f40daa · 2020-03-03T11:07:53.000+01:00
Update the script for extracting southern countours accordingly.
diff --git a/conda_package/mpas_tools/ocean/moc.py b/conda_package/mpas_tools/ocean/moc.py
@@ -3,37 +3,188 @@
 
 import xarray
 import numpy
-
-
-def add_moc_southern_boundary_transects(in_filename, mesh_filename,
-                                        out_filename):
+import logging
+import sys
+import copy
+import shapely.ops
+import shapely.geometry
+from geometric_features import GeometricFeatures, FeatureCollection
+
+import mpas_tools.conversion
+from mpas_tools.io import write_netcdf
+
+
+def make_moc_basins_and_transects(gf, mesh_filename,
+                                  mask_and_transect_filename,
+                                  geojson_filename=None,
+                                  mask_filename=None,
+                                  logger=None):
     """
+    Builds features defining the ocean basins and southern transects used in
+    computing the meridional overturning circulation (MOC)
     Parameters
     ----------
-    in_filename : str
-        A file containing MOC region masks
+    gf : ``GeometricFeatures``
+        An object that knows how to download and read geometric featuers
 
     mesh_filename : str
         A file with MPAS mesh information
 
-    out_filename : str
+    mask_and_transect_filename : str
         A file to write the MOC region masks and southern-boundary transects to
+
+    geojson_filename : str, optional
+        A file to write MOC regions to
+
+    mask_filename : str, optional
+        A file to write MOC region masks to
+
+    logger : ``logging.Logger``, optional
+        A logger for the output if not stdout
+
+    Returns
+    -------
+    fc : ``FeatureCollection``
+        The new feature collection
+    """
+    # Authors
+    # -------
+    # Xylar Asay-Davis
+
+    fcMOC = build_moc_basins(gf, logger)
+
+    if geojson_filename is not None:
+        fcMOC.to_geojson(geojson_filename)
+
+    dsMesh = xarray.open_dataset(mesh_filename)
+    dsMasks = mpas_tools.conversion.mask(dsMesh=dsMesh, fcMask=fcMOC,
+                                         logger=logger)
+
+    if mask_filename is not None:
+        write_netcdf(dsMasks, mask_filename)
+
+    dsMasksAndTransects = add_moc_southern_boundary_transects(dsMasks, dsMesh,
+                                                              logger=logger)
+    write_netcdf(dsMasksAndTransects, mask_and_transect_filename)
+
+
+def build_moc_basins(gf, logger=None):
     """
-    mocMask = xarray.open_dataset(in_filename)
-    mesh = xarray.open_dataset(mesh_filename)
+    Builds features defining the ocean basins used in computing the meridional
+    overturning circulation (MOC)
+    Parameters
+    ----------
+    gf : ``GeometricFeatures``
+        An object that knows how to download and read geometric featuers
+
+    logger : ``logging.Logger``, optional
+        A logger for the output if not stdout
+
+    Returns
+    -------
+    fc : ``FeatureCollection``
+        The new feature collection
+    """
+    # Authors
+    # -------
+    # Xylar Asay-Davis
+
+    useStdout = logger is None
+    if useStdout:
+        logger = logging.getLogger()
+        logger.addHandler(logging.StreamHandler(sys.stdout))
+        logger.setLevel(logging.INFO)
+
+    MOCSubBasins = {'Atlantic': ['Atlantic', 'Mediterranean'],
+                    'IndoPacific': ['Pacific', 'Indian'],
+                    'Pacific': ['Pacific'],
+                    'Indian': ['Indian']}
+
+    MOCSouthernBoundary = {'Atlantic': '34S',
+                           'IndoPacific': '34S',
+                           'Pacific': '6S',
+                           'Indian': '6S'}
+
+    fc = FeatureCollection()
+    fc.set_group_name(groupName='MOCBasinRegionsGroup')
+
+    # build MOC basins from regions with the appropriate tags
+    for basinName in MOCSubBasins:
+        logger.info('{} MOC'.format(basinName))
+
+        logger.info(' * merging features')
+        tags = ['{}_Basin'.format(basin) for basin in
+                MOCSubBasins[basinName]]
+
+        fcBasin = gf.read(componentName='ocean', objectType='region',
+                          tags=tags, allTags=False)
+
+        logger.info(' * combining features')
+        fcBasin = fcBasin.combine(featureName='{}_MOC'.format(basinName))
 
-    southernBoundaryEdges, southernBounderyEdgeSigns, \
+        logger.info(' * masking out features south of MOC region')
+        maskName = 'MOC mask {}'.format(MOCSouthernBoundary[basinName])
+        fcMask = gf.read(componentName='ocean', objectType='region',
+                         featureNames=[maskName])
+        # mask out the region covered by the mask
+        fcBasin = fcBasin.difference(fcMask)
+
+        # remove various small polygons that are not part of the main MOC
+        # basin shapes.  Most are tiny but one below Australia is about 20
+        # deg^2, so make the threshold 100 deg^2 to be on the safe side.
+        fcBasin = _remove_small_polygons(fcBasin, minArea=100., logger=logger)
+
+        # add this basin to the full feature collection
+        fc.merge(fcBasin)
+
+    if useStdout:
+        logger.handlers = []
+
+    return fc
+
+
+def add_moc_southern_boundary_transects(dsMask, dsMesh, logger=None):
+    """
+    Parameters
+    ----------
+    dsMask : ``xarray.Dataset``
+        Region masks defining MOC basins
+
+    dsMesh : ``xarray.Dataset``, optional
+        An MPAS mesh on which the masks should be created
+
+    logger : ``logging.Logger``, optional
+        A logger for the output if not stdout
+
+    Returns
+    -------
+    dsMask : ``xarray.Dataset``
+        Region masks defining MOC basins and the corresponding southern-boundary
+        transects
+    """
+
+    useStdout = logger is None
+    if useStdout:
+        logger = logging.getLogger()
+        logger.addHandler(logging.StreamHandler(sys.stdout))
+        logger.setLevel(logging.INFO)
+
+    southernBoundaryEdges, southernBoiundaryEdgeSigns, \
         southernBoundaryVertices = \
-        _extract_southern_boundary(mesh, mocMask, latBuffer=3.*numpy.pi/180.)
+        _extract_southern_boundary(dsMesh, dsMask, latBuffer=3.*numpy.pi/180.,
+                                   logger=logger)
 
-    _add_transects_to_moc(mesh, mocMask, southernBoundaryEdges,
-                          southernBounderyEdgeSigns,
+    _add_transects_to_moc(dsMesh, dsMask, southernBoundaryEdges,
+                          southernBoiundaryEdgeSigns,
                           southernBoundaryVertices)
 
-    mocMask.to_netcdf(out_filename)
+    if useStdout:
+        logger.handlers = []
 
+    return dsMask
 
-def _extract_southern_boundary(mesh, mocMask, latBuffer):
+
+def _extract_southern_boundary(mesh, mocMask, latBuffer, logger):
     # Extrcts the southern boundary of each region mask in mocMask.  Mesh info
     # is taken from mesh.  latBuffer is a number of radians above the southern-
     # most point that should be considered to definitely be in the southern
@@ -56,12 +207,12 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
                                            cellsOnEdge < nCells)
 
     southernBoundaryEdges = []
-    southernBounderyEdgeSigns = []
+    southernBoiundaryEdgeSigns = []
     southernBoundaryVertices = []
 
     for iRegion in range(nRegions):
         name = mocMask.regionNames[iRegion].values.astype('U')
-        print(name)
+        logger.info(name)
         cellMask = mocMask.variables['regionCellMasks'][:, iRegion].values
 
         # land cells are outside not in the MOC region
@@ -70,7 +221,7 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
         cellsOnEdgeMask[cellsOnEdgeInRange] = \
             cellMask[cellsOnEdge[cellsOnEdgeInRange]] == 1
 
-        print('  computing edge sign...')
+        logger.info('  computing edge sign...')
         edgeSign = numpy.zeros(nEdges)
         # positive sign if the first cell on edge is in the region
         mask = numpy.logical_and(cellsOnEdgeMask[:, 0],
@@ -82,13 +233,13 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
         edgeSign[mask] = 1.
         isMOCBoundaryEdge = edgeSign != 0.
         edgesMOCBoundary = numpy.arange(nEdges)[isMOCBoundaryEdge]
-        print('  done.')
+        logger.info('  done.')
 
         startEdge = numpy.argmin(latEdge[isMOCBoundaryEdge])
         startEdge = edgesMOCBoundary[startEdge]
         minLat = latEdge[startEdge]
 
-        print('  getting edge sequence...')
+        logger.info('  getting edge sequence...')
         # follow the boundary from this point to get a loop of edges
         # Note: it is possible but unlikely that the southern-most point is
         # not within bulk region of the MOC mask if the region is not a single
@@ -97,7 +248,7 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
             _get_edge_sequence_on_boundary(startEdge, edgeSign, edgesOnVertex,
                                            verticesOnEdge)
 
-        print('  done: {} edges in transect.'.format(len(edgeSequence)))
+        logger.info('  done: {} edges in transect.'.format(len(edgeSequence)))
 
         aboveSouthernBoundary = latEdge[edgeSequence] > minLat + latBuffer
 
@@ -121,7 +272,7 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
         if len(startIndices) == 0:
             # the whole sequence is the southern boundary
             southernBoundaryEdges.append(edgeSequence)
-            southernBounderyEdgeSigns.append(edgeSequenceSigns)
+            southernBoiundaryEdgeSigns.append(edgeSequenceSigns)
             southernBoundaryVertices.append(vertexSequence)
             continue
 
@@ -136,7 +287,7 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
         indices = numpy.mod(indices, len(edgeSequence))
 
         southernBoundaryEdges.append(edgeSequence[indices])
-        southernBounderyEdgeSigns.append(edgeSequenceSigns[indices])
+        southernBoiundaryEdgeSigns.append(edgeSequenceSigns[indices])
 
         # we want one extra vertex in the vertex sequence
         indices = numpy.arange(endIndices[longest],
@@ -145,12 +296,12 @@ def _extract_southern_boundary(mesh, mocMask, latBuffer):
 
         southernBoundaryVertices.append(vertexSequence[indices])
 
-    return (southernBoundaryEdges, southernBounderyEdgeSigns,
+    return (southernBoundaryEdges, southernBoiundaryEdgeSigns,
             southernBoundaryVertices)
 
 
 def _add_transects_to_moc(mesh, mocMask, southernBoundaryEdges,
-                          southernBounderyEdgeSigns, southernBoundaryVertices):
+                          southernBoiundaryEdgeSigns, southernBoundaryVertices):
     # Creates transect fields in mocMask from the edges, edge signs and
     # vertices defining the southern boundaries.  Mesh info (nEdges and
     # nVertices) is taken from the mesh file.
@@ -182,7 +333,7 @@ def _add_transects_to_moc(mesh, mocMask, southernBoundaryEdges,
         transectEdgeMasks[southernBoundaryEdges[iTransect], iTransect] = 1
 
         transectEdgeMaskSigns[southernBoundaryEdges[iTransect], iTransect] \
-            = southernBounderyEdgeSigns[iTransect]
+            = southernBoiundaryEdgeSigns[iTransect]
 
         transectCount = len(southernBoundaryEdges[iTransect])
         transectEdgeGlobalIDs[iTransect, 0:transectCount] \
@@ -267,3 +418,61 @@ def _get_edge_sequence_on_boundary(startEdge, edgeSign, edgesOnVertex,
     vertexSequence = numpy.array(vertexSequence)
 
     return edgeSequence, edgeSequenceSigns, vertexSequence
+
+
+def _remove_small_polygons(fc, minArea, logger):
+    """
+    A helper function to remove small polygons from a feature collection
+    Parameters
+    ----------
+    fc : ``FeatureCollection``
+        The feature collection to remove polygons from
+    minArea : float
+        The minimum area (in square degrees) below which polygons should be
+        removed
+    Returns
+    -------
+    fcOut : ``FeatureCollection``
+        The new feature collection with small polygons removed
+    """
+    # Authors
+    # -------
+    # Xylar Asay-Davis
+
+    fcOut = FeatureCollection()
+
+    removedCount = 0
+    for feature in fc.features:
+        geom = feature['geometry']
+        add = False
+        if geom['type'] not in ['Polygon', 'MultiPolygon']:
+            # no area to check, so just add it
+            fcOut.add_feature(copy.deepcopy(feature))
+        else:
+            featureShape = shapely.geometry.shape(geom)
+            if featureShape.type == 'Polygon':
+                if featureShape.area > minArea:
+                    add = True
+                else:
+                    removedCount += 1
+            else:
+                # a MultiPolygon
+                outPolygons = []
+                for polygon in featureShape:
+                    if polygon.area > minArea:
+                        outPolygons.append(polygon)
+                    else:
+                        removedCount += 1
+                if len(outPolygons) > 0:
+                    outShape = shapely.ops.cascaded_union(outPolygons)
+                    feature['geometry'] = shapely.geometry.mapping(outShape)
+                    add = True
+        if add:
+            fcOut.add_feature(copy.deepcopy(feature))
+        else:
+            logger.info("{} has been removed.".format(
+                feature['properties']['name']))
+
+    logger.info(' * Removed {} small polygons'.format(removedCount))
+
+    return fcOut
diff --git a/ocean/moc_southern_boundary_extractor/moc_southern_boundary_extractor.py b/ocean/moc_southern_boundary_extractor/moc_southern_boundary_extractor.py
@@ -15,8 +15,11 @@
 from __future__ import absolute_import, division, print_function, \
     unicode_literals
 
+import xarray
 import argparse
 
+from mpas_tools.io import \
+    write_netcdf
 from mpas_tools.ocean.moc import \
     add_moc_southern_boundary_transects
 
@@ -38,5 +41,9 @@
                         required=True)
     args = parser.parse_args()
 
-    add_moc_southern_boundary_transects(args.in_file, args.mesh_file,
-                                        args.out_file)
+    dsMasks = xarray.open_dataset(args.in_file)
+    dsMesh = xarray.open_dataset(args.mesh_file)
+
+    dsMasksAndTransects = add_moc_southern_boundary_transects(dsMasks, dsMesh)
+
+    write_netcdf(dsMasksAndTransects, args.out_file)
