From 693f01a0ce41f08d05395832b540fcc4f7dcff43 Mon Sep 17 00:00:00 2001
From: martclanor <mart.dayniel.clanor@witteveenbos.com>
Date: Mon, 16 Sep 2024 17:43:03 +0200
Subject: [PATCH] fix(array3d_export): fix exporting of array3d to shp (#2310)

Fix the issue raised in discussion #2308

- fix(MultiList): increment index per layer in build_list
- fix(array3d_export): fix case for disu grid
- add tests
---
 autotest/test_export.py | 59 +++++++++++++++++++++++++++++++++++++++++
 flopy/export/utils.py   | 26 +++++++++++-------
 flopy/utils/datautil.py |  4 ++-
 3 files changed, 78 insertions(+), 11 deletions(-)

diff --git a/autotest/test_export.py b/autotest/test_export.py
index 0028ada34..a4ac22f78 100644
--- a/autotest/test_export.py
+++ b/autotest/test_export.py
@@ -646,6 +646,65 @@ def test_export_array2(function_tmpdir):
     assert os.path.isfile(filename), "did not create array shapefile"
 
 
+@pytest.mark.mf6
+@requires_pkg("pyshp", name_map={"pyshp": "shapefile"})
+def test_array3d_export_structured(function_tmpdir):
+    from shapefile import Reader
+
+    xll, yll = 468970, 3478635
+    xur, yur = 681010, 3716462
+    spacing = 20000
+    ncol = int((xur - xll) / spacing)
+    nrow = int((yur - yll) / spacing)
+    sim = flopy.mf6.MFSimulation("sim", sim_ws=function_tmpdir)
+    gwf = flopy.mf6.ModflowGwf(
+        sim,
+        modelname="array3d_export_unstructured",
+    )
+    flopy.mf6.ModflowGwfdis(
+        gwf,
+        nlay=3,
+        top=5,
+        botm=[4, 3, 2],
+        delr=spacing,
+        delc=spacing,
+        nrow=nrow,
+        ncol=ncol,
+    )
+
+    shp_file = os.path.join(function_tmpdir, "dis_botm.shp")
+    gwf.dis.botm.export(shp_file)
+
+    with Reader(shp_file) as shp:
+        assert list(shp.shapeRecord(-1).record) == [
+            110,  # node
+            11,  # row
+            10,  # column
+            4.0,  # botm_1
+            3.0,  # botm_2
+            2.0,  # botm_3
+        ]
+
+
+@requires_pkg("pyshp", name_map={"pyshp": "shapefile"})
+def test_array3d_export_unstructured(function_tmpdir):
+    from shapefile import Reader
+
+    name = "array3d_export_unstructured"
+    sim = disu_sim(name, function_tmpdir)
+    gwf = sim.get_model(name)
+
+    shp_file = function_tmpdir / "disu_bot.shp"
+    gwf.disu.bot.export(shp_file)
+
+    with Reader(shp_file) as shp:
+        assert list(shp.shapeRecord(-1).record) == [
+            1770,  # node
+            3,  # layer
+            0.0,  # bot
+        ]
+
+
 @requires_pkg("pyshp", "shapely", name_map={"pyshp": "shapefile"})
 def test_export_array_contours_structured(function_tmpdir):
     nrow = 7
diff --git a/flopy/export/utils.py b/flopy/export/utils.py
index d0b577398..ef446b765 100644
--- a/flopy/export/utils.py
+++ b/flopy/export/utils.py
@@ -1207,16 +1207,22 @@ def array3d_export(f: Union[str, os.PathLike], u3d, fmt=None, **kwargs):
         f
     ).suffix.lower() == ".shp":
         array_dict = {}
-        for ilay in range(modelgrid.nlay):
-            u2d = u3d[ilay]
-            if isinstance(u2d, np.ndarray):
-                dname = u3d.name
-                array = u2d
-            else:
-                dname = u2d.name
-                array = u2d.array
-            name = f"{shapefile_utils.shape_attr_name(dname)}_{ilay + 1}"
-            array_dict[name] = array
+        array_shape = u3d.array.shape
+
+        if len(array_shape) == 1:
+            name = shapefile_utils.shape_attr_name(u3d.name)
+            array_dict[name] = u3d.array
+        else:
+            for ilay in range(array_shape[0]):
+                u2d = u3d[ilay]
+                if isinstance(u2d, np.ndarray):
+                    dname = u3d.name
+                    array = u2d
+                else:
+                    dname = u2d.name
+                    array = u2d.array
+                name = f"{shapefile_utils.shape_attr_name(dname)}_{ilay + 1}"
+                array_dict[name] = array
         shapefile_utils.write_grid_shapefile(f, modelgrid, array_dict)
 
     elif isinstance(f, NetCdf) or isinstance(f, dict):
diff --git a/flopy/utils/datautil.py b/flopy/utils/datautil.py
index aa59a0976..0fba1bdb4 100644
--- a/flopy/utils/datautil.py
+++ b/flopy/utils/datautil.py
@@ -606,7 +606,9 @@ def build_list(self, callback):
                             (entry_point[0][-1], new_location)
                         )
                     else:
-                        entry_point[0].append(callback(entry_point[1]))
+                        entry_point[0].append(
+                            callback(tuple(i + val for i in entry_point[1]))
+                        )
             entry_points = new_entry_points
 
     def first_item(self):