diff --git a/flopy4/mf6/codec/writer/__init__.py b/flopy4/mf6/codec/writer/__init__.py index 811db456..df030400 100644 --- a/flopy4/mf6/codec/writer/__init__.py +++ b/flopy4/mf6/codec/writer/__init__.py @@ -13,7 +13,8 @@ ) _JINJA_ENV.filters["field_type"] = filters.field_type _JINJA_ENV.filters["array_how"] = filters.array_how -_JINJA_ENV.filters["array_chunks"] = filters.array_chunks +_JINJA_ENV.filters["array2const"] = filters.array2const +_JINJA_ENV.filters["array2chunks"] = filters.array2chunks _JINJA_ENV.filters["array2string"] = filters.array2string _JINJA_ENV.filters["data2list"] = filters.data2list _JINJA_TEMPLATE_NAME = "blocks.jinja" diff --git a/flopy4/mf6/codec/writer/filters.py b/flopy4/mf6/codec/writer/filters.py index 2e3c4815..25c89e88 100644 --- a/flopy4/mf6/codec/writer/filters.py +++ b/flopy4/mf6/codec/writer/filters.py @@ -1,11 +1,12 @@ from collections.abc import Hashable, Mapping from io import StringIO -from typing import Any +from typing import Any, Literal import numpy as np import xarray as xr from modflow_devtools.dfn.schema.v2 import FieldType from numpy.typing import NDArray +from xattree import Scalar from flopy4.mf6.constants import FILL_DNODATA @@ -32,14 +33,34 @@ def field_type(value: Any) -> FieldType: raise ValueError(f"Unsupported field type: {type(value)}") -def array_how(value: xr.DataArray) -> str: - # TODO - # - detect constant arrays? - # - above certain size, use external? +ArrayHow = Literal["constant", "internal", "external"] + + +def array_how(value: xr.DataArray) -> ArrayHow: + """ + Determine how an array should be represented in MF6 input. + Options are "constant", "internal", or "external". If the + array dask-backed, assumed it's big and return "external". + Otherwise there is no materialization cost to check if all + values are the same, so return "constant" or "internal" as + appropriate. + """ + if hasattr(value.data, "blocks"): + return "external" + if value.max() == value.min(): + return "constant" return "internal" -def array_chunks(value: xr.DataArray, chunks: Mapping[Hashable, int] | None = None): +def array2const(value: xr.DataArray) -> Scalar: + if np.issubdtype(value.dtype, np.integer): + return value.max().item() + if np.issubdtype(value.dtype, np.floating): + return f"{value.max().item():.8f}" + return value.ravel()[0] + + +def array2chunks(value: xr.DataArray, chunks: Mapping[Hashable, int] | None = None): """ Yield chunks from a dask-backed array of up to 3 dimensions. If it's not already chunked, split it into chunks of the diff --git a/flopy4/mf6/codec/writer/templates/macros.jinja b/flopy4/mf6/codec/writer/templates/macros.jinja index 2c8c5747..29c1369a 100644 --- a/flopy4/mf6/codec/writer/templates/macros.jinja +++ b/flopy4/mf6/codec/writer/templates/macros.jinja @@ -32,14 +32,14 @@ {{ inset ~ name.upper() }}{% if "layered" in how %} LAYERED{% endif %} {% if how == "constant" %} -CONSTANT {{ value.item() }} +CONSTANT {{ value|array2const }} {% elif how == "layered constant" %} {% for layer in value -%} -CONSTANT {{ layer.item() }} +CONSTANT {{ layer|array2const }} {%- endfor %} {% elif how == "internal" %} INTERNAL -{% for chunk in value|array_chunks -%} +{% for chunk in value|array2chunks -%} {{ (2 * inset) ~ chunk|array2string }} {%- endfor %} {% elif how == "external" %}