Source code for mpas_analysis.ocean.time_series_transport

# This software is open source software available under the BSD-3 license.
#
# Copyright (c) 2022 Triad National Security, LLC. All rights reserved.
# Copyright (c) 2022 Lawrence Livermore National Security, LLC. All rights
# reserved.
# Copyright (c) 2022 UT-Battelle, LLC. All rights reserved.
#
# Additional copyright and license information can be found in the LICENSE file
# distributed with this code, or at
# https://raw.githubusercontent.com/MPAS-Dev/MPAS-Analysis/main/LICENSE
import os
import xarray
import numpy
import matplotlib.pyplot as plt

from geometric_features import FeatureCollection, read_feature_collection

from mpas_analysis.shared.analysis_task import AnalysisTask

from mpas_analysis.shared.constants import constants

from mpas_analysis.shared.plot import timeseries_analysis_plot, savefig, \
    add_inset

from mpas_analysis.shared.io import open_mpas_dataset, write_netcdf_with_fill

from mpas_analysis.shared.io.utility import build_config_full_path, \
    get_files_year_month, decode_strings

from mpas_analysis.shared.html import write_image_xml

from mpas_analysis.shared.transects import ComputeTransectMasksSubtask


[docs] class TimeSeriesTransport(AnalysisTask): """ Extract and plot time series of transport through transects on the MPAS mesh. """ # Authors # ------- # Xylar Asay-Davis, Stephen Price
[docs] def __init__(self, config, controlConfig=None): """ Construct the analysis task. Parameters ---------- config : mpas_tools.config.MpasConfigParser Configuration options controlconfig : mpas_tools.config.MpasConfigParser, optional Configuration options for a control run (if any) """ # Authors # ------- # Xylar Asay-Davis # first, call the constructor from the base class (AnalysisTask) super().__init__( config=config, taskName='timeSeriesTransport', componentName='ocean', tags=['timeSeries', 'transport']) startYear = config.getint('timeSeries', 'startYear') endYear = config.getint('timeSeries', 'endYear') transportGroups = config.getexpression('timeSeriesTransport', 'transportGroups') years = [year for year in range(startYear, endYear + 1)] for transportGroup in transportGroups: groupSuffix = transportGroup.replace(' ', '') transectsToPlot = config.getexpression(f'timeSeries{groupSuffix}', 'transectsToPlot') if len(transectsToPlot) == 0: return masksSubtask = ComputeTransectMasksSubtask( parentTask=self, transectGroup=transportGroup) transectsToPlot = masksSubtask.expand_transect_names(transectsToPlot) transportTransectFileName = masksSubtask.geojsonFileName self.add_subtask(masksSubtask) # in the end, we'll combine all the time series into one, but we # create this task first so it's easier to tell it to run after all # the compute tasks combineSubtask = CombineTransportSubtask( self, startYears=years, endYears=years, groupSuffix=groupSuffix) # run one subtask per year for year in years: computeSubtask = ComputeTransportSubtask( self, startYear=year, endYear=year, masksSubtask=masksSubtask, transectsToPlot=transectsToPlot, groupSuffix=groupSuffix) self.add_subtask(computeSubtask) computeSubtask.run_after(masksSubtask) combineSubtask.run_after(computeSubtask) for index, transect in enumerate(transectsToPlot): plotTransportSubtask = PlotTransportSubtask( self, transect, index, controlConfig, transportTransectFileName, transportGroup) plotTransportSubtask.run_after(combineSubtask) self.add_subtask(plotTransportSubtask)
class ComputeTransportSubtask(AnalysisTask): """ Computes time-series of transport through transects. Attributes ---------- startYear, endYear : int The beginning and end of the time series to compute masksSubtask : ``ComputeRegionMasksSubtask`` A task for creating mask files for each ice shelf to plot transectsToPlot : list of str A list of transects to plot groupSuffix : str standard transects vs Arctic transects """ # Authors # ------- # Xylar Asay-Davis, Stephen Price def __init__(self, parentTask, startYear, endYear, masksSubtask, transectsToPlot, groupSuffix): """ Construct the analysis task. Parameters ---------- parentTask : ``AnalysisTask`` The parent task, used to get the ``taskName``, ``config`` and ``componentName`` startYear, endYear : int The beginning and end of the time series to compute masksSubtask : ``ComputeRegionMasksSubtask`` A task for creating mask files for each ice shelf to plot transectsToPlot : list of str A list of transects to plot groupSuffix : str standard transects vs Arctic transects """ # Authors # ------- # Xylar Asay-Davis subtaskName = f'compute{groupSuffix}_{startYear:04d}-{endYear:04d}' # first, call the constructor from the base class (AnalysisTask) super().__init__( config=parentTask.config, taskName=parentTask.taskName, componentName=parentTask.componentName, tags=parentTask.tags, subtaskName=subtaskName) self.subprocessCount = self.config.getint(f'timeSeries{groupSuffix}', 'subprocessCount') self.startYear = startYear self.endYear = endYear self.masksSubtask = masksSubtask self.run_after(masksSubtask) self.transectsToPlot = transectsToPlot self.restartFileName = None self.groupSuffix = groupSuffix def setup_and_check(self): """ Perform steps to set up the analysis and check for errors in the setup. Raises ------ IOError If a restart file is not present ValueError If ``config_land_ice_flux_mode`` is not one of ``standalone`` or ``coupled`` """ # Authors # ------- # Xylar Asay-Davis # first, call setup_and_check from the base class (AnalysisTask), # which will perform some common setup, including storing: # self.inDirectory, self.plotsDirectory, self.namelist, self.streams # self.calendar super(ComputeTransportSubtask, self).setup_and_check() self.check_analysis_enabled( analysisOptionName='config_am_timeseriesstatsmonthly_enable', raiseException=True) # Load mesh related variables try: self.restartFileName = self.runStreams.readpath('restart')[0] except ValueError: raise IOError('No MPAS-O restart file found: need at least one ' 'restart file for transport calculations') 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 = f'{outputDirectory}/{self.groupSuffix}_{self.startYear:04d}-{self.endYear:04d}.nc' 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.sizes['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) dsMesh = dsMesh[['dvEdge', 'cellsOnEdge']] dsMesh.load() dvEdge = dsMesh.dvEdge cellsOnEdge = dsMesh.cellsOnEdge - 1 transectIndices = [] transectData = [] self.logger.info(' Caching transect data...') 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]) dsMask.load() edgeIndices = numpy.flatnonzero(dsMask.transectEdgeMasks.values) edgeIndices = edgeIndices[edgeIndices >= 0].astype(int) edgeSign = dsMask.transectEdgeMaskSigns.isel( nEdges=edgeIndices) dv = dvEdge.isel(nEdges=edgeIndices) coe = cellsOnEdge.isel(nEdges=edgeIndices) transectData.append({'edgeIndices': edgeIndices, 'edgeSign': edgeSign, 'dv': dv, 'coe': coe, 'transect': transect}) 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) dsTimeSlice.load() transectDatasets = [] for data in transectData: self.logger.info(' transect: {}'.format(data['transect'])) edgeIndices = data['edgeIndices'] coe = data['coe'] edgeSign = data['edgeSign'] dv = data['dv'] dsIn = dsTimeSlice.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=['nEdges', '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_with_fill(dsOut, outFileName) class CombineTransportSubtask(AnalysisTask): """ Combine individual time series into a single data set """ # Authors # ------- # Xylar Asay-Davis def __init__(self, parentTask, startYears, endYears, groupSuffix): """ Construct the analysis task. Parameters ---------- parentTask : ``TimeSeriesOceanRegions`` The main task of which this is a subtask startYears, endYears : list of int The beginning and end of each time series to combine groupSuffix : str standard transects vs Arctic transects """ # Authors # ------- # Xylar Asay-Davis # first, call the constructor from the base class (AnalysisTask) super(CombineTransportSubtask, self).__init__( config=parentTask.config, taskName=parentTask.taskName, componentName=parentTask.componentName, tags=parentTask.tags, subtaskName=f'combine{groupSuffix}TimeSeries') #print(self.taskName, self.subtaskName) self.startYears = startYears self.endYears = endYears self.groupSuffix = groupSuffix def run_task(self): """ Combine the time series """ # Authors # ------- # Xylar Asay-Davis groupSuffix = self.groupSuffix outputDirectory = '{}/transport/'.format( build_config_full_path(self.config, 'output', 'timeseriesSubdirectory')) outFileName = f'{outputDirectory}/{groupSuffix}_{self.startYears[0]:04d}-{self.endYears[-1]:04d}.nc' if not os.path.exists(outFileName): inFileNames = [] for startYear, endYear in zip(self.startYears, self.endYears): inFileName = f'{outputDirectory}/{groupSuffix}_{startYear:04d}-{endYear:04d}.nc' inFileNames.append(inFileName) ds = xarray.open_mfdataset(inFileNames, combine='nested', concat_dim='Time', decode_times=False) ds.load() write_netcdf_with_fill(ds, outFileName) class PlotTransportSubtask(AnalysisTask): """ Plots time-series output of transport through transects. Attributes ---------- transect : str Name of the transect to plot transectIndex : int The index into the dimension ``nTransects`` of the transect to plot controlConfig : mpas_tools.config.MpasConfigParser The configuration options for the control run (if any) transportGroup : str (with spaces) standard transects (``Transport Transects``) vs Arctic transects (``Arctic Transport Transects``) """ # Authors # ------- # Xylar Asay-Davis, Stephen Price def __init__(self, parentTask, transect, transectIndex, controlConfig, transportTransectFileName, transportGroup): """ Construct the analysis task. Parameters ---------- parentTask : ``AnalysisTask`` The parent task, used to get the ``taskName``, ``config`` and ``componentName`` transect : str Name of the transect to plot transectIndex : int The index into the dimension ``nTransects`` of the transect to plot controlconfig : mpas_tools.config.MpasConfigParser, optional Configuration options for a control run (if any) transportGroup : str (with spaces) standard transects (``Transport Transects``) vs Arctic transects (``Arctic Transport Transects``) """ # Authors # ------- # Xylar Asay-Davis # first, call the constructor from the base class (AnalysisTask) transectKey = transect.replace(' ', '_') super(PlotTransportSubtask, self).__init__( config=parentTask.config, taskName=parentTask.taskName, componentName=parentTask.componentName, tags=parentTask.tags, subtaskName=f'plotTransport_{transectKey}') self.transportTransectFileName = transportTransectFileName self.transect = transect self.transectIndex = transectIndex self.controlConfig = controlConfig self.transportGroup = transportGroup def setup_and_check(self): """ Perform steps to set up the analysis and check for errors in the setup. Raises ------ IOError If files are not present """ # Authors # ------- # Xylar Asay-Davis # first, call setup_and_check from the base class (AnalysisTask), # which will perform some common setup, including storing: # self.inDirectory, self.plotsDirectory, self.namelist, self.streams # self.calendar super(PlotTransportSubtask, self).setup_and_check() self.xmlFileNames = ['{}/transport_{}.xml'.format( self.plotsDirectory, self.transect.replace(' ', '_'))] def run_task(self): """ Plots time-series output of transport through transects. """ # Authors # ------- # Xylar Asay-Davis, Stephen Price self.logger.info("\nPlotting time series of transport through " "{}...".format(self.transect)) self.logger.info(' Load transport data...') obsDict = {'Drake Passage': [120, 175], 'Tasmania-Ant': [147, 167], 'Africa-Ant': None, 'Antilles Inflow': [-23.1, -13.7], 'Mona Passage': [-3.8, -1.4], 'Windward Passage': [-7.2, -6.8], 'Florida-Cuba': [30, 33], 'Florida-Bahamas': [30, 33], 'Indonesian Throughflow': [-21, -11], 'Agulhas': [-90, -50], 'Mozambique Channel': [-20, -8], 'Bering Strait': [0.6, 1.0], 'Lancaster Sound': [-1.0, -0.5], 'Fram Strait': [-4.7, 0.7], 'Davis Strait': [-1.6, -3.6], 'Barents Sea Opening': [1.4, 2.6], 'Nares Strait': [-1.8, 0.2], 'OSNAP section East': [15.6 - 0.8, 15.6 + 0.8], 'OSNAP section West': [2.1 - 0.3, 2.1 + 0.3], 'Denmark Strait': None, 'Iceland-Faroe-Scotland': None} config = self.config calendar = self.calendar fcAll = read_feature_collection(self.transportTransectFileName) fc = FeatureCollection() for feature in fcAll.features: if feature['properties']['name'] == self.transect: fc.add_feature(feature) break transport, trans_mean, trans_std = self._load_transport(config) if self.transect in obsDict: bounds = obsDict[self.transect] else: bounds = None plotControl = self.controlConfig is not None mainRunName = config.get('runs', 'mainRunName') groupSuffix = self.transportGroup.replace(' ','') movingAveragePoints = config.getint(f'timeSeries{groupSuffix}', 'movingAveragePoints') self.logger.info(' Plotting...') transectName = self.transect.replace('_', ' ') title = transectName thumbnailDescription = transectName xLabel = 'Time (yr)' yLabel = 'Transport (Sv)' filePrefix = 'transport_{}'.format(self.transect.replace(' ', '_')) outFileName = '{}/{}.png'.format(self.plotsDirectory, filePrefix) fields = [transport] lineColors = [config.get('timeSeries', 'mainColor')] lineWidths = [2.5] meanString = 'mean={:.2f} $\\pm$ {:.2f}'.format(trans_mean, trans_std) if plotControl: controlRunName = self.controlConfig.get('runs', 'mainRunName') ref_transport, ref_mean, ref_std = \ self._load_transport(self.controlConfig) refMeanString = f'mean={ref_mean:.2f} $\\pm$ {ref_std:.2f}' fields.append(ref_transport) lineColors.append(config.get('timeSeries', 'controlColor')) lineWidths.append(1.2) legendText = ['{} ({})'.format(mainRunName, meanString), '{} ({})'.format(controlRunName, refMeanString)] else: legendText = [mainRunName] title = '{} ({})'.format(title, meanString) if config.has_option(self.taskName, 'firstYearXTicks'): firstYearXTicks = config.getint(self.taskName, 'firstYearXTicks') else: firstYearXTicks = None if config.has_option(self.taskName, 'yearStrideXTicks'): yearStrideXTicks = config.getint(self.taskName, 'yearStrideXTicks') else: yearStrideXTicks = None fig = timeseries_analysis_plot(config, fields, calendar=calendar, title=title, xlabel=xLabel, ylabel=yLabel, movingAveragePoints=movingAveragePoints, lineColors=lineColors, lineWidths=lineWidths, legendText=legendText, firstYearXTicks=firstYearXTicks, yearStrideXTicks=yearStrideXTicks) if bounds is not None: t = transport.Time.values plt.gca().fill_between(t, bounds[0] * numpy.ones_like(t), bounds[1] * numpy.ones_like(t), alpha=0.3, label='observations') plt.legend(loc='lower left') # 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=2.0, height=2.0) savefig(outFileName, config) caption = 'Transport through the {} Transect'.format(transectName) write_image_xml( config=config, filePrefix=filePrefix, componentName='Ocean', componentSubdirectory='ocean', galleryGroup='Transport Time Series', groupLink='transporttime', thumbnailDescription=thumbnailDescription, imageDescription=caption, imageCaption=caption, gallery=self.transportGroup) def _load_transport(self, config): """ Reads transport time series for this transect """ # Authors # ------- # Xylar Asay-Davis groupSuffix = self.transportGroup.replace(' ', '') baseDirectory = build_config_full_path( config, 'output', 'timeSeriesSubdirectory') startYear = config.getint('timeSeries', 'startYear') endYear = config.getint('timeSeries', 'endYear') inFileName = f'{baseDirectory}/transport/{groupSuffix}_{startYear:04d}-{endYear:04d}.nc' dsIn = xarray.open_dataset(inFileName) transport = dsIn.transport.isel(nTransects=self.transectIndex) mean = transport.mean().values std = transport.std().values return transport, mean, std