jinja improvements (#211)

jinja template/filter improvements. introduce an array2const filter for constant arrays. and revise the array_how filter to determine how to write arrays (later it will be possible for the user to set this at the simulation or subcomponent level and override this automatic determination).

Co-authored-by: mjreno <[email protected]>

Author: wpbonelli

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"

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

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" %}