fix(modelgrid): add more support for mf6 surface water models (modflo…

…wpy#2295)

* fix(modelgrid): add more support for mf6 surface water models

* ruff

flopy/discretization/structuredgrid.py

@@ -192,14 +192,18 @@ def __init__(
         if top is not None:
             assert self.__nrow * self.__ncol == len(np.ravel(top))
         if botm is not None:
-            assert self.__nrow * self.__ncol == len(np.ravel(botm[0]))
-            if nlay is not None:
-                self.__nlay = nlay
-            else:
-                if laycbd is not None:
-                    self.__nlay = len(botm) - np.count_nonzero(laycbd)
+            if botm.ndim == 3:
+                assert self.__nrow * self.__ncol == len(np.ravel(botm[0]))
+                if nlay is not None:
+                    self.__nlay = nlay
                 else:
-                    self.__nlay = len(botm)
+                    if laycbd is not None:
+                        self.__nlay = len(botm) - np.count_nonzero(laycbd)
+                    else:
+                        self.__nlay = len(botm)
+            elif botm.ndim == 2:
+                assert botm.shape == (self.__nrow, self.__ncol)
+                self.__nlay = 1
         else:
             self.__nlay = nlay
         if laycbd is not None:

flopy/discretization/vertexgrid.py

@@ -140,7 +140,10 @@ def ncpl(self):
         if self._cell1d is not None:
             return len(self._cell1d)
         if self._botm is not None:
-            return len(self._botm[0])
+            if self._botm.ndim == 2:  # (nlay, ncpl)
+                return self._botm.shape[1]
+            elif self._botm.ndim == 1:  # (ncpl,)
+                return self._botm.shape[0]
         if self._cell2d is not None and self._nlay is None:
             return len(self._cell2d)
         else:

flopy/mf6/mfmodel.py

@@ -508,15 +508,53 @@ def modelgrid(self):
                         angrot=self._modelgrid.angrot,
                     )
             else:
-                botm = dis.botm.array
+                botm = dis.bottom.array
                 idomain = dis.idomain.array
                 if idomain is None:
                     force_resync = True
                     idomain = self._resolve_idomain(idomain, botm)
                 self._modelgrid = VertexGrid(
                     vertices=dis.vertices.array,
                     cell1d=dis.cell1d.array,
-                    top=dis.top.array,
+                    top=None,
+                    botm=botm,
+                    idomain=idomain,
+                    lenuni=dis.length_units.array,
+                    crs=self._modelgrid.crs,
+                    xoff=self._modelgrid.xoffset,
+                    yoff=self._modelgrid.yoffset,
+                    angrot=self._modelgrid.angrot,
+                )
+        elif self.get_grid_type() == DiscretizationType.DIS2D:
+            dis = self.get_package("dis2d")
+            if not hasattr(dis, "_init_complete"):
+                if not hasattr(dis, "delr"):
+                    # dis package has not yet been initialized
+                    return self._modelgrid
+                else:
+                    # dis package has been partially initialized
+                    self._modelgrid = StructuredGrid(
+                        delc=dis.delc.array,
+                        delr=dis.delr.array,
+                        top=None,
+                        botm=None,
+                        idomain=None,
+                        lenuni=None,
+                        crs=self._modelgrid.crs,
+                        xoff=self._modelgrid.xoffset,
+                        yoff=self._modelgrid.yoffset,
+                        angrot=self._modelgrid.angrot,
+                    )
+            else:
+                botm = dis.bottom.array
+                idomain = dis.idomain.array
+                if idomain is None:
+                    force_resync = True
+                    idomain = self._resolve_idomain(idomain, botm)
+                self._modelgrid = StructuredGrid(
+                    delc=dis.delc.array,
+                    delr=dis.delr.array,
+                    top=None,
                     botm=botm,
                     idomain=idomain,
                     lenuni=dis.length_units.array,
@@ -546,15 +584,15 @@ def modelgrid(self):
                         angrot=self._modelgrid.angrot,
                     )
             else:
-                botm = dis.botm.array
+                botm = dis.bottom.array
                 idomain = dis.idomain.array
                 if idomain is None:
                     force_resync = True
                     idomain = self._resolve_idomain(idomain, botm)
                 self._modelgrid = VertexGrid(
                     vertices=dis.vertices.array,
                     cell2d=dis.cell2d.array,
-                    top=dis.top.array,
+                    top=None,
                     botm=botm,
                     idomain=idomain,
                     lenuni=dis.length_units.array,
@@ -1247,6 +1285,13 @@ def get_grid_type(self):
             is not None
         ):
             return DiscretizationType.DISV1D
+        elif (
+            package_recarray.search_data(
+                f"dis2d{structure.get_version_string()}", 0
+            )
+            is not None
+        ):
+            return DiscretizationType.DIS2D
         elif (
             package_recarray.search_data(
                 f"disv2d{structure.get_version_string()}", 0