8053365388

Committed 26 Feb 2024 05:51PM UTC coverage: 62.974% (-0.3%) from 63.244%

Build # 8053365388

Build Type

push

github

Committed by

ghiggi

Commit Message

Drop support for python 3.8

Run Details

3723 of 5912 relevant lines covered (62.97%)

0.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

14.14

/gpm_api/visualization/orbit.py

# -----------------------------------------------------------------------------.
# MIT License

# Copyright (c) 2024 GPM-API developers
#
# This file is part of GPM-API.

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# -----------------------------------------------------------------------------.
"""This module contains functions to visualize GPM-API ORBIT data."""
import functools

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import numpy as np

from gpm_api.checks import check_is_spatial_2d
from gpm_api.utils.checks import (
    check_contiguous_scans,
    get_slices_regular,
)
from gpm_api.utils.utils_cmap import get_colorbar_settings
from gpm_api.visualization.facetgrid import CartopyFacetGrid, ImageFacetGrid
from gpm_api.visualization.plot import (
    _plot_cartopy_pcolormesh,
    #  _plot_mpl_imshow,
    _plot_xr_imshow,
    _preprocess_figure_args,
    _preprocess_subplot_kwargs,
    plot_cartopy_background,
)


def plot_swath(
    ds, ax=None, facecolor="orange", edgecolor="black", alpha=0.4, add_background=True, **kwargs
):
    """Plot GPM orbit granule."""
    from shapely import Polygon

    # TODO: swath_def.plot() one day ...
    # - iterate by descending/ascending blocks
    # - ensure ccw boundary

    # Retrieve polygon
    swath_def = ds.gpm_api.pyresample_area
    boundary = swath_def.boundary(force_clockwise=True)
    polygon = Polygon(boundary.vertices[::-1])

    # - Initialize figure
    subplot_kwargs = kwargs.get("subplot_kwargs", {})
    fig_kwargs = kwargs.get("fig_kwargs", {})
    if ax is None:
        subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)
        fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)

    # - Add cartopy background
    if add_background:
        ax = plot_cartopy_background(ax)

    p = ax.add_geometries(
        [polygon],
        crs=ccrs.Geodetic(),
        facecolor=facecolor,
        edgecolor=edgecolor,
        alpha=alpha,
        **kwargs,
    )
    return p


def plot_swath_lines(ds, ax=None, x="lon", y="lat", linestyle="--", color="k", **kwargs):
    """Plot GPM orbit granule swath lines."""
    # - 0.0485 to account for 2.5 km from pixel center
    # TODO: adapt based on bin length (changing for each sensor) --> FUNCTION
    # ds.gpm_api.pyresample_area ...

    # - Initialize figure
    subplot_kwargs = kwargs.get("subplot_kwargs", {})
    fig_kwargs = kwargs.get("fig_kwargs", {})
    if ax is None:
        subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)
        fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)
        # - Add cartopy background
        ax = plot_cartopy_background(ax)

    # - Plot swath line
    lon = ds[x].transpose("cross_track", "along_track").data
    lat = ds[y].transpose("cross_track", "along_track").data
    p = ax.plot(
        lon[0, :] + 0.0485,
        lat[0, :],
        transform=ccrs.Geodetic(),
        linestyle=linestyle,
        color=color,
        **kwargs,
    )
    p = ax.plot(
        lon[-1, :] - 0.0485,
        lat[-1, :],
        transform=ccrs.Geodetic(),
        linestyle=linestyle,
        color=color,
        **kwargs,
    )
    return p[0]


def infill_invalid_coords(xr_obj, x="lon", y="lat", mask_variables=True):
    """Replace invalid coordinates with closer valid location.

    This operation is required to plot with pcolormesh.
    If mask_variables is True (the default) sets invalid pixel variables to NaN.

    Return tuple with 'sanitized' xr_obj and a mask with the valid coordinates.
    """
    from gpm_api.utils.checks import _is_valid_geolocation

    # Copy object
    xr_obj = xr_obj.copy()

    # Retrieve pixel with valid/invalid geolocation
    xr_valid_mask = _is_valid_geolocation(xr_obj, x=x)  # True=Valid, False=Invalid
    xr_valid_mask.name = "valid_geolocation_mask"

    np_valid_mask = xr_valid_mask.data  # True=Valid, False=Invalid
    np_unvalid_mask = ~np_valid_mask  # True=Invalid, False=Valid

    # If there are invalid pixels, replace invalid coordinates with closer valid values
    if np.any(np_unvalid_mask):
        lon = np.asanyarray(xr_obj[x].data)
        lat = np.asanyarray(xr_obj[y].data)
        lon_dummy = lon.copy()
        lon_dummy[np_unvalid_mask] = np.interp(
            np.flatnonzero(np_unvalid_mask), np.flatnonzero(np_valid_mask), lon[np_valid_mask]
        )
        lat_dummy = lat.copy()
        lat_dummy[np_unvalid_mask] = np.interp(
            np.flatnonzero(np_unvalid_mask), np.flatnonzero(np_valid_mask), lat[np_valid_mask]
        )
        xr_obj[x].data = lon_dummy
        xr_obj[y].data = lat_dummy

    # Mask variables if asked
    if mask_variables:
        xr_obj = xr_obj.where(xr_valid_mask)

    return xr_obj, xr_valid_mask


# TODO: plot swath polygon
# def plot_swath(ds, ax=None):

# da.gpm_api.pyresample_area.boundary
# da.gpm_api.pyresample_area.outer_boundary.polygon
# da.gpm_api.pyresample_area.outer_boundary.sides ..


