Golnesa variable friction

threedigrid_builder/base/lines.py

@@ -38,6 +38,8 @@ class Line:
     frict_type2: int
     frict_value1: float  # For 1D-2Dgw: resistance factor-out (hydr. cond / thickness)
     frict_value2: float  # For 1D-2Dgw: resistance factor-in (hydr. cond / thickness)
+    veg_coef1: float
+    veg_coef2: float
     cross_weight: float  # For 1D-2Dgw: the exchange length (tables: cross_width)
     discharge_coefficient_positive: float
     discharge_coefficient_negative: float

threedigrid_builder/grid/cross_section_definitions.py

@@ -17,6 +17,12 @@ class CrossSectionDefinition:
     shape: CrossSectionShape
     height: str  # space-separated list of floats
     width: str  # space-separated list of floats
+    friction_values: str  # space-separated list of floats
+    vegetation_stem_densities: str  # space-separated list of floats
+    vegetation_stem_diameters: str  # space-separated list of floats
+    vegetation_heights: str  # space-separated list of floats
+    vegetation_drag_coefficients: str  # space-separated list of floats
+
 
 
 class CrossSectionDefinitions(Array[CrossSectionDefinition]):
@@ -43,36 +49,47 @@ def convert(self, ids):
             content_pk=ids,
             code=self.code[idx],
             count=0,
+            count_yz = 0,
         )
         if len(result) == 0:
             return result
 
-        offset = 0
         tables = []
+        tables_yz = []
+
 
         # Numpy array views on width/height based on idx
         width_idx = self.width[idx]
         height_idx = self.height[idx]
 
         for i, shape in enumerate(self.shape[idx]):
             tabulator = tabulators[shape]
-            result.shape[i], result.width_1d[i], result.height_1d[i], table = tabulator(
+            result.shape[i], result.width_1d[i], result.height_1d[i], table, yz = tabulator(
                 shape, width_idx[i], height_idx[i]
             )
             if table is not None:
                 result.count[i] = len(table)
-                result.offset[i] = offset
-                offset += len(table)
                 tables.append(table)
+            if yz is not None:
+                result.count_yz[i] = len(yz)
+                tables_yz.append(yz)
+
 
         result.offset[:] = np.roll(np.cumsum(result.count), 1)
         result.offset[0] = 0
+        result.offset_yz[:] = np.roll(np.cumsum(result.count_yz), 1)
+        result.offset_yz[0] = 0
 
         if len(tables) > 0:
             result.tables = np.concatenate(tables, axis=0)
         else:
             result.tables = np.empty((0, 2))
 
+        if len(tables_yz) > 0:
+            result.tables_yz = np.concatenate(tables_yz, axis=0)
+        else:
+            result.tables_yz = np.empty((0, 4))    
+
         return result
 
 
@@ -85,11 +102,14 @@ class CrossSection:
     height_1d: float
     offset: int
     count: int
+    offset_yz: int
+    count_yz: int
     # tables: Tuple[float, float] has different length so is specified on CrossSections
 
 
 class CrossSections(Array[CrossSection]):
     tables = None
+    tables_yz = None
 
 
 def tabulate_builtin(shape, width, height):
@@ -112,7 +132,7 @@ def tabulate_builtin(shape, width, height):
             f"Unable to parse cross section definition width (got: '{width}')."
         )
 
-    return shape, width, None, None
+    return shape, width, None, None, None
 
 
 def tabulate_egg(shape, width, height):
@@ -151,17 +171,17 @@ def tabulate_egg(shape, width, height):
     widths = np.sqrt(p / q) * 2
 
     table = np.array([heights, widths]).T
-    return CrossSectionShape.TABULATED_TRAPEZIUM, width, height, table
+    return CrossSectionShape.TABULATED_TRAPEZIUM, width, height, table, None
 
 
 def tabulate_inverted_egg(shape, width, height):
     """Tabulate the egg shape, upside down.
 
     See tabulate_egg.
     """
