anemoi-datasets 0.5.20__py3-none-any.whl → 0.5.22__py3-none-any.whl

anemoi/datasets/create/sources/accumulations.py

@@ -68,9 +68,11 @@ class Accumulation:
         number: int,
         step: List[int],
         frequency: int,
+        accumulations_reset_frequency: Optional[int] = None,
+        user_date: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        """Initializes an Accumulation instance.
+        """Initialises an Accumulation instance.
 
         Parameters
         ----------
@@ -88,6 +90,10 @@ class Accumulation:
             List of steps.
         frequency : int
             Frequency of accumulation.
+        accumulations_reset_frequency : Optional[int], optional
+            Frequency at which accumulations reset. Defaults to None.
+        user_date : Optional[str], optional
+            User-defined date. Defaults to None.
         **kwargs : Any
             Additional keyword arguments.
         """
@@ -103,7 +109,9 @@ class Accumulation:
         self.endStep: Optional[int] = None
         self.done = False
         self.frequency = frequency
+        self.accumulations_reset_frequency = accumulations_reset_frequency
         self._check = None
+        self.user_date = user_date
 
     @property
     def key(self) -> Tuple[str, int, int, List[int], int]:
@@ -161,14 +169,28 @@ class Accumulation:
                 f"Negative values when computing accumutation for {self.param} ({self.date} {self.time}): min={np.amin(self.values)} max={np.amax(self.values)}"
             )
 
-        self.out.write(
-            self.values,
-            template=template,
-            stepType="accum",
-            startStep=self.startStep,
-            endStep=self.endStep,
-            check_nans=True,
-        )
+        # In GRIB1, is the step is greater that 254 (one byte), we cannot use a range, because both P1 and P2 values
+        # are used to store the end step
+
+        edition = template.metadata("edition")
+
+        if edition == 1 and self.endStep > 254:
+            self.out.write(
+                self.values,
+                template=template,
+                stepType="instant",
+                step=self.endStep,
+                check_nans=True,
+            )
+        else:
+            self.out.write(
+                self.values,
+                template=template,
+                stepType="accum",
+                startStep=self.startStep,
+                endStep=self.endStep,
+                check_nans=True,
+            )
         self.values = None
         self.done = True
 
@@ -179,7 +201,7 @@ class Accumulation:
         ----------
         field : Any
             The field containing the values.
-        values : np.ndarray
+        values : NDArray[Any]
             The values to add.
         """
         self.check(field)
@@ -189,7 +211,7 @@ class Accumulation:
             return
 
         if not np.all(values >= 0):
-            warnings.warn(f"Negative values for {field}: {np.amin(values)} {np.amax(values)}")
+            warnings.warn(f"Negative values for {field}: {np.nanmin(values)} {np.nanmax(values)}")
 
         assert not self.done, (self.key, step)
         assert step not in self.seen, (self.key, step)
@@ -197,8 +219,8 @@ class Accumulation:
         startStep = field.metadata("startStep")
         endStep = field.metadata("endStep")
 
-        if self.buggy_steps and startStep == endStep:
-            startStep = 0
+        if startStep == endStep:
+            startStep, endStep = self.adjust_steps(startStep, endStep)
 
         assert step == endStep, (startStep, endStep, step)
 
