anemoi-datasets 0.5.28__py3-none-any.whl → 0.5.30__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.
Files changed (28) hide show
  1. anemoi/datasets/_version.py +2 -2
  2. anemoi/datasets/create/__init__.py +4 -12
  3. anemoi/datasets/create/config.py +50 -53
  4. anemoi/datasets/create/input/result/field.py +1 -3
  5. anemoi/datasets/create/sources/accumulate.py +517 -0
  6. anemoi/datasets/create/sources/accumulate_utils/__init__.py +8 -0
  7. anemoi/datasets/create/sources/accumulate_utils/covering_intervals.py +221 -0
  8. anemoi/datasets/create/sources/accumulate_utils/field_to_interval.py +153 -0
  9. anemoi/datasets/create/sources/accumulate_utils/interval_generators.py +321 -0
  10. anemoi/datasets/create/sources/grib_index.py +79 -51
  11. anemoi/datasets/create/sources/mars.py +56 -27
  12. anemoi/datasets/create/sources/xarray_support/__init__.py +1 -0
  13. anemoi/datasets/create/sources/xarray_support/coordinates.py +1 -4
  14. anemoi/datasets/create/sources/xarray_support/flavour.py +2 -2
  15. anemoi/datasets/create/sources/xarray_support/patch.py +178 -5
  16. anemoi/datasets/data/complement.py +26 -17
  17. anemoi/datasets/data/dataset.py +6 -0
  18. anemoi/datasets/data/masked.py +74 -13
  19. anemoi/datasets/data/missing.py +5 -0
  20. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/METADATA +8 -7
  21. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/RECORD +25 -23
  22. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/WHEEL +1 -1
  23. anemoi/datasets/create/sources/accumulations.py +0 -1042
  24. anemoi/datasets/create/sources/accumulations2.py +0 -618
  25. anemoi/datasets/create/sources/tendencies.py +0 -171
  26. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/entry_points.txt +0 -0
  27. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/licenses/LICENSE +0 -0
  28. {anemoi_datasets-0.5.28.dist-info → anemoi_datasets-0.5.30.dist-info}/top_level.txt +0 -0
