From 722224c4b438116528b83a77338e34d6f9807ef7 Mon Sep 17 00:00:00 2001
From: Chris Nicol <37314969+cnicol-gwlogic@users.noreply.github.com>
Date: Sat, 4 Sep 2021 10:30:25 +1000
Subject: [PATCH] fix(shapefile_utils) unstructured shapefile export (#1220)
 (#1222)

* Fix errors encountered when exporting to shapefile from unstructured models. See https://github.com/modflowpy/flopy/issues/1220
* add unstructured clauses to export/shapefile_utils.py for static arrays (e.g. DISU, LPF, BAS6). More work needed for boundary conditions (MFList and transientlist data types - eg RIV/RCH/EVT packages)
* add test to t506.py to export mfusg aunstructured model shapefiles
* close #1220
---
 autotest/t506_test.py           | 18 ++++++++++++
 flopy/export/shapefile_utils.py | 52 +++++++++++++++++++++++----------
 flopy/modflow/mfdisu.py         |  2 +-
 3 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/autotest/t506_test.py b/autotest/t506_test.py
index c40e61df6..b2d1b9349 100644
--- a/autotest/t506_test.py
+++ b/autotest/t506_test.py
@@ -471,6 +471,24 @@ def test_mfusg():
             and "NOCVCORRECTION" in lpf2.options.upper()
         ), msg
 
+    # test disu, bas6, lpf shapefile export for mfusg unstructured models
+    try:
+        m.disu.export(os.path.join(ws, f"{name}_disu.shp"))
+    except:
+        raise AssertionError("Error exporting mfusg disu to shapefile.")
+    try:
+        m.bas6.export(os.path.join(ws, f"{name}_bas6.shp"))
+    except:
+        raise AssertionError("Error exporting mfusg bas6 to shapefile.")
+    try:
+        m.lpf.export(os.path.join(ws, f"{name}_lpf.shp"))
+    except:
+        raise AssertionError("Error exporting mfusg lpf to shapefile.")
+    try:
+        m.export(os.path.join(ws, f"{name}.shp"))
+    except:
+        raise AssertionError("Error exporting mfusg model to shapefile.")
+
     return
 
 
diff --git a/flopy/export/shapefile_utils.py b/flopy/export/shapefile_utils.py
index c3b30e69b..fd9bf5871 100755
--- a/flopy/export/shapefile_utils.py
+++ b/flopy/export/shapefile_utils.py
@@ -179,15 +179,35 @@ def write_grid_shapefile(
         names = enforce_10ch_limit(names)
 
     elif mg.grid_type == "unstructured":
-        names = ["node"] + list(array_dict.keys())
-        dtypes = [("node", np.dtype("int"))] + [
-            (enforce_10ch_limit([name])[0], array_dict[name].dtype)
-            for name in names[1:]
-        ]
-        node = list(range(1, mg.nnodes + 1))
-        at = np.vstack(
-            [node] + [array_dict[name].ravel() for name in names[1:]]
-        ).transpose()
+        if mg.nlay is None:
+            names = ["node"] + list(array_dict.keys())
+            dtypes = [("node", np.dtype("int"))] + [
+                (enforce_10ch_limit([name])[0], array_dict[name].dtype)
+                for name in names[1:]
+            ]
+            node = list(range(1, mg.nnodes + 1))
+            at = np.vstack(
+                [node] + [array_dict[name].ravel() for name in names[1:]]
+            ).transpose()
+        else:
+            names = ["node", "layer"] + list(array_dict.keys())
+            dtypes = [
+                ("node", np.dtype("int")),
+                ("layer", np.dtype("int")),
+            ] + [
+                (enforce_10ch_limit([name])[0], array_dict[name].dtype)
+                for name in names[2:]
+            ]
+            node = list(range(1, mg.nnodes + 1))
+            layer = np.zeros(mg.nnodes)
+            for ilay in range(mg.nlay):
+                istart, istop = mg.get_layer_node_range(ilay)
+                layer[istart:istop] = ilay + 1
+            at = np.vstack(
+                [node]
+                + [layer]
+                + [array_dict[name].ravel() for name in names[2:]]
+            ).transpose()
 
         names = enforce_10ch_limit(names)
 
@@ -274,11 +294,6 @@ def model_attributes_to_shapefile(
     else:
         grid = ml.modelgrid
 
-    if grid.grid_type == "USG-Unstructured":
-        raise Exception(
-            "Flopy does not support exporting to shapefile from "
-            "and MODFLOW-USG unstructured grid."
-        )
     horz_shape = grid.get_plottable_layer_shape()
     for pname in package_names:
         pak = ml.get_package(pname)
@@ -317,7 +332,14 @@ def model_attributes_to_shapefile(
                         continue
 
                     if a.array.shape == horz_shape:
-                        array_dict[a.name] = a.array
+                        if hasattr(a, "shape"):
+                            if a.shape[1] is None:  # usg unstructured Util3d
+                                # return a flattened array, with a.name[0] (a per-layer list)
+                                array_dict[a.name[0]] = a.array
+                            else:
+                                array_dict[a.name] = a.array
+                        else:
+                            array_dict[a.name] = a.array
                     else:
                         # array is not the same shape as the layer shape
                         for ilay in range(a.array.shape[0]):
diff --git a/flopy/modflow/mfdisu.py b/flopy/modflow/mfdisu.py
index 20f77fa3f..1506586d1 100644
--- a/flopy/modflow/mfdisu.py
+++ b/flopy/modflow/mfdisu.py
@@ -509,7 +509,7 @@ def zcentroids(self):
 
         """
         z = np.empty((self.nodes))
-        z[:] = (self.top[:] - self.bot[:]) / 2.0
+        z[:] = (self.top.array - self.bot.array) / 2.0
         return z
 
     @property