def _write_mpas_t_s(self, config): # {{{ climatologyName = 'TS_{}_{}'.format(self.prefix, self.season) outFileName = get_masked_mpas_climatology_file_name(config, self.season, self.componentName, climatologyName, op='avg') if os.path.exists(outFileName): ds = xarray.open_dataset(outFileName) zmin, zmax = ds.zbounds.values return zmin, zmax with dask.config.set(schedular='threads', pool=ThreadPool(self.daskThreads)): self.logger.info(' Extracting T and S in the region...') sectionName = self.sectionName cellsChunk = 32768 chunk = {'nCells': cellsChunk} try: restartFileName = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one' ' restart file to plot T-S diagrams') dsRestart = xarray.open_dataset(restartFileName) dsRestart = dsRestart.isel(Time=0).chunk(chunk) regionMaskFileName = self.mpasMasksSubtask.maskFileName dsRegionMask = xarray.open_dataset(regionMaskFileName) maskRegionNames = decode_strings(dsRegionMask.regionNames) regionIndex = maskRegionNames.index(self.regionName) dsMask = dsRegionMask.isel(nRegions=regionIndex).chunk(chunk) cellMask = dsMask.regionCellMasks == 1 if 'landIceMask' in dsRestart: # only the region outside of ice-shelf cavities cellMask = numpy.logical_and(cellMask, dsRestart.landIceMask == 0) if config.has_option(sectionName, 'zmin'): zmin = config.getfloat(sectionName, 'zmin') else: if 'zminRegions' in dsMask: zmin = dsMask.zminRegions.values else: # the old naming convention, used in some pre-generated # mask files zmin = dsMask.zmin.values if config.has_option(sectionName, 'zmax'): zmax = config.getfloat(sectionName, 'zmax') else: if 'zmaxRegions' in dsMask: zmax = dsMask.zmaxRegions.values else: # the old naming convention, used in some pre-generated # mask files zmax = dsMask.zmax.values inFileName = get_unmasked_mpas_climatology_file_name( config, self.season, self.componentName, op='avg') ds = xarray.open_dataset(inFileName) variableList = [ 'timeMonthly_avg_activeTracers_temperature', 'timeMonthly_avg_activeTracers_salinity', 'timeMonthly_avg_layerThickness' ] ds = ds[variableList] ds['zMid'] = compute_zmid(dsRestart.bottomDepth, dsRestart.maxLevelCell, dsRestart.layerThickness) ds['volume'] = (dsRestart.areaCell * ds['timeMonthly_avg_layerThickness']) ds = ds.where(cellMask, drop=True) self.logger.info("Don't worry about the following dask " "warnings.") depthMask = numpy.logical_and(ds.zMid >= zmin, ds.zMid <= zmax) depthMask.compute() self.logger.info("Dask warnings should be done.") ds['depthMask'] = depthMask for var in variableList: ds[var] = ds[var].where(depthMask) T = ds['timeMonthly_avg_activeTracers_temperature'].values.ravel() mask = numpy.isfinite(T) T = T[mask] S = ds['timeMonthly_avg_activeTracers_salinity'].values.ravel() S = S[mask] zMid = ds['zMid'].values.ravel()[mask] volume = ds['volume'].values.ravel()[mask] dsOut = xarray.Dataset() dsOut['T'] = ('nPoints', T) dsOut['S'] = ('nPoints', S) dsOut['z'] = ('nPoints', zMid) dsOut['volume'] = ('nPoints', volume) dsOut['zbounds'] = ('nBounds', [zmin, zmax]) write_netcdf(dsOut, outFileName) return zmin, zmax # }}}
def _write_obs_t_s(self, obsDict, zmin, zmax): # {{{ obsSection = '{}Observations'.format(self.componentName) climatologyDirectory = build_config_full_path( config=self.config, section='output', relativePathOption='climatologySubdirectory', relativePathSection=obsSection) outFileName = '{}/TS_{}_{}_{}.nc'.format(climatologyDirectory, obsDict['suffix'], self.prefix, self.season) if os.path.exists(outFileName): return with dask.config.set(schedular='threads', pool=ThreadPool(self.daskThreads)): chunk = {obsDict['latVar']: 400, obsDict['lonVar']: 400} regionMaskFileName = obsDict['maskTask'].maskFileName dsRegionMask = \ xarray.open_dataset(regionMaskFileName).chunk(chunk).stack( nCells=(obsDict['latVar'], obsDict['lonVar'])) dsRegionMask = dsRegionMask.reset_index('nCells').drop_vars( [obsDict['latVar'], obsDict['lonVar']]) maskRegionNames = decode_strings(dsRegionMask.regionNames) regionIndex = maskRegionNames.index(self.regionName) dsMask = dsRegionMask.isel(nRegions=regionIndex) cellMask = dsMask.regionCellMasks == 1 TVarName = obsDict['TVar'] SVarName = obsDict['SVar'] zVarName = obsDict['zVar'] volVarName = obsDict['volVar'] obsFileName = obsDict['climatologyTask'][self.season].fileName ds = xarray.open_dataset(obsFileName, chunks=chunk) ds = ds.stack(nCells=(obsDict['latVar'], obsDict['lonVar'])) ds = ds.reset_index('nCells').drop_vars( [obsDict['latVar'], obsDict['lonVar']]) ds = ds.where(cellMask, drop=True) cellsChunk = 32768 chunk = {'nCells': cellsChunk} ds = ds.chunk(chunk) depthMask = numpy.logical_and(ds[zVarName] >= zmin, ds[zVarName] <= zmax) ds = ds.where(depthMask) ds.compute() T = ds[TVarName].values.ravel() mask = numpy.isfinite(T) T = T[mask] S = ds[SVarName].values.ravel()[mask] z = ds['zBroadcast'].values.ravel()[mask] volume = ds[volVarName].values.ravel()[mask] dsOut = xarray.Dataset() dsOut['T'] = ('nPoints', T) dsOut['S'] = ('nPoints', S) dsOut['z'] = ('nPoints', z) dsOut['volume'] = ('nPoints', volume) dsOut['zbounds'] = ('nBounds', [zmin, zmax]) write_netcdf(dsOut, outFileName)
def _compute_time_series_with_ncrcat(self): # {{{ ''' Uses ncrcat to extact time series from timeSeriesMonthlyOutput files Raises ------ OSError If ``ncrcat`` is not in the system path. Author ------ Xylar Asay-Davis ''' if find_executable('ncrcat') is None: raise OSError('ncrcat not found. Make sure the latest nco ' 'package is installed: \n' 'conda install nco\n' 'Note: this presumes use of the conda-forge ' 'channel.') inputFiles = self.inputFiles append = False if os.path.exists(self.outputFile): # make sure all the necessary variables are also present with xr.open_dataset(self.outputFile) as ds: if ds.sizes['Time'] == 0: updateSubset = False else: updateSubset = True for variableName in self.variableList: if variableName not in ds.variables: updateSubset = False break if updateSubset: # add only input files wiht times that aren't already in # the output file append = True fileNames = sorted(self.inputFiles) inYears, inMonths = get_files_year_month( fileNames, self.historyStreams, 'timeSeriesStatsMonthlyOutput') inYears = numpy.array(inYears) inMonths = numpy.array(inMonths) totalMonths = 12 * inYears + inMonths dates = decode_strings(ds.xtime_startMonthly) lastDate = dates[-1] lastYear = int(lastDate[0:4]) lastMonth = int(lastDate[5:7]) lastTotalMonths = 12 * lastYear + lastMonth inputFiles = [] for index, inputFile in enumerate(fileNames): if totalMonths[index] > lastTotalMonths: inputFiles.append(inputFile) if len(inputFiles) == 0: # nothing to do return else: # there is an output file but it has the wrong variables # so we need ot delete it. self.logger.warning('Warning: deleting file {} because ' 'it is empty or some variables were ' 'missing'.format(self.outputFile)) os.remove(self.outputFile) variableList = self.variableList + ['xtime_startMonthly', 'xtime_endMonthly'] args = ['ncrcat', '-4', '--no_tmp_fl', '-v', ','.join(variableList)] if append: args.append('--record_append') printCommand = '{} {} ... {} {}'.format(' '.join(args), inputFiles[0], inputFiles[-1], self.outputFile) args.extend(inputFiles) args.append(self.outputFile) self.logger.info('running: {}'.format(printCommand)) for handler in self.logger.handlers: handler.flush() process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr = process.communicate() if stdout: stdout = stdout.decode('utf-8') for line in stdout.split('\n'): self.logger.info(line) if stderr: stderr = stderr.decode('utf-8') for line in stderr.split('\n'): self.logger.error(line) if process.returncode != 0: raise subprocess.CalledProcessError(process.returncode, ' '.join(args))
def run_task(self): # {{{ """ Compute the regional-mean time series """ # Authors # ------- # Xylar Asay-Davis config = self.config self.logger.info("\nCompute depth mask for regional means...") regionGroup = self.regionGroup sectionSuffix = regionGroup[0].upper() + \ regionGroup[1:].replace(' ', '') timeSeriesName = sectionSuffix[0].lower() + sectionSuffix[1:] sectionName = 'timeSeries{}'.format(sectionSuffix) outputDirectory = '{}/{}/'.format( build_config_full_path(config, 'output', 'timeseriesSubdirectory'), timeSeriesName) try: os.makedirs(outputDirectory) except OSError: pass outFileName = '{}/depthMasks{}.nc'.format(outputDirectory, timeSeriesName) if os.path.exists(outFileName): self.logger.info(' Mask file exists -- Done.') return # Load mesh related variables try: restartFileName = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for ocean region time series') if config.has_option(sectionName, 'zmin'): config_zmin = config.getfloat(sectionName, 'zmin') else: config_zmin = None if config.has_option(sectionName, 'zmax'): config_zmax = config.getfloat(sectionName, 'zmax') else: config_zmax = None dsRestart = xarray.open_dataset(restartFileName).isel(Time=0) zMid = compute_zmid(dsRestart.bottomDepth, dsRestart.maxLevelCell, dsRestart.layerThickness) areaCell = dsRestart.areaCell if 'landIceMask' in dsRestart: # only the region outside of ice-shelf cavities openOceanMask = dsRestart.landIceMask == 0 else: openOceanMask = None regionMaskFileName = self.masksSubtask.maskFileName dsRegionMask = xarray.open_dataset(regionMaskFileName) maskRegionNames = decode_strings(dsRegionMask.regionNames) regionIndices = [] for regionName in self.regionNames: for index, otherName in enumerate(maskRegionNames): if regionName == otherName: regionIndices.append(index) break # select only those regions we want to plot dsRegionMask = dsRegionMask.isel(nRegions=regionIndices) nRegions = dsRegionMask.sizes['nRegions'] datasets = [] for regionIndex in range(nRegions): self.logger.info(' region: {}'.format( self.regionNames[regionIndex])) dsRregion = dsRegionMask.isel(nRegions=regionIndex) cellMask = dsRregion.regionCellMasks == 1 if openOceanMask is not None: cellMask = numpy.logical_and(cellMask, openOceanMask) totalArea = areaCell.where(cellMask).sum() self.logger.info(' totalArea: {} mil. km^2'.format( 1e-12 * totalArea.values)) if config_zmin is None: if 'zminRegions' in dsRregion: zmin = dsRregion.zminRegions else: # the old naming convention, used in some pre-generated # mask files zmin = dsRregion.zmin else: zmin = (('nRegions', ), config_zmin) if config_zmax is None: if 'zmaxRegions' in dsRregion: zmax = dsRregion.zmaxRegions else: # the old naming convention, used in some pre-generated # mask files zmax = dsRregion.zmax else: zmax = (('nRegions', ), config_zmax) depthMask = numpy.logical_and(zMid >= zmin, zMid <= zmax) dsOut = xarray.Dataset() dsOut['zmin'] = zmin dsOut['zmax'] = zmax dsOut['totalArea'] = totalArea dsOut['cellMask'] = cellMask dsOut['depthMask'] = depthMask datasets.append(dsOut) dsOut = xarray.concat(objs=datasets, dim='nRegions') zbounds = numpy.zeros((nRegions, 2)) zbounds[:, 0] = dsOut.zmin.values zbounds[:, 1] = dsOut.zmax.values dsOut['zbounds'] = (('nRegions', 'nbounds'), zbounds) dsOut['areaCell'] = areaCell dsOut['regionNames'] = dsRegionMask.regionNames write_netcdf(dsOut, outFileName)
def run_task(self): # {{{ """ Compute time series of regional profiles """ # Authors # ------- # Milena Veneziani, Mark Petersen, Phillip J. Wolfram, Xylar Asay-Davis self.logger.info("\nCompute time series of regional profiles...") startDate = '{:04d}-01-01_00:00:00'.format(self.startYear) endDate = '{:04d}-12-31_23:59:59'.format(self.endYear) timeSeriesName = self.masksSubtask.regionGroup.replace(' ', '') outputDirectory = '{}/{}/'.format( build_config_full_path(self.config, 'output', 'timeseriesSubdirectory'), timeSeriesName) try: os.makedirs(outputDirectory) except OSError: pass outputFileName = '{}/regionalProfiles_{}_{:04d}-{:04d}.nc'.format( outputDirectory, timeSeriesName, self.startYear, self.endYear) inputFiles = sorted(self.historyStreams.readpath( 'timeSeriesStatsMonthlyOutput', startDate=startDate, endDate=endDate, calendar=self.calendar)) years, months = get_files_year_month(inputFiles, self.historyStreams, 'timeSeriesStatsMonthlyOutput') variableList = [field['mpas'] for field in self.fields] outputExists = os.path.exists(outputFileName) outputValid = outputExists if outputExists: with open_mpas_dataset(fileName=outputFileName, calendar=self.calendar, timeVariableNames=None, variableList=None, startDate=startDate, endDate=endDate) as dsIn: for inIndex in range(dsIn.dims['Time']): mask = np.logical_and( dsIn.year[inIndex].values == years, dsIn.month[inIndex].values == months) if np.count_nonzero(mask) == 0: outputValid = False break if outputValid: self.logger.info(' Time series exists -- Done.') return # get areaCell restartFileName = \ self.runStreams.readpath('restart')[0] dsRestart = xr.open_dataset(restartFileName) dsRestart = dsRestart.isel(Time=0) areaCell = dsRestart.areaCell nVertLevels = dsRestart.sizes['nVertLevels'] vertIndex = \ xr.DataArray.from_dict({'dims': ('nVertLevels',), 'data': np.arange(nVertLevels)}) vertMask = vertIndex < dsRestart.maxLevelCell # get region masks regionMaskFileName = self.masksSubtask.maskFileName dsRegionMask = xr.open_dataset(regionMaskFileName) # figure out the indices of the regions to plot regionNames = decode_strings(dsRegionMask.regionNames) regionIndices = [] for regionToPlot in self.regionNames: for index, regionName in enumerate(regionNames): if regionToPlot == regionName: regionIndices.append(index) break # select only those regions we want to plot dsRegionMask = dsRegionMask.isel(nRegions=regionIndices) cellMasks = dsRegionMask.regionCellMasks regionNamesVar = dsRegionMask.regionNames totalArea = (cellMasks * areaCell * vertMask).sum('nCells') datasets = [] for timeIndex, fileName in enumerate(inputFiles): dsLocal = open_mpas_dataset( fileName=fileName, calendar=self.calendar, variableList=variableList, startDate=startDate, endDate=endDate) dsLocal = dsLocal.isel(Time=0) time = dsLocal.Time.values date = days_to_datetime(time, calendar=self.calendar) self.logger.info(' date: {:04d}-{:02d}'.format(date.year, date.month)) # for each region and variable, compute area-weighted sum and # squared sum for field in self.fields: variableName = field['mpas'] prefix = field['prefix'] self.logger.info(' {}'.format(field['titleName'])) var = dsLocal[variableName].where(vertMask) meanName = '{}_mean'.format(prefix) dsLocal[meanName] = \ (cellMasks * areaCell * var).sum('nCells') / totalArea meanSquaredName = '{}_meanSquared'.format(prefix) dsLocal[meanSquaredName] = \ (cellMasks * areaCell * var**2).sum('nCells') / totalArea # drop the original variables dsLocal = dsLocal.drop_vars(variableList) datasets.append(dsLocal) # combine data sets into a single data set dsOut = xr.concat(datasets, 'Time') dsOut.coords['regionNames'] = regionNamesVar dsOut['totalArea'] = totalArea dsOut.coords['year'] = (('Time',), years) dsOut['year'].attrs['units'] = 'years' dsOut.coords['month'] = (('Time',), months) dsOut['month'].attrs['units'] = 'months' # Note: restart file, not a mesh file because we need refBottomDepth, # not in a mesh file try: restartFile = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for plotting time series vs. depth') with xr.open_dataset(restartFile) as dsRestart: depths = dsRestart.refBottomDepth.values z = np.zeros(depths.shape) z[0] = -0.5 * depths[0] z[1:] = -0.5 * (depths[0:-1] + depths[1:]) dsOut.coords['z'] = (('nVertLevels',), z) dsOut['z'].attrs['units'] = 'meters' write_netcdf(dsOut, outputFileName)
def run_task(self): # {{{ """ Plot a depth profile with variability """ # Authors # ------- # Xylar Asay-Davis config = self.config startYear = self.startYear endYear = self.endYear regionMaskFile = self.masksSubtask.geojsonFileName fcAll = read_feature_collection(regionMaskFile) fc = FeatureCollection() for feature in fcAll.features: if feature['properties']['name'] == self.regionName: fc.add_feature(feature) break inDirectory = build_config_full_path(config, 'output', 'profilesSubdirectory') timeSeriesName = self.timeSeriesName inFileName = '{}/{}_{}_{:04d}-{:04d}.nc'.format( inDirectory, timeSeriesName, self.season, self.startYear, self.endYear) regionGroup = self.masksSubtask.regionGroup regionGroupSection = 'profiles{}'.format( regionGroup.replace(' ', '')) ds = xr.open_dataset(inFileName) allRegionNames = decode_strings(ds.regionNames) regionIndex = allRegionNames.index(self.regionName) ds = ds.isel(nRegions=regionIndex) meanFieldName = '{}_mean'.format(self.field['prefix']) stdFieldName = '{}_std'.format(self.field['prefix']) mainRunName = config.get('runs', 'mainRunName') profileGalleryGroup = config.get(regionGroupSection, 'profileGalleryGroup') titleFieldName = self.field['titleName'] regionName = self.regionName.replace('_', ' ') xLabel = '{} ({})'.format(titleFieldName, self.field['units']) yLabel = 'depth (m)' outFileName = '{}/{}.png'.format(self.plotsDirectory, self.filePrefix) lineColors = ['k'] lineWidths = [1.6] zArrays = [ds.z.values] fieldArrays = [ds[meanFieldName].values] errArrays = [ds[stdFieldName].values] if self.controlConfig is None: title = '{} {}, years {:04d}-{:04d}\n{}'.format( regionName, self.season, startYear, endYear, mainRunName) legendText = [None] else: controlStartYear = self.controlConfig.getint('climatology', 'startYear') controlEndYear = self.controlConfig.getint('climatology', 'endYear') controlRunName = self.controlConfig.get('runs', 'mainRunName') if controlStartYear == startYear and controlEndYear == endYear: title = '{} {}, years {:04d}-{:04d}'.format( regionName, self.season, startYear, endYear) legendText = [mainRunName, controlRunName] elif mainRunName == controlRunName: title = '{} {}\n{}'.format( regionName, self.season, mainRunName) legendText = ['{:04d}-{:04d}'.format(startYear, endYear), '{:04d}-{:04d}'.format(controlStartYear, controlEndYear)] else: title = '{} {} '.format(regionName, self.season) legendText = ['{} {:04d}-{:04d}'.format(mainRunName, startYear, endYear), '{} {:04d}-{:04d}'.format(controlRunName, controlStartYear, controlEndYear)] controlDirectory = build_config_full_path( self.controlConfig, 'output', 'profilesSubdirectory') controlFileName = \ '{}/{}_{}_{:04d}-{:04d}.nc'.format( controlDirectory, timeSeriesName, self.season, controlStartYear, controlEndYear) dsControl = xr.open_dataset(controlFileName) allRegionNames = decode_strings(dsControl.regionNames) regionIndex = allRegionNames.index(self.regionName) dsControl = dsControl.isel(nRegions=regionIndex) lineColors.append('r') lineWidths.append(1.2) zArrays.append(dsControl.z.values) fieldArrays.append(dsControl[meanFieldName].values) errArrays.append(dsControl[stdFieldName].values) depthRange = config.getExpression(regionGroupSection, 'depthRange') if len(depthRange) == 0: depthRange = None fig = self.plot(zArrays, fieldArrays, errArrays, lineColors=lineColors, lineWidths=lineWidths, legendText=legendText, title=title, xLabel=xLabel, yLabel=yLabel, yLim=depthRange) # do this before the inset because otherwise it moves the inset # and cartopy doesn't play too well with tight_layout anyway plt.tight_layout() add_inset(fig, fc, width=1.0, height=1.0) savefig(outFileName, tight=False) caption = '{} {} vs depth'.format(regionName, titleFieldName) write_image_xml( config=config, filePrefix=self.filePrefix, componentName='Ocean', componentSubdirectory='ocean', galleryGroup=profileGalleryGroup, groupLink='ocnregprofs', imageDescription=caption, imageCaption=caption, gallery=titleFieldName, thumbnailDescription='{} {}'.format(regionName, self.season))
def run_task(self): # {{{ """ Computes time-series of Antarctic sub-ice-shelf melt rates. """ # Authors # ------- # Xylar Asay-Davis, Stephen Price self.logger.info("Computing Antarctic melt rate time series...") mpasTimeSeriesTask = self.mpasTimeSeriesTask config = self.config baseDirectory = build_config_full_path( config, 'output', 'timeSeriesSubdirectory') outFileName = '{}/iceShelfAggregatedFluxes.nc'.format(baseDirectory) # Load data: inputFile = mpasTimeSeriesTask.outputFile dsIn = open_mpas_dataset(fileName=inputFile, calendar=self.calendar, variableList=self.variableList, startDate=self.startDate, endDate=self.endDate) try: if os.path.exists(outFileName): # The file already exists so load it dsOut = xarray.open_dataset(outFileName) if numpy.all(dsOut.Time.values == dsIn.Time.values): return else: self.logger.warning('File {} is incomplete. Deleting ' 'it.'.format(outFileName)) os.remove(outFileName) except OSError: # something is potentailly wrong with the file, so let's delete # it and try again self.logger.warning('Problems reading file {}. Deleting ' 'it.'.format(outFileName)) os.remove(outFileName) restartFileName = \ mpasTimeSeriesTask.runStreams.readpath('restart')[0] dsRestart = xarray.open_dataset(restartFileName) areaCell = \ dsRestart.landIceFraction.isel(Time=0) * dsRestart.areaCell regionMaskFileName = self.masksSubtask.maskFileName dsRegionMask = xarray.open_dataset(regionMaskFileName) # figure out the indices of the regions to plot regionNames = decode_strings(dsRegionMask.regionNames) regionIndices = [] for iceShelf in self.iceShelvesToPlot: for index, regionName in enumerate(regionNames): if iceShelf == regionName: regionIndices.append(index) break # select only those regions we want to plot dsRegionMask = dsRegionMask.isel(nRegions=regionIndices) datasets = [] nTime = dsIn.sizes['Time'] for tIndex in range(nTime): self.logger.info(' {}/{}'.format(tIndex+1, nTime)) freshwaterFlux = \ dsIn.timeMonthly_avg_landIceFreshwaterFlux.isel(Time=tIndex) nRegions = dsRegionMask.sizes['nRegions'] meltRates = numpy.zeros((nRegions,)) totalMeltFluxes = numpy.zeros((nRegions,)) for regionIndex in range(nRegions): cellMask = \ dsRegionMask.regionCellMasks.isel(nRegions=regionIndex) # convert from kg/s to kg/yr totalMeltFlux = constants.sec_per_year * \ (cellMask * areaCell * freshwaterFlux).sum(dim='nCells') totalArea = (cellMask * areaCell).sum(dim='nCells') # from kg/m^2/yr to m/yr meltRates[regionIndex] = ((1. / constants.rho_fw) * (totalMeltFlux / totalArea)) # convert from kg/yr to GT/yr totalMeltFlux /= constants.kg_per_GT totalMeltFluxes[regionIndex] = totalMeltFlux dsOut = xarray.Dataset() dsOut.coords['Time'] = dsIn.Time.isel(Time=tIndex) dsOut['totalMeltFlux'] = (('nRegions'), totalMeltFluxes) dsOut['meltRates'] = (('nRegions'), meltRates) datasets.append(dsOut) dsOut = xarray.concat(objs=datasets, dim='Time') dsOut['regionNames'] = dsRegionMask.regionNames dsOut.totalMeltFlux.attrs['units'] = 'GT a$^{-1}$' dsOut.totalMeltFlux.attrs['description'] = \ 'Total melt flux summed over each ice shelf or region' dsOut.meltRates.attrs['units'] = 'm a$^{-1}$' dsOut.meltRates.attrs['description'] = \ 'Melt rate averaged over each ice shelf or region' write_netcdf(dsOut, outFileName)
def run_task(self): # {{{ """ Compute time-series output of properties in an ocean region. """ # Authors # ------- # Xylar Asay-Davis self.logger.info("\nAveraging T and S for {}...".format( self.regionName)) obsDict = self.obsDict config = self.config regionGroup = self.regionGroup timeSeriesName = regionGroup[0].lower() + \ regionGroup[1:].replace(' ', '') sectionSuffix = regionGroup[0].upper() + \ regionGroup[1:].replace(' ', '') sectionName = 'timeSeries{}'.format(sectionSuffix) outputDirectory = '{}/{}/'.format( build_config_full_path(self.config, 'output', 'timeseriesSubdirectory'), timeSeriesName) try: os.makedirs(outputDirectory) except OSError: pass outFileName = '{}/TS_{}_{}.nc'.format(outputDirectory, obsDict['suffix'], self.prefix) if os.path.exists(outFileName): return regionMaskFileName = obsDict['maskTask'].maskFileName print(regionMaskFileName) print(xarray.open_dataset(regionMaskFileName)) dsRegionMask = \ xarray.open_dataset(regionMaskFileName).stack( nCells=(obsDict['latVar'], obsDict['lonVar'])) dsRegionMask = dsRegionMask.reset_index('nCells').drop_vars( [obsDict['latVar'], obsDict['lonVar']]) maskRegionNames = decode_strings(dsRegionMask.regionNames) regionIndex = maskRegionNames.index(self.regionName) dsMask = dsRegionMask.isel(nRegions=regionIndex) cellMask = dsMask.regionCellMasks == 1 if config.has_option(sectionName, 'zmin'): zmin = config.getfloat(sectionName, 'zmin') else: zmin = dsMask.zminRegions.values if config.has_option(sectionName, 'zmax'): zmax = config.getfloat(sectionName, 'zmax') else: zmax = dsMask.zmaxRegions.values TVarName = obsDict['TVar'] SVarName = obsDict['SVar'] zVarName = obsDict['zVar'] lonVarName = obsDict['lonVar'] latVarName = obsDict['latVar'] volVarName = obsDict['volVar'] tDim = obsDict['tDim'] obsFileName = build_obs_path(config, component=self.componentName, relativePath=obsDict['TFileName']) self.logger.info(' Reading from {}...'.format(obsFileName)) ds = xarray.open_dataset(obsFileName) if obsDict['SFileName'] != obsDict['TFileName']: obsFileName = build_obs_path(config, component=self.componentName, relativePath=obsDict['SFileName']) self.logger.info(' Reading from {}...'.format(obsFileName)) dsS = xarray.open_dataset(obsFileName) ds[SVarName] = dsS[SVarName] if obsDict['volFileName'] is None: # compute volume from lat, lon, depth bounds self.logger.info(' Computing volume...'.format(obsFileName)) latBndsName = ds[latVarName].attrs['bounds'] lonBndsName = ds[lonVarName].attrs['bounds'] zBndsName = ds[zVarName].attrs['bounds'] latBnds = ds[latBndsName] lonBnds = ds[lonBndsName] zBnds = ds[zBndsName] dLat = numpy.deg2rad(latBnds[:, 1] - latBnds[:, 0]) dLon = numpy.deg2rad(lonBnds[:, 1] - lonBnds[:, 0]) lat = numpy.deg2rad(ds[latVarName]) dz = zBnds[:, 1] - zBnds[:, 0] radius = 6378137.0 area = radius**2 * numpy.cos(lat) * dLat * dLon volume = dz * area ds[volVarName] = volume elif obsDict['volFileName'] != obsDict['TFileName']: obsFileName = build_obs_path(config, component=self.componentName, relativePath=obsDict['volFileName']) self.logger.info(' Reading from {}...'.format(obsFileName)) dsVol = xarray.open_dataset(obsFileName) ds[volVarName] = dsVol[volVarName] if 'positive' in ds[zVarName].attrs and \ ds[zVarName].attrs['positive'] == 'down': attrs = ds[zVarName].attrs ds[zVarName] = -ds[zVarName] ds[zVarName].attrs = attrs ds[zVarName].attrs['positive'] = 'up' TMean = numpy.zeros(ds.sizes[tDim]) SMean = numpy.zeros(ds.sizes[tDim]) depthMask = numpy.logical_and(ds[zVarName] >= zmin, ds[zVarName] <= zmax) for tIndex in range(ds.sizes[tDim]): dsMonth = ds.isel({tDim: tIndex}) dsMonth = dsMonth.stack(nCells=(obsDict['latVar'], obsDict['lonVar'])) dsMonth = dsMonth.reset_index('nCells').drop_vars( [obsDict['latVar'], obsDict['lonVar']]) dsMonth = dsMonth.where(cellMask, drop=True) dsMonth = dsMonth.where(depthMask) mask = dsMonth[TVarName].notnull() TSum = (dsMonth[TVarName] * dsMonth[volVarName]).sum(dim=('nCells', zVarName)) volSum = (mask * dsMonth[volVarName]).sum(dim=('nCells', zVarName)) TMean[tIndex] = TSum / volSum mask = dsMonth[SVarName].notnull() SSum = (dsMonth[SVarName] * dsMonth[volVarName]).sum(dim=('nCells', zVarName)) volSum = (mask * dsMonth[volVarName]).sum(dim=('nCells', zVarName)) SMean[tIndex] = SSum / volSum dsOut = xarray.Dataset() dsOut['temperature'] = ('Time', TMean) dsOut['salinity'] = ('Time', SMean) dsOut['zbounds'] = ('nBounds', [zmin, zmax]) dsOut['month'] = ('Time', numpy.array(ds.month.values, dtype=float)) dsOut['year'] = ('Time', numpy.ones(ds.sizes[tDim])) write_netcdf(dsOut, outFileName)
def run_task(self): # {{{ """ Make the Hovmoller plot from the time series. """ # Authors # ------- # Xylar Asay-Davis, Milena Veneziani, Greg Streletz self.logger.info("\nPlotting {} time series vs. depth...".format( self.fieldNameInTitle)) config = self.config mainRunName = config.get('runs', 'mainRunName') self.logger.info(' Load ocean data...') ds = xr.open_dataset(self.inFileName) if 'regionNames' in ds.coords: allRegionNames = decode_strings(ds.regionNames) regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = self.regionName.replace('_', ' ') regionDim = ds.regionNames.dims[0] else: plotTitles = config.getExpression('regions', 'plotTitles') allRegionNames = config.getExpression('regions', 'regions') regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = plotTitles[regionIndex] regionDim = 'nOceanRegionsTmp' ds = ds.isel(**{regionDim: regionIndex}) # Note: restart file, not a mesh file because we need refBottomDepth, # not in a mesh file try: restartFile = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for plotting time series vs. depth') # Define/read in general variables self.logger.info(' Read in depth...') with xr.open_dataset(restartFile) as dsRestart: # reference depth [m] depths = dsRestart.refBottomDepth.values z = np.zeros(depths.shape) z[0] = -0.5 * depths[0] z[1:] = -0.5 * (depths[0:-1] + depths[1:]) Time = ds.Time.values field = ds[self.mpasFieldName].values.transpose() xLabel = 'Time (years)' yLabel = 'Depth (m)' title = '{}, {}'.format(self.fieldNameInTitle, regionNameInTitle) outFileName = '{}/{}.png'.format(self.plotsDirectory, self.filePrefix) if config.has_option(self.sectionName, 'firstYearXTicks'): firstYearXTicks = config.getint(self.sectionName, 'firstYearXTicks') else: firstYearXTicks = None if config.has_option(self.sectionName, 'yearStrideXTicks'): yearStrideXTicks = config.getint(self.sectionName, 'yearStrideXTicks') else: yearStrideXTicks = None movingAverageMonths = config.getWithDefault(self.sectionName, 'movingAverageMonths', 1) if config.has_option(self.sectionName, 'yLim'): yLim = config.getExpression(self.sectionName, 'yLim') else: yLim = None if self.controlConfig is None: refField = None diff = None refTitle = None diffTitle = None else: controlConfig = self.controlConfig dsRef = xr.open_dataset(self.controlFileName) if 'regionNames' in dsRef.coords: allRegionNames = decode_strings(dsRef.regionNames) regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = self.regionName.replace('_', ' ') regionDim = dsRef.regionNames.dims[0] else: plotTitles = controlConfig.getExpression( 'regions', 'plotTitles') allRegionNames = controlConfig.getExpression( 'regions', 'regions') regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = plotTitles[regionIndex] regionDim = 'nOceanRegionsTmp' dsRef = dsRef.isel(**{regionDim: regionIndex}) refField = dsRef[self.mpasFieldName].values.transpose() assert (field.shape == refField.shape) diff = field - refField refTitle = self.controlConfig.get('runs', 'mainRunName') diffTitle = 'Main - Control' fig, _, suptitle = plot_vertical_section_comparison( config, Time, z, field, refField, diff, self.sectionName, colorbarLabel=self.unitsLabel, title=title, modelTitle=mainRunName, refTitle=refTitle, diffTitle=diffTitle, xlabel=xLabel, ylabel=yLabel, lineWidth=1, xArrayIsTime=True, movingAveragePoints=movingAverageMonths, calendar=self.calendar, firstYearXTicks=firstYearXTicks, yearStrideXTicks=yearStrideXTicks, yLim=yLim, invertYAxis=False) if self.regionMaskFile is not None: # shift the super-title a little to the left to make room for the # inset pos = suptitle.get_position() suptitle.set_position((pos[0] - 0.05, pos[1])) fcAll = read_feature_collection(self.regionMaskFile) fc = FeatureCollection() for feature in fcAll.features: if feature['properties']['name'] == self.regionName: fc.add_feature(feature) break add_inset(fig, fc, width=1.0, height=1.0, xbuffer=0.1, ybuffer=0.1) savefig(outFileName, tight=False) else: savefig(outFileName) write_image_xml(config=config, filePrefix=self.filePrefix, componentName='Ocean', componentSubdirectory='ocean', galleryGroup=self.galleryGroup, groupSubtitle=self.groupSubtitle, groupLink=self.groupLink, gallery=self.galleryName, thumbnailDescription='{} {}'.format( regionNameInTitle, self.thumbnailSuffix), imageDescription=self.imageCaption, imageCaption=self.imageCaption)
climoYear1 = 148 climoYear2 = 177 # #meshfile = '/global/project/projectdirs/e3sm/inputdata/ocn/mpas-o/oEC60to30v3/oEC60to30v3_60layer.170506.nc' #maskfile = '/global/project/projectdirs/e3sm/milena/mpas-region_masks/oEC60to30v3_oceanSubBasins20210315.nc' #casename = 'E3SM-LR-OSI' # no spaces #modeldir = '/global/project/projectdirs/m1199/milena/analysis/mpas/E3SM60to30/clim/mpas/avg/unmasked_oEC60to30v3' #climoYear1 = 148 #climoYear2 = 177 #seasons = ['ANN', 'JFM', 'JAS'] seasons = ['ANN'] if os.path.exists(maskfile): dsRegionMask = xr.open_dataset(maskfile) regionNames = decode_strings(dsRegionMask.regionNames) regionsToPlot = [ 'Southern Ocean Atlantic Sector', 'South Atlantic Basin', 'North Atlantic Basin', 'Arctic Ocean Basin' ] else: raise IOError('No regional mask file found') figdir = './verticalSections/{}'.format(casename) if not os.path.isdir(figdir): os.makedirs(figdir) figsize = (32, 5) figdpi = 300 colorIndices0 = [ 0, 10, 28, 57, 85, 113, 125, 142, 155, 170, 198, 227, 242, 255 ]
def run_task(self): # {{{ """ Computes and plots table of Antarctic sub-ice-shelf melt rates. """ # Authors # ------- # Xylar Asay-Davis self.logger.info("Computing Antarctic melt rate table...") config = self.config sectionName = self.taskName iceShelvesInTable = config.getExpression(sectionName, 'iceShelvesInTable') if len(iceShelvesInTable) == 0: return iceShelvesInTable = self.masksSubtask.expand_region_names( iceShelvesInTable) meltRateFileName = get_masked_mpas_climatology_file_name( config, self.season, self.componentName, climatologyName='antarcticMeltTable') if not os.path.exists(meltRateFileName): with dask.config.set(schedular='threads', pool=ThreadPool(1)): # Load data: inFileName = self.mpasClimatologyTask.get_file_name( self.season) mpasFieldName = 'timeMonthly_avg_landIceFreshwaterFlux' dsIn = xr.open_dataset(inFileName) freshwaterFlux = dsIn[mpasFieldName] if 'Time' in freshwaterFlux.dims: freshwaterFlux = freshwaterFlux.isel(Time=0) regionMaskFileName = self.masksSubtask.maskFileName dsRegionMask = xr.open_dataset(regionMaskFileName) # figure out the indices of the regions to plot regionNames = decode_strings(dsRegionMask.regionNames) regionIndices = [] for iceShelf in iceShelvesInTable: for index, regionName in enumerate(regionNames): if iceShelf == regionName: regionIndices.append(index) break # select only those regions we want to plot dsRegionMask = dsRegionMask.isel(nRegions=regionIndices) cellMasks = dsRegionMask.regionCellMasks.chunk( {'nRegions': 10}) restartFileName = \ self.runStreams.readpath('restart')[0] dsRestart = xr.open_dataset(restartFileName) areaCell = \ dsRestart.landIceFraction.isel(Time=0) * dsRestart.areaCell # convert from kg/s to kg/yr totalMeltFlux = constants.sec_per_year * \ (cellMasks * areaCell * freshwaterFlux).sum(dim='nCells') totalMeltFlux.compute() totalArea = (cellMasks * areaCell).sum(dim='nCells') # from kg/m^2/yr to m/yr meltRates = ((1. / constants.rho_fw) * (totalMeltFlux / totalArea)) meltRates.compute() # convert from kg/yr to GT/yr totalMeltFlux /= constants.kg_per_GT ds = xr.Dataset() ds['totalMeltFlux'] = totalMeltFlux ds.totalMeltFlux.attrs['units'] = 'GT a$^{-1}$' ds.totalMeltFlux.attrs['description'] = \ 'Total melt flux summed over each ice shelf or region' ds['meltRates'] = meltRates ds.meltRates.attrs['units'] = 'm a$^{-1}$' ds.meltRates.attrs['description'] = \ 'Melt rate averaged over each ice shelf or region' ds['area'] = 1e-6 * totalArea ds.meltRates.attrs['units'] = 'km$^2$' ds.meltRates.attrs['description'] = \ 'Region or ice shelf area' ds['regionNames'] = dsRegionMask.regionNames write_netcdf(ds, meltRateFileName) else: ds = xr.open_dataset(meltRateFileName) mainRunName = config.get('runs', 'mainRunName') fieldNames = ['Region', 'Area', mainRunName] controlConfig = self.controlConfig if controlConfig is not None: controlFileName = get_masked_mpas_climatology_file_name( controlConfig, self.season, self.componentName, climatologyName='antarcticMeltTable') dsControl = xr.open_dataset(controlFileName) controlRunName = controlConfig.get('runs', 'mainRunName') fieldNames.append(controlRunName) else: dsControl = None controlRunName = None regionNames = decode_strings(ds.regionNames) outDirectory = '{}/antarcticMelt/'.format( build_config_full_path(config, 'output', 'tablesSubdirectory')) try: os.makedirs(outDirectory) except OSError: pass tableFileName = '{}/antarcticMeltRateTable_{}.csv'.format( outDirectory, self.season) with open(tableFileName, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldNames) writer.writeheader() for index, regionName in enumerate(regionNames): row = { 'Region': regionName, 'Area': '{}'.format(ds.area[index].values), mainRunName: '{}'.format(ds.meltRates[index].values) } if dsControl is not None: row[controlRunName] = \ '{}'.format(dsControl.meltRates[index].values) writer.writerow(row) tableFileName = '{}/antarcticMeltFluxTable_{}.csv'.format( outDirectory, self.season) with open(tableFileName, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldNames) writer.writeheader() for index, regionName in enumerate(regionNames): row = { 'Region': regionName, 'Area': '{}'.format(ds.area[index].values), mainRunName: '{}'.format(ds.totalMeltFlux[index].values) } if dsControl is not None: row[controlRunName] = \ '{}'.format(dsControl.totalMeltFlux[index].values) writer.writerow(row)
def run_task(self): # {{{ """ Computes time-series of transport through transects. """ # Authors # ------- # Xylar Asay-Davis, Stephen Price self.logger.info("Computing time series of transport through " "transects...") config = self.config startDate = '{:04d}-01-01_00:00:00'.format(self.startYear) endDate = '{:04d}-12-31_23:59:59'.format(self.endYear) outputDirectory = '{}/transport/'.format( build_config_full_path(config, 'output', 'timeseriesSubdirectory')) try: os.makedirs(outputDirectory) except OSError: pass outFileName = '{}/transport_{:04d}-{:04d}.nc'.format( outputDirectory, self.startYear, self.endYear) inputFiles = sorted( self.historyStreams.readpath('timeSeriesStatsMonthlyOutput', startDate=startDate, endDate=endDate, calendar=self.calendar)) years, months = get_files_year_month(inputFiles, self.historyStreams, 'timeSeriesStatsMonthlyOutput') variableList = ['timeMonthly_avg_layerThickness'] with open_mpas_dataset(fileName=inputFiles[0], calendar=self.calendar, startDate=startDate, endDate=endDate) as dsIn: if 'timeMonthly_avg_normalTransportVelocity' in dsIn: variableList.append('timeMonthly_avg_normalTransportVelocity') elif 'timeMonthly_avg_normalGMBolusVelocity' in dsIn: variableList = variableList + \ ['timeMonthly_avg_normalVelocity', 'timeMonthly_avg_normalGMBolusVelocity'] else: self.logger.warning('Cannot compute transport velocity. ' 'Using advection velocity.') variableList.append('timeMonthly_avg_normalVelocity') outputExists = os.path.exists(outFileName) outputValid = outputExists if outputExists: with open_mpas_dataset(fileName=outFileName, calendar=self.calendar, timeVariableNames=None, variableList=None, startDate=startDate, endDate=endDate) as dsOut: for inIndex in range(dsOut.dims['Time']): mask = numpy.logical_and( dsOut.year[inIndex].values == years, dsOut.month[inIndex].values == months) if numpy.count_nonzero(mask) == 0: outputValid = False break if outputValid: self.logger.info(' Time series exists -- Done.') return transectMaskFileName = self.masksSubtask.maskFileName dsTransectMask = xarray.open_dataset(transectMaskFileName) # figure out the indices of the transects to plot maskTransectNames = decode_strings(dsTransectMask.transectNames) dsMesh = xarray.open_dataset(self.restartFileName) dvEdge = dsMesh.dvEdge cellsOnEdge = dsMesh.cellsOnEdge - 1 timeDatasets = [] self.logger.info(' Computing transport...') for fileName in inputFiles: self.logger.info(' input file: {}'.format(fileName)) dsTimeSlice = open_mpas_dataset(fileName=fileName, calendar=self.calendar, variableList=variableList, startDate=startDate, endDate=endDate) transectDatasets = [] transectIndices = [] for transect in self.transectsToPlot: self.logger.info(' transect: {}'.format(transect)) try: transectIndex = maskTransectNames.index(transect) except ValueError: self.logger.warning(' Not found in masks. ' 'Skipping.') continue transectIndices.append(transectIndex) # select the current transect dsMask = dsTransectMask.isel(nTransects=[transectIndex]) edgeIndices = dsMask.transectEdgeGlobalIDs - 1 edgeIndices = edgeIndices.where(edgeIndices >= 0, drop=True).astype(int) edgeSign = dsMask.transectEdgeMaskSigns.isel( nEdges=edgeIndices) dsIn = dsTimeSlice.isel(nEdges=edgeIndices) dv = dvEdge.isel(nEdges=edgeIndices) coe = cellsOnEdge.isel(nEdges=edgeIndices) # work on data from simulations if 'timeMonthly_avg_normalTransportVelocity' in dsIn: vel = dsIn.timeMonthly_avg_normalTransportVelocity elif 'timeMonthly_avg_normalGMBolusVelocity' in dsIn: vel = (dsIn.timeMonthly_avg_normalVelocity + dsIn.timeMonthly_avg_normalGMBolusVelocity) else: vel = dsIn.timeMonthly_avg_normalVelocity # get layer thickness on edges by averaging adjacent cells h = 0.5 * dsIn.timeMonthly_avg_layerThickness.isel( nCells=coe).sum(dim='TWO') edgeTransport = edgeSign * vel * h * dv # convert from m^3/s to Sv transport = (constants.m3ps_to_Sv * edgeTransport.sum( dim=['maxEdgesInTransect', 'nVertLevels'])) dsOut = xarray.Dataset() dsOut['transport'] = transport dsOut.transport.attrs['units'] = 'Sv' dsOut.transport.attrs['description'] = \ 'Transport through transects' transectDatasets.append(dsOut) dsOut = xarray.concat(transectDatasets, 'nTransects') timeDatasets.append(dsOut) # combine data sets into a single data set dsOut = xarray.concat(timeDatasets, 'Time') dsOut.coords['transectNames'] = dsTransectMask.transectNames.isel( nTransects=transectIndices) dsOut.coords['year'] = (('Time'), years) dsOut['year'].attrs['units'] = 'years' dsOut.coords['month'] = (('Time'), months) dsOut['month'].attrs['units'] = 'months' write_netcdf(dsOut, outFileName)
def run_task(self): # {{{ """ Make the Hovmoller plot from the time series. """ # Authors # ------- # Xylar Asay-Davis, Milena Veneziani, Greg Streletz self.logger.info("\nPlotting {} time series vs. depth...".format( self.fieldNameInTitle)) config = self.config mainRunName = config.get('runs', 'mainRunName') self.logger.info(' Load ocean data...') ds = xr.open_dataset(self.inFileName) if 'regionNames' in ds.coords: allRegionNames = decode_strings(ds.regionNames) regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = self.regionName.replace('_', ' ') regionDim = ds.regionNames.dims[0] else: plotTitles = config.getExpression('regions', 'plotTitles') allRegionNames = config.getExpression('regions', 'regions') regionIndex = allRegionNames.index(self.regionName) regionNameInTitle = plotTitles[regionIndex] regionDim = 'nOceanRegionsTmp' ds = ds.isel(**{regionDim: regionIndex}) # Note: restart file, not a mesh file because we need refBottomDepth, # not in a mesh file try: restartFile = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for plotting time series vs. depth') # Define/read in general variables self.logger.info(' Read in depth...') with xr.open_dataset(restartFile) as dsRestart: # reference depth [m] depths = dsRestart.refBottomDepth.values z = np.zeros(depths.shape) z[0] = -0.5 * depths[0] z[1:] = -0.5 * (depths[0:-1] + depths[1:]) Time = ds.Time.values field = ds[self.mpasFieldName].values.transpose() xLabel = 'Time (years)' yLabel = 'Depth (m)' title = '{}, {} \n {}'.format(self.fieldNameInTitle, regionNameInTitle, mainRunName) outFileName = '{}/{}.png'.format(self.plotsDirectory, self.filePrefix) if config.has_option(self.sectionName, 'firstYearXTicks'): firstYearXTicks = config.getint(self.sectionName, 'firstYearXTicks') else: firstYearXTicks = None if config.has_option(self.sectionName, 'yearStrideXTicks'): yearStrideXTicks = config.getint(self.sectionName, 'yearStrideXTicks') else: yearStrideXTicks = None if config.has_option(self.sectionName, 'yLim'): yLim = config.getExpression(self.sectionName, 'yLim') else: yLim = None plot_vertical_section(config, Time, z, field, self.sectionName, suffix='', colorbarLabel=self.unitsLabel, title=title, xlabel=xLabel, ylabel=yLabel, lineWidth=1, xArrayIsTime=True, calendar=self.calendar, firstYearXTicks=firstYearXTicks, yearStrideXTicks=yearStrideXTicks, yLim=yLim, invertYAxis=False) savefig(outFileName) write_image_xml(config=config, filePrefix=self.filePrefix, componentName='Ocean', componentSubdirectory='ocean', galleryGroup=self.galleryGroup, groupSubtitle=self.groupSubtitle, groupLink=self.groupLink, gallery=self.galleryName, thumbnailDescription='{} {}'.format( regionNameInTitle, self.thumbnailSuffix), imageDescription=self.imageCaption, imageCaption=self.imageCaption)
def run_task(self): # {{{ ''' Compute the regional-mean time series ''' # Authors # ------- # Xylar Asay-Davis config = self.config self.logger.info("\nCompute time series of regional means...") startDate = '{:04d}-01-01_00:00:00'.format(self.startYear) endDate = '{:04d}-12-31_23:59:59'.format(self.endYear) regionGroup = self.regionGroup sectionSuffix = regionGroup[0].upper() + \ regionGroup[1:].replace(' ', '') timeSeriesName = sectionSuffix[0].lower() + sectionSuffix[1:] sectionName = 'timeSeries{}'.format(sectionSuffix) outputDirectory = '{}/{}/'.format( build_config_full_path(config, 'output', 'timeseriesSubdirectory'), timeSeriesName) try: os.makedirs(outputDirectory) except OSError: pass outFileName = '{}/{}_{:04d}-{:04d}.nc'.format(outputDirectory, timeSeriesName, self.startYear, self.endYear) inputFiles = sorted( self.historyStreams.readpath('timeSeriesStatsMonthlyOutput', startDate=startDate, endDate=endDate, calendar=self.calendar)) years, months = get_files_year_month(inputFiles, self.historyStreams, 'timeSeriesStatsMonthlyOutput') variables = config.getExpression(sectionName, 'variables') variableList = [var['mpas'] for var in variables] + \ ['timeMonthly_avg_layerThickness'] outputExists = os.path.exists(outFileName) outputValid = outputExists if outputExists: with open_mpas_dataset(fileName=outFileName, calendar=self.calendar, timeVariableNames=None, variableList=None, startDate=startDate, endDate=endDate) as dsOut: for inIndex in range(dsOut.dims['Time']): mask = numpy.logical_and( dsOut.year[inIndex].values == years, dsOut.month[inIndex].values == months) if numpy.count_nonzero(mask) == 0: outputValid = False break if outputValid: self.logger.info(' Time series exists -- Done.') return # Load mesh related variables try: restartFileName = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for ocean region time series') cellsChunk = 32768 timeChunk = 1 datasets = [] for timeIndex, fileName in enumerate(inputFiles): dsTimeSlice = open_mpas_dataset(fileName=fileName, calendar=self.calendar, variableList=variableList, startDate=startDate, endDate=endDate) datasets.append(dsTimeSlice) chunk = {'Time': timeChunk, 'nCells': cellsChunk} if config.has_option(sectionName, 'zmin'): config_zmin = config.getfloat(sectionName, 'zmin') else: config_zmin = None if config.has_option(sectionName, 'zmax'): config_zmax = config.getfloat(sectionName, 'zmax') else: config_zmax = None with dask.config.set(schedular='threads', pool=ThreadPool(self.daskThreads)): # combine data sets into a single data set dsIn = xarray.concat(datasets, 'Time').chunk(chunk) chunk = {'nCells': cellsChunk} dsRestart = xarray.open_dataset(restartFileName) dsRestart = dsRestart.isel(Time=0).chunk(chunk) dsIn['areaCell'] = dsRestart.areaCell if 'landIceMask' in dsRestart: # only the region outside of ice-shelf cavities dsIn['openOceanMask'] = dsRestart.landIceMask == 0 dsIn['zMid'] = compute_zmid(dsRestart.bottomDepth, dsRestart.maxLevelCell, dsRestart.layerThickness) regionMaskFileName = self.masksSubtask.maskFileName dsRegionMask = xarray.open_dataset(regionMaskFileName) maskRegionNames = decode_strings(dsRegionMask.regionNames) datasets = [] regionIndices = [] for regionName in self.regionNames: self.logger.info(' region: {}'.format(regionName)) regionIndex = maskRegionNames.index(regionName) regionIndices.append(regionIndex) chunk = {'nCells': cellsChunk} dsMask = dsRegionMask.isel(nRegions=regionIndex).chunk(chunk) cellMask = dsMask.regionCellMasks == 1 if 'openOceanMask' in dsIn: cellMask = numpy.logical_and(cellMask, dsIn.openOceanMask) dsRegion = dsIn.where(cellMask, drop=True) totalArea = dsRegion['areaCell'].sum() self.logger.info(' totalArea: {} mil. km^2'.format( 1e-12 * totalArea.values)) self.logger.info("Don't worry about the following dask " "warnings.") if config_zmin is None: zmin = dsMask.zmin else: zmin = config_zmin if config_zmax is None: zmax = dsMask.zmax else: zmax = config_zmax depthMask = numpy.logical_and(dsRegion.zMid >= zmin, dsRegion.zMid <= zmax) depthMask.compute() self.logger.info("Dask warnings should be done.") dsRegion['depthMask'] = depthMask layerThickness = dsRegion.timeMonthly_avg_layerThickness dsRegion['volCell'] = (dsRegion.areaCell * layerThickness).where(depthMask) totalVol = dsRegion.volCell.sum(dim='nVertLevels').sum( dim='nCells') totalVol.compute() self.logger.info(' totalVol (mil. km^3): {}'.format( 1e-15 * totalVol.values)) dsRegion = dsRegion.transpose('Time', 'nCells', 'nVertLevels') dsOut = xarray.Dataset() dsOut['totalVol'] = totalVol dsOut.totalVol.attrs['units'] = 'm^3' dsOut['totalArea'] = totalArea dsOut.totalArea.attrs['units'] = 'm^2' dsOut['zbounds'] = ('nbounds', [zmin, zmax]) dsOut.zbounds.attrs['units'] = 'm' for var in variables: outName = var['name'] self.logger.info(' {}'.format(outName)) mpasVarName = var['mpas'] timeSeries = dsRegion[mpasVarName] units = timeSeries.units description = timeSeries.long_name is3d = 'nVertLevels' in timeSeries.dims if is3d: timeSeries = \ (dsRegion.volCell*timeSeries.where(depthMask)).sum( dim='nVertLevels').sum(dim='nCells') / totalVol else: timeSeries = \ (dsRegion.areaCell*timeSeries).sum( dim='nCells') / totalArea timeSeries.compute() dsOut[outName] = timeSeries dsOut[outName].attrs['units'] = units dsOut[outName].attrs['description'] = description dsOut[outName].attrs['is3d'] = str(is3d) datasets.append(dsOut) # combine data sets into a single data set dsOut = xarray.concat(datasets, 'nRegions') dsOut.coords['regionNames'] = dsRegionMask.regionNames.isel( nRegions=regionIndices) dsOut.coords['year'] = (('Time'), years) dsOut['year'].attrs['units'] = 'years' dsOut.coords['month'] = (('Time'), months) dsOut['month'].attrs['units'] = 'months' write_netcdf(dsOut, outFileName)
# cnormT = mpl.colors.BoundaryNorm(clevelsT, colormapT.N) cnormS = mpl.colors.BoundaryNorm(clevelsS, colormapS.N) cnormV = mpl.colors.BoundaryNorm(clevelsV, colormapV.N) #sigma2contours = [35, 36, 36.5, 36.8, 37, 37.1, 37.2, 37.25, 37.44, 37.52, 37.6] sigma2contours = None #sigma0contours = np.arange(26.0, 28.0, 0.2) # Good for OSNAP, but not for all arcticSections sigma0contours = [24.0, 25.0, 26.0, 27.0, 27.2, 27.4, 27.6, 27.8, 28.0] #sigma0contours = None # Load in MPAS mesh and transect mask file mesh = xr.open_dataset(meshfile) mask = xr.open_dataset(maskfile) allTransects = decode_strings(mask.transectNames) if transectNames[0]=='all' or transectNames[0]=='StandardTransportSectionsRegionsGroup': transectNames = allTransects # Get depth z = mesh.refBottomDepth.values nlevels = len(z) nTransects = len(transectNames) maxEdges = mask.dims['maxEdgesInTransect'] for n in range(nTransects): # Identify transect transectName = transectNames[n] transectIndex = allTransects.index(transectName) print('Plotting sections for transect: ', transectName)