def _remove_invalid_outer_cross_track(xr_obj, x="lon"):
    """Remove outer crosstrack scans if geolocation is always missing."""
    lon = np.asanyarray(xr_obj[x].transpose("cross_track", "along_track"))
    isna = np.all(np.isnan(lon), axis=1)
    if isna[0] or isna[-1]:
        is_valid = ~isna
        # Find the index where coordinates start to be valid
        start_index = np.argmax(is_valid)
        # Find the index where the first False value occurs (from the end)
        end_index = len(is_valid) - np.argmax(is_valid[::-1])
        # Define slice
        slc = slice(start_index, end_index)
        # Subset object
        xr_obj = xr_obj.isel({"cross_track": slc})
    return xr_obj


def _call_over_contiguous_scans(function):
    """Decorator to call the plotting function multiple times only over contiguous scans intervals."""

    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        # Assumption: only da and ax are passed as args

        # Get data array (first position)
        da = args[0] if len(args) > 0 else kwargs.get("da")

        # Get axis
        ax = args[1] if len(args) > 1 else kwargs.get("ax")

        # Check data validity
        # TODO: improve for rgb=True
        # TODO: Make independent of name of lon/lat, cross_track, along_track
        rgb = kwargs.get("rgb", False)
        if not rgb:
            # - Check data array
            check_is_spatial_2d(da)

            # - Get slices with contiguous scans and valid geolocation
            list_slices = get_slices_regular(da, min_size=2, min_n_scans=2)
            if len(list_slices) == 0:
                raise ValueError("No regular scans available. Impossible to plot.")
        else:
            list_slices = [slice(0, None)]

        # - Define kwargs
        user_kwargs = kwargs.copy()
        p = None
        x = user_kwargs["x"]
        y = user_kwargs["y"]

        # - Call the function over each slice
        for i, slc in enumerate(list_slices):
            if not rgb:
                # Retrieve contiguous data array
                tmp_da = da.isel({"along_track": slc})

                # Remove outer cross-track indices if all without coordinates
                tmp_da = _remove_invalid_outer_cross_track(tmp_da, x=x)
            else:
                tmp_da = da

            # Replace invalid coordinate with closer value
            # - This might be necessary for some products
            #   having all the outer swath invalid coordinates
            # - An example is the 2B-GPM-CORRA
            tmp_da, tmp_da_valid_mask = infill_invalid_coords(tmp_da, x=x, y=y, mask_variables=True)

            # Define temporary kwargs
            tmp_kwargs = user_kwargs.copy()
            tmp_kwargs["da"] = tmp_da
            if i == 0:
                tmp_kwargs["ax"] = ax

            else:
                tmp_kwargs["ax"] = p.axes
                tmp_kwargs["add_background"] = False

            # Define alpha to make invalid coordinates transparent
            # --> cartopy.pcolormesh currently bug when providing a alpha array :(
            # TODO: Open an issue on that !

            # tmp_valid_mask = tmp_da_valid_mask.data
            # if not np.all(tmp_valid_mask):
            #     alpha = tmp_valid_mask.astype(int)  # 0s and 1s
            #     if "alpha" in tmp_kwargs:
            #         alpha = tmp_kwargs["alpha"] * alpha
            #     tmp_kwargs["alpha"] = alpha

            # Set colorbar to False for all except last iteration
            # --> Avoid drawing multiple colorbars
            if i != len(list_slices) - 1 and "add_colorbar" in user_kwargs:
                tmp_kwargs["add_colorbar"] = False

            # Before function call
            p = function(**tmp_kwargs)
            # p.set_alpha(alpha)

        return p

    return wrapper


