Switch analysis scripts to open_multifile_datset

xylar · xylar · commit c42fb0e8dc8e · 2017-02-16T23:03:25.000+01:00
diff --git a/mpas_analysis/ocean/ocean_modelvsobs.py b/mpas_analysis/ocean/ocean_modelvsobs.py
@@ -11,20 +11,21 @@
 import matplotlib.pyplot as plt
 import matplotlib.colors as cols
 
-import numpy as np
 import xarray as xr
+import numpy as np
 import datetime
-from netCDF4 import Dataset as netcdf_dataset
+import netCDF4
 
-from ..shared.mpas_xarray.mpas_xarray import preprocess_mpas, \
-    remove_repeated_time_index
 from ..shared.plot.plotting import plot_global_comparison
 from ..shared.interpolation.interpolate import interp_fields, init_tree
 from ..shared.constants import constants
 
 from ..shared.io import NameList, StreamsFile
 from ..shared.io.utility import buildConfigFullPath
 
+from ..shared.generalized_reader.generalized_reader \
+    import open_multifile_dataset
+
 
 def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
 
@@ -85,7 +86,7 @@ def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
     sectionName = 'regridded{}'.format(field.upper())
     outputTimes = config.getExpression(sectionName, 'comparisonTimes')
 
-    ncFile = netcdf_dataset(restartFile, mode='r')
+    ncFile = netCDF4.Dataset(restartFile, mode='r')
     lonCell = ncFile.variables["lonCell"][:]
     latCell = ncFile.variables["latCell"][:]
     ncFile.close()
@@ -94,7 +95,7 @@ def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
 
     if field == 'mld':
 