@@ -212,12 +234,15 @@ class Accumulation:
     @classmethod
     def mars_date_time_steps(
         cls,
+        *,
         dates: List[datetime.datetime],
         step1: int,
         step2: int,
         frequency: Optional[int],
         base_times: List[int],
         adjust_step: bool,
+        accumulations_reset_frequency: Optional[int],
+        user_date: Optional[str],
     ) -> Generator[Tuple[int, int, Tuple[int, ...]], None, None]:
         """Generates MARS date-time steps.
 
@@ -235,6 +260,10 @@ class Accumulation:
             List of base times.
         adjust_step : bool
             Whether to adjust the step.
+        accumulations_reset_frequency : Optional[int], optional
+            Frequency at which accumulations reset. Defaults to None.
+        user_date : Optional[str], optional
+            User-defined date. Defaults to None.
 
         Returns
         -------
@@ -244,8 +273,16 @@ class Accumulation:
         # assert step1 > 0, (step1, step2, frequency)
 
         for valid_date in dates:
-            base_date = valid_date - datetime.timedelta(hours=step2)
             add_step = 0
+            base_date = valid_date - datetime.timedelta(hours=step2)
+            if user_date is not None:
+                assert user_date == "????-??-01", user_date
+                new_base_date = base_date.replace(day=1)
+                assert new_base_date <= base_date, (new_base_date, base_date)
+                add_step = int((base_date - new_base_date).total_seconds() // 3600)
+
+                base_date = new_base_date
+
             if base_date.hour not in base_times:
                 if not adjust_step:
                     raise ValueError(
@@ -257,23 +294,93 @@ class Accumulation:
                     base_date -= datetime.timedelta(hours=1)
                     add_step += 1
 
-            yield cls._mars_date_time_step(base_date, step1, step2, add_step, frequency)
+            yield cls._mars_date_time_step(
+                base_date=base_date,
+                step1=step1,
+                step2=step2,
+                add_step=add_step,
+                frequency=frequency,
+                accumulations_reset_frequency=accumulations_reset_frequency,
+                user_date=user_date,
+                requested_date=valid_date,
+            )
+
+    def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
+        """Computes the accumulation.
 
-    def __repr__(self) -> str:
-        """Returns a string representation of the Accumulation instance.
+        Parameters
+        ----------
+        values : NDArray[Any]
+            The values to accumulate.
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+        """
+        pass
+
+    @classmethod
+    def _mars_date_time_step(
+        cls,
+        *,
+        base_date: datetime.datetime,
+        step1: int,
+        step2: int,
+        add_step: int,
+        frequency: Optional[int],
+        accumulations_reset_frequency: Optional[int],
+        user_date: Optional[str],
+        requested_date: Optional[datetime.datetime] = None,
+    ) -> Tuple[int, int, Tuple[int, ...]]:
+        """Generates a MARS date-time step.
+
+        Parameters
+        ----------
+        base_date : datetime.datetime
+            The base date.
+        step1 : int
+            First step.
+        step2 : int
+            Second step.
+        add_step : int
+            Additional step.
+        frequency : Optional[int]
+            Frequency of accumulation.
+        accumulations_reset_frequency : Optional[int], optional
+            Frequency at which accumulations reset. Defaults to None.
+        user_date : Optional[str], optional
+            User-defined date. Defaults to None.
+        requested_date : Optional[datetime.datetime], optional
+            Requested date. Defaults to None.
 
         Returns
         -------
-        str
-            String representation of the Accumulation instance.
+        Tuple[int, int, Tuple[int, ...]]
+            A tuple representing the MARS date-time step.
         """
-        return f"{self.__class__.__name__}({self.key})"
+        pass
 
 
 class AccumulationFromStart(Accumulation):
     """Class to handle data accumulation from the start of the forecast."""
 
-    buggy_steps = True
+    def adjust_steps(self, startStep: int, endStep: int) -> Tuple[int, int]:
+        """Adjusts the start and end steps.
+
+        Parameters
+        ----------
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+
+        Returns
+        -------
+        Tuple[int, int]
+            The adjusted start and end steps.
+        """
+        assert endStep == startStep
+        return (0, endStep)
 
     def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
         """Computes the accumulation from the start.
@@ -314,7 +421,16 @@ class AccumulationFromStart(Accumulation):
 
     @classmethod
     def _mars_date_time_step(
-        cls, base_date: datetime.datetime, step1: int, step2: int, add_step: int, frequency: Optional[int]
+        cls,
+        *,
+        base_date: datetime.datetime,
+        step1: int,
+        step2: int,
+        add_step: int,
+        frequency: Optional[int],
+        accumulations_reset_frequency: Optional[int],
+        user_date: Optional[str],
+        requested_date: Optional[datetime.datetime] = None,
     ) -> Tuple[int, int, Tuple[int, ...]]:
         """Generates a MARS date-time step.
 
@@ -330,12 +446,19 @@ class AccumulationFromStart(Accumulation):
             Additional step.
         frequency : Optional[int]
             Frequency of accumulation.