@_call_over_contiguous_scans
def _plot_orbit_map_cartopy(
    da,
    ax=None,
    x="lon",
    y="lat",
    add_colorbar=True,
    add_swath_lines=True,
    add_background=True,
    rgb=False,
    fig_kwargs={},
    subplot_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot GPM orbit granule in a cartographic map."""
    # - Check inputs
    if not rgb:
        check_is_spatial_2d(da)
    _preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)

    # - Initialize figure
    if ax is None:
        subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)
        fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)

    # - Add cartopy background
    if add_background:
        ax = plot_cartopy_background(ax)

    # - Sanitize plot_kwargs passed by FacetGrid
    facet_grid_args = ["levels", "extend", "add_labels", "_is_facetgrid"]
    _ = [plot_kwargs.pop(arg, None) for arg in facet_grid_args]

    # - If not specified, retrieve/update plot_kwargs and cbar_kwargs as function of variable name
    variable = da.name
    plot_kwargs, cbar_kwargs = get_colorbar_settings(
        name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
    )
    # - Specify colorbar label
    if "label" not in cbar_kwargs:
        unit = da.attrs.get("units", "-")
        cbar_kwargs["label"] = f"{variable} [{unit}]"

    # - Add swath lines
    # --> Fail if not cross_track, along_track dimension currently
    if add_swath_lines and not rgb:
        p = plot_swath_lines(da, ax=ax, linestyle="--", color="black")

    # - Add variable field with cartopy
    p = _plot_cartopy_pcolormesh(
        ax=ax,
        da=da,
        x=x,
        y=y,
        plot_kwargs=plot_kwargs,
        cbar_kwargs=cbar_kwargs,
        add_colorbar=add_colorbar,
        rgb=rgb,
    )
    # - Return mappable
    return p


@_call_over_contiguous_scans
def plot_orbit_mesh(
    da,
    ax=None,
    x="lon",
    y="lat",
    edgecolors="k",
    linewidth=0.1,
    add_background=True,
    fig_kwargs={},
    subplot_kwargs={},
    **plot_kwargs,
):
    """Plot GPM orbit granule mesh in a cartographic map."""
    # - Check inputs
    _preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)

    # - Initialize figure
    if ax is None:
        subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)
        fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)

    # - Add cartopy background
    if add_background:
        ax = plot_cartopy_background(ax)

    # - Define plot_kwargs to display only the mesh
    plot_kwargs["facecolor"] = "none"
    plot_kwargs["alpha"] = 1
    plot_kwargs["edgecolors"] = (edgecolors,)  # Introduce bugs in Cartopy !
    plot_kwargs["linewidth"] = (linewidth,)
    plot_kwargs["antialiased"] = True

    # - Add variable field with cartopy
    p = _plot_cartopy_pcolormesh(
        da=da,
        ax=ax,
        x=x,
        y=y,
        plot_kwargs=plot_kwargs,
        add_colorbar=False,
    )
    # - Return mappable
    return p


def _plot_orbit_image(
    da,
    x=None,
    y=None,
    ax=None,
    add_colorbar=True,
    interpolation="nearest",
    fig_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot GPM orbit granule as in image."""
    # NOTE:
    # - Code is almost equal to plot_grid_image
    # - Refactor after developed test units

    # ----------------------------------------------
    # - Check inputs
    check_contiguous_scans(da)
    _preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs)

    # Initialize figure
    if ax is None:
        if "rgb" not in plot_kwargs:
            check_is_spatial_2d(da)
        fig, ax = plt.subplots(**fig_kwargs)

    # - Sanitize plot_kwargs passed by FacetGrid
    facet_grid_args = ["levels", "extend", "add_labels", "_is_facetgrid"]
    _ = [plot_kwargs.pop(arg, None) for arg in facet_grid_args]

    # - If not specified, retrieve/update plot_kwargs and cbar_kwargs as function of product name
    plot_kwargs, cbar_kwargs = get_colorbar_settings(
        name=da.name, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
    )

    # - Plot with xarray
    p = _plot_xr_imshow(
        ax=ax,
        da=da,
        x=x,
        y=y,
        interpolation=interpolation,
        add_colorbar=add_colorbar,
        plot_kwargs=plot_kwargs,
        cbar_kwargs=cbar_kwargs,
    )

    # Add axis labels
    p.axes.set_xlabel("Along-Track")
    p.axes.set_ylabel("Cross-Track")

    # - Return mappable
    return p


def plot_orbit_map(
    da,
    ax=None,
    x="lon",
    y="lat",
    add_colorbar=True,
    add_swath_lines=True,
    add_background=True,
    rgb=False,
    fig_kwargs={},
    subplot_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot DataArray 2D field with cartopy."""
    # Plot FacetGrid
    if "col" in plot_kwargs or "row" in plot_kwargs:
        p = _plot_orbit_map_facetgrid(
            da=da,
            x=x,
            y=y,
            ax=ax,
            add_colorbar=add_colorbar,
            add_swath_lines=add_swath_lines,
            add_background=add_background,
            rgb=rgb,
            fig_kwargs=fig_kwargs,
            subplot_kwargs=subplot_kwargs,
            cbar_kwargs=cbar_kwargs,
            **plot_kwargs,
        )
    # Plot with cartopy imshow
    else:
        da = da.squeeze()  # remove time if dim=1
        p = _plot_orbit_map_cartopy(
            da=da,
            x=x,
            y=y,
            ax=ax,
            add_colorbar=add_colorbar,
            add_swath_lines=add_swath_lines,
            add_background=add_background,
            rgb=rgb,
            fig_kwargs=fig_kwargs,
            subplot_kwargs=subplot_kwargs,
            cbar_kwargs=cbar_kwargs,
            **plot_kwargs,
        )
    # - Return mappable
    return p


def _plot_orbit_map_facetgrid(
    da,
    x="lon",
    y="lat",
    ax=None,
    add_colorbar=True,
    add_swath_lines=True,
    add_background=True,
    rgb=False,
    fig_kwargs={},
    subplot_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot 2D fields with FacetGrid."""
    # - Check inputs
    if ax is not None:
        raise ValueError("When plotting with FacetGrid, do not specify the 'ax'.")
    _preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)
    subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)

    # Retrieve GPM-API defaults cmap and cbar kwargs
    variable = da.name
    plot_kwargs, cbar_kwargs = get_colorbar_settings(
        name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
    )
    # Retrieve projection
    projection = subplot_kwargs.get("projection", None)
    if projection is None:
        raise ValueError("Please specify a Cartopy projection in subplot_kwargs['projection'].")

    # Create FacetGrid
    optimize_layout = plot_kwargs.pop("optimize_layout", True)
    fc = CartopyFacetGrid(
        data=da.compute(),
        projection=projection,
        col=plot_kwargs.pop("col", None),
        row=plot_kwargs.pop("row", None),
        col_wrap=plot_kwargs.pop("col_wrap", None),
        axes_pad=plot_kwargs.pop("axes_pad", None),
        add_colorbar=add_colorbar,
        cbar_kwargs=cbar_kwargs,
        fig_kwargs=fig_kwargs,
        facet_height=plot_kwargs.pop("facet_height", 3),
        facet_aspect=plot_kwargs.pop("facet_aspect", 1),
    )

    # Plot the maps
    fc = fc.map_dataarray(
        _plot_orbit_map_cartopy,
        x=x,
        y=y,
        add_colorbar=False,
        add_background=add_background,
        add_swath_lines=add_swath_lines,
        rgb=rgb,
        cbar_kwargs=cbar_kwargs,
        **plot_kwargs,
    )

    # Remove duplicated gridline labels
    fc.remove_duplicated_axis_labels()

    # Add colorbar
    if add_colorbar:
        fc.add_colorbar(**cbar_kwargs)

    # Optimize layout
    if optimize_layout:
        fc.optimize_layout()

    return fc


