From ca359349e8c91a27c2a571eaf34a5da9e06402ed Mon Sep 17 00:00:00 2001
From: mcflugen <mcflugen@gmail.com>
Date: Mon, 29 Jan 2024 17:30:01 -0700
Subject: [PATCH 1/4] add get_extension method

---
 heat/bmi_heat.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/heat/bmi_heat.py b/heat/bmi_heat.py
index 310855a..22751fa 100644
--- a/heat/bmi_heat.py
+++ b/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 ("bmi_geospatial@heat.bmi_heat_geo:BmiHeatGeo",)
+
     def __init__(self):
         """Create a BmiHeat model that is ready for initialization."""
         self._model = None

From 8b7dfa41f3bb441c6cc5d3a7ae049b2bf3801341 Mon Sep 17 00:00:00 2001
From: mcflugen <mcflugen@gmail.com>
Date: Mon, 29 Jan 2024 17:31:17 -0700
Subject: [PATCH 2/4] add abc for the geospatial extension

---
 heat/bmi_geo.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 heat/bmi_geo.py

diff --git a/heat/bmi_geo.py b/heat/bmi_geo.py
new file mode 100644
index 0000000..fd04fba
--- /dev/null
+++ b/heat/bmi_geo.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+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:
+        ...

From 86db3217da2d710e08b04b36c762f39809ffc3aa Mon Sep 17 00:00:00 2001
From: mcflugen <mcflugen@gmail.com>
Date: Mon, 29 Jan 2024 17:31:44 -0700
Subject: [PATCH 3/4] implement geospatial extension for bmi_heat

---
 heat/bmi_heat_geo.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 heat/bmi_heat_geo.py

diff --git a/heat/bmi_heat_geo.py b/heat/bmi_heat_geo.py
new file mode 100644
index 0000000..2fdfe60
--- /dev/null
+++ b/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"

From b6f9254ee04807330433a3b414463f1a8a8db2d8 Mon Sep 17 00:00:00 2001
From: mcflugen <mcflugen@gmail.com>
Date: Mon, 29 Jan 2024 22:27:03 -0700
Subject: [PATCH 4/4] remove some lint

---
 heat/bmi_geo.py      | 1 -
 heat/bmi_heat_geo.py | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/heat/bmi_geo.py b/heat/bmi_geo.py
index fd04fba..43b7fe4 100644
--- a/heat/bmi_geo.py
+++ b/heat/bmi_geo.py
@@ -2,7 +2,6 @@
 
 from abc import ABC
 from abc import abstractmethod
-from typing import Any
 
 import numpy as np
 from numpy.typing import NDArray
diff --git a/heat/bmi_heat_geo.py b/heat/bmi_heat_geo.py
index 2fdfe60..095ac44 100644
--- a/heat/bmi_heat_geo.py
+++ b/heat/bmi_heat_geo.py
@@ -4,7 +4,7 @@
 from numpy.typing import NDArray
 
 from heat.bmi_geo import BmiGeo
-from heat.heat import Heat
+from heat.bmi_heat import BmiHeat
 
 
 class BmiHeatGeo(BmiGeo):