docstring + log erros

src/anemoi/graphs/edges/__init__.py

@@ -0,0 +1,4 @@
+from .connections import CutOffEdgeBuilder
+from .connections import KNNEdgeBuilder
+
+__all__ = ["KNNEdgeBuilder", "CutOffEdgeBuilder"]

src/anemoi/graphs/nodes/__init__.py

@@ -0,0 +1,4 @@
+from .nodes import NPZNodes
+from .nodes import ZarrNodes
+
+__all__ = ["ZarrNodes", "NPZNodes"]

src/anemoi/graphs/nodes/nodes.py

@@ -11,30 +11,10 @@
 from torch_geometric.data import HeteroData
 
 logger = logging.getLogger(__name__)
-earth_radius = 6371.0  # km
-
-
-def latlon_to_radians(coords: np.ndarray) -> np.ndarray:
-    return np.deg2rad(coords)
-
-
-def rad_to_latlon(coords: np.ndarray) -> np.ndarray:
-    """Converts coordinates from radians to degrees.
-
-    Parameters
-    ----------
-    coords : np.ndarray
-        Coordinates in radians.
-
-    Returns
-    -------
-    np.ndarray
-        _description_
-    """
-    return np.rad2deg(coords)
 
 
 class BaseNodeBuilder(ABC):
+    """Base class for node builders."""
 
     def register_nodes(self, graph: HeteroData, name: str) -> None:
         graph[name].x = self.get_coordinates()
@@ -52,7 +32,6 @@ def get_coordinates(self) -> np.ndarray: ...
     def reshape_coords(self, latitudes: np.ndarray, longitudes: np.ndarray) -> np.ndarray:
         coords = np.stack([latitudes, longitudes], axis=-1).reshape((-1, 2))
         coords = np.deg2rad(coords)
-        # TODO: type needs to be variable?
         return torch.tensor(coords, dtype=torch.float32)
 
     def transform(self, graph: HeteroData, name: str, attr_config: DotDict) -> HeteroData:

src/anemoi/graphs/nodes/weights.py

@@ -23,18 +23,19 @@ def __init__(self, norm: Optional[str] = None):
     @abstractmethod
     def compute(self, nodes: NodeStorage, *args, **kwargs): ...
 
-    def get_weights(self, *args, **kwargs):
+    def get_weights(self, *args, **kwargs) -> torch.Tensor:
         weights = self.compute(*args, **kwargs)
         if weights.ndim == 1:
             weights = weights[:, np.newaxis]
-        return self.normalize(weights)
+        norm_weights = self.normalize(weights)
+        return torch.tensor(norm_weights, dtype=torch.float32)
 
 
 class UniformWeights(BaseWeights):
     """Implements a uniform weight for the nodes."""
 
     def compute(self, nodes: NodeStorage) -> np.ndarray:
-        return torch.ones(nodes.num_nodes)
+        return np.ones(nodes.num_nodes)
 
 
 class AreaWeights(BaseWeights):

src/anemoi/graphs/normalizer.py

@@ -19,5 +19,11 @@ def normalize(self, values: np.ndarray) -> np.ndarray:
         if self.norm == "unit-sum":
             return values / np.sum(values)
         if self.norm == "unit-std":
-            return values / np.std(values)
-        raise ValueError("Weight normalization must be 'l1', 'l2', 'unit-max' 'unit-sum' or 'unit-std'.")
+            std = np.std(values)
+            if std == 0:
+                logger.warning(f"Std. dev. of the {self.__class__.__name__} is 0. Cannot normalize.")
+                return values
+            return values / std
+        raise ValueError(
+            f"Weight normalization \"{values}\" is not valid. Options are: 'l1', 'l2', 'unit-max' 'unit-sum' or 'unit-std'."
+        )