def compute_barotropic_streamfunction(dsMesh, ds, folder): ''' compute the barotropic streamfunction for the given mesh and monthly-mean data set ''' bsfFileName = '{}/barotropicStreamfunction.nc'.format(folder) if file_complete(ds, bsfFileName): return bsfVertex = _compute_barotropic_streamfunction_vertex(dsMesh, ds) bsfCell = _compute_barotropic_streamfunction_cell(dsMesh, bsfVertex) dsBSF = xarray.Dataset() dsBSF['xtime_startMonthly'] = ds.xtime_startMonthly dsBSF['xtime_endMonthly'] = ds.xtime_endMonthly dsBSF['bsfVertex'] = bsfVertex dsBSF.bsfVertex.attrs['units'] = 'Sv' dsBSF.bsfVertex.attrs['description'] = 'barotropic streamfunction ' \ 'on vertices' dsBSF['bsfCell'] = bsfCell dsBSF.bsfCell.attrs['units'] = 'Sv' dsBSF.bsfCell.attrs['description'] = 'barotropic streamfunction ' \ 'on cells' dsBSF = dsBSF.transpose('Time', 'nCells', 'nVertices') write_netcdf(dsBSF, bsfFileName)
def _vertical_cumsum_horizontal_transport(ds, outFileName): ''' compute the cumsum in the vertical of the horizontal transport ''' if file_complete(ds, outFileName): return chunks = {'Time': 1, 'nInternalEdges': 32768} ds = ds.chunk(chunks) nTime = ds.sizes['Time'] nInternalEdges = ds.sizes['nInternalEdges'] nz = ds.sizes['nz'] transport = ds.transport.rename({'nzM1': 'nz'}) transportSumTop = xarray.DataArray(numpy.zeros((nTime, nInternalEdges, 1)), dims=('Time', 'nInternalEdges', 'nz')).chunk(chunks) # zeros on top and then the cumsum for the rest transportSum = xarray.concat( [transportSumTop, transport.cumsum(dim='nz')], dim='nz') # mask out locations on the output where no input-grid layers overlap # with either the output layer above or the one below mask = ds.mask.rename({'nzM1': 'nz'}) maskTop = mask.isel(nz=0) maskBot = mask.isel(nz=nz - 2) outMask = xarray.concat( [maskTop, numpy.logical_or(mask[:, 0:-1, :], mask[:, 1:, :]), maskBot], dim='nz') dsOut = xarray.Dataset() dsOut['xtime_startMonthly'] = ds.xtime_startMonthly dsOut['xtime_endMonthly'] = ds.xtime_endMonthly dsOut['z'] = ds.z dsOut['transportSum'] = transportSum dsOut['mask'] = outMask dsOut = dsOut.transpose('Time', 'nz', 'nInternalEdges') print('compute and caching vertical transport sum on z-level grid:') write_netcdf(dsOut, outFileName, progress=True)
def compute_haney_number(dsMesh, ds, folder): ''' compute the Haney number rx1 for each edge, and interpolate it to cells. ''' haneyFileName = '{}/haney.nc'.format(folder) if file_complete(ds, haneyFileName): return haneyEdge, haneyCell = _compute_haney_number(dsMesh, ds) dsHaney = xarray.Dataset() dsHaney['xtime_startMonthly'] = ds.xtime_startMonthly dsHaney['xtime_endMonthly'] = ds.xtime_endMonthly dsHaney['haneyEdge'] = haneyEdge dsHaney.haneyEdge.attrs['units'] = 'unitless' dsHaney.haneyEdge.attrs['description'] = 'Haney number on edges' dsHaney['haneyCell'] = haneyCell dsHaney.haneyCell.attrs['units'] = 'unitless' dsHaney.haneyCell.attrs['description'] = 'Haney number on cells' dsHaney = dsHaney.transpose('Time', 'nCells', 'nEdges', 'nVertLevels') write_netcdf(dsHaney, haneyFileName)
def _horizontally_bin_overturning_streamfunction(ds, dsMesh, x, osfFileName, cacheFileName): ''' bin and sum the vertically cumsummed horizontal transport on the z-level grid to get the OSF. ''' chunks = {'Time': 1} ds = ds.chunk(chunks) nTime = ds.sizes['Time'] nz = ds.sizes['nz'] nx = len(x) x = xarray.DataArray.from_dict({'dims': ('nx', ), 'data': x}) widgets = [ 'bin overturning streamfunction: ', progressbar.Percentage(), ' ', progressbar.Bar(), ' ', progressbar.ETA() ] bar = progressbar.ProgressBar(widgets=widgets, maxval=nx).start() # sum up the transport into the region bounded by each x value # on the output grid fileNames = [] for xIndex in range(nx): fileName = cacheFileName.replace('.nc', '_{}.nc'.format(xIndex)) fileNames.append(fileName) if file_complete(ds, fileName): continue cellMask = dsMesh.xCell >= x[xIndex] edgeIndices, edgeSigns = _compute_region_boundary_edges( dsMesh, cellMask) if len(edgeIndices) == 0: localOSF = numpy.nan * xarray.DataArray(numpy.ones((nTime, nz)), dims=('Time', 'nz')) else: # convert to Sv transportSum = 1e-6 * \ edgeSigns*ds.transportSum.isel(nInternalEdges=edgeIndices) localOSF = transportSum.sum(dim='nInternalEdges') localMask = ds.mask.isel(nInternalEdges=edgeIndices).sum( dim='nInternalEdges') > 0 localOSF = localOSF.where(localMask) dsOSF = xarray.Dataset() dsOSF['osf'] = localOSF write_netcdf(dsOSF, fileName, progress=False) bar.update(xIndex + 1) bar.finish() dsOSF = xarray.open_mfdataset(fileNames, concat_dim='nx') dsOSF['xtime_startMonthly'] = ds.xtime_startMonthly dsOSF['xtime_endMonthly'] = ds.xtime_endMonthly dsOSF['x'] = x dsOSF['z'] = ds.z dsOSF.osf.attrs['units'] = 'Sv' dsOSF.osf.attrs['description'] = 'overturning streamfunction ' dsOSF = dsOSF.transpose('Time', 'nz', 'nx') write_netcdf(dsOSF, osfFileName)
def _interpolate_horizontal_transport_zlevel(ds, z, outFileName): ''' interpolate the horizontal transport through edges onto a z-level grid. ''' if file_complete(ds, outFileName): return ds = ds.chunk({ 'Time': 1, 'nInternalEdges': None, 'nVertLevels': 1, 'nVertLevelsP1': 1 }) nz = len(z) z = xarray.DataArray.from_dict({'dims': ('nz', ), 'data': z}) # make sure we don't miss anything z[0] = max(z[0].values, ds.zInterfaceEdge.max()) z[-1] = min(z[-1].values, ds.zInterfaceEdge.min()) z0 = z[0:-1].rename({'nz': 'nzM1'}) z1 = z[1:].rename({'nz': 'nzM1'}) nTime = ds.sizes['Time'] nInternalEdges = ds.sizes['nInternalEdges'] nVertLevels = ds.sizes['nVertLevels'] widgets = [ 'interpolating tansport on z-level grid: ', progressbar.Percentage(), ' ', progressbar.Bar(), ' ', progressbar.ETA() ] bar = progressbar.ProgressBar(widgets=widgets, maxval=nTime).start() fileNames = [] for tIndex in range(nTime): fileName = outFileName.replace('.nc', '_{}.nc'.format(tIndex)) fileNames.append(fileName) if os.path.exists(fileName): continue outTransport = xarray.DataArray(numpy.zeros((nInternalEdges, nz - 1)), dims=('nInternalEdges', 'nzM1')) dzSum = xarray.DataArray(numpy.zeros((nInternalEdges, nz - 1)), dims=('nInternalEdges', 'nzM1')) dsIn = ds.isel(Time=tIndex) for inZIndex in range(nVertLevels): zTop = dsIn.zInterfaceEdge.isel(nVertLevelsP1=inZIndex) zBot = dsIn.zInterfaceEdge.isel(nVertLevelsP1=inZIndex + 1) inTransportPerDepth = \ dsIn.transportPerDepth.isel(nVertLevels=inZIndex) zt = numpy.minimum(zTop, z0) zb = numpy.maximum(zBot, z1) dz = numpy.maximum(zt - zb, 0.) outTransport = outTransport + dz * inTransportPerDepth dzSum = dzSum + dz outTransport.compute() dzSum.compute() dsOut = xarray.Dataset() dsOut['mask'] = dzSum > 0 dsOut['transport'] = outTransport dsOut['transportVertSum'] = outTransport.sum('nzM1') dsOut['transportVertSumCheck'] = \ dsIn.transportVertSum - dsOut.transportVertSum dsOut = dsOut.transpose('nzM1', 'nInternalEdges') write_netcdf(dsOut, fileName, progress=False) assert (numpy.abs(dsOut.transportVertSumCheck).max().values < 1e-9) bar.update(tIndex + 1) bar.finish() dsOut = xarray.open_mfdataset(fileNames, concat_dim='Time') dsOut['xtime_startMonthly'] = ds.xtime_startMonthly dsOut['xtime_endMonthly'] = ds.xtime_endMonthly dsOut['z'] = z dsOut = dsOut.transpose('Time', 'nzM1', 'nz', 'nInternalEdges') print('caching transport on z-level grid:') write_netcdf(dsOut, outFileName, progress=True)
def _compute_horizontal_transport_mpas(ds, dsMesh, outFileName): ''' compute the horizontal transport through edges on the native MPAS grid. ''' if file_complete(ds, outFileName): return nVertLevels = dsMesh.sizes['nVertLevels'] cellsOnEdge = dsMesh.cellsOnEdge - 1 maxLevelCell = dsMesh.maxLevelCell - 1 cell0 = cellsOnEdge[:, 0] cell1 = cellsOnEdge[:, 1] internalEdgeIndices = xarray.DataArray(numpy.nonzero( numpy.logical_and(cell0.values >= 0, cell1.values >= 0))[0], dims=('nInternalEdges', )) cell0 = cell0[internalEdgeIndices] cell1 = cell1[internalEdgeIndices] bottomDepth = dsMesh.bottomDepth maxLevelEdgeTop = maxLevelCell[cell0] mask = numpy.logical_or(cell0 == -1, maxLevelCell[cell1] < maxLevelEdgeTop) maxLevelEdgeTop[mask] = maxLevelCell[cell1][mask] nVertLevels = dsMesh.sizes['nVertLevels'] vertIndex = \ xarray.DataArray.from_dict({'dims': ('nVertLevels',), 'data': numpy.arange(nVertLevels)}) ds = ds.chunk({'Time': 1}) chunks = {'nInternalEdges': 1024} maxLevelEdgeTop = maxLevelEdgeTop.chunk(chunks) dvEdge = dsMesh.dvEdge[internalEdgeIndices].chunk(chunks) bottomDepthEdge = 0.5 * (bottomDepth[cell0] + bottomDepth[cell1]).chunk(chunks) chunks = {'Time': 1, 'nInternalEdges': 1024} normalVelocity = ds.timeMonthly_avg_normalVelocity.isel( nEdges=internalEdgeIndices).chunk(chunks) layerThickness = ds.timeMonthly_avg_layerThickness.chunk() layerThicknessEdge = 0.5 * (layerThickness.isel( nCells=cell0) + layerThickness.isel(nCells=cell1)).chunk(chunks) layerThicknessEdge = layerThicknessEdge.where(vertIndex <= maxLevelEdgeTop, other=0.) thicknessSum = layerThicknessEdge.sum(dim='nVertLevels') thicknessCumSum = layerThicknessEdge.cumsum(dim='nVertLevels') zSurface = thicknessSum - bottomDepthEdge zInterfaceEdge = -thicknessCumSum + zSurface zInterfaceEdge = xarray.concat([ zSurface.expand_dims(dim='nVertLevelsP1', axis=2), zInterfaceEdge.rename({'nVertLevels': 'nVertLevelsP1'}) ], dim='nVertLevelsP1') transportPerDepth = dvEdge * normalVelocity dsOut = xarray.Dataset() dsOut['xtime_startMonthly'] = ds.xtime_startMonthly dsOut['xtime_endMonthly'] = ds.xtime_endMonthly dsOut['zInterfaceEdge'] = zInterfaceEdge dsOut['layerThicknessEdge'] = layerThicknessEdge dsOut['transportPerDepth'] = transportPerDepth dsOut['transportVertSum'] = \ (transportPerDepth*layerThicknessEdge).sum(dim='nVertLevels') dsOut = dsOut.transpose('Time', 'nInternalEdges', 'nVertLevels', 'nVertLevelsP1') print('compute and caching transport on MPAS grid:') write_netcdf(dsOut, outFileName, progress=True)