add block spec functions (#134)

analogous to fields and fields_dict add blocks and blocks_dict functions which return the block specification

Author: wpbonelli

flopy4/mf6/gwf/__init__.py

@@ -7,13 +7,13 @@
 from flopy.discretization.grid import Grid
 from xattree import xattree
 
-from flopy4.mf6.decorators import field
 from flopy4.mf6.gwf.chd import Chd
 from flopy4.mf6.gwf.dis import Dis
 from flopy4.mf6.gwf.ic import Ic
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
 from flopy4.mf6.model import Model
+from flopy4.mf6.spec import field
 from flopy4.mf6.utils import open_cbc, open_hds
 
 __all__ = ["Gwf", "Chd", "Dis", "Ic", "Npf", "Oc"]

flopy4/mf6/gwf/chd.py

@@ -7,8 +7,8 @@
 from xattree import xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, field
 
 
 @xattree

flopy4/mf6/gwf/dis.py

@@ -5,8 +5,8 @@
 from xattree import xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, dim, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, dim, field
 
 
 @xattree

flopy4/mf6/gwf/ic.py

@@ -4,8 +4,8 @@
 from xattree import xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, field
 
 
 @xattree

flopy4/mf6/gwf/npf.py

@@ -7,8 +7,8 @@
 from xattree import xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, field
 
 
 @xattree

flopy4/mf6/gwf/oc.py

@@ -7,8 +7,8 @@
 from xattree import xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, field
 from flopy4.utils import to_path
 
 

flopy4/mf6/ims.py

@@ -3,8 +3,8 @@
 
 from xattree import xattree
 
-from flopy4.mf6.decorators import field
 from flopy4.mf6.solution import Solution
+from flopy4.mf6.spec import field
 
 
 @xattree

flopy4/mf6/decorators.py renamed to flopy4/mf6/spec.py

@@ -1,4 +1,9 @@
-from attrs import NOTHING
+"""
+Wrap `xattree` and `attrs` specification utilities for MF6.
+These include field decorators and introspection functions.
+"""
+
+from attrs import NOTHING, Attribute, fields_dict
 from xattree import array as xattree_array
 from xattree import coord as xattree_coord
 from xattree import dim as xattree_dim
@@ -15,7 +20,7 @@ def field(
     metadata=None,
     block: str | None = None,
 ):
-    """Create a field."""
+    """Define a field."""
     if block:
         metadata = metadata or {}
         metadata["block"] = block
@@ -40,7 +45,7 @@ def dim(
     metadata=None,
     block: str | None = None,
 ):
-    """Create a dimension field."""
+    """Define a dimension field."""
     if block:
         metadata = metadata or {}
         metadata["block"] = block
@@ -63,7 +68,7 @@ def coord(
     metadata=None,
     block: str | None = None,
 ):
-    """Create a coordinate field."""
+    """Define a coordinate field."""
     if block:
         metadata = metadata or {}
         metadata["block"] = block
@@ -87,7 +92,7 @@ def array(
     metadata=None,
     block: str | None = None,
 ):
-    """Create an array field."""
+    """Define an array field."""
     if block:
         metadata = metadata or {}
         metadata["block"] = block
@@ -101,3 +106,44 @@ def array(
         eq=eq,
         metadata=metadata,
     )
+
+
+Block = dict[str, Attribute]
+
+
+def _block_sort_key(item) -> int:
+    k, _ = item
+    if k == "options":
+        return 0
+    elif k == "dimensions":
+        return 1
+    elif k == "griddata":
+        return 2
+    elif k == "packagedata":
+        return 3
+    elif k == "perioddata":
+        return 4
+    else:
+        return 5
+
+
+def blocks(cls) -> list[list[Attribute]]:
+    """Return an ordered list of blocks for a component class."""
+    return [list(v.values()) for v in blocks_dict(cls).values()]
+
+
+def blocks_dict(cls) -> dict[str, Block]:
+    """
+    Return an ordered dictionary of blocks for a component class,
+    whose keys are block names. Each block is a map from variable
+    (field) name to `attrs.Attribute`.
+    """
+    fields = fields_dict(cls)
+    fields = {k: v for k, v in fields.items() if "block" in v.metadata}
+    blocks: dict[str, Block] = {}
+    for k, v in fields.items():
+        block = v.metadata["block"]
+        if block not in blocks:
+            blocks[block] = {}
+        blocks[block][k] = v
+    return dict(sorted(blocks.items(), key=_block_sort_key))

flopy4/mf6/tdis.py

@@ -8,8 +8,8 @@
 from xattree import ROOT, xattree
 
 from flopy4.mf6.converters import convert_array
-from flopy4.mf6.decorators import array, dim, field
 from flopy4.mf6.package import Package
+from flopy4.mf6.spec import array, dim, field
 
 
 @xattree

flopy4/mf6/utils/cbc_reader.py

@@ -135,9 +135,7 @@ def open_cbc(
 
     """
     grid = StructuredGridWrapper.from_binary_grid_file(grb_path)
-    cbc = _open_cbc_dis(
-        cbc_path, grid, flowja, simulation_start_time, time_unit
-    )
+    cbc = _open_cbc_dis(cbc_path, grid, flowja, simulation_start_time, time_unit)
     return xr.merge([cbc])
 
 
@@ -267,17 +265,13 @@ def read_cbc_headers(
             if header["imeth"] == 1:
                 # Multiply by -1 because ndim3 is stored as a negative for some
                 # reason. (ndim3 is the integer size of the third dimension)
-                datasize = (
-                    header["ndim1"] * header["ndim2"] * header["ndim3"] * -1
-                ) * 8
+                datasize = (header["ndim1"] * header["ndim2"] * header["ndim3"] * -1) * 8
                 header["pos"] = f.tell()
                 key = header["text"]
                 headers[key].append(Imeth1Header(**header))
             elif header["imeth"] == 6:
                 imeth6_header = read_imeth6_header(f)
-                datasize = imeth6_header["nlist"] * (
-                    8 + imeth6_header["ndat"] * 8
-                )
+                datasize = imeth6_header["nlist"] * (8 + imeth6_header["ndat"] * 8)
                 header["pos"] = f.tell()
                 # key-format:
                 # "package type"-"optional_package_variable"_"package name"
@@ -287,11 +281,7 @@ def read_cbc_headers(
                 # npf-key can be present multiple times in cases of saved
                 # saturation + specific discharge
                 if header["text"].startswith("data-"):
-                    key = (
-                        imeth6_header["txt2id2"]
-                        + "_"
-                        + header["text"].replace("data-", "")
-                    )
+                    key = imeth6_header["txt2id2"] + "_" + header["text"].replace("data-", "")
                 headers[key].append(Imeth6Header(**header, **imeth6_header))
             else:
                 raise ValueError(
@@ -332,9 +322,7 @@ def read_imeth6_header(f: BinaryIO) -> dict[str, Any]:
     content["txt2id2"] = f.read(16).decode("utf-8").strip().lower()
     ndat = struct.unpack("i", f.read(4))[0]
     content["ndat"] = ndat
-    content["auxtxt"] = [
-        f.read(16).decode("utf-8").strip().lower() for _ in range(ndat - 1)
-    ]
+    content["auxtxt"] = [f.read(16).decode("utf-8").strip().lower() for _ in range(ndat - 1)]
     content["nlist"] = struct.unpack("i", f.read(4))[0]
     return content
 
@@ -345,14 +333,9 @@ def assign_datetime_coords(
     time_unit: str | None = "d",
 ) -> xr.DataArray:
     if "time" not in da.coords:
-        raise ValueError(
-            "cannot convert time column, "
-            "because a time column could not be found"
-        )
+        raise ValueError("cannot convert time column, because a time column could not be found")
 
-    time = pd.Timestamp(simulation_start_time) + pd.to_timedelta(
-        da["time"], unit=time_unit
-    )
+    time = pd.Timestamp(simulation_start_time) + pd.to_timedelta(da["time"], unit=time_unit)
     return da.assign_coords(time=time)
 
 
@@ -411,9 +394,7 @@ def open_imeth6_budgets(
     coords = get_coords(grid)
     coords["time"] = time
     name = header_list[0].text
-    return xr.DataArray(
-        daskarr, coords, ("time", "layer", "y", "x"), name=name
-    )
+    return xr.DataArray(daskarr, coords, ("time", "layer", "y", "x"), name=name)
 
 
 def read_imeth6_budgets_dense(
@@ -470,9 +451,7 @@ def read_imeth6_budgets_dense(
     return out.reshape(shape)
 
 
-def read_imeth6_budgets(
-    cbc_path: Path, count: int, dtype: np.dtype, pos: int
-) -> Any:
+def read_imeth6_budgets(cbc_path: Path, count: int, dtype: np.dtype, pos: int) -> Any:
     """
     Read the data for an imeth==6 budget section for a single timestep.
 
@@ -549,9 +528,7 @@ def open_imeth1_budgets(
     )
 
 
-def cbc_open_imeth1_budgets(
-    cbc_path: Path, header_list: list[Imeth1Header]
-) -> xr.DataArray:
+def cbc_open_imeth1_budgets(cbc_path: Path, header_list: list[Imeth1Header]) -> xr.DataArray:
     """
     Open the data for an imeth==1 budget section. Data is read lazily per
     timestep. The cell data is not spatially labelled.
@@ -625,19 +602,15 @@ def dis_open_face_budgets(
     front: xr.DataArray of floats with dims ("time", "layer", "y", "x")
     lower: xr.DataArray of floats with dims ("time", "layer", "y", "x")
     """
-    right_index, front_index, lower_index = dis_to_right_front_lower_indices(
-        grid
-    )
+    right_index, front_index, lower_index = dis_to_right_front_lower_indices(grid)
     budgets = cbc_open_imeth1_budgets(cbc_path, header_list)
     right = dis_extract_face_budgets(budgets, right_index)
     front = dis_extract_face_budgets(budgets, front_index)
     lower = dis_extract_face_budgets(budgets, lower_index)
     return right, front, lower
 
 
-def dis_extract_face_budgets(
-    budgets: xr.DataArray, index: xr.DataArray
-) -> xr.DataArray:
+def dis_extract_face_budgets(budgets: xr.DataArray, index: xr.DataArray) -> xr.DataArray:
     """
     Grab right, front, or lower face flows from the flow-ja-face array.