Source code for mpas_tools.seaice.mask

from netCDF4 import Dataset
import numpy as np
import math

from mpas_tools.cime.constants import constants


[docs]def extend_seaice_mask(filenameMesh, filenamePresence, extendDistance, unitSphere=False): """ Add a field ``icePresenceExtended`` to ``filenamePresence`` if it doesn't already exist. This field is the ``icePresence`` field extended by a distance of ``extendDistance``. Parameters ---------- filenameMesh : str The filename of the MPAS-Seaice mesh filenamePresence : str A filename for a file containing an ``icePresence`` field to be extended and to which a new ``icePresenceExtended`` will be added extendDistance : float The distance in km to expand (no expansion is performed if ``extendDistance=0.0``) unitSphere : bool, optional Whether the mesh is provided on the unit sphere, rather than one with the radius of the Earth """ # mesh print("Load mesh...") fileMesh = Dataset(filenameMesh, "r") nCells = len(fileMesh.dimensions["nCells"]) nEdgesOnCell = fileMesh.variables["nEdgesOnCell"][:] cellsOnCell = fileMesh.variables["cellsOnCell"][:] cellsOnCell[:] = cellsOnCell[:] - 1 xCell = fileMesh.variables["xCell"][:] yCell = fileMesh.variables["yCell"][:] zCell = fileMesh.variables["zCell"][:] fileMesh.close() # presence print("Load ice presence...") filePresence = Dataset(filenamePresence, "r") if "icePresenceExtended" in filePresence.variables: # we're already done, probably from a previous call filePresence.close() return icePresence = filePresence.variables["icePresence"][:] filePresence.close() # ice edge cells print("Get ice edge cells...") iceEdgeCell = np.zeros(nCells, dtype="i") for iCell in range(0, nCells): # if (iCell % 100000 == 0): # print(iCell, " of ", nCells, " cells...") for iCellOnCell in range(0, nEdgesOnCell[iCell]): iCell2 = cellsOnCell[iCell, iCellOnCell] if icePresence[iCell] == 1 and icePresence[iCell2] == 0: iceEdgeCell[iCell] = 1 # only edge cells nEdgeCells = np.sum(iceEdgeCell) print("nEdgeCells: ", nEdgeCells) print("Get edge cell vector...") iCellEdge = np.zeros(nEdgeCells, dtype="i") iEdgeCell = 0 for iCell in range(0, nCells): if iceEdgeCell[iCell] == 1: iCellEdge[iEdgeCell] = iCell iEdgeCell = iEdgeCell + 1 del iceEdgeCell # get edge positions print("Get edge positions...") xCellEdgeCell = np.zeros(nEdgeCells) yCellEdgeCell = np.zeros(nEdgeCells) zCellEdgeCell = np.zeros(nEdgeCells) for iEdgeCell in range(0, nEdgeCells): iCell = iCellEdge[iEdgeCell] xCellEdgeCell[iEdgeCell] = xCell[iCell] yCellEdgeCell[iEdgeCell] = yCell[iCell] zCellEdgeCell[iEdgeCell] = zCell[iCell] # find extended mask print("Find new extended mask...") extendDistance = extendDistance * 1000.0 # convert from km to m if unitSphere: earthRadius = constants['SHR_CONST_REARTH'] extendDistance = extendDistance / earthRadius distanceLimit = math.pow(extendDistance, 2) icePresenceNew = icePresence.copy() if extendDistance > 0.0: distances = np.zeros(nEdgeCells) for iCell in range(0, nCells): if iCell % 100000 == 0: print(iCell, " of ", nCells, " cells...") distances = np.multiply( (xCell[iCell] - xCellEdgeCell), (xCell[iCell] - xCellEdgeCell)) + np.multiply( (yCell[iCell] - yCellEdgeCell), (yCell[iCell] - yCellEdgeCell)) + np.multiply( (zCell[iCell] - zCellEdgeCell), (zCell[iCell] - zCellEdgeCell)) if distances[distances.argmin()] <= distanceLimit: icePresenceNew[iCell] = 1 print("Load ice presence...") filePresence = Dataset(filenamePresence, "a") icePresenceVar = filePresence.createVariable( "icePresenceExtended", "d", dimensions=["nCells"]) icePresenceVar[:] = icePresenceNew[:] filePresence.close()