jinja improvements

also update pixi section in pyproject.toml

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