+        accumulations_reset_frequency : Optional[int], optional
+            Frequency at which accumulations reset. Defaults to None.
+        user_date : Optional[str], optional
+            User-defined date. Defaults to None.
+        requested_date : Optional[datetime.datetime], optional
+            Requested date. Defaults to None.
 
         Returns
         -------
         Tuple[int, int, Tuple[int, ...]]
             A tuple representing the MARS date-time step.
         """
+        assert user_date is None, user_date
         assert not frequency, frequency
 
         steps = (step1 + add_step, step2 + add_step)
@@ -352,8 +475,6 @@ class AccumulationFromStart(Accumulation):
 class AccumulationFromLastStep(Accumulation):
     """Class to handle data accumulation from the last step of the forecast."""
 
-    buggy_steps = False
-
     def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
         """Computes the accumulation from the last step.
 
@@ -389,7 +510,16 @@ class AccumulationFromLastStep(Accumulation):
 
     @classmethod
     def _mars_date_time_step(
-        cls, base_date: datetime.datetime, step1: int, step2: int, add_step: int, frequency: int
+        cls,
+        *,
+        base_date: datetime.datetime,
+        step1: int,
+        step2: int,
+        add_step: int,
+        frequency: int,
+        accumulations_reset_frequency: Optional[int],
+        user_date: Optional[str] = None,
+        requested_date: Optional[datetime.datetime] = None,
     ) -> Tuple[int, int, Tuple[int, ...]]:
         """Generates a MARS date-time step.
 
@@ -405,18 +535,221 @@ class AccumulationFromLastStep(Accumulation):
             Additional step.
         frequency : int
             Frequency of accumulation.
+        accumulations_reset_frequency : Optional[int], optional
+            Frequency at which accumulations reset. Defaults to None.
+        user_date : Optional[str], optional
+            User-defined date. Defaults to None.
+        requested_date : Optional[datetime.datetime], optional
+            Requested date. Defaults to None.
 
         Returns
         -------
         Tuple[int, int, Tuple[int, ...]]
             A tuple representing the MARS date-time step.
         """
+
+        assert user_date is None, user_date
+
         assert frequency > 0, frequency
         # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
 
         steps = []
         for step in range(step1 + frequency, step2 + frequency, frequency):
             steps.append(step + add_step)
