feat: add geometry utilities for global domains (spherical coords, area weights, KNN graph)

neural_lam/datastore/base.py

@@ -630,3 +630,34 @@ def num_grid_points(self) -> int:
 
         """
         return self.grid_shape_state.x * self.grid_shape_state.y
+
+def get_area_weights(self) -> "xr.DataArray":
+    """Return latitude-based area weights for global domains.
+
+    Returns:
+        DataArray of shape (grid_index,) with cos(lat) weights normalized
+        to have mean 1.0.
+    """
+    lat_lon = self.get_lat_lon()
+    # Extract latitude values (first element of stacked dim)
+    lat = lat_lon.isel(grid_index=0).lat.values
+
+    from neural_lam.geometry import get_area_weights as compute_weights
+    weights = compute_weights(lat)
+
+    import xarray as xr
+    return xr.DataArray(weights, dims=["grid_index"])
+
+def get_cartesian_coords(self) -> np.ndarray:
+    """Return cartesian coordinates on unit sphere.
+
+    Returns:
+        Array of shape (grid_index, 3) with (x, y, z) coordinates
+        on the unit sphere.
+    """
+    lat_lon = self.get_lat_lon()
+    lon = lat_lon.isel(grid_index=0).lon.values
+    lat = lat_lon.isel(grid_index=0).lat.values
+
+    from neural_lam.geometry import lat_lon_to_cartesian
+    return lat_lon_to_cartesian(lon, lat)

neural_lam/geometry.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+def lat_lon_to_cartesian(lon: np.ndarray, lat: np.ndarray) -> np.ndarray:
+    """Convert (lon, lat) in degrees to (x, y, z) on unit sphere."""
+    lon_rad = np.radians(lon)
+    lat_rad = np.radians(lat)
+    x = np.cos(lat_rad) * np.cos(lon_rad)
+    y = np.cos(lat_rad) * np.sin(lon_rad)
+    z = np.sin(lat_rad)
+    return np.stack([x, y, z], axis=-1)
+
+def get_area_weights(lat: np.ndarray) -> np.ndarray:
+    """Return cos(latitude) weights normalized by mean weight."""
+    weights = np.cos(np.radians(lat))
+    return weights / weights.mean()

neural_lam/graph_builder.py

@@ -0,0 +1,37 @@
+import numpy as np
+from typing import Tuple
+
+def build_spherical_knn_graph(
+    coords: np.ndarray,
+    k: int,
+    backend: str = "scipy",
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Build KNN graph on sphere using cartesian coordinates.
+
+    Args:
+        coords: (N, 3) cartesian coordinates on unit sphere
+        k: Number of neighbors per node
+        backend: "scipy" (more backends to come)
+
+    Returns:
+        edge_index: (2, E) adjacency list
+        edge_attr: (E,) distances
+    """
+    if backend == "scipy":
+        from scipy.spatial import KDTree
+        tree = KDTree(coords)
+        distances, indices = tree.query(coords, k=k+1)
+
+        # Remove self-loop (first neighbor is self)
+        indices = indices[:, 1:]
+        distances = distances[:, 1:]
+
+        n_nodes = len(coords)
+        src = np.repeat(np.arange(n_nodes), k)
+        dst = indices.flatten()
+        edge_index = np.stack([src, dst])
+        edge_attr = distances.flatten()
+
+        return edge_index, edge_attr
+    else:
+        raise ValueError(f"Backend '{backend}' not implemented")