Unidata / MetPy

# Distributed under the terms of the BSD 3-Clause License.

# SPDX-License-Identifier: BSD-3-Clause

"""Create Station-model plots."""

try:

    from enum import Enum

except ImportError:

    from enum34 import Enum

import numpy as np

from .wx_symbols import (current_weather, high_clouds, low_clouds, mid_clouds,

                         pressure_tendency, sky_cover, wx_symbol_font)

from ..cbook import is_string_like

from ..package_tools import Exporter

exporter = Exporter(globals())

@exporter.export

class StationPlot(object):

    """Make a standard meteorological station plot.

    Plots values, symbols, or text spaced around a central location. Can also plot wind

    barbs as the center of the location.

    """

    location_names = {'C': (0, 0), 'N': (0, 1), 'NE': (1, 1), 'E': (1, 0), 'SE': (1, -1),

                      'S': (0, -1), 'SW': (-1, -1), 'W': (-1, 0), 'NW': (-1, 1)}

    def __init__(self, ax, x, y, fontsize=10, spacing=None, transform=None):

        """Initialize the StationPlot with items that do not change.

        This sets up the axes and station locations. The `fontsize` and `spacing`

        are also specified here to ensure that they are consistent between individual

        station elements.

        Parameters

        ----------

        ax : matplotlib.axes.Axes

            The :class:`~matplotlib.axes.Axes` for plotting

        x : array_like

            The x location of the stations in the plot

        y : array_like

            The y location of the stations in the plot

        fontsize : int

            The fontsize to use for drawing text

        spacing : int

            The spacing, in points, that corresponds to a single increment between

            station plot elements.

        transform : matplotlib.transforms.Transform (or compatible)

            The default transform to apply to the x and y positions when plotting.

        """

        self.ax = ax

        self.x = x

        self.y = y

        self.fontsize = fontsize

        self.spacing = fontsize if spacing is None else spacing

        self.transform = transform

        self.items = {}

        self.barbs = None

    def plot_symbol(self, location, codes, symbol_mapper, **kwargs):

        """At the specified location in the station model plot a set of symbols.

        This specifies that at the offset `location`, the data in `codes` should be

        converted to unicode characters (for our :data:`wx_symbol_font`) using `symbol_mapper`,

        and plotted.

        Additional keyword arguments given will be passed onto the actual plotting

        code; this is useful for specifying things like color or font properties.

        If something has already been plotted at this location, it will be replaced.

        Parameters

        ----------

        location : str or tuple[float, float]

            The offset (relative to center) to plot this parameter. If str, should be one of

            'C', 'N', 'NE', 'E', 'SE', 'S', 'SW', 'W', or 'NW'. Otherwise, should be a tuple

            specifying the number of increments in the x and y directions; increments

            are multiplied by `spacing` to give offsets in x and y relative to the center.

        codes : array_like

            The numeric values that should be converted to unicode characters for plotting.

        symbol_mapper : callable

            Controls converting data values to unicode code points for the

            :data:`wx_symbol_font` font. This should take a value and return a single unicode

            character. See :mod:`metpy.plots.wx_symbols` for included mappers.

        kwargs

            Additional keyword arguments to use for matplotlib's plotting functions.

        .. plot::

           import matplotlib.pyplot as plt

           import numpy as np

           from math import ceil

           from metpy.plots import StationPlot

           from metpy.plots.wx_symbols import current_weather, current_weather_auto

           from metpy.plots.wx_symbols import low_clouds, mid_clouds, high_clouds

           from metpy.plots.wx_symbols import sky_cover, pressure_tendency

           def plot_symbols(mapper, name, nwrap=12, figsize=(10, 1.4)):

               # Determine how many symbols there are and layout in rows of nwrap

               # if there are more than nwrap symbols

               num_symbols = len(mapper)

               codes = np.arange(len(mapper))

               ncols = nwrap

               if num_symbols <= nwrap:

                   nrows = 1

                   x = np.linspace(0, 1, len(mapper))

                   y = np.ones_like(x)

                   ax_height = 0.8

               else:

                   nrows = int(ceil(num_symbols / ncols))

                   x = np.tile(np.linspace(0, 1, ncols), nrows)[:num_symbols]

                   y = np.repeat(np.arange(nrows, 0, -1), ncols)[:num_symbols]

                   figsize = (10, 1 * nrows + 0.4)

                   ax_height = 0.8 + 0.018 * nrows

               fig = plt.figure(figsize=figsize,  dpi=300)

               ax = fig.add_axes([0, 0, 1, ax_height])

               ax.set_title(name, size=20)

               ax.xaxis.set_ticks([])

               ax.yaxis.set_ticks([])

               ax.set_frame_on(False)

               # Plot

               sp = StationPlot(ax, x, y, fontsize=36)

               sp.plot_symbol('C', codes, mapper)

               sp.plot_parameter((0, -1), codes, fontsize=18)

               ax.set_ylim(-0.05, nrows + 0.5)

               plt.show()

           plot_symbols(current_weather, "Current Weather Symbols")

           plot_symbols(current_weather_auto, "Current Weather Auto Reported Symbols")

           plot_symbols(low_clouds, "Low Cloud Symbols")

           plot_symbols(mid_clouds, "Mid Cloud Symbols")

           plot_symbols(high_clouds, "High Cloud Symbols")

           plot_symbols(sky_cover, "Sky Cover Symbols")

           plot_symbols(pressure_tendency, "Pressure Tendency Symbols")

        See Also

        --------

        plot_barb, plot_parameter, plot_text

        """

        # Make sure we use our font for symbols

        kwargs['fontproperties'] = wx_symbol_font.copy()

        return self.plot_parameter(location, codes, symbol_mapper, **kwargs)

    def plot_parameter(self, location, parameter, formatter='.0f', **kwargs):

        """At the specified location in the station model plot a set of values.

        This specifies that at the offset `location`, the data in `parameter` should be

        plotted. The conversion of the data values to a string is controlled by `formatter`.

        Additional keyword arguments given will be passed onto the actual plotting

        code; this is useful for specifying things like color or font properties.

        If something has already been plotted at this location, it will be replaced.

        Parameters

        ----------

        location : str or tuple[float, float]

            The offset (relative to center) to plot this parameter. If str, should be one of

            'C', 'N', 'NE', 'E', 'SE', 'S', 'SW', 'W', or 'NW'. Otherwise, should be a tuple

            specifying the number of increments in the x and y directions; increments

            are multiplied by `spacing` to give offsets in x and y relative to the center.

        parameter : array_like

            The numeric values that should be plotted

        formatter : str or callable, optional

            How to format the data as a string for plotting. If a string, it should be

            compatible with the :func:`format` builtin. If a callable, this should take a

            value and return a string. Defaults to '0.f'.

        kwargs

            Additional keyword arguments to use for matplotlib's plotting functions.

        See Also

        --------

        plot_barb, plot_symbol, plot_text

        """

        text = self._to_string_list(parameter, formatter)

        return self.plot_text(location, text, **kwargs)

    def plot_text(self, location, text, **kwargs):

        """At the specified location in the station model plot a collection of text.

        This specifies that at the offset `location`, the strings in `text` should be

        plotted.

        Additional keyword arguments given will be passed onto the actual plotting

        code; this is useful for specifying things like color or font properties.

        If something has already been plotted at this location, it will be replaced.

        Parameters

        ----------

        location : str or tuple[float, float]

            The offset (relative to center) to plot this parameter. If str, should be one of

            'C', 'N', 'NE', 'E', 'SE', 'S', 'SW', 'W', or 'NW'. Otherwise, should be a tuple

            specifying the number of increments in the x and y directions; increments

            are multiplied by `spacing` to give offsets in x and y relative to the center.

        text : list (or array) of strings

            The strings that should be plotted

        kwargs

            Additional keyword arguments to use for matplotlib's plotting functions.

        See Also

        --------

        plot_barb, plot_parameter, plot_symbol

        """

        location = self._handle_location(location)

        # Pass transform if necessary

        if 'transform' not in kwargs and self.transform:

            kwargs['transform'] = self.transform

        text_collection = self.ax.scattertext(self.x, self.y, text, loc=location,

                                              size=kwargs.pop('fontsize', self.fontsize),

                                              **kwargs)

        if location in self.items:

            self.items[location].remove()

        self.items[location] = text_collection

        return text_collection

    def plot_barb(self, u, v, **kwargs):

        r"""At the center of the station model plot wind barbs.

        Additional keyword arguments given will be passed onto matplotlib's

        :meth:`~matplotlib.axes.Axes.barbs` function; this is useful for specifying things

        like color or line width.

        Parameters

        ----------

        u : array-like

            The data to use for the u-component of the barbs.

        v : array-like

            The data to use for the v-component of the barbs.

        kwargs

            Additional keyword arguments to pass to matplotlib's

            :meth:`~matplotlib.axes.Axes.barbs` function.

        See Also

        --------

        plot_parameter, plot_symbol, plot_text

        """

        # Handle transforming our center points. CartoPy doesn't like 1D barbs

        # TODO: This can be removed for cartopy > 0.14.3

        if hasattr(self.ax, 'projection') and (self.transform or 'transform' in kwargs):

            trans = kwargs.pop('transform', None) or self.transform

            x, y, _ = self. ax.projection.transform_points(trans, self.x, self.y).T

        else:

            x, y = self.x, self.y

        # Empirically determined

        pivot = 0.51 * np.sqrt(self.fontsize)

        length = 1.95 * np.sqrt(self.fontsize)

        defaults = {'sizes': {'spacing': .15, 'height': 0.5, 'emptybarb': 0.35},

                    'length': length, 'pivot': pivot}

        defaults.update(kwargs)

        # Remove old barbs

        if self.barbs:

            self.barbs.remove()

        self.barbs = self.ax.barbs(x, y, u, v, **defaults)

    @staticmethod

    def _to_string_list(vals, fmt):

        """Convert a sequence of values to a list of strings."""

        if not callable(fmt):

            def formatter(s):

                """Turn a format string into a callable."""

                if hasattr(s, 'units'):

                    s = np.asscalar(s)

                return format(s, fmt)

        else:

            formatter = fmt

        return [formatter(v) if np.isfinite(v) else '' for v in vals]

    def _handle_location(self, location):

        """Process locations to get a consistent set of tuples for location."""

        if is_string_like(location):

            location = self.location_names[location]

        xoff, yoff = location

        return xoff * self.spacing, yoff * self.spacing

