back to original dict->array converter

Author: wpbonelli

flopy4/mf6/codec/__init__.py

@@ -8,13 +8,7 @@
 from jinja2 import Environment, PackageLoader
 
 from flopy4.mf6 import filters
-from flopy4.mf6.codec.converter import (
-    unstructure_array,
-    unstructure_chd,
-    unstructure_component,
-    unstructure_oc,
-    unstructure_tdis,
-)
+from flopy4.mf6.codec.converter import unstructure_component
 
 _JINJA_ENV = Environment(
     loader=PackageLoader("flopy4.mf6"),
@@ -28,6 +22,8 @@
 _JINJA_ENV.filters["array_how"] = filters.array_how
 _JINJA_ENV.filters["array_chunks"] = filters.array_chunks
 _JINJA_ENV.filters["array2string"] = filters.array2string
+_JINJA_ENV.filters["to_sparse_dict"] = filters.to_sparse_dict
+_JINJA_ENV.filters["to_period_records"] = filters.to_period_records
 
 _JINJA_TEMPLATE_NAME = "blocks.jinja"
 
@@ -39,21 +35,12 @@
 
 
 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?
-
+    """Create a simple converter that just handles structure conversion."""
     from flopy4.mf6.component import Component
-    from flopy4.mf6.gwf.chd import Chd
-    from flopy4.mf6.gwf.oc import Oc
-    from flopy4.mf6.tdis import Tdis
 
     converter = Converter()
     converter.register_unstructure_hook_factory(xattree.has, lambda _: xattree.asdict)
     converter.register_unstructure_hook(Component, unstructure_component)
-    converter.register_unstructure_hook(Tdis, unstructure_tdis)
-    converter.register_unstructure_hook(Chd, unstructure_chd)
-    converter.register_unstructure_hook(Oc, unstructure_oc)
     return converter
 
 
@@ -72,19 +59,20 @@ def load(path: str | PathLike) -> Any:
 
 def dumps(data) -> str:
     template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
+    unstructured_data = _CONVERTER.unstructure(data)
     with np.printoptions(**_PRINT_OPTIONS):  # type: ignore
-        return template.render(dfn=type(data).dfn, data=_CONVERTER.unstructure(data))
+        return template.render(dfn=type(data).dfn, data=unstructured_data, component=data)
 
 
 def dump(data, path: str | PathLike) -> None:
     template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
-    iterator = template.generate(dfn=type(data).dfn, data=_CONVERTER.unstructure(data))
+    unstructured_data = _CONVERTER.unstructure(data)
+    iterator = template.generate(dfn=type(data).dfn, data=unstructured_data, component=data)
     with np.printoptions(**_PRINT_OPTIONS), open(path, "w") as f:  # type: ignore
         f.writelines(iterator)
 
 
 __all__ = [
-    "unstructure_array",
     "loads",
     "load",
     "dumps",

flopy4/mf6/codec/converter.py

@@ -1,154 +1,15 @@
 from typing import Any
 
-import numpy as np
-import sparse
 import xattree
-from xarray import DataArray
 
-from flopy4.adapters import get_cellid
 from flopy4.mf6.component import Component
-from flopy4.mf6.constants import FILL_DNODATA
-from flopy4.mf6.spec import get_blocks, is_list_field
-
-
-def unstructure_array(value: DataArray) -> dict:
-    """
-    Convert a dense numpy array or a sparse COO array to a sparse
-    dictionary representation suitable for serialization into the
-    MF6 list-based input format.
-
-    The input array must have a time dimension named 'nper', i.e.
-    it must be stress period data for some MODFLOW 6 component.
-
-    Returns:
-        dict: {kper: {spatial indices: value, ...}, ...}
-    """
-    if (time_dim := "nper") not in value.dims:
-        raise ValueError(f"Array must have dimension '{time_dim}'")
-    if isinstance(value.data, sparse.COO):
-        coords = value.coords
-        data = value.data
-    else:
-        coords = np.array(np.where(value.data != FILL_DNODATA)).T  # type: ignore
-        data = value.data[tuple(coords.T)]  # type: ignore
-    if not coords.size:  # type: ignore
-        return {}
-    result = {}
-    match value.ndim:
-        case 1:
-            # Only kper, no spatial dims
-            for kper, v in zip(coords[:, 0], data):
-                result[int(kper)] = v
-        case _:
-            # kper + spatial dims
-            for row, v in zip(coords, data):
-                kper = int(row[0])  # type: ignore
-                spatial = tuple(int(x) for x in row[1:])  # type: ignore
-                if kper not in result:
-                    result[kper] = {}
-                # flatten spatial index if only one spatial dim
-                key = spatial[0] if len(spatial) == 1 else spatial
-                result[kper][key] = v
-    return result
 
 
 def unstructure_component(value: Component) -> dict[str, Any]:
-    data = xattree.asdict(value)
-    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])
-    return data
-
-
-def unstructure_tdis(value: Any) -> dict[str, Any]:
-    data = xattree.asdict(value)
-    blocks = get_blocks(value.dfn)
-    for block_name, block in blocks.items():
-        if block_name == "perioddata":
-            arrs_d = {}
-            periods = set()  # type: ignore
-            for field_name in block.keys():
-                arr = data.get(field_name, None)
-                arr_d = {} if arr is None else unstructure_array(arr)
-                arrs_d[field_name] = arr_d
-                periods.update(arr_d.keys())
-            periods = sorted(periods)  # type: ignore
-            perioddata = {}  # type: ignore
-            for kper in periods:
-                line = []
-                if kper not in perioddata:
-                    perioddata[kper] = []  # type: ignore
-                for arr_d in arrs_d.values():
-                    if val := arr_d.get(kper, None):
-                        line.append(val)
-                perioddata[kper] = tuple(line)
-            data["perioddata"] = perioddata
-    return data
-
-
-def unstructure_chd(value: Any) -> dict[str, Any]:
-    if (parent := value.parent) is None:
-        raise ValueError(
-            "CHD cannot be unstructured without a parent "
-            "model and corresponding grid discretization."
-        )
-    grid = parent.grid
-    data = xattree.asdict(value)
-    blocks = get_blocks(value.dfn)
-    for block_name, block in blocks.items():
-        if block_name == "period":
-            arrs_d = {}
-            periods = set()  # type: ignore
-            for field_name in block.keys():
-                arr = data.get(field_name, None)
-                arr_d = {} if arr is None else unstructure_array(arr)
-                arrs_d[field_name] = arr_d
-                periods.update(arr_d.keys())
-            periods = sorted(periods)  # type: ignore
-            perioddata = {}  # type: ignore
-            for kper in periods:
-                line = []
-                if kper not in perioddata:
-                    perioddata[kper] = []  # type: ignore
-                for arr_d in arrs_d.values():
-                    if val := arr_d.get(kper, None):
-                        for nn, v in val.items():
-                            cellid = get_cellid(nn, grid)
-                            line.append((*cellid, v))
-                perioddata[kper] = tuple(line)
-            data["period"] = perioddata
-    return data
+    """Simple converter: component -> dict with validation."""
+    return xattree.asdict(value)
 
 
-def unstructure_oc(value: Any) -> dict[str, Any]:
-    data = xattree.asdict(value)
-    blocks = get_blocks(value.dfn)
-    for block_name, block in blocks.items():
-        if block_name == "period":
-            fields = []
-            for field_name, field in block.items():
-                action, rtype = field_name.split("_")
-                fields.append((action, rtype, field_name))
-            arrs_d = {}
-            periods = set()  # type: ignore
-            for action, rtype, field_name in fields:
-                arr = data.get(field_name, None)
-                arr_d = {} if arr is None else unstructure_array(arr)
-                arrs_d[(action, rtype)] = arr_d
-                periods.update(arr_d.keys())
-            periods = sorted(periods)  # type: ignore
-            perioddata = {}  # type: ignore
-            for kper in periods:
-                if kper not in perioddata:
-                    perioddata[kper] = []
-                for (action, rtype), arr_d in arrs_d.items():
-                    if arr := arr_d.get(kper, None):
-                        perioddata[kper].append((action, rtype, arr))
-            data["period"] = perioddata
-        else:
-            for field_name, field in block.items():
-                if is_list_field(field):
-                    data[field_name] = unstructure_array(data[field_name])
-    return data
+def structure_component(cls: type[Component], data: dict[str, Any]) -> Component:
+    """Simple converter: dict -> component with validation."""
+    return xattree.structure(data, cls)