-    type_, width, height, table = tabulate_egg(shape, width, height)
+    type_, width, height, table, yz = tabulate_egg(shape, width, height)
     table[:, 1] = table[::-1, 1]
-    return type_, width, height, table
+    return type_, width, height, table, None
 
 
 def tabulate_closed_rectangle(shape, width, height):
@@ -185,7 +205,7 @@ def tabulate_closed_rectangle(shape, width, height):
             f"(got: '{width}', '{height}')."
         )
     table = np.array([[0.0, width], [height, 0.0]], order="F")
-    return CrossSectionShape.TABULATED_RECTANGLE, width, height, table
+    return CrossSectionShape.TABULATED_RECTANGLE, width, height, table, None
 
 
 def _parse_tabulated(width, height):
@@ -229,10 +249,10 @@ def tabulate_tabulated(shape, width, height):
             f"(got: {height})."
         )
 
-    return shape, np.max(widths), np.max(heights), np.array([heights, widths]).T
+    return shape, np.max(widths), np.max(heights), np.array([heights, widths]).T, None
 
 
-def tabulate_yz(shape, width, height):
+def tabulate_yz(shape, width, height, friction_values, vegetation_stem_densities, vegetation_stem_diameters, vegetation_heights, vegetation_drag_coefficients):
     """Tabulate an (open or closed) YZ profile
 
     Args:
@@ -275,6 +295,19 @@ def tabulate_yz(shape, width, height):
     if is_closed:
         ys = ys[:-1]
         zs = zs[:-1]
+        yz = None
+    else:
+        yz = np.zeros((len(ys), 4), dtype=float)
+        yz[:, 0] = ys
+        yz[:, 1] = zs
+        fric = np.array([float(x) for x in friction_values.split(" ")])
+        veg_stemden = np.array([float(x) for x in vegetation_stem_densities.split(" ")])
+        veg_stemdia = np.array([float(x) for x in vegetation_stem_diameters.split(" ")])
+        veg_hght = np.array([float(x) for x in vegetation_heights.split(" ")])
+        veg_drag = np.array([float(x) for x in vegetation_drag_coefficients.split(" ")])
+        yz[:-1, 2] = fric
+        yz[:-1, 3] = veg_stemden * veg_stemdia * veg_hght * veg_drag
+
 
     # Adapt non-unique height coordinates. Why?
     # Because if a segment of the profile is exactly horizontal, we need 2 widths
@@ -318,7 +351,7 @@ def tabulate_yz(shape, width, height):
         table = table[1:]
 
     height_1d, width_1d = table.max(axis=0).tolist()
-    return CrossSectionShape.TABULATED_TRAPEZIUM, width_1d, height_1d, table
+    return CrossSectionShape.TABULATED_TRAPEZIUM, width_1d, height_1d, table, yz
 
 
 tabulators = {

threedigrid_builder/grid/cross_section_locations.py

@@ -17,6 +17,10 @@ class CrossSectionLocation:
     bank_level: float
     friction_type: FrictionType
     friction_value: float
+    vegetation_stem_density: float
+    vegetation_stem_diameter: float
+    vegetation_height: float
+    vegetation_drag_coefficient: float
 
 
 class CrossSectionLocations(Array[CrossSectionLocation]):
@@ -64,6 +68,8 @@ def apply_to_lines(self, lines, channels, extrapolate=False):
         lines.frict_type2 = self.friction_type[idx2]
         lines.frict_value1 = self.friction_value[idx1]
         lines.frict_value2 = self.friction_value[idx2]
+        lines.veg_coef1 = self.vegetation_stem_density[idx1] * self.vegetation_stem_diameter[idx1] * self.vegetation_height[idx1] * self.vegetation_drag_coefficient[idx1]
+        lines.veg_coef2 = self.vegetation_stem_density[idx2] * self.vegetation_stem_diameter[idx2] * self.vegetation_height[idx2] * self.vegetation_drag_coefficient[idx2]
 
         # Compute invert levels and start and end
         lines.invert_level_start_point = compute_bottom_level(

threedigrid_builder/interface/db.py

@@ -468,6 +468,11 @@ def get_cross_section_definitions(self) -> CrossSectionDefinitions:
                     models.CrossSectionDefinition.shape,
                     models.CrossSectionDefinition.width,
                     models.CrossSectionDefinition.height,
+                    models.CrossSectionDefinition.friction_values,
+                    models.CrossSectionDefinition.vegetation_stem_densities,
+                    models.CrossSectionDefinition.vegetation_stem_diameters,
+                    models.CrossSectionDefinition.vegetation_heights,
+                    models.CrossSectionDefinition.vegetation_drag_coefficients,
                 )
                 .order_by(models.CrossSectionDefinition.id)
                 .as_structarray()
@@ -497,6 +502,10 @@ def get_cross_section_locations(self) -> CrossSectionLocations:
                     models.CrossSectionLocation.bank_level,
                     models.CrossSectionLocation.friction_type,
                     models.CrossSectionLocation.friction_value,
+                    models.CrossSectionLocation.vegetation_stem_density,
+                    models.CrossSectionLocation.vegetation_stem_diameter,
+                    models.CrossSectionLocation.vegetation_height,
+                    models.CrossSectionLocation.vegetation_drag_coefficient,
                 )
                 .order_by(models.CrossSectionLocation.id)
                 .as_structarray()

threedigrid_builder/interface/gridadmin.py

@@ -442,6 +442,8 @@ def write_lines(self, lines: Lines, cross_sections):
         self.write_dataset(group, "frict_type2", lines.frict_type2)
         self.write_dataset(group, "frict_value1", lines.frict_value1)
         self.write_dataset(group, "frict_value2", lines.frict_value2)
+        self.write_dataset(group, "veg_coef1", lines.veg_coef1)
+        self.write_dataset(group, "veg_coef2", lines.veg_coef2)
         self.write_dataset(group, "cross_weight", lines.cross_weight)
         self.write_dataset(
             group, "hydraulic_resistance_in", lines.hydraulic_resistance_in
@@ -690,6 +692,10 @@ def write_cross_sections(self, cross_sections: CrossSections):
         self.write_dataset(group, "offset", cross_sections.offset)
         self.write_dataset(group, "count", cross_sections.count)
         self.write_dataset(group, "tables", cross_sections.tables.T, insert_dummy=False)
+        if cross_sections.tables_yz is not None:
+            self.write_dataset(group, "offset_yz", cross_sections.offset_yz)
+            self.write_dataset(group, "count_yz", cross_sections.count_yz)
+            self.write_dataset(group, "tables_yz", cross_sections.tables_yz.T, insert_dummy=False)
 
     def write_obstacles(self, obstacles: Obstacles):
         """For backwards compat, the group is named 'levees'"""

threedigrid_builder/tests/test_cross_section_definitions.py

@@ -1,3 +1,4 @@
+from cmath import nan
 from unittest import mock
 
 import numpy as np
@@ -36,9 +37,9 @@ def test_convert_multiple(cross_section_definitions):
     with mock.patch.dict(
         "threedigrid_builder.grid.cross_section_definitions.tabulators",
         {
-            SHP.CIRCLE: mock.Mock(return_value=(1, 0.1, None, None)),
-            SHP.TABULATED_TRAPEZIUM: mock.Mock(return_value=(5, 15.0, 2.0, table_1)),
-            SHP.TABULATED_RECTANGLE: mock.Mock(return_value=(6, 11.0, 2.0, table_2)),
+            SHP.CIRCLE: mock.Mock(return_value=(1, 0.1, None, None, None)),
+            SHP.TABULATED_TRAPEZIUM: mock.Mock(return_value=(5, 15.0, 2.0, table_1, None)),
+            SHP.TABULATED_RECTANGLE: mock.Mock(return_value=(6, 11.0, 2.0, table_2, None)),
         },
     ):
         actual = cross_section_definitions.convert([1, 3, 9])
@@ -59,7 +60,7 @@ def test_convert_multiple_filtered(cross_section_definitions):
     with mock.patch.dict(
         "threedigrid_builder.grid.cross_section_definitions.tabulators",
         {
-            SHP.CIRCLE: mock.Mock(return_value=(1, 0.1, 0.1, None)),
+            SHP.CIRCLE: mock.Mock(return_value=(1, 0.1, 0.1, None, None)),
         },
     ):
         actual = cross_section_definitions.convert([1])
@@ -85,26 +86,28 @@ def test_convert_empty(cross_section_definitions):
 
 def test_tabulate_builtin():
     actual = tabulate_builtin("my-shape", "1.52", "1.33")
-    assert actual == ("my-shape", 1.52, None, None)
+    assert actual == ("my-shape", 1.52, None, None, None)
 
 
 def test_tabulate_closed_rectangle():
-    shape, width_1d, height_1d, table = tabulate_closed_rectangle(
+    shape, width_1d, height_1d, table, yz = tabulate_closed_rectangle(
         "my-shape", "1.52", "5.2"
     )
 
     assert shape == CrossSectionShape.TABULATED_RECTANGLE
     assert width_1d == 1.52
     assert height_1d == 5.2
     assert_almost_equal(table, np.array([[0.0, 1.52], [5.2, 0.0]]))
+    assert yz is None
 
 
 def test_tabulate_egg():
-    shape, width_1d, height_1d, table = tabulate_egg("my-shape", "1.52", "ignored")
+    shape, width_1d, height_1d, table, yz = tabulate_egg("my-shape", "1.52", "ignored")
 
     assert shape == CrossSectionShape.TABULATED_TRAPEZIUM
     assert width_1d == 1.52
     assert height_1d == 1.52 * 1.5
+    assert yz is None
 
     # the expected table is exactly what inpy returns for a width of 1.52
     expected_table = np.array(
@@ -131,13 +134,13 @@ def test_tabulate_egg():
 
 
 def test_tabulate_tabulated():
-    shape, width_1d, height_1d, table = tabulate_tabulated("my-shape", "1 2 3", "0 1 2")
+    shape, width_1d, height_1d, table, yz = tabulate_tabulated("my-shape", "1 2 3", "0 1 2")
 
     assert shape == "my-shape"
     assert width_1d == 3.0  # the max
     assert height_1d == 2.0
     assert_almost_equal(table, np.array([[0, 1], [1, 2], [2, 3]], dtype=float))
-
+    assert yz is None
 
 @pytest.mark.parametrize(
     "width,height,match",
@@ -156,13 +159,14 @@ def test_tabulate_tabulated_err(width, height, match):
 
 
 def test_tabulate_inverted_egg():
-    shape, width_1d, height_1d, table = tabulate_inverted_egg(
+    shape, width_1d, height_1d, table, yz = tabulate_inverted_egg(
         "my-shape", "1.52", "ignored"
     )
 
     assert shape == CrossSectionShape.TABULATED_TRAPEZIUM
     assert width_1d == 1.52
     assert height_1d == 1.52 * 1.5
+    assert yz is None
 
     # the expected table is exactly what inpy returns for a width of 1.52
     expected_table = np.array(
@@ -189,41 +193,61 @@ def test_tabulate_inverted_egg():
 
 
 @pytest.mark.parametrize(
-    "width,height,exp_width,exp_height,exp_table",
+    "width,height,friction_values,vegetation_stem_densities,vegetation_stem_diameters,vegetation_heights,vegetation_drag_coefficients,exp_width,exp_height,exp_table,exp_yz",
     [
-        ("0 0.5 1 1.5", "0.5 0 0 0.5", 1.5, 0.5, [[0, 0.5], [0.5, 1.5]]),
+        ("0 0.5 1 1.5", "0.5 0 0 0.5", "1 1 1", "1 1 1", "1 1 1", "1 1 1", "1 1 1", 1.5, 0.5, [[0, 0.5], [0.5, 1.5]], [[0, 0.5, 1, 1], [0.5, 0, 1, 1], [1, 0, 1, 1], [1.5, 0.5, 0, 0]]),
         (
             "0 0.5 1 1.5",
             "0.5 0 0 0.25",
+            "1 1 1",
+            "0.5 1 1",
+            "1 1 1",
+            "3 1 1",
+            "2 1 1",
             1.5,
             0.5,
             [[0, 0.5], [0.25, 1.25], [0.5, 1.5]],
+            [[0, 0.5, 1, 3], [0.5, 0, 1, 1], [1, 0, 1, 1], [1.5, 0.25, 0, 0]]
         ),
         (
             "0 1 2 3 4 5",
             "1 0 0.5 0.5 0 1",
+            "1 1 1 1 1",
+            "1 1 1 1 1",
+            "1 1 1 1 1",
+            "1 1 1 1 1",
+            "1 1 1 1 1",
             5,
             1,
             [[0, 0], [0.5, 3], [0.5, 4], [1, 5]],
+            [[0, 1, 1, 1], [1, 0, 1, 1], [2, 0.5, 1, 1], [3, 0.5, 1, 1], [4, 0, 1, 1], [5, 1, 0, 0]]
         ),
         (
             "0 1 2 2 0 0",
             "0.5 0 0.5 1.5 1.5 0.5",
+            "1 1 1 1",
+            "1 1 1 1",
+            "1 1 1 1",
+            "1 1 1 1",
+            "1 1 1 1",
             2.0,
             1.5,
             [[0, 0], [0.5, 2.0], [1.5, 2.0], [1.5, 0.0]],
+            None
         ),
-        ("0 0.5 0.75 1.0 1.5", "0.5 0 0 0 0.5", 1.5, 0.5, [[0, 0.5], [0.5, 1.5]]),
-        ("0 1 0 1 0", "0 1 1 0 0", 1, 1, [[0, 1], [0.5, 0], [1, 1], [1, 0]]),
+        ("0 0.5 0.75 1.0 1.5", "0.5 0 0 0 0.5", "1 1 1 1", "1 1 1 1", "1 1 1 1", "1 1 1 1", "1 1 1 1", 1.5, 0.5, [[0, 0.5], [0.5, 1.5]], [[0, 0.5, 1, 1], [0.5, 0, 1, 1], [0.75, 0, 1, 1], [1, 0, 1, 1], [1.5, 0.5, 0, 0]]),
+        ("0 1 0 1 0", "0 1 1 0 0", "1 1 1 1", "1 1 1 1", "1 1 1 1", "1 1 1 1", "1 1 1 1", 1, 1, [[0, 1], [0.5, 0], [1, 1], [1, 0]], None),
     ],
 )
-def test_tabulate_yz(width, height, exp_width, exp_height, exp_table):
-    shape, width_1d, height_1d, table = tabulate_yz("my-shape", width, height)
+def test_tabulate_yz(width, height, friction_values,vegetation_stem_densities,vegetation_stem_diameters,vegetation_heights,vegetation_drag_coefficients,exp_width, exp_height, exp_table, exp_yz):
+    shape, width_1d, height_1d, table, yz = tabulate_yz("my-shape", width, height,friction_values,vegetation_stem_densities,vegetation_stem_diameters,vegetation_heights,vegetation_drag_coefficients)
 
     assert shape == CrossSectionShape.TABULATED_TRAPEZIUM
     assert width_1d == exp_width
     assert height_1d == exp_height
     assert_almost_equal(table, np.array(exp_table, dtype=float))
+    if yz is not None:
+        assert_almost_equal(yz, np.array(exp_yz, dtype=float))
 
 
 @pytest.mark.parametrize(
@@ -247,4 +271,4 @@ def test_tabulate_yz(width, height, exp_width, exp_height, exp_table):
 )
 def test_tabulate_yz_err(width, height, match):
     with pytest.raises(SchematisationError, match=match):
-        tabulate_yz("my-shape", width, height)
+        tabulate_yz("my-shape", width, height, None, None, None, None, None)