+
+        return (
+            base_date.year * 10000 + base_date.month * 100 + base_date.day,
+            base_date.hour * 100 + base_date.minute,
+            tuple(steps),
+        )
+
+
+class AccumulationFromLastReset(Accumulation):
+    """Class to handle data accumulation from the last step of the forecast."""
+
+    def adjust_steps(self, startStep: int, endStep: int) -> Tuple[int, int]:
+        """Adjusts the start and end steps.
+
+        Parameters
+        ----------
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+
+        Returns
+        -------
+        Tuple[int, int]
+            The adjusted start and end steps.
+        """
+        return self.__class__._adjust_steps(startStep, endStep, self.frequency, self.accumulations_reset_frequency)
+
+    @classmethod
+    def _adjust_steps(
+        self, startStep: int, endStep: int, frequency: int, accumulations_reset_frequency: int
+    ) -> Tuple[int, int]:
+        """Adjusts the start and end steps.
+
+        Parameters
+        ----------
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+        frequency : int
+            Frequency of accumulation.
+        accumulations_reset_frequency : int
+            Frequency at which accumulations reset.
+
+        Returns
+        -------
+        Tuple[int, int]
+            The adjusted start and end steps.
+        """
+
+        assert frequency == 1, (frequency, startStep, endStep)
+        assert endStep - startStep <= accumulations_reset_frequency, (startStep, endStep)
+
+        return ((startStep // accumulations_reset_frequency) * accumulations_reset_frequency, endStep)
+
+    @classmethod
+    def _steps(
+        cls,
+        valid_date: datetime.datetime,
+        base_date: datetime.datetime,
+        frequency: int,
+        accumulations_reset_frequency: int,
+    ) -> Tuple[int, int]:
+        """Calculates the steps for accumulation.
+
+        Parameters
+        ----------
+        valid_date : datetime.datetime
+            The valid date.
+        base_date : datetime.datetime
+            The base date.
+        frequency : int
+            Frequency of accumulation.
+        accumulations_reset_frequency : int
+            Frequency at which accumulations reset.
+
+        Returns
+        -------
+        Tuple[int, int]
+            A tuple representing the steps for accumulation.
+        """
+
+        assert base_date.day == 1, (base_date, valid_date)
+
+        step = (valid_date - base_date).total_seconds()
+        assert int(step) == step, (valid_date, base_date, step)
+        assert int(step) % 3600 == 0, (valid_date, base_date, step)
+        step = int(step // 3600)
+        start, end = cls._adjust_steps(step, step, frequency, accumulations_reset_frequency)
+
+        return start + frequency, end
+
+    def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
+        """Computes the accumulation from the last step.
+
+        Parameters
+        ----------
+        values : NDArray[Any]
+            The values to accumulate.
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+        """
+
+        assert self.frequency == 1
+
+        assert startStep % self.accumulations_reset_frequency == 0, (
+            startStep,
+            endStep,
+            self.accumulations_reset_frequency,
+        )
+
+        if self.values is None:
+
+            self.values = np.copy(values)
+            self.startStep = startStep
+            self.endStep = endStep
+
+            if len(self.steps) == 1:
+                assert self.startStep == self.endStep - self.frequency, (self.startStep, self.endStep)
+
+        else:
+            assert endStep != self.endStep, (self.endStep, endStep)
+
+            if endStep > self.endStep:
+                # assert endStep - self.endStep == self.stepping, (self.endStep, endStep, self.stepping)
+                self.values = values - self.values
+                self.startStep = self.endStep
+                self.endStep = endStep
+            else:
+                # assert self.endStep - endStep == self.stepping, (self.endStep, endStep, self.stepping)
+                self.values = self.values - values
+                self.startStep = endStep
+
+            assert self.endStep - self.startStep <= self.accumulations_reset_frequency, (self.startStep, startStep)
+
+    @classmethod
+    def _mars_date_time_step(
+        cls,
+        *,
+        base_date: datetime.datetime,
+        step1: int,
+        step2: int,
+        add_step: int,
+        frequency: int,
+        accumulations_reset_frequency: Optional[int],
+        user_date: Optional[str],
+        requested_date: Optional[datetime.datetime] = None,
+    ) -> Tuple[int, int, Tuple[int, ...]]:
+        """Generates a MARS date-time step.
+
+        Parameters
+        ----------
+        base_date : datetime.datetime
+            The base date.
+        step1 : int
+            First step.
+        step2 : int
+            Second step.
+        add_step : int
+            Additional step.
+        frequency : int
+            Frequency of accumulation.
+        accumulations_reset_frequency : Optional[int]
+            Frequency at which accumulations reset.
+        user_date : Optional[str]
+            User-defined date.
+        requested_date : Optional[datetime.datetime], optional
+            Requested date. Defaults to None.
+
+        Returns
+        -------
+        Tuple[int, int, Tuple[int, ...]]
+            A tuple representing the MARS date-time step.
+        """
+        # assert frequency > 0, frequency
+        # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
+
+        step1 += add_step
+        step2 += add_step
+
+        assert step2 - step1 == frequency, (step1, step2, frequency)
+
+        adjust_step1 = cls._adjust_steps(step1, step1, frequency, accumulations_reset_frequency)
+        adjust_step2 = cls._adjust_steps(step2, step2, frequency, accumulations_reset_frequency)
+
+        if adjust_step1[1] % accumulations_reset_frequency == 0:
+            # First step of a new accumulation
+            steps = (adjust_step2[1],)
+        else:
+            steps = (adjust_step1[1], adjust_step2[1])
+
         return (
             base_date.year * 10000 + base_date.month * 100 + base_date.day,
             base_date.hour * 100 + base_date.minute,
@@ -446,6 +779,8 @@ def _compute_accumulations(
     request: Dict[str, Any],
     user_accumulation_period: Union[int, Tuple[int, int]] = 6,
     data_accumulation_period: Optional[int] = None,
+    accumulations_reset_frequency: Optional[int] = None,
+    user_date: Optional[str] = None,
     patch: Any = _identity,
     base_times: Optional[List[int]] = None,
     use_cdsapi_dataset: Optional[str] = None,
@@ -464,6 +799,10 @@ def _compute_accumulations(
         User-defined accumulation period. Defaults to 6.
     data_accumulation_period : Optional[int], optional
         Data accumulation period. Defaults to None.
+    accumulations_reset_frequency : Optional[int], optional
+        Frequency at which accumulations reset. Defaults to None.
+    user_date : Optional[str], optional
+        User-defined date. Defaults to None.
     patch : Any, optional
         Patch function. Defaults to _identity.
     base_times : Optional[List[int]], optional
@@ -485,20 +824,36 @@ def _compute_accumulations(
     step1, step2 = user_accumulation_period
     assert step1 < step2, user_accumulation_period
 
+    if data_accumulation_period is None:
+        data_accumulation_period = user_accumulation_period[1] - user_accumulation_period[0]
+
     if base_times is None:
-        base_times = [0, 6, 12, 18]
+        if "time" in request:
+            time = request.pop("time")
+            if time > 100:
+                time = time // 100
+            base_times = [time]
+        else:
+            base_times = [0, 6, 12, 18]
 
     base_times = [t // 100 if t > 100 else t for t in base_times]
 
-    AccumulationClass = AccumulationFromStart if data_accumulation_period in (0, None) else AccumulationFromLastStep
+    if accumulations_reset_frequency is not None:
+        AccumulationClass = AccumulationFromLastReset
+    else:
+        AccumulationClass = AccumulationFromStart if data_accumulation_period in (0, None) else AccumulationFromLastStep
+
+    LOG.info(f"XXXXXXXXXXX {step1=}, {step2=}, {data_accumulation_period=}, {base_times=}, {adjust_step=}")
 
     mars_date_time_steps = AccumulationClass.mars_date_time_steps(
-        dates,
-        step1,
-        step2,
-        data_accumulation_period,
-        base_times,
-        adjust_step,
+        dates=dates,
+        step1=step1,
+        step2=step2,
+        frequency=data_accumulation_period,
+        base_times=base_times,
+        adjust_step=adjust_step,
+        accumulations_reset_frequency=accumulations_reset_frequency,
+        user_date=user_date,
     )
 
     request = deepcopy(request)
@@ -527,18 +882,25 @@ def _compute_accumulations(
     accumulations = {}
 
     for date, time, steps in mars_date_time_steps:
+        LOG.info(f"Accumulation request: { date, time, steps}")
         for p in param:
             for n in number:
                 r = dict(request, param=p, date=date, time=time, step=sorted(steps), number=n)
 
                 requests.append(patch(r))
 
+    for r in requests:
+        LOG.info(f"Accumulation request: {r}")
+
     ds = mars(
         context, dates, *requests, request_already_using_valid_datetime=True, use_cdsapi_dataset=use_cdsapi_dataset
     )
 
     accumulations = {}
-    for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:
+    for a in [
+        AccumulationClass(out, frequency=frequency, accumulations_reset_frequency=accumulations_reset_frequency, **r)
+        for r in requests
+    ]:
         for s in a.steps:
             key = (a.param, a.date, a.time, s, a.number)
             accumulations.setdefault(key, []).append(a)
@@ -654,6 +1016,8 @@ def accumulations(
     stream = request.get("stream", "oper")
 
     user_accumulation_period = request.pop("accumulation_period", 6)
+    accumulations_reset_frequency = request.pop("accumulations_reset_frequency", None)
+    user_date = request.pop("date", None)
 
     # If `data_accumulation_period` is not set, this means that the accumulations are from the start
     # of the forecast.
@@ -677,7 +1041,9 @@ def accumulations(
         dates,
         request,
         user_accumulation_period=user_accumulation_period,
+        accumulations_reset_frequency=accumulations_reset_frequency,
         use_cdsapi_dataset=use_cdsapi_dataset,
+        user_date=user_date,
         **kwargs,
     )
 

anemoi/datasets/create/sources/forcings.py

@@ -25,7 +25,7 @@ def forcings(context: Any, dates: List[str], template: str, param: str) -> Any:
         The context in which the function is executed.
     dates : list
         List of dates for which data is to be loaded.
-    template : str
+    template : FieldList
         Template for the data source.
     param : str
         Parameter for the data source.