From bdda005cb5d8db2737f2446c95abbb84930277c7 Mon Sep 17 00:00:00 2001
From: Aleksei Rutkovskii <70371908+rutkovskii@users.noreply.github.com>
Date: Wed, 18 Sep 2024 04:30:19 -0400
Subject: [PATCH] Support for custom colormap in `precipfield.py` and
 `plot_custom_precipitation_range.py` example (#433)

* [MODIFY] Add support for custom colormap configuration in `precipfield.py`

This update allows user-defined ranges and colors for plots. It supports intensity (tested) and depth (untested), but not probability.

[ADD] Add `plot_custom_precipitation_range.py` example demonstrating how to create a custom config and use it for plotting.

* [MODIFY] removed file saving logic from
---
 .pre-commit-config.yaml                     |   2 +-
 examples/plot_custom_precipitation_range.py | 143 ++++++++++++++++++++
 pysteps/visualization/precipfields.py       |  70 ++++++++--
 3 files changed, 200 insertions(+), 15 deletions(-)
 create mode 100644 examples/plot_custom_precipitation_range.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6421a7ef1..6ffe446bf 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/psf/black
-    rev: 24.4.2
+    rev: 24.8.0
     hooks:
     - id: black
       language_version: python3
diff --git a/examples/plot_custom_precipitation_range.py b/examples/plot_custom_precipitation_range.py
new file mode 100644
index 000000000..4f1a98df1
--- /dev/null
+++ b/examples/plot_custom_precipitation_range.py
@@ -0,0 +1,143 @@
+#!/bin/env python
+"""
+Plot precipitation using custom colormap
+=============
+
+This tutorial shows how to plot data using a custom colormap with a specific
+range of precipitation values.
+
+"""
+
+import os
+from datetime import datetime
+import matplotlib.pyplot as plt
+
+import pysteps
+from pysteps import io, rcparams
+from pysteps.utils import conversion
+from pysteps.visualization import plot_precip_field
+from pysteps.datasets import download_pysteps_data, create_default_pystepsrc
+
+
+###############################################################################
+# Download the data if it is not available
+# ----------------------------------------
+#
+# The following code block downloads datasets from the pysteps-data repository
+# if it is not available on the disk. The dataset is used to demonstrate the
+# plotting of precipitation data using a custom colormap.
+
+# Check if the pysteps-data repository is available (it would be pysteps-data in pysteps)
+# Implies that you are running this script from the `pysteps/examples` folder
+
+if not os.path.exists(rcparams.data_sources["mrms"]["root_path"]):
+    download_pysteps_data("pysteps_data")
+    config_file_path = create_default_pystepsrc("pysteps_data")
+    print(f"Configuration file has been created at {config_file_path}")
+
+
+###############################################################################
+# Read precipitation field
+# ------------------------
+#
+# First thing, load a frame from Multi-Radar Multi-Sensor dataset and convert it
+# to precipitation rate in mm/h.
+
+# Define the dataset and the date for which you want to load data
+data_source = pysteps.rcparams.data_sources["mrms"]
+date = datetime(2019, 6, 10, 0, 2, 0)  # Example date
+
+# Extract the parameters from the data source
+root_path = data_source["root_path"]
+path_fmt = data_source["path_fmt"]
+fn_pattern = data_source["fn_pattern"]
+fn_ext = data_source["fn_ext"]
+importer_name = data_source["importer"]
+importer_kwargs = data_source["importer_kwargs"]
+timestep = data_source["timestep"]
+
+# Find the frame in the archive for the specified date
+fns = io.find_by_date(
+    date, root_path, path_fmt, fn_pattern, fn_ext, timestep, num_prev_files=1
+)
+
+# Read the frame from the archive
+importer = io.get_method(importer_name, "importer")
+R, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+
+# Convert the reflectivity data to rain rate
+R, metadata = conversion.to_rainrate(R, metadata)
+
+# Plot the first rainfall field from the loaded data
+plt.figure(figsize=(10, 5), dpi=300)
+plt.axis("off")
+plot_precip_field(R[0, :, :], geodata=metadata, axis="off")
+
+plt.tight_layout()
+plt.show()
+
+###############################################################################
+# Define the custom colormap
+# --------------------------
+#
+# Assume that the default colormap does not represent the precipitation values
+# in the desired range. In this case, you can define a custom colormap that will
+# be used to plot the precipitation data and pass the class instance to the
+# `plot_precip_field` function.
+#
+# It essential for the custom colormap to have the following attributes:
+#
+# - `cmap`: The colormap object.
+# - `norm`: The normalization object.
+# - `clevs`: The color levels for the colormap.
+#
+# `plot_precip_field` can handle each of the classes defined in the `matplotlib.colors`
+# https://matplotlib.org/stable/api/colors_api.html#colormaps
+# There must be as many colors in the colormap as there are levels in the color levels.
+
+
+# Define the custom colormap
+
+from matplotlib import colors
+
+
+class ColormapConfig:
+    def __init__(self):
+        self.cmap = None
+        self.norm = None
+        self.clevs = None
+
+        self.build_colormap()
+
+    def build_colormap(self):
+        # Define the colormap boundaries and colors
+        # color_list = ['lightgrey', 'lightskyblue', 'blue', 'yellow', 'orange', 'red', 'darkred']
+        color_list = ["blue", "navy", "yellow", "orange", "green", "brown", "red"]
+
+        self.clevs = [0.1, 0.5, 1.5, 2.5, 4, 6, 10]  # mm/hr
+
+        # Create a ListedColormap object with the defined colors
+        self.cmap = colors.ListedColormap(color_list)
+        self.cmap.name = "Custom Colormap"
+
+        # Set the color for values above the maximum level
+        self.cmap.set_over("darkmagenta")
+        # Set the color for values below the minimum level
+        self.cmap.set_under("none")
+        # Set the color for missing values
+        self.cmap.set_bad("gray", alpha=0.5)
+
+        # Create a BoundaryNorm object to normalize the data values to the colormap boundaries
+        self.norm = colors.BoundaryNorm(self.clevs, self.cmap.N)
+
+
+# Create an instance of the ColormapConfig class
+config = ColormapConfig()
+
+# Plot the precipitation field using the custom colormap
+plt.figure(figsize=(10, 5), dpi=300)
+plt.axis("off")
+plot_precip_field(R[0, :, :], geodata=metadata, axis="off", colormap_config=config)
+
+plt.tight_layout()
+plt.show()
diff --git a/pysteps/visualization/precipfields.py b/pysteps/visualization/precipfields.py
index da231ffb4..8b03ed807 100644
--- a/pysteps/visualization/precipfields.py
+++ b/pysteps/visualization/precipfields.py
@@ -43,6 +43,7 @@ def plot_precip_field(
     axis="on",
     cax=None,
     map_kwargs=None,
+    colormap_config=None,
 ):
     """
     Function to plot a precipitation intensity or probability field with a
@@ -114,6 +115,11 @@ def plot_precip_field(
     cax : Axes_ object, optional
         Axes into which the colorbar will be drawn. If no axes is provided
         the colorbar axes are created next to the plot.
+    colormap_config : ColormapConfig, optional
+        Custom colormap configuration. If provided, this will override the
+        colorscale parameter.
+        The ColormapConfig class must have the following attributes: cmap,
+        norm, clevs.
 
     Other parameters
     ----------------
@@ -158,20 +164,24 @@ def plot_precip_field(
     ax = get_basemap_axis(extent, ax=ax, geodata=geodata, map_kwargs=map_kwargs)
 
     precip = np.ma.masked_invalid(precip)
-    # plot rainfield
+
+    # Handle colormap configuration
+    if colormap_config is None:
+        cmap, norm, clevs, clevs_str = get_colormap(ptype, units, colorscale)
+    else:
+        cmap, norm, clevs = _validate_colormap_config(colormap_config, ptype)
+        clevs_str = _dynamic_formatting_floats(clevs)
+
+    # Plot the precipitation field
     if regular_grid:
-        im = _plot_field(precip, ax, ptype, units, colorscale, extent, origin=origin)
+        im = _plot_field(precip, ax, extent, cmap, norm, origin=origin)
     else:
-        im = _plot_field(
-            precip, ax, ptype, units, colorscale, extent, x_grid=x_grid, y_grid=y_grid
-        )
+        im = _plot_field(precip, ax, extent, cmap, norm, x_grid=x_grid, y_grid=y_grid)
 
     plt.title(title)
 
-    # add colorbar
+    # Add colorbar
     if colorbar:
-        # get colormap and color levels
-        _, _, clevs, clevs_str = get_colormap(ptype, units, colorscale)
         if ptype in ["intensity", "depth"]:
             extend = "max"
         else:
@@ -202,14 +212,9 @@ def plot_precip_field(
     return ax
 
 
-def _plot_field(
-    precip, ax, ptype, units, colorscale, extent, origin=None, x_grid=None, y_grid=None
-):
+def _plot_field(precip, ax, extent, cmap, norm, origin=None, x_grid=None, y_grid=None):
     precip = precip.copy()
 
-    # Get colormap and color levels
-    cmap, norm, _, _ = get_colormap(ptype, units, colorscale)
-
     if (x_grid is None) or (y_grid is None):
         im = ax.imshow(
             precip,
@@ -510,3 +515,40 @@ def _dynamic_formatting_floats(float_array, colorscale="pysteps"):
             labels.append(str(int(label)))
 
     return labels
+
+
+def _validate_colormap_config(colormap_config, ptype):
+    """Validate the colormap configuration provided by the user."""
+
+    # Ensure colormap_config has the necessary attributes
+    required_attrs = ["cmap", "norm", "clevs"]
+    missing_attrs = [
+        attr for attr in required_attrs if not hasattr(colormap_config, attr)
+    ]
+    if missing_attrs:
+        raise ValueError(
+            f"colormap_config is missing required attributes: {', '.join(missing_attrs)}"
+        )
+
+    # Ensure that ptype is appropriate when colormap_config is provided
+    if ptype not in ["intensity", "depth"]:
+        raise ValueError(
+            "colormap_config is only supported for ptype='intensity' or 'depth'"
+        )
+
+    cmap = colormap_config.cmap
+    clevs = colormap_config.clevs
+
+    # Validate that the number of colors matches len(clevs)
+    if isinstance(cmap, colors.ListedColormap):
+        num_colors = len(cmap.colors)
+    else:
+        num_colors = cmap.N
+
+    expected_colors = len(clevs)
+    if num_colors != expected_colors:
+        raise ValueError(
+            f"Number of colors in colormap (N={num_colors}) does not match len(clevs) (N={expected_colors})."
+        )
+
+    return colormap_config.cmap, colormap_config.norm, colormap_config.clevs