Break horiz transect data into segments and nodes

This is helpful for defining transect cells and nodes later, once
we have a vertical coordinate.

conda_package/mpas_tools/viz/transects.py

@@ -137,6 +137,8 @@ def find_transect_cells_and_weights(lonTransect, latTransect, dsTris, dsMesh,
     interpCells = None
     cellWeights = None
 
+    nWeights = 6
+
     first = True
 
     dStart = 0.
@@ -253,8 +255,8 @@ def find_transect_cells_and_weights(lonTransect, latTransect, dsTris, dsMesh,
         nodeWeights = (_angular_distance(first=intersections, second=n1Inter) /
                        _angular_distance(first=n0Inter, second=n1Inter))
 
-        weights = numpy.zeros((len(trisInter), 6))
-        cellIndices = numpy.zeros((len(trisInter), 6), int)
+        weights = numpy.zeros((len(trisInter), nWeights))
+        cellIndices = numpy.zeros((len(trisInter), nWeights), int)
         for index in range(3):
             weights[:, index] = (nodeWeights *
                                  nodeCellWeights[trisInter, node0Inter, index])
@@ -282,7 +284,7 @@ def find_transect_cells_and_weights(lonTransect, latTransect, dsTris, dsMesh,
     # we need to figure out if the end points of the transect are in a triangle
     # and add tris, nodes, interpCells and cellWeights if so
 
-    if len(tris) > 2 and tris[0] != tris[1]:
+    if len(tris) >= 2 and tris[0] != tris[1]:
         # the starting point is in a triangle so we need to duplicate the first
         # entry in several fields to handle this
         tris = numpy.concatenate((tris[0:1], tris))
@@ -299,7 +301,7 @@ def find_transect_cells_and_weights(lonTransect, latTransect, dsTris, dsMesh,
         yOut = yOut[1:]
         zOut = zOut[1:]
 
-    if len(tris) > 2 and tris[-1] != tris[-2]:
+    if len(tris) >= 2 and tris[-1] != tris[-2]:
         # the end point is in a triangle so we need to add final entries or
         # duplicate the last entry in several fields to handle this
         tris = numpy.concatenate((tris, tris[-1:]))
@@ -316,44 +318,64 @@ def find_transect_cells_and_weights(lonTransect, latTransect, dsTris, dsMesh,
 
     assert(numpy.all(tris[0::2] == tris[1::2]))
 
+    tris = tris[0::2]
+
     lonOut, latOut = _cartesian_to_lon_lat(xOut, yOut, zOut, earth_radius,
                                            degrees)
 
-    dsOut = xarray.Dataset()
-    dsOut['xNode'] = ('nNodes', xOut)
-    dsOut['yNode'] = ('nNodes', yOut)
-    dsOut['zNode'] = ('nNodes', zOut)
-    dsOut['dNode'] = ('nNodes', dNode)
-    dsOut['lonNode'] = ('nNodes', lonOut)
-    dsOut['latNode'] = ('nNodes', latOut)
+    nSegments = len(xOut)//2
+    nBounds = 2
 
     cellIndices = dsTris.triCellIndices.values[tris]
+    nodes = nodes.reshape((nSegments, nBounds))
+    dNode = dNode.reshape((nSegments, nBounds))
 
-    dsOut['triangleIndices'] = ('nNodes', tris)
-    dsOut['triangleNodeIndices'] = ('nNodes', nodes)
-    dsOut['cellIndices'] = ('nNodes', cellIndices)
-    dsOut['interpCellIndices'] = (('nNodes', 'nWeights'), interpCells)
-    dsOut['interpCellWeights'] = (('nNodes', 'nWeights'), cellWeights)
-
-    transectIndicesOnNode = numpy.zeros(nodes.shape, int)
-    transectWeightsOnNode = numpy.zeros(nodes.shape)
+    dsOut = xarray.Dataset()
+    dsOut['xNode'] = (('nSegments', 'nBounds'),
+                      xOut.reshape((nSegments, nBounds)))
+    dsOut['yNode'] = (('nSegments', 'nBounds'),
+                      yOut.reshape((nSegments, nBounds)))
+    dsOut['zNode'] = (('nSegments', 'nBounds'),
+                      zOut.reshape((nSegments, nBounds)))
+    dsOut['dNode'] = (('nSegments', 'nBounds'), dNode)
+    dsOut['lonNode'] = (('nSegments', 'nBounds'),
+                        lonOut.reshape((nSegments, nBounds)))
+    dsOut['latNode'] = (('nSegments', 'nBounds'),
+                        latOut.reshape((nSegments, nBounds)))
+
+    dsOut['triangleIndices'] = ('nSegments', tris)
+    dsOut['cellIndices'] = ('nSegments', cellIndices)
+    dsOut['triangleNodeIndices'] = (('nSegments', 'nBounds'), nodes)
+    dsOut['interpCellIndices'] = \
+        (('nSegments', 'nBounds', 'nWeights'),
+         interpCells.reshape((nSegments, nBounds, nWeights)))
+    dsOut['interpCellWeights'] = \
+        (('nSegments', 'nBounds', 'nWeights'),
+         cellWeights.reshape((nSegments, nBounds, nWeights)))
+
+    transectIndicesOnNode = numpy.zeros(dNode.shape, int)
+    transectWeightsOnNode = numpy.zeros(dNode.shape)
     for segIndex in range(len(dTransect)-1):
         d0 = dTransect[segIndex]
         d1 = dTransect[segIndex+1]
         mask = numpy.logical_and(dNode >= d0, dNode < d1)
         transectIndicesOnNode[mask] = segIndex
         transectWeightsOnNode[mask] = (d1 - dNode[mask])/(d1 - d0)
-    transectIndicesOnNode[-1] = len(dTransect)-1
-    transectWeightsOnNode[-1] = 0.0
+    # last index will get missed by the mask and needs to be handled as a
+    # special case
+    transectIndicesOnNode[-1, 1] = len(dTransect)-1
+    transectWeightsOnNode[-1, 1] = 0.0
 
     dsOut['lonTransect'] = lonTransect
     dsOut['latTransect'] = latTransect
     dsOut['xTransect'] = x
     dsOut['yTransect'] = y
     dsOut['zTransect'] = z
     dsOut['dTransect'] = (lonTransect.dims, dTransect)
-    dsOut['transectIndicesOnNode'] = ('nNodes', transectIndicesOnNode)
-    dsOut['transectWeightsOnNode'] = ('nNodes', transectWeightsOnNode)
+    dsOut['transectIndicesOnNode'] = (('nSegments', 'nBounds'),
+                                      transectIndicesOnNode)
+    dsOut['transectWeightsOnNode'] = (('nSegments', 'nBounds'),
+                                      transectWeightsOnNode)
 
     return dsOut