Skip to content

Commit

Permalink
add new Arctic transects from OSNAP + separate all Arctic transects f…
Browse files Browse the repository at this point in the history 
…rom standard
  • Loading branch information
@anirban89
anirban89 committed Sep 27, 2022
1 parent 169ada5 commit a9f8814
Showing 1 changed file with 82 additions and 61 deletions.
143 changes: 82 additions & 61 deletions mpas_analysis/ocean/time_series_transport.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,14 +60,16 @@ def __init__(self, config, controlConfig=None):
# Xylar Asay-Davis

# first, call the constructor from the base class (AnalysisTask)
super(TimeSeriesTransport, self).__init__(
super().__init__(
config=config,
taskName='timeSeriesTransport',
componentName='ocean',
tags=['timeSeries', 'transport'])

startYear = config.getint('timeSeries', 'startYear')
endYear = config.get('timeSeries', 'endYear')
transportGroups = config.getexpression('timeSeriesTransport', 'transportGroups')

if endYear == 'end':
# a valid end year wasn't found, so likely the run was not found,
# perhaps because we're just listing analysis tasks
Expand All @@ -76,40 +78,42 @@ def __init__(self, config, controlConfig=None):
endYear = int(endYear)

years = [year for year in range(startYear, endYear + 1)]

transectsToPlot = config.getexpression('timeSeriesTransport',
'transectsToPlot')
if len(transectsToPlot) == 0:
return

masksSubtask = ComputeTransectMasksSubtask(
parentTask=self, transectGroup='Transport Transects')

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)

# run one subtask per year
for year in years:
computeSubtask = ComputeTransportSubtask(
self, startYear=year, endYear=year, masksSubtask=masksSubtask,
transectsToPlot=transectsToPlot)
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)
plotTransportSubtask.run_after(combineSubtask)
self.add_subtask(plotTransportSubtask)
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):
Expand All @@ -126,14 +130,17 @@ class ComputeTransportSubtask(AnalysisTask):
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):
masksSubtask, transectsToPlot, groupSuffix):
"""
Construct the analysis task.
Expand All @@ -151,21 +158,23 @@ def __init__(self, parentTask, startYear, endYear,
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(ComputeTransportSubtask, self).__init__(
super().__init__(
config=parentTask.config,
taskName=parentTask.taskName,
componentName=parentTask.componentName,
tags=parentTask.tags,
subtaskName='computeTransport_{:04d}-{:04d}'.format(startYear,
endYear))

self.subprocessCount = self.config.getint('timeSeriesTransport',
subtaskName=subtaskName)

self.subprocessCount = self.config.getint(f'timeSeries{groupSuffix}',
'subprocessCount')
self.startYear = startYear
self.endYear = endYear
Expand All @@ -175,6 +184,7 @@ def __init__(self, parentTask, startYear, endYear,

self.transectsToPlot = transectsToPlot
self.restartFileName = None
self.groupSuffix = groupSuffix

def setup_and_check(self):
"""
Expand Down Expand Up @@ -233,8 +243,7 @@ def run_task(self):
except OSError:
pass

outFileName = '{}/transport_{:04d}-{:04d}.nc'.format(
outputDirectory, self.startYear, self.endYear)
outFileName = f'{outputDirectory}/{self.groupSuffix}_{self.startYear:04d}-{self.endYear:04d}.nc'

inputFiles = sorted(self.historyStreams.readpath(
'timeSeriesStatsMonthlyOutput', startDate=startDate,
Expand Down Expand Up @@ -395,7 +404,7 @@ class CombineTransportSubtask(AnalysisTask):
# -------
# Xylar Asay-Davis

def __init__(self, parentTask, startYears, endYears):
def __init__(self, parentTask, startYears, endYears, groupSuffix):
"""
Construct the analysis task.
Expand All @@ -407,21 +416,25 @@ def __init__(self, parentTask, startYears, endYears):
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='combineTimeSeries')

subtaskName=f'combine{groupSuffix}TimeSeries')
#print(self.taskName, self.subtaskName)
self.startYears = startYears
self.endYears = endYears
self.groupSuffix = groupSuffix

def run_task(self):
"""
Expand All @@ -430,19 +443,17 @@ def run_task(self):
# 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'

outFileName = '{}/transport_{:04d}-{:04d}.nc'.format(
outputDirectory, self.startYears[0], self.endYears[-1])

if not os.path.exists(outFileName):
inFileNames = []
for startYear, endYear in zip(self.startYears, self.endYears):
inFileName = '{}/transport_{:04d}-{:04d}.nc'.format(
outputDirectory, startYear, endYear)
inFileName = f'{outputDirectory}/{groupSuffix}_{startYear:04d}-{endYear:04d}.nc'
inFileNames.append(inFileName)

ds = xarray.open_mfdataset(inFileNames, combine='nested',
Expand All @@ -466,14 +477,18 @@ class PlotTransportSubtask(AnalysisTask):
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):
transportTransectFileName, transportGroup):

"""
Construct the analysis task.
Expand All @@ -492,23 +507,29 @@ def __init__(self, parentTask, transect, transectIndex, controlConfig,
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='plotTransport_{}'.format(transect.replace(' ', '_')))
subtaskName=f'plotTransport_{transectKey}')

self.transportTransectFileName = transportTransectFileName
self.transect = transect
self.transectIndex = transectIndex
self.controlConfig = controlConfig
self.transportGroup = transportGroup

def setup_and_check(self):
"""
Expand Down Expand Up @@ -586,7 +607,8 @@ def run_task(self):
plotControl = self.controlConfig is not None

mainRunName = config.get('runs', 'mainRunName')
movingAveragePoints = config.getint('timeSeriesTransport',
groupSuffix = self.transportGroup.replace(' ','')
movingAveragePoints = config.getint(f'timeSeries{groupSuffix}',
'movingAveragePoints')

self.logger.info(' Plotting...')
Expand Down Expand Up @@ -667,7 +689,8 @@ def run_task(self):
groupLink='transporttime',
thumbnailDescription=thumbnailDescription,
imageDescription=caption,
imageCaption=caption)
imageCaption=caption,
gallery=self.transportGroup)

def _load_transport(self, config):
"""
Expand All @@ -676,15 +699,13 @@ def _load_transport(self, config):
# 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 = '{}/transport/transport_{:04d}-{:04d}.nc'.format(
baseDirectory, startYear, endYear)
inFileName = f'{baseDirectory}/transport/{groupSuffix}_{startYear:04d}-{endYear:04d}.nc'

dsIn = xarray.open_dataset(inFileName)
transport = dsIn.transport.isel(nTransects=self.transectIndex)
Expand Down

0 comments on commit a9f8814

Please sign in to comment.