Labels on SkewT lines and legend in hodograph

adb_graphics/figures/skewt.py

@@ -11,13 +11,15 @@
 import matplotlib.font_manager as fm
 import matplotlib.pyplot as plt
 from matplotlib.ticker import FixedLocator
+from matplotlib.lines import Line2D
 import metpy.calc as mpcalc
 from metpy.plots import Hodograph, SkewT
 from metpy.units import units
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 
 import adb_graphics.datahandler.grib as grib
 import adb_graphics.errors as errors
+import adb_graphics.utils as utils
 
 class SkewTDiagram(grib.profileData):
 
@@ -63,13 +65,13 @@ def _add_thermo_inset(self, skew):
 
             # Sure would have been nice to use a variable in the f string to
             # denote the format per variable.
-            line = f"{name.upper():<7s}: {str(value):>10} {items['units']}"
+            line = f"{name.upper():<7s} {str(value):>6} {items['units']}"
             lines.append(line)
 
         contents = '\n'.join(lines)
 
         # Draw the text box
-        skew.ax.text(0.70, 0.98, contents,
+        skew.ax.text(0.78, 0.98, contents,
                      bbox=dict(facecolor='white', edgecolor='black', alpha=0.7),
                      fontproperties=fm.FontProperties(family='monospace'),
                      size=8,
@@ -170,40 +172,44 @@ def _plot_hodograph(self, skew):
         #
         agl = np.copy(self.atmo_profiles.get('gh', {}).get('data')).to('km')
 
-        heights = [10, 3, 1]
-        agl_arr = agl.magnitude
-        for i, height in enumerate(heights):
-
-            mag_top = height
-            mag_bottom = 0 if i >= len(heights) - 1 else heights[i+1]
-
-            # Use exclude later to remove values above 10km
-            if i == 0:
-                exclude = -np.sum(agl_arr > mag_top)
-
-            # Check for the values between two levels
-            condition = np.logical_and(agl_arr <= mag_top, agl_arr > mag_bottom)
-            agl.magnitude[condition] = mag_top
-
-        # Note: agl is now an array with values corresponding to the heights
-        # array
-
         # Retrieve the wind data profiles
         u_wind = self.atmo_profiles.get('u', {}).get('data')
         v_wind = self.atmo_profiles.get('v', {}).get('data')
 
-        # Drop the points above 10 km
-        u_wind = u_wind.magnitude[:exclude] * u_wind.units
-        v_wind = v_wind.magnitude[:exclude] * v_wind.units
-
         # Create an inset axes object that is 28% width and height of the
         # figure and put it in the upper left hand corner.
         ax = inset_axes(skew.ax, '25%', '25%', loc=2)
         h = Hodograph(ax, component_range=80.)
         h.add_grid(increment=20, linewidth=0.5)
 
+        intervals = [0, 1, 3, 10] * agl.units
+        colors = ['xkcd:salmon', 'xkcd:aquamarine', 'xkcd:navy blue']
+        line_width = 1.5
+
         # Plot the line colored by height AGL only up to the 10km level
-        h.plot_colormapped(u_wind, v_wind, agl[:exclude], linewidth=2)
+        lines = h.plot_colormapped(u_wind, v_wind, agl,
+                                   colors=colors,
+                                   intervals=intervals,
+                                   linewidth=line_width,
+                                   )
+
+        # Local function to create a proxy line object for creating a legend on
+        # a LineCollection returned from plot_colormapped. Using lines and
+        # colors from outside scope.
+        def make_proxy(zval, idx=None, **kwargs):
+            color = colors[idx] if idx < len(colors) else lines.cmap(zval-1)
+            return Line2D([0, 1], [0, 1], color=color, linewidth=line_width, **kwargs)
+
+        # Make a list of proxies
+        proxies = [make_proxy(item, idx=i) for i, item in
+                   enumerate(intervals.magnitude)]
+
+        # Draw the legend
+        ax.legend(proxies[:-1],
+                  ['0-1 km', '1-3 km', '3-10 km', ''],
+                  fontsize='small',
+                  loc='lower left',
+                  )
 
     @staticmethod
     def _plot_labels(skew):
@@ -289,13 +295,19 @@ def _setup_diagram(self):
                             which='major',
                             )
 
-        # Add the relevant special lines
+        # Add the relevant special lines with their labels
         dry_adiabats = np.arange(-40, 210, 10) * units.degC
         skew.plot_dry_adiabats(dry_adiabats,
                                colors='tan',
                                linestyles='solid',
                                linewidth=0.7,
                                )
+        utils.label_lines(ax=skew.ax,
+                          lines=skew.dry_adiabats,
+                          labels=dry_adiabats.magnitude,
+                          end='top',
+                          offset=1,
+                          )
 
         moist_adiabats = np.arange(8, 36, 4) * units.degC
         moist_pr = np.arange(1001, 220, -10) * units.hPa
@@ -305,6 +317,11 @@ def _setup_diagram(self):
                                  linestyles='solid',
                                  linewidth=0.7,
                                  )
+        utils.label_lines(ax=skew.ax,
+                          lines=skew.moist_adiabats,
+                          labels=moist_adiabats.magnitude,
+                          end='top',
+                          )
 
         mixing_lines = np.array([1, 2, 3, 5, 8, 12, 16, 20]).reshape(-1, 1)  / 1000
         mix_pr = np.arange(1001, 400, -50) * units.hPa
