Add StormCast model

CHANGELOG.md

@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add StormCast model to prognostic models
+
+### Changed
+
 - Interpolation between arbitrary lat-lon grids
 
 ### Changed

earth2studio/data/rand.py

@@ -39,10 +39,18 @@ def __init__(
     ):
         self.domain_coords = domain_coords
 
+        # Check for regular vs. curvilinear coordinates
+        _, value = list(self.domain_coords.items()).pop()
+        value = np.array(value)
+        self.curv = len(value.shape) > 1
+        if self.curv:
+            self.domain_coord_shape = value.shape
+
     def __call__(
         self,
         time: datetime | list[datetime] | TimeArray,
         variable: str | list[str] | VariableArray,
+        lead_time: np.array = None,
     ) -> xr.DataArray:
         """Retrieve random gaussian data.
 
@@ -63,12 +71,24 @@ def __call__(
 
         shape = [len(time), len(variable)]
         coords = {"time": time, "variable": variable}
-        for key, value in self.domain_coords.items():
-            shape.append(len(value))
-            coords[key] = value
 
-        da = xr.DataArray(
-            data=np.random.randn(*shape), dims=list(coords), coords=coords
-        )
+        if self.curv:
+            shape.extend(self.domain_coord_shape)
+            dims = ["time", "variable", "y", "x"]
+            coords = coords | {
+                "lat": (("y", "x"), self.domain_coords["lat"]),
+                "lon": (("y", "x"), self.domain_coords["lon"]),
+            }
+            da = xr.DataArray(data=np.random.randn(*shape), dims=dims, coords=coords)
+
+        else:
+
+            for key, value in self.domain_coords.items():
+                shape.append(len(value))
+                coords[key] = value
+
+            da = xr.DataArray(
+                data=np.random.randn(*shape), dims=list(coords), coords=coords
+            )
 
         return da

earth2studio/data/utils.py

@@ -43,8 +43,12 @@ def fetch_data(
     variable: VariableArray,
     lead_time: LeadTimeArray = np.array([np.timedelta64(0, "h")]),
     device: torch.device = "cpu",
+    interp_to: CoordSystem = None,
+    interp_method: str = "nearest",
 ) -> tuple[torch.Tensor, CoordSystem]:
     """Utility function to fetch data for models and load data on the target device.
+    If desired, xarray interpolation/regridding in the spatial domain can be used
+    by passing a target coordinate system via the optional `interp_to` argument.
 
     Parameters
     ----------
@@ -59,6 +63,11 @@ def fetch_data(
         np.array(np.timedelta64(0, "h"))
     device : torch.device, optional
         Torch devive to load data tensor to, by default "cpu"
+    interp_to : CoordSystem, optional
+        If provided, the fetched data will be interpolated to the coordinates
+        specified by lat/lon arrays in this CoordSystem
+    interp_method : str
+        Interpolation method to use with xarray (by default 'nearest')
 
     Returns
     -------
@@ -75,34 +84,81 @@ def fetch_data(
         da0 = da0.assign_coords(time=time)
         da.append(da0)
 
-    return prep_data_array(xr.concat(da, "lead_time"), device=device)
+    return prep_data_array(
+        xr.concat(da, "lead_time"),
+        device=device,
+        interp_to=interp_to,
+        interp_method=interp_method,
+    )
 
 
 def prep_data_array(
     da: xr.DataArray,
     device: torch.device = "cpu",
+    interp_to: CoordSystem = None,
+    interp_method: str = "nearest",
 ) -> tuple[torch.Tensor, CoordSystem]:
     """Prepares a data array from a data source for inference workflows by converting
     the data array to a torch tensor and the coordinate system to an OrderedDict.
 
+    If desired, xarray interpolation/regridding in the spatial domain can be used
+    by passing a target coordinate system via the optional `interp_to` argument.
+
     Parameters
     ----------
     da : xr.DataArray
         Input data array
     device : torch.device, optional
         Torch devive to load data tensor to, by default "cpu"
+    interp_to : CoordSystem, optional
+        If provided, the fetched data will be interpolated to the coordinates
+        specified by lat/lon arrays in this CoordSystem
+    interp_method : str
+        Interpolation method to use with xarray (by default 'nearest')
 
     Returns
     -------
     tuple[torch.Tensor, CoordSystem]
         Tuple containing output tensor and coordinate OrderedDict
     """
 
+    if interp_to is not None:
+        if len(interp_to["lat"].shape) > 1 or len(interp_to["lon"].shape) > 1:
+            if len(interp_to["lat"].shape) != len(interp_to["lon"].shape):
+                raise ValueError(
+                    "Discrepancy in interpolation coordinates: latitude has different shape than longitude"
+                )
+
+            # Curvilinear coordinates: define internal dims y, x
+            target_lat = xr.DataArray(interp_to["lat"], dims=["y", "x"])
+            target_lon = xr.DataArray(interp_to["lon"], dims=["y", "x"])
+        else:
+
+            target_lat = xr.DataArray(interp_to["lat"], dims=["lat"])
+            target_lon = xr.DataArray(interp_to["lon"], dims=["lon"])
+        target_grid = xr.Dataset({"latitude": target_lat, "longitude": target_lon})
+
+        da = da.interp(
+            lat=target_grid["latitude"],
+            lon=target_grid["longitude"],
+            method=interp_method,
+        )
+
     out = torch.Tensor(da.values).to(device)
 
     out_coords = OrderedDict()
-    for dim in da.coords.dims:
-        out_coords[dim] = np.array(da.coords[dim])
+
+    curvilinear = "lat" in da.coords and len(da.coords["lat"].shape) > 1
+    if curvilinear:
+        # Assume the lat, lon array are defined as DataArray coordinates
+        for dim in da.coords.dims:
+            if dim in ["time", "variable", "lead_time"]:
+                out_coords[dim] = np.array(da.coords[dim])
+        out_coords["lat"] = np.array(da.coords["lat"])
+        out_coords["lon"] = np.array(da.coords["lon"])
+    else:
+        for dim in da.coords.dims:
+            out_coords[dim] = np.array(da.coords[dim])
 
     return out, out_coords
 

earth2studio/models/batch.py

@@ -164,15 +164,28 @@ def _batch_wrap(self, func: Callable) -> Callable:
         # TODO: Better typing for model object
         @functools.wraps(func)
         def _wrapper(
-            model: Any, x: torch.Tensor, coords: CoordSystem
+            model: Any,
+            x: torch.Tensor,
+            coords: CoordSystem,
+            conditioning: torch.Tensor | None = None,
+            conditioning_coords: CoordSystem | None = None,
         ) -> tuple[torch.Tensor, CoordSystem]:
 
             x, flatten_coords, batched_coords, batched_shape = self._compress_batch(
                 model, x, coords
             )
 
             # Model forward
-            out, out_coords = func(model, x, flatten_coords)
+            if conditioning is not None:
+                out, out_coords = func(
+                    model,
+                    x,
+                    flatten_coords,
+                    conditioning=conditioning,
+                    conditioning_coords=conditioning_coords,
+                )
+            else:
+                out, out_coords = func(model, x, flatten_coords)
             out, out_coords = self._decompress_batch(
                 out, out_coords, batched_coords, batched_shape
             )

earth2studio/models/px/__init__.py

@@ -22,3 +22,4 @@
 from earth2studio.models.px.pangu import Pangu3, Pangu6, Pangu24
 from earth2studio.models.px.persistence import Persistence
 from earth2studio.models.px.sfno import SFNO
+from earth2studio.models.px.stormcast import StormCast