anemoi-datasets 0.5.15__py3-none-any.whl → 0.5.17__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.
- anemoi/datasets/__init__.py +4 -1
- anemoi/datasets/__main__.py +12 -2
- anemoi/datasets/_version.py +9 -4
- anemoi/datasets/commands/cleanup.py +17 -2
- anemoi/datasets/commands/compare.py +18 -2
- anemoi/datasets/commands/copy.py +196 -14
- anemoi/datasets/commands/create.py +50 -7
- anemoi/datasets/commands/finalise-additions.py +17 -2
- anemoi/datasets/commands/finalise.py +17 -2
- anemoi/datasets/commands/init-additions.py +17 -2
- anemoi/datasets/commands/init.py +16 -2
- anemoi/datasets/commands/inspect.py +283 -62
- anemoi/datasets/commands/load-additions.py +16 -2
- anemoi/datasets/commands/load.py +16 -2
- anemoi/datasets/commands/patch.py +17 -2
- anemoi/datasets/commands/publish.py +17 -2
- anemoi/datasets/commands/scan.py +31 -3
- anemoi/datasets/compute/recentre.py +47 -11
- anemoi/datasets/create/__init__.py +612 -85
- anemoi/datasets/create/check.py +142 -20
- anemoi/datasets/create/chunks.py +64 -4
- anemoi/datasets/create/config.py +185 -21
- anemoi/datasets/create/filter.py +50 -0
- anemoi/datasets/create/filters/__init__.py +33 -0
- anemoi/datasets/create/filters/empty.py +37 -0
- anemoi/datasets/create/filters/legacy.py +93 -0
- anemoi/datasets/create/filters/noop.py +37 -0
- anemoi/datasets/create/filters/orog_to_z.py +58 -0
- anemoi/datasets/create/{functions/filters → filters}/pressure_level_relative_humidity_to_specific_humidity.py +33 -10
- anemoi/datasets/create/{functions/filters → filters}/pressure_level_specific_humidity_to_relative_humidity.py +32 -8
- anemoi/datasets/create/filters/rename.py +205 -0
- anemoi/datasets/create/{functions/filters → filters}/rotate_winds.py +43 -28
- anemoi/datasets/create/{functions/filters → filters}/single_level_dewpoint_to_relative_humidity.py +32 -9
- anemoi/datasets/create/{functions/filters → filters}/single_level_relative_humidity_to_dewpoint.py +33 -9
- anemoi/datasets/create/{functions/filters → filters}/single_level_relative_humidity_to_specific_humidity.py +55 -7
- anemoi/datasets/create/{functions/filters → filters}/single_level_specific_humidity_to_relative_humidity.py +98 -37
- anemoi/datasets/create/filters/speeddir_to_uv.py +95 -0
- anemoi/datasets/create/{functions/filters → filters}/sum.py +24 -27
- anemoi/datasets/create/filters/transform.py +53 -0
- anemoi/datasets/create/{functions/filters → filters}/unrotate_winds.py +27 -18
- anemoi/datasets/create/filters/uv_to_speeddir.py +94 -0
- anemoi/datasets/create/{functions/filters → filters}/wz_to_w.py +51 -33
- anemoi/datasets/create/input/__init__.py +76 -5
- anemoi/datasets/create/input/action.py +149 -13
- anemoi/datasets/create/input/concat.py +81 -10
- anemoi/datasets/create/input/context.py +39 -4
- anemoi/datasets/create/input/data_sources.py +72 -6
- anemoi/datasets/create/input/empty.py +21 -3
- anemoi/datasets/create/input/filter.py +60 -12
- anemoi/datasets/create/input/function.py +154 -37
- anemoi/datasets/create/input/join.py +86 -14
- anemoi/datasets/create/input/misc.py +67 -17
- anemoi/datasets/create/input/pipe.py +33 -6
- anemoi/datasets/create/input/repeated_dates.py +189 -41
- anemoi/datasets/create/input/result.py +202 -87
- anemoi/datasets/create/input/step.py +119 -22
- anemoi/datasets/create/input/template.py +100 -13
- anemoi/datasets/create/input/trace.py +62 -7
- anemoi/datasets/create/patch.py +52 -4
- anemoi/datasets/create/persistent.py +134 -17
- anemoi/datasets/create/size.py +15 -1
- anemoi/datasets/create/source.py +51 -0
- anemoi/datasets/create/sources/__init__.py +36 -0
- anemoi/datasets/create/{functions/sources → sources}/accumulations.py +296 -30
- anemoi/datasets/create/{functions/sources → sources}/constants.py +27 -2
- anemoi/datasets/create/{functions/sources → sources}/eccc_fstd.py +7 -3
- anemoi/datasets/create/sources/empty.py +37 -0
- anemoi/datasets/create/{functions/sources → sources}/forcings.py +25 -1
- anemoi/datasets/create/sources/grib.py +297 -0
- anemoi/datasets/create/{functions/sources → sources}/hindcasts.py +38 -4
- anemoi/datasets/create/sources/legacy.py +93 -0
- anemoi/datasets/create/{functions/sources → sources}/mars.py +168 -20
- anemoi/datasets/create/sources/netcdf.py +42 -0
- anemoi/datasets/create/sources/opendap.py +43 -0
- anemoi/datasets/create/{functions/sources/__init__.py → sources/patterns.py} +35 -4
- anemoi/datasets/create/sources/recentre.py +150 -0
- anemoi/datasets/create/{functions/sources → sources}/source.py +27 -5
- anemoi/datasets/create/{functions/sources → sources}/tendencies.py +64 -7
- anemoi/datasets/create/sources/xarray.py +92 -0
- anemoi/datasets/create/sources/xarray_kerchunk.py +36 -0
- anemoi/datasets/create/sources/xarray_support/README.md +1 -0
- anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/__init__.py +109 -8
- anemoi/datasets/create/sources/xarray_support/coordinates.py +442 -0
- anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/field.py +94 -16
- anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/fieldlist.py +90 -25
- anemoi/datasets/create/sources/xarray_support/flavour.py +1036 -0
- anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/grid.py +92 -31
- anemoi/datasets/create/sources/xarray_support/metadata.py +395 -0
- anemoi/datasets/create/sources/xarray_support/patch.py +91 -0
- anemoi/datasets/create/sources/xarray_support/time.py +391 -0
- anemoi/datasets/create/sources/xarray_support/variable.py +331 -0
- anemoi/datasets/create/sources/xarray_zarr.py +41 -0
- anemoi/datasets/create/{functions/sources → sources}/zenodo.py +34 -5
- anemoi/datasets/create/statistics/__init__.py +233 -44
- anemoi/datasets/create/statistics/summary.py +52 -6
- anemoi/datasets/create/testing.py +76 -0
- anemoi/datasets/create/{functions/filters/noop.py → typing.py} +6 -3
- anemoi/datasets/create/utils.py +97 -6
- anemoi/datasets/create/writer.py +26 -4
- anemoi/datasets/create/zarr.py +170 -23
- anemoi/datasets/data/__init__.py +51 -4
- anemoi/datasets/data/complement.py +191 -40
- anemoi/datasets/data/concat.py +141 -16
- anemoi/datasets/data/dataset.py +552 -61
- anemoi/datasets/data/debug.py +197 -26
- anemoi/datasets/data/ensemble.py +93 -8
- anemoi/datasets/data/fill_missing.py +165 -18
- anemoi/datasets/data/forwards.py +428 -56
- anemoi/datasets/data/grids.py +323 -97
- anemoi/datasets/data/indexing.py +112 -19
- anemoi/datasets/data/interpolate.py +92 -12
- anemoi/datasets/data/join.py +158 -19
- anemoi/datasets/data/masked.py +129 -15
- anemoi/datasets/data/merge.py +137 -23
- anemoi/datasets/data/misc.py +172 -16
- anemoi/datasets/data/missing.py +233 -29
- anemoi/datasets/data/rescale.py +111 -10
- anemoi/datasets/data/select.py +168 -26
- anemoi/datasets/data/statistics.py +67 -6
- anemoi/datasets/data/stores.py +149 -64
- anemoi/datasets/data/subset.py +159 -25
- anemoi/datasets/data/unchecked.py +168 -57
- anemoi/datasets/data/xy.py +168 -25
- anemoi/datasets/dates/__init__.py +191 -16
- anemoi/datasets/dates/groups.py +189 -47
- anemoi/datasets/grids.py +270 -31
- anemoi/datasets/testing.py +28 -1
- {anemoi_datasets-0.5.15.dist-info → anemoi_datasets-0.5.17.dist-info}/METADATA +10 -7
- anemoi_datasets-0.5.17.dist-info/RECORD +137 -0
- {anemoi_datasets-0.5.15.dist-info → anemoi_datasets-0.5.17.dist-info}/WHEEL +1 -1
- {anemoi_datasets-0.5.15.dist-info → anemoi_datasets-0.5.17.dist-info/licenses}/LICENSE +1 -1
- anemoi/datasets/create/functions/__init__.py +0 -66
- anemoi/datasets/create/functions/filters/__init__.py +0 -9
- anemoi/datasets/create/functions/filters/empty.py +0 -17
- anemoi/datasets/create/functions/filters/orog_to_z.py +0 -58
- anemoi/datasets/create/functions/filters/rename.py +0 -79
- anemoi/datasets/create/functions/filters/speeddir_to_uv.py +0 -78
- anemoi/datasets/create/functions/filters/uv_to_speeddir.py +0 -56
- anemoi/datasets/create/functions/sources/empty.py +0 -15
- anemoi/datasets/create/functions/sources/grib.py +0 -150
- anemoi/datasets/create/functions/sources/netcdf.py +0 -15
- anemoi/datasets/create/functions/sources/opendap.py +0 -15
- anemoi/datasets/create/functions/sources/recentre.py +0 -60
- anemoi/datasets/create/functions/sources/xarray/coordinates.py +0 -255
- anemoi/datasets/create/functions/sources/xarray/flavour.py +0 -472
- anemoi/datasets/create/functions/sources/xarray/metadata.py +0 -148
- anemoi/datasets/create/functions/sources/xarray/patch.py +0 -44
- anemoi/datasets/create/functions/sources/xarray/time.py +0 -177
- anemoi/datasets/create/functions/sources/xarray/variable.py +0 -188
- anemoi/datasets/create/functions/sources/xarray_kerchunk.py +0 -42
- anemoi/datasets/create/functions/sources/xarray_zarr.py +0 -15
- anemoi/datasets/utils/fields.py +0 -47
- anemoi_datasets-0.5.15.dist-info/RECORD +0 -129
- {anemoi_datasets-0.5.15.dist-info → anemoi_datasets-0.5.17.dist-info}/entry_points.txt +0 -0
- {anemoi_datasets-0.5.15.dist-info → anemoi_datasets-0.5.17.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,1036 @@
|
|
|
1
|
+
# (C) Copyright 2024 Anemoi contributors.
|
|
2
|
+
#
|
|
3
|
+
# This software is licensed under the terms of the Apache Licence Version 2.0
|
|
4
|
+
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
|
|
5
|
+
#
|
|
6
|
+
# In applying this licence, ECMWF does not waive the privileges and immunities
|
|
7
|
+
# granted to it by virtue of its status as an intergovernmental organisation
|
|
8
|
+
# nor does it submit to any jurisdiction.
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
import logging
|
|
12
|
+
from abc import ABC
|
|
13
|
+
from abc import abstractmethod
|
|
14
|
+
from typing import Any
|
|
15
|
+
from typing import Dict
|
|
16
|
+
from typing import Hashable
|
|
17
|
+
from typing import Optional
|
|
18
|
+
from typing import Tuple
|
|
19
|
+
|
|
20
|
+
import xarray as xr
|
|
21
|
+
from anemoi.utils.config import DotDict
|
|
22
|
+
|
|
23
|
+
from .coordinates import Coordinate
|
|
24
|
+
from .coordinates import DateCoordinate
|
|
25
|
+
from .coordinates import EnsembleCoordinate
|
|
26
|
+
from .coordinates import LatitudeCoordinate
|
|
27
|
+
from .coordinates import LevelCoordinate
|
|
28
|
+
from .coordinates import LongitudeCoordinate
|
|
29
|
+
from .coordinates import ScalarCoordinate
|
|
30
|
+
from .coordinates import StepCoordinate
|
|
31
|
+
from .coordinates import TimeCoordinate
|
|
32
|
+
from .coordinates import UnsupportedCoordinate
|
|
33
|
+
from .coordinates import XCoordinate
|
|
34
|
+
from .coordinates import YCoordinate
|
|
35
|
+
from .coordinates import is_scalar
|
|
36
|
+
from .grid import Grid
|
|
37
|
+
from .grid import MeshedGrid
|
|
38
|
+
from .grid import MeshProjectionGrid
|
|
39
|
+
from .grid import UnstructuredGrid
|
|
40
|
+
from .grid import UnstructuredProjectionGrid
|
|
41
|
+
|
|
42
|
+
LOG = logging.getLogger(__name__)
|
|
43
|
+
|
|
44
|
+
# CoordinateAttributes = namedtuple("CoordinateAttributes", ["axis", "name", "long_name", "standard_name", "units"])
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class CoordinateAttributes(DotDict):
|
|
48
|
+
pass
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class CoordinateGuesser(ABC):
|
|
52
|
+
"""Class to guess the type of coordinates in a dataset."""
|
|
53
|
+
|
|
54
|
+
def __init__(self, ds: xr.Dataset) -> None:
|
|
55
|
+
"""Initializes the CoordinateGuesser.
|
|
56
|
+
|
|
57
|
+
Parameters
|
|
58
|
+
----------
|
|
59
|
+
ds : xr.Dataset
|
|
60
|
+
The dataset to guess coordinates from.
|
|
61
|
+
"""
|
|
62
|
+
self.ds = ds
|
|
63
|
+
self._coordinate_cache: Dict[Hashable, Coordinate] = {}
|
|
64
|
+
self._grid_cache: Dict[Hashable, Grid] = {}
|
|
65
|
+
|
|
66
|
+
@classmethod
|
|
67
|
+
def from_flavour(cls, ds: xr.Dataset, flavour: Optional[Dict[str, Any]]) -> "CoordinateGuesser":
|
|
68
|
+
"""Creates a CoordinateGuesser from a flavour.
|
|
69
|
+
|
|
70
|
+
Parameters
|
|
71
|
+
----------
|
|
72
|
+
ds : xr.Dataset
|
|
73
|
+
The dataset to guess coordinates from.
|
|
74
|
+
flavour : Optional[Dict[str, Any]]
|
|
75
|
+
The flavour to use for guessing.
|
|
76
|
+
|
|
77
|
+
Returns
|
|
78
|
+
-------
|
|
79
|
+
CoordinateGuesser
|
|
80
|
+
The created CoordinateGuesser.
|
|
81
|
+
"""
|
|
82
|
+
if flavour is None:
|
|
83
|
+
return DefaultCoordinateGuesser(ds)
|
|
84
|
+
else:
|
|
85
|
+
return FlavourCoordinateGuesser(ds, flavour)
|
|
86
|
+
|
|
87
|
+
def guess(self, c: xr.DataArray, coord: Hashable) -> Coordinate:
|
|
88
|
+
"""Guesses the type of a coordinate.
|
|
89
|
+
|
|
90
|
+
Parameters
|
|
91
|
+
----------
|
|
92
|
+
c : xr.DataArray
|
|
93
|
+
The coordinate to guess.
|
|
94
|
+
coord : Hashable
|
|
95
|
+
The name of the coordinate.
|
|
96
|
+
|
|
97
|
+
Returns
|
|
98
|
+
-------
|
|
99
|
+
Coordinate
|
|
100
|
+
The guessed coordinate type.
|
|
101
|
+
"""
|
|
102
|
+
if coord not in self._coordinate_cache:
|
|
103
|
+
self._coordinate_cache[coord] = self._guess(c, coord)
|
|
104
|
+
return self._coordinate_cache[coord]
|
|
105
|
+
|
|
106
|
+
def _guess(self, coordinate: xr.DataArray, coord: Hashable) -> Coordinate:
|
|
107
|
+
"""Internal method to guess the type of a coordinate.
|
|
108
|
+
|
|
109
|
+
Parameters
|
|
110
|
+
----------
|
|
111
|
+
coordinate : xr.DataArray
|
|
112
|
+
The coordinate to guess.
|
|
113
|
+
coord : Hashable
|
|
114
|
+
The name of the coordinate.
|
|
115
|
+
|
|
116
|
+
Returns
|
|
117
|
+
-------
|
|
118
|
+
Coordinate
|
|
119
|
+
The guessed coordinate type.
|
|
120
|
+
"""
|
|
121
|
+
name = coordinate.name
|
|
122
|
+
standard_name = getattr(coordinate, "standard_name", "").lower()
|
|
123
|
+
axis = getattr(coordinate, "axis", "")
|
|
124
|
+
long_name = getattr(coordinate, "long_name", "").lower()
|
|
125
|
+
units = getattr(coordinate, "units", "")
|
|
126
|
+
|
|
127
|
+
attributes = CoordinateAttributes(
|
|
128
|
+
axis=axis,
|
|
129
|
+
name=name,
|
|
130
|
+
long_name=long_name,
|
|
131
|
+
standard_name=standard_name,
|
|
132
|
+
units=units,
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
d: Optional[Coordinate] = None
|
|
136
|
+
|
|
137
|
+
d = self._is_longitude(coordinate, attributes)
|
|
138
|
+
if d is not None:
|
|
139
|
+
return d
|
|
140
|
+
|
|
141
|
+
d = self._is_latitude(coordinate, attributes)
|
|
142
|
+
if d is not None:
|
|
143
|
+
return d
|
|
144
|
+
|
|
145
|
+
d = self._is_x(coordinate, attributes)
|
|
146
|
+
if d is not None:
|
|
147
|
+
return d
|
|
148
|
+
|
|
149
|
+
d = self._is_y(coordinate, attributes)
|
|
150
|
+
if d is not None:
|
|
151
|
+
return d
|
|
152
|
+
|
|
153
|
+
d = self._is_time(coordinate, attributes)
|
|
154
|
+
if d is not None:
|
|
155
|
+
return d
|
|
156
|
+
|
|
157
|
+
d = self._is_step(coordinate, attributes)
|
|
158
|
+
if d is not None:
|
|
159
|
+
return d
|
|
160
|
+
|
|
161
|
+
d = self._is_date(coordinate, attributes)
|
|
162
|
+
if d is not None:
|
|
163
|
+
return d
|
|
164
|
+
|
|
165
|
+
d = self._is_level(coordinate, attributes)
|
|
166
|
+
if d is not None:
|
|
167
|
+
return d
|
|
168
|
+
|
|
169
|
+
d = self._is_number(coordinate, attributes)
|
|
170
|
+
if d is not None:
|
|
171
|
+
return d
|
|
172
|
+
|
|
173
|
+
if coordinate.shape in ((1,), tuple()):
|
|
174
|
+
return ScalarCoordinate(coordinate)
|
|
175
|
+
|
|
176
|
+
LOG.warning(
|
|
177
|
+
f"Coordinate {coord} not supported\n{axis=}, {name=},"
|
|
178
|
+
f" {long_name=}, {standard_name=}, units\n\n{coordinate}\n\n{type(coordinate.values)} {coordinate.shape}"
|
|
179
|
+
)
|
|
180
|
+
|
|
181
|
+
return UnsupportedCoordinate(coordinate)
|
|
182
|
+
|
|
183
|
+
def grid(self, coordinates: Any, variable: Any) -> Any:
|
|
184
|
+
"""Determines the grid type for the given coordinates and variable.
|
|
185
|
+
|
|
186
|
+
Parameters
|
|
187
|
+
----------
|
|
188
|
+
coordinates : Any
|
|
189
|
+
The coordinates to determine the grid from.
|
|
190
|
+
variable : Any
|
|
191
|
+
The variable to determine the grid from.
|
|
192
|
+
|
|
193
|
+
Returns
|
|
194
|
+
-------
|
|
195
|
+
Any
|
|
196
|
+
The determined grid type.
|
|
197
|
+
"""
|
|
198
|
+
lat = [c for c in coordinates if c.is_lat]
|
|
199
|
+
lon = [c for c in coordinates if c.is_lon]
|
|
200
|
+
|
|
201
|
+
if len(lat) == 1 and len(lon) == 1:
|
|
202
|
+
return self._lat_lon_provided(lat, lon, variable)
|
|
203
|
+
|
|
204
|
+
x = [c for c in coordinates if c.is_x]
|
|
205
|
+
y = [c for c in coordinates if c.is_y]
|
|
206
|
+
|
|
207
|
+
if len(x) == 1 and len(y) == 1:
|
|
208
|
+
return self._x_y_provided(x, y, variable)
|
|
209
|
+
|
|
210
|
+
raise NotImplementedError(f"Cannot establish grid {coordinates}")
|
|
211
|
+
|
|
212
|
+
def _check_dims(self, variable: Any, x_or_lon: Any, y_or_lat: Any) -> Tuple[Any, bool]:
|
|
213
|
+
"""Checks the dimensions of the variable against the coordinates.
|
|
214
|
+
|
|
215
|
+
Parameters
|
|
216
|
+
----------
|
|
217
|
+
variable : Any
|
|
218
|
+
The variable to check.
|
|
219
|
+
x_or_lon : Any
|
|
220
|
+
The x or longitude coordinate.
|
|
221
|
+
y_or_lat : Any
|
|
222
|
+
The y or latitude coordinate.
|
|
223
|
+
|
|
224
|
+
Returns
|
|
225
|
+
-------
|
|
226
|
+
Tuple[Any, bool]
|
|
227
|
+
The checked dimensions and a flag indicating if the grid is unstructured.
|
|
228
|
+
"""
|
|
229
|
+
x_dims = set(x_or_lon.variable.dims)
|
|
230
|
+
y_dims = set(y_or_lat.variable.dims)
|
|
231
|
+
variable_dims = set(variable.dims)
|
|
232
|
+
|
|
233
|
+
if not (x_dims <= variable_dims) or not (y_dims <= variable_dims):
|
|
234
|
+
raise ValueError(
|
|
235
|
+
f"Dimensions do not match {variable.name}{variable.dims} !="
|
|
236
|
+
f" {x_or_lon.name}{x_or_lon.variable.dims} and {y_or_lat.name}{y_or_lat.variable.dims}"
|
|
237
|
+
)
|
|
238
|
+
|
|
239
|
+
variable_dims = tuple(v for v in variable.dims if v in (x_dims | y_dims))
|
|
240
|
+
if x_dims == y_dims:
|
|
241
|
+
# It's unstructured
|
|
242
|
+
return variable_dims, True
|
|
243
|
+
|
|
244
|
+
if len(x_dims) == 1 and len(y_dims) == 1:
|
|
245
|
+
# It's a mesh
|
|
246
|
+
return variable_dims, False
|
|
247
|
+
|
|
248
|
+
raise ValueError(
|
|
249
|
+
f"Cannot establish grid for {variable.name}{variable.dims},"
|
|
250
|
+
f" {x_or_lon.name}{x_or_lon.variable.dims},"
|
|
251
|
+
f" {y_or_lat.name}{y_or_lat.variable.dims}"
|
|
252
|
+
)
|
|
253
|
+
|
|
254
|
+
def _lat_lon_provided(self, lat: Any, lon: Any, variable: Any) -> Any:
|
|
255
|
+
"""Determines the grid type when latitude and longitude are provided.
|
|
256
|
+
|
|
257
|
+
Parameters
|
|
258
|
+
----------
|
|
259
|
+
lat : Any
|
|
260
|
+
The latitude coordinate.
|
|
261
|
+
lon : Any
|
|
262
|
+
The longitude coordinate.
|
|
263
|
+
variable : Any
|
|
264
|
+
The variable to determine the grid from.
|
|
265
|
+
|
|
266
|
+
Returns
|
|
267
|
+
-------
|
|
268
|
+
Any
|
|
269
|
+
The determined grid type.
|
|
270
|
+
"""
|
|
271
|
+
lat = lat[0]
|
|
272
|
+
lon = lon[0]
|
|
273
|
+
|
|
274
|
+
dim_vars, unstructured = self._check_dims(variable, lon, lat)
|
|
275
|
+
|
|
276
|
+
if (lat.name, lon.name, dim_vars) in self._grid_cache:
|
|
277
|
+
return self._grid_cache[(lat.name, lon.name, dim_vars)]
|
|
278
|
+
|
|
279
|
+
grid: Grid = UnstructuredGrid(lat, lon, dim_vars) if unstructured else MeshedGrid(lat, lon, dim_vars)
|
|
280
|
+
|
|
281
|
+
self._grid_cache[(lat.name, lon.name, dim_vars)] = grid
|
|
282
|
+
|
|
283
|
+
return grid
|
|
284
|
+
|
|
285
|
+
def _x_y_provided(self, x: Any, y: Any, variable: Any) -> Any:
|
|
286
|
+
"""Determines the grid type when x and y coordinates are provided.
|
|
287
|
+
|
|
288
|
+
Parameters
|
|
289
|
+
----------
|
|
290
|
+
x : Any
|
|
291
|
+
The x coordinate.
|
|
292
|
+
y : Any
|
|
293
|
+
The y coordinate.
|
|
294
|
+
variable : Any
|
|
295
|
+
The variable to determine the grid from.
|
|
296
|
+
|
|
297
|
+
Returns
|
|
298
|
+
-------
|
|
299
|
+
Any
|
|
300
|
+
The determined grid type.
|
|
301
|
+
"""
|
|
302
|
+
x = x[0]
|
|
303
|
+
y = y[0]
|
|
304
|
+
|
|
305
|
+
dim_vars, unstructured = self._check_dims(variable, x, y)
|
|
306
|
+
|
|
307
|
+
if (x.name, y.name, dim_vars) in self._grid_cache:
|
|
308
|
+
return self._grid_cache[(x.name, y.name, dim_vars)]
|
|
309
|
+
|
|
310
|
+
grid_mapping = variable.attrs.get("grid_mapping", None)
|
|
311
|
+
if grid_mapping is not None:
|
|
312
|
+
print(f"grid_mapping: {grid_mapping}")
|
|
313
|
+
print(self.ds[grid_mapping])
|
|
314
|
+
|
|
315
|
+
if grid_mapping is None:
|
|
316
|
+
LOG.warning(f"No 'grid_mapping' attribute provided for '{variable.name}'")
|
|
317
|
+
LOG.warning("Trying to guess...")
|
|
318
|
+
|
|
319
|
+
PROBE = {
|
|
320
|
+
"prime_meridian_name",
|
|
321
|
+
"reference_ellipsoid_name",
|
|
322
|
+
"crs_wkt",
|
|
323
|
+
"horizontal_datum_name",
|
|
324
|
+
"semi_major_axis",
|
|
325
|
+
"spatial_ref",
|
|
326
|
+
"inverse_flattening",
|
|
327
|
+
"semi_minor_axis",
|
|
328
|
+
"geographic_crs_name",
|
|
329
|
+
"GeoTransform",
|
|
330
|
+
"grid_mapping_name",
|
|
331
|
+
"longitude_of_prime_meridian",
|
|
332
|
+
}
|
|
333
|
+
candidate = None
|
|
334
|
+
for v in self.ds.variables:
|
|
335
|
+
var = self.ds[v]
|
|
336
|
+
if not is_scalar(var):
|
|
337
|
+
continue
|
|
338
|
+
|
|
339
|
+
if PROBE.intersection(var.attrs.keys()):
|
|
340
|
+
if candidate:
|
|
341
|
+
raise ValueError(f"Multiple candidates for 'grid_mapping': {candidate} and {v}")
|
|
342
|
+
candidate = v
|
|
343
|
+
|
|
344
|
+
if candidate:
|
|
345
|
+
LOG.warning(f"Using '{candidate}' as 'grid_mapping'")
|
|
346
|
+
grid_mapping = candidate
|
|
347
|
+
else:
|
|
348
|
+
LOG.warning("Could not fine a candidate for 'grid_mapping'")
|
|
349
|
+
|
|
350
|
+
if grid_mapping is None:
|
|
351
|
+
if "crs" in self.ds[variable].attrs:
|
|
352
|
+
grid_mapping = self.ds[variable].attrs["crs"]
|
|
353
|
+
LOG.warning(f"Using CRS {grid_mapping} from variable '{variable.name}' attributes")
|
|
354
|
+
|
|
355
|
+
if grid_mapping is None:
|
|
356
|
+
if "crs" in self.ds.attrs:
|
|
357
|
+
grid_mapping = self.ds.attrs["crs"]
|
|
358
|
+
LOG.warning(f"Using CRS {grid_mapping} from global attributes")
|
|
359
|
+
|
|
360
|
+
grid: Optional[Grid] = None
|
|
361
|
+
if grid_mapping is not None:
|
|
362
|
+
|
|
363
|
+
grid_mapping = dict(self.ds[grid_mapping].attrs)
|
|
364
|
+
|
|
365
|
+
if unstructured:
|
|
366
|
+
grid = UnstructuredProjectionGrid(x, y, grid_mapping)
|
|
367
|
+
else:
|
|
368
|
+
grid = MeshProjectionGrid(x, y, grid_mapping)
|
|
369
|
+
|
|
370
|
+
if grid is not None:
|
|
371
|
+
self._grid_cache[(x.name, y.name, dim_vars)] = grid
|
|
372
|
+
return grid
|
|
373
|
+
|
|
374
|
+
LOG.error("Could not fine a candidate for 'grid_mapping'")
|
|
375
|
+
raise NotImplementedError(f"Unstructured grid {x.name} {y.name}")
|
|
376
|
+
|
|
377
|
+
@abstractmethod
|
|
378
|
+
def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
|
|
379
|
+
"""Checks if the coordinate is a longitude.
|
|
380
|
+
|
|
381
|
+
Parameters
|
|
382
|
+
----------
|
|
383
|
+
c : xr.DataArray
|
|
384
|
+
The coordinate to check.
|
|
385
|
+
attributes : CoordinateAttributes
|
|
386
|
+
The attributes of the coordinate.
|
|
387
|
+
|
|
388
|
+
Returns
|
|
389
|
+
-------
|
|
390
|
+
Optional[LongitudeCoordinate]
|
|
391
|
+
The LongitudeCoordinate if matched, else None.
|
|
392
|
+
"""
|
|
393
|
+
pass
|
|
394
|
+
|
|
395
|
+
@abstractmethod
|
|
396
|
+
def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
|
|
397
|
+
"""Checks if the coordinate is a latitude.
|
|
398
|
+
|
|
399
|
+
Parameters
|
|
400
|
+
----------
|
|
401
|
+
c : xr.DataArray
|
|
402
|
+
The coordinate to check.
|
|
403
|
+
attributes : CoordinateAttributes
|
|
404
|
+
The attributes of the coordinate.
|
|
405
|
+
|
|
406
|
+
Returns
|
|
407
|
+
-------
|
|
408
|
+
Optional[LatitudeCoordinate]
|
|
409
|
+
The LatitudeCoordinate if matched, else None.
|
|
410
|
+
"""
|
|
411
|
+
pass
|
|
412
|
+
|
|
413
|
+
@abstractmethod
|
|
414
|
+
def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
|
|
415
|
+
"""Checks if the coordinate is an x coordinate.
|
|
416
|
+
|
|
417
|
+
Parameters
|
|
418
|
+
----------
|
|
419
|
+
c : xr.DataArray
|
|
420
|
+
The coordinate to check.
|
|
421
|
+
attributes : CoordinateAttributes
|
|
422
|
+
The attributes of the coordinate.
|
|
423
|
+
|
|
424
|
+
Returns
|
|
425
|
+
-------
|
|
426
|
+
Optional[XCoordinate]
|
|
427
|
+
The XCoordinate if matched, else None.
|
|
428
|
+
"""
|
|
429
|
+
pass
|
|
430
|
+
|
|
431
|
+
@abstractmethod
|
|
432
|
+
def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
|
|
433
|
+
"""Checks if the coordinate is a y coordinate.
|
|
434
|
+
|
|
435
|
+
Parameters
|
|
436
|
+
----------
|
|
437
|
+
c : xr.DataArray
|
|
438
|
+
The coordinate to check.
|
|
439
|
+
attributes : CoordinateAttributes
|
|
440
|
+
The attributes of the coordinate.
|
|
441
|
+
|
|
442
|
+
Returns
|
|
443
|
+
-------
|
|
444
|
+
Optional[YCoordinate]
|
|
445
|
+
The YCoordinate if matched, else None.
|
|
446
|
+
"""
|
|
447
|
+
pass
|
|
448
|
+
|
|
449
|
+
@abstractmethod
|
|
450
|
+
def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
|
|
451
|
+
"""Checks if the coordinate is a time coordinate.
|
|
452
|
+
|
|
453
|
+
Parameters
|
|
454
|
+
----------
|
|
455
|
+
c : xr.DataArray
|
|
456
|
+
The coordinate to check.
|
|
457
|
+
attributes : CoordinateAttributes
|
|
458
|
+
The attributes of the coordinate.
|
|
459
|
+
|
|
460
|
+
Returns
|
|
461
|
+
-------
|
|
462
|
+
Optional[TimeCoordinate]
|
|
463
|
+
The TimeCoordinate if matched, else None.
|
|
464
|
+
"""
|
|
465
|
+
pass
|
|
466
|
+
|
|
467
|
+
@abstractmethod
|
|
468
|
+
def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
|
|
469
|
+
"""Checks if the coordinate is a date coordinate.
|
|
470
|
+
|
|
471
|
+
Parameters
|
|
472
|
+
----------
|
|
473
|
+
c : xr.DataArray
|
|
474
|
+
The coordinate to check.
|
|
475
|
+
attributes : CoordinateAttributes
|
|
476
|
+
The attributes of the coordinate.
|
|
477
|
+
|
|
478
|
+
Returns
|
|
479
|
+
-------
|
|
480
|
+
Optional[DateCoordinate]
|
|
481
|
+
The DateCoordinate if matched, else None.
|
|
482
|
+
"""
|
|
483
|
+
pass
|
|
484
|
+
|
|
485
|
+
@abstractmethod
|
|
486
|
+
def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
|
|
487
|
+
"""Checks if the coordinate is a step coordinate.
|
|
488
|
+
|
|
489
|
+
Parameters
|
|
490
|
+
----------
|
|
491
|
+
c : xr.DataArray
|
|
492
|
+
The coordinate to check.
|
|
493
|
+
attributes : CoordinateAttributes
|
|
494
|
+
The attributes of the coordinate.
|
|
495
|
+
|
|
496
|
+
Returns
|
|
497
|
+
-------
|
|
498
|
+
Optional[StepCoordinate]
|
|
499
|
+
The StepCoordinate if matched, else None.
|
|
500
|
+
"""
|
|
501
|
+
pass
|
|
502
|
+
|
|
503
|
+
@abstractmethod
|
|
504
|
+
def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
|
|
505
|
+
"""Checks if the coordinate is a level coordinate.
|
|
506
|
+
|
|
507
|
+
Parameters
|
|
508
|
+
----------
|
|
509
|
+
c : xr.DataArray
|
|
510
|
+
The coordinate to check.
|
|
511
|
+
attributes : CoordinateAttributes
|
|
512
|
+
The attributes of the coordinate.
|
|
513
|
+
|
|
514
|
+
Returns
|
|
515
|
+
-------
|
|
516
|
+
Optional[LevelCoordinate]
|
|
517
|
+
The LevelCoordinate if matched, else None.
|
|
518
|
+
"""
|
|
519
|
+
pass
|
|
520
|
+
|
|
521
|
+
@abstractmethod
|
|
522
|
+
def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
|
|
523
|
+
"""Checks if the coordinate is an ensemble coordinate.
|
|
524
|
+
|
|
525
|
+
Parameters
|
|
526
|
+
----------
|
|
527
|
+
c : xr.DataArray
|
|
528
|
+
The coordinate to check.
|
|
529
|
+
attributes : CoordinateAttributes
|
|
530
|
+
The attributes of the coordinate.
|
|
531
|
+
|
|
532
|
+
Returns
|
|
533
|
+
-------
|
|
534
|
+
Optional[EnsembleCoordinate]
|
|
535
|
+
The EnsembleCoordinate if matched, else None.
|
|
536
|
+
"""
|
|
537
|
+
pass
|
|
538
|
+
|
|
539
|
+
|
|
540
|
+
class DefaultCoordinateGuesser(CoordinateGuesser):
|
|
541
|
+
"""Default implementation of CoordinateGuesser."""
|
|
542
|
+
|
|
543
|
+
def __init__(self, ds: xr.Dataset) -> None:
|
|
544
|
+
"""Initializes the DefaultCoordinateGuesser.
|
|
545
|
+
|
|
546
|
+
Parameters
|
|
547
|
+
----------
|
|
548
|
+
ds : xr.Dataset
|
|
549
|
+
The dataset to guess coordinates from.
|
|
550
|
+
"""
|
|
551
|
+
super().__init__(ds)
|
|
552
|
+
|
|
553
|
+
def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
|
|
554
|
+
"""Checks if the coordinate is a longitude.
|
|
555
|
+
|
|
556
|
+
Parameters
|
|
557
|
+
----------
|
|
558
|
+
c : xr.DataArray
|
|
559
|
+
The coordinate to check.
|
|
560
|
+
attributes : CoordinateAttributes
|
|
561
|
+
The attributes of the coordinate.
|
|
562
|
+
|
|
563
|
+
Returns
|
|
564
|
+
-------
|
|
565
|
+
Optional[LongitudeCoordinate]
|
|
566
|
+
The LongitudeCoordinate if matched, else None.
|
|
567
|
+
"""
|
|
568
|
+
if attributes.standard_name == "longitude":
|
|
569
|
+
return LongitudeCoordinate(c)
|
|
570
|
+
|
|
571
|
+
if attributes.long_name == "longitude" and attributes.units == "degrees_east":
|
|
572
|
+
return LongitudeCoordinate(c)
|
|
573
|
+
|
|
574
|
+
if attributes.name == "longitude": # WeatherBench
|
|
575
|
+
return LongitudeCoordinate(c)
|
|
576
|
+
|
|
577
|
+
return None
|
|
578
|
+
|
|
579
|
+
def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
|
|
580
|
+
"""Checks if the coordinate is a latitude.
|
|
581
|
+
|
|
582
|
+
Parameters
|
|
583
|
+
----------
|
|
584
|
+
c : xr.DataArray
|
|
585
|
+
The coordinate to check.
|
|
586
|
+
attributes : CoordinateAttributes
|
|
587
|
+
The attributes of the coordinate.
|
|
588
|
+
|
|
589
|
+
Returns
|
|
590
|
+
-------
|
|
591
|
+
Optional[LatitudeCoordinate]
|
|
592
|
+
The LatitudeCoordinate if matched, else None.
|
|
593
|
+
"""
|
|
594
|
+
if attributes.standard_name == "latitude":
|
|
595
|
+
return LatitudeCoordinate(c)
|
|
596
|
+
|
|
597
|
+
if attributes.long_name == "latitude" and attributes.units == "degrees_north":
|
|
598
|
+
return LatitudeCoordinate(c)
|
|
599
|
+
|
|
600
|
+
if attributes.name == "latitude": # WeatherBench
|
|
601
|
+
return LatitudeCoordinate(c)
|
|
602
|
+
|
|
603
|
+
return None
|
|
604
|
+
|
|
605
|
+
def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
|
|
606
|
+
"""Checks if the coordinate is an x coordinate.
|
|
607
|
+
|
|
608
|
+
Parameters
|
|
609
|
+
----------
|
|
610
|
+
c : xr.DataArray
|
|
611
|
+
The coordinate to check.
|
|
612
|
+
attributes : CoordinateAttributes
|
|
613
|
+
The attributes of the coordinate.
|
|
614
|
+
|
|
615
|
+
Returns
|
|
616
|
+
-------
|
|
617
|
+
Optional[XCoordinate]
|
|
618
|
+
The XCoordinate if matched, else None.
|
|
619
|
+
"""
|
|
620
|
+
if attributes.standard_name in ["projection_x_coordinate", "grid_longitude"]:
|
|
621
|
+
return XCoordinate(c)
|
|
622
|
+
|
|
623
|
+
if attributes.name == "x":
|
|
624
|
+
return XCoordinate(c)
|
|
625
|
+
|
|
626
|
+
return None
|
|
627
|
+
|
|
628
|
+
def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
|
|
629
|
+
"""Checks if the coordinate is a y coordinate.
|
|
630
|
+
|
|
631
|
+
Parameters
|
|
632
|
+
----------
|
|
633
|
+
c : xr.DataArray
|
|
634
|
+
The coordinate to check.
|
|
635
|
+
attributes : CoordinateAttributes
|
|
636
|
+
The attributes of the coordinate.
|
|
637
|
+
|
|
638
|
+
Returns
|
|
639
|
+
-------
|
|
640
|
+
Optional[YCoordinate]
|
|
641
|
+
The YCoordinate if matched, else None.
|
|
642
|
+
"""
|
|
643
|
+
if attributes.standard_name in ["projection_y_coordinate", "grid_latitude"]:
|
|
644
|
+
return YCoordinate(c)
|
|
645
|
+
|
|
646
|
+
if attributes.name == "y":
|
|
647
|
+
return YCoordinate(c)
|
|
648
|
+
|
|
649
|
+
return None
|
|
650
|
+
|
|
651
|
+
def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
|
|
652
|
+
"""Checks if the coordinate is a time coordinate.
|
|
653
|
+
|
|
654
|
+
Parameters
|
|
655
|
+
----------
|
|
656
|
+
c : xr.DataArray
|
|
657
|
+
The coordinate to check.
|
|
658
|
+
attributes : CoordinateAttributes
|
|
659
|
+
The attributes of the coordinate.
|
|
660
|
+
|
|
661
|
+
Returns
|
|
662
|
+
-------
|
|
663
|
+
Optional[TimeCoordinate]
|
|
664
|
+
The TimeCoordinate if matched, else None.
|
|
665
|
+
"""
|
|
666
|
+
if attributes.standard_name == "time":
|
|
667
|
+
return TimeCoordinate(c)
|
|
668
|
+
|
|
669
|
+
if attributes.name == "time" and attributes.standard_name != "forecast_reference_time":
|
|
670
|
+
return TimeCoordinate(c)
|
|
671
|
+
|
|
672
|
+
return None
|
|
673
|
+
|
|
674
|
+
def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
|
|
675
|
+
"""Checks if the coordinate is a date coordinate.
|
|
676
|
+
|
|
677
|
+
Parameters
|
|
678
|
+
----------
|
|
679
|
+
c : xr.DataArray
|
|
680
|
+
The coordinate to check.
|
|
681
|
+
attributes : CoordinateAttributes
|
|
682
|
+
The attributes of the coordinate.
|
|
683
|
+
|
|
684
|
+
Returns
|
|
685
|
+
-------
|
|
686
|
+
Optional[DateCoordinate]
|
|
687
|
+
The DateCoordinate if matched, else None.
|
|
688
|
+
"""
|
|
689
|
+
if attributes.standard_name == "forecast_reference_time":
|
|
690
|
+
return DateCoordinate(c)
|
|
691
|
+
|
|
692
|
+
if attributes.name == "forecast_reference_time":
|
|
693
|
+
return DateCoordinate(c)
|
|
694
|
+
|
|
695
|
+
return None
|
|
696
|
+
|
|
697
|
+
def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
|
|
698
|
+
"""Checks if the coordinate is a step coordinate.
|
|
699
|
+
|
|
700
|
+
Parameters
|
|
701
|
+
----------
|
|
702
|
+
c : xr.DataArray
|
|
703
|
+
The coordinate to check.
|
|
704
|
+
attributes : CoordinateAttributes
|
|
705
|
+
The attributes of the coordinate.
|
|
706
|
+
|
|
707
|
+
Returns
|
|
708
|
+
-------
|
|
709
|
+
Optional[StepCoordinate]
|
|
710
|
+
The StepCoordinate if matched, else None.
|
|
711
|
+
"""
|
|
712
|
+
if attributes.standard_name == "forecast_period":
|
|
713
|
+
return StepCoordinate(c)
|
|
714
|
+
|
|
715
|
+
if attributes.long_name == "time elapsed since the start of the forecast":
|
|
716
|
+
return StepCoordinate(c)
|
|
717
|
+
|
|
718
|
+
if attributes.name == "prediction_timedelta": # WeatherBench
|
|
719
|
+
return StepCoordinate(c)
|
|
720
|
+
|
|
721
|
+
return None
|
|
722
|
+
|
|
723
|
+
def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
|
|
724
|
+
"""Checks if the coordinate is a level coordinate.
|
|
725
|
+
|
|
726
|
+
Parameters
|
|
727
|
+
----------
|
|
728
|
+
c : xr.DataArray
|
|
729
|
+
The coordinate to check.
|
|
730
|
+
attributes : CoordinateAttributes
|
|
731
|
+
The attributes of the coordinate.
|
|
732
|
+
|
|
733
|
+
Returns
|
|
734
|
+
-------
|
|
735
|
+
Optional[LevelCoordinate]
|
|
736
|
+
The LevelCoordinate if matched, else None.
|
|
737
|
+
"""
|
|
738
|
+
if attributes.standard_name == "atmosphere_hybrid_sigma_pressure_coordinate":
|
|
739
|
+
return LevelCoordinate(c, "ml")
|
|
740
|
+
|
|
741
|
+
if attributes.long_name == "height" and attributes.units == "m":
|
|
742
|
+
return LevelCoordinate(c, "height")
|
|
743
|
+
|
|
744
|
+
if attributes.standard_name == "air_pressure" and attributes.units == "hPa":
|
|
745
|
+
return LevelCoordinate(c, "pl")
|
|
746
|
+
|
|
747
|
+
if attributes.name == "level":
|
|
748
|
+
return LevelCoordinate(c, "pl")
|
|
749
|
+
|
|
750
|
+
if attributes.name == "vertical" and attributes.units == "hPa":
|
|
751
|
+
return LevelCoordinate(c, "pl")
|
|
752
|
+
|
|
753
|
+
if attributes.standard_name == "depth":
|
|
754
|
+
return LevelCoordinate(c, "depth")
|
|
755
|
+
|
|
756
|
+
if attributes.name == "vertical" and attributes.units == "hPa":
|
|
757
|
+
return LevelCoordinate(c, "pl")
|
|
758
|
+
|
|
759
|
+
return None
|
|
760
|
+
|
|
761
|
+
def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
|
|
762
|
+
"""Checks if the coordinate is an ensemble coordinate.
|
|
763
|
+
|
|
764
|
+
Parameters
|
|
765
|
+
----------
|
|
766
|
+
c : xr.DataArray
|
|
767
|
+
The coordinate to check.
|
|
768
|
+
attributes : CoordinateAttributes
|
|
769
|
+
The attributes of the coordinate.
|
|
770
|
+
|
|
771
|
+
Returns
|
|
772
|
+
-------
|
|
773
|
+
Optional[EnsembleCoordinate]
|
|
774
|
+
The EnsembleCoordinate if matched, else None.
|
|
775
|
+
"""
|
|
776
|
+
if attributes.name in ("realization", "number"):
|
|
777
|
+
return EnsembleCoordinate(c)
|
|
778
|
+
|
|
779
|
+
return None
|
|
780
|
+
|
|
781
|
+
|
|
782
|
+
class FlavourCoordinateGuesser(CoordinateGuesser):
|
|
783
|
+
"""Implementation of CoordinateGuesser that uses a flavour for guessing."""
|
|
784
|
+
|
|
785
|
+
def __init__(self, ds: xr.Dataset, flavour: Dict[str, Any]) -> None:
|
|
786
|
+
"""Initializes the FlavourCoordinateGuesser.
|
|
787
|
+
|
|
788
|
+
Parameters
|
|
789
|
+
----------
|
|
790
|
+
ds : xr.Dataset
|
|
791
|
+
The dataset to guess coordinates from.
|
|
792
|
+
flavour : Dict[str, Any]
|
|
793
|
+
The flavour to use for guessing.
|
|
794
|
+
"""
|
|
795
|
+
super().__init__(ds)
|
|
796
|
+
self.flavour = flavour
|
|
797
|
+
|
|
798
|
+
def _match(self, c: xr.DataArray, key: str, attributes: CoordinateAttributes) -> Optional[Dict[str, Any]]:
|
|
799
|
+
"""Matches the coordinate against the flavour rules.
|
|
800
|
+
|
|
801
|
+
Parameters
|
|
802
|
+
----------
|
|
803
|
+
c : xr.DataArray
|
|
804
|
+
The coordinate to match.
|
|
805
|
+
key : str
|
|
806
|
+
The key to match in the flavour rules.
|
|
807
|
+
attributes : CoordinateAttributes
|
|
808
|
+
The values to match against.
|
|
809
|
+
|
|
810
|
+
Returns
|
|
811
|
+
-------
|
|
812
|
+
Optional[Dict[str, Any]]
|
|
813
|
+
The matched rule if any, else None.
|
|
814
|
+
"""
|
|
815
|
+
if key not in self.flavour["rules"]:
|
|
816
|
+
return None
|
|
817
|
+
|
|
818
|
+
rules = self.flavour["rules"][key]
|
|
819
|
+
|
|
820
|
+
if not isinstance(rules, list):
|
|
821
|
+
rules = [rules]
|
|
822
|
+
|
|
823
|
+
for rule in rules:
|
|
824
|
+
ok = True
|
|
825
|
+
for k, v in rule.items():
|
|
826
|
+
if isinstance(v, str) and attributes.get(k) != v:
|
|
827
|
+
ok = False
|
|
828
|
+
if ok:
|
|
829
|
+
return rule
|
|
830
|
+
|
|
831
|
+
return None
|
|
832
|
+
|
|
833
|
+
def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
|
|
834
|
+
"""Checks if the coordinate is a longitude using the flavour rules.
|
|
835
|
+
|
|
836
|
+
Parameters
|
|
837
|
+
----------
|
|
838
|
+
c : xr.DataArray
|
|
839
|
+
The coordinate to check.
|
|
840
|
+
attributes : CoordinateAttributes
|
|
841
|
+
The attributes of the coordinate.
|
|
842
|
+
|
|
843
|
+
Returns
|
|
844
|
+
-------
|
|
845
|
+
Optional[LongitudeCoordinate]
|
|
846
|
+
The LongitudeCoordinate if matched, else None.
|
|
847
|
+
"""
|
|
848
|
+
if self._match(c, "longitude", attributes):
|
|
849
|
+
return LongitudeCoordinate(c)
|
|
850
|
+
|
|
851
|
+
return None
|
|
852
|
+
|
|
853
|
+
def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
|
|
854
|
+
"""Checks if the coordinate is a latitude using the flavour rules.
|
|
855
|
+
|
|
856
|
+
Parameters
|
|
857
|
+
----------
|
|
858
|
+
c : xr.DataArray
|
|
859
|
+
The coordinate to check.
|
|
860
|
+
attributes : CoordinateAttributes
|
|
861
|
+
The attributes of the coordinate.
|
|
862
|
+
|
|
863
|
+
Returns
|
|
864
|
+
-------
|
|
865
|
+
Optional[LatitudeCoordinate]
|
|
866
|
+
The LatitudeCoordinate if matched, else None.
|
|
867
|
+
"""
|
|
868
|
+
if self._match(c, "latitude", attributes):
|
|
869
|
+
return LatitudeCoordinate(c)
|
|
870
|
+
|
|
871
|
+
return None
|
|
872
|
+
|
|
873
|
+
def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
|
|
874
|
+
"""Checks if the coordinate is an x coordinate using the flavour rules.
|
|
875
|
+
|
|
876
|
+
Parameters
|
|
877
|
+
----------
|
|
878
|
+
c : xr.DataArray
|
|
879
|
+
The coordinate to check.
|
|
880
|
+
attributes : CoordinateAttributes
|
|
881
|
+
The attributes of the coordinate.
|
|
882
|
+
|
|
883
|
+
Returns
|
|
884
|
+
-------
|
|
885
|
+
Optional[XCoordinate]
|
|
886
|
+
The XCoordinate if matched, else None.
|
|
887
|
+
"""
|
|
888
|
+
if self._match(c, "x", attributes):
|
|
889
|
+
return XCoordinate(c)
|
|
890
|
+
|
|
891
|
+
return None
|
|
892
|
+
|
|
893
|
+
def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
|
|
894
|
+
"""Checks if the coordinate is a y coordinate using the flavour rules.
|
|
895
|
+
|
|
896
|
+
Parameters
|
|
897
|
+
----------
|
|
898
|
+
c : xr.DataArray
|
|
899
|
+
The coordinate to check.
|
|
900
|
+
attributes : CoordinateAttributes
|
|
901
|
+
The attributes of the coordinate.
|
|
902
|
+
|
|
903
|
+
Returns
|
|
904
|
+
-------
|
|
905
|
+
Optional[YCoordinate]
|
|
906
|
+
The YCoordinate if matched, else None.
|
|
907
|
+
"""
|
|
908
|
+
if self._match(c, "y", attributes):
|
|
909
|
+
return YCoordinate(c)
|
|
910
|
+
|
|
911
|
+
return None
|
|
912
|
+
|
|
913
|
+
def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
|
|
914
|
+
"""Checks if the coordinate is a time coordinate using the flavour rules.
|
|
915
|
+
|
|
916
|
+
Parameters
|
|
917
|
+
----------
|
|
918
|
+
c : xr.DataArray
|
|
919
|
+
The coordinate to check.
|
|
920
|
+
attributes : CoordinateAttributes
|
|
921
|
+
The attributes of the coordinate.
|
|
922
|
+
|
|
923
|
+
Returns
|
|
924
|
+
-------
|
|
925
|
+
Optional[TimeCoordinate]
|
|
926
|
+
The TimeCoordinate if matched, else None.
|
|
927
|
+
"""
|
|
928
|
+
if self._match(c, "time", attributes):
|
|
929
|
+
return TimeCoordinate(c)
|
|
930
|
+
|
|
931
|
+
return None
|
|
932
|
+
|
|
933
|
+
def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
|
|
934
|
+
"""Checks if the coordinate is a step coordinate using the flavour rules.
|
|
935
|
+
|
|
936
|
+
Parameters
|
|
937
|
+
----------
|
|
938
|
+
c : xr.DataArray
|
|
939
|
+
The coordinate to check.
|
|
940
|
+
attributes : CoordinateAttributes
|
|
941
|
+
The attributes of the coordinate.
|
|
942
|
+
|
|
943
|
+
Returns
|
|
944
|
+
-------
|
|
945
|
+
Optional[StepCoordinate]
|
|
946
|
+
The StepCoordinate if matched, else None.
|
|
947
|
+
"""
|
|
948
|
+
if self._match(c, "step", attributes):
|
|
949
|
+
return StepCoordinate(c)
|
|
950
|
+
|
|
951
|
+
return None
|
|
952
|
+
|
|
953
|
+
def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
|
|
954
|
+
"""Checks if the coordinate is a date coordinate using the flavour rules.
|
|
955
|
+
|
|
956
|
+
Parameters
|
|
957
|
+
----------
|
|
958
|
+
c : xr.DataArray
|
|
959
|
+
The coordinate to check.
|
|
960
|
+
attributes : CoordinateAttributes
|
|
961
|
+
The attributes of the coordinate.
|
|
962
|
+
|
|
963
|
+
Returns
|
|
964
|
+
-------
|
|
965
|
+
Optional[DateCoordinate]
|
|
966
|
+
The DateCoordinate if matched, else None.
|
|
967
|
+
"""
|
|
968
|
+
if self._match(c, "date", attributes):
|
|
969
|
+
return DateCoordinate(c)
|
|
970
|
+
|
|
971
|
+
return None
|
|
972
|
+
|
|
973
|
+
def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
|
|
974
|
+
"""Checks if the coordinate is a level coordinate using the flavour rules.
|
|
975
|
+
|
|
976
|
+
Parameters
|
|
977
|
+
----------
|
|
978
|
+
c : xr.DataArray
|
|
979
|
+
The coordinate to check.
|
|
980
|
+
attributes : CoordinateAttributes
|
|
981
|
+
The attributes of the coordinate.
|
|
982
|
+
|
|
983
|
+
Returns
|
|
984
|
+
-------
|
|
985
|
+
Optional[LevelCoordinate]
|
|
986
|
+
The LevelCoordinate if matched, else None.
|
|
987
|
+
"""
|
|
988
|
+
rule = self._match(c, "level", attributes)
|
|
989
|
+
if rule:
|
|
990
|
+
# assert False, rule
|
|
991
|
+
return LevelCoordinate(
|
|
992
|
+
c,
|
|
993
|
+
self._levtype(c, attributes),
|
|
994
|
+
)
|
|
995
|
+
|
|
996
|
+
return None
|
|
997
|
+
|
|
998
|
+
def _levtype(self, c: xr.DataArray, attributes: CoordinateAttributes) -> str:
|
|
999
|
+
"""Determines the level type for the coordinate.
|
|
1000
|
+
|
|
1001
|
+
Parameters
|
|
1002
|
+
----------
|
|
1003
|
+
c : xr.DataArray
|
|
1004
|
+
The coordinate to check.
|
|
1005
|
+
attributes : CoordinateAttributes
|
|
1006
|
+
The attributes of the coordinate.
|
|
1007
|
+
|
|
1008
|
+
Returns
|
|
1009
|
+
-------
|
|
1010
|
+
str
|
|
1011
|
+
The level type.
|
|
1012
|
+
"""
|
|
1013
|
+
if "levtype" in self.flavour:
|
|
1014
|
+
return self.flavour["levtype"]
|
|
1015
|
+
|
|
1016
|
+
raise NotImplementedError(f"levtype for {c=}")
|
|
1017
|
+
|
|
1018
|
+
def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
|
|
1019
|
+
"""Checks if the coordinate is an ensemble coordinate using the flavour rules.
|
|
1020
|
+
|
|
1021
|
+
Parameters
|
|
1022
|
+
----------
|
|
1023
|
+
c : xr.DataArray
|
|
1024
|
+
The coordinate to check.
|
|
1025
|
+
attributes : CoordinateAttributes
|
|
1026
|
+
The attributes of the coordinate.
|
|
1027
|
+
|
|
1028
|
+
Returns
|
|
1029
|
+
-------
|
|
1030
|
+
Optional[EnsembleCoordinate]
|
|
1031
|
+
The EnsembleCoordinate if matched, else None.
|
|
1032
|
+
"""
|
|
1033
|
+
if self._match(c, "number", attributes):
|
|
1034
|
+
return EnsembleCoordinate(c)
|
|
1035
|
+
|
|
1036
|
+
return None
|