@@ -1,1042 +0,0 @@
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
- import datetime
11
- import logging
12
- import warnings
13
- from collections.abc import Generator
14
- from copy import deepcopy
15
- from typing import Any
16
-
17
- import earthkit.data as ekd
18
- import numpy as np
19
- from earthkit.data.core.temporary import temp_file
20
- from earthkit.data.readers.grib.output import new_grib_output
21
- from numpy.typing import NDArray
22
-
23
- from anemoi.datasets.create.sources import source_registry
24
-
25
- from .legacy import LegacySource
26
- from .mars import mars
27
-
28
- LOG = logging.getLogger(__name__)
29
- MISSING_VALUE = 1e-38
30
-
31
-
32
- def _member(field: Any) -> int:
33
- """Retrieves the member number from the field metadata.
34
-
35
- Parameters
36
- ----------
37
- field : Any
38
- The field from which to retrieve the member number.
39
-
40
- Returns
41
- -------
42
- int
43
- The member number.
44
- """
45
- # Bug in eccodes has number=0 randomly
46
- number = field.metadata("number", default=0)
47
- if number is None:
48
- number = 0
49
- return number
50
-
51
-
52
- class Accumulation:
53
- """Class to handle data accumulation for a specific parameter, date, time, and member."""
54
-
55
- buggy_steps: bool = False
56
-
57
- def __init__(
58
- self,
59
- out: Any,
60
- /,
61
- param: str,
62
- date: int,
63
- time: int,
64
- number: int,
65
- step: list[int],
66
- frequency: int,
67
- accumulations_reset_frequency: int | None = None,
68
- user_date: str | None = None,
69
- **kwargs: Any,
70
- ) -> None:
71
- """Initialises an Accumulation instance.
72
-
73
- Parameters
74
- ----------
75
- out : Any
76
- Output object for writing data.
77
- param : str
78
- Parameter name.
79
- date : int
80
- Date of the accumulation.
81
- time : int
82
- Time of the accumulation.
83
- number : int
84
- Member number.
85
- step : List[int]
86
- List of steps.
87
- frequency : int
88
- Frequency of accumulation.
89
- accumulations_reset_frequency : Optional[int], optional
90
- Frequency at which accumulations reset. Defaults to None.
91
- user_date : Optional[str], optional
92
- User-defined date. Defaults to None.
93
- **kwargs : Any
94
- Additional keyword arguments.
95
- """
96
- self.out = out
97
- self.param = param
98
- self.date = date
99
- self.time = time
100
- self.steps = step
101
- self.number = number
102
- self.values: NDArray[None] | None = None
103
- self.seen = set()
104
- self.startStep: int | None = None
105
- self.endStep: int | None = None
106
- self.done = False
107
- self.frequency = frequency
108
- self.accumulations_reset_frequency = accumulations_reset_frequency
109
- self._check = None
110
- self.user_date = user_date
111
-
112
- @property
113
- def key(self) -> tuple[str, int, int, list[int], int]:
114
- """Returns the key for the accumulation."""
115
- return (self.param, self.date, self.time, self.steps, self.number)
116
-
117
- def check(self, field: Any) -> None:
118
- """Checks the field metadata against the accumulation parameters.
119
-
120
- Parameters
121
- ----------
122
- field : Any
123
- The field to check.
124
- """
125
- if self._check is None:
126
- self._check = field.metadata(namespace="mars")
127
-
128
- assert self.param == field.metadata("param"), (
129
- self.param,
130
- field.metadata("param"),
131
- )
132
- assert self.date == field.metadata("date"), (
133
- self.date,
134
- field.metadata("date"),
135
- )
136
- assert self.time == field.metadata("time"), (
137
- self.time,
138
- field.metadata("time"),
139
- )
140
- assert self.number == _member(field), (self.number, _member(field))
141
-
142
- return
143
-
144
- mars = field.metadata(namespace="mars")
145
- keys1 = sorted(self._check.keys())
146
- keys2 = sorted(mars.keys())
147
-
148
- assert keys1 == keys2, (keys1, keys2)
149
-
150
- for k in keys1:
151
- if k not in ("step",):
152
- assert self._check[k] == mars[k], (k, self._check[k], mars[k])
153
-
154
- def write(self, template: Any) -> None:
155
- """Writes the accumulated values to the output.
156
-
157
- Parameters
158
- ----------
159
- template : Any
160
- Template for writing the output.
161
- """
162
- assert self.startStep != self.endStep, (self.startStep, self.endStep)
163
- if np.all(self.values < 0):
164
- LOG.warning(
165
- f"Negative values when computing accumutation for {self.param} ({self.date} {self.time}): min={np.amin(self.values)} max={np.amax(self.values)}"
166
- )
167
-
168
- # In GRIB1, is the step is greater that 254 (one byte), we cannot use a range, because both P1 and P2 values
169
- # are used to store the end step
170
-
171
- edition = template.metadata("edition")
172
- assert np.all(self.values != MISSING_VALUE)
173
-
174
- if edition == 1 and self.endStep > 254:
175
- self.out.write(
176
- self.values,
177
- template=template,
178
- stepType="instant",
179
- step=self.endStep,
180
- check_nans=True,
181
- missing_value=MISSING_VALUE,
182
- )
183
- else:
184
- self.out.write(
185
- self.values,
186
- template=template,
187
- stepType="accum",
188
- startStep=self.startStep,
189
- endStep=self.endStep,
190
- check_nans=True,
191
- missing_value=MISSING_VALUE,
192
- )
193
- self.values = None
194
- self.done = True
195
-
196
- def add(self, field: Any, values: NDArray[Any]) -> None:
197
- """Adds values to the accumulation.
198
-
199
- Parameters
200
- ----------
201
- field : Any
202
- The field containing the values.
203
- values : NDArray[Any]
204
- The values to add.
205
- """
206
- self.check(field)
207
-
208
- step = field.metadata("step")
209
- if step not in self.steps:
210
- return
211
-
212
- assert not self.done, (self.key, step)
213
- assert step not in self.seen, (self.key, step)
214
-
215
- startStep = field.metadata("startStep")
216
- endStep = field.metadata("endStep")
217
-
218
- if startStep == endStep:
219
- startStep, endStep = self.adjust_steps(startStep, endStep)
220
-
221
- assert step == endStep, (startStep, endStep, step)
222
-
223
- self.compute(values, startStep, endStep)
224
-
225
- self.seen.add(step)
226
-
227
- if len(self.seen) == len(self.steps):
228
- self.write(template=field)
229
-
230
- @classmethod
231
- def mars_date_time_steps(
232
- cls,
233
- *,
234
- dates: list[datetime.datetime],
235
- step1: int,
236
- step2: int,
237
- frequency: int | None,
238
- base_times: list[int],
239
- adjust_step: bool,
240
- accumulations_reset_frequency: int | None,
241
- user_date: str | None,
242
- ) -> Generator[tuple[int, int, tuple[int, ...]], None, None]:
243
- """Generates MARS date-time steps.
244
-
245
- Parameters
246
- ----------
247
- dates : List[datetime.datetime]
248
- List of dates.
249
- step1 : int
250
- First step.
251
- step2 : int
252
- Second step.
253
- frequency : Optional[int]
254
- Frequency of accumulation.
255
- base_times : List[int]
256
- List of base times.
257
- adjust_step : bool
258
- Whether to adjust the step.
259
- accumulations_reset_frequency : Optional[int], optional
260
- Frequency at which accumulations reset. Defaults to None.
261
- user_date : Optional[str], optional
262
- User-defined date. Defaults to None.
263
-
264
- Returns
265
- -------
266
- Generator[Tuple[int, int, Tuple[int, ...]], None, None]
267
- A generator of MARS date-time steps.
268
- """
269
- # assert step1 > 0, (step1, step2, frequency)
270
-
271
- for valid_date in dates:
272
- add_step = 0
273
- base_date = valid_date - datetime.timedelta(hours=step2)
274
- if user_date is not None:
275
- assert user_date == "????-??-01", user_date
276
- new_base_date = base_date.replace(day=1)
277
- assert new_base_date <= base_date, (new_base_date, base_date)
278
- add_step = int((base_date - new_base_date).total_seconds() // 3600)
279
-
280
- base_date = new_base_date
281
-
282
- if base_date.hour not in base_times:
283
- if not adjust_step:
284
- raise ValueError(
285
- f"Cannot find a base time in {base_times} that validates on {valid_date.isoformat()} for step={step2}"
286
- )
287
-
288
- while base_date.hour not in base_times:
289
- # print(f'{base_date=}, {base_times=}, {add_step=} {frequency=}')
290
- base_date -= datetime.timedelta(hours=1)
291
- add_step += 1
292
-
293
- yield cls._mars_date_time_step(
294
- base_date=base_date,
295
- step1=step1,
296
- step2=step2,
297
- add_step=add_step,
298
- frequency=frequency,
299
- accumulations_reset_frequency=accumulations_reset_frequency,
300
- user_date=user_date,
301
- requested_date=valid_date,
302
- )
303
-
304
- def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
305
- """Computes the accumulation.
306
-
307
- Parameters
308
- ----------
309
- values : NDArray[Any]
310
- The values to accumulate.
311
- startStep : int
312
- The start step.
313
- endStep : int
314
- The end step.
315
- """
316
- pass
317
-
318
- @classmethod
319
- def _mars_date_time_step(
320
- cls,
321
- *,
322
- base_date: datetime.datetime,
323
- step1: int,
324
- step2: int,
325
- add_step: int,
326
- frequency: int | None,
327
- accumulations_reset_frequency: int | None,
328
- user_date: str | None,
329
- requested_date: datetime.datetime | None = None,
330
- ) -> tuple[int, int, tuple[int, ...]]:
331
- """Generates a MARS date-time step.
332
-
333
- Parameters
334
- ----------
335
- base_date : datetime.datetime
336
- The base date.
337
- step1 : int
338
- First step.
339
- step2 : int
340
- Second step.
341
- add_step : int
342
- Additional step.
343
- frequency : Optional[int]
344
- Frequency of accumulation.
345
- accumulations_reset_frequency : Optional[int], optional
346
- Frequency at which accumulations reset. Defaults to None.
347
- user_date : Optional[str], optional
348
- User-defined date. Defaults to None.
349
- requested_date : Optional[datetime.datetime], optional
350
- Requested date. Defaults to None.
351
-
352
- Returns
353
- -------
354
- Tuple[int, int, Tuple[int, ...]]
355
- A tuple representing the MARS date-time step.
356
- """
357
- pass
358
-
359
-
360
- class AccumulationFromStart(Accumulation):
361
- """Class to handle data accumulation from the start of the forecast."""
362
-
363
- def adjust_steps(self, startStep: int, endStep: int) -> tuple[int, int]:
364
- """Adjusts the start and end steps.
365
-
366
- Parameters
367
- ----------
368
- startStep : int
369
- The start step.
370
- endStep : int
371
- The end step.
372
-
373
- Returns
374
- -------
375
- Tuple[int, int]
376
- The adjusted start and end steps.
377
- """
378
- assert endStep == startStep
379
- return (0, endStep)
380
-
381
- def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
382
- """Computes the accumulation from the start.
383
-
384
- Parameters
385
- ----------
386
- values : np.ndarray
387
- The values to accumulate.
388
- startStep : int
389
- The start step.
390
- endStep : int
391
- The end step.
392
- """
393
- assert startStep == 0, startStep
394
-
395
- if self.values is None:
396
-
397
- self.values = np.copy(values)
398
- self.startStep = 0
399
- self.endStep = endStep
400
-
401
- else:
402
- assert endStep != self.endStep, (self.endStep, endStep)
403
-
404
- if endStep > self.endStep:
405
- # assert endStep - self.endStep == self.stepping, (self.endStep, endStep, self.stepping)
406
- self.values = values - self.values
407
- self.startStep = self.endStep
408
- self.endStep = endStep
409
- else:
410
- # assert self.endStep - endStep == self.stepping, (self.endStep, endStep, self.stepping)
411
- self.values = self.values - values
412
- self.startStep = endStep
413
-
414
- if not np.all(self.values >= 0):
415
- warnings.warn(f"Negative values for {self.param}: {np.amin(self.values)} {np.amax(self.values)}")
416
- self.values = np.maximum(self.values, 0)
417
-
418
- @classmethod
419
- def _mars_date_time_step(
420
- cls,
421
- *,
422
- base_date: datetime.datetime,
423
- step1: int,
424
- step2: int,
425
- add_step: int,
426
- frequency: int | None,
427
- accumulations_reset_frequency: int | None,
428
- user_date: str | None,
429
- requested_date: datetime.datetime | None = None,
430
- ) -> tuple[int, int, tuple[int, ...]]:
431
- """Generates a MARS date-time step.
432
-
433
- Parameters
434
- ----------
435
- base_date : datetime.datetime
436
- The base date.
437
- step1 : int
438
- First step.
439
- step2 : int
440
- Second step.
441
- add_step : int
442
- Additional step.
443
- frequency : Optional[int]
444
- Frequency of accumulation.
445
- accumulations_reset_frequency : Optional[int], optional
446
- Frequency at which accumulations reset. Defaults to None.
447
- user_date : Optional[str], optional
448
- User-defined date. Defaults to None.
449
- requested_date : Optional[datetime.datetime], optional
450
- Requested date. Defaults to None.
451
-
452
- Returns
453
- -------
454
- Tuple[int, int, Tuple[int, ...]]
455
- A tuple representing the MARS date-time step.
456
- """
457
- assert user_date is None, user_date
458
-
459
- steps = (step1 + add_step, step2 + add_step)
460
- if steps[0] == 0:
461
- steps = (steps[1],)
462
-
463
- assert frequency == 0 or frequency == (step2 - step1), frequency
464
-
465
- return (
466
- base_date.year * 10000 + base_date.month * 100 + base_date.day,
467
- base_date.hour * 100 + base_date.minute,
468
- steps,
469
- )
470
-
471
-
472
- class AccumulationFromLastStep(Accumulation):
473
- """Class to handle data accumulation from the last step of the forecast."""
474
-
475
- def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
476
- """Computes the accumulation from the last step.
477
-
478
- Parameters
479
- ----------
480
- values : np.ndarray
481
- The values to accumulate.
482
- startStep : int
483
- The start step.
484
- endStep : int
485
- The end step.
486
- """
487
- assert endStep - startStep == self.frequency, (
488
- startStep,
489
- endStep,
490
- self.frequency,
491
- )
492
-
493
- if self.startStep is None:
494
- self.startStep = startStep
495
- else:
496
- self.startStep = min(self.startStep, startStep)
497
-
498
- if self.endStep is None:
499
- self.endStep = endStep
500
- else:
501
- self.endStep = max(self.endStep, endStep)
502
-
503
- if self.values is None:
504
- self.values = np.zeros_like(values)
505
-
506
- self.values += values
507
-
508
- @classmethod
509
- def _mars_date_time_step(
510
- cls,
511
- *,
512
- base_date: datetime.datetime,
513
- step1: int,
514
- step2: int,
515
- add_step: int,
516
- frequency: int,
517
- accumulations_reset_frequency: int | None,
518
- user_date: str | None = None,
519
- requested_date: datetime.datetime | None = None,
520
- ) -> tuple[int, int, tuple[int, ...]]:
521
- """Generates a MARS date-time step.
522
-
523
- Parameters
524
- ----------
525
- base_date : datetime.datetime
526
- The base date.
527
- step1 : int
528
- First step.
529
- step2 : int
530
- Second step.
531
- add_step : int
532
- Additional step.
533
- frequency : int
534
- Frequency of accumulation.
535
- accumulations_reset_frequency : Optional[int], optional
536
- Frequency at which accumulations reset. Defaults to None.
537
- user_date : Optional[str], optional
538
- User-defined date. Defaults to None.
539
- requested_date : Optional[datetime.datetime], optional
540
- Requested date. Defaults to None.
541
-
542
- Returns
543
- -------
544
- Tuple[int, int, Tuple[int, ...]]
545
- A tuple representing the MARS date-time step.
546
- """
547
-
548
- assert user_date is None, user_date
549
-
550
- assert frequency > 0, frequency
551
- # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
552
-
553
- steps = []
554
- for step in range(step1 + frequency, step2 + frequency, frequency):
555
- steps.append(step + add_step)
556
-
557
- return (
558
- base_date.year * 10000 + base_date.month * 100 + base_date.day,
559
- base_date.hour * 100 + base_date.minute,
560
- tuple(steps),
561
- )
562
-
563
-
564
- class AccumulationFromLastReset(Accumulation):
565
- """Class to handle data accumulation from the last step of the forecast."""
566
-
567
- def adjust_steps(self, startStep: int, endStep: int) -> tuple[int, int]:
568
- """Adjusts the start and end steps.
569
-
570
- Parameters
571
- ----------
572
- startStep : int
573
- The start step.
574
- endStep : int
575
- The end step.
576
-
577
- Returns
578
- -------
579
- Tuple[int, int]
580
- The adjusted start and end steps.
581
- """
582
- return self.__class__._adjust_steps(startStep, endStep, self.frequency, self.accumulations_reset_frequency)
583
-
584
- @classmethod
585
- def _adjust_steps(
586
- self, startStep: int, endStep: int, frequency: int, accumulations_reset_frequency: int
587
- ) -> tuple[int, int]:
588
- """Adjusts the start and end steps.
589
-
590
- Parameters
591
- ----------
592
- startStep : int
593
- The start step.
594
- endStep : int
595
- The end step.
596
- frequency : int
597
- Frequency of accumulation.
598
- accumulations_reset_frequency : int
599
- Frequency at which accumulations reset.
600
-
601
- Returns
602
- -------
603
- Tuple[int, int]
604
- The adjusted start and end steps.
605
- """
606
-
607
- assert frequency == 1, (frequency, startStep, endStep)
608
- assert endStep - startStep <= accumulations_reset_frequency, (startStep, endStep)
609
-
610
- return ((startStep // accumulations_reset_frequency) * accumulations_reset_frequency, endStep)
611
-
612
- @classmethod
613
- def _steps(
614
- cls,
615
- valid_date: datetime.datetime,
616
- base_date: datetime.datetime,
617
- frequency: int,
618
- accumulations_reset_frequency: int,
619
- ) -> tuple[int, int]:
620
- """Calculates the steps for accumulation.
621
-
622
- Parameters
623
- ----------
624
- valid_date : datetime.datetime
625
- The valid date.
626
- base_date : datetime.datetime
627
- The base date.
628
- frequency : int
629
- Frequency of accumulation.
630
- accumulations_reset_frequency : int
631
- Frequency at which accumulations reset.
632
-
633
- Returns
634
- -------
635
- Tuple[int, int]
636
- A tuple representing the steps for accumulation.
637
- """
638
-
639
- assert base_date.day == 1, (base_date, valid_date)
640
-
641
- step = (valid_date - base_date).total_seconds()
642
- assert int(step) == step, (valid_date, base_date, step)
643
- assert int(step) % 3600 == 0, (valid_date, base_date, step)
644
- step = int(step // 3600)
645
- start, end = cls._adjust_steps(step, step, frequency, accumulations_reset_frequency)
646
-
647
- return start + frequency, end
648
-
649
- def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
650
- """Computes the accumulation from the last step.
651
-
652
- Parameters
653
- ----------
654
- values : NDArray[Any]
655
- The values to accumulate.
656
- startStep : int
657
- The start step.
658
- endStep : int
659
- The end step.
660
- """
661
-
662
- assert self.frequency == 1
663
-
664
- assert startStep % self.accumulations_reset_frequency == 0, (
665
- startStep,
666
- endStep,
667
- self.accumulations_reset_frequency,
668
- )
669
-
670
- if self.values is None:
671
-
672
- self.values = np.copy(values)
673
- self.startStep = startStep
674
- self.endStep = endStep
675
-
676
- if len(self.steps) == 1:
677
- assert self.startStep == self.endStep - self.frequency, (self.startStep, self.endStep)
678
-
679
- else:
680
- assert endStep != self.endStep, (self.endStep, endStep)
681
-
682
- if endStep > self.endStep:
683
- # assert endStep - self.endStep == self.stepping, (self.endStep, endStep, self.stepping)
684
- self.values = values - self.values
685
- self.startStep = self.endStep
686
- self.endStep = endStep
687
- else:
688
- # assert self.endStep - endStep == self.stepping, (self.endStep, endStep, self.stepping)
689
- self.values = self.values - values
690
- self.startStep = endStep
691
-
692
- assert self.endStep - self.startStep <= self.accumulations_reset_frequency, (self.startStep, startStep)
693
-
694
- @classmethod
695
- def _mars_date_time_step(
696
- cls,
697
- *,
698
- base_date: datetime.datetime,
699
- step1: int,
700
- step2: int,
701
- add_step: int,
702
- frequency: int,
703
- accumulations_reset_frequency: int | None,
704
- user_date: str | None,
705
- requested_date: datetime.datetime | None = None,
706
- ) -> tuple[int, int, tuple[int, ...]]:
707
- """Generates a MARS date-time step.
708
-
709
- Parameters
710
- ----------
711
- base_date : datetime.datetime
712
- The base date.
713
- step1 : int
714
- First step.
715
- step2 : int
716
- Second step.
717
- add_step : int
718
- Additional step.
719
- frequency : int
720
- Frequency of accumulation.
721
- accumulations_reset_frequency : Optional[int]
722
- Frequency at which accumulations reset.
723
- user_date : Optional[str]
724
- User-defined date.
725
- requested_date : Optional[datetime.datetime], optional
726
- Requested date. Defaults to None.
727
-
728
- Returns
729
- -------
730
- Tuple[int, int, Tuple[int, ...]]
731
- A tuple representing the MARS date-time step.
732
- """
733
- # assert frequency > 0, frequency
734
- # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
735
-
736
- step1 += add_step
737
- step2 += add_step
738
-
739
- assert step2 - step1 == frequency, (step1, step2, frequency)
740
-
741
- adjust_step1 = cls._adjust_steps(step1, step1, frequency, accumulations_reset_frequency)
742
- adjust_step2 = cls._adjust_steps(step2, step2, frequency, accumulations_reset_frequency)
743
-
744
- if adjust_step1[1] % accumulations_reset_frequency == 0:
745
- # First step of a new accumulation
746
- steps = (adjust_step2[1],)
747
- else:
748
- steps = (adjust_step1[1], adjust_step2[1])
749
-
750
- return (
751
- base_date.year * 10000 + base_date.month * 100 + base_date.day,
752
- base_date.hour * 100 + base_date.minute,
753
- tuple(steps),
754
- )
755
-
756
-
757
- def _identity(x: Any) -> Any:
758
- """Identity function that returns the input as is.
759
-
760
- Parameters
761
- ----------
762
- x : Any
763
- Input value.
764
-
765
- Returns
766
- -------
767
- Any
768
- The input value.
769
- """
770
- return x
771
-
772
-
773
- def _compute_accumulations(
774
- context: Any,
775
- dates: list[datetime.datetime],
776
- request: dict[str, Any],
777
- user_accumulation_period: int | tuple[int, int] = 6,
778
- data_accumulation_period: int | None = None,
779
- accumulations_reset_frequency: int | None = None,
780
- user_date: str | None = None,
781
- patch: Any = _identity,
782
- base_times: list[int] | None = None,
783
- use_cdsapi_dataset: str | None = None,
784
- ) -> Any:
785
- """Computes accumulations based on the provided parameters.
786
-
787
- Parameters
788
- ----------
789
- context : Any
790
- Context for the computation.
791
- dates : List[datetime.datetime]
792
- List of dates.
793
- request : Dict[str, Any]
794
- Request parameters.
795
- user_accumulation_period : Union[int, Tuple[int, int]], optional
796
- User-defined accumulation period. Defaults to 6.
797
- data_accumulation_period : Optional[int], optional
798
- Data accumulation period. Defaults to None.
799
- accumulations_reset_frequency : Optional[int], optional
800
- Frequency at which accumulations reset. Defaults to None.
801
- user_date : Optional[str], optional
802
- User-defined date. Defaults to None.
803
- patch : Any, optional
804
- Patch function. Defaults to _identity.
805
- base_times : Optional[List[int]], optional
806
- List of base times. Defaults to None.
807
- use_cdsapi_dataset : Optional[str], optional
808
- CDSAPI dataset to use. Defaults to None.
809
-
810
- Returns
811
- -------
812
- Any
813
- The computed accumulations.
814
- """
815
- adjust_step = isinstance(user_accumulation_period, int)
816
-
817
- if not isinstance(user_accumulation_period, (list, tuple)):
818
- user_accumulation_period = (0, user_accumulation_period)
819
-
820
- assert len(user_accumulation_period) == 2, user_accumulation_period
821
- step1, step2 = user_accumulation_period
822
- assert step1 < step2, user_accumulation_period
823
-
824
- if accumulations_reset_frequency is not None:
825
- AccumulationClass = AccumulationFromLastReset
826
- else:
827
- AccumulationClass = AccumulationFromStart if data_accumulation_period in (0, None) else AccumulationFromLastStep
828
-
829
- if data_accumulation_period is None:
830
- data_accumulation_period = user_accumulation_period[1] - user_accumulation_period[0]
831
-
832
- if base_times is None:
833
- if "time" in request:
834
- time = request.pop("time")
835
- if time > 100:
836
- time = time // 100
837
- base_times = [time]
838
- else:
839
- base_times = [0, 6, 12, 18]
840
-
841
- base_times = [t // 100 if t > 100 else t for t in base_times]
842
-
843
- mars_date_time_steps = AccumulationClass.mars_date_time_steps(
844
- dates=dates,
845
- step1=step1,
846
- step2=step2,
847
- frequency=data_accumulation_period,
848
- base_times=base_times,
849
- adjust_step=adjust_step,
850
- accumulations_reset_frequency=accumulations_reset_frequency,
851
- user_date=user_date,
852
- )
853
-
854
- request = deepcopy(request)
855
-
856
- param = request["param"]
857
- if not isinstance(param, (list, tuple)):
858
- param = [param]
859
-
860
- number = request.get("number", [0])
861
- assert isinstance(number, (list, tuple))
862
-
863
- frequency = data_accumulation_period
864
-
865
- type_ = request.get("type", "an")
866
- if type_ == "an":
867
- type_ = "fc"
868
-
869
- request.update({"type": type_, "levtype": "sfc"})
870
-
871
- tmp = temp_file()
872
- path = tmp.path
873
- out = new_grib_output(path)
874
-
875
- requests = []
876
-
877
- accumulations = {}
878
-
879
- for date, time, steps in mars_date_time_steps:
880
- for p in param:
881
- for n in number:
882
- r = dict(request, param=p, date=date, time=time, step=sorted(steps), number=n)
883
-
884
- requests.append(patch(r))
885
-
886
- ds = mars(
887
- context, dates, *requests, request_already_using_valid_datetime=True, use_cdsapi_dataset=use_cdsapi_dataset
888
- )
889
-
890
- accumulations = {}
891
- for a in [
892
- AccumulationClass(out, frequency=frequency, accumulations_reset_frequency=accumulations_reset_frequency, **r)
893
- for r in requests
894
- ]:
895
- for s in a.steps:
896
- key = (a.param, a.date, a.time, s, a.number)
897
- accumulations.setdefault(key, []).append(a)
898
-
899
- for field in ds:
900
- key = (
901
- field.metadata("param"),
902
- field.metadata("date"),
903
- field.metadata("time"),
904
- field.metadata("step"),
905
- _member(field),
906
- )
907
- values = field.values # optimisation
908
- if key not in accumulations:
909
- raise ValueError(f"Key not found: {key}. Is it an accumulation field?")
910
-
911
- for a in accumulations[key]:
912
- a.add(field, values)
913
-
914
- for acc in accumulations.values():
915
- for a in acc:
916
- assert a.done, (a.key, a.seen, a.steps)
917
-
918
- out.close()
919
-
920
- ds = ekd.from_source("file", path)
921
-
922
- assert len(ds) / len(param) / len(number) == len(dates), (
923
- len(ds),
924
- len(param),
925
- len(dates),
926
- )
927
- ds._tmp = tmp
928
-
929
- return ds
930
-
931
-
932
- def _to_list(x: list[Any] | tuple[Any] | Any) -> list[Any]:
933
- """Converts the input to a list if it is not already a list or tuple.
934
-
935
- Parameters
936
- ----------
937
- x : Union[List[Any], Tuple[Any], Any]
938
- Input value.
939
-
940
- Returns
941
- -------
942
- List[Any]
943
- The input value as a list.
944
- """
945
- if isinstance(x, (list, tuple)):
946
- return x
947
- return [x]
948
-
949
-
950
- def _scda(request: dict[str, Any]) -> dict[str, Any]:
951
- """Modifies the request stream based on the time.
952
-
953
- Parameters
954
- ----------
955
- request : Dict[str, Any]
956
- Request parameters.
957
-
958
- Returns
959
- -------
960
- Dict[str, Any]
961
- The modified request parameters.
962
- """
963
- if request["time"] in (6, 18, 600, 1800):
964
- request["stream"] = "scda"
965
- else:
966
- request["stream"] = "oper"
967
- return request
968
-
969
-
970
- @source_registry.register("accumulations")
971
- class AccumulationsSource(LegacySource):
972
-
973
- @staticmethod
974
- def _execute(
975
- context: Any, dates: list[datetime.datetime], use_cdsapi_dataset: str | None = None, **request: Any
976
- ) -> Any:
977
- """Computes accumulations based on the provided context, dates, and request parameters.
978
-
979
- Parameters
980
- ----------
981
- context : Any
982
- Context for the computation.
983
- dates : List[datetime.datetime]
984
- List of dates.
985
- use_cdsapi_dataset : Optional[str], optional
986
- CDSAPI dataset to use. Defaults to None.
987
- **request : Any
988
- Additional request parameters.
989
-
990
- Returns
991
- -------
992
- Any
993
- The computed accumulations.
994
- """
995
-
996
- if (
997
- request.get("class") == "ea"
998
- and request.get("stream", "oper") == "oper"
999
- and request.get("accumulation_period") == 24
1000
- ):
1001
- from .accumulations2 import Accumulations2Source
1002
-
1003
- LOG.warning(
1004
- "🧪️ Experimental features: Using accumulations2, because class=ea stream=oper and accumulation_period=24"
1005
- )
1006
- return Accumulations2Source._execute(context, dates, **request)
1007
-
1008
- _to_list(request["param"])
1009
- class_ = request.get("class", "od")
1010
- stream = request.get("stream", "oper")
1011
-
1012
- user_accumulation_period = request.pop("accumulation_period", 6)
1013
- accumulations_reset_frequency = request.pop("accumulations_reset_frequency", None)
1014
- user_date = request.pop("date", None)
1015
-
1016
- # If `data_accumulation_period` is not set, this means that the accumulations are from the start
1017
- # of the forecast.
1018
-
1019
- KWARGS = {
1020
- ("od", "oper"): dict(patch=_scda),
1021
- ("od", "elda"): dict(base_times=(6, 18)),
1022
- ("od", "enfo"): dict(base_times=(0, 6, 12, 18)),
1023
- ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
1024
- ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
1025
- ("rr", "oper"): dict(base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
1026
- ("l5", "oper"): dict(data_accumulation_period=1, base_times=(0,)),
1027
- }
1028
-
1029
- kwargs = KWARGS.get((class_, stream), {})
1030
-
1031
- context.trace("🌧️", f"accumulations {request} {user_accumulation_period} {kwargs}")
1032
-
1033
- return _compute_accumulations(
1034
- context,
1035
- dates,
1036
- request,
1037
- user_accumulation_period=user_accumulation_period,
1038
- accumulations_reset_frequency=accumulations_reset_frequency,
1039
- use_cdsapi_dataset=use_cdsapi_dataset,
1040
- user_date=user_date,
1041
- **kwargs,
1042
- )