def _plot_orbit_image_facetgrid(
    da,
    x="lon",
    y="lat",
    ax=None,
    add_colorbar=True,
    interpolation="nearest",
    fig_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot 2D fields with FacetGrid."""
    # - Check inputs
    if ax is not None:
        raise ValueError("When plotting with FacetGrid, do not specify the 'ax'.")
    _preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs)

    # Retrieve GPM-API defaults cmap and cbar kwargs
    variable = da.name
    plot_kwargs, cbar_kwargs = get_colorbar_settings(
        name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
    )

    # Create FacetGrid
    fc = ImageFacetGrid(
        data=da.compute(),
        col=plot_kwargs.pop("col", None),
        row=plot_kwargs.pop("row", None),
        col_wrap=plot_kwargs.pop("col_wrap", None),
        axes_pad=plot_kwargs.pop("axes_pad", None),
        fig_kwargs=fig_kwargs,
        cbar_kwargs=cbar_kwargs,
        add_colorbar=add_colorbar,
        aspect=plot_kwargs.pop("aspect", False),
        facet_height=plot_kwargs.pop("facet_height", 3),
        facet_aspect=plot_kwargs.pop("facet_aspect", 1),
    )

    # Plot the maps
    fc = fc.map_dataarray(
        _plot_orbit_image,
        x=x,
        y=y,
        add_colorbar=False,
        interpolation=interpolation,
        cbar_kwargs=cbar_kwargs,
        **plot_kwargs,
    )

    fc.remove_duplicated_axis_labels()

    # Add colorbar
    if add_colorbar:
        fc.add_colorbar(**cbar_kwargs)

    return fc


def plot_orbit_image(
    da,
    x="lon",
    y="lat",
    ax=None,
    add_colorbar=True,
    interpolation="nearest",
    fig_kwargs={},
    cbar_kwargs={},
    **plot_kwargs,
):
    """Plot DataArray 2D field with cartopy."""
    # Plot FacetGrid with xarray imshow
    if "col" in plot_kwargs or "row" in plot_kwargs:
        p = _plot_orbit_image_facetgrid(
            da=da,
            x=x,
            y=y,
            ax=ax,
            add_colorbar=add_colorbar,
            interpolation=interpolation,
            fig_kwargs=fig_kwargs,
            cbar_kwargs=cbar_kwargs,
            **plot_kwargs,
        )

    # Plot with cartopy imshow
    else:
        da = da.squeeze()  # remove time if dim=1
        p = _plot_orbit_image(
            da=da,
            x=x,
            y=y,
            ax=ax,
            add_colorbar=add_colorbar,
            interpolation=interpolation,
            fig_kwargs=fig_kwargs,
            cbar_kwargs=cbar_kwargs,
            **plot_kwargs,
        )
    # - Return imagepable
    return p

1	# -----------------------------------------------------------------------------.
2	# MIT License
3
4	# Copyright (c) 2024 GPM-API developers
5	#
6	# This file is part of GPM-API.
7
8	# Permission is hereby granted, free of charge, to any person obtaining a copy
9	# of this software and associated documentation files (the "Software"), to deal
10	# in the Software without restriction, including without limitation the rights
11	# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12	# copies of the Software, and to permit persons to whom the Software is
13	# furnished to do so, subject to the following conditions:
14	#
15	# The above copyright notice and this permission notice shall be included in all
16	# copies or substantial portions of the Software.
17	#
18	# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19	# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20	# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21	# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22	# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23	# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24	# SOFTWARE.
25
26	# -----------------------------------------------------------------------------.
27	"""This module contains functions to visualize GPM-API ORBIT data."""	1✔
28	import functools	1✔
29
30	import cartopy.crs as ccrs	1✔
31	import matplotlib.pyplot as plt	1✔
32	import numpy as np	1✔
33
34	from gpm_api.checks import check_is_spatial_2d	1✔
35	from gpm_api.utils.checks import (	1✔
36	check_contiguous_scans,
37	get_slices_regular,
38	)
39	from gpm_api.utils.utils_cmap import get_colorbar_settings	1✔
40	from gpm_api.visualization.facetgrid import CartopyFacetGrid, ImageFacetGrid	1✔
41	from gpm_api.visualization.plot import (	1✔
42	_plot_cartopy_pcolormesh,
43	# _plot_mpl_imshow,
44	_plot_xr_imshow,
45	_preprocess_figure_args,
46	_preprocess_subplot_kwargs,
47	plot_cartopy_background,
48	)
49
50
51	def plot_swath(	1✔
52	ds, ax=None, facecolor="orange", edgecolor="black", alpha=0.4, add_background=True, **kwargs
53	):
54	"""Plot GPM orbit granule."""
55	from shapely import Polygon	×
56
57	# TODO: swath_def.plot() one day ...
58	# - iterate by descending/ascending blocks
59	# - ensure ccw boundary
60
61	# Retrieve polygon
62	swath_def = ds.gpm_api.pyresample_area	×
63	boundary = swath_def.boundary(force_clockwise=True)	×
64	polygon = Polygon(boundary.vertices[::-1])	×
65
66	# - Initialize figure
67	subplot_kwargs = kwargs.get("subplot_kwargs", {})	×
68	fig_kwargs = kwargs.get("fig_kwargs", {})	×
69	if ax is None:	×
70	subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)	×
71	fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)	×
72
73	# - Add cartopy background
74	if add_background:	×
75	ax = plot_cartopy_background(ax)	×
76
77	p = ax.add_geometries(	×
78	[polygon],
79	crs=ccrs.Geodetic(),
80	facecolor=facecolor,
81	edgecolor=edgecolor,
82	alpha=alpha,
83	**kwargs,
84	)
85	return p	×
86
87
88	def plot_swath_lines(ds, ax=None, x="lon", y="lat", linestyle="--", color="k", **kwargs):	1✔
89	"""Plot GPM orbit granule swath lines."""
90	# - 0.0485 to account for 2.5 km from pixel center
91	# TODO: adapt based on bin length (changing for each sensor) --> FUNCTION
92	# ds.gpm_api.pyresample_area ...
93
94	# - Initialize figure
95	subplot_kwargs = kwargs.get("subplot_kwargs", {})	×
96	fig_kwargs = kwargs.get("fig_kwargs", {})	×
97	if ax is None:	×
98	subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)	×
99	fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)	×
100	# - Add cartopy background
101	ax = plot_cartopy_background(ax)	×
102
103	# - Plot swath line
104	lon = ds[x].transpose("cross_track", "along_track").data	×
105	lat = ds[y].transpose("cross_track", "along_track").data	×
106	p = ax.plot(	×
107	lon[0, :] + 0.0485,
108	lat[0, :],
109	transform=ccrs.Geodetic(),
110	linestyle=linestyle,
111	color=color,
112	**kwargs,
113	)
114	p = ax.plot(	×
115	lon[-1, :] - 0.0485,
116	lat[-1, :],
117	transform=ccrs.Geodetic(),
118	linestyle=linestyle,
119	color=color,
120	**kwargs,
121	)
122	return p[0]	×
123
124
125	def infill_invalid_coords(xr_obj, x="lon", y="lat", mask_variables=True):	1✔
126	"""Replace invalid coordinates with closer valid location.
127
128	This operation is required to plot with pcolormesh.
129	If mask_variables is True (the default) sets invalid pixel variables to NaN.
130
131	Return tuple with 'sanitized' xr_obj and a mask with the valid coordinates.
132	"""
133	from gpm_api.utils.checks import _is_valid_geolocation	×
134
135	# Copy object
136	xr_obj = xr_obj.copy()	×
137
138	# Retrieve pixel with valid/invalid geolocation
139	xr_valid_mask = _is_valid_geolocation(xr_obj, x=x) # True=Valid, False=Invalid	×
140	xr_valid_mask.name = "valid_geolocation_mask"	×
141
142	np_valid_mask = xr_valid_mask.data # True=Valid, False=Invalid	×
143	np_unvalid_mask = ~np_valid_mask # True=Invalid, False=Valid	×
144
145	# If there are invalid pixels, replace invalid coordinates with closer valid values
146	if np.any(np_unvalid_mask):	×
147	lon = np.asanyarray(xr_obj[x].data)	×
148	lat = np.asanyarray(xr_obj[y].data)	×
149	lon_dummy = lon.copy()	×
150	lon_dummy[np_unvalid_mask] = np.interp(	×
151	np.flatnonzero(np_unvalid_mask), np.flatnonzero(np_valid_mask), lon[np_valid_mask]
152	)
153	lat_dummy = lat.copy()	×
154	lat_dummy[np_unvalid_mask] = np.interp(	×
155	np.flatnonzero(np_unvalid_mask), np.flatnonzero(np_valid_mask), lat[np_valid_mask]
156	)
157	xr_obj[x].data = lon_dummy	×
158	xr_obj[y].data = lat_dummy	×
159
160	# Mask variables if asked
161	if mask_variables:	×
162	xr_obj = xr_obj.where(xr_valid_mask)	×
163
164	return xr_obj, xr_valid_mask	×
165
166
167	# TODO: plot swath polygon
168	# def plot_swath(ds, ax=None):
169
170	# da.gpm_api.pyresample_area.boundary
171	# da.gpm_api.pyresample_area.outer_boundary.polygon
172	# da.gpm_api.pyresample_area.outer_boundary.sides ..
173
174
175	def _remove_invalid_outer_cross_track(xr_obj, x="lon"):	1✔
176	"""Remove outer crosstrack scans if geolocation is always missing."""
177	lon = np.asanyarray(xr_obj[x].transpose("cross_track", "along_track"))	×
178	isna = np.all(np.isnan(lon), axis=1)	×
179	if isna[0] or isna[-1]:	×
180	is_valid = ~isna	×
181	# Find the index where coordinates start to be valid
182	start_index = np.argmax(is_valid)	×
183	# Find the index where the first False value occurs (from the end)
184	end_index = len(is_valid) - np.argmax(is_valid[::-1])	×
185	# Define slice
186	slc = slice(start_index, end_index)	×
187	# Subset object
188	xr_obj = xr_obj.isel({"cross_track": slc})	×
189	return xr_obj	×
190
191
192	def _call_over_contiguous_scans(function):	1✔
193	"""Decorator to call the plotting function multiple times only over contiguous scans intervals."""
194
195	@functools.wraps(function)	1✔
196	def wrapper(args, *kwargs):	1✔
197	# Assumption: only da and ax are passed as args
198
199	# Get data array (first position)
200	da = args[0] if len(args) > 0 else kwargs.get("da")	×
201
202	# Get axis
203	ax = args[1] if len(args) > 1 else kwargs.get("ax")	×
204
205	# Check data validity
206	# TODO: improve for rgb=True
207	# TODO: Make independent of name of lon/lat, cross_track, along_track
208	rgb = kwargs.get("rgb", False)	×
209	if not rgb:	×
210	# - Check data array
211	check_is_spatial_2d(da)	×
212
213	# - Get slices with contiguous scans and valid geolocation
214	list_slices = get_slices_regular(da, min_size=2, min_n_scans=2)	×
215	if len(list_slices) == 0:	×
216	raise ValueError("No regular scans available. Impossible to plot.")	×
217	else:
218	list_slices = [slice(0, None)]	×
219
220	# - Define kwargs
221	user_kwargs = kwargs.copy()	×
222	p = None	×
223	x = user_kwargs["x"]	×
224	y = user_kwargs["y"]	×
225
226	# - Call the function over each slice
227	for i, slc in enumerate(list_slices):	×
228	if not rgb:	×
229	# Retrieve contiguous data array
230	tmp_da = da.isel({"along_track": slc})	×
231
232	# Remove outer cross-track indices if all without coordinates
233	tmp_da = _remove_invalid_outer_cross_track(tmp_da, x=x)	×
234	else:
235	tmp_da = da	×
236
237	# Replace invalid coordinate with closer value
238	# - This might be necessary for some products
239	# having all the outer swath invalid coordinates
240	# - An example is the 2B-GPM-CORRA
241	tmp_da, tmp_da_valid_mask = infill_invalid_coords(tmp_da, x=x, y=y, mask_variables=True)	×
242
243	# Define temporary kwargs
244	tmp_kwargs = user_kwargs.copy()	×
245	tmp_kwargs["da"] = tmp_da	×
246	if i == 0:	×
247	tmp_kwargs["ax"] = ax	×
248
249	else:
250	tmp_kwargs["ax"] = p.axes	×
251	tmp_kwargs["add_background"] = False	×
252
253	# Define alpha to make invalid coordinates transparent
254	# --> cartopy.pcolormesh currently bug when providing a alpha array :(
255	# TODO: Open an issue on that !
256
257	# tmp_valid_mask = tmp_da_valid_mask.data
258	# if not np.all(tmp_valid_mask):
259	# alpha = tmp_valid_mask.astype(int) # 0s and 1s
260	# if "alpha" in tmp_kwargs:
261	# alpha = tmp_kwargs["alpha"] * alpha
262	# tmp_kwargs["alpha"] = alpha
263
264	# Set colorbar to False for all except last iteration
265	# --> Avoid drawing multiple colorbars
266	if i != len(list_slices) - 1 and "add_colorbar" in user_kwargs:	×
267	tmp_kwargs["add_colorbar"] = False	×
268
269	# Before function call
270	p = function(**tmp_kwargs)	×
271	# p.set_alpha(alpha)
272
273	return p	×
274
275	return wrapper	1✔
276
277
278	@_call_over_contiguous_scans	1✔
279	def _plot_orbit_map_cartopy(	1✔
280	da,
281	ax=None,
282	x="lon",
283	y="lat",
284	add_colorbar=True,
285	add_swath_lines=True,
286	add_background=True,
287	rgb=False,
288	fig_kwargs={},
289	subplot_kwargs={},
290	cbar_kwargs={},
291	**plot_kwargs,
292	):
293	"""Plot GPM orbit granule in a cartographic map."""
294	# - Check inputs
295	if not rgb:	×
296	check_is_spatial_2d(da)	×
297	_preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)	×
298
299	# - Initialize figure
300	if ax is None:	×
301	subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)	×
302	fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)	×
303
304	# - Add cartopy background
305	if add_background:	×
306	ax = plot_cartopy_background(ax)	×
307
308	# - Sanitize plot_kwargs passed by FacetGrid
309	facet_grid_args = ["levels", "extend", "add_labels", "_is_facetgrid"]	×
310	_ = [plot_kwargs.pop(arg, None) for arg in facet_grid_args]	×
311
312	# - If not specified, retrieve/update plot_kwargs and cbar_kwargs as function of variable name
313	variable = da.name	×
314	plot_kwargs, cbar_kwargs = get_colorbar_settings(	×
315	name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
316	)
317	# - Specify colorbar label
318	if "label" not in cbar_kwargs:	×
319	unit = da.attrs.get("units", "-")	×
320	cbar_kwargs["label"] = f"{variable} [{unit}]"	×
321
322	# - Add swath lines
323	# --> Fail if not cross_track, along_track dimension currently
324	if add_swath_lines and not rgb:	×
325	p = plot_swath_lines(da, ax=ax, linestyle="--", color="black")	×
326
327	# - Add variable field with cartopy
328	p = _plot_cartopy_pcolormesh(	×
329	ax=ax,
330	da=da,
331	x=x,
332	y=y,
333	plot_kwargs=plot_kwargs,
334	cbar_kwargs=cbar_kwargs,
335	add_colorbar=add_colorbar,
336	rgb=rgb,
337	)
338	# - Return mappable
339	return p	×
340
341
342	@_call_over_contiguous_scans	1✔
343	def plot_orbit_mesh(	1✔
344	da,
345	ax=None,
346	x="lon",
347	y="lat",
348	edgecolors="k",
349	linewidth=0.1,
350	add_background=True,
351	fig_kwargs={},
352	subplot_kwargs={},
353	**plot_kwargs,
354	):
355	"""Plot GPM orbit granule mesh in a cartographic map."""
356	# - Check inputs
357	_preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)	×
358
359	# - Initialize figure
360	if ax is None:	×
361	subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)	×
362	fig, ax = plt.subplots(subplot_kw=subplot_kwargs, **fig_kwargs)	×
363
364	# - Add cartopy background
365	if add_background:	×
366	ax = plot_cartopy_background(ax)	×
367
368	# - Define plot_kwargs to display only the mesh
369	plot_kwargs["facecolor"] = "none"	×
370	plot_kwargs["alpha"] = 1	×
371	plot_kwargs["edgecolors"] = (edgecolors,) # Introduce bugs in Cartopy !	×
372	plot_kwargs["linewidth"] = (linewidth,)	×
373	plot_kwargs["antialiased"] = True	×
374
375	# - Add variable field with cartopy
376	p = _plot_cartopy_pcolormesh(	×
377	da=da,
378	ax=ax,
379	x=x,
380	y=y,
381	plot_kwargs=plot_kwargs,
382	add_colorbar=False,
383	)
384	# - Return mappable
385	return p	×
386
387
388	def _plot_orbit_image(	1✔
389	da,
390	x=None,
391	y=None,
392	ax=None,
393	add_colorbar=True,
394	interpolation="nearest",
395	fig_kwargs={},
396	cbar_kwargs={},
397	**plot_kwargs,
398	):
399	"""Plot GPM orbit granule as in image."""
400	# NOTE:
401	# - Code is almost equal to plot_grid_image
402	# - Refactor after developed test units
403
404	# ----------------------------------------------
405	# - Check inputs
406	check_contiguous_scans(da)	×
407	_preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs)	×
408
409	# Initialize figure
410	if ax is None:	×
411	if "rgb" not in plot_kwargs:	×
412	check_is_spatial_2d(da)	×
413	fig, ax = plt.subplots(**fig_kwargs)	×
414
415	# - Sanitize plot_kwargs passed by FacetGrid
416	facet_grid_args = ["levels", "extend", "add_labels", "_is_facetgrid"]	×
417	_ = [plot_kwargs.pop(arg, None) for arg in facet_grid_args]	×
418
419	# - If not specified, retrieve/update plot_kwargs and cbar_kwargs as function of product name
420	plot_kwargs, cbar_kwargs = get_colorbar_settings(	×
421	name=da.name, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
422	)
423
424	# - Plot with xarray
425	p = _plot_xr_imshow(	×
426	ax=ax,
427	da=da,
428	x=x,
429	y=y,
430	interpolation=interpolation,
431	add_colorbar=add_colorbar,
432	plot_kwargs=plot_kwargs,
433	cbar_kwargs=cbar_kwargs,
434	)
435
436	# Add axis labels
437	p.axes.set_xlabel("Along-Track")	×
438	p.axes.set_ylabel("Cross-Track")	×
439
440	# - Return mappable
441	return p	×
442
443
444	def plot_orbit_map(	1✔
445	da,
446	ax=None,
447	x="lon",
448	y="lat",
449	add_colorbar=True,
450	add_swath_lines=True,
451	add_background=True,
452	rgb=False,
453	fig_kwargs={},
454	subplot_kwargs={},
455	cbar_kwargs={},
456	**plot_kwargs,
457	):
458	"""Plot DataArray 2D field with cartopy."""
459	# Plot FacetGrid
460	if "col" in plot_kwargs or "row" in plot_kwargs:	×
461	p = _plot_orbit_map_facetgrid(	×
462	da=da,
463	x=x,
464	y=y,
465	ax=ax,
466	add_colorbar=add_colorbar,
467	add_swath_lines=add_swath_lines,
468	add_background=add_background,
469	rgb=rgb,
470	fig_kwargs=fig_kwargs,
471	subplot_kwargs=subplot_kwargs,
472	cbar_kwargs=cbar_kwargs,
473	**plot_kwargs,
474	)
475	# Plot with cartopy imshow
476	else:
477	da = da.squeeze() # remove time if dim=1	×
478	p = _plot_orbit_map_cartopy(	×
479	da=da,
480	x=x,
481	y=y,
482	ax=ax,
483	add_colorbar=add_colorbar,
484	add_swath_lines=add_swath_lines,
485	add_background=add_background,
486	rgb=rgb,
487	fig_kwargs=fig_kwargs,
488	subplot_kwargs=subplot_kwargs,
489	cbar_kwargs=cbar_kwargs,
490	**plot_kwargs,
491	)
492	# - Return mappable
493	return p	×
494
495
496	def _plot_orbit_map_facetgrid(	1✔
497	da,
498	x="lon",
499	y="lat",
500	ax=None,
501	add_colorbar=True,
502	add_swath_lines=True,
503	add_background=True,
504	rgb=False,
505	fig_kwargs={},
506	subplot_kwargs={},
507	cbar_kwargs={},
508	**plot_kwargs,
509	):
510	"""Plot 2D fields with FacetGrid."""
511	# - Check inputs
512	if ax is not None:	×
513	raise ValueError("When plotting with FacetGrid, do not specify the 'ax'.")	×
514	_preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs, subplot_kwargs=subplot_kwargs)	×
515	subplot_kwargs = _preprocess_subplot_kwargs(subplot_kwargs)	×
516
517	# Retrieve GPM-API defaults cmap and cbar kwargs
518	variable = da.name	×
519	plot_kwargs, cbar_kwargs = get_colorbar_settings(	×
520	name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
521	)
522	# Retrieve projection
523	projection = subplot_kwargs.get("projection", None)	×
524	if projection is None:	×
525	raise ValueError("Please specify a Cartopy projection in subplot_kwargs['projection'].")	×
526
527	# Create FacetGrid
528	optimize_layout = plot_kwargs.pop("optimize_layout", True)	×
529	fc = CartopyFacetGrid(	×
530	data=da.compute(),
531	projection=projection,
532	col=plot_kwargs.pop("col", None),
533	row=plot_kwargs.pop("row", None),
534	col_wrap=plot_kwargs.pop("col_wrap", None),
535	axes_pad=plot_kwargs.pop("axes_pad", None),
536	add_colorbar=add_colorbar,
537	cbar_kwargs=cbar_kwargs,
538	fig_kwargs=fig_kwargs,
539	facet_height=plot_kwargs.pop("facet_height", 3),
540	facet_aspect=plot_kwargs.pop("facet_aspect", 1),
541	)
542
543	# Plot the maps
544	fc = fc.map_dataarray(	×
545	_plot_orbit_map_cartopy,
546	x=x,
547	y=y,
548	add_colorbar=False,
549	add_background=add_background,
550	add_swath_lines=add_swath_lines,
551	rgb=rgb,
552	cbar_kwargs=cbar_kwargs,
553	**plot_kwargs,
554	)
555
556	# Remove duplicated gridline labels
557	fc.remove_duplicated_axis_labels()	×
558
559	# Add colorbar
560	if add_colorbar:	×
561	fc.add_colorbar(**cbar_kwargs)	×
562
563	# Optimize layout
564	if optimize_layout:	×
565	fc.optimize_layout()	×
566
567	return fc	×
568
569
570	def _plot_orbit_image_facetgrid(	1✔
571	da,
572	x="lon",
573	y="lat",
574	ax=None,
575	add_colorbar=True,
576	interpolation="nearest",
577	fig_kwargs={},
578	cbar_kwargs={},
579	**plot_kwargs,
580	):
581	"""Plot 2D fields with FacetGrid."""
582	# - Check inputs
583	if ax is not None:	×
584	raise ValueError("When plotting with FacetGrid, do not specify the 'ax'.")	×
585	_preprocess_figure_args(ax=ax, fig_kwargs=fig_kwargs)	×
586
587	# Retrieve GPM-API defaults cmap and cbar kwargs
588	variable = da.name	×
589	plot_kwargs, cbar_kwargs = get_colorbar_settings(	×
590	name=variable, plot_kwargs=plot_kwargs, cbar_kwargs=cbar_kwargs
591	)
592
593	# Create FacetGrid
594	fc = ImageFacetGrid(	×
595	data=da.compute(),
596	col=plot_kwargs.pop("col", None),
597	row=plot_kwargs.pop("row", None),
598	col_wrap=plot_kwargs.pop("col_wrap", None),
599	axes_pad=plot_kwargs.pop("axes_pad", None),
600	fig_kwargs=fig_kwargs,
601	cbar_kwargs=cbar_kwargs,
602	add_colorbar=add_colorbar,
603	aspect=plot_kwargs.pop("aspect", False),
604	facet_height=plot_kwargs.pop("facet_height", 3),
605	facet_aspect=plot_kwargs.pop("facet_aspect", 1),
606	)
607
608	# Plot the maps
609	fc = fc.map_dataarray(	×
610	_plot_orbit_image,
611	x=x,
612	y=y,
613	add_colorbar=False,
614	interpolation=interpolation,
615	cbar_kwargs=cbar_kwargs,
616	**plot_kwargs,
617	)
618
619	fc.remove_duplicated_axis_labels()	×
620
621	# Add colorbar
622	if add_colorbar:	×
623	fc.add_colorbar(**cbar_kwargs)	×
624
625	return fc	×
626
627
628	def plot_orbit_image(	1✔
629	da,
630	x="lon",
631	y="lat",
632	ax=None,
633	add_colorbar=True,
634	interpolation="nearest",
635	fig_kwargs={},
636	cbar_kwargs={},
637	**plot_kwargs,
638	):
639	"""Plot DataArray 2D field with cartopy."""
640	# Plot FacetGrid with xarray imshow
641	if "col" in plot_kwargs or "row" in plot_kwargs:	×
642	p = _plot_orbit_image_facetgrid(	×
643	da=da,
644	x=x,
645	y=y,
646	ax=ax,
647	add_colorbar=add_colorbar,
648	interpolation=interpolation,
649	fig_kwargs=fig_kwargs,
650	cbar_kwargs=cbar_kwargs,
651	**plot_kwargs,
652	)
653
654	# Plot with cartopy imshow
655	else:
656	da = da.squeeze() # remove time if dim=1	×
657	p = _plot_orbit_image(	×
658	da=da,
659	x=x,
660	y=y,
661	ax=ax,
662	add_colorbar=add_colorbar,
663	interpolation=interpolation,
664	fig_kwargs=fig_kwargs,
665	cbar_kwargs=cbar_kwargs,
666	**plot_kwargs,
667	)
668	# - Return imagepable
669	return p	×

ghiggi / gpm_api / 8053365388

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous