Skip to content

Commit

Permalink
Merge pull request #35 from martinvonk/dev
Browse files Browse the repository at this point in the history
Update main to v0.4.1
  • Loading branch information
martinvonk authored Mar 28, 2024
2 parents 877cc3e + 9fa5949 commit 2549eab
Show file tree
Hide file tree
Showing 9 changed files with 338 additions and 479 deletions.
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
[![PyPI](https://img.shields.io/pypi/v/spei?style=flat-square)](https://pypi.org/project/spei/)
[![PyPi Supported Python Versions](https://img.shields.io/pypi/pyversions/spei?style=flat-square)](https://pypi.org/project/spei/)
[![Code Size](https://img.shields.io/github/languages/code-size/martinvonk/spei?style=flat-square)](https://pypi.org/project/spei/)
[![PyPi Downloads](https://img.shields.io/pypi/dm/spei?style=flat-square)
![License](https://img.shields.io/pypi/l/spei?style=flat-square)](https://pypi.org/project/spei/)
[![PyPi Downloads](https://img.shields.io/pypi/dm/spei?style=flat-square)](https://pypi.org/project/spei/)
[![License](https://img.shields.io/pypi/l/spei?style=flat-square)](https://pypi.org/project/spei/)
[![DOI](https://img.shields.io/badge/DOI-10.5281/zenodo.10816741-blue?style=flat-square)](https://doi.org/10.5281/zenodo.10816741)

[![Tests](https://img.shields.io/github/actions/workflow/status/martinvonk/spei/tests.yml?style=flat-square)](https://github.com/martinvonk/SPEI/actions/workflows/tests.yml)
Expand Down
47 changes: 19 additions & 28 deletions doc/examples/example03_drought_NL.ipynb

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion src/spei/_version.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from importlib import metadata
from platform import python_version

__version__ = "0.4.0"
__version__ = "0.4.1"


def show_versions() -> str:
Expand Down
325 changes: 309 additions & 16 deletions src/spei/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@
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 pandas import Series
from matplotlib.dates import date2num
from numpy import array, concatenate, linspace, meshgrid, reshape
from pandas import DatetimeIndex, Series
from scipy.stats import gaussian_kde

from ._typing import Axes
Expand All @@ -15,6 +17,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 +29,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 +46,40 @@ 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"):
colormap = roma(_r=True if "_r" in cmap else False)
else:
colormap = 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=colormap, 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:
datetime = DatetimeIndex(si.index).to_pydatetime()
x = date2num(datetime)
y = si.values.astype(float)
points = array([x, y]).T.reshape(-1, 1, 2)
segments = concatenate([points[:-1], points[1:]], axis=1)
lc = mpl.collections.LineCollection(
segments, cmap=colormap, norm=plt.Normalize(nmin, nmax)
)
lc.set_array(y)
lc.set_linewidth(1.2)
_ = ax.add_collection(lc)

ax.set_ylim(nmin, nmax)

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

return ax


def roma(_r: bool = False) -> mpl.colors.Colormap:
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
Loading

0 comments on commit 2549eab

Please sign in to comment.