Skip to content

Commit bdda005

Browse files
authored
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
1 parent ee60fa6 commit bdda005

File tree

3 files changed

+200
-15
lines changed

3 files changed

+200
-15
lines changed

.pre-commit-config.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
repos:
22
- repo: https://github.com/psf/black
3-
rev: 24.4.2
3+
rev: 24.8.0
44
hooks:
55
- id: black
66
language_version: python3
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
#!/bin/env python
2+
"""
3+
Plot precipitation using custom colormap
4+
=============
5+
6+
This tutorial shows how to plot data using a custom colormap with a specific
7+
range of precipitation values.
8+
9+
"""
10+
11+
import os
12+
from datetime import datetime
13+
import matplotlib.pyplot as plt
14+
15+
import pysteps
16+
from pysteps import io, rcparams
17+
from pysteps.utils import conversion
18+
from pysteps.visualization import plot_precip_field
19+
from pysteps.datasets import download_pysteps_data, create_default_pystepsrc
20+
21+
22+
###############################################################################
23+
# Download the data if it is not available
24+
# ----------------------------------------
25+
#
26+
# The following code block downloads datasets from the pysteps-data repository
27+
# if it is not available on the disk. The dataset is used to demonstrate the
28+
# plotting of precipitation data using a custom colormap.
29+
30+
# Check if the pysteps-data repository is available (it would be pysteps-data in pysteps)
31+
# Implies that you are running this script from the `pysteps/examples` folder
32+
33+
if not os.path.exists(rcparams.data_sources["mrms"]["root_path"]):
34+
download_pysteps_data("pysteps_data")
35+
config_file_path = create_default_pystepsrc("pysteps_data")
36+
print(f"Configuration file has been created at {config_file_path}")
37+
38+
39+
###############################################################################
40+
# Read precipitation field
41+
# ------------------------
42+
#
43+
# First thing, load a frame from Multi-Radar Multi-Sensor dataset and convert it
44+
# to precipitation rate in mm/h.
45+
46+
# Define the dataset and the date for which you want to load data
47+
data_source = pysteps.rcparams.data_sources["mrms"]
48+
date = datetime(2019, 6, 10, 0, 2, 0) # Example date
49+
50+
# Extract the parameters from the data source
51+
root_path = data_source["root_path"]
52+
path_fmt = data_source["path_fmt"]
53+
fn_pattern = data_source["fn_pattern"]
54+
fn_ext = data_source["fn_ext"]
55+
importer_name = data_source["importer"]
56+
importer_kwargs = data_source["importer_kwargs"]
57+
timestep = data_source["timestep"]
58+
59+
# Find the frame in the archive for the specified date
60+
fns = io.find_by_date(
61+
date, root_path, path_fmt, fn_pattern, fn_ext, timestep, num_prev_files=1
62+
)
63+
64+
# Read the frame from the archive
65+
importer = io.get_method(importer_name, "importer")
66+
R, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
67+
68+
# Convert the reflectivity data to rain rate
69+
R, metadata = conversion.to_rainrate(R, metadata)
70+
71+
# Plot the first rainfall field from the loaded data
72+
plt.figure(figsize=(10, 5), dpi=300)
73+
plt.axis("off")
74+
plot_precip_field(R[0, :, :], geodata=metadata, axis="off")
75+
76+
plt.tight_layout()
77+
plt.show()
78+
79+
###############################################################################
80+
# Define the custom colormap
81+
# --------------------------
82+
#
83+
# Assume that the default colormap does not represent the precipitation values
84+
# in the desired range. In this case, you can define a custom colormap that will
85+
# be used to plot the precipitation data and pass the class instance to the
86+
# `plot_precip_field` function.
87+
#
88+
# It essential for the custom colormap to have the following attributes:
89+
#
90+
# - `cmap`: The colormap object.
91+
# - `norm`: The normalization object.
92+
# - `clevs`: The color levels for the colormap.
93+
#
94+
# `plot_precip_field` can handle each of the classes defined in the `matplotlib.colors`
95+
# https://matplotlib.org/stable/api/colors_api.html#colormaps
96+
# There must be as many colors in the colormap as there are levels in the color levels.
97+
98+
99+
# Define the custom colormap
100+
101+
from matplotlib import colors
102+
103+
104+
class ColormapConfig:
105+
def __init__(self):
106+
self.cmap = None
107+
self.norm = None
108+
self.clevs = None
109+
110+
self.build_colormap()
111+
112+
def build_colormap(self):
113+
# Define the colormap boundaries and colors
114+
# color_list = ['lightgrey', 'lightskyblue', 'blue', 'yellow', 'orange', 'red', 'darkred']
115+
color_list = ["blue", "navy", "yellow", "orange", "green", "brown", "red"]
116+
117+
self.clevs = [0.1, 0.5, 1.5, 2.5, 4, 6, 10] # mm/hr
118+
119+
# Create a ListedColormap object with the defined colors
120+
self.cmap = colors.ListedColormap(color_list)
121+
self.cmap.name = "Custom Colormap"
122+
123+
# Set the color for values above the maximum level
124+
self.cmap.set_over("darkmagenta")
125+
# Set the color for values below the minimum level
126+
self.cmap.set_under("none")
127+
# Set the color for missing values
128+
self.cmap.set_bad("gray", alpha=0.5)
129+
130+
# Create a BoundaryNorm object to normalize the data values to the colormap boundaries
131+
self.norm = colors.BoundaryNorm(self.clevs, self.cmap.N)
132+
133+
134+
# Create an instance of the ColormapConfig class
135+
config = ColormapConfig()
136+
137+
# Plot the precipitation field using the custom colormap
138+
plt.figure(figsize=(10, 5), dpi=300)
139+
plt.axis("off")
140+
plot_precip_field(R[0, :, :], geodata=metadata, axis="off", colormap_config=config)
141+
142+
plt.tight_layout()
143+
plt.show()

