cleanup (#158)

misc cleanup, also pin llvmlite to avoid issues with numba

Author: wpbonelli

flopy4/adapters.py

@@ -1,6 +1,9 @@
 import os
 
-from flopy.discretization import StructuredGrid
+from flopy.discretization.grid import Grid
+from flopy.discretization.structuredgrid import StructuredGrid
+from flopy.discretization.unstructuredgrid import UnstructuredGrid
+from flopy.discretization.vertexgrid import VertexGrid
 
 
 class StructuredGridWrapper(StructuredGrid):
@@ -108,3 +111,34 @@ def from_binary_grid_file(cls, file_path, verbose=False):
             ia=grb_obj.ia,
             ja=grb_obj.ja,
         )
+
+
+def get_kij(nn: int, nlay: int, nrow: int, ncol: int) -> tuple[int, int, int]:
+    nodes = nlay * nrow * ncol
+    if nn < 0 or nn >= nodes:
+        raise ValueError(f"Node number {nn} is out of bounds (1 to {nodes})")
+    k = (nn - 1) / (ncol * nrow) + 1
+    ij = nn - (k - 1) * ncol * nrow
+    i = (ij - 1) / ncol + 1
+    j = ij - (i - 1) * ncol
+    return int(k), int(i), int(j)
+
+
+def get_jk(nn: int, ncpl: int) -> tuple[int, int]:
+    if nn < 0 or nn >= ncpl:
+        raise ValueError(f"Node number {nn} is out of bounds (1 to {ncpl})")
+    k = (nn - 1) / ncpl + 1
+    j = nn - (k - 1) * ncpl
+    return int(j), int(k)
+
+
+def get_cellid(nn: int, grid: Grid) -> tuple[int, ...]:
+    match grid:
+        case StructuredGrid():
+            return get_kij(nn, *grid.shape)
+        case VertexGrid():
+            return get_jk(nn, grid.ncpl)
+        case UnstructuredGrid():
+            return (nn,)
+        case _:
+            raise TypeError(f"Unsupported grid type: {type(grid)}")

flopy4/mf6/codec/__init__.py

@@ -40,6 +40,10 @@
 
 
 def _make_converter() -> Converter:
+    # TODO: document what is converter's responsibility vs Jinja's
+    # TODO: how can we make sure writing remains lazy for list input?
+    # don't eagerly unstructure to dict, lazily access from the template?
+
     from flopy4.mf6.component import Component
     from flopy4.mf6.gwf.chd import Chd
     from flopy4.mf6.gwf.oc import Oc

flopy4/mf6/codec/converter.py

@@ -3,14 +3,11 @@
 import numpy as np
 import sparse
 import xattree
-from flopy.discretization.grid import Grid
-from flopy.discretization.structuredgrid import StructuredGrid
-from flopy.discretization.unstructuredgrid import UnstructuredGrid
-from flopy.discretization.vertexgrid import VertexGrid
 from numpy.typing import NDArray
 from xarray import DataArray
 from xattree import get_xatspec
 
+from flopy4.adapters import get_cellid
 from flopy4.mf6.component import Component
 from flopy4.mf6.config import SPARSE_THRESHOLD
 from flopy4.mf6.constants import FILL_DNODATA
@@ -147,7 +144,8 @@ def unstructure_array(value: DataArray) -> dict:
 
 def unstructure_component(value: Component) -> dict[str, Any]:
     data = xattree.asdict(value)
-    for block in get_blocks(value.dfn).values():
+    blocks = get_blocks(value.dfn)
+    for block in blocks.values():
         for field_name, field in block.items():
             if is_list_field(field):
                 data[field_name] = unstructure_array(data[field_name])
@@ -180,37 +178,6 @@ def unstructure_tdis(value: Any) -> dict[str, Any]:
     return data
 
 
-def get_kij(nn: int, nlay: int, nrow: int, ncol: int) -> tuple[int, int, int]:
-    nodes = nlay * nrow * ncol
-    if nn < 0 or nn >= nodes:
-        raise ValueError(f"Node number {nn} is out of bounds (1 to {nodes})")
-    k = (nn - 1) / (ncol * nrow) + 1
-    ij = nn - (k - 1) * ncol * nrow
-    i = (ij - 1) / ncol + 1
-    j = ij - (i - 1) * ncol
-    return int(k), int(i), int(j)
-
-
-def get_jk(nn: int, ncpl: int) -> tuple[int, int]:
-    if nn < 0 or nn >= ncpl:
-        raise ValueError(f"Node number {nn} is out of bounds (1 to {ncpl})")
-    k = (nn - 1) / ncpl + 1
-    j = nn - (k - 1) * ncpl
-    return int(j), int(k)
-
-
-def get_cellid(nn: int, grid: Grid) -> tuple[int, ...]:
-    match grid:
-        case StructuredGrid():
-            return get_kij(nn, *grid.shape)
-        case VertexGrid():
-            return get_jk(nn, grid.ncpl)
-        case UnstructuredGrid():
-            return (nn,)
-        case _:
-            raise TypeError(f"Unsupported grid type: {type(grid)}")
-
-
 def unstructure_chd(value: Any) -> dict[str, Any]:
     if (parent := value.parent) is None:
         raise ValueError(

flopy4/mf6/component.py

@@ -23,10 +23,6 @@ class Component(ABC, MutableMapping):
     Notes
     -----
     All subclasses of `Component` must be decorated with `xattree`.
-
-    We use the `children` attribute provided by `xattree`. We know
-    children are also `Component`s, but mypy does not. TODO: fix??
-    Then we can remove the `# type: ignore` comments.
     """
 
     _load = IO(Loader)  # type: ignore
@@ -60,6 +56,9 @@ def __attrs_init_subclass__(cls):
         cls.dfn = cls.get_dfn()
 
     def __getitem__(self, key):
+        # We use `children` from `xattree` to implement MutableMapping.
+        # children are also `Component`s, but mypy doesn't know this..
+        # TODO fix, then we can remove the `# type: ignore` comments.
         return self.children[key]  # type: ignore
 
     def __setitem__(self, key, value):
@@ -76,7 +75,7 @@ def __len__(self):
 
     @classmethod
     def get_dfn(cls) -> Dfn:
-        """Generate the component's MODFLOW 6 definition."""
+        """Get the component's definition (i.e. specification)."""
         fields = {field_name: to_dfn_field(field) for field_name, field in fields_dict(cls).items()}
         blocks: dict[str, dict[str, Field]] = {}
         for field_name, field_ in fields.items():
@@ -95,7 +94,7 @@ def get_dfn(cls) -> Dfn:
         )
 
     def _preio(self, format: str) -> None:
-        """Place for any pre-IO setup"""
+        # prep for io operations
         if not self.filename:
             self.filename = self.default_filename()