New PlotSurfaceAnalysis class to plot WPC surface analysis bulletins

src/metpy/plots/declarative.py

@@ -20,8 +20,9 @@
 
 from . import ctables, wx_symbols
 from ._mpl import TextCollection
+from .patheffects import ColdFront, OccludedFront, WarmFront, StationaryFront
 from .cartopy_utils import import_cartopy
 from .station_plot import StationPlot
 from ..calc import reduce_point_density, smooth_n_point, zoom_xarray
 from ..package_tools import Exporter
 from ..units import units
@@ -1986,3 +1987,282 @@
 
                 # Finally, draw the label
                 self._draw_label(label, lon, lat, fontcolor, fontoutline, offset)
+
+
+@exporter.export
+class PlotSurfaceAnalysis(MetPyHasTraits):
+    """Plot Surface Analysis Features.
+
+    This class visualizes Surface Analysis features, including the parsed WPC Surface
+    Analysis bulletins processed by the `parse_wpc_surface_bulletin()` function.
+    """
+
+    parent = Instance(Panel)
+    _need_redraw = Bool(default_value=True)
+
+    geometry = Instance(collections.abc.Iterable, allow_none=False)
+    geometry.__doc__ = """A collection of Shapely objects to plot.
+
+    A collection of Shapely objects, such as the 'geometry' column from a bulletin parsed
+    with `parse_wpc_surface_bulletin()`. Acceptable Shapely objects are
+    ``shapely.LineString``, and ``shapely.Point``.
+    """
+
+    feature = Union([Instance(collections.abc.Iterable), Unicode()],
+                    allow_none=False)
+    feature.__doc__ = """Collection of names of features to be plotted.
+
+    Collection of strings, each corresponding one-to-one with geometries, such as the
+    'features' column from a bulletin parsed with `parse_wpc_surface_bulletin()`.
+    Acceptable feature names include: 'HIGH', 'LOW', 'WARM', 'COLD', 'OCFNT', 'STNRY', 'TROF'.
+    """
+
+    strength = Union([Instance(collections.abc.Iterable), Float()], default_value=[],
+                     allow_none=True)
+    strength.__doc__ = """Collection of strengths corresponding to pressure systems.
+
+    Collection of floats, each corresponding one-to-one with pressure system features. Such
+    as the 'strength' column from a bulletin parsed with `parse_wpc_surface_bulletin()`.
+    """
+
+    HIGH_color = Union([Unicode()], default_value='blue', allow_none=True)
+    HIGH_color.__doc__ = """Color for plotting high-pressure systems.
+
+    A single string (color name or hex code) used to plot label of high-pressure system and
+    their strength, if provided. Default value is 'blue'.
+    """
+
+    LOW_color = Union([Unicode()], default_value='red', allow_none=True)
+    LOW_color.__doc__ = """Color for plotting low-pressure systems.
+
+    A single string (color name or hex code) used to plot label of low-pressure system and
+    their strength, if provided. Default value is 'red'.
+    """
+
+    WARM_color = Union([Unicode()], default_value='red', allow_none=True)
+    WARM_color.__doc__ = """Color for plotting warm fronts.
+
+    A single string (color name or hex code) used to plot warm fronts. Default
+    color is 'red', which is used by `WarmFront()` class. `WARM_color` alternates
+    with `COLD_color` to plot stationary fronts.
+    """
+
+    COLD_color = Union([Unicode()], default_value='blue', allow_none=True)
+    COLD_color.__doc__ = """Color for plotting cold fronts.
+
+    A single string (color name or hex code) used to plot cold fronts. Default
+    color is 'blue', which is used by `ColdFront()` class. `COLD_color` alternates
+    with `WARM_color` to plot stationary fronts.
+    """
+
+    OCFNT_color = Union([Unicode()], default_value='purple', allow_none=True)
+    OCFNT_color.__doc__ = """Color for plotting occluded fronts.
+
+    A single string (color name or hex code) used to plot Occluded fronts. Default
+    color is 'purple', which is used by `OccludedFront()` class.
+    """
+
+    TROF_color = Union([Unicode()], default_value='darkorange', allow_none=True)
+    TROF_color.__doc__ = """Color for plotting trough lines.
+
+    A single string (color name or hex code) used to plot trough lines. Default
+    color is 'darkorange'.
+    """
+
+    HIGH_label = Union([Unicode()], default_value='H', allow_none=True)
+    HIGH_label.__doc__ = """Label used to plot high-pressure systems.
+
+    Single string used as marker to plot high-pressure systems. Default value is 'H'.
+    """
+
+    LOW_label = Union([Unicode()], default_value='L', allow_none=True)
+    LOW_label.__doc__ = """Label used to plot low-pressure systems.
+
+    Single string used as marker to plot low-pressure systems. Default value is 'L'.
+    """
+
+    TROF_linestyle = Union([Unicode()], default_value='dashed',
+                           allow_none=True)
+    TROF_linestyle.__doc__ = """Linestyle of Trough lines.
+
+    Single string, default value is 'dashed'.
+    Accept matplotlib linestyles: 'solid', 'dotted', 'dashdot'.
+    """
+
+    label_fontsize = Union([Int(), Float(), Unicode()], default_value=10, allow_none=True)
+    label_fontsize.__doc__ = """Font sizes of pressure systems labels.
+
+    Accepts size in points or relative size. Allowed relative sizes are those of Matplotlib:
+    'xx-small', 'x-small', 'small', 'medium', 'large', 'x-large', 'xx-large'.
+    """
+
+    TROF_linewidth = Union([Float()], default_value=2,
+                           allow_none=True)
+    TROF_linewidth.__doc__ = """Stroke width for trough lines.
+
+    A single integer or floating point value representing the size of the stroke width.
+    """
+
+    FRONT_linewidth = Union([Float()], default_value=1,
+                            allow_none=True)
+    TROF_linewidth.__doc__ = """Stroke width for front lines.
+
+    A single floating point value representing the size of the stroke width.
+    """
+
+    FRONT_markersize = Union([Int(), Float(), Unicode()], default_value=3, allow_none=True)
+    FRONT_markersize.__doc__ = """Size of symbols in front lines.
+
+    Accepts size in points or relative size. Default value is 3. Allowed relative sizes are
+    those of Matplotlib: 'xx-small', 'x-small', 'small', 'medium', 'large',
+    'x-large', 'xx-large'.
+    """
+
+    strength_offset = Union([Tuple()], default_value=(0, -1), allow_none=True)
+    strength_offset.__doc__ = """Offset between label of pressure system and its
+    corresponding strength.
+
+    Tuple representing the relative position of strength value with respect to label of
+    pressure system. Default value is (0,-1). Scaled by multiplying times 80% of
+    label_fontsize value.
+    """
+
+    def _effect_map(self):
+        return {
+            'WARM': [WarmFront(size=self.FRONT_markersize, color=self.WARM_color)],
+            'COLD': [ColdFront(size=self.FRONT_markersize, color=self.COLD_color)],
+            'OCFNT': [OccludedFront(size=self.FRONT_markersize, color=self.OCFNT_color)],
+            'STNRY': [StationaryFront(size=self.FRONT_markersize,
+                                      colors=(self.WARM_color, self.COLD_color))],
+            'TROF': None
+        }
+
+    def _color_map(self):
+        return {
+            'HIGH': self.HIGH_color,
+            'LOW': self.LOW_color,
+            'TROF': self.TROF_color
+        }
+
+    def _linewidth_map(self):
+        return {
+            'WARM': self.FRONT_linewidth,
+            'COLD': self.FRONT_linewidth,
+            'OCFNT': self.FRONT_linewidth,
+            'STNRY': self.FRONT_linewidth,
+            'TROF': self.TROF_linewidth
+        }
+
+    def _label_map(self):
+        return {
+            'HIGH': self.HIGH_label,
+            'LOW': self.LOW_label
+        }
+
+    @property
+    def name(self):
+        """Generate a name for the plot."""
+        # Unlike Plots2D and PlotObs, there are no other attributes (such as 'fields' or
+        # 'levels') from which to name the plot. A generic name is returned here in case the
+        # user does not provide their own title, in which case MapPanel.draw() looks here.
+        return 'Surface Analysis Plot'
+
+    def _draw_strengths(self, text, lon, lat, color, offset=None):
+        """Draw strengths in the plot.
+
+        Parameters
+        ----------
+        text : str
+            The strength's value
+        lon : float
+            Longitude at which to position the label
+        lat : float
+            Latitude at which to position the label
+        color : str
+            Name or hex code for the color of the text
+        offset : tuple (default: (0, 0))
+            A tuple containing the x- and y-offset of the label, respectively
+        """
+        if offset is None:
+            offset = tuple(x * self.label_fontsize * 0.8 for x in self.strength_offset)
+
+        self.parent.ax.add_artist(TextCollection([lon], [lat], [str(text)],
+                                                 va='center',
+                                                 ha='center',
+                                                 color=color,
+                                                 offset=offset,
+                                                 weight='demi',
+                                                 size=int(self.label_fontsize * 0.7),
+                                                 transform=ccrs.PlateCarree()))
+
+    def _draw_labels(self, text, lon, lat, color, offset=(0, 0)):
+        """Draw labels in the plot.
+
+        Parameters
+        ----------
+        text : str
+            The label's text
+        lon : float
+            Longitude at which to position the label
+        lat : float
+            Latitude at which to position the label
+        color : str
+            Name or hex code for the color of the text
+        offset : tuple (default: (0, 0))
+            A tuple containing the x- and y-offset of the label, respectively
+        """
+        self.parent.ax.add_artist(TextCollection([lon], [lat], [str(text)],
+                                                 va='center',
+                                                 ha='center',
+                                                 color=color,
+                                                 offset=offset,
+                                                 weight='demi',
+                                                 size=self.label_fontsize,
+                                                 transform=ccrs.PlateCarree()))
+
+    def draw(self):
+        """Draw the plot."""
+        if self._need_redraw:
+            if getattr(self, 'handles', None) is None:
+                self._build()
+            self._need_redraw = False
+
+    def copy(self):
+        """Return a copy of the plot."""
+        return copy.copy(self)
+
+    def _build(self):
+        """Build the plot by calling needed plotting methods as necessary."""
+        from shapely.geometry import (LineString, Point)
+
+        # Ensure strength is a valid iterable
+        strengths = self.strength if len(self.strength) > 0 else cycle([None])
+
+        # Map plotting parameters
+        effect_map = self._effect_map()
+        color_map = self._color_map()
+        linewidth_map = self._linewidth_map()
+        label_map = self._label_map()
+
+        # Each Shapely object is plotted separately with its corresponding strength
+        # and customizable parameters
+        for geo_obj, strengthvalues, feature in zip(
+                self.geometry, strengths, self.feature):
+            kwargs = self.mpl_args.copy()
+            # Plot the Shapely object with the appropriate method and style
+            if isinstance(geo_obj, (LineString)):
+                kwargs.setdefault('linewidths', linewidth_map[feature])
+                kwargs.setdefault('facecolor', 'none')
+                kwargs.setdefault('crs', ccrs.PlateCarree())
+                kwargs.setdefault('path_effects', effect_map[feature])
+                if feature == 'TROF':
+                    kwargs.setdefault('edgecolor', color_map[feature])
+                    kwargs.setdefault('linestyle', self.TROF_linestyle)
+                self.parent.ax.add_geometries([geo_obj], **kwargs)
+            elif isinstance(geo_obj, Point):
+                kwargs.setdefault('color', color_map[feature])
+                lon, lat = geo_obj.coords[0]
+                self._draw_labels(label_map[feature], lon, lat, **kwargs)
+                # Plot strengths if provided
+                if strengthvalues is not None:
+                    self._draw_strengths(strengthvalues, lon, lat, **kwargs)