@@ -313,6 +330,11 @@ def _setup_diagram(self):
                                linestyles=(0, (5, 10)),
                                linewidth=0.7,
                                )
+        utils.label_lines(ax=skew.ax,
+                          lines=skew.mixing_lines,
+                          labels=mixing_lines * 1000,
+                          )
+
         return skew
 
     @property

adb_graphics/utils.py

@@ -56,47 +56,52 @@ def get_func(val: str):
 
 
 # pylint: disable=invalid-name, too-many-locals
-def label_line(line, x, label=None, align=True, **kwargs):
+def label_line(ax, label, segment, align=True, end='bottom', offset=0, **kwargs):
 
-    ''' Label line with line2D label data '''
+    '''
+    Label a single line with line2D label data.
+
+    Input:
 
-    ax = line.axes
-    xdata = line.get_xdata()
-    ydata = line.get_ydata()
+    Input:
 
-    if (x < xdata[0]) or (x > xdata[-1]):
-        print('x label location is outside data range!')
-        return
+      ax        the SkewT object axis
+      label     label to be used for the current line
+      segment   a list (array) of values for the current line
+      align     optional bool to enable the rotation of the label to line angle
+      end       the end of the line at which to put the label. 'bottom' or 'top'
+      offset    index to use for the "end" of the array
 
-    #Find corresponding y co-ordinate and angle of the line
-    ip = 1
-    for i, xd in enumerate(xdata):
-        if x < xd:
-            ip = i
-            break
+    Key Word Arguments
 
-    y = ydata[ip-1] + (ydata[ip]-ydata[ip-1])*(x-xdata[ip-1])/(xdata[ip]-xdata[ip-1])
+      Any kwargs accepted by matplotlib's text box.
+    '''
 
-    if not label:
-        label = line.get_label()
+    # Label location
+    if end == 'bottom':
+        x, y = segment[0 + offset, :]
+        ip = 1
+    elif end == 'top':
+        x, y = segment[-1 - offset, :]
+        ip = -1
 
     if align:
         #Compute the slope
-        dx = xdata[ip] - xdata[ip-1]
-        dy = ydata[ip] - ydata[ip-1]
+        dx = segment[ip, 0] - segment[ip-1, 0]
+        dy = segment[ip, 1] - segment[ip-1, 1]
         ang = degrees(atan2(dy, dx))
 
         #Transform to screen co-ordinates
         pt = np.array([x, y]).reshape((1, 2))
         trans_angle = ax.transData.transform_angles(np.array((ang, )), pt)[0]
 
+        if end == 'top':
+            trans_angle -= 180
+
     else:
         trans_angle = 0
 
     #Set a bunch of keyword arguments
-    if 'color' not in kwargs:
-        kwargs['color'] = line.get_color()
-
     if ('horizontalalignment' not in kwargs) and ('ha' not in kwargs):
         kwargs['ha'] = 'center'
 
@@ -109,36 +114,45 @@ def label_line(line, x, label=None, align=True, **kwargs):
     if 'clip_on' not in kwargs:
         kwargs['clip_on'] = True
 
+    if 'fontsize' not in kwargs:
+        kwargs['fontsize'] = 'larger'
+
+    if 'fontweight' not in kwargs:
+        kwargs['fontweight'] = 'bold'
+
+    # Larger value (e.g., 2.0) to move box in front of other diagram elements
     if 'zorder' not in kwargs:
-        kwargs['zorder'] = 2.5
+        kwargs['zorder'] = 1.50
 
+    # Place the text box label on the line.
     ax.text(x, y, label, rotation=trans_angle, **kwargs)
 
-def label_lines(lines, align=True, xvals=None, **kwargs):
+def label_lines(ax, lines, labels, offset=0, **kwargs):
 
     '''
-    retrieved from
-    https://stackoverflow.com/questions/16992038/inline-labels-in-matplotlib
-    '''
+    Plots labels on a set of lines from SkewT.
+
+    Input:
+
+      ax      the SkewT object axis
+      lines   the SkewT object special lines
+      labels  list of labels to be used 
+      offset  index to use for the "end" of the array
 
-    #ax = lines[0].axes
-    ax = kwargs.get('ax', lines.axes)
-    labLines = []
-    labels = []
+    Key Word Arguments
 
-    #Take only the lines which have labels other than the default ones
-    for line in lines.get_segments():
-        label = line.get_label()
-        if "_line" not in label:
-            labLines.append(line)
-            labels.append(label)
+      color   line color
 
-    if xvals is None:
-        xmin, xmax = ax.get_xlim()
-        xvals = np.linspace(xmin, xmax, len(labLines)+2)[1:-1]
+      Along with any other kwargs accepted by matplotlib's text box.
+    '''
+
+    if 'color' not in kwargs:
+        kwargs['color'] = lines.get_color()[0]
 
-    for line, x, label in zip(labLines, xvals, labels):
-        label_line(line, x, label, align, **kwargs)
+    for i, line in enumerate(lines.get_segments()):
+        label = int(labels[i])
+        label = label[0] if isinstance(label, list) else label
+        label_line(ax, label, line, align=True, offset=offset, **kwargs)
 
 def timer(func):
 

tests/test_common.py

@@ -9,7 +9,7 @@
 
 To run the tests, type the following in the top level repo directory:
 
-    python -m pytest --grib-file [path/to/gribfile]
+    python -m pytest --nat-file [path/to/gribfile] --prs-file [path/to/gribfile]
 
 '''