Change add_land_locked_cells_to_mask to use xarray

Author: xylar

conda_package/mpas_tools/ocean/coastline_alteration.py

@@ -2,9 +2,7 @@
     unicode_literals
 
 import numpy
-from netCDF4 import Dataset
-import os
-import shutil
+import xarray
 
 
 def add_critical_land_blockages(dsMask, dsBlockages):
@@ -72,100 +70,88 @@ def widen_transect_edge_masks(dsMask, dsMesh, latitude_threshold=43.0):
     return dsMask
 
 
-def add_land_locked_cells_to_mask(input_mask_filename, output_mask_filename,
-                                  mesh_filename, latitude_threshold=43.0,
+def add_land_locked_cells_to_mask(dsMask, dsMesh, latitude_threshold=43.0,
                                   nSweeps=10):
     '''
     Find ocean cells that are land-locked, and alter the cell mask so that they
     are counted as land cells.
 
     Parameters
     ----------
-    input_mask_filename : str
-        Mask file that includes cell and edge masks.
+    dsMask : ``xarray.Dataset``
+        A land-mask data set
 
-    output_mask_filename : str
-        Mask file that includes cell and edge masks.
-
-    mesh_filename : str
-        MPAS Mesh filename.
+    dsMesh : ``xarray.Dataset``
+        MPAS Mesh data set
 
     latitude_threshold : float, optional
-        Minimum latitude, in degrees, for transect widening.
+        Minimum latitude, in degrees, for transect widening
 
     nSweeps : int, optional
-        Maximum number of sweeps to search for land-locked cells.
+        Maximum number of sweeps to search for land-locked cells
+
+    Returns
+    -------
+    dsMask : ``xarray.Dataset``
+        A copy of the land-mask data set with land-locked cells added to the
+        mask for the first region
     '''
 
-    # Obtain mesh variables
-    meshFile = Dataset(mesh_filename, "r")
-    nCells = len(meshFile.dimensions["nCells"])
-    cellsOnCell = meshFile.variables["cellsOnCell"][:, :] - 1
-    nEdgesOnCell = meshFile.variables["nEdgesOnCell"][:]
-    latCell = numpy.rad2deg(meshFile.variables["latCell"][:])
-    lonCell = numpy.rad2deg(meshFile.variables["lonCell"][:])
-    meshFile.close()
+    dsMask = xarray.Dataset(dsMask)
+    dsMesh = dsMesh.copy(deep=True)
+
+    landMask = dsMask.regionCellMasks.max(dim='nRegions') > 0
 
-    _remove_file(output_mask_filename)
-    shutil.copyfile(input_mask_filename, output_mask_filename)
+    dsMask['landMaskDiagnostic'] = xarray.where(landMask, 1, 0)
 
-    # Obtain original cell mask from input file
-    inputMaskFile = Dataset(input_mask_filename, "r")
-    regionCellMasks = inputMaskFile.variables["regionCellMasks"][:, :]
-    # set landMask = flattened regionCellMasks
-    landMask = numpy.amax(regionCellMasks, axis=1)
-    inputMaskFile.close()
+    print("Running add_land_locked_cells_to_mask.py.  Total number of cells: "
+          "{}".format(dsMesh.sizes['nCells']))
 
-    # Open output file
-    outputMaskFile = Dataset(output_mask_filename, "a")
-    landMaskDiagnostic = outputMaskFile.createVariable(
-        "landMaskDiagnostic", "i", dimensions=("nCells"))
+    cellsOnCell = dsMesh.cellsOnCell - 1
+    nEdgesOnCell = dsMesh.nEdgesOnCell
 
-    regionCellMasks = outputMaskFile['regionCellMasks']
+    nextCellsOnCell = cellsOnCell.copy(deep=True)
+    prevCellsOnCell = cellsOnCell.copy(deep=True)
+    for iEdgeOnCell in range(nextCellsOnCell.shape[1]):
+        iP1 = numpy.mod(iEdgeOnCell + 1, nEdgesOnCell)
+        nextCellsOnCell[:, iEdgeOnCell] = cellsOnCell[:, iP1]
+        iM1 = numpy.mod(iEdgeOnCell - 1, nEdgesOnCell)
+        prevCellsOnCell[:, iEdgeOnCell] = cellsOnCell[:, iM1]
 
-    print("Running add_land_locked_cells_to_mask.py.  Total number of cells: "
-          "{}".format(nCells))
+    dsMesh['cellsOnCell'] = cellsOnCell
+    dsMesh['nextCellsOnCell'] = nextCellsOnCell
+    dsMesh['prevCellsOnCell'] = prevCellsOnCell
+    dsMesh['latCell'] = numpy.rad2deg(dsMesh.latCell)
+    dsMesh['lonCell'] = numpy.rad2deg(dsMesh.lonCell)
 
     landMask, removable = _remove_cells_with_isolated_edges1(
-        landMask, landMaskDiagnostic, regionCellMasks, latCell, nEdgesOnCell,
-        cellsOnCell, nCells, latitude_threshold)
+        dsMask, dsMesh, landMask, latitude_threshold)
     landMask = _remove_cells_with_isolated_edges2(
-        landMask, landMaskDiagnostic, regionCellMasks, removable,
-        nEdgesOnCell, cellsOnCell, nCells, nSweeps)
-    oceanMask = _flood_fill(landMask, landMaskDiagnostic, removable, lonCell,
-                            latCell, nEdgesOnCell, cellsOnCell, nCells)
+        dsMask, dsMesh, landMask, removable, nSweeps)
+    oceanMask = _flood_fill(dsMask, dsMesh, landMask, removable)
     landMask = _revert_cells_with_connected_edges(
-        oceanMask, landMask, landMaskDiagnostic, regionCellMasks, removable,
-        nEdgesOnCell, cellsOnCell, nCells, nSweeps)
-    outputMaskFile.close()
+        dsMask, dsMesh, oceanMask, landMask, removable, nSweeps)
+
+    return dsMask
 
 
-def _remove_cells_with_isolated_edges1(landMask, landMaskDiagnostic,
-                                       regionCellMasks, latCell, nEdgesOnCell,
-                                       cellsOnCell, nCells,
+def _remove_cells_with_isolated_edges1(dsMask, dsMesh, landMask,
                                        latitude_threshold):
     print("Step 1: Searching for land-locked cells.  Remove cells that only "
           "have isolated active edges.")
 
-    landMaskNew = numpy.array(landMask)
-
-    landMaskDiagnostic[:] = landMask
-
-    removable = numpy.logical_and(numpy.abs(latCell) >= latitude_threshold,
-                                  landMask == 0)
-
-    nextCellsOnCell = numpy.zeros(cellsOnCell.shape, int)
-
-    for iEdgeOnCell in range(nextCellsOnCell.shape[1]):
-        iP1 = numpy.mod(iEdgeOnCell + 1, nEdgesOnCell)
-        nextCellsOnCell[:, iEdgeOnCell] = \
-            cellsOnCell[numpy.arange(nCells), iP1]
-
-    valid = numpy.logical_and(removable.reshape(nCells, 1),
-                              cellsOnCell >= 0)
+    landMaskNew = landMask.copy(deep=True)
 
     active = numpy.logical_not(landMask)
+    removable = numpy.logical_and(
+        numpy.abs(dsMesh.latCell) >= latitude_threshold, active)
+
+    cellsOnCell = dsMesh.cellsOnCell
+    valid = numpy.logical_and(removable, cellsOnCell >= 0)
     activeEdge = numpy.logical_and(valid, active[cellsOnCell])
+
+    nextCellsOnCell = dsMesh.nextCellsOnCell
+    valid = numpy.logical_and(removable, nextCellsOnCell >= 0)
     activeNextEdge = numpy.logical_and(valid, active[nextCellsOnCell])
 
     # which vertices have adjacent active edges on this cell?
@@ -178,50 +164,37 @@ def _remove_cells_with_isolated_edges1(landMask, landMaskDiagnostic,
     landMaskNew[noActiveAdjacentEdges] = 1
     landLockedCounter = numpy.count_nonzero(noActiveAdjacentEdges)
 
-    regionCellMasks[:, 0] = numpy.maximum(regionCellMasks[:, 0],
-                                          noActiveAdjacentEdges)
+    dsMask.regionCellMasks[:, 0] = numpy.maximum(dsMask.regionCellMasks[:, 0],
+                                                 1*noActiveAdjacentEdges)
 
-    landMaskDiagnostic[noActiveAdjacentEdges] = 2
+    dsMask.landMaskDiagnostic[noActiveAdjacentEdges] = 2
 
     print("  Number of landLocked cells: {}".format(landLockedCounter))
 
     return landMaskNew, removable
 
 
-def _remove_cells_with_isolated_edges2(landMask, landMaskDiagnostic,
-                                       regionCellMasks, removable,
-                                       nEdgesOnCell, cellsOnCell, nCells,
+def _remove_cells_with_isolated_edges2(dsMask, dsMesh, landMask, removable,
                                        nSweeps):
     print("Step 2: Searching for land-locked cells. Remove cells that have "
           "any isolated active edges.")
 
-    nextCellsOnCell = numpy.zeros(cellsOnCell.shape, int)
-    prevCellsOnCell = numpy.zeros(cellsOnCell.shape, int)
-
-    for iEdgeOnCell in range(nextCellsOnCell.shape[1]):
-        iP1 = numpy.mod(iEdgeOnCell + 1, nEdgesOnCell)
-        nextCellsOnCell[:, iEdgeOnCell] = \
-            cellsOnCell[numpy.arange(nCells), iP1]
-        iM1 = numpy.mod(iEdgeOnCell - 1, nEdgesOnCell)
-        prevCellsOnCell[:, iEdgeOnCell] = \
-            cellsOnCell[numpy.arange(nCells), iM1]
+    cellsOnCell = dsMesh.cellsOnCell
+    nextCellsOnCell = dsMesh.nextCellsOnCell
+    prevCellsOnCell = dsMesh.prevCellsOnCell
 
     for iSweep in range(nSweeps):
         landLockedCounter = 0
-        landMaskNew = numpy.array(landMask)
-
-        mask = numpy.logical_and(removable,
-                                 landMask == 0)
+        landMaskNew = landMask.copy(deep=True)
 
         active = numpy.logical_not(landMask)
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  cellsOnCell >= 0)
+        mask = numpy.logical_and(removable, active)
+
+        valid = numpy.logical_and(mask, cellsOnCell >= 0)
         activeEdge = numpy.logical_and(valid, active[cellsOnCell])
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  nextCellsOnCell >= 0)
+        valid = numpy.logical_and(mask, nextCellsOnCell >= 0)
         activeNextEdge = numpy.logical_and(valid, active[nextCellsOnCell])
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  prevCellsOnCell >= 0)
+        valid = numpy.logical_and(mask, prevCellsOnCell >= 0)
         activePrevEdge = numpy.logical_and(valid, active[prevCellsOnCell])
 
         # an edge is land-locked if it is active but neither neighbor is active
