From 3c335313cb2f61ed2c49f666086f4297d3564e74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Fri, 17 Nov 2023 21:57:18 +0100
Subject: [PATCH] Add tests from plot.flopy, gwt and hfb

---
 nlmod/gwf/horizontal_flow_barrier.py |  3 +++
 nlmod/plot/flopy.py                  |  4 +++
 tests/test_012_plot.py               | 23 +++++++++++++++++
 tests/test_022_gwt.py                |  8 ++++++
 tests/test_023_hfb.py                | 37 ++++++++++++++++++++++++++++
 tests/util.py                        |  9 +++++++
 6 files changed, 84 insertions(+)
 create mode 100644 tests/test_022_gwt.py
 create mode 100644 tests/test_023_hfb.py

diff --git a/nlmod/gwf/horizontal_flow_barrier.py b/nlmod/gwf/horizontal_flow_barrier.py
index d6e9b557..0c6cc07a 100644
--- a/nlmod/gwf/horizontal_flow_barrier.py
+++ b/nlmod/gwf/horizontal_flow_barrier.py
@@ -46,6 +46,9 @@ def get_hfb_spd(gwf, linestrings, hydchr=1 / 100, depth=None, elevation=None):
 
     cells = line2hfb(linestrings, gwf)
 
+    # drop cells on the edge of the model
+    cells = [x for x in cells if len(x) > 1]
+
     spd = []
 
     # hydchr = 1 / 100  # resistance of 100 days
diff --git a/nlmod/plot/flopy.py b/nlmod/plot/flopy.py
index b6be4b43..0efc2f90 100644
--- a/nlmod/plot/flopy.py
+++ b/nlmod/plot/flopy.py
@@ -450,6 +450,8 @@ def facet_plot(
                 if iper is None:
                     raise ValueError("Pass 'period' to select timestep to plot.")
                 a = arr[iper]
+            else:
+                a = arr
         elif plot_dim == "time":
             ilay = layer
             iper = i
@@ -457,6 +459,8 @@ def facet_plot(
                 if ilay is None:
                     raise ValueError("Pass 'layer' to select layer to plot.")
                 a = arr[iper]
+            else:
+                a = arr
         else:
             raise ValueError("'plot_dim' must be one of ['layer', 'time']")
 
diff --git a/tests/test_012_plot.py b/tests/test_012_plot.py
index 23d14271..7dd08e80 100644
--- a/tests/test_012_plot.py
+++ b/tests/test_012_plot.py
@@ -29,3 +29,26 @@ def test_plot_data_array_vertex():
 
 def test_plot_get_map():
     nlmod.plot.get_map([100000, 101000, 400000, 401000], background=True, figsize=3)
+
+
+def test_map_array():
+    ds = util.get_ds_structured()
+    gwf = util.get_gwf(ds)
+    nlmod.plot.flopy.map_array(ds["kh"], gwf, ilay=0)
+
+
+def test_flopy_contour_array():
+    ds = util.get_ds_structured()
+    gwf = util.get_gwf(ds)
+    nlmod.plot.flopy.contour_array(ds["kh"], gwf, ilay=0)
+
+
+def test_flopy_animate_map():
+    # no test implemented yet
+    pass
+
+
+def test_flopy_facet_plot():
+    ds = util.get_ds_structured()
+    gwf = util.get_gwf(ds)
+    nlmod.plot.flopy.facet_plot(gwf, ds["kh"])
diff --git a/tests/test_022_gwt.py b/tests/test_022_gwt.py
new file mode 100644
index 00000000..aaa017b2
--- /dev/null
+++ b/tests/test_022_gwt.py
@@ -0,0 +1,8 @@
+import nlmod
+
+
+def test_prepare():
+    ds = nlmod.get_ds([0, 1000, 0, 2000])
+    ds = nlmod.gwt.prepare.set_default_transport_parameters(
+        ds, transport_type="chloride"
+    )
diff --git a/tests/test_023_hfb.py b/tests/test_023_hfb.py
new file mode 100644
index 00000000..f7e2a73f
--- /dev/null
+++ b/tests/test_023_hfb.py
@@ -0,0 +1,37 @@
+from shapely.geometry import LineString, Polygon
+import geopandas as gpd
+import flopy
+import nlmod
+import util
+
+
+def test_get_hfb_spd():
+    # this test also tests line2hfb
+    ds = util.get_ds_vertex()
+    ds = nlmod.time.set_ds_time(ds, "2023", time="2024")
+    gwf = util.get_gwf(ds)
+
+    coords = [(0, 1000), (1000, 0)]
+    gdf = gpd.GeoDataFrame({"geometry": [LineString(coords)]})
+
+    spd = nlmod.gwf.horizontal_flow_barrier.get_hfb_spd(gwf, gdf, depth=5.0)
+    hfb = flopy.mf6.ModflowGwfhfb(gwf, stress_period_data={0: spd})
+
+    # also test the plot method
+    ax = gdf.plot()
+    nlmod.gwf.horizontal_flow_barrier.plot_hfb(hfb, gwf, ax=ax)
+
+
+def test_polygon_to_hfb():
+    ds = util.get_ds_vertex()
+    ds = nlmod.time.set_ds_time(ds, "2023", time="2024")
+    gwf = util.get_gwf(ds)
+
+    coords = [(135, 230), (568, 170), (778, 670), (260, 786)]
+    gdf = gpd.GeoDataFrame({"geometry": [Polygon(coords)]}).reset_index()
+
+    hfb = nlmod.gwf.horizontal_flow_barrier.polygon_to_hfb(gdf, ds, gwf=gwf)
+
+    # also test the plot method
+    ax = gdf.plot()
+    nlmod.gwf.horizontal_flow_barrier.plot_hfb(hfb, gwf, ax=ax)
diff --git a/tests/util.py b/tests/util.py
index 67957a42..3d2c3265 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -21,3 +21,12 @@ def get_ds_vertex(extent=None, line=None, **kwargs):
     refinement_features = [([LineString(line)], "line", 1)]
     ds = nlmod.grid.refine(ds, model_ws, refinement_features=refinement_features)
     return ds
+
+
+def get_gwf(ds):
+    sim = nlmod.sim.sim(ds)
+    if "time" in ds.variables:
+        nlmod.sim.tdis(ds, sim)
+    gwf = nlmod.gwf.gwf(ds, sim)
+    nlmod.gwf.dis(ds, gwf)
+    return gwf