with exporter:

    #: :desc: Simple station plot layout

    simple_layout = StationPlotLayout()

    simple_layout.add_barb('eastward_wind', 'northward_wind', 'knots')

    simple_layout.add_value('NW', 'air_temperature', units='degC')

    simple_layout.add_value('SW', 'dew_point_temperature', units='degC')

    simple_layout.add_value('NE', 'air_pressure_at_sea_level', units='mbar',

                            fmt=lambda v: format(10 * v, '03.0f')[-3:])

    simple_layout.add_symbol('C', 'cloud_coverage', sky_cover)

    simple_layout.add_symbol('W', 'present_weather', current_weather)

    #: Full NWS station plot `layout`__

    #:

    #: __ http://oceanservice.noaa.gov/education/yos/resource/JetStream/synoptic/wxmaps.htm

    nws_layout = StationPlotLayout()

    nws_layout.add_value((-1, 1), 'air_temperature', units='degF')

    nws_layout.add_symbol((0, 2), 'high_cloud_type', high_clouds)

    nws_layout.add_symbol((0, 1), 'medium_cloud_type', mid_clouds)

    nws_layout.add_symbol((0, -1), 'low_cloud_type', low_clouds)

    nws_layout.add_value((1, 1), 'air_pressure_at_sea_level', units='mbar',

                         fmt=lambda v: format(10 * v, '03.0f')[-3:])

    nws_layout.add_value((-2, 0), 'visibility_in_air', fmt='.0f', units='miles')

    nws_layout.add_symbol((-1, 0), 'present_weather', current_weather)

    nws_layout.add_symbol((0, 0), 'cloud_coverage', sky_cover)

    nws_layout.add_value((1, 0), 'tendency_of_air_pressure', units='mbar',

                         fmt=lambda v: ('-' if v < 0 else '') + format(10 * abs(v), '02.0f'))

    nws_layout.add_symbol((2, 0), 'tendency_of_air_pressure_symbol', pressure_tendency)

    nws_layout.add_barb('eastward_wind', 'northward_wind', units='knots')

    nws_layout.add_value((-1, -1), 'dew_point_temperature', units='degF')

    # TODO: Fix once we have the past weather symbols converted

    nws_layout.add_symbol((1, -1), 'past_weather', current_weather)