@@ -235,8 +208,8 @@ def _remove_cells_with_isolated_edges2(landMask, landMaskDiagnostic,
         landLockedCounter = numpy.count_nonzero(landLockedCells)
         if landLockedCounter > 0:
             landMaskNew[landLockedCells] = 1
-            regionCellMasks[landLockedCells, 0] = 1
-            landMaskDiagnostic[landLockedCells] = 3
+            dsMask.regionCellMasks[landLockedCells, 0] = 1
+            dsMask.landMaskDiagnostic[landLockedCells] = 3
 
         landMask = landMaskNew
         print("  Sweep: {} Number of landLocked cells removed: {}".format(
@@ -247,20 +220,21 @@ def _remove_cells_with_isolated_edges2(landMask, landMaskDiagnostic,
     return landMask
 
 
-def _flood_fill(landMask, landMaskDiagnostic, removable, lonCell, latCell,
-                nEdgesOnCell, cellsOnCell, nCells):
+def _flood_fill(dsMask, dsMesh, landMask, removable):
     print("Step 3: Perform flood fill, starting from open ocean.")
 
     # init flood fill to 0 for water, -1 for land, 1 for known open ocean
-    floodFill = -1*numpy.ones(nCells, dtype="i")
-    mask = numpy.logical_and(removable, landMask == 0)
-    floodFill[mask] = 0
+    floodFill = xarray.where(
+        numpy.logical_and(removable, numpy.logical_not(landMask)), 0, -1)
 
-    openOceanMask = numpy.zeros(nCells, bool)
+    latCell = dsMesh.latCell
+    lonCell = dsMesh.lonCell
+
+    cellsOnCell = dsMesh.cellsOnCell
 
     # North Pole
     mask = latCell > 84.0
-    openOceanMask = numpy.logical_or(openOceanMask, mask)
+    openOceanMask = mask
 
     # Arctic
     mask = numpy.logical_and(
@@ -312,20 +286,22 @@ def _flood_fill(landMask, landMaskDiagnostic, removable, lonCell, latCell,
     nFloodableCells = numpy.count_nonzero(floodFill == 0)
     print("  Initial number of flood cells: {}".format(nFloodableCells))
 
-    landMaskDiagnostic[floodFill == 1] = 5
+    dsMask.landMaskDiagnostic[floodFill == 1] = 5
 
     # sweep over neighbors of known open ocean points
-    for iSweep in range(0, nCells):
+    for iSweep in range(dsMesh.sizes['nCells']):
 
         newFloodCellsThisSweep = 0
         mask = floodFill == 0
+        cellIndices = numpy.nonzero(mask.values)[0]
         for iCellOnCell in range(cellsOnCell.shape[1]):
-            neighbors = cellsOnCell[:, iCellOnCell]
-            fill = numpy.logical_and(
-                mask,
-                numpy.logical_and(neighbors >= 0, floodFill[neighbors] == 1))
-            floodFill[fill] = 1
-            newFloodCellsThisSweep += numpy.count_nonzero(fill)
+            neighbors = cellsOnCell[cellIndices, iCellOnCell]
+            filledNeighbors = numpy.logical_and(neighbors >= 0,
+                                                floodFill[neighbors] == 1)
+            fillIndices = cellIndices[filledNeighbors.values]
+            if(len(fillIndices) > 0):
+                floodFill[fillIndices] = 1
+                newFloodCellsThisSweep += len(fillIndices)
 
         print("  Sweep {} new flood cells this sweep: {}".format(
             iSweep, newFloodCellsThisSweep))
@@ -340,42 +316,29 @@ def _flood_fill(landMask, landMaskDiagnostic, removable, lonCell, latCell,
     return oceanMask
 
 
-def _revert_cells_with_connected_edges(oceanMask, landMask, landMaskDiagnostic,
-                                       regionCellMasks, removable,
-                                       nEdgesOnCell, cellsOnCell, nCells,
-                                       nSweeps):
+def _revert_cells_with_connected_edges(dsMask, dsMesh, oceanMask, landMask,
+                                       removable, nSweeps):
     print("Step 4: Searching for land-locked cells, step 3: revert cells with "
           "connected active edges")
 
-    nextCellsOnCell = numpy.zeros(cellsOnCell.shape, int)
-    prevCellsOnCell = numpy.zeros(cellsOnCell.shape, int)
-
-    for iEdgeOnCell in range(nextCellsOnCell.shape[1]):
-        iP1 = numpy.mod(iEdgeOnCell + 1, nEdgesOnCell)
-        nextCellsOnCell[:, iEdgeOnCell] = \
-            cellsOnCell[numpy.arange(nCells), iP1]
-        iM1 = numpy.mod(iEdgeOnCell - 1, nEdgesOnCell)
-        prevCellsOnCell[:, iEdgeOnCell] = \
-            cellsOnCell[numpy.arange(nCells), iM1]
+    cellsOnCell = dsMesh.cellsOnCell
+    nextCellsOnCell = dsMesh.nextCellsOnCell
+    prevCellsOnCell = dsMesh.prevCellsOnCell
 
     for iSweep in range(nSweeps):
         landMaskNew = numpy.array(landMask)
 
         # only remove a cell that was added in Step 2,
         # _remove_cells_with_isolated_edges2
-        mask = numpy.logical_and(removable, landMaskDiagnostic[:] == 3)
+        mask = numpy.logical_and(removable, dsMask.landMaskDiagnostic == 3)
 
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  cellsOnCell >= 0)
+        notLand = numpy.logical_not(landMask)
+        valid = numpy.logical_and(mask, cellsOnCell >= 0)
         oceanEdge = numpy.logical_and(valid, oceanMask[cellsOnCell])
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  nextCellsOnCell >= 0)
-        activeNextEdge = numpy.logical_and(valid,
-                                           landMask[nextCellsOnCell] == 0)
-        valid = numpy.logical_and(mask.reshape(nCells, 1),
-                                  prevCellsOnCell >= 0)
-        activePrevEdge = numpy.logical_and(valid,
-                                           landMask[prevCellsOnCell] == 0)
+        valid = numpy.logical_and(mask, nextCellsOnCell >= 0)
+        activeNextEdge = numpy.logical_and(valid, notLand[nextCellsOnCell])
+        valid = numpy.logical_and(mask, prevCellsOnCell >= 0)
+        activePrevEdge = numpy.logical_and(valid, notLand[prevCellsOnCell])
 
         reactivate = numpy.any(
             numpy.logical_and(
@@ -385,9 +348,9 @@ def _revert_cells_with_connected_edges(oceanMask, landMask, landMaskDiagnostic,
         landLockedCounter = numpy.count_nonzero(reactivate)
         if landLockedCounter > 0:
             landMaskNew[reactivate] = 0
-            regionCellMasks[reactivate, 0] = 0
+            dsMask.regionCellMasks[reactivate, 0] = 0
             oceanMask[reactivate] = 1
-            landMaskDiagnostic[reactivate] = 4
+            dsMask.landMaskDiagnostic[reactivate] = 4
 
         landMask = landMaskNew
         print("  Sweep: {} Number of land-locked cells returned: {}".format(
@@ -396,10 +359,3 @@ def _revert_cells_with_connected_edges(oceanMask, landMask, landMaskDiagnostic,
             break
 
     return landMask
-
-
-def _remove_file(fileName):
-    try:
-        os.remove(fileName)
-    except OSError:
-        pass