Skip to content

Commit

Permalink
Change add_land_locked_cells_to_mask to use xarray
Browse files Browse the repository at this point in the history
  • Loading branch information
@xylar
xylar committed May 6, 2019
1 parent 16dff5b commit f49a851
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 166 deletions.
240 changes: 98 additions & 142 deletions conda_package/mpas_tools/ocean/coastline_alteration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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?
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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))
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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
Loading

0 comments on commit f49a851

Please sign in to comment.