diff --git a/src/anemoi/inference/outputs/apply_mask.py b/src/anemoi/inference/outputs/apply_mask.py deleted file mode 100644 index 317723e5..00000000 --- a/src/anemoi/inference/outputs/apply_mask.py +++ /dev/null @@ -1,60 +0,0 @@ -# (C) Copyright 2024 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - -import logging - -from anemoi.inference.config import Configuration -from anemoi.inference.context import Context -from anemoi.inference.types import ProcessorConfig - -from . import output_registry -from .masked import MaskedOutput - -LOG = logging.getLogger(__name__) - - -@output_registry.register("apply_mask") -class ApplyMaskOutput(MaskedOutput): - """Apply mask output class.""" - - def __init__( - self, - context: Context, - *, - mask: str, - output: Configuration, - variables: list[str] | None = None, - post_processors: list[ProcessorConfig] | None = None, - output_frequency: int | None = None, - write_initial_state: bool | None = None, - ) -> None: - """Parameters - ---------- - context : dict - The context dictionary. - mask : str - The mask identifier. - output : dict - The output configuration dictionary. - post_processors : Optional[List[ProcessorConfig]], default None - Post-processors to apply to the input - output_frequency : int, optional - The frequency of output, by default None. - write_initial_state : bool, optional - Whether to write the initial state, by default None. - """ - super().__init__( - context, - mask=context.checkpoint.load_supporting_array(mask), - output=output, - variables=variables, - post_processors=post_processors, - output_frequency=output_frequency, - write_initial_state=write_initial_state, - ) diff --git a/src/anemoi/inference/outputs/assign_mask.py b/src/anemoi/inference/outputs/assign_mask.py deleted file mode 100644 index c2b0e64a..00000000 --- a/src/anemoi/inference/outputs/assign_mask.py +++ /dev/null @@ -1,114 +0,0 @@ -# (C) Copyright 2025 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - -import logging -import warnings - -import numpy as np - -from anemoi.inference.config import Configuration -from anemoi.inference.context import Context -from anemoi.inference.types import ProcessorConfig -from anemoi.inference.types import State - -from ..output import ForwardOutput -from . import output_registry - -LOG = logging.getLogger(__name__) - - -@output_registry.register("assign_mask") -class AssignMask(ForwardOutput): - """Assign the current output to larger area using a mask. - - This operation can be seen as the opposite of "extract_mask". - Instead of extracting a smaller area from a larger one, - it assigns the current output to a larger area using a mask. - This is useful when you want to restore the original state of the model - after applying a mask to it. The portion of the state that is not - covered by the mask will be set to a fill value (NaN by default). - - Parameters - ---------- - context : dict - The context dictionary. - output : dict - The output configuration dictionary. - mask : str - The mask supporting array name. - fill_value : float, optional - The fill value to use for the masked area, by default np.nan. - post_processors : Optional[List[ProcessorConfig]], default None - Post-processors to apply to the input - output_frequency : int, optional - The frequency of output, by default None. - write_initial_state : bool, optional - Whether to write the initial state, by default None. - """ - - def __init__( - self, - context: Context, - *, - output: Configuration, - mask: str, - fill_value: float = np.nan, - post_processors: list[ProcessorConfig] | None = None, - output_frequency: int | None = None, - write_initial_state: bool | None = None, - ) -> None: - super().__init__( - context, output, post_processors, output_frequency=output_frequency, write_initial_state=write_initial_state - ) - warnings.warn( - "The AssignMask output is deprecated and will be removed in a future release. " - "Use the `assign_mask` post-processor instead.", - DeprecationWarning, - ) - - if mask not in self.checkpoint.supporting_arrays: - raise ValueError(f"Assignment mask '{mask}' not found in supporting arrays.") - - self.mask = self.checkpoint.load_supporting_array(mask) - self.fill_value = fill_value - - def __repr__(self) -> str: - """Return a string representation of the ExtractLamOutput object.""" - return f"AssignMask({self.output})" - - def modify_state(self, state: State) -> State: - """Assign the state to the mask. - - Parameters - ---------- - state : State - The state dictionary. - - Returns - ------- - State - The masked state dictionary. - """ - state = state.copy() - - state["latitudes"] = self._assign_mask(state["latitudes"]) - state["longitudes"] = self._assign_mask(state["longitudes"]) - for field in state["fields"]: - state["fields"][field] = self._assign_mask(state["fields"][field]) - - return state - - def _assign_mask(self, array: np.ndarray): - shape = array.shape[:-1] + self.mask.shape - res = np.full(shape, self.fill_value, dtype=array.dtype) - if array.ndim == 1: - res[self.mask] = array - else: - res[..., self.mask] = array - return res diff --git a/src/anemoi/inference/outputs/extract_lam.py b/src/anemoi/inference/outputs/extract_lam.py deleted file mode 100644 index a5169ade..00000000 --- a/src/anemoi/inference/outputs/extract_lam.py +++ /dev/null @@ -1,81 +0,0 @@ -# (C) Copyright 2024 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - -import logging - -import numpy as np - -from anemoi.inference.config import Configuration -from anemoi.inference.context import Context -from anemoi.inference.types import ProcessorConfig - -from . import output_registry -from .masked import MaskedOutput - -LOG = logging.getLogger(__name__) - - -@output_registry.register("extract_lam") -class ExtractLamOutput(MaskedOutput): - """Extract LAM output class.""" - - def __init__( - self, - context: Context, - *, - output: Configuration, - lam: str = "lam_0", - variables: list[str] | None = None, - post_processors: list[ProcessorConfig] | None = None, - output_frequency: int | None = None, - write_initial_state: bool | None = None, - ) -> None: - """Parameters - ---------- - context : dict - The context dictionary. - output : dict - The output configuration dictionary. - lam : str, optional - The LAM identifier, by default "lam_0". - variables : list, optional - The list of variables to extract, by default None. - post_processors : Optional[List[ProcessorConfig]], default None - Post-processors to apply to the input - output_frequency : int, optional - The frequency of output, by default None. - write_initial_state : bool, optional - Whether to write the initial state, by default None. - """ - - if "cutout_mask" in context.checkpoint.supporting_arrays: - # Backwards compatibility - mask = context.checkpoint.load_supporting_array("cutout_mask") - points = slice(None, -np.sum(mask)) - else: - if "lam_0" not in lam: - raise NotImplementedError("Only lam_0 is supported") - - if "lam_1/cutout_mask" in context.checkpoint.supporting_arrays: - raise NotImplementedError("Only lam_0 is supported") - - mask = context.checkpoint.load_supporting_array(f"{lam}/cutout_mask") - - assert len(mask) == np.sum(mask) - points = slice(None, np.sum(mask)) - - super().__init__( - context, - mask=points, - output=output, - variables=variables, - post_processors=post_processors, - output_frequency=output_frequency, - write_initial_state=write_initial_state, - ) diff --git a/src/anemoi/inference/outputs/masked.py b/src/anemoi/inference/outputs/masked.py deleted file mode 100644 index f08efaa5..00000000 --- a/src/anemoi/inference/outputs/masked.py +++ /dev/null @@ -1,102 +0,0 @@ -# (C) Copyright 2024 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - -import logging -import warnings -from typing import Any - -from anemoi.inference.config import Configuration -from anemoi.inference.context import Context -from anemoi.inference.types import ProcessorConfig -from anemoi.inference.types import State - -from ..output import ForwardOutput - -LOG = logging.getLogger(__name__) - - -class MaskedOutput(ForwardOutput): - """Apply mask output class. - - Parameters - ---------- - context : dict - The context dictionary. - mask : Any - The mask. - output : dict - The output configuration dictionary. - variables : list, optional - The list of variables to extract, by default None. - post_processors : Optional[List[ProcessorConfig]], default None - Post-processors to apply to the input - output_frequency : int, optional - The frequency of output, by default None. - write_initial_state : bool, optional - Whether to write the initial state, by default None. - """ - - def __init__( - self, - context: Context, - *, - mask: Any, - output: Configuration, - variables: list[str] | None = None, - post_processors: list[ProcessorConfig] | None = None, - output_frequency: int | None = None, - write_initial_state: bool | None = None, - ) -> None: - - warnings.warn( - f"The {self.__class__.__name__} is deprecated and will be removed in a future release. " - "Use extraction post-processors instead.", - DeprecationWarning, - ) - super().__init__( - context, - output, - variables=variables, - post_processors=post_processors, - output_frequency=output_frequency, - write_initial_state=write_initial_state, - ) - self.mask = mask - - def modify_state(self, state: State) -> State: - """Apply the mask to the state. - - Parameters - ---------- - state : State - The state dictionary. - - Returns - ------- - State - The masked state dictionary. - """ - state = state.copy() - state["fields"] = state["fields"].copy() - state["latitudes"] = state["latitudes"][self.mask] - state["longitudes"] = state["longitudes"][self.mask] - - for field in state["fields"]: - data = state["fields"][field] - if data.ndim == 1: - data = data[self.mask] - else: - data = data[..., self.mask] - state["fields"][field] = data - - return state - - def __repr__(self) -> str: - """Return a string representation of the object.""" - return f"{self.__class__.__name__}({self.mask}, {self.output})" diff --git a/src/anemoi/inference/runners/crps.py b/src/anemoi/inference/runners/crps.py deleted file mode 100644 index 506f09c6..00000000 --- a/src/anemoi/inference/runners/crps.py +++ /dev/null @@ -1,45 +0,0 @@ -# (C) Copyright 2025 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - -import logging -import warnings -from typing import Any - -from . import runner_registry -from .default import DefaultRunner - -LOG = logging.getLogger(__name__) - - -@runner_registry.register("crps") -class CrpsRunner(DefaultRunner): - """Runner for CRPS (Continuous Ranked Probability Score). - - Inherits from DefaultRunner. - """ - - def predict_step(self, model: Any, input_tensor_torch: Any, **kwargs: Any) -> Any: - """Perform a prediction step using the model. - - Parameters - ---------- - model : Any - The model to use for prediction. - input_tensor_torch : torch.Tensor - The input tensor for the model. - **kwargs : Any - Additional keyword arguments. - - Returns - ------- - Any - The prediction result. - """ - warnings.warn("CRPS runner is deprecated, use DefaultRunner instead") - return model.predict_step(input_tensor_torch, kwargs["fcstep"]) diff --git a/src/anemoi/inference/runners/cutout.py b/src/anemoi/inference/runners/cutout.py deleted file mode 100644 index 282fbf28..00000000 --- a/src/anemoi/inference/runners/cutout.py +++ /dev/null @@ -1,68 +0,0 @@ -# (C) Copyright 2024 Anemoi contributors. -# -# This software is licensed under the terms of the Apache Licence Version 2.0 -# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. - - -import logging -from typing import Any - -from ..context import Context -from ..input import Input -from ..runner import Runner -from . import runner_registry - -LOG = logging.getLogger(__name__) - - -class CutoutInput(Input): - """An Input object that combines two inputs.""" - - def __init__(self, context: Context, lam: Input, globe: Input) -> None: - """Initialize CutoutInput. - - Parameters - ---------- - context : Context - The context for the input. - lam : Input - The LAM input. - globe : Input - The globe input. - """ - super().__init__(context) - self.lam = lam - self.globe = globe - - def create_input_state(self, *, date: str | None = None) -> None: - """Create the input state. - - Parameters - ---------- - date : Optional[str], optional - The date for the input state, by default None - """ - state1 = self.lam.create_input_state(date=date) - state2 = self.globe.create_input_state(date=date) - - assert False, (state1, state2) - - -@runner_registry.register("cutout") -class CutoutRunner(Runner): - """A Runner that for LAMs.""" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - """Initialize CutoutRunner.""" - super().__init__(*args, **kwargs) - - sources = self.checkpoint.sources - - if len(sources) != 2: - raise ValueError(f"CutoutRunner expects two source, found {len(sources)}.") - - self.lam, self.globe = sources