ClimateGraph 0.4.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ClimateGraph/__init__.py +9 -0
- ClimateGraph/appkernel.py +148 -0
- ClimateGraph/cli.py +63 -0
- ClimateGraph/data/__init__.py +8 -0
- ClimateGraph/data/data.py +508 -0
- ClimateGraph/data/point_surface.py +91 -0
- ClimateGraph/data/regular_grid.py +25 -0
- ClimateGraph/data/satellite_swath.py +21 -0
- ClimateGraph/domain/__init__.py +11 -0
- ClimateGraph/domain/domain.py +181 -0
- ClimateGraph/domain/domains.py +279 -0
- ClimateGraph/plot/__init__.py +11 -0
- ClimateGraph/plot/plot.py +402 -0
- ClimateGraph/plot/plots.py +1054 -0
- ClimateGraph/plot/primitive.py +99 -0
- ClimateGraph/plot/primitives.py +160 -0
- ClimateGraph/reader/__init__.py +15 -0
- ClimateGraph/reader/reader/__init__.py +14 -0
- ClimateGraph/reader/reader/point_surface/__init__.py +7 -0
- ClimateGraph/reader/reader/point_surface/csv_base.py +135 -0
- ClimateGraph/reader/reader/point_surface/default.py +36 -0
- ClimateGraph/reader/reader/point_surface/dmc.py +10 -0
- ClimateGraph/reader/reader/point_surface/sinca.py +10 -0
- ClimateGraph/reader/reader/point_surface/single_file.py +60 -0
- ClimateGraph/reader/reader/point_surface/station_per_file.py +55 -0
- ClimateGraph/reader/reader/point_surface/variable_per_file.py +91 -0
- ClimateGraph/reader/reader/reader.py +353 -0
- ClimateGraph/reader/reader/regular_grid/__init__.py +7 -0
- ClimateGraph/reader/reader/regular_grid/chimere.py +24 -0
- ClimateGraph/reader/reader/regular_grid/default.py +51 -0
- ClimateGraph/reader/reader/regular_grid/wrf.py +53 -0
- ClimateGraph/utils/__init__.py +1 -0
- ClimateGraph/utils/control_model.py +292 -0
- ClimateGraph/utils/dataset_utils.py +287 -0
- ClimateGraph/utils/general_utils.py +259 -0
- ClimateGraph/utils/parser.py +269 -0
- ClimateGraph/utils/registry.py +81 -0
- ClimateGraph/utils/resample_engine.py +156 -0
- climategraph-0.4.0.dist-info/METADATA +177 -0
- climategraph-0.4.0.dist-info/RECORD +44 -0
- climategraph-0.4.0.dist-info/WHEEL +5 -0
- climategraph-0.4.0.dist-info/entry_points.txt +2 -0
- climategraph-0.4.0.dist-info/licenses/LICENSE +21 -0
- climategraph-0.4.0.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,1054 @@
|
|
|
1
|
+
"""The concrete plot types and their config models."""
|
|
2
|
+
|
|
3
|
+
import contextlib
|
|
4
|
+
import math
|
|
5
|
+
from typing import Literal
|
|
6
|
+
|
|
7
|
+
import cartopy.feature as cfeature
|
|
8
|
+
import matplotlib as mpl
|
|
9
|
+
import matplotlib.pyplot as plt
|
|
10
|
+
import numpy as np
|
|
11
|
+
import xarray as xr
|
|
12
|
+
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
|
13
|
+
|
|
14
|
+
mpl.use("Agg")
|
|
15
|
+
|
|
16
|
+
from ClimateGraph.data import PointSurface, RegularGrid
|
|
17
|
+
from ClimateGraph.utils.dataset_utils import change_unit, dim_reduction, time_resampling
|
|
18
|
+
from ClimateGraph.utils.general_utils import (
|
|
19
|
+
CRSEnum,
|
|
20
|
+
ReductionMethodEnum,
|
|
21
|
+
TimeBucketEnum,
|
|
22
|
+
TimestepEnum,
|
|
23
|
+
manage_time_interval,
|
|
24
|
+
normalize_time,
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
from . import primitives as _primitives # noqa: F401 registers primitive configs
|
|
28
|
+
from .plot import Plot
|
|
29
|
+
from .primitive import Primitive
|
|
30
|
+
from .primitives import _drop_nan_points
|
|
31
|
+
|
|
32
|
+
# Discriminated union over the registered primitive configs (Series/ContourFill/
|
|
33
|
+
# Points). Built here — after importing `primitives` — so every primitive config
|
|
34
|
+
# is registered before CustomPlotConfig references it.
|
|
35
|
+
PrimitiveModel = Primitive.build_config_union()
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _pad_extent(
|
|
39
|
+
lon_min: float,
|
|
40
|
+
lon_max: float,
|
|
41
|
+
lat_min: float,
|
|
42
|
+
lat_max: float,
|
|
43
|
+
padding: float,
|
|
44
|
+
) -> tuple[float, float, float, float]:
|
|
45
|
+
"""Grow a (lon_min, lon_max, lat_min, lat_max) extent outward by
|
|
46
|
+
``padding`` (a fraction of each axis span) so auto-computed map bounds don't
|
|
47
|
+
clip markers sitting on the edge. A zero-width span (e.g. a single point)
|
|
48
|
+
falls back to a fixed 0.5-degree pad so ``set_extent`` stays valid.
|
|
49
|
+
"""
|
|
50
|
+
lon_span = lon_max - lon_min
|
|
51
|
+
lat_span = lat_max - lat_min
|
|
52
|
+
lon_pad = lon_span * padding if lon_span else 0.5
|
|
53
|
+
lat_pad = lat_span * padding if lat_span else 0.5
|
|
54
|
+
return (
|
|
55
|
+
lon_min - lon_pad,
|
|
56
|
+
lon_max + lon_pad,
|
|
57
|
+
lat_min - lat_pad,
|
|
58
|
+
lat_max + lat_pad,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _unit_label(name: str, unit: str | None) -> str:
|
|
63
|
+
"""Format an axis/colorbar label as ``name [unit]``.
|
|
64
|
+
|
|
65
|
+
Single style for every plot (square brackets), and omits the unit entirely
|
|
66
|
+
when it's unknown (``None``) so labels never read ``"Temperatura [None]"``.
|
|
67
|
+
"""
|
|
68
|
+
return f"{name} [{unit}]" if unit else name
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
class GridConfig(BaseModel):
|
|
72
|
+
"""Dashed reference lines that make an axis's resolution legible.
|
|
73
|
+
|
|
74
|
+
``x``/``y`` accept either ``"ticks"`` (one line per major tick — the default)
|
|
75
|
+
or an integer N (exactly N evenly-spaced interior lines, none on the spines).
|
|
76
|
+
``offset`` shifts every line by a *fraction of the spacing*, so ``0.5`` puts
|
|
77
|
+
them midway between neighbours; shifted lines that fall outside the axis are
|
|
78
|
+
dropped. Set an axis to ``None`` to leave it undecorated.
|
|
79
|
+
"""
|
|
80
|
+
|
|
81
|
+
x: int | Literal["ticks"] | None = Field(default=None)
|
|
82
|
+
y: int | Literal["ticks"] | None = Field(default=None)
|
|
83
|
+
offset: float = Field(default=0.0)
|
|
84
|
+
color: str = Field(default="grey")
|
|
85
|
+
alpha: float = Field(default=0.4)
|
|
86
|
+
linestyle: str = Field(default="--")
|
|
87
|
+
linewidth: float = Field(default=0.8)
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
class BasePlotConfig(BaseModel):
|
|
91
|
+
"""Base configuration as for all plots, Pydantic Model. Used to manage common arguments."""
|
|
92
|
+
|
|
93
|
+
model_config = ConfigDict(extra="allow")
|
|
94
|
+
|
|
95
|
+
filename: str | None = Field(default=None)
|
|
96
|
+
domains: list[str] = Field(default_factory=list)
|
|
97
|
+
vars: list[str] | dict[str, str]
|
|
98
|
+
dim_reduce: dict[str, str | dict] | None = Field(default=None)
|
|
99
|
+
# Axis decoration, applied to every plot type via Plot._finalize.
|
|
100
|
+
grid: GridConfig | None = Field(default=None)
|
|
101
|
+
xticks: dict[float, str] | None = Field(default=None)
|
|
102
|
+
yticks: dict[float, str] | None = Field(default=None)
|
|
103
|
+
|
|
104
|
+
@field_validator("vars", mode="before")
|
|
105
|
+
@classmethod
|
|
106
|
+
def _wrap_single_var(cls, v):
|
|
107
|
+
"""Coerce a lone ``vars: T2`` string into ``["T2"]`` so plot methods
|
|
108
|
+
iterate over variable names, not the characters of a single name."""
|
|
109
|
+
if isinstance(v, str):
|
|
110
|
+
return [v]
|
|
111
|
+
return v
|
|
112
|
+
|
|
113
|
+
@field_validator("grid", mode="before")
|
|
114
|
+
@classmethod
|
|
115
|
+
def _coerce_grid(cls, v):
|
|
116
|
+
"""Accept the shorthand forms: ``grid: true`` (one line per major tick),
|
|
117
|
+
``grid: false`` (off) and ``grid: 12`` (12 lines on both axes).
|
|
118
|
+
|
|
119
|
+
The shorthands name both axes explicitly, because the dict form's ``x``/``y``
|
|
120
|
+
default to ``None`` — so ``grid: {x: 12}`` decorates only x and leaves y clean.
|
|
121
|
+
"""
|
|
122
|
+
# bool must be checked first — in Python `True` is also an `int`.
|
|
123
|
+
if isinstance(v, bool):
|
|
124
|
+
return {"x": "ticks", "y": "ticks"} if v else None
|
|
125
|
+
if isinstance(v, int):
|
|
126
|
+
return {"x": v, "y": v}
|
|
127
|
+
return v
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
class TimeSeriesConfig(BasePlotConfig):
|
|
131
|
+
"""Timeseries plot configuration as Pydantic Model."""
|
|
132
|
+
|
|
133
|
+
type: Literal["timeseries", "ts", "time-series"]
|
|
134
|
+
data: str | list[str]
|
|
135
|
+
time: str | list[str] | None = Field(default=None)
|
|
136
|
+
timestep: TimestepEnum | None = Field(default=None)
|
|
137
|
+
reduction_method: ReductionMethodEnum = Field(default=ReductionMethodEnum.mean)
|
|
138
|
+
|
|
139
|
+
@field_validator("data", mode="before")
|
|
140
|
+
@classmethod
|
|
141
|
+
def _wrap_single_data(cls, v):
|
|
142
|
+
"""Accept a single dataset name as well as a list."""
|
|
143
|
+
return [v] if isinstance(v, str) else v
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
class Timeseries(Plot):
|
|
147
|
+
"""Plot class. Implements all particular operations for Timeseries plot creation."""
|
|
148
|
+
|
|
149
|
+
config = TimeSeriesConfig
|
|
150
|
+
aliases = ["ts", "time-series"]
|
|
151
|
+
|
|
152
|
+
def plot(self):
|
|
153
|
+
"""Iterate over self.plot_config.time entries, rendering once per entry."""
|
|
154
|
+
for time_interval in normalize_time(self.plot_config.time):
|
|
155
|
+
self._plot_one(time_interval)
|
|
156
|
+
|
|
157
|
+
def _plot_one(self, time_interval: str | None):
|
|
158
|
+
"""Timeseries plotting method for a single time entry.
|
|
159
|
+
The process goes as follows:
|
|
160
|
+
1) Process arguments.
|
|
161
|
+
2) Process base data: Time resampling and aligning, and unit conversion.
|
|
162
|
+
3) Process other data: Space and Time resampling. Time alignment and unit conversion.
|
|
163
|
+
4) Iterate through Domains and Vars.
|
|
164
|
+
4.1) Apply domain to data objects.
|
|
165
|
+
4.2) Reduce variable to time dimension.
|
|
166
|
+
4.3) Plot variable for each data object.
|
|
167
|
+
5) Save figure.
|
|
168
|
+
"""
|
|
169
|
+
# TODO: remove time intervals with nan values.
|
|
170
|
+
|
|
171
|
+
vars = self.plot_config.vars
|
|
172
|
+
domains = self.resolve_domains()
|
|
173
|
+
timestep = self.plot_config.timestep
|
|
174
|
+
data_names = self.plot_config.data
|
|
175
|
+
first = self.data[data_names[0]]
|
|
176
|
+
|
|
177
|
+
for dom_name, dom in domains.items():
|
|
178
|
+
# Apply the domain to each raw Data object (Data -> Data; resample
|
|
179
|
+
# domains reproject here), then time-filter. Canonical var names — no
|
|
180
|
+
# {var}__{dataset} mangling; each dataset is handled on its own.
|
|
181
|
+
prepared = {}
|
|
182
|
+
for ds_name in data_names:
|
|
183
|
+
d = self.data[ds_name]
|
|
184
|
+
d = dom.apply(d) if dom is not None else d
|
|
185
|
+
obj = time_resampling(
|
|
186
|
+
d.obj, timestep=timestep, time_interval=time_interval
|
|
187
|
+
)
|
|
188
|
+
prepared[ds_name] = (d, obj)
|
|
189
|
+
|
|
190
|
+
for variable in vars:
|
|
191
|
+
# dict vars -> convert to the requested unit; list vars -> keep
|
|
192
|
+
# native (dst None) and label with the first dataset's declared unit.
|
|
193
|
+
dst_unit = vars[variable] if isinstance(vars, dict) else None
|
|
194
|
+
label_unit = (
|
|
195
|
+
vars[variable]
|
|
196
|
+
if isinstance(vars, dict)
|
|
197
|
+
else first.var_unit(variable)
|
|
198
|
+
)
|
|
199
|
+
|
|
200
|
+
figure = plt.figure(**self.figure_kwargs())
|
|
201
|
+
ax = figure.add_subplot(1, 1, 1)
|
|
202
|
+
|
|
203
|
+
for ds_name, (d, obj) in prepared.items():
|
|
204
|
+
da = obj[variable]
|
|
205
|
+
if self.plot_config.dim_reduce:
|
|
206
|
+
da = dim_reduction(da, self.plot_config.dim_reduce)
|
|
207
|
+
da = da.reduce(
|
|
208
|
+
self.plot_config.reduction_method.func,
|
|
209
|
+
list(set(da.dims) - {"time"}),
|
|
210
|
+
)
|
|
211
|
+
da = change_unit(da, d.var_unit(variable), dst_unit)
|
|
212
|
+
line = da.plot.line(ax=ax)
|
|
213
|
+
line[0].set_label(ds_name)
|
|
214
|
+
|
|
215
|
+
ax.legend()
|
|
216
|
+
|
|
217
|
+
title = self.plot_kwargs.get(
|
|
218
|
+
"title", f"Timeseries comparison of {variable}"
|
|
219
|
+
)
|
|
220
|
+
xlabel = self.plot_kwargs.get("xlabel", "Time")
|
|
221
|
+
ylabel = self.plot_kwargs.get(
|
|
222
|
+
"ylabel", _unit_label(variable, label_unit)
|
|
223
|
+
)
|
|
224
|
+
|
|
225
|
+
ax.set_xlabel(xlabel)
|
|
226
|
+
ax.set_ylabel(ylabel)
|
|
227
|
+
figure.suptitle(title)
|
|
228
|
+
|
|
229
|
+
start, end = manage_time_interval(time_interval)
|
|
230
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
231
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
232
|
+
format = self.plot_kwargs.get("format", "jpg")
|
|
233
|
+
filename = (
|
|
234
|
+
f"ts-{dom_name}-{variable}-{start}_{end}.{format}"
|
|
235
|
+
if self.plot_config.filename is None
|
|
236
|
+
else self.plot_config.filename
|
|
237
|
+
)
|
|
238
|
+
|
|
239
|
+
self._finalize(figure, filename)
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
class ScatterConfig(BasePlotConfig):
|
|
243
|
+
"""Scatter plot configuration Pydantic model."""
|
|
244
|
+
|
|
245
|
+
type: Literal["scatter", "sc"]
|
|
246
|
+
# Exactly two datasets — [x-axis, y-axis]. Co-location (pairing per site)
|
|
247
|
+
# is expressed with a `resample_to` domain; pairing per time needs none.
|
|
248
|
+
data: list[str]
|
|
249
|
+
time: str | list[str] | None = Field(default=None)
|
|
250
|
+
dimension: str = Field(default="time")
|
|
251
|
+
timestep: TimestepEnum | None = Field(default=None)
|
|
252
|
+
reduction_method: ReductionMethodEnum = Field(default=ReductionMethodEnum.mean)
|
|
253
|
+
|
|
254
|
+
@field_validator("data")
|
|
255
|
+
@classmethod
|
|
256
|
+
def _exactly_two(cls, v):
|
|
257
|
+
"""Require exactly two datasets: the x axis and the y axis."""
|
|
258
|
+
if len(v) != 2:
|
|
259
|
+
raise ValueError(
|
|
260
|
+
"scatter 'data' must list exactly two datasets: [x-axis, y-axis]."
|
|
261
|
+
)
|
|
262
|
+
return v
|
|
263
|
+
|
|
264
|
+
|
|
265
|
+
class Scatter(Plot):
|
|
266
|
+
"""Plot class. Implements all particular operations for scatter plot creation.
|
|
267
|
+
|
|
268
|
+
Parameters
|
|
269
|
+
----------
|
|
270
|
+
Plot : _type_
|
|
271
|
+
_description_
|
|
272
|
+
"""
|
|
273
|
+
|
|
274
|
+
config = ScatterConfig
|
|
275
|
+
aliases = ["sc"]
|
|
276
|
+
|
|
277
|
+
def plot(self):
|
|
278
|
+
"""Iterate over self.plot_config.time entries, rendering once per entry."""
|
|
279
|
+
for time_interval in normalize_time(self.plot_config.time):
|
|
280
|
+
self._plot_one(time_interval)
|
|
281
|
+
|
|
282
|
+
def _plot_one(self, time_interval: str | None):
|
|
283
|
+
"""Scatter plotting method for a single time entry.
|
|
284
|
+
|
|
285
|
+
Pairs two datasets value-for-value: ``data[0]`` on x, ``data[1]`` on y.
|
|
286
|
+
Per domain, each dataset gets the domain applied as a RAW Data object (a
|
|
287
|
+
resample_to domain co-locates them here), is time-filtered, then reduced
|
|
288
|
+
over every dim except ``dimension`` so the two align for pairing.
|
|
289
|
+
"""
|
|
290
|
+
vars = self.plot_config.vars
|
|
291
|
+
timestep = self.plot_config.timestep
|
|
292
|
+
dimension = self.plot_config.dimension
|
|
293
|
+
domains = self.resolve_domains()
|
|
294
|
+
x_name, y_name = self.plot_config.data
|
|
295
|
+
|
|
296
|
+
for dom_name, dom in domains.items():
|
|
297
|
+
prepared = {}
|
|
298
|
+
for ds_name in (x_name, y_name):
|
|
299
|
+
d = self.data[ds_name]
|
|
300
|
+
d = dom.apply(d) if dom is not None else d
|
|
301
|
+
obj = time_resampling(
|
|
302
|
+
d.obj, timestep=timestep, time_interval=time_interval
|
|
303
|
+
)
|
|
304
|
+
prepared[ds_name] = (d, obj)
|
|
305
|
+
|
|
306
|
+
for variable in vars:
|
|
307
|
+
unit = vars[variable] if isinstance(vars, dict) else None
|
|
308
|
+
reduced = {}
|
|
309
|
+
for ds_name, (d, obj) in prepared.items():
|
|
310
|
+
da = obj[variable]
|
|
311
|
+
if self.plot_config.dim_reduce:
|
|
312
|
+
da = dim_reduction(da, self.plot_config.dim_reduce)
|
|
313
|
+
da = da.reduce(
|
|
314
|
+
self.plot_config.reduction_method.func,
|
|
315
|
+
list(set(da.dims) - {dimension}),
|
|
316
|
+
)
|
|
317
|
+
da = change_unit(da, d.var_unit(variable), unit)
|
|
318
|
+
reduced[ds_name] = da
|
|
319
|
+
|
|
320
|
+
x_var, y_var = reduced[x_name], reduced[y_name]
|
|
321
|
+
# Scatter pairs the two datasets element-wise along `dimension`.
|
|
322
|
+
# resample_vars is now purely spatial (it no longer snaps time), so
|
|
323
|
+
# align the two on the shared axis here — the correct home for the
|
|
324
|
+
# time-alignment concern. inner-join keeps only matching coords.
|
|
325
|
+
if dimension in x_var.dims and dimension in y_var.dims:
|
|
326
|
+
x_var, y_var = xr.align(x_var, y_var, join="inner")
|
|
327
|
+
|
|
328
|
+
figure = plt.figure(**self.figure_kwargs())
|
|
329
|
+
ax = figure.add_subplot(1, 1, 1)
|
|
330
|
+
min_val, max_val = (
|
|
331
|
+
math.floor(np.nanmin([np.nanmin(x_var), np.nanmin(y_var)])),
|
|
332
|
+
math.ceil(np.nanmax([np.nanmax(x_var), np.nanmax(y_var)])),
|
|
333
|
+
)
|
|
334
|
+
|
|
335
|
+
title = self.plot_kwargs.get(
|
|
336
|
+
"title",
|
|
337
|
+
f"Scatter comparison of {variable} between {x_name} and {y_name}",
|
|
338
|
+
)
|
|
339
|
+
xlabel = self.plot_kwargs.get(
|
|
340
|
+
"xlabel", f"{_unit_label(variable, unit)}, {x_name}"
|
|
341
|
+
)
|
|
342
|
+
ylabel = self.plot_kwargs.get(
|
|
343
|
+
"ylabel", f"{_unit_label(variable, unit)}, {y_name}"
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
ax.set_xlabel(xlabel)
|
|
347
|
+
ax.set_ylabel(ylabel)
|
|
348
|
+
|
|
349
|
+
figure.suptitle(title)
|
|
350
|
+
ax.set_xlim(min_val, max_val)
|
|
351
|
+
ax.set_ylim(min_val, max_val)
|
|
352
|
+
|
|
353
|
+
ax.scatter(x_var.values, y_var.values)
|
|
354
|
+
|
|
355
|
+
x = [min_val + x * (max_val - min_val) / 5 for x in range(5 + 1)]
|
|
356
|
+
ax.plot(x, x)
|
|
357
|
+
|
|
358
|
+
start, end = manage_time_interval(time_interval)
|
|
359
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
360
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
361
|
+
format = self.plot_kwargs.get("format", "jpg")
|
|
362
|
+
filename = (
|
|
363
|
+
f"scatter-{dom_name}-{variable}-{start}_{end}.{format}"
|
|
364
|
+
if self.plot_config.filename is None
|
|
365
|
+
else self.plot_config.filename
|
|
366
|
+
)
|
|
367
|
+
self._finalize(figure, filename)
|
|
368
|
+
|
|
369
|
+
|
|
370
|
+
class SpatialOverlayConfig(BasePlotConfig):
|
|
371
|
+
"""SpatialOverlay Spatial-Overlay plot configuration Pydantic model."""
|
|
372
|
+
|
|
373
|
+
type: Literal["spatial-overlay", "spatialoverlay", "so"]
|
|
374
|
+
base: str
|
|
375
|
+
superposed: str
|
|
376
|
+
time: str | list[str] | None = Field(default=None)
|
|
377
|
+
levels: int = Field(default=10)
|
|
378
|
+
reduction_method: ReductionMethodEnum = Field(default=ReductionMethodEnum.mean)
|
|
379
|
+
crs: CRSEnum | None = Field(default=None)
|
|
380
|
+
coastlines: bool = Field(default=True)
|
|
381
|
+
borders: bool = Field(default=True)
|
|
382
|
+
cmap: str = Field(default="viridis")
|
|
383
|
+
bbox: list[float | int] = Field(default=None)
|
|
384
|
+
padding: float = Field(default=0.05)
|
|
385
|
+
drop_nans: bool = Field(default=False)
|
|
386
|
+
|
|
387
|
+
|
|
388
|
+
class SpatialOverlay(Plot):
|
|
389
|
+
"""Plot class. Implements all particular operations for spatial-overlay plot creation."""
|
|
390
|
+
|
|
391
|
+
aliases = ["spatial-overlay", "spatialoverlay", "so"]
|
|
392
|
+
config = SpatialOverlayConfig
|
|
393
|
+
|
|
394
|
+
def plot(self):
|
|
395
|
+
"""Iterate over self.plot_config.time entries, rendering once per entry."""
|
|
396
|
+
for time_interval in normalize_time(self.plot_config.time):
|
|
397
|
+
self._plot_one(time_interval)
|
|
398
|
+
|
|
399
|
+
def _plot_one(self, time_interval: str | None):
|
|
400
|
+
"""Spatial Overlay plotting method for a single time entry.
|
|
401
|
+
The process goes as follows:
|
|
402
|
+
1) Process arguments.
|
|
403
|
+
2) Iterate through Domains.
|
|
404
|
+
2.1) Apply domains to data objects.
|
|
405
|
+
2.2) Time alignment.
|
|
406
|
+
2.3) Reduce variable to latitude and longitude.
|
|
407
|
+
2.4) Iterate through Variables
|
|
408
|
+
2.4.1) Contourf spatially distributed data variable plot from "base".
|
|
409
|
+
2.4.2) Scatter point surface data variable plot from "superposed".
|
|
410
|
+
2.4.3) Save figure.
|
|
411
|
+
"""
|
|
412
|
+
# Get relevant data from the config
|
|
413
|
+
base: RegularGrid = self.data[self.plot_config.base]
|
|
414
|
+
superposed: PointSurface = self.data[self.plot_config.superposed]
|
|
415
|
+
vars = self.plot_config.vars
|
|
416
|
+
crs = (
|
|
417
|
+
base.crs.crs if self.plot_config.crs is None else self.plot_config.crs.crs
|
|
418
|
+
) # TODO: make this simpler haha
|
|
419
|
+
domains = {
|
|
420
|
+
name: dom
|
|
421
|
+
for name, dom in self.domains.items()
|
|
422
|
+
if name in self.plot_config.domains
|
|
423
|
+
}
|
|
424
|
+
if not domains:
|
|
425
|
+
domains = {"": None}
|
|
426
|
+
|
|
427
|
+
for dom_name, dom in domains.items():
|
|
428
|
+
base_dom = dom.apply(base) if dom is not None else base
|
|
429
|
+
superposed_dom = dom.apply(superposed) if dom is not None else superposed
|
|
430
|
+
for var in vars:
|
|
431
|
+
base_var = base_dom.obj[var]
|
|
432
|
+
superposed_var = superposed_dom.obj[var]
|
|
433
|
+
|
|
434
|
+
unit = base.var_unit(var) if isinstance(vars, list | str) else vars[var]
|
|
435
|
+
|
|
436
|
+
# Time alignment
|
|
437
|
+
base_var = time_resampling(base_var, time_interval=time_interval)
|
|
438
|
+
superposed_var = time_resampling(
|
|
439
|
+
superposed_var, time_interval=time_interval
|
|
440
|
+
)
|
|
441
|
+
|
|
442
|
+
if self.plot_config.dim_reduce:
|
|
443
|
+
base_var = dim_reduction(base_var, self.plot_config.dim_reduce)
|
|
444
|
+
superposed_var = dim_reduction(
|
|
445
|
+
superposed_var, self.plot_config.dim_reduce
|
|
446
|
+
)
|
|
447
|
+
|
|
448
|
+
base_reduction_dims = [x for x in ["time", "z"] if x in base_var.dims]
|
|
449
|
+
base_var = base_var.reduce(
|
|
450
|
+
self.plot_config.reduction_method.func, base_reduction_dims
|
|
451
|
+
)
|
|
452
|
+
superposed_var = superposed_var.reduce(
|
|
453
|
+
self.plot_config.reduction_method.func, "time"
|
|
454
|
+
)
|
|
455
|
+
|
|
456
|
+
# unit conversion
|
|
457
|
+
base_var = change_unit(base_var, base.var_unit(var), unit)
|
|
458
|
+
superposed_var = change_unit(
|
|
459
|
+
superposed_var, superposed.var_unit(var), unit
|
|
460
|
+
)
|
|
461
|
+
|
|
462
|
+
# Plotting
|
|
463
|
+
figure = plt.figure(**self.figure_kwargs())
|
|
464
|
+
|
|
465
|
+
ax = figure.add_subplot(1, 1, 1, projection=crs())
|
|
466
|
+
|
|
467
|
+
if self.plot_config.coastlines:
|
|
468
|
+
ax.coastlines()
|
|
469
|
+
if self.plot_config.borders:
|
|
470
|
+
ax.add_feature(cfeature.BORDERS)
|
|
471
|
+
|
|
472
|
+
figure.suptitle(
|
|
473
|
+
f"Spatial overlay comparison of {var} between {base.name} and {superposed.name}"
|
|
474
|
+
)
|
|
475
|
+
|
|
476
|
+
vmin, vmax = (
|
|
477
|
+
np.nanmin([base_var.min(), superposed_var.min()]),
|
|
478
|
+
np.nanmax([base_var.max(), superposed_var.max()]),
|
|
479
|
+
)
|
|
480
|
+
|
|
481
|
+
# Superposed (point) coords/values, optionally dropping the
|
|
482
|
+
# stations whose reduced observation is NaN so empty sites
|
|
483
|
+
# neither draw nor stretch the auto extent.
|
|
484
|
+
sup_lons = superposed_var["longitude"].values
|
|
485
|
+
sup_lats = superposed_var["latitude"].values
|
|
486
|
+
sup_vals = superposed_var.values
|
|
487
|
+
if self.plot_config.drop_nans:
|
|
488
|
+
sup_lons, sup_lats, sup_vals = _drop_nan_points(
|
|
489
|
+
sup_lons, sup_lats, sup_vals
|
|
490
|
+
)
|
|
491
|
+
|
|
492
|
+
if self.plot_config.bbox is not None:
|
|
493
|
+
lon_min, lat_min, lon_max, lat_max = self.plot_config.bbox
|
|
494
|
+
else:
|
|
495
|
+
lon_min = np.nanmax(
|
|
496
|
+
[np.nanmin(base_var["longitude"]), np.nanmin(sup_lons)]
|
|
497
|
+
)
|
|
498
|
+
lat_min = np.nanmax(
|
|
499
|
+
[np.nanmin(base_var["latitude"]), np.nanmin(sup_lats)]
|
|
500
|
+
)
|
|
501
|
+
lon_max = np.nanmin(
|
|
502
|
+
[np.nanmax(base_var["longitude"]), np.nanmax(sup_lons)]
|
|
503
|
+
)
|
|
504
|
+
lat_max = np.nanmin(
|
|
505
|
+
[np.nanmax(base_var["latitude"]), np.nanmax(sup_lats)]
|
|
506
|
+
)
|
|
507
|
+
lon_min, lon_max, lat_min, lat_max = _pad_extent(
|
|
508
|
+
lon_min, lon_max, lat_min, lat_max, self.plot_config.padding
|
|
509
|
+
)
|
|
510
|
+
|
|
511
|
+
norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
|
|
512
|
+
|
|
513
|
+
ax.contourf(
|
|
514
|
+
base_var["longitude"].values,
|
|
515
|
+
base_var["latitude"].values,
|
|
516
|
+
base_var.values,
|
|
517
|
+
transform=base.crs.crs(),
|
|
518
|
+
cmap=self.plot_config.cmap,
|
|
519
|
+
norm=norm,
|
|
520
|
+
levels=self.plot_config.levels,
|
|
521
|
+
)
|
|
522
|
+
|
|
523
|
+
ax.scatter(
|
|
524
|
+
sup_lons,
|
|
525
|
+
sup_lats,
|
|
526
|
+
c=sup_vals,
|
|
527
|
+
transform=superposed.crs.crs(),
|
|
528
|
+
cmap=self.plot_config.cmap,
|
|
529
|
+
norm=norm,
|
|
530
|
+
edgecolor="k",
|
|
531
|
+
)
|
|
532
|
+
|
|
533
|
+
ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=crs())
|
|
534
|
+
|
|
535
|
+
sm = mpl.cm.ScalarMappable(norm=norm, cmap=self.plot_config.cmap)
|
|
536
|
+
|
|
537
|
+
figure.colorbar(
|
|
538
|
+
sm, ax=ax, orientation="vertical", label=_unit_label(var, unit)
|
|
539
|
+
)
|
|
540
|
+
|
|
541
|
+
start, end = manage_time_interval(time_interval)
|
|
542
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
543
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
544
|
+
format = self.plot_kwargs.get("format", "jpg")
|
|
545
|
+
filename = (
|
|
546
|
+
f"spatial_overlay-{dom_name}-{var}-{base.name}-{superposed.name}-{start}_{end}.{format}"
|
|
547
|
+
if self.plot_config.filename is None
|
|
548
|
+
else self.plot_config.filename
|
|
549
|
+
)
|
|
550
|
+
|
|
551
|
+
self._finalize(figure, filename)
|
|
552
|
+
|
|
553
|
+
|
|
554
|
+
class SpatialMapConfig(BasePlotConfig):
|
|
555
|
+
"""SpatialMap single-dataset map plot configuration Pydantic model."""
|
|
556
|
+
|
|
557
|
+
type: Literal["spatial-map", "spatialmap", "map", "sm"]
|
|
558
|
+
data: str
|
|
559
|
+
time: str | list[str] | None = Field(default=None)
|
|
560
|
+
levels: int = Field(default=10)
|
|
561
|
+
reduction_method: ReductionMethodEnum = Field(default=ReductionMethodEnum.mean)
|
|
562
|
+
crs: CRSEnum | None = Field(default=None)
|
|
563
|
+
coastlines: bool = Field(default=True)
|
|
564
|
+
borders: bool = Field(default=True)
|
|
565
|
+
cmap: str = Field(default="viridis")
|
|
566
|
+
bbox: list[float | int] | None = Field(default=None)
|
|
567
|
+
markersize: float = Field(default=40.0)
|
|
568
|
+
padding: float = Field(default=0.05)
|
|
569
|
+
drop_nans: bool = Field(default=False)
|
|
570
|
+
|
|
571
|
+
|
|
572
|
+
class SpatialMap(Plot):
|
|
573
|
+
"""Plot class. Plots a single dataset over a map.
|
|
574
|
+
|
|
575
|
+
The rendering style is chosen from the dataset topology: spatially
|
|
576
|
+
distributed data (``RegularGrid``) is drawn as a filled contour
|
|
577
|
+
(``contourf``); in-situ data (``PointSurface`` and any other topology)
|
|
578
|
+
is drawn as a coloured ``scatter`` of points. It is, in essence, one
|
|
579
|
+
half of ``SpatialOverlay`` applied to a single dataset.
|
|
580
|
+
"""
|
|
581
|
+
|
|
582
|
+
aliases = ["spatial-map", "spatialmap", "map", "sm"]
|
|
583
|
+
config = SpatialMapConfig
|
|
584
|
+
|
|
585
|
+
def plot(self):
|
|
586
|
+
"""Iterate over self.plot_config.time entries, rendering once per entry."""
|
|
587
|
+
for time_interval in normalize_time(self.plot_config.time):
|
|
588
|
+
self._plot_one(time_interval)
|
|
589
|
+
|
|
590
|
+
def _plot_one(self, time_interval: str | None):
|
|
591
|
+
"""Spatial Map plotting method for a single time entry.
|
|
592
|
+
The process goes as follows:
|
|
593
|
+
1) Process arguments.
|
|
594
|
+
2) Iterate through Domains.
|
|
595
|
+
2.1) Apply domain to the data object.
|
|
596
|
+
2.2) Time alignment.
|
|
597
|
+
2.3) Reduce variable to latitude and longitude.
|
|
598
|
+
2.4) Iterate through Variables
|
|
599
|
+
2.4.1) Contourf for RegularGrid data, scatter otherwise.
|
|
600
|
+
2.4.2) Save figure.
|
|
601
|
+
"""
|
|
602
|
+
# Get relevant data from the config
|
|
603
|
+
data: RegularGrid | PointSurface = self.data[self.plot_config.data]
|
|
604
|
+
vars = self.plot_config.vars
|
|
605
|
+
crs = data.crs.crs if self.plot_config.crs is None else self.plot_config.crs.crs
|
|
606
|
+
domains = {
|
|
607
|
+
name: dom
|
|
608
|
+
for name, dom in self.domains.items()
|
|
609
|
+
if name in self.plot_config.domains
|
|
610
|
+
}
|
|
611
|
+
if not domains:
|
|
612
|
+
domains = {"": None}
|
|
613
|
+
|
|
614
|
+
# RegularGrid renders as a filled contour, everything else as points.
|
|
615
|
+
is_grid = isinstance(data, RegularGrid)
|
|
616
|
+
|
|
617
|
+
for dom_name, dom in domains.items():
|
|
618
|
+
data_dom = dom.apply(data) if dom is not None else data
|
|
619
|
+
for var in vars:
|
|
620
|
+
data_var = data_dom.obj[var]
|
|
621
|
+
|
|
622
|
+
unit = data.var_unit(var) if isinstance(vars, list | str) else vars[var]
|
|
623
|
+
|
|
624
|
+
# Time alignment
|
|
625
|
+
data_var = time_resampling(data_var, time_interval=time_interval)
|
|
626
|
+
|
|
627
|
+
if self.plot_config.dim_reduce:
|
|
628
|
+
data_var = dim_reduction(data_var, self.plot_config.dim_reduce)
|
|
629
|
+
|
|
630
|
+
reduction_dims = [x for x in ["time", "z"] if x in data_var.dims]
|
|
631
|
+
data_var = data_var.reduce(
|
|
632
|
+
self.plot_config.reduction_method.func, reduction_dims
|
|
633
|
+
)
|
|
634
|
+
|
|
635
|
+
# Unit conversion
|
|
636
|
+
data_var = change_unit(data_var, data.var_unit(var), unit)
|
|
637
|
+
|
|
638
|
+
# Plotting
|
|
639
|
+
figure = plt.figure(**self.figure_kwargs())
|
|
640
|
+
|
|
641
|
+
ax = figure.add_subplot(1, 1, 1, projection=crs())
|
|
642
|
+
|
|
643
|
+
if self.plot_config.coastlines:
|
|
644
|
+
ax.coastlines()
|
|
645
|
+
if self.plot_config.borders:
|
|
646
|
+
ax.add_feature(cfeature.BORDERS)
|
|
647
|
+
|
|
648
|
+
figure.suptitle(
|
|
649
|
+
self.plot_kwargs.get(
|
|
650
|
+
"title", f"Spatial map of {var} for {data.name}"
|
|
651
|
+
)
|
|
652
|
+
)
|
|
653
|
+
|
|
654
|
+
vmin, vmax = np.nanmin(data_var), np.nanmax(data_var)
|
|
655
|
+
norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
|
|
656
|
+
|
|
657
|
+
lons = data_var["longitude"].values
|
|
658
|
+
lats = data_var["latitude"].values
|
|
659
|
+
|
|
660
|
+
if is_grid:
|
|
661
|
+
ax.contourf(
|
|
662
|
+
lons,
|
|
663
|
+
lats,
|
|
664
|
+
data_var.values,
|
|
665
|
+
transform=data.crs.crs(),
|
|
666
|
+
cmap=self.plot_config.cmap,
|
|
667
|
+
norm=norm,
|
|
668
|
+
levels=self.plot_config.levels,
|
|
669
|
+
)
|
|
670
|
+
else:
|
|
671
|
+
vals = data_var.values
|
|
672
|
+
# Drop stations whose reduced observation is NaN so empty
|
|
673
|
+
# sites neither draw nor stretch the auto extent.
|
|
674
|
+
if self.plot_config.drop_nans:
|
|
675
|
+
lons, lats, vals = _drop_nan_points(lons, lats, vals)
|
|
676
|
+
ax.scatter(
|
|
677
|
+
lons,
|
|
678
|
+
lats,
|
|
679
|
+
c=vals,
|
|
680
|
+
s=self.plot_config.markersize,
|
|
681
|
+
transform=data.crs.crs(),
|
|
682
|
+
cmap=self.plot_config.cmap,
|
|
683
|
+
norm=norm,
|
|
684
|
+
edgecolor="k",
|
|
685
|
+
)
|
|
686
|
+
|
|
687
|
+
if self.plot_config.bbox is not None:
|
|
688
|
+
lon_min, lat_min, lon_max, lat_max = self.plot_config.bbox
|
|
689
|
+
else:
|
|
690
|
+
lon_min, lon_max = np.nanmin(lons), np.nanmax(lons)
|
|
691
|
+
lat_min, lat_max = np.nanmin(lats), np.nanmax(lats)
|
|
692
|
+
lon_min, lon_max, lat_min, lat_max = _pad_extent(
|
|
693
|
+
lon_min, lon_max, lat_min, lat_max, self.plot_config.padding
|
|
694
|
+
)
|
|
695
|
+
|
|
696
|
+
ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=crs())
|
|
697
|
+
|
|
698
|
+
sm = mpl.cm.ScalarMappable(norm=norm, cmap=self.plot_config.cmap)
|
|
699
|
+
figure.colorbar(
|
|
700
|
+
sm, ax=ax, orientation="vertical", label=_unit_label(var, unit)
|
|
701
|
+
)
|
|
702
|
+
|
|
703
|
+
start, end = manage_time_interval(time_interval)
|
|
704
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
705
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
706
|
+
format = self.plot_kwargs.get("format", "jpg")
|
|
707
|
+
filename = (
|
|
708
|
+
f"spatial_map-{dom_name}-{var}-{data.name}-{start}_{end}.{format}"
|
|
709
|
+
if self.plot_config.filename is None
|
|
710
|
+
else self.plot_config.filename
|
|
711
|
+
)
|
|
712
|
+
|
|
713
|
+
self._finalize(figure, filename)
|
|
714
|
+
|
|
715
|
+
|
|
716
|
+
class TimeCycleConfig(BasePlotConfig):
|
|
717
|
+
"""TimeCycle plot configuration as Pydantic Model."""
|
|
718
|
+
|
|
719
|
+
type: Literal["timecycle", "time cycle", "cycle"]
|
|
720
|
+
# One or more datasets; the first is the reference and gets the ± std band.
|
|
721
|
+
data: str | list[str]
|
|
722
|
+
time: str | list[str] | None = Field(default=None)
|
|
723
|
+
timestep: TimestepEnum | None = Field(default=None)
|
|
724
|
+
time_buckets: TimeBucketEnum = Field(default=TimeBucketEnum.day)
|
|
725
|
+
reduction_method: ReductionMethodEnum = Field(default=ReductionMethodEnum.mean)
|
|
726
|
+
|
|
727
|
+
@field_validator("data", mode="before")
|
|
728
|
+
@classmethod
|
|
729
|
+
def _wrap_single_data(cls, v):
|
|
730
|
+
"""Accept a single dataset name as well as a list."""
|
|
731
|
+
return [v] if isinstance(v, str) else v
|
|
732
|
+
|
|
733
|
+
@model_validator(mode="after")
|
|
734
|
+
def check_timestep_vs_bucket(self) -> "TimeCycleConfig":
|
|
735
|
+
"""Reject a timestep coarser than the bucket it is grouped into."""
|
|
736
|
+
if self.timestep is None:
|
|
737
|
+
return self
|
|
738
|
+
# map both to hours for comparison
|
|
739
|
+
bucket_hours = {
|
|
740
|
+
"hour": 1,
|
|
741
|
+
"day": 24,
|
|
742
|
+
"month": 720,
|
|
743
|
+
"dayofyear": 24,
|
|
744
|
+
"season": 2160,
|
|
745
|
+
}
|
|
746
|
+
timestep_hours = {
|
|
747
|
+
TimestepEnum.hourly: 1,
|
|
748
|
+
TimestepEnum.daily: 24,
|
|
749
|
+
TimestepEnum.monthly: 720,
|
|
750
|
+
# extend as needed
|
|
751
|
+
}
|
|
752
|
+
if timestep_hours[self.timestep] > bucket_hours[self.time_buckets.value]:
|
|
753
|
+
raise ValueError(
|
|
754
|
+
f"timestep '{self.timestep}' is coarser than "
|
|
755
|
+
f"time_bucket '{self.time_buckets.value}' — std bands will be meaningless"
|
|
756
|
+
)
|
|
757
|
+
return self
|
|
758
|
+
|
|
759
|
+
|
|
760
|
+
class TimeCycle(Plot):
|
|
761
|
+
"""Average cycle over a time bucket.
|
|
762
|
+
|
|
763
|
+
Groups the series by the configured bucket (hour, day, season, ...) and draws
|
|
764
|
+
the mean of each group. The first dataset is the reference and also gets a
|
|
765
|
+
standard-deviation band.
|
|
766
|
+
"""
|
|
767
|
+
|
|
768
|
+
config = TimeCycleConfig
|
|
769
|
+
aliases = ["time cycle", "cycle"]
|
|
770
|
+
|
|
771
|
+
def plot(self):
|
|
772
|
+
"""Iterate over self.plot_config.time entries, rendering once per entry."""
|
|
773
|
+
for time_interval in normalize_time(self.plot_config.time):
|
|
774
|
+
self._plot_one(time_interval)
|
|
775
|
+
|
|
776
|
+
def _plot_one(self, time_interval: str | None):
|
|
777
|
+
"""Render one figure for a single time entry."""
|
|
778
|
+
vars = self.plot_config.vars
|
|
779
|
+
domains = self.resolve_domains()
|
|
780
|
+
timestep = self.plot_config.timestep
|
|
781
|
+
time_bucket = (
|
|
782
|
+
self.plot_config.time_buckets.value
|
|
783
|
+
) # e.g. "hour", "month", "dayofyear"
|
|
784
|
+
data_names = self.plot_config.data
|
|
785
|
+
first = self.data[data_names[0]]
|
|
786
|
+
reference = data_names[0] # gets the ± std band
|
|
787
|
+
|
|
788
|
+
for dom_name, dom in domains.items():
|
|
789
|
+
# Domain applied to each raw Data (resample domains reproject here);
|
|
790
|
+
# optional timestep pre-aggregation is applied uniformly so each period
|
|
791
|
+
# contributes one value per bucket, then groupby builds the cycle.
|
|
792
|
+
prepared = {}
|
|
793
|
+
for ds_name in data_names:
|
|
794
|
+
d = self.data[ds_name]
|
|
795
|
+
d = dom.apply(d) if dom is not None else d
|
|
796
|
+
obj = time_resampling(
|
|
797
|
+
d.obj, timestep=timestep, time_interval=time_interval
|
|
798
|
+
)
|
|
799
|
+
prepared[ds_name] = (d, obj)
|
|
800
|
+
|
|
801
|
+
for variable in vars:
|
|
802
|
+
dst_unit = vars[variable] if isinstance(vars, dict) else None
|
|
803
|
+
label_unit = (
|
|
804
|
+
vars[variable]
|
|
805
|
+
if isinstance(vars, dict)
|
|
806
|
+
else first.var_unit(variable)
|
|
807
|
+
)
|
|
808
|
+
|
|
809
|
+
figure = plt.figure(**self.figure_kwargs(figsize=(8, 5)))
|
|
810
|
+
ax = figure.add_subplot(1, 1, 1)
|
|
811
|
+
|
|
812
|
+
xticklabels = None
|
|
813
|
+
|
|
814
|
+
for ds_name, (d, obj) in prepared.items():
|
|
815
|
+
da = obj[variable]
|
|
816
|
+
if self.plot_config.dim_reduce:
|
|
817
|
+
da = dim_reduction(da, self.plot_config.dim_reduce)
|
|
818
|
+
da = da.reduce(
|
|
819
|
+
self.plot_config.reduction_method.func,
|
|
820
|
+
list(set(da.dims) - {"time"}),
|
|
821
|
+
)
|
|
822
|
+
da = change_unit(da, d.var_unit(variable), dst_unit)
|
|
823
|
+
|
|
824
|
+
grouped = da.groupby(f"time.{time_bucket}")
|
|
825
|
+
mean = grouped.mean("time", skipna=True)
|
|
826
|
+
std = grouped.std("time", skipna=True)
|
|
827
|
+
|
|
828
|
+
bucket_vals = mean[time_bucket].values
|
|
829
|
+
xticklabels = (
|
|
830
|
+
xticklabels if xticklabels is not None else bucket_vals
|
|
831
|
+
)
|
|
832
|
+
|
|
833
|
+
ax.plot(bucket_vals, mean.values, label=ds_name)
|
|
834
|
+
# ± std band only for the reference (first) dataset.
|
|
835
|
+
if ds_name == reference:
|
|
836
|
+
ax.fill_between(
|
|
837
|
+
bucket_vals,
|
|
838
|
+
(mean - std).values,
|
|
839
|
+
(mean + std).values,
|
|
840
|
+
alpha=0.2,
|
|
841
|
+
)
|
|
842
|
+
|
|
843
|
+
title = self.plot_kwargs.get(
|
|
844
|
+
"title", f"{time_bucket.capitalize()} cycle of {variable}"
|
|
845
|
+
)
|
|
846
|
+
xlabel = self.plot_kwargs.get("xlabel", time_bucket.capitalize())
|
|
847
|
+
ylabel = self.plot_kwargs.get(
|
|
848
|
+
"ylabel", _unit_label(variable, label_unit)
|
|
849
|
+
)
|
|
850
|
+
|
|
851
|
+
ax.set_xlabel(xlabel)
|
|
852
|
+
ax.set_xticks(list(range(len(xticklabels))))
|
|
853
|
+
ax.set_xticklabels(xticklabels)
|
|
854
|
+
ax.set_ylabel(ylabel)
|
|
855
|
+
ax.legend()
|
|
856
|
+
figure.suptitle(title)
|
|
857
|
+
|
|
858
|
+
start, end = manage_time_interval(time_interval)
|
|
859
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
860
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
861
|
+
fmt = self.plot_kwargs.get("format", "jpg")
|
|
862
|
+
filename = (
|
|
863
|
+
f"cycle-{time_bucket}-{dom_name}-{variable}-{start}_{end}.{fmt}"
|
|
864
|
+
if self.plot_config.filename is None
|
|
865
|
+
else self.plot_config.filename
|
|
866
|
+
)
|
|
867
|
+
|
|
868
|
+
self._finalize(figure, filename)
|
|
869
|
+
|
|
870
|
+
|
|
871
|
+
class CustomPlotConfig(BasePlotConfig):
|
|
872
|
+
"""Free-form composition of primitives on shared axes.
|
|
873
|
+
|
|
874
|
+
Var-scope is mutually exclusive: EITHER ``vars`` is set at the plot level
|
|
875
|
+
(the whole composition re-renders once per var, every subplot sees that
|
|
876
|
+
var) OR each subplot declares its own ``var`` (a fixed composition — e.g. a
|
|
877
|
+
NO2 contour overlaid with a CO line). Never both.
|
|
878
|
+
"""
|
|
879
|
+
|
|
880
|
+
type: Literal["custom"]
|
|
881
|
+
# Optional here (unlike BasePlotConfig where vars is required): omitted means
|
|
882
|
+
# every subplot supplies its own var.
|
|
883
|
+
vars: list[str] | dict[str, str] | None = Field(default=None)
|
|
884
|
+
time: str | list[str] | None = Field(default=None)
|
|
885
|
+
timestep: TimestepEnum | None = Field(default=None)
|
|
886
|
+
subplots: list[PrimitiveModel] = Field(min_length=1)
|
|
887
|
+
|
|
888
|
+
@model_validator(mode="after")
|
|
889
|
+
def check_var_scope(self) -> "CustomPlotConfig":
|
|
890
|
+
"""Require plot-level ``vars`` or per-subplot ``var``, never both."""
|
|
891
|
+
plot_vars_set = self.vars is not None
|
|
892
|
+
subplot_vars = [getattr(sp, "var", None) for sp in self.subplots]
|
|
893
|
+
any_subplot_var = any(v is not None for v in subplot_vars)
|
|
894
|
+
all_subplot_vars = all(v is not None for v in subplot_vars)
|
|
895
|
+
if plot_vars_set and any_subplot_var:
|
|
896
|
+
raise ValueError(
|
|
897
|
+
"custom plot: set 'vars' at the plot level OR a 'var' on every "
|
|
898
|
+
"subplot, not both."
|
|
899
|
+
)
|
|
900
|
+
if not plot_vars_set and not all_subplot_vars:
|
|
901
|
+
raise ValueError(
|
|
902
|
+
"custom plot: when plot-level 'vars' is omitted, every subplot "
|
|
903
|
+
"must declare its own 'var'."
|
|
904
|
+
)
|
|
905
|
+
return self
|
|
906
|
+
|
|
907
|
+
|
|
908
|
+
class Custom(Plot):
|
|
909
|
+
"""Plot: overlay primitives on one shared set of axes.
|
|
910
|
+
|
|
911
|
+
The orchestrator owns all data preparation; primitives only draw. For each
|
|
912
|
+
(time, domain, var) context it prepares each subplot's data — resolve
|
|
913
|
+
dataset/var, apply domain, filter time, per-dim reduce (plot-level merged
|
|
914
|
+
with subplot-level, subplot wins), blanket-reduce the remaining non-axis
|
|
915
|
+
dims, convert units — then calls ``primitive.render`` and assembles the
|
|
916
|
+
figure-level legend/colorbar.
|
|
917
|
+
"""
|
|
918
|
+
|
|
919
|
+
config = CustomPlotConfig
|
|
920
|
+
aliases = ["custom"]
|
|
921
|
+
|
|
922
|
+
def plot(self):
|
|
923
|
+
"""Render one figure per (time, domain, var) context."""
|
|
924
|
+
domains = self.resolve_domains()
|
|
925
|
+
plot_vars = self.plot_config.vars
|
|
926
|
+
# dict form ({var: unit}) fans out over its keys; None → single None slot.
|
|
927
|
+
var_list = list(plot_vars) if plot_vars is not None else None
|
|
928
|
+
for time_interval, dom_name, dom, var in self.iterate_contexts(
|
|
929
|
+
self.plot_config.time, domains, var_list
|
|
930
|
+
):
|
|
931
|
+
self._render_composition(time_interval, dom_name, dom, var)
|
|
932
|
+
|
|
933
|
+
@staticmethod
|
|
934
|
+
def _resolve_keep_dims(data, x: str | None, y: str | None) -> set[str]:
|
|
935
|
+
"""Dims the x/y coordinate names depend on — the axes to preserve.
|
|
936
|
+
|
|
937
|
+
Names a *coordinate* (which may be N-D, e.g. RegularGrid ``longitude``
|
|
938
|
+
over ``(y, x)``), falling back to treating the name as a literal dim.
|
|
939
|
+
Everything not returned here gets reduced away before rendering.
|
|
940
|
+
"""
|
|
941
|
+
keep: set[str] = set()
|
|
942
|
+
for coord in (x, y):
|
|
943
|
+
if coord is None:
|
|
944
|
+
continue
|
|
945
|
+
if coord in data.coords:
|
|
946
|
+
keep.update(data[coord].dims)
|
|
947
|
+
elif coord in data.dims:
|
|
948
|
+
keep.add(coord)
|
|
949
|
+
return keep
|
|
950
|
+
|
|
951
|
+
def _prepare_subplot_data(self, subplot, dom, var, time_interval):
|
|
952
|
+
"""Full data-prep pipeline for one subplot; returns the drawable DataArray."""
|
|
953
|
+
data = self.data[subplot.dataset]
|
|
954
|
+
data = dom.apply(data) if dom is not None else data
|
|
955
|
+
da = data.obj[var]
|
|
956
|
+
|
|
957
|
+
# Per-subplot time wins over the plot-level iteration time (lets different
|
|
958
|
+
# windows overlay on one figure); falls back to the iteration otherwise.
|
|
959
|
+
effective_time = subplot.time if subplot.time is not None else time_interval
|
|
960
|
+
da = time_resampling(
|
|
961
|
+
da, timestep=self.plot_config.timestep, time_interval=effective_time
|
|
962
|
+
)
|
|
963
|
+
|
|
964
|
+
# Per-dim reductions: plot-level defaults merged with subplot overrides.
|
|
965
|
+
merged_dim_reduce = {}
|
|
966
|
+
if self.plot_config.dim_reduce:
|
|
967
|
+
merged_dim_reduce.update(self.plot_config.dim_reduce)
|
|
968
|
+
if subplot.dim_reduce:
|
|
969
|
+
merged_dim_reduce.update(subplot.dim_reduce)
|
|
970
|
+
if merged_dim_reduce:
|
|
971
|
+
da = dim_reduction(da, merged_dim_reduce, name=var)
|
|
972
|
+
|
|
973
|
+
# Keep the axis dims (resolved after dim_reduce so dim-dropping selections
|
|
974
|
+
# are reflected); blanket-reduce whatever is left.
|
|
975
|
+
keep = self._resolve_keep_dims(da, subplot.x, subplot.y)
|
|
976
|
+
reduction_dims = [d for d in da.dims if d not in keep]
|
|
977
|
+
if reduction_dims:
|
|
978
|
+
da = da.reduce(subplot.reduction_method.func, reduction_dims)
|
|
979
|
+
|
|
980
|
+
# Unit precedence: subplot.unit wins over a plot-level vars-dict unit.
|
|
981
|
+
plot_vars = self.plot_config.vars
|
|
982
|
+
dst_unit = subplot.unit
|
|
983
|
+
if dst_unit is None and isinstance(plot_vars, dict):
|
|
984
|
+
dst_unit = plot_vars.get(var)
|
|
985
|
+
da = change_unit(da, data.var_unit(var), dst_unit)
|
|
986
|
+
return da, dst_unit
|
|
987
|
+
|
|
988
|
+
def _render_composition(self, time_interval, dom_name, dom, plot_var):
|
|
989
|
+
"""Draw every primitive of the composition onto the shared axes."""
|
|
990
|
+
figure = plt.figure(**self.figure_kwargs())
|
|
991
|
+
ax = figure.add_subplot(1, 1, 1)
|
|
992
|
+
|
|
993
|
+
colorbar = None # (mappable, label) for the first colour primitive
|
|
994
|
+
has_legend = False
|
|
995
|
+
used_vars = []
|
|
996
|
+
|
|
997
|
+
for subplot in self.plot_config.subplots:
|
|
998
|
+
var = plot_var if plot_var is not None else subplot.var
|
|
999
|
+
used_vars.append(var)
|
|
1000
|
+
|
|
1001
|
+
da, dst_unit = self._prepare_subplot_data(subplot, dom, var, time_interval)
|
|
1002
|
+
|
|
1003
|
+
primitive = Primitive.get_primitive_class(subplot.type)(self.name, subplot)
|
|
1004
|
+
label = subplot.label if subplot.label is not None else var
|
|
1005
|
+
artist = primitive.render(ax, da, x=subplot.x, y=subplot.y, label=label)
|
|
1006
|
+
|
|
1007
|
+
if primitive.wants_colorbar and colorbar is None:
|
|
1008
|
+
colorbar = (artist, _unit_label(var, dst_unit))
|
|
1009
|
+
if primitive.wants_legend:
|
|
1010
|
+
has_legend = True
|
|
1011
|
+
|
|
1012
|
+
if colorbar is not None:
|
|
1013
|
+
mappable, clabel = colorbar
|
|
1014
|
+
figure.colorbar(mappable, ax=ax, orientation="vertical", label=clabel)
|
|
1015
|
+
if has_legend:
|
|
1016
|
+
ax.legend()
|
|
1017
|
+
|
|
1018
|
+
# Label the y axis with the variable, as every other plot type does.
|
|
1019
|
+
# Only when nothing drew a colorbar: for contourf/points the unit lives
|
|
1020
|
+
# on the colorbar and y is a real coordinate (z, latitude, ...), so
|
|
1021
|
+
# naming it after the variable would be wrong.
|
|
1022
|
+
if colorbar is None:
|
|
1023
|
+
ylabel = self.plot_kwargs.get("ylabel")
|
|
1024
|
+
if ylabel is None:
|
|
1025
|
+
names = "+".join(dict.fromkeys(str(v) for v in used_vars))
|
|
1026
|
+
ylabel = _unit_label(names, dst_unit)
|
|
1027
|
+
ax.set_ylabel(ylabel)
|
|
1028
|
+
|
|
1029
|
+
start, end = manage_time_interval(time_interval)
|
|
1030
|
+
start = "start" if start is None else start.strftime("%d-%m-%Y")
|
|
1031
|
+
end = "end" if end is None else end.strftime("%d-%m-%Y")
|
|
1032
|
+
fmt = self.plot_kwargs.get("format", "jpg")
|
|
1033
|
+
# var token: the single plot-level var, or the distinct subplot vars
|
|
1034
|
+
# (declaration order) joined with "+", e.g. "NO2+CO".
|
|
1035
|
+
var_token = (
|
|
1036
|
+
plot_var
|
|
1037
|
+
if plot_var is not None
|
|
1038
|
+
else "+".join(dict.fromkeys(str(v) for v in used_vars))
|
|
1039
|
+
)
|
|
1040
|
+
|
|
1041
|
+
# A title may reference the fan-out context it is rendered for, e.g.
|
|
1042
|
+
# "Temperatura en {domain}". Titles without placeholders pass through
|
|
1043
|
+
# untouched, and an unknown or malformed one is left as written rather
|
|
1044
|
+
# than aborting the render.
|
|
1045
|
+
title = self.plot_kwargs.get("title", self.name)
|
|
1046
|
+
with contextlib.suppress(KeyError, IndexError, ValueError):
|
|
1047
|
+
title = title.format(domain=dom_name, var=var_token)
|
|
1048
|
+
figure.suptitle(title)
|
|
1049
|
+
filename = (
|
|
1050
|
+
f"custom-{dom_name}-{var_token}-{start}_{end}.{fmt}"
|
|
1051
|
+
if self.plot_config.filename is None
|
|
1052
|
+
else self.plot_config.filename
|
|
1053
|
+
)
|
|
1054
|
+
self._finalize(figure, filename)
|