tests/plots/test_declarative.py

@@ -17,11 +17,12 @@
 
 from metpy.calc import wind_speed
 from metpy.cbook import get_test_data
-from metpy.io import GiniFile
+from metpy.io import (GiniFile, parse_wpc_surface_bulletin)
 from metpy.io.metar import parse_metar_file
 from metpy.plots import (ArrowPlot, BarbPlot, ContourPlot, FilledContourPlot, ImagePlot,
-                         MapPanel, PanelContainer, PlotGeometry, PlotObs, RasterPlot)
+                         MapPanel, PanelContainer, PlotGeometry, PlotObs, RasterPlot,
+                        PlotSurfaceAnalysis)
 from metpy.testing import needs_cartopy, version_check
 from metpy.units import units
 
 
@@ -2245,3 +2246,75 @@
         panel = MapPanel()
         panel.galaxy = 'Andromeda'
     assert 'Perhaps you meant' not in str(excinfo.value)
+
+
+@pytest.mark.mpl_image_compare(remove_text=False)
+@needs_cartopy
+def test_declarative_plot_surface_analysis_default():
+    """Test that `PlotSurfaceAnalysis` correctly plots features and strengths."""
+    # WPC Surface Analysis Bulletin to plot
+    df = parse_wpc_surface_bulletin(get_test_data('WPC_sfc_fronts_20210628_1800.txt'))
+
+    # Plot geometries and strengths
+    ps = PlotSurfaceAnalysis()
+    ps.geometry = df['geometry']
+    ps.feature = df['feature']
+    ps.strength = df['strength']
+
+    # Place plot in a panel and container
+    panel = MapPanel()
+    panel.area = [-120, -80, 30, 70]
+    panel.projection = 'lcc'
+    panel.layers = ['lakes', 'land', 'ocean',
+                    'states', 'coastline', 'borders']
+    panel.plots = [ps]
+
+    pc = PanelContainer()
+    pc.size = (12,8)
+    pc.panels = [panel]
+    pc.draw()
+
+    return pc.figure
+
+@pytest.mark.mpl_image_compare(remove_text=False)
+@needs_cartopy
+def test_declarative_plot_surface_analysis_custom():
+    """Test customization traits of `PlotSurfaceAnalysis`."""
+    # WPC Surface Analysis Bulletin to plot
+    df = parse_wpc_surface_bulletin(get_test_data('WPC_sfc_fronts_20210628_1800.txt'))
+
+    # Plot geometries and strengths
+    ps = PlotSurfaceAnalysis()
+    ps.geometry = df['geometry']
+    ps.feature = df['feature']
+    ps.strength = df['strength']
+    #customize
+    ps.HIGH_color = '#377eb8'
+    ps.COLD_color = '#ff7f00'
+    ps.TROF_color = '#4daf4a'
+    ps.OCFNT_color = '#f781bf'
+    ps.WARM_color = '#a65628'
+    ps.LOW_color = '#984ea3'
+    ps.FRONT_markersize = 4
+    ps.FRONT_linewidth = 3
+    ps.HIGH_label = 'HIGH'
+    ps.LOW_label = 'LOW'
+    ps.TROF_linestyle = 'dotted'
+    ps.label_fontsize = 15
+    ps.TROF_linewidth = 5
+    ps.strength_offset = (0,1)
+
+    # Place plot in a panel and container
+    panel = MapPanel()
+    panel.area = [-120, -80, 30, 70]
+    panel.projection = 'lcc'
+    panel.layers = ['lakes', 'land', 'ocean',
+                    'states', 'coastline', 'borders']
+    panel.plots = [ps]
+
+    pc = PanelContainer()
+    pc.size = (12,8)
+    pc.panels = [panel]
+    pc.draw()
+
+    return pc.figure