Skip to content

Commit

Permalink
add roma colormap and plot si as colored line
Browse files Browse the repository at this point in the history
  • Loading branch information
@martinvonk
martinvonk committed Mar 22, 2024
1 parent 1752dd2 commit 2b57e2f
Show file tree
Hide file tree
Showing 4 changed files with 331 additions and 45 deletions.
47 changes: 19 additions & 28 deletions doc/examples/example03_drought_NL.ipynb
View file

Large diffs are not rendered by default.

320 changes: 305 additions & 15 deletions src/spei/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@
from itertools import cycle
from typing import List, Optional

import matplotlib as mpl
import matplotlib.pyplot as plt
from numpy import array, linspace, meshgrid, reshape
from numpy import array, concatenate, linspace, meshgrid, reshape
from pandas import Series
from scipy.stats import gaussian_kde

Expand All @@ -15,6 +16,8 @@ def si(
si: Series,
ybound: float = 3.0,
figsize: tuple = (6.5, 4),
cmap: str = "seismic_r",
background: bool = True,
ax: Optional[Axes] = None,
) -> Axes:
"""Plot the standardized index values as a time series.
Expand All @@ -25,6 +28,10 @@ def si(
Series of the standardized index
ybound : int, optional
Maximum and minimum ylim of plot
cmap: str, optional
Colormap for the background or line fill
background: bool, optional
Color the background if True, else color the line
figsize : tuple, optional
Figure size, by default (8, 4)
ax : matplotlib.Axes, optional
Expand All @@ -38,22 +45,38 @@ def si(
if ax is None:
_, ax = plt.subplots(figsize=figsize)

ax.plot(si.index, si.values, linewidth=1.0, color="k")
ax.axhline(0, linestyle="--", linewidth=1.0, color="k")

nmin = -ybound
nmax = ybound
droughts = si.to_numpy(dtype=float, copy=True)
droughts[droughts > 0] = 0
nodroughts = si.to_numpy(dtype=float, copy=True)
nodroughts[nodroughts < 0] = 0

x, y = meshgrid(si.index, linspace(nmin, nmax, 100))
ax.contourf(
x, y, y, cmap=plt.get_cmap("seismic_r"), levels=linspace(nmin, nmax, 100)
)
ax.fill_between(x=si.index, y1=droughts, y2=nmin, color="w")
ax.fill_between(x=si.index, y1=nodroughts, y2=nmax, color="w")

if cmap in ("roma", "roma_r"):
cmap = roma(_r=True if "_r" in cmap else False)
else:
cmap = plt.get_cmap(cmap)

if background:
ax.plot(si.index, si.values, linewidth=0.8, color="k")
ax.axhline(0, linestyle="--", linewidth=0.5, color="k")

droughts = si.to_numpy(dtype=float, copy=True)
droughts[droughts > 0] = 0
nodroughts = si.to_numpy(dtype=float, copy=True)
nodroughts[nodroughts < 0] = 0

x, y = meshgrid(si.index, linspace(nmin, nmax, 100))
ax.contourf(x, y, y, cmap=cmap, levels=linspace(nmin, nmax, 100))
ax.fill_between(x=si.index, y1=droughts, y2=nmin, color="w")
ax.fill_between(x=si.index, y1=nodroughts, y2=nmax, color="w")
else:
x = mpl.dates.date2num(si.index.to_pydatetime())
points = array([x, si.values]).T.reshape(-1, 1, 2)
segments = concatenate([points[:-1], points[1:]], axis=1)
lc = mpl.collections.LineCollection(
segments, cmap=cmap, norm=plt.Normalize(nmin, nmax)
)
lc.set_array(si.values)
lc.set_linewidth(1.2)
_ = ax.add_collection(lc)

ax.set_ylim(nmin, nmax)

return ax
Expand Down Expand Up @@ -122,3 +145,270 @@ def monthly_density(
ax.grid(True)

return ax


def roma(_r: bool = False) -> mpl.colors.LinearSegmentedColormap:
colors = [
[0.492325, 0.090787, 7.6e-05],
[0.49673, 0.102802, 0.003675],
[0.501125, 0.114034, 0.007134],
[0.505473, 0.124685, 0.010421],
[0.509813, 0.13489, 0.013817],
[0.514125, 0.144643, 0.016841],
[0.518397, 0.154036, 0.01972],
[0.522634, 0.163193, 0.022451],
[0.52685, 0.172041, 0.025034],
[0.531016, 0.180682, 0.027528],
[0.535142, 0.189147, 0.030132],
[0.539225, 0.197418, 0.032869],
[0.543266, 0.205524, 0.035925],
[0.547254, 0.213477, 0.038888],
[0.551203, 0.221318, 0.041994],
[0.5551, 0.229015, 0.045012],
[0.558954, 0.236607, 0.047967],
[0.562776, 0.244053, 0.051012],
[0.56654, 0.251456, 0.053998],
[0.570257, 0.258727, 0.057033],
[0.57394, 0.265922, 0.060051],
[0.577577, 0.273024, 0.063001],
[0.581178, 0.280048, 0.065873],
[0.584739, 0.286993, 0.068856],
[0.588249, 0.29388, 0.071712],
[0.591737, 0.300697, 0.074564],
[0.595174, 0.30745, 0.077376],
[0.598577, 0.314151, 0.080252],
[0.601948, 0.320797, 0.083076],
[0.605282, 0.327364, 0.085853],
[0.608593, 0.333908, 0.088711],
[0.611855, 0.340387, 0.091525],
[0.615092, 0.346826, 0.094279],
[0.618309, 0.353217, 0.096979],
[0.621487, 0.359572, 0.099753],
[0.62465, 0.365896, 0.102506],
[0.627785, 0.372174, 0.105218],
[0.630901, 0.378439, 0.107956],
[0.633994, 0.38467, 0.110736],
[0.637068, 0.390894, 0.11348],
[0.640129, 0.397078, 0.116219],
[0.643171, 0.403273, 0.11896],
[0.646208, 0.409445, 0.121674],
[0.649242, 0.415612, 0.12447],
[0.652261, 0.421776, 0.127267],
[0.655276, 0.427961, 0.130152],
[0.658297, 0.434132, 0.132974],
[0.661315, 0.440331, 0.135825],
[0.664334, 0.446543, 0.138777],
[0.667364, 0.45278, 0.141759],
[0.670402, 0.459031, 0.144786],
[0.67344, 0.465303, 0.147848],
[0.676494, 0.471628, 0.150966],
[0.679554, 0.477967, 0.15416],
[0.682637, 0.484345, 0.157428],
[0.685735, 0.490782, 0.160789],
[0.688849, 0.497239, 0.164248],
[0.691977, 0.503744, 0.167748],
[0.69512, 0.510297, 0.171371],
[0.698285, 0.516904, 0.175115],
[0.701472, 0.523556, 0.178956],
[0.704681, 0.530258, 0.182899],
[0.707908, 0.537021, 0.18701],
[0.711162, 0.543825, 0.191228],
[0.714426, 0.550692, 0.195605],
[0.717725, 0.557617, 0.200115],
[0.72103, 0.564592, 0.204829],
[0.724366, 0.57163, 0.209678],
[0.727717, 0.578727, 0.214708],
[0.731087, 0.585886, 0.219946],
[0.73448, 0.593098, 0.225346],
[0.737879, 0.600364, 0.230969],
[0.741297, 0.607699, 0.236808],
[0.744731, 0.615074, 0.242822],
[0.748175, 0.622516, 0.249103],
[0.751624, 0.630013, 0.255599],
[0.755074, 0.637558, 0.262334],
[0.758523, 0.645147, 0.269325],
[0.761963, 0.652786, 0.276547],
[0.765406, 0.66047, 0.28401],
[0.76882, 0.668197, 0.291755],
[0.772224, 0.675947, 0.299742],
[0.775598, 0.683727, 0.307981],
[0.778939, 0.691534, 0.316505],
[0.782242, 0.699353, 0.325257],
[0.785493, 0.707188, 0.334299],
[0.788687, 0.715019, 0.343564],
[0.791812, 0.722841, 0.353091],
[0.794864, 0.730654, 0.362848],
[0.797828, 0.738438, 0.372842],
[0.800699, 0.746187, 0.383072],
[0.803461, 0.753893, 0.393496],
[0.806107, 0.761537, 0.404131],
[0.808624, 0.769113, 0.41495],
[0.810997, 0.776617, 0.425936],
[0.813218, 0.784025, 0.437088],
[0.815281, 0.791322, 0.448363],
[0.817167, 0.798508, 0.459757],
[0.818865, 0.805562, 0.471253],
[0.820371, 0.812478, 0.482813],
[0.821675, 0.81924, 0.49444],
[0.822757, 0.825842, 0.50608],
[0.823613, 0.83227, 0.517753],
[0.824237, 0.838506, 0.5294],
[0.82462, 0.844553, 0.541003],
[0.824755, 0.850392, 0.552561],
[0.824634, 0.856022, 0.564037],
[0.824253, 0.861429, 0.575407],
[0.823606, 0.866611, 0.586656],
[0.822692, 0.871564, 0.597769],
[0.821502, 0.876273, 0.608734],
[0.820031, 0.880744, 0.619513],
[0.818285, 0.884972, 0.630101],
[0.816266, 0.888952, 0.64049],
[0.813955, 0.892681, 0.65066],
[0.811371, 0.89616, 0.6606],
[0.808502, 0.899389, 0.670314],
[0.805347, 0.902364, 0.679761],
[0.801918, 0.905092, 0.688973],
[0.79821, 0.907568, 0.697906],
[0.794228, 0.909794, 0.706586],
[0.789963, 0.911771, 0.714983],
[0.785431, 0.913506, 0.723105],
[0.780623, 0.914991, 0.730953],
[0.775551, 0.916237, 0.738521],
[0.770217, 0.917249, 0.745803],
[0.764624, 0.918014, 0.7528],
[0.758773, 0.918542, 0.759517],
[0.752665, 0.918838, 0.765957],
[0.746318, 0.918903, 0.772108],
[0.739731, 0.918739, 0.777981],
[0.732897, 0.91835, 0.783574],
[0.725834, 0.917737, 0.788893],
[0.718537, 0.916896, 0.793938],
[0.71103, 0.915828, 0.798708],
[0.7033, 0.914552, 0.803212],
[0.69536, 0.913054, 0.807458],
[0.687217, 0.911341, 0.811442],
[0.678874, 0.909414, 0.815168],
[0.670351, 0.90728, 0.818643],
[0.661635, 0.904933, 0.82188],
[0.652735, 0.902379, 0.824863],
[0.643676, 0.899624, 0.827614],
[0.634464, 0.896664, 0.830135],
[0.625086, 0.8935, 0.832423],
[0.615567, 0.890145, 0.834491],
[0.605914, 0.886598, 0.836336],
[0.596145, 0.882852, 0.83797],
[0.58626, 0.878925, 0.839394],
[0.576278, 0.874814, 0.840613],
[0.56619, 0.87052, 0.841632],
[0.556024, 0.866058, 0.842458],
[0.545813, 0.86142, 0.84309],
[0.535527, 0.856619, 0.843538],
[0.525217, 0.851658, 0.843811],
[0.514877, 0.846543, 0.843909],
[0.504532, 0.841286, 0.843838],
[0.49418, 0.835884, 0.843605],
[0.483842, 0.83035, 0.843216],
[0.473543, 0.824685, 0.842677],
[0.463293, 0.818902, 0.84199],
[0.453115, 0.81301, 0.84116],
[0.443005, 0.807018, 0.840201],
[0.432982, 0.80092, 0.839108],
[0.423057, 0.794738, 0.8379],
[0.413269, 0.788474, 0.836571],
[0.403599, 0.782136, 0.835134],
[0.394062, 0.775727, 0.833595],
[0.384684, 0.769261, 0.831959],
[0.375472, 0.762745, 0.830227],
[0.366446, 0.756186, 0.828409],
[0.357593, 0.749581, 0.826514],
[0.348916, 0.742958, 0.82454],
[0.340447, 0.736291, 0.822504],
[0.332185, 0.729618, 0.820394],
[0.324113, 0.722919, 0.81823],
[0.316282, 0.716224, 0.816018],
[0.308653, 0.709514, 0.813746],
[0.301225, 0.702804, 0.811437],
[0.294036, 0.696099, 0.809086],
[0.287067, 0.689401, 0.806696],
[0.280322, 0.682707, 0.804272],
[0.273816, 0.676034, 0.801816],
[0.267484, 0.669377, 0.799343],
[0.261414, 0.662739, 0.796836],
[0.255543, 0.65612, 0.79432],
[0.24986, 0.649522, 0.791778],
[0.244403, 0.642944, 0.789229],
[0.239162, 0.636402, 0.786665],
[0.234103, 0.629874, 0.78409],
[0.229233, 0.623385, 0.781503],
[0.22453, 0.616915, 0.778909],
[0.220062, 0.610472, 0.776315],
[0.215727, 0.604065, 0.77371],
[0.211566, 0.597675, 0.771114],
[0.207553, 0.591317, 0.768502],
[0.203709, 0.584989, 0.765902],
[0.199968, 0.578679, 0.763292],
[0.196442, 0.572403, 0.760691],
[0.192988, 0.566143, 0.758086],
[0.189677, 0.559913, 0.755482],
[0.186487, 0.553693, 0.752873],
[0.183398, 0.547501, 0.750276],
[0.180424, 0.541321, 0.747679],
[0.177586, 0.535164, 0.745074],
[0.174797, 0.529026, 0.742478],
[0.172082, 0.522885, 0.73988],
[0.169502, 0.516755, 0.737273],
[0.166959, 0.510628, 0.734673],
[0.164485, 0.504518, 0.732066],
[0.162089, 0.498397, 0.729453],
[0.159699, 0.492262, 0.726836],
[0.157413, 0.486128, 0.724215],
[0.15518, 0.479992, 0.721583],
[0.15296, 0.473828, 0.718948],
[0.150805, 0.467661, 0.716306],
[0.148654, 0.461471, 0.713644],
[0.146549, 0.45526, 0.710977],
[0.144488, 0.449025, 0.708292],
[0.142409, 0.44276, 0.705595],
[0.140344, 0.436466, 0.702886],
[0.13831, 0.430142, 0.700156],
[0.13621, 0.423782, 0.697413],
[0.134205, 0.417407, 0.694652],
[0.132133, 0.410984, 0.691884],
[0.130086, 0.404521, 0.689094],
[0.127967, 0.398042, 0.686279],
[0.125859, 0.391517, 0.683451],
[0.123663, 0.384945, 0.6806],
[0.121475, 0.378353, 0.677747],
[0.119286, 0.371714, 0.67486],
[0.117003, 0.365049, 0.671962],
[0.114652, 0.358343, 0.669046],
[0.112324, 0.351597, 0.666107],
[0.10989, 0.344819, 0.663159],
[0.107324, 0.337994, 0.660182],
[0.104641, 0.33113, 0.657195],
[0.101951, 0.324237, 0.65418],
[0.099119, 0.317305, 0.651154],
[0.096135, 0.310338, 0.648101],
[0.093031, 0.303295, 0.645033],
[0.089832, 0.29624, 0.64195],
[0.086378, 0.28914, 0.638844],
[0.082771, 0.281987, 0.635717],
[0.078888, 0.274774, 0.632572],
[0.074823, 0.267513, 0.629402],
[0.070429, 0.260237, 0.626217],
[0.065707, 0.252891, 0.623014],
[0.060588, 0.245475, 0.619791],
[0.054957, 0.238038, 0.616544],
[0.048861, 0.230521, 0.613271],
[0.041963, 0.22298, 0.609992],
[0.034076, 0.215343, 0.606696],
[0.026246, 0.207675, 0.603381],
[0.018222, 0.199913, 0.600048],
[0.009824, 0.192129, 0.596704],
]
cmap = mpl.colors.LinearSegmentedColormap.from_list(
name="roma",
colors=list(reversed(colors)) if _r else colors,
N=len(colors),
)
return cmap
4 changes: 4 additions & 0 deletions tests/test_plots.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,5 +11,9 @@ def test_plot_si(si: Series) -> None:
_ = plot_si(si)


def test_plot_si_no_background(si: Series) -> None:
_ = plot_si(si, cmap="roma_r", background=False)


def test_plot_monthly_density(si: Series) -> None:
_ = monthly_density(si, years=[2011], months=[1, 2, 3, 4, 5])
5 changes: 3 additions & 2 deletions tests/test_si.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
from pandas import Series, Timestamp
from scipy.stats import norm

from spei import SI, sgi, spei, spi, ssfi
from scipy.stats import norm


def test_spi(prec: Series) -> None:
precr = prec.rolling("30D", min_periods=30).sum().dropna()
Expand Down Expand Up @@ -31,4 +32,4 @@ def test_SI(prec: Series) -> None:
si.fit_distribution()
si.pdf()
dist = si.get_dist(Timestamp("2010-01-01"))
dist.ks_test()
dist.ks_test()

0 comments on commit 2b57e2f

Please sign in to comment.