pysteps/visualization/precipfields.py

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def plot_precip_field(
4343
axis="on",
4444
cax=None,
4545
map_kwargs=None,
46+
colormap_config=None,
4647
):
4748
"""
4849
Function to plot a precipitation intensity or probability field with a
@@ -114,6 +115,11 @@ def plot_precip_field(
114115
cax : Axes_ object, optional
115116
Axes into which the colorbar will be drawn. If no axes is provided
116117
the colorbar axes are created next to the plot.
118+
colormap_config : ColormapConfig, optional
119+
Custom colormap configuration. If provided, this will override the
120+
colorscale parameter.
121+
The ColormapConfig class must have the following attributes: cmap,
122+
norm, clevs.
117123
118124
Other parameters
119125
----------------
@@ -158,20 +164,24 @@ def plot_precip_field(
158164
ax = get_basemap_axis(extent, ax=ax, geodata=geodata, map_kwargs=map_kwargs)
159165

160166
precip = np.ma.masked_invalid(precip)
161-
# plot rainfield
167+
168+
# Handle colormap configuration
169+
if colormap_config is None:
170+
cmap, norm, clevs, clevs_str = get_colormap(ptype, units, colorscale)
171+
else:
172+
cmap, norm, clevs = _validate_colormap_config(colormap_config, ptype)
173+
clevs_str = _dynamic_formatting_floats(clevs)
174+
175+
# Plot the precipitation field
162176
if regular_grid:
163-
im = _plot_field(precip, ax, ptype, units, colorscale, extent, origin=origin)
177+
im = _plot_field(precip, ax, extent, cmap, norm, origin=origin)
164178
else:
165-
im = _plot_field(
166-
precip, ax, ptype, units, colorscale, extent, x_grid=x_grid, y_grid=y_grid
167-
)
179+
im = _plot_field(precip, ax, extent, cmap, norm, x_grid=x_grid, y_grid=y_grid)
168180

169181
plt.title(title)
170182

171-
# add colorbar
183+
# Add colorbar
172184
if colorbar:
173-
# get colormap and color levels
174-
_, _, clevs, clevs_str = get_colormap(ptype, units, colorscale)
175185
if ptype in ["intensity", "depth"]:
176186
extend = "max"
177187
else:
@@ -202,14 +212,9 @@ def plot_precip_field(
202212
return ax
203213

204214

205-
def _plot_field(
206-
precip, ax, ptype, units, colorscale, extent, origin=None, x_grid=None, y_grid=None
207-
):
215+
def _plot_field(precip, ax, extent, cmap, norm, origin=None, x_grid=None, y_grid=None):
208216
precip = precip.copy()
209217

210-
# Get colormap and color levels
211-
cmap, norm, _, _ = get_colormap(ptype, units, colorscale)
212-
213218
if (x_grid is None) or (y_grid is None):
214219
im = ax.imshow(
215220
precip,
@@ -510,3 +515,40 @@ def _dynamic_formatting_floats(float_array, colorscale="pysteps"):
510515
labels.append(str(int(label)))
511516

512517
return labels
518+
519+
520+
def _validate_colormap_config(colormap_config, ptype):
521+
"""Validate the colormap configuration provided by the user."""
522+
523+
# Ensure colormap_config has the necessary attributes
524+
required_attrs = ["cmap", "norm", "clevs"]
525+
missing_attrs = [
526+
attr for attr in required_attrs if not hasattr(colormap_config, attr)
527+
]
528+
if missing_attrs:
529+
raise ValueError(
530+
f"colormap_config is missing required attributes: {', '.join(missing_attrs)}"
531+
)
532+
533+
# Ensure that ptype is appropriate when colormap_config is provided
534+
if ptype not in ["intensity", "depth"]:
535+
raise ValueError(
536+
"colormap_config is only supported for ptype='intensity' or 'depth'"
537+
)
538+
539+
cmap = colormap_config.cmap
540+
clevs = colormap_config.clevs
541+
542+
# Validate that the number of colors matches len(clevs)
543+
if isinstance(cmap, colors.ListedColormap):
544+
num_colors = len(cmap.colors)
545+
else:
546+
num_colors = cmap.N
547+
548+
expected_colors = len(clevs)
549+
if num_colors != expected_colors:
550+
raise ValueError(
551+
f"Number of colors in colormap (N={num_colors}) does not match len(clevs) (N={expected_colors})."
552+
)
553+
554+
return colormap_config.cmap, colormap_config.norm, colormap_config.clevs

0 commit comments

Comments
 (0)