Add in-memory output providers (ItziModel#129)

* move DrainageNodeCouplingData to data_containers

* get_node_list() no longer dependent on grass_interface

* analytic tests now use in-memory results and use generated arrays as input

Author: lrntct

src/itzi/data_containers.py

@@ -13,12 +13,29 @@
 GNU General Public License for more details.
 """
 
-from typing import Dict, Tuple
+from typing import Dict, Tuple, TYPE_CHECKING
 import dataclasses
 from datetime import datetime
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from itzi.drainage import DrainageNode
+
+
+@dataclasses.dataclass(frozen=True)
+class DrainageNodeCouplingData:
+    """Store the translation between coordinates and array location for a given drainage node."""
+
+    node_id: str  # Name of the drainage node
+    node_object: "DrainageNode"
+    # Location in the coordinate system
+    x: float | None
+    y: float | None
+    # Location in the array
+    row: int | None
+    col: int | None
+
 
 @dataclasses.dataclass(frozen=True)
 class DrainageAttributes:

src/itzi/itzi.py

@@ -110,9 +110,9 @@ def initialize(self, conf_file):
             raster_mask_id=self.conf.grass_params["mask"],
         )
         # Instantiate Simulation object and initialize it
-        from itzi.simulation_factories import create_simulation
+        from itzi.simulation_factories import create_grass_simulation
 
-        self.sim, self.tarr = create_simulation(
+        self.sim, self.tarr = create_grass_simulation(
             sim_times=self.conf.sim_times,
             stats_file=self.conf.stats_file,
             input_maps=self.conf.input_map_names,

src/itzi/providers/__init__.py

@@ -1,2 +0,0 @@
-from itzi.providers.grass_output import GrassRasterOutputProvider as GrassRasterOutputProvider
-from itzi.providers.grass_output import GrassVectorOutputProvider as GrassVectorOutputProvider

src/itzi/providers/base.py

@@ -30,9 +30,7 @@ def initialize(self, config: Dict) -> Self:
         pass
 
     @abstractmethod
-    def write_array(
-        self, array: np.ndarray, map_name: str, map_key: str, sim_time: datetime | timedelta
-    ) -> None:
+    def write_array(self, array: np.ndarray, map_key: str, sim_time: datetime | timedelta) -> None:
         """Write simulation data for current time step."""
         pass
 
@@ -58,6 +56,6 @@ def write_vector(
         pass
 
     @abstractmethod
-    def finalize(self) -> None:
+    def finalize(self, drainage_data: DrainageNetworkData) -> None:
         """Finalize outputs and cleanup."""
         pass

src/itzi/providers/memory_output.py

@@ -0,0 +1,62 @@
+# coding=utf8
+"""
+Copyright (C) 2025 Laurent Courty
+
+This program is free software; you can redistribute it and/or
+modify it under the terms of the GNU General Public License
+as published by the Free Software Foundation; either version 2
+of the License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+"""
+
+from datetime import datetime, timedelta
+from typing import Dict, Self
+from copy import deepcopy
+
+import numpy as np
+
+from itzi.providers.base import RasterOutputProvider, VectorOutputProvider
+from itzi.data_containers import SimulationData, DrainageNetworkData
+
+
+class MemoryRasterOutputProvider(RasterOutputProvider):
+    """Save rasters in memory as numpy arrays."""
+
+    def initialize(self, config: Dict) -> Self:
+        """Initialize output provider with simulation configuration."""
+        self.out_map_names = config["out_map_names"]
+
+        self.output_maps_dict = {k: [] for k in self.out_map_names.keys()}
+        return self
+
+    def write_array(self, array: np.ndarray, map_key: str, sim_time: datetime | timedelta) -> None:
+        """Save simulation data for current time step."""
+        self.output_maps_dict[map_key].append((deepcopy(sim_time), array.copy()))
+
+    def finalize(self, final_data: SimulationData) -> None:
+        """Finalize outputs and cleanup."""
+        pass
+
+
+class MemoryVectorOutputProvider(VectorOutputProvider):
+    """Save drainage simulation outputs in memory."""
+
+    def initialize(self, config: Dict | None = None) -> Self:
+        """Initialize output provider with simulation configuration."""
+        self.drainage_data = []
+
+        return self
+
+    def write_vector(
+        self, drainage_data: DrainageNetworkData, sim_time: datetime | timedelta
+    ) -> None:
+        """Save simulation data for current time step."""
+        self.drainage_data.append((deepcopy(sim_time), deepcopy(drainage_data)))
+
+    def finalize(self, drainage_data: DrainageNetworkData) -> None:
+        """Finalize outputs and cleanup."""
+        pass

src/itzi/rasterdomain.py

@@ -14,10 +14,57 @@
 """
 
 from datetime import datetime
+from typing import Tuple
 import numpy as np
 
 import itzi.flow as flow
-from itzi import rastermetrics
+
+
+class DomainData:
+    """Store raster domain information. Alike GRASS region."""
+
+    def __init__(self, north: float, south: float, east: float, west: float, rows: int, cols: int):
+        self.north = north
+        self.south = south
+        self.east = east
+        self.west = west
+        self.rows = rows
+        self.cols = cols
+
+        if self.north < self.south:
+            raise ValueError(f"north must be superior to south. {self.north=}, {self.south=}")
+        if self.east < self.west:
+            raise ValueError(f"east must be superior to west. {self.east=}, {self.west=}")
+
+        self.nsres = (self.north - self.south) / self.rows
+        self.ewres = (self.east - self.west) / self.cols
+        self.cell_area = self.ewres * self.nsres
+        self.cell_shape = (self.ewres, self.nsres)
+        self.shape = (self.rows, self.cols)
+        self.cells = self.rows * self.cols
+
+    def is_in_domain(self, *, x: float, y: float) -> bool:
+        """For a given coordinate pair(x, y),
+        return True is inside the domain, False otherwise.
+        """
+        bool_x = self.west < x < self.east
+        bool_y = self.south < y < self.north
+        return bool(bool_x and bool_y)
+
+    def coordinates_to_pixel(self, *, x: float, y: float) -> Tuple[float, float] | None:
+        """For a given coordinate pair(x, y),
+        return True is inside the domain, False otherwise.
+        """
+        if not self.is_in_domain(x=x, y=y):
+            return None
+        else:
+            norm_row = (y - self.south) / (self.north - self.south)
+            row = int(np.round((1 - norm_row) * (self.rows - 1)))
+
+            norm_col = (x - self.west) / (self.east - self.west)
+            col = int(np.round(norm_col * (self.cols - 1)))
+
+            return (row, col)
 
 
 class TimedArray:

src/itzi/simulation.py

@@ -14,21 +14,23 @@
 """
 
 from datetime import datetime, timedelta
-from typing import Self
+from typing import Self, Union, TYPE_CHECKING
 import copy
 
 import numpy as np
 
-from itzi.surfaceflow import SurfaceFlowSimulation
-import itzi.rasterdomain as rasterdomain
-from itzi.report import Report
 from itzi.data_containers import ContinuityData, SimulationData
-from itzi.drainage import DrainageSimulation
 import itzi.messenger as msgr
-from itzi.hydrology import Hydrology
 from itzi.itzi_error import NullError, MassBalanceError
 from itzi import rastermetrics
 
+if TYPE_CHECKING:
+    from itzi.drainage import DrainageSimulation
+    from itzi.hydrology import Hydrology
+    from itzi.surfaceflow import SurfaceFlowSimulation
+    from itzi.rasterdomain import RasterDomain
+    from itzi.report import Report
+
 
 class Simulation:
     """ """
@@ -72,12 +74,12 @@ def __init__(
         self,
         start_time: datetime,
         end_time: datetime,
-        raster_domain: rasterdomain.RasterDomain,
-        hydrology_model: Hydrology,
-        surface_flow: SurfaceFlowSimulation,
-        drainage_model: DrainageSimulation | None,
+        raster_domain: "RasterDomain",
+        hydrology_model: "Hydrology",
+        surface_flow: "SurfaceFlowSimulation",
+        drainage_model: Union["DrainageSimulation", None],
         nodes_list: list | None,
-        report: Report,
+        report: "Report",
         mass_balance_error_threshold: float,
     ):
         self.raster_domain = raster_domain
@@ -123,7 +125,7 @@ def __init__(
             self.next_ts["drainage"] = self.end_time
         else:
             self.node_id_to_loc = {
-                n.id: (n.row, n.col) for n in self.nodes_list if n.object.is_coupled()
+                n.node_id: (n.row, n.col) for n in self.nodes_list if n.node_object.is_coupled()
             }
         # Grid spacing (for BMI)
         self.spacing = (self.raster_domain.dy, self.raster_domain.dx)