fix: argument order and explicit types

ecmwf · May 28, 2024 · 92ad881 · 92ad881
1 parent 9500a13
commit 92ad881
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/anemoi/models/preprocessing/imputer.py b/src/anemoi/models/preprocessing/imputer.py
@@ -13,6 +13,7 @@
 
 import torch
 
+from anemoi.models.data_indices.collection import IndexCollection
 from anemoi.models.preprocessing import BasePreprocessor
 
 LOGGER = logging.getLogger(__name__)
@@ -24,8 +25,8 @@ class BaseImputer(BasePreprocessor, ABC):
     def __init__(
         self,
         config=None,
+        data_indices: Optional[IndexCollection] = None,
         statistics: Optional[dict] = None,
-        data_indices: Optional[dict] = None,
     ) -> None:
         """Initialize the imputer.
 
@@ -176,7 +177,7 @@ def __init__(
         statistics: Optional[dict] = None,
         data_indices: Optional[dict] = None,
     ) -> None:
-        super().__init__(config, statistics, data_indices)
+        super().__init__(config, data_indices, statistics)
 
         self._create_imputation_indices(statistics)
 
@@ -199,8 +200,10 @@ class ConstantImputer(BaseImputer):
     ```
     """
 
-    def __init__(self, config=None, statistics: Optional[dict] = None, data_indices: Optional[dict] = None) -> None:
-        super().__init__(config, statistics, data_indices)
+    def __init__(
+        self, config=None, statistics: Optional[dict] = None, data_indices: Optional[IndexCollection] = None
+    ) -> None:
+        super().__init__(config, data_indices, statistics)
 
         self._create_imputation_indices()
 

diff --git a/src/anemoi/models/preprocessing/normalizer.py b/src/anemoi/models/preprocessing/normalizer.py
@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 
+from anemoi.models.data_indices.collection import IndexCollection
 from anemoi.models.preprocessing import BasePreprocessor
 
 LOGGER = logging.getLogger(__name__)
@@ -25,7 +26,7 @@ class InputNormalizer(BasePreprocessor):
     def __init__(
         self,
         config=None,
-        data_indices: Optional[dict] = None,
+        data_indices: Optional[IndexCollection] = None,
         statistics: Optional[dict] = None,
     ) -> None:
         """Initialize the normalizer.