diff --git a/flopy4/mf6/gwe/__init__.py b/flopy4/mf6/gwe/__init__.py index 230fbbad..da65b54a 100644 --- a/flopy4/mf6/gwe/__init__.py +++ b/flopy4/mf6/gwe/__init__.py @@ -1,12 +1,15 @@ 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 @@ -14,13 +17,26 @@ 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() diff --git a/flopy4/mf6/gwe/dis.py b/flopy4/mf6/gwe/dis.py index 0c41fa48..defd204b 100644 --- a/flopy4/mf6/gwe/dis.py +++ b/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) diff --git a/flopy4/mf6/gwe/disv.py b/flopy4/mf6/gwe/disv.py new file mode 100644 index 00000000..36f764b6 --- /dev/null +++ b/flopy4/mf6/gwe/disv.py @@ -0,0 +1,266 @@ +from pathlib import Path +from typing import ClassVar, Optional + +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.schema import Column, Schema +from flopy4.mf6.utils.grid import VertexGrid +from flopy4.mf6.utl.ncf import Ncf + + +@attrs.define(kw_only=True, slots=False) +class Disv(DisBase): + @attrs.define(slots=False) + class Cell2dRecord: + icell2d: int = attrs.field() + xc: float = attrs.field() + yc: float = attrs.field() + ncvert: int = attrs.field() + icvert: tuple[int, ...] = attrs.field() + + class _VerticesSchema(Schema): + iv = Column("iv", role="value", dfn_type="integer") + xv = Column("xv", role="value", dfn_type="double") + yv = Column("yv", role="value", dfn_type="double") + + __vertices_schema__: ClassVar[type[Schema]] = _VerticesSchema + + 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: Optional[float] = attrs.field( + default=None, + metadata={"dfn_block": "options", "dfn_type": "double", "optional": True}, + ) + yorigin: Optional[float] = attrs.field( + default=None, + 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=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + ncpl: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + nvert: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + top: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("ncpl",), + "layered": False, + "netcdf": True, + }, + ) # type: ignore[assignment] + botm: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("nodes",), + "layered": True, + "netcdf": True, + }, + ) # type: ignore[assignment] + idomain: Optional[NDArray[np.int64]] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "integer", + "shape": ("nodes",), + "layered": True, + "netcdf": True, + }, + ) # type: ignore[assignment] + # User-facing parallel arrays for vertices. + iv: Optional[NDArray[np.int64]] = attrs.field(default=None) # type: ignore[assignment] + xv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + yv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + # Combined vertices recarray for the codec (built in __attrs_post_init__). + vertices: Optional[np.recarray] = attrs.field( + default=None, + metadata={"dfn_block": "vertices", "schema": "__vertices_schema__"}, + ) # type: ignore[assignment] + # Cell2d data — list of Cell2dRecord objects (user-facing). + cell2ddata: Optional[list] = attrs.field(default=None) + # Pre-formatted cell2d rows for the codec (built in __attrs_post_init__). + cell2d: Optional[list] = attrs.field( + default=None, + init=False, + metadata={"dfn_block": "cell2d"}, + ) + + def __attrs_post_init__(self): + # Coerce list inputs to numpy arrays for vertices. + if self.iv is not None and not isinstance(self.iv, np.ndarray): + object.__setattr__(self, "iv", np.asarray(self.iv, dtype=np.int64)) + if self.xv is not None and not isinstance(self.xv, np.ndarray): + object.__setattr__(self, "xv", np.asarray(self.xv, dtype=np.float64)) + if self.yv is not None and not isinstance(self.yv, np.ndarray): + object.__setattr__(self, "yv", np.asarray(self.yv, dtype=np.float64)) + # Build combined vertices recarray for the codec. + if self.iv is not None and self.xv is not None and self.yv is not None: + dtype = np.dtype([("iv", np.int64), ("xv", np.float64), ("yv", np.float64)]) + n = len(self.iv) + arr = np.zeros(n, dtype=dtype) + arr["iv"] = self.iv + 1 # MF6 uses 1-based vertex IDs + arr["xv"] = self.xv + arr["yv"] = self.yv + object.__setattr__(self, "vertices", arr.view(np.recarray)) + # Build cell2d list of tuples for the codec. + if self.cell2ddata is not None: + rows = [] + for rec in self.cell2ddata: + row = (rec.icell2d + 1, rec.xc, rec.yc, rec.ncvert) + tuple( + v + 1 for v in rec.icvert + ) + rows.append(row) + object.__setattr__(self, "cell2d", rows) + # Set derived dimensions. + self.nodes = self.ncpl * self.nlay + self.nrow = 0 + self.ncol = 0 + self._coerce_griddata() + super().__attrs_post_init__() + + def get_dims(self) -> dict[str, int]: + """Get all dimensions.""" + return { + "nlay": self.nlay, + "ncpl": self.ncpl, + "nvert": self.nvert, + "nodes": self.nlay * self.ncpl, + } + + def to_grid(self) -> VertexGrid: + """Convert the discretization to a `VertexGrid`.""" + vertices = [] + for i in range(len(self.iv)): # type: ignore[arg-type] + vertices.append([self.iv[i], self.xv[i], self.yv[i]]) # type: ignore[index] + # VertexGrid expects top as 1D (ncpl,) and botm as 2D (nlay, ncpl). + botm = ( + self.botm.reshape(self.nlay, self.ncpl) + if isinstance(self.botm, np.ndarray) + else self.botm + ) + idomain = ( + self.idomain.reshape(self.nlay, self.ncpl) + if isinstance(self.idomain, np.ndarray) and self.idomain is not None + else self.idomain + ) + return VertexGrid( + length_units=self.length_units, + xoff=self.xorigin, + yoff=self.yorigin, + angrot=self.angrot, + crs=self.crs, + nlay=self.nlay, + ncpl=self.ncpl, + top=self.top, + botm=botm, + idomain=idomain, + vertices=vertices, + cell2d=Disv.disv_to_grid_cell2d(self.cell2ddata), + ) + + @classmethod + def from_grid(cls, grid: VertexGrid) -> "Disv": + """Create a discretization from a `VertexGrid`.""" + _lenunits = {1: "FEET", 2: "METERS", 3: "CENTIMETERS"} + kwargs = { + "xorigin": grid.xoffset, + "yorigin": grid.yoffset, + "nlay": grid.nlay, + "ncpl": grid.ncpl, + "nvert": grid.nvert, + "top": grid.top, + "botm": grid.botm, + "idomain": np.asarray(grid.idomain).reshape(grid.nlay, grid.ncpl) + if grid.idomain is not None + else None, + "iv": np.array([v[0] for v in grid._vertices], dtype=int), + "xv": grid.verts[:, 0].ravel(), + "yv": grid.verts[:, 1].ravel(), + "cell2ddata": Disv.grid_to_disv_cell2d(grid.cell2d), + } + 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 Disv(**kwargs) + + @staticmethod + def disv_to_grid_cell2d(cell2ddata) -> list: + cell2d = [] + iverts = [] + xcenters = [] + ycenters = [] + for rec in cell2ddata: + iverts.append(list(rec.icvert)) + xcenters.append(rec.xc) + ycenters.append(rec.yc) + for n in range(len(iverts)): + cell2d_n = [ + n, + xcenters[n], + ycenters[n], + ] + iverts[n] + cell2d.append(cell2d_n) + return cell2d + + @staticmethod + def grid_to_disv_cell2d(cell2d): + cell2ddata = [] + for cell in cell2d: + verts = cell[4:] + if verts[0] != verts[-1]: + verts.append(cell[4]) + rec = Disv.Cell2dRecord( + cell[0], + cell[1], + cell[2], + len(verts), + tuple(verts), + ) + cell2ddata.append(rec) + return cell2ddata diff --git a/flopy4/mf6/gwf/dis.py b/flopy4/mf6/gwf/dis.py index 4cf7c8d2..defd204b 100644 --- a/flopy4/mf6/gwf/dis.py +++ b/flopy4/mf6/gwf/dis.py @@ -154,12 +154,14 @@ def to_grid(self) -> StructuredGrid: 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, @@ -172,6 +174,10 @@ def from_grid(cls, grid: StructuredGrid) -> "Dis": "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) diff --git a/flopy4/mf6/gwf/disv.py b/flopy4/mf6/gwf/disv.py index 5bc7129a..36f764b6 100644 --- a/flopy4/mf6/gwf/disv.py +++ b/flopy4/mf6/gwf/disv.py @@ -190,8 +190,10 @@ def to_grid(self) -> VertexGrid: length_units=self.length_units, xoff=self.xorigin, yoff=self.yorigin, + angrot=self.angrot, crs=self.crs, nlay=self.nlay, + ncpl=self.ncpl, top=self.top, botm=botm, idomain=idomain, @@ -202,22 +204,30 @@ def to_grid(self) -> VertexGrid: @classmethod def from_grid(cls, grid: VertexGrid) -> "Disv": """Create a discretization from a `VertexGrid`.""" - return Disv( - xorigin=grid.xoffset, - yorigin=grid.yoffset, - nlay=grid.nlay, - ncpl=grid.ncpl, - nvert=grid.nvert, - top=grid.top, - botm=grid.botm, - idomain=np.asarray(grid.idomain).reshape(grid.nlay, grid.ncpl) + _lenunits = {1: "FEET", 2: "METERS", 3: "CENTIMETERS"} + kwargs = { + "xorigin": grid.xoffset, + "yorigin": grid.yoffset, + "nlay": grid.nlay, + "ncpl": grid.ncpl, + "nvert": grid.nvert, + "top": grid.top, + "botm": grid.botm, + "idomain": np.asarray(grid.idomain).reshape(grid.nlay, grid.ncpl) if grid.idomain is not None else None, - iv=np.array([v[0] for v in grid._vertices], dtype=int), - xv=grid.verts[:, 0].ravel(), - yv=grid.verts[:, 1].ravel(), - cell2ddata=Disv.grid_to_disv_cell2d(grid.cell2d), - ) + "iv": np.array([v[0] for v in grid._vertices], dtype=int), + "xv": grid.verts[:, 0].ravel(), + "yv": grid.verts[:, 1].ravel(), + "cell2ddata": Disv.grid_to_disv_cell2d(grid.cell2d), + } + 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 Disv(**kwargs) @staticmethod def disv_to_grid_cell2d(cell2ddata) -> list: diff --git a/flopy4/mf6/gwt/__init__.py b/flopy4/mf6/gwt/__init__.py index 3098e186..3688b042 100644 --- a/flopy4/mf6/gwt/__init__.py +++ b/flopy4/mf6/gwt/__init__.py @@ -1,12 +1,16 @@ 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.gwf.disbase import DisBase from flopy4.mf6.gwt.adv import Adv from flopy4.mf6.gwt.api import Api from flopy4.mf6.gwt.cnc import Cnc from flopy4.mf6.gwt.dis import Dis +from flopy4.mf6.gwt.disv import Disv from flopy4.mf6.gwt.dsp import Dsp from flopy4.mf6.gwt.ic import Ic from flopy4.mf6.gwt.lkt import Lkt @@ -19,9 +23,21 @@ 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__ = [ "Gwt", "Dis", + "Disv", "Adv", "Api", "Cnc", @@ -52,7 +68,7 @@ class Gwt(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) @@ -64,3 +80,8 @@ class Gwt(Model): ssm: Ssm | None = field(block="packages", default=None) mvt: Mvt | None = field(block="packages", default=None) api: Api | None = field(block="packages", default=None) + + @property + def grid(self): + if self.dis is not None: + return self.dis.to_grid() diff --git a/flopy4/mf6/gwt/dis.py b/flopy4/mf6/gwt/dis.py index 0c41fa48..defd204b 100644 --- a/flopy4/mf6/gwt/dis.py +++ b/flopy4/mf6/gwt/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) diff --git a/flopy4/mf6/gwt/disv.py b/flopy4/mf6/gwt/disv.py new file mode 100644 index 00000000..36f764b6 --- /dev/null +++ b/flopy4/mf6/gwt/disv.py @@ -0,0 +1,266 @@ +from pathlib import Path +from typing import ClassVar, Optional + +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.schema import Column, Schema +from flopy4.mf6.utils.grid import VertexGrid +from flopy4.mf6.utl.ncf import Ncf + + +@attrs.define(kw_only=True, slots=False) +class Disv(DisBase): + @attrs.define(slots=False) + class Cell2dRecord: + icell2d: int = attrs.field() + xc: float = attrs.field() + yc: float = attrs.field() + ncvert: int = attrs.field() + icvert: tuple[int, ...] = attrs.field() + + class _VerticesSchema(Schema): + iv = Column("iv", role="value", dfn_type="integer") + xv = Column("xv", role="value", dfn_type="double") + yv = Column("yv", role="value", dfn_type="double") + + __vertices_schema__: ClassVar[type[Schema]] = _VerticesSchema + + 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: Optional[float] = attrs.field( + default=None, + metadata={"dfn_block": "options", "dfn_type": "double", "optional": True}, + ) + yorigin: Optional[float] = attrs.field( + default=None, + 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=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + ncpl: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + nvert: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + top: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("ncpl",), + "layered": False, + "netcdf": True, + }, + ) # type: ignore[assignment] + botm: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("nodes",), + "layered": True, + "netcdf": True, + }, + ) # type: ignore[assignment] + idomain: Optional[NDArray[np.int64]] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "integer", + "shape": ("nodes",), + "layered": True, + "netcdf": True, + }, + ) # type: ignore[assignment] + # User-facing parallel arrays for vertices. + iv: Optional[NDArray[np.int64]] = attrs.field(default=None) # type: ignore[assignment] + xv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + yv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + # Combined vertices recarray for the codec (built in __attrs_post_init__). + vertices: Optional[np.recarray] = attrs.field( + default=None, + metadata={"dfn_block": "vertices", "schema": "__vertices_schema__"}, + ) # type: ignore[assignment] + # Cell2d data — list of Cell2dRecord objects (user-facing). + cell2ddata: Optional[list] = attrs.field(default=None) + # Pre-formatted cell2d rows for the codec (built in __attrs_post_init__). + cell2d: Optional[list] = attrs.field( + default=None, + init=False, + metadata={"dfn_block": "cell2d"}, + ) + + def __attrs_post_init__(self): + # Coerce list inputs to numpy arrays for vertices. + if self.iv is not None and not isinstance(self.iv, np.ndarray): + object.__setattr__(self, "iv", np.asarray(self.iv, dtype=np.int64)) + if self.xv is not None and not isinstance(self.xv, np.ndarray): + object.__setattr__(self, "xv", np.asarray(self.xv, dtype=np.float64)) + if self.yv is not None and not isinstance(self.yv, np.ndarray): + object.__setattr__(self, "yv", np.asarray(self.yv, dtype=np.float64)) + # Build combined vertices recarray for the codec. + if self.iv is not None and self.xv is not None and self.yv is not None: + dtype = np.dtype([("iv", np.int64), ("xv", np.float64), ("yv", np.float64)]) + n = len(self.iv) + arr = np.zeros(n, dtype=dtype) + arr["iv"] = self.iv + 1 # MF6 uses 1-based vertex IDs + arr["xv"] = self.xv + arr["yv"] = self.yv + object.__setattr__(self, "vertices", arr.view(np.recarray)) + # Build cell2d list of tuples for the codec. + if self.cell2ddata is not None: + rows = [] + for rec in self.cell2ddata: + row = (rec.icell2d + 1, rec.xc, rec.yc, rec.ncvert) + tuple( + v + 1 for v in rec.icvert + ) + rows.append(row) + object.__setattr__(self, "cell2d", rows) + # Set derived dimensions. + self.nodes = self.ncpl * self.nlay + self.nrow = 0 + self.ncol = 0 + self._coerce_griddata() + super().__attrs_post_init__() + + def get_dims(self) -> dict[str, int]: + """Get all dimensions.""" + return { + "nlay": self.nlay, + "ncpl": self.ncpl, + "nvert": self.nvert, + "nodes": self.nlay * self.ncpl, + } + + def to_grid(self) -> VertexGrid: + """Convert the discretization to a `VertexGrid`.""" + vertices = [] + for i in range(len(self.iv)): # type: ignore[arg-type] + vertices.append([self.iv[i], self.xv[i], self.yv[i]]) # type: ignore[index] + # VertexGrid expects top as 1D (ncpl,) and botm as 2D (nlay, ncpl). + botm = ( + self.botm.reshape(self.nlay, self.ncpl) + if isinstance(self.botm, np.ndarray) + else self.botm + ) + idomain = ( + self.idomain.reshape(self.nlay, self.ncpl) + if isinstance(self.idomain, np.ndarray) and self.idomain is not None + else self.idomain + ) + return VertexGrid( + length_units=self.length_units, + xoff=self.xorigin, + yoff=self.yorigin, + angrot=self.angrot, + crs=self.crs, + nlay=self.nlay, + ncpl=self.ncpl, + top=self.top, + botm=botm, + idomain=idomain, + vertices=vertices, + cell2d=Disv.disv_to_grid_cell2d(self.cell2ddata), + ) + + @classmethod + def from_grid(cls, grid: VertexGrid) -> "Disv": + """Create a discretization from a `VertexGrid`.""" + _lenunits = {1: "FEET", 2: "METERS", 3: "CENTIMETERS"} + kwargs = { + "xorigin": grid.xoffset, + "yorigin": grid.yoffset, + "nlay": grid.nlay, + "ncpl": grid.ncpl, + "nvert": grid.nvert, + "top": grid.top, + "botm": grid.botm, + "idomain": np.asarray(grid.idomain).reshape(grid.nlay, grid.ncpl) + if grid.idomain is not None + else None, + "iv": np.array([v[0] for v in grid._vertices], dtype=int), + "xv": grid.verts[:, 0].ravel(), + "yv": grid.verts[:, 1].ravel(), + "cell2ddata": Disv.grid_to_disv_cell2d(grid.cell2d), + } + 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 Disv(**kwargs) + + @staticmethod + def disv_to_grid_cell2d(cell2ddata) -> list: + cell2d = [] + iverts = [] + xcenters = [] + ycenters = [] + for rec in cell2ddata: + iverts.append(list(rec.icvert)) + xcenters.append(rec.xc) + ycenters.append(rec.yc) + for n in range(len(iverts)): + cell2d_n = [ + n, + xcenters[n], + ycenters[n], + ] + iverts[n] + cell2d.append(cell2d_n) + return cell2d + + @staticmethod + def grid_to_disv_cell2d(cell2d): + cell2ddata = [] + for cell in cell2d: + verts = cell[4:] + if verts[0] != verts[-1]: + verts.append(cell[4]) + rec = Disv.Cell2dRecord( + cell[0], + cell[1], + cell[2], + len(verts), + tuple(verts), + ) + cell2ddata.append(rec) + return cell2ddata diff --git a/flopy4/mf6/prt/__init__.py b/flopy4/mf6/prt/__init__.py index c3dc9406..180d0099 100644 --- a/flopy4/mf6/prt/__init__.py +++ b/flopy4/mf6/prt/__init__.py @@ -1,18 +1,34 @@ from typing import Optional +from flopy.discretization.structuredgrid import StructuredGrid +from flopy.discretization.vertexgrid import VertexGrid from xattree import xattree +from flopy4.mf6.gwf.disbase import DisBase from flopy4.mf6.model import Model from flopy4.mf6.prt.dis import Dis +from flopy4.mf6.prt.disv import Disv from flopy4.mf6.prt.fmi import Fmi from flopy4.mf6.prt.mip import Mip from flopy4.mf6.prt.oc import Oc from flopy4.mf6.prt.prp import Prp from flopy4.mf6.spec import field + +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__ = [ "Prt", "Dis", + "Disv", "Fmi", "Mip", "Oc", @@ -26,8 +42,13 @@ class Prt(Model): print_input: bool = field(block="options", default=False) print_flows: bool = field(block="options", default=False) save_flows: bool = field(block="options", default=False) - dis: Dis | None = field(block="packages", default=None) + dis: DisBase | None = field(converter=convert_grid, block="packages", default=None) fmi: Fmi | None = field(block="packages", default=None) mip: Mip | None = field(block="packages", default=None) oc: Oc | None = field(block="packages", default=None) prp: list[Prp] = field(block="packages") + + @property + def grid(self): + if self.dis is not None: + return self.dis.to_grid() diff --git a/flopy4/mf6/prt/dis.py b/flopy4/mf6/prt/dis.py index afb98e7c..78120b3c 100644 --- a/flopy4/mf6/prt/dis.py +++ b/flopy4/mf6/prt/dis.py @@ -57,7 +57,7 @@ class Dis(DisBase): "dfn_type": "double", "shape": ("ncol",), "layered": False, - "netcdf": True, + "netcdf": False, }, ) # type: ignore[assignment] delc: NDArray[np.float64] = attrs.field( @@ -67,7 +67,7 @@ class Dis(DisBase): "dfn_type": "double", "shape": ("nrow",), "layered": False, - "netcdf": True, + "netcdf": False, }, ) # type: ignore[assignment] top: NDArray[np.float64] = attrs.field( @@ -77,7 +77,7 @@ class Dis(DisBase): "dfn_type": "double", "shape": ("ncpl",), "layered": False, - "netcdf": True, + "netcdf": False, }, ) # type: ignore[assignment] botm: NDArray[np.float64] = attrs.field( @@ -87,7 +87,7 @@ class Dis(DisBase): "dfn_type": "double", "shape": ("nodes",), "layered": True, - "netcdf": True, + "netcdf": False, }, ) # type: ignore[assignment] idomain: Optional[NDArray[np.int64]] = attrs.field( @@ -97,7 +97,7 @@ class Dis(DisBase): "dfn_type": "integer", "shape": ("nodes",), "layered": True, - "netcdf": True, + "netcdf": False, }, ) # type: ignore[assignment] @@ -120,18 +120,12 @@ def get_dims(self) -> dict[str, int]: def to_grid(self) -> StructuredGrid: """Convert the discretization to a `StructuredGrid`.""" - top = ( - self.top.reshape(self.nrow, self.ncol) if isinstance(self.top, np.ndarray) else self.top - ) - botm = ( - self.botm.reshape(self.nlay, self.nrow, self.ncol) - if isinstance(self.botm, np.ndarray) - else self.botm - ) + top = np.asarray(self.top).reshape(self.nrow, self.ncol) + botm = np.asarray(self.botm).reshape(self.nlay, self.nrow, self.ncol) idomain = ( - self.idomain.reshape(self.nlay, self.nrow, self.ncol) - if isinstance(self.idomain, np.ndarray) - else self.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, @@ -140,17 +134,19 @@ def to_grid(self) -> StructuredGrid: nlay=self.nlay, nrow=self.nrow, ncol=self.ncol, - delr=self.delr, - delc=self.delc, + 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, @@ -163,6 +159,10 @@ def from_grid(cls, grid: StructuredGrid) -> "Dis": "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) diff --git a/flopy4/mf6/prt/disv.py b/flopy4/mf6/prt/disv.py new file mode 100644 index 00000000..c793bc61 --- /dev/null +++ b/flopy4/mf6/prt/disv.py @@ -0,0 +1,251 @@ +from typing import ClassVar, Optional + +import attrs +import numpy as np +from numpy.typing import NDArray + +from flopy4.mf6.gwf.disbase import DisBase +from flopy4.mf6.schema import Column, Schema +from flopy4.mf6.utils.grid import VertexGrid + + +@attrs.define(kw_only=True, slots=False) +class Disv(DisBase): + @attrs.define(slots=False) + class Cell2dRecord: + icell2d: int = attrs.field() + xc: float = attrs.field() + yc: float = attrs.field() + ncvert: int = attrs.field() + icvert: tuple[int, ...] = attrs.field() + + class _VerticesSchema(Schema): + iv = Column("iv", role="value", dfn_type="integer") + xv = Column("xv", role="value", dfn_type="double") + yv = Column("yv", role="value", dfn_type="double") + + __vertices_schema__: ClassVar[type[Schema]] = _VerticesSchema + + 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: Optional[float] = attrs.field( + default=None, + metadata={"dfn_block": "options", "dfn_type": "double", "optional": True}, + ) + yorigin: Optional[float] = attrs.field( + default=None, + 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}, + ) + nlay: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + ncpl: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + nvert: int = attrs.field( + default=0, + metadata={"dfn_block": "dimensions", "dfn_type": "integer"}, + ) + top: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("ncpl",), + "layered": False, + "netcdf": False, + }, + ) # type: ignore[assignment] + botm: NDArray[np.float64] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "double", + "shape": ("nodes",), + "layered": True, + "netcdf": False, + }, + ) # type: ignore[assignment] + idomain: Optional[NDArray[np.int64]] = attrs.field( + default=None, + metadata={ + "dfn_block": "griddata", + "dfn_type": "integer", + "shape": ("nodes",), + "layered": True, + "netcdf": False, + }, + ) # type: ignore[assignment] + # User-facing parallel arrays for vertices. + iv: Optional[NDArray[np.int64]] = attrs.field(default=None) # type: ignore[assignment] + xv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + yv: Optional[NDArray[np.float64]] = attrs.field(default=None) # type: ignore[assignment] + # Combined vertices recarray for the codec (built in __attrs_post_init__). + vertices: Optional[np.recarray] = attrs.field( + default=None, + metadata={"dfn_block": "vertices", "schema": "__vertices_schema__"}, + ) # type: ignore[assignment] + # Cell2d data — list of Cell2dRecord objects (user-facing). + cell2ddata: Optional[list] = attrs.field(default=None) + # Pre-formatted cell2d rows for the codec (built in __attrs_post_init__). + cell2d: Optional[list] = attrs.field( + default=None, + init=False, + metadata={"dfn_block": "cell2d"}, + ) + + def __attrs_post_init__(self): + # Coerce list inputs to numpy arrays for vertices. + if self.iv is not None and not isinstance(self.iv, np.ndarray): + object.__setattr__(self, "iv", np.asarray(self.iv, dtype=np.int64)) + if self.xv is not None and not isinstance(self.xv, np.ndarray): + object.__setattr__(self, "xv", np.asarray(self.xv, dtype=np.float64)) + if self.yv is not None and not isinstance(self.yv, np.ndarray): + object.__setattr__(self, "yv", np.asarray(self.yv, dtype=np.float64)) + # Build combined vertices recarray for the codec. + if self.iv is not None and self.xv is not None and self.yv is not None: + dtype = np.dtype([("iv", np.int64), ("xv", np.float64), ("yv", np.float64)]) + n = len(self.iv) + arr = np.zeros(n, dtype=dtype) + arr["iv"] = self.iv + 1 # MF6 uses 1-based vertex IDs + arr["xv"] = self.xv + arr["yv"] = self.yv + object.__setattr__(self, "vertices", arr.view(np.recarray)) + # Build cell2d list of tuples for the codec. + if self.cell2ddata is not None: + rows = [] + for rec in self.cell2ddata: + row = (rec.icell2d + 1, rec.xc, rec.yc, rec.ncvert) + tuple( + v + 1 for v in rec.icvert + ) + rows.append(row) + object.__setattr__(self, "cell2d", rows) + # Set derived dimensions. + self.nodes = self.ncpl * self.nlay + self.nrow = 0 + self.ncol = 0 + self._coerce_griddata() + super().__attrs_post_init__() + + def get_dims(self) -> dict[str, int]: + """Get all dimensions.""" + return { + "nlay": self.nlay, + "ncpl": self.ncpl, + "nvert": self.nvert, + "nodes": self.nlay * self.ncpl, + } + + def to_grid(self) -> VertexGrid: + """Convert the discretization to a `VertexGrid`.""" + vertices = [] + for i in range(len(self.iv)): # type: ignore[arg-type] + vertices.append([self.iv[i], self.xv[i], self.yv[i]]) # type: ignore[index] + botm = ( + self.botm.reshape(self.nlay, self.ncpl) + if isinstance(self.botm, np.ndarray) + else self.botm + ) + idomain = ( + self.idomain.reshape(self.nlay, self.ncpl) + if isinstance(self.idomain, np.ndarray) and self.idomain is not None + else self.idomain + ) + return VertexGrid( + length_units=self.length_units, + xoff=self.xorigin, + yoff=self.yorigin, + angrot=self.angrot, + crs=self.crs, + nlay=self.nlay, + ncpl=self.ncpl, + top=self.top, + botm=botm, + idomain=idomain, + vertices=vertices, + cell2d=Disv.disv_to_grid_cell2d(self.cell2ddata), + ) + + @classmethod + def from_grid(cls, grid: VertexGrid) -> "Disv": + """Create a discretization from a `VertexGrid`.""" + _lenunits = {1: "FEET", 2: "METERS", 3: "CENTIMETERS"} + kwargs = { + "xorigin": grid.xoffset, + "yorigin": grid.yoffset, + "nlay": grid.nlay, + "ncpl": grid.ncpl, + "nvert": grid.nvert, + "top": grid.top, + "botm": grid.botm, + "idomain": np.asarray(grid.idomain).reshape(grid.nlay, grid.ncpl) + if grid.idomain is not None + else None, + "iv": np.array([v[0] for v in grid._vertices], dtype=int), + "xv": grid.verts[:, 0].ravel(), + "yv": grid.verts[:, 1].ravel(), + "cell2ddata": Disv.grid_to_disv_cell2d(grid.cell2d), + } + 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 Disv(**kwargs) + + @staticmethod + def disv_to_grid_cell2d(cell2ddata) -> list: + cell2d = [] + iverts = [] + xcenters = [] + ycenters = [] + for rec in cell2ddata: + iverts.append(list(rec.icvert)) + xcenters.append(rec.xc) + ycenters.append(rec.yc) + for n in range(len(iverts)): + cell2d_n = [ + n, + xcenters[n], + ycenters[n], + ] + iverts[n] + cell2d.append(cell2d_n) + return cell2d + + @staticmethod + def grid_to_disv_cell2d(cell2d): + cell2ddata = [] + for cell in cell2d: + verts = cell[4:] + if verts[0] != verts[-1]: + verts.append(cell[4]) + rec = Disv.Cell2dRecord( + cell[0], + cell[1], + cell[2], + len(verts), + tuple(verts), + ) + cell2ddata.append(rec) + return cell2ddata diff --git a/flopy4/mf6/utils/grid.py b/flopy4/mf6/utils/grid.py index 02ca0369..1cbe2f84 100644 --- a/flopy4/mf6/utils/grid.py +++ b/flopy4/mf6/utils/grid.py @@ -140,6 +140,7 @@ def from_dis(cls, dis): length_units=dis.length_units, xoff=dis.xorigin, yoff=dis.yorigin, + angrot=getattr(dis, "angrot", None), crs=dis.crs, nlay=dis.nlay, nrow=dis.nrow, @@ -900,6 +901,7 @@ def from_dis(cls, dis, **kwargs): length_units=dis.length_units, xoff=dis.xorigin, yoff=dis.yorigin, + angrot=getattr(dis, "angrot", None), crs=dis.crs, nlay=dis.nlay, ncpl=dis.ncpl, diff --git a/test/test_mf6_component.py b/test/test_mf6_component.py index 74ffba6a..55f2e957 100644 --- a/test/test_mf6_component.py +++ b/test/test_mf6_component.py @@ -1537,3 +1537,144 @@ def test_npf_to_xarray_via_parent_chain(): assert "strt" in ic_ds assert ic_ds.strt.shape == (1, 2, 2) assert np.allclose(ic_ds.strt.values, 10.0) + + +# --------------------------------------------------------------------------- +# Dis/Disv class identity across models +# --------------------------------------------------------------------------- + + +def test_dis_class_identity(): + """Each model's Dis is a distinct class, not a re-export of gwf.Dis.""" + from flopy4.mf6 import gwe, gwf, gwt, prt + + assert gwt.Dis is not gwf.Dis + assert gwe.Dis is not gwf.Dis + assert prt.Dis is not gwf.Dis + assert gwt.Dis.__module__ == "flopy4.mf6.gwt.dis" + assert gwe.Dis.__module__ == "flopy4.mf6.gwe.dis" + assert prt.Dis.__module__ == "flopy4.mf6.prt.dis" + + +def test_disv_class_identity(): + """Each model's Disv is a distinct class, not a re-export of gwf.Disv.""" + from flopy4.mf6 import gwe, gwf, gwt, prt + + assert gwt.Disv is not gwf.Disv + assert gwe.Disv is not gwf.Disv + assert prt.Disv is not gwf.Disv + assert gwt.Disv.__module__ == "flopy4.mf6.gwt.disv" + assert gwe.Disv.__module__ == "flopy4.mf6.gwe.disv" + assert prt.Disv.__module__ == "flopy4.mf6.prt.disv" + + +def test_prt_dis_no_ncf(): + """prt.Dis and prt.Disv must not expose NCF fields.""" + import attrs + + from flopy4.mf6 import prt + + dis_field_names = {f.name for f in attrs.fields(prt.Dis)} + assert "ncf6_filerecord" not in dis_field_names + assert "ncf" not in dis_field_names + + disv_field_names = {f.name for f in attrs.fields(prt.Disv)} + assert "ncf6_filerecord" not in disv_field_names + assert "ncf" not in disv_field_names + + +def test_gwt_gwe_disv_instantiate(): + """gwt.Disv and gwe.Disv can be instantiated with basic grid dims.""" + from flopy4.mf6 import gwe, gwt + + for cls in (gwt.Disv, gwe.Disv): + d = cls(nlay=1, ncpl=4) + assert d.nlay == 1 + assert d.ncpl == 4 + assert hasattr(d, "ncf6_filerecord") + + +# --------------------------------------------------------------------------- +# convert_grid and grid property for gwt, gwe, prt +# --------------------------------------------------------------------------- + + +def test_gwt_convert_grid_structured(): + from flopy4.mf6.gwt import Dis as GwtDis + from flopy4.mf6.gwt import Gwt + + grid = StructuredGrid.uniform(nlay=1, nrow=5, ncol=5) + gwt = Gwt(dis=grid) + assert isinstance(gwt.dis, GwtDis) + assert gwt.dis.nrow == 5 + assert isinstance(gwt.grid, StructuredGrid) + + +def test_gwt_convert_grid_vertex(vgrid): + from flopy4.mf6.gwt import Disv as GwtDisv + from flopy4.mf6.gwt import Gwt + + grid = VertexGrid( + nlay=vgrid["nlay"], + ncpl=vgrid["ncpl"], + vertices=vgrid["vertices"], + cell2d=vgrid["cell2d"], + ) + gwt = Gwt(dis=grid) + assert isinstance(gwt.dis, GwtDisv) + assert gwt.dis.ncpl == vgrid["ncpl"] + assert isinstance(gwt.grid, VertexGrid) + + +def test_gwe_convert_grid_structured(): + from flopy4.mf6.gwe import Dis as GweDis + from flopy4.mf6.gwe import Gwe + + grid = StructuredGrid.uniform(nlay=1, nrow=5, ncol=5) + gwe = Gwe(dis=grid) + assert isinstance(gwe.dis, GweDis) + assert gwe.dis.nrow == 5 + assert isinstance(gwe.grid, StructuredGrid) + + +def test_gwe_convert_grid_vertex(vgrid): + from flopy4.mf6.gwe import Disv as GweDisv + from flopy4.mf6.gwe import Gwe + + grid = VertexGrid( + nlay=vgrid["nlay"], + ncpl=vgrid["ncpl"], + vertices=vgrid["vertices"], + cell2d=vgrid["cell2d"], + ) + gwe = Gwe(dis=grid) + assert isinstance(gwe.dis, GweDisv) + assert gwe.dis.ncpl == vgrid["ncpl"] + assert isinstance(gwe.grid, VertexGrid) + + +def test_prt_convert_grid_structured(): + from flopy4.mf6.prt import Dis as PrtDis + from flopy4.mf6.prt import Prt + + grid = StructuredGrid.uniform(nlay=1, nrow=5, ncol=5) + prt = Prt(dis=grid) + assert isinstance(prt.dis, PrtDis) + assert prt.dis.nrow == 5 + assert isinstance(prt.grid, StructuredGrid) + + +def test_prt_convert_grid_vertex(vgrid): + from flopy4.mf6.prt import Disv as PrtDisv + from flopy4.mf6.prt import Prt + + grid = VertexGrid( + nlay=vgrid["nlay"], + ncpl=vgrid["ncpl"], + vertices=vgrid["vertices"], + cell2d=vgrid["cell2d"], + ) + prt = Prt(dis=grid) + assert isinstance(prt.dis, PrtDisv) + assert prt.dis.ncpl == vgrid["ncpl"] + assert isinstance(prt.grid, VertexGrid)