Update transects to mpas_tools 1.0.0

Author: xylar

mpas_analysis/ocean/compute_transects_subtask.py

@@ -18,11 +18,15 @@
 from mpas_tools.viz import mesh_to_triangles
 from mpas_tools.transects import subdivide_great_circle, \
     cartesian_to_great_circle_distance
-from mpas_tools.viz.transects import find_transect_cells_and_weights, \
+from mpas_tools.viz.transect.horiz import (
+    find_spherical_transect_cells_and_weights,
     make_triangle_tree
-from mpas_tools.ocean.transects import find_transect_levels_and_weights, \
-    interp_mpas_to_transect_triangle_nodes, \
-    interp_transect_grid_to_transect_triangle_nodes
+)
+from mpas_tools.ocean.viz.transect.vert import (
+    find_transect_levels_and_weights,
+    interp_mpas_to_transect_nodes,
+    interp_transect_grid_to_transect_nodes
+)
 
 from mpas_analysis.shared.climatology import RemapMpasClimatologySubtask, \
     get_climatology_op_directory
@@ -161,8 +165,8 @@ def __init__(self, mpasClimatologyTask, parentTask, climatologyName,
         self.remap = self.obsDatasets.horizontalResolution != 'mpas'
         if self.obsDatasets.horizontalResolution == 'mpas' and \
                 self.verticalComparisonGridName != 'mpas':
-            raise ValueError('If the horizontal comparison grid is "mpas", the '
-                             'vertical grid must also be "mpas".')
+            raise ValueError('If the horizontal comparison grid is "mpas", '
+                             'the vertical grid must also be "mpas".')
 
     def setup_and_check(self):
         """
@@ -471,6 +475,14 @@ def _compute_mpas_transects(self, dsMesh):
 
         dsTris = mesh_to_triangles(dsMesh)
 
+        layerThickness = dsMesh.layerThickness
+        bottomDepth = dsMesh.bottomDepth
+        maxLevelCell = dsMesh.maxLevelCell - 1
+        if 'minLevelCell' in dsMesh:
+            minLevelCell = dsMesh.minLevelCell - 1
+        else:
+            minLevelCell = xr.zeros_like(maxLevelCell)
+
         triangleTree = make_triangle_tree(dsTris)
 
         for transectName in transectNames:
@@ -493,22 +505,30 @@ def _compute_mpas_transects(self, dsMesh):
                 else:
                     transectZ = None
 
-                dsMpasTransect = find_transect_cells_and_weights(
-                    dsTransect.lon, dsTransect.lat, dsTris, dsMesh,
-                    triangleTree, degrees=True)
+                dsMpasTransect = find_spherical_transect_cells_and_weights(
+                    lon_transect=dsTransect.lon,
+                    lat_transect=dsTransect.lat,
+                    ds_tris=dsTris,
+                    ds_mesh=dsMesh,
+                    tree=triangleTree,
+                    degrees=True)
 
                 dsMpasTransect = find_transect_levels_and_weights(
-                    dsMpasTransect, dsMesh.layerThickness,
-                    dsMesh.bottomDepth, dsMesh.maxLevelCell - 1,
-                    transectZ)
+                    ds_horiz_transect=dsMpasTransect,
+                    layer_thickness=layerThickness,
+                    bottom_depth=bottomDepth,
+                    min_level_cell=minLevelCell,
+                    max_level_cell=maxLevelCell,
+                    z_transect=transectZ)
 
                 if 'landIceFraction' in dsMesh:
                     interpCellIndices = dsMpasTransect.interpHorizCellIndices
                     interpCellWeights = dsMpasTransect.interpHorizCellWeights
                     landIceFraction = dsMesh.landIceFraction.isel(
                         nCells=interpCellIndices)
-                    landIceFraction = (landIceFraction * interpCellWeights).sum(
-                        dim='nHorizWeights')
+                    landIceFraction = (
+                        landIceFraction * interpCellWeights).sum(
+                            dim='nHorizWeights')
                     dsMpasTransect['landIceFraction'] = landIceFraction
 
                 # use to_netcdf rather than write_netcdf_with_fill because
@@ -517,9 +537,7 @@ def _compute_mpas_transects(self, dsMesh):
                 dsMpasTransect.to_netcdf(transectInfoFileName)
 
                 dsTransectOnMpas = xr.Dataset(dsMpasTransect)
-                dsTransectOnMpas['x'] = dsMpasTransect.dNode.isel(
-                    nSegments=dsMpasTransect.segmentIndices,
-                    nHorizBounds=dsMpasTransect.nodeHorizBoundsIndices)
+                dsTransectOnMpas['x'] = dsMpasTransect.dNode
 
                 dsTransectOnMpas['z'] = dsMpasTransect.zTransectNode
 
@@ -545,9 +563,11 @@ def _compute_mpas_transects(self, dsMesh):
                         dims = dsMask[var].dims
                         if 'nCells' in dims and 'nVertLevels' in dims:
                             dsOnMpas[var] = \
-                                interp_mpas_to_transect_triangle_nodes(
+                                interp_mpas_to_transect_nodes(
                                     dsMpasTransect, dsMask[var])
 
+                    dsOnMpas = self._transpose(dsOnMpas)
+
                     outFileName = self.get_remapped_file_name(
                         season, comparisonGridName=transectName)
                     dsOnMpas.to_netcdf(outFileName)
@@ -558,13 +578,37 @@ def _interp_obs_to_mpas(self, da, dsMpasTransect, threshold=0.1):
         """
         daMask = da.notnull()
         da = da.where(daMask, 0.)
-        da = interp_transect_grid_to_transect_triangle_nodes(
-            dsMpasTransect, da)
-        daMask = interp_transect_grid_to_transect_triangle_nodes(
-            dsMpasTransect, daMask)
+        da = interp_transect_grid_to_transect_nodes(
+            ds_transect=dsMpasTransect,
+            da=da)
+        daMask = interp_transect_grid_to_transect_nodes(
+            ds_transect=dsMpasTransect,
+            da=daMask)
         da = (da / daMask).where(daMask > threshold)
         return da
 
+    @staticmethod
+    def _transpose(dsOnMpas):
+        """
+        Transpose the data set to have the expected dimension order
+        """
+        dims = dsOnMpas.dims
+        dimsTransposed = ['nPoints', 'nz',
+                          'nSegments', 'nHalfLevels',
+                          'nHorizLevels', 'nVertLevels',
+                          'nHorizWeights', 'nVertWeights']
+
+        # drop any dimensions not in the dataset
+        dimsTransposed = [dim for dim in dimsTransposed if dim in
+                          dims]
+        # add any other dimensions at the end
+        for dim in dims:
+            if dim not in dimsTransposed:
+                dimsTransposed.append(dim)
+        dsOnMpas = dsOnMpas.transpose(*dimsTransposed)
+
+        return dsOnMpas
+
 
 class TransectsObservations(object):
     """
@@ -611,8 +655,9 @@ def __init__(self, config, obsFileNames, horizontalResolution,
             observations for a transect
 
         horizontalResolution : str
-            'obs' for the obs as they are, 'mpas' for the native MPAS mesh, or a
-            size in km if subdivision of the observational transect is desired.
+            'obs' for the obs as they are, 'mpas' for the native MPAS mesh, or
+            a size in km if subdivision of the observational transect is
+            desired.
 
         transectCollectionName : str
             A name that describes the collection of transects (e.g. the name

mpas_analysis/ocean/plot_transect_subtask.py

@@ -23,8 +23,6 @@
 
 from geometric_features import FeatureCollection
 
-from mpas_tools.ocean.transects import get_outline_segments
-
 from mpas_analysis.shared.plot import plot_vertical_section_comparison, \
     savefig, add_inset
 
@@ -397,18 +395,12 @@ def _plot_transect(self, remappedModelClimatology, remappedRefClimatology):
 
         else:
             x = 1e-3*remappedModelClimatology.dNode
-            z = None
+            z = remappedModelClimatology.zTransectNode
             lon = remappedModelClimatology.lonNode
             lat = remappedModelClimatology.latNode
 
             remappedModelClimatology['dNode'] = x
 
-            # flatten the x, lon and lat arrays because this is what
-            # vertical_section is expecting
-            x = xr.DataArray(data=x.values.ravel(), dims=('nx',))
-            lon = xr.DataArray(data=lon.values.ravel(), dims=('nx',))
-            lat = xr.DataArray(data=lat.values.ravel(), dims=('nx',))
-
         # This will do strange things at the antemeridian but there's little
         # we can do about that.
         lon_pm180 = numpy.mod(lon + 180., 360.) - 180.
@@ -437,12 +429,6 @@ def _plot_transect(self, remappedModelClimatology, remappedRefClimatology):
 
         modelOutput = remappedModelClimatology[self.mpasFieldName]
 
-        if remap:
-            triangulation_args = None
-        else:
-            triangulation_args = self._get_ds_triangulation(
-                remappedModelClimatology)
-
         if remappedRefClimatology is None:
             refOutput = None
             bias = None
@@ -578,12 +564,11 @@ def _plot_transect(self, remappedModelClimatology, remappedRefClimatology):
             configSectionName,
             xCoords=xs,
             zCoord=z,
-            triangulation_args=triangulation_args,
             colorbarLabel=self.unitsLabel,
             xlabels=xLabels,
             ylabel=yLabel,
             title=title,
-            modelTitle='{}'.format(mainRunName),
+            modelTitle=mainRunName,
             refTitle=self.refTitleLabel,
             diffTitle=self.diffTitleLabel,
             numUpperTicks=numUpperTicks,
@@ -694,7 +679,7 @@ def _lat_greater_extent(self, lat, lon):
             maxes = []
             last_idx = 0
 
-            while(len(lon_r) > 0 and len(lon_l) > 0):
+            while len(lon_r) > 0 and len(lon_l) > 0:
                 if lon_r[0] < lon_l[0]:
                     mins.append(numpy.min(lon[last_idx:lon_r[0]]))
                     last_idx = lon_r[0]
@@ -741,7 +726,8 @@ def _strictly_monotonic(self, coord):
         # Greg Streletz, Xylar Asay-Davis
 
         coord_diff = numpy.diff(coord.values)
-        coord_diff = numpy.where(coord_diff > 180, coord_diff - 360, coord_diff)
+        coord_diff = numpy.where(coord_diff > 180, coord_diff - 360,
+                                 coord_diff)
         coord_diff = numpy.where(coord_diff < -180, coord_diff + 360,
                                  coord_diff)
         return numpy.all(coord_diff > 0) or numpy.all(coord_diff < 0)
@@ -838,24 +824,6 @@ def _lat_fewest_direction_changes(self, lat, lon):
         else:
             return False
 
-    def _get_ds_triangulation(self, dsTransectTriangles):
-        """get matplotlib Triangulation from triangulation dataset"""
-
-        nTransectTriangles = dsTransectTriangles.sizes['nTransectTriangles']
-        dNode = dsTransectTriangles.dNode.isel(
-            nSegments=dsTransectTriangles.segmentIndices,
-            nHorizBounds=dsTransectTriangles.nodeHorizBoundsIndices)
-        x = dNode.values.ravel()
-
-        zTransectNode = dsTransectTriangles.zTransectNode
-        y = zTransectNode.values.ravel()
-
-        tris = numpy.arange(3 * nTransectTriangles).reshape(
-            (nTransectTriangles, 3))
-        triangulation_args = dict(x=x, y=y, triangles=tris)
-
-        return triangulation_args
-
     @staticmethod
     def _get_contour_colormap():
         # https://stackoverflow.com/a/18926541/7728169