update model dis/disv

flopy4/mf6/gwe/__init__.py

@@ -1,26 +1,42 @@
 from pathlib import Path
 from typing import Optional
 
+from flopy.discretization.structuredgrid import StructuredGrid
+from flopy.discretization.vertexgrid import VertexGrid
 from xattree import xattree
 
 from flopy4.mf6.gwe.adv import Adv
 from flopy4.mf6.gwe.cnd import Cnd
 from flopy4.mf6.gwe.ctp import Ctp
 from flopy4.mf6.gwe.dis import Dis
+from flopy4.mf6.gwe.disv import Disv
 from flopy4.mf6.gwe.esl import Esl
 from flopy4.mf6.gwe.est import Est
 from flopy4.mf6.gwe.ic import Ic
 from flopy4.mf6.gwe.lke import Lke
 from flopy4.mf6.gwe.mve import Mve
 from flopy4.mf6.gwe.oc import Oc
 from flopy4.mf6.gwe.ssm import Ssm
+from flopy4.mf6.gwf.disbase import DisBase
 from flopy4.mf6.model import Model
 from flopy4.mf6.spec import field, path
 from flopy4.utils import to_path
 
+
+def convert_grid(value):
+    if isinstance(value, StructuredGrid):
+        return Dis.from_grid(value)
+    if isinstance(value, VertexGrid):
+        return Disv.from_grid(value)
+    if isinstance(value, (Dis, Disv)) or value is None:
+        return value
+    raise TypeError(f"Expected Grid or Dis/Disv, got {type(value)}")
+
+
 __all__ = [
     "Gwe",
     "Dis",
+    "Disv",
     "Adv",
     "Cnd",
     "Ctp",
@@ -50,7 +66,7 @@ class Gwe(Model):
     netcdf_input_file: Optional[Path] = path(
         block="options", default=None, converter=to_path, inout="filein"
     )
-    dis: Dis | None = field(block="packages", default=None)
+    dis: DisBase | None = field(converter=convert_grid, block="packages", default=None)
     ic: Ic | None = field(block="packages", default=None)
     oc: Oc | None = field(block="packages", default=None)
     adv: Adv | None = field(block="packages", default=None)
@@ -61,3 +77,8 @@ class Gwe(Model):
     lke: list[Lke] = field(block="packages")
     ssm: Ssm | None = field(block="packages", default=None)
     mve: Mve | None = field(block="packages", default=None)
+
+    @property
+    def grid(self):
+        if self.dis is not None:
+            return self.dis.to_grid()

flopy4/mf6/gwe/dis.py

@@ -1,3 +1,183 @@
-from flopy4.mf6.gwf.dis import Dis  # noqa: F401
+from pathlib import Path
+from typing import Optional
 
-__all__ = ["Dis"]
+import attrs
+import numpy as np
+from numpy.typing import NDArray
+
+from flopy4.mf6._types import _optional_path
+from flopy4.mf6.gwf.disbase import DisBase
+from flopy4.mf6.utils.grid import StructuredGrid
+from flopy4.mf6.utl.ncf import Ncf
+
+
+@attrs.define(kw_only=True, slots=False)
+class Dis(DisBase):
+    length_units: Optional[str] = attrs.field(
+        default=None,
+        metadata={"dfn_block": "options", "dfn_type": "string", "optional": True},
+    )
+    nogrb: bool = attrs.field(
+        default=False,
+        metadata={"dfn_block": "options", "dfn_type": "keyword", "optional": True},
+    )
+    xorigin: float = attrs.field(
+        default=0.0,
+        metadata={"dfn_block": "options", "dfn_type": "double", "optional": True},
+    )
+    yorigin: float = attrs.field(
+        default=0.0,
+        metadata={"dfn_block": "options", "dfn_type": "double", "optional": True},
+    )
+    angrot: Optional[float] = attrs.field(
+        default=None,
+        metadata={"dfn_block": "options", "dfn_type": "double", "optional": True},
+    )
+    export_array_netcdf: bool = attrs.field(
+        default=False,
+        metadata={"dfn_block": "options", "dfn_type": "keyword", "optional": True},
+    )
+    crs: Optional[str] = attrs.field(
+        default=None,
+        metadata={"dfn_block": "options", "dfn_type": "string", "optional": True},
+    )
+    ncf6_filerecord: Optional[Path] = attrs.field(
+        default=None,
+        converter=_optional_path,
+        metadata={
+            "dfn_block": "options",
+            "dfn_type": "record",
+            "optional": True,
+            "inout": "filein",
+        },
+    )
+    ncf: Optional[Ncf] = attrs.field(default=None)
+    nlay: int = attrs.field(
+        default=1,
+        metadata={"dfn_block": "dimensions", "dfn_type": "integer"},
+    )
+    ncol: int = attrs.field(
+        default=2,
+        metadata={"dfn_block": "dimensions", "dfn_type": "integer"},
+    )
+    nrow: int = attrs.field(
+        default=2,
+        metadata={"dfn_block": "dimensions", "dfn_type": "integer"},
+    )
+    delr: NDArray[np.float64] = attrs.field(
+        default=1.0,
+        metadata={
+            "dfn_block": "griddata",
+            "dfn_type": "double",
+            "shape": ("ncol",),
+            "layered": False,
+            "netcdf": True,
+        },
+    )  # type: ignore[assignment]
+    delc: NDArray[np.float64] = attrs.field(
+        default=1.0,
+        metadata={
+            "dfn_block": "griddata",
+            "dfn_type": "double",
+            "shape": ("nrow",),
+            "layered": False,
+            "netcdf": True,
+        },
+    )  # type: ignore[assignment]
+    top: NDArray[np.float64] = attrs.field(
+        default=1.0,
+        metadata={
+            "dfn_block": "griddata",
+            "dfn_type": "double",
+            "shape": ("ncpl",),
+            "layered": False,
+            "netcdf": True,
+        },
+    )  # type: ignore[assignment]
+    botm: NDArray[np.float64] = attrs.field(
+        default=0.0,
+        metadata={
+            "dfn_block": "griddata",
+            "dfn_type": "double",
+            "shape": ("nodes",),
+            "layered": True,
+            "netcdf": True,
+        },
+    )  # type: ignore[assignment]
+    idomain: Optional[NDArray[np.int64]] = attrs.field(
+        default=1,
+        metadata={
+            "dfn_block": "griddata",
+            "dfn_type": "integer",
+            "shape": ("nodes",),
+            "layered": True,
+            "netcdf": True,
+        },
+    )  # type: ignore[assignment]
+
+    def __attrs_post_init__(self):
+        self.nodes = self.ncol * self.nrow * self.nlay
+        self.ncpl = self.ncol * self.nrow
+        self.nvert = (self.ncol + 1) * (self.nrow + 1)
+        self._coerce_griddata()
+        super().__attrs_post_init__()
+
+    def get_dims(self) -> dict[str, int]:
+        """Get all dimensions."""
+        return {
+            "nlay": self.nlay,
+            "nrow": self.nrow,
+            "ncol": self.ncol,
+            "nodes": self.nlay * self.nrow * self.ncol,
+            "ncpl": self.nrow * self.ncol,
+        }
+
+    def to_grid(self) -> StructuredGrid:
+        """Convert the discretization to a `StructuredGrid`."""
+        # Reshape flat arrays to grid shape for StructuredGrid constructor.
+        top = np.asarray(self.top).reshape(self.nrow, self.ncol)
+        botm = np.asarray(self.botm).reshape(self.nlay, self.nrow, self.ncol)
+        idomain = (
+            np.asarray(self.idomain).reshape(self.nlay, self.nrow, self.ncol)
+            if self.idomain is not None
+            else None
+        )
+        return StructuredGrid(
+            length_units=self.length_units,
+            xoff=self.xorigin,
+            yoff=self.yorigin,
+            nlay=self.nlay,
+            nrow=self.nrow,
+            ncol=self.ncol,
+            delr=np.asarray(self.delr),
+            delc=np.asarray(self.delc),
+            top=top,
+            botm=botm,
+            idomain=idomain,
+            angrot=self.angrot,
+            crs=self.crs,
+        )
+
+    @classmethod
+    def from_grid(cls, grid: StructuredGrid) -> "Dis":
+        """Create a discretization from a `StructuredGrid`."""
+        _lenunits = {1: "FEET", 2: "METERS", 3: "CENTIMETERS"}
+        kwargs = {
+            "xorigin": grid.xoffset,
+            "yorigin": grid.yoffset,
+            "nlay": grid.nlay,
+            "nrow": grid.nrow,
+            "ncol": grid.ncol,
+            "delr": grid.delr,
+            "delc": grid.delc,
+            "top": grid.top,
+            "botm": grid.botm,
+            "idomain": grid.idomain,
+        }
+        if grid.lenuni in _lenunits:
+            kwargs["length_units"] = _lenunits[grid.lenuni]
+        if grid.angrot:
+            kwargs["angrot"] = grid.angrot
+        if grid.crs is not None:
+            kwargs["crs"] = f"EPSG:{grid.crs.to_epsg()}"
+        return Dis(**kwargs)