Add a geospatial extension

heat/bmi_geo.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class BmiGeo(ABC):
+
+    @abstractmethod
+    def initialize(self, filename: str=None) -> None:
+        ...
+
+    @abstractmethod
+    def get_grid_coordinate_names(self, grid: int) -> tuple[str, ...]:
+        ...
+
+    @abstractmethod
+    def get_grid_coordinate_units(self, grid: int) -> tuple[str, ...]:
+        ...
+
+    @abstractmethod
+    def get_grid_coordinate(self, grid: int, coordinate: str, values: NDArray[np.float64]) -> None:
+        ...
+
+    @abstractmethod
+    def get_grid_crs(self, grid: int) -> str:
+        ...

heat/bmi_heat.py

@@ -15,6 +15,9 @@ class BmiHeat(Bmi):
     _input_var_names = ("plate_surface__temperature",)
     _output_var_names = ("plate_surface__temperature",)
 
+    def get_extensions(self):
+        return ("[email protected]_heat_geo:BmiHeatGeo",)
+
     def __init__(self):
         """Create a BmiHeat model that is ready for initialization."""
         self._model = None

heat/bmi_heat_geo.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+from heat.bmi_geo import BmiGeo
+from heat.bmi_heat import BmiHeat
+
+
+class BmiHeatGeo(BmiGeo):
+    def __init__(self, bmi_heat: BmiHeat):
+        self._bmi_heat = bmi_heat
+
+    def initialize(self, filename: str = None) -> None:
+        pass
+
+    def get_grid_coordinate_names(self, grid: int) -> tuple[str, ...]:
+        return ("y", "x")
+
+    def get_grid_coordinate_units(self, grid: int) -> tuple[str, ...]:
+        return ("m", "m")
+
+    def get_grid_coordinate(
+        self, grid: int, coordinate: str, values: NDArray[np.float64]
+    ) -> None:
+        coords = np.meshgrid(
+            range(self._bmi_heat._model.shape[0]),
+            range(self._bmi_heat._model.shape[1]),
+            indexing="ij"
+        )
+        if coordinate == "y":
+            dim = 0
+        elif coordinate == "x":
+            dim = 1
+        else:
+            raise RuntimeError(f"{coordinate}: unknown coordinate")
+
+        values[:] = coords[dim].reshape(-1) * self._bmi_heat._model.spacing[dim] + self._bmi_heat._model.origin[dim]
+
+    def get_grid_crs(self, grid: int) -> str:
+        return "none"