-        selvals = None
+        iselvals = None
 
         # Load MLD observational data
         obsFileName = "{}/holtetalley_mld_climatology.nc".format(
@@ -122,7 +123,7 @@ def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
 
     elif field == 'sst':
 
-        selvals = {'nVertLevels': 0}
+        iselvals = {'nVertLevels': 0}
 
         obsFileName = \
             "{}/MODEL.SST.HAD187001-198110.OI198111-201203.nc".format(
@@ -154,7 +155,7 @@ def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
 
     elif field == 'sss':
 
-        selvals = {'nVertLevels': 0}
+        iselvals = {'nVertLevels': 0}
 
         obsFileName = "{}/Aquarius_V3_SSS_Monthly.nc".format(
                 observationsDirectory)
@@ -183,14 +184,15 @@ def ocn_modelvsobs(config, field, streamMap=None, variableMap=None):
         outFileLabel = 'sssAquarius'
         unitsLabel = 'PSU'
 
-    ds = xr.open_mfdataset(
-        inputFiles,
-        preprocess=lambda x: preprocess_mpas(x, yearoffset=yearOffset,
-                                             timestr='Time',
-                                             onlyvars=varList,
-                                             selvals=selvals,
-                                             varmap=variableMap))
-    ds = remove_repeated_time_index(ds)
+    ds = open_multifile_dataset(fileNames=inputFiles,
+                                calendar=calendar,
+                                timeVariableName='Time',
+                                variableList=varList,
+                                isel_values=iselvals,
+                                variableMap=variableMap,
+                                startDate=startDate,
+                                endDate=endDate,
+                                yearOffset=yearOffset)
 
     timeStart = datetime.datetime(yearOffset+startYear, 1, 1)
     timeEnd = datetime.datetime(yearOffset+endYear, 12, 31)
diff --git a/mpas_analysis/ocean/ohc_timeseries.py b/mpas_analysis/ocean/ohc_timeseries.py
@@ -1,18 +1,16 @@
 import numpy as np
 import netCDF4
-from netCDF4 import Dataset as netcdf_dataset
-import xarray as xr
 import pandas as pd
 import datetime
 
-from ..shared.mpas_xarray.mpas_xarray import preprocess_mpas, \
-    remove_repeated_time_index
-
 from ..shared.plot.plotting import timeseries_analysis_plot
 
 from ..shared.io import NameList, StreamsFile
 from ..shared.io.utility import buildConfigFullPath
 
+from ..shared.generalized_reader.generalized_reader \
+    import open_multifile_dataset
+
 from ..shared.timekeeping.utility import stringToDatetime, \
     clampToNumpyDatetime64
 
@@ -33,7 +31,7 @@ def ohc_timeseries(config, streamMap=None, variableMap=None):
     to their mpas_analysis counterparts.
 
     Author: Xylar Asay-Davis, Milena Veneziani
-    Last Modified: 02/02/2017
+    Last Modified: 02/08/2017
     """
 
     inDirectory = config.get('input', 'baseDirectory')
@@ -82,13 +80,13 @@ def ohc_timeseries(config, streamMap=None, variableMap=None):
     startDate = config.get('timeSeries', 'startDate')
     endDate = config.get('timeSeries', 'endDate')
     streamName = streams.find_stream(streamMap['timeSeriesStats'])
-    inFiles = streams.readpath(streamName, startDate=startDate,
-                               endDate=endDate, calendar=calendar)
-    print 'Reading files {} through {}'.format(inFiles[0], inFiles[-1])
+    fileNames = streams.readpath(streamName, startDate=startDate,
+                                 endDate=endDate, calendar=calendar)
+    print 'Reading files {} through {}'.format(fileNames[0], fileNames[-1])
 
     # Define/read in general variables
     print '  Read in depth and compute specific depth indexes...'
-    ncFile = netcdf_dataset(restartFile, mode='r')
+    ncFile = netCDF4.Dataset(restartFile, mode='r')
     # reference depth [m]
     depth = ncFile.variables['refBottomDepth'][:]
     # simulation start time
@@ -113,20 +111,17 @@ def ohc_timeseries(config, streamMap=None, variableMap=None):
                     'sumLayerMaskValue',
                     'avgLayerArea',
                     'avgLayerThickness']
-    ds = xr.open_mfdataset(
-        inFiles,
-        preprocess=lambda x: preprocess_mpas(x,
-                                             yearoffset=yearOffset,
-                                             timestr='Time',
-                                             onlyvars=variableList,
-                                             varmap=variableMap))
-
-    ds = remove_repeated_time_index(ds)
+    ds = open_multifile_dataset(fileNames=fileNames,
+                                calendar=calendar,
+                                timeVariableName='Time',
+                                variableList=variableList,
+                                variableMap=variableMap,
+                                startDate=startDate,
+                                endDate=endDate,
+                                yearOffset=yearOffset)
 
     timeStart = clampToNumpyDatetime64(stringToDatetime(startDate), yearOffset)
     timeEnd = clampToNumpyDatetime64(stringToDatetime(endDate), yearOffset)
-    # select only the data in the specified range of years
-    ds = ds.sel(Time=slice(timeStart, timeEnd))
 
     # Select year-1 data and average it (for later computing anomalies)
     timeStartFirstYear = clampToNumpyDatetime64(
@@ -139,15 +134,14 @@ def ohc_timeseries(config, streamMap=None, variableMap=None):
                                           startDate=startDateFirstYear,
                                           endDate=endDateFirstYear,
                                           calendar=calendar)
-        dsFirstYear = xr.open_mfdataset(
-            filesFirstYear,
-            preprocess=lambda x: preprocess_mpas(x,
-                                                 yearoffset=yearOffset,
-                                                 timestr='Time',
-                                                 onlyvars=variableList,
-                                                 varmap=variableMap))
-
-        dsFirstYear = remove_repeated_time_index(dsFirstYear)
+        dsFirstYear = open_multifile_dataset(fileNames=filesFirstYear,
+                                             calendar=calendar,
+                                             timeVariableName='Time',
+                                             variableList=variableList,
+                                             variableMap=variableMap,
+                                             startDate=startDateFirstYear,
+                                             endDate=endDateFirstYear,
+                                             yearOffset=yearOffset)
         firstYear += yearOffset
     else:
         dsFirstYear = ds
@@ -174,12 +168,12 @@ def ohc_timeseries(config, streamMap=None, variableMap=None):
 
     if preprocessedReferenceRunName != 'None':
         print '  Load in OHC from preprocessed reference run...'
-        inFilesPreprocesses = '{}/OHC.{}.year*.nc'.format(
+        inFilesPreprocessed = '{}/OHC.{}.year*.nc'.format(
             preprocessedInputDirectory, preprocessedReferenceRunName)
-        dsPreprocessed = xr.open_mfdataset(
-            inFilesPreprocesses,
-            preprocess=lambda x: preprocess_mpas(x, yearoffset=yearOffset))
-        dsPreprocessed = remove_repeated_time_index(dsPreprocessed)
+        dsPreprocessed = open_multifile_dataset(fileNames=inFilesPreprocessed,
+                                                calendar=calendar,
+                                                timeVariableName='xtime',
+                                                yearOffset=yearOffset)
         yearEndPreprocessed = \
             (pd.to_datetime(dsPreprocessed.Time.max().values)).year
         if yearStart <= yearEndPreprocessed:
diff --git a/mpas_analysis/ocean/sst_timeseries.py b/mpas_analysis/ocean/sst_timeseries.py
@@ -1,15 +1,14 @@
-import xarray as xr
 import pandas as pd
 import datetime
 
-from ..shared.mpas_xarray.mpas_xarray import preprocess_mpas, \
-    remove_repeated_time_index
-
 from ..shared.plot.plotting import timeseries_analysis_plot
 
 from ..shared.io import NameList, StreamsFile
 from ..shared.io.utility import buildConfigFullPath
 
+from ..shared.generalized_reader.generalized_reader \
+    import open_multifile_dataset
+
 from ..shared.timekeeping.utility import stringToDatetime, \
     clampToNumpyDatetime64
 
@@ -50,9 +49,9 @@ def sst_timeseries(config, streamMap=None, variableMap=None):
     startDate = config.get('timeSeries', 'startDate')
     endDate = config.get('timeSeries', 'endDate')
     streamName = streams.find_stream(streamMap['timeSeriesStats'])
-    inFiles = streams.readpath(streamName, startDate=startDate,
-                               endDate=endDate,  calendar=calendar)
-    print 'Reading files {} through {}'.format(inFiles[0], inFiles[-1])
+    fileNames = streams.readpath(streamName, startDate=startDate,
+                                 endDate=endDate,  calendar=calendar)
+    print 'Reading files {} through {}'.format(fileNames[0], fileNames[-1])
 
     mainRunName = config.get('runs', 'mainRunName')
     preprocessedReferenceRunName = config.get('runs',
@@ -73,13 +72,14 @@ def sst_timeseries(config, streamMap=None, variableMap=None):
 
     # Load data:
     varList = ['avgSurfaceTemperature']
-    ds = xr.open_mfdataset(
-        inFiles,
-        preprocess=lambda x: preprocess_mpas(x, yearoffset=yearOffset,
-                                             timestr='Time',
-                                             onlyvars=varList,
-                                             varmap=variableMap))
-    ds = remove_repeated_time_index(ds)
+    ds = open_multifile_dataset(fileNames=fileNames,
+                                calendar=calendar,
+                                timeVariableName='Time',
+                                variableList=varList,
+                                variableMap=variableMap,
+                                startDate=startDate,
+                                endDate=endDate,
+                                yearOffset=yearOffset)
 
     timeStart = clampToNumpyDatetime64(stringToDatetime(startDate), yearOffset)
     timeEnd = clampToNumpyDatetime64(stringToDatetime(endDate), yearOffset)
@@ -97,10 +97,10 @@ def sst_timeseries(config, streamMap=None, variableMap=None):
         print '  Load in SST for a preprocesses reference run...'
         inFilesPreprocessed = '{}/SST.{}.year*.nc'.format(
             preprocessedInputDirectory, preprocessedReferenceRunName)
-        dsPreprocessed = xr.open_mfdataset(
-            inFilesPreprocessed,
-            preprocess=lambda x: preprocess_mpas(x, yearoffset=yearOffset))
-        dsPreprocessed = remove_repeated_time_index(dsPreprocessed)
+        dsPreprocessed = open_multifile_dataset(fileNames=inFilesPreprocessed,
+                                                calendar=calendar,
+                                                timeVariableName='xtime',
+                                                yearOffset=yearOffset)
         yearEndPreprocessed = \
             (pd.to_datetime(dsPreprocessed.Time.max().values)).year
         if yearStart <= yearEndPreprocessed:
diff --git a/mpas_analysis/sea_ice/modelvsobs.py b/mpas_analysis/sea_ice/modelvsobs.py
@@ -6,18 +6,18 @@
 
 import numpy as np
 import numpy.ma as ma
-import xarray as xr
 import datetime
 
-from netCDF4 import Dataset as netcdf_dataset
+import netCDF4
 
-from ..shared.mpas_xarray.mpas_xarray import preprocess_mpas, \
-    remove_repeated_time_index
 from ..shared.plot.plotting import plot_polar_comparison
 
 from ..shared.io import NameList, StreamsFile
 from ..shared.io.utility import buildConfigFullPath
 
+from ..shared.generalized_reader.generalized_reader \
+    import open_multifile_dataset
+
 
 def seaice_modelvsobs(config, streamMap=None, variableMap=None):
     """
@@ -53,9 +53,9 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
     startDate = config.get('climatology', 'startDate')
     endDate = config.get('climatology', 'endDate')
     streamName = streams.find_stream(streamMap['timeSeriesStats'])
-    infiles = streams.readpath(streamName, startDate=startDate,
-                               endDate=endDate, calendar=calendar)
-    print 'Reading files {} through {}'.format(infiles[0], infiles[-1])
+    fileNames = streams.readpath(streamName, startDate=startDate,
+                                 endDate=endDate, calendar=calendar)
+    print 'Reading files {} through {}'.format(fileNames[0], fileNames[-1])
 
     plotsDirectory = buildConfigFullPath(config, 'output', 'plotsSubdirectory')
 
@@ -158,14 +158,15 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
 
     # Load data
     print "  Load sea-ice data..."
-    ds = xr.open_mfdataset(
-        infiles,
-        preprocess=lambda x: preprocess_mpas(x, yearoffset=yearOffset,
-                                             timestr='Time',
-                                             onlyvars=['iceAreaCell',
-                                                       'iceVolumeCell'],
-                                             varmap=variableMap))
-    ds = remove_repeated_time_index(ds)
+    ds = open_multifile_dataset(fileNames=fileNames,
+                                calendar=calendar,
+                                timeVariableName='Time',
+                                variableList=['iceAreaCell',
+                                              'iceVolumeCell'],
+                                variableMap=variableMap,
+                                startDate=startDate,
+                                endDate=endDate,
+                                yearOffset=yearOffset)
 
     # Compute climatologies (first motnhly and then seasonally)
     print "  Compute seasonal climatologies..."
@@ -254,7 +255,7 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
             # ice concentrations
             fileName = "{}/{}".format(climatologyRegriddedDirectory,
                                       climatologyFiles[climName])
-            ncFile = netcdf_dataset(fileName, mode='r')
+            ncFile = netCDF4.Dataset(fileName, mode='r')
             iceConcentration = ncFile.variables["iceAreaCell"][:]
             if(first):
                 lons = ncFile.variables["lon"][:]
@@ -271,7 +272,7 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
 
                 fileName = obsIceConcFileNames[
                     '{}_{}'.format(climName, obsName)]
-                ncFile = netcdf_dataset(fileName, mode='r')
+                ncFile = netCDF4.Dataset(fileName, mode='r')
                 obsIceConcentration = ncFile.variables["AICE"][:]
                 ncFile.close()
 
@@ -319,7 +320,7 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
         # ice concentrations
         fileName = "{}/{}".format(climatologyRegriddedDirectory,
                                   climatologyFiles[climName])
-        ncFile = netcdf_dataset(fileName, mode='r')
+        ncFile = netCDF4.Dataset(fileName, mode='r')
         iceThickness = ncFile.variables["iceVolumeCell"][:]
         ncFile.close()
 
@@ -362,7 +363,7 @@ def seaice_modelvsobs(config, streamMap=None, variableMap=None):
 
             fileName = obsIceThickFileNames['{}{}'.format(climName,
                                                           hemisphere)]
-            ncFile = netcdf_dataset(fileName, mode='r')
+            ncFile = netCDF4.Dataset(fileName, mode='r')
             obsIceThickness = ncFile.variables["HI"][:]
             ncFile.close()
             # Mask thickness fields
diff --git a/mpas_analysis/sea_ice/timeseries.py b/mpas_analysis/sea_ice/timeseries.py
