store drainage data attributes in a dataclass (ItziModel#127)

Author: lrntct

src/itzi/data_containers.py

@@ -13,91 +13,87 @@
 GNU General Public License for more details.
 """
 
-from typing import Dict, Tuple, ClassVar
-from dataclasses import dataclass
+from typing import Dict, Tuple
+import dataclasses
 from datetime import datetime
 
 import numpy as np
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
+class DrainageAttributes:
+    """A base class for drainage data attributes."""
+
+    def get_columns_definition(self) -> list[tuple[str, str]]:
+        """Return a list of tuples to create DB columns"""
+        type_corresp = {str: "TEXT", int: "INT", float: "REAL"}
+        db_columns_def = [("cat", "INTEGER PRIMARY KEY")]
+        for f in dataclasses.fields(self):
+            db_field = (f.name, type_corresp[f.type])
+            db_columns_def.append(db_field)
+        return db_columns_def
+
+
+@dataclasses.dataclass(frozen=True)
+class DrainageLinkAttributes(DrainageAttributes):
+    link_id: str
+    link_type: str
+    flow: float
+    depth: float
+    volume: float
+    inlet_offset: float
+    outlet_offset: float
+    froude: float
+
+
+@dataclasses.dataclass(frozen=True)
 class DrainageLinkData:
     """Store the instantaneous state of a node during a drainage simulation"""
 
     vertices: Tuple[Tuple[float, float], ...]
-    attributes: Tuple  # one values for each columns, minus "cat"
-    columns_definition: ClassVar[Tuple[Tuple[str, str], ...]] = (
-        ("cat", "INTEGER PRIMARY KEY"),
-        ("link_id", "TEXT"),
-        ("type", "TEXT"),
-        ("flow", "REAL"),
-        ("depth", "REAL"),
-        #    (u'velocity', 'REAL'),
-        ("volume", "REAL"),
-        ("offset1", "REAL"),
-        ("offset2", "REAL"),
-        #    (u'yFull', 'REAL'),
-        ("froude", "REAL"),
-    )
-
-    def __post_init__(self):
-        """Validate attributes length after initialization."""
-        expected_len = len(self.columns_definition) - 1
-        if len(self.attributes) != expected_len:
-            raise ValueError(
-                f"DrainageLinkData: Incorrect number of attributes. "
-                f"Expected {expected_len}, got {len(self.attributes)}"
-            )
-
-
-@dataclass(frozen=True)
+    attributes: DrainageLinkAttributes
+
+
+@dataclasses.dataclass(frozen=True)
+class DrainageNodeAttributes(DrainageAttributes):
+    node_id: str
+    node_type: str
+    coupling_type: str
+    coupling_flow: float
+    inflow: float
+    outflow: float
+    lateral_inflow: float
+    losses: float
+    overflow: float
+    depth: float
+    head: float
+    # crownElev: float
+    crest_elevation: float
+    invert_elevation: float
+    initial_depth: float
+    full_depth: float
+    surcharge_depth: float
+    ponding_area: float
+    # degree: int
+    volume: float
+    full_volume: float
+
+
+@dataclasses.dataclass(frozen=True)
 class DrainageNodeData:
     """Store the instantaneous state of a node during a drainage simulation"""
 
     coordinates: Tuple[float, float]
-    attributes: Tuple  # one values for each columns, minus "cat"
-    columns_definition: ClassVar[Tuple[Tuple[str, str], ...]] = (
-        ("cat", "INTEGER PRIMARY KEY"),
-        ("node_id", "TEXT"),
-        ("type", "TEXT"),
-        ("linkage_type", "TEXT"),
-        ("linkage_flow", "REAL"),
-        ("inflow", "REAL"),
-        ("outflow", "REAL"),
-        ("latFlow", "REAL"),
-        ("losses", "REAL"),
-        ("overflow", "REAL"),
-        ("depth", "REAL"),
-        ("head", "REAL"),
-        #    (u'crownElev', 'REAL'),
-        ("crestElev", "REAL"),
-        ("invertElev", "REAL"),
-        ("initDepth", "REAL"),
-        ("fullDepth", "REAL"),
-        ("surDepth", "REAL"),
-        ("pondedArea", "REAL"),
-        #    (u'degree', 'INT'),
-        ("newVolume", "REAL"),
-        ("fullVolume", "REAL"),
-    )
-
-    def __post_init__(self):
-        """Validate attributes length after initialization."""
-        expected_len = len(self.columns_definition) - 1
-        if len(self.attributes) != expected_len:
-            raise ValueError(
-                f"DrainageNodeData: Incorrect number of attributes. "
-                f"Expected {expected_len}, got {len(self.attributes)}"
-            )
-
-
-@dataclass(frozen=True)
+    attributes: DrainageNodeAttributes
+
+
+@dataclasses.dataclass(frozen=True)
 class DrainageNetworkData:
     nodes: Tuple[DrainageNodeData, ...]
     links: Tuple[DrainageLinkData, ...]
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class ContinuityData:
     """Store information about simulation continuity"""
 
@@ -107,7 +103,7 @@ class ContinuityData:
     continuity_error: float
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class SimulationData:
     """Immutable data container for passing raw simulation state to Report.
 
@@ -127,7 +123,7 @@ class SimulationData:
     drainage_network_data: DrainageNetworkData | None
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class MassBalanceData:
     """Contains the fields written to the mass balance file"""
 

src/itzi/drainage.py

@@ -22,7 +22,13 @@
 
 from itzi import DefaultValues
 from itzi import messenger as msgr
-from itzi.data_containers import DrainageNodeData, DrainageLinkData, DrainageNetworkData
+from itzi.data_containers import (
+    DrainageNodeData,
+    DrainageLinkData,
+    DrainageNetworkData,
+    DrainageLinkAttributes,
+    DrainageNodeAttributes,
+)
 
 
 class CouplingTypes(StrEnum):
@@ -168,37 +174,34 @@ def is_coupled(self):
         """return True if the node is coupled to the 2D domain"""
         return self.coupling_type != CouplingTypes.NOT_COUPLED
 
-    def get_attrs(self):
-        """return a list of node data in the right DB order
-        TODO: put the burden of DB order to the code actually writing the DB
-        """
-        attrs = [
-            self.node_id,
-            self.node_type,
-            self.coupling_type.value,
-            self.coupling_flow,
-            self.pyswmm_node.total_inflow,
-            self.pyswmm_node.total_outflow,
-            self.pyswmm_node.lateral_inflow,
-            self.pyswmm_node.losses,
-            self.get_overflow(),
-            self.pyswmm_node.depth,
-            self.pyswmm_node.head,
-            #  values['crown_elev'],
-            self.get_crest_elev(),
-            self.pyswmm_node.invert_elevation,
-            self.pyswmm_node.initial_depth,
-            self.pyswmm_node.full_depth,
-            self.pyswmm_node.surcharge_depth,
-            self.pyswmm_node.ponding_area,
-            #  values['degree'],
-            self.pyswmm_node.volume,
-            self.get_full_volume(),
-        ]
-        return attrs
+    def get_attrs(self) -> DrainageNodeAttributes:
+        """ """
+        return DrainageNodeAttributes(
+            node_id=self.node_id,
+            node_type=self.node_type,
+            coupling_type=self.coupling_type.value,
+            coupling_flow=self.coupling_flow,
+            inflow=self.pyswmm_node.total_inflow,
+            outflow=self.pyswmm_node.total_outflow,
+            lateral_inflow=self.pyswmm_node.lateral_inflow,
+            losses=self.pyswmm_node.losses,
+            overflow=self.get_overflow(),
+            depth=self.pyswmm_node.depth,
+            head=self.pyswmm_node.head,
+            # crownElev=values['crown_elev'],
+            crest_elevation=self.get_crest_elev(),
+            invert_elevation=self.pyswmm_node.invert_elevation,
+            initial_depth=self.pyswmm_node.initial_depth,
+            full_depth=self.pyswmm_node.full_depth,
+            surcharge_depth=self.pyswmm_node.surcharge_depth,
+            ponding_area=self.pyswmm_node.ponding_area,
+            # degree=values['degree'],
+            volume=self.pyswmm_node.volume,
+            full_volume=self.get_full_volume(),
+        )
 
     def get_data(self) -> DrainageNodeData:
-        return DrainageNodeData(coordinates=self.coordinates, attributes=tuple(self.get_attrs()))
+        return DrainageNodeData(coordinates=self.coordinates, attributes=self.get_attrs())
 
     def apply_coupling(self, z, h, dt_drainage, cell_surf):
         """Apply the coupling to the node"""
@@ -333,23 +336,20 @@ def _get_link_type(self):
             raise ValueError(f"Unknown link type for link {self.link_id}")
         return link_type
 
-    def get_attrs(self):
-        """return a list of link data in the right DB order
-        TODO: put the burden of DB order to the code actually writing the DB
-        """
-        attrs = [
-            self.link_id,
-            self.link_type,
-            self.pyswmm_link.flow,
-            self.pyswmm_link.depth,
-            #  values['velocity'],
-            self.pyswmm_link.volume,
-            self.pyswmm_link.inlet_offset,
-            self.pyswmm_link.outlet_offset,
-            #  values['full_depth'],
-            self.pyswmm_link.froude,
-        ]
-        return attrs
+    def get_attrs(self) -> DrainageLinkAttributes:
+        """ """
+        return DrainageLinkAttributes(
+            link_id=self.link_id,
+            link_type=self.link_type,
+            flow=self.pyswmm_link.flow,
+            depth=self.pyswmm_link.depth,
+            # velocity=values['velocity'],
+            volume=self.pyswmm_link.volume,
+            inlet_offset=self.pyswmm_link.inlet_offset,
+            outlet_offset=self.pyswmm_link.outlet_offset,
+            # full_depth=values['full_depth'],
+            froude=self.pyswmm_link.froude,
+        )
 
     def get_data(self) -> DrainageLinkData:
-        return DrainageLinkData(vertices=self.vertices, attributes=tuple(self.get_attrs()))
+        return DrainageLinkData(vertices=self.vertices, attributes=self.get_attrs())

src/itzi/providers/grass_interface.py

@@ -17,6 +17,7 @@
 from pathlib import Path
 from collections import namedtuple
 from datetime import datetime, timedelta
+import dataclasses
 
 # from multiprocessing import Process, JoinableQueue
 from threading import Thread, Lock
@@ -547,8 +548,8 @@ def write_vector_map(self, drainage_data: DrainageNetworkData, map_name: str):
 
         with VectorTopo(map_name, mode="w", overwrite=self.overwrite) as vect_map:
             # create db links and tables
-            node_columns_def = drainage_data.nodes[0].columns_definition
-            link_columns_def = drainage_data.links[0].columns_definition
+            node_columns_def = drainage_data.nodes[0].attributes.get_columns_definition()
+            link_columns_def = drainage_data.links[0].attributes.get_columns_definition()
             linking_elements = {
                 "node": self.LayerDescr(
                     table_suffix="_node", cols=node_columns_def, layer_number=1
@@ -573,7 +574,10 @@ def write_vector_map(self, drainage_data: DrainageNetworkData, map_name: str):
                     map_layer, dbtable = dblinks["node"]
                     self.write_vector_geometry(vect_map, point, cat_num, map_layer)
                     # Get DB attributes
-                    attrs = (cat_num,) + node.attributes
+                    node_attributes = tuple(
+                        value for _, value in dataclasses.asdict(node.attributes).items()
+                    )
+                    attrs = (cat_num,) + node_attributes
                     db_info["node"].append(attrs)
                     # bump cat
                     cat_num += 1
@@ -587,7 +591,10 @@ def write_vector_map(self, drainage_data: DrainageNetworkData, map_name: str):
                     map_layer, dbtable = dblinks["link"]
                     self.write_vector_geometry(vect_map, line_object, cat_num, map_layer)
                     # keep DB info
-                    attrs = (cat_num,) + link.attributes
+                    link_attributes = tuple(
+                        value for _, value in dataclasses.asdict(link.attributes).items()
+                    )
+                    attrs = (cat_num,) + link_attributes
                     db_info["link"].append(attrs)
                     # bump cat
                     cat_num += 1

tests/test_ea8b.py

@@ -149,12 +149,12 @@ def ea_test8b_sim(ea_test8b, test_data_path, test_data_temp_path):
 
 @pytest.fixture(scope="module")
 def ea8b_itzi_drainage_results(ea_test8b_sim):
-    """Extract linkage flow from the drainage network"""
+    """Extract coupling flow from the drainage network"""
     current_mapset = gscript.read_command("g.mapset", flags="p").rstrip()
     assert current_mapset == "ea8b"
-    select_col = ["start_time", "linkage_flow"]
+    select_col = ["start_time", "coupling_flow"]
     itzi_results = gscript.read_command("t.vect.db.select", input="out_drainage")
-    # translate to Pandas dataframe and keep only linkage_flow with start_time over 3000
+    # translate to Pandas dataframe and keep only coupling_flow with start_time over 3000
     df_itzi_results = pd.read_csv(StringIO(itzi_results), sep="|")[select_col]
     df_itzi_results = df_itzi_results[df_itzi_results.start_time >= 3000]
     df_itzi_results.set_index("start_time", drop=True, inplace=True, verify_integrity=True)

tests/test_tutorial.py

@@ -152,7 +152,7 @@ def test_tutorial_drainage(itzi_tutorial, test_data_path, test_data_temp_path, h
         sim_runner.initialize(config_file)
         sim_runner.run().finalize()
         # Check the results
-        select_cols = ["start_time", "linkage_flow"]
+        select_cols = ["start_time", "coupling_flow"]
         drainage_results = gscript.read_command(
             "t.vect.db.select",
             input="nc_itzi_tutorial_drainage",