implement geospatial extension for bmi_heat

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.heat import Heat
+
+
+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"