configurable write context (#276)

First pass at #256. WriteContext class inspired by imod-python's, to control how MODFLOW 6 input files are written, including float precision, binary vs ASCII format, path handling, etc.

There are 2 usage patterns:

1. Pass to `write()`: `sim.write(context=WriteContext(float_precision=4))`
2. Context manager: `with WriteContext(float_precision=8): sim.write()`

Nested contexts are supported, with inner contexts overwriting settings from outer ones.

The context manager approach is thread-local so should support multi-threaded writing later on. Precision defaults to 8 decimal places to match the current default in the writer jinja filters.

Author: wpbonelli

flopy4/mf6/__init__.py

@@ -40,11 +40,16 @@ def _load_toml(path: Path) -> Component:
         return structure(load_toml(fp), path)
 
 
-def _write_mf6(component: Component) -> None:
+def _write_mf6(component: Component, context=None, **kwargs) -> None:
+    from flopy4.mf6.write_context import WriteContext
+
+    # Use provided context or default
+    ctx = context if context is not None else WriteContext.default()
+
     with open(component.path, "w") as fp:
         data = unstructure(component)
         try:
-            dump_mf6(data, fp)
+            dump_mf6(data, fp, context=ctx)
         except Exception as e:
             raise WriteError(
                 f"Failed to write MF6 format file for component '{component.name}' "  # type: ignore

flopy4/mf6/codec/writer/__init__.py

@@ -19,21 +19,69 @@
 _JINJA_ENV.filters["array2string"] = writer_filters.array2string
 _JINJA_ENV.filters["data2list"] = writer_filters.data2list
 _JINJA_TEMPLATE_NAME = "blocks.jinja"
-_PRINT_OPTIONS = {
-    "precision": 4,
-    "linewidth": sys.maxsize,
-    "threshold": sys.maxsize,
-}
 
 
-def dumps(data) -> str:
+def _get_print_options(context=None):
+    """Get numpy print options from WriteContext."""
+    if context is not None:
+        return context.to_numpy_printoptions()
+    # Default options
+    return {
+        "precision": 4,
+        "linewidth": sys.maxsize,
+        "threshold": sys.maxsize,
+    }
+
+
+def dumps(data, context=None) -> str:
+    """
+    Serialize data to MF6 format string.
+
+    Parameters
+    ----------
+    data : dict
+        Data to serialize
+    context : WriteContext, optional
+        Configuration context for writing
+
+    Returns
+    -------
+    str
+        Serialized MF6 format string
+    """
+    from flopy4.mf6.write_context import WriteContext
+
+    if context is None:
+        context = WriteContext.default()
+
     template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
-    with np.printoptions(**_PRINT_OPTIONS):  # type: ignore
-        return template.render(blocks=data)
+    print_opts = _get_print_options(context)
+    with np.printoptions(**print_opts):  # type: ignore
+        result = template.render(blocks=data, context=context)
+
+    return result
+
+
+def dump(data, fp: IO[str], context=None) -> None:
+    """
+    Serialize data to MF6 format and write to file.
+
+    Parameters
+    ----------
+    data : dict
+        Data to serialize
+    fp : IO[str]
+        File pointer to write to
+    context : WriteContext, optional
+        Configuration context for writing
+    """
+    from flopy4.mf6.write_context import WriteContext
 
+    if context is None:
+        context = WriteContext.default()
 
-def dump(data, fp: IO[str]) -> None:
     template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
-    iterator = template.generate(blocks=data)
-    with np.printoptions(**_PRINT_OPTIONS):  # type: ignore
+    iterator = template.generate(blocks=data, context=context)
+    print_opts = _get_print_options(context)
+    with np.printoptions(**print_opts):  # type: ignore
         fp.writelines(iterator)

flopy4/mf6/codec/writer/filters.py

@@ -41,11 +41,21 @@ def array_how(value: xr.DataArray) -> ArrayHow:
     raise ValueError(f"Arrays with ndim > 3 are not supported, got ndim={value.ndim}")
 
 
-def array2const(value: xr.DataArray) -> Scalar:
+def array2const(value: xr.DataArray, precision: int = 8) -> Scalar:
+    """
+    Convert array to constant scalar value.
+
+    Parameters
+    ----------
+    value : xr.DataArray
+        Array to convert
+    precision : int, optional
+        Number of decimal places for float output. Default is 8.
+    """
     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 f"{value.max().item():.{precision}f}"
     return value.ravel()[0]
 
 
@@ -95,14 +105,21 @@ def array2chunks(value: xr.DataArray, chunks: Mapping[Hashable, int] | None = No
         yield np.squeeze(value.values)
 
 
-def array2string(value: NDArray) -> str:
+def array2string(value: NDArray, precision: int = 9) -> str:
     """
     Convert an array to a string. The array can be 1D or 2D.
     If the array is 1D, it is converted to a 1-line string,
     with elements separated by whitespace. If the array is
     2D, each row becomes a line in the string.
 
     Used for writing array-based input to MF6 input files.
+
+    Parameters
+    ----------
+    value : NDArray
+        Array to convert
+    precision : int, optional
+        Number of decimal places for float output. Default is 9.
     """
     buffer = StringIO()
     value = np.asarray(value)
@@ -112,13 +129,14 @@ def array2string(value: NDArray) -> str:
         # add an axis to 1d arrays so np.savetxt writes elements on 1 line
         value = value[None]
     value = np.atleast_1d(value)
-    format = (
-        "%d"
-        if np.issubdtype(value.dtype, np.integer)
-        else "%.9e"
-        if np.issubdtype(value.dtype, np.floating)
-        else "%s"
-    )
+
+    if np.issubdtype(value.dtype, np.floating):
+        format = f"%.{precision}e"
+    elif np.issubdtype(value.dtype, np.integer):
+        format = "%d"
+    else:
+        format = "%s"
+
     np.savetxt(buffer, value, fmt=format, delimiter=" ")
     return buffer.getvalue().strip()
 

flopy4/mf6/codec/writer/templates/macros.jinja

@@ -32,23 +32,23 @@
 {{ inset ~ name.upper() }}{% if "layered" in how %} LAYERED{% endif %}
 
 {% if how == "constant" %}
-{{ inset }}CONSTANT {{ value|array2const }}
+{{ inset }}CONSTANT {{ value|array2const(context.float_precision) }}
 {% elif how == "layered constant" %}
 {% for layer in value -%}
-{{ "\n" ~ inset }}CONSTANT {{ layer|array2const }}
+{{ "\n" ~ inset }}CONSTANT {{ layer|array2const(context.float_precision) }}
 {%- endfor %}
 {% elif how == "layered internal" %}
 {% for layer in value %}
 
 {{ inset }}INTERNAL
 {% for chunk in layer|array2chunks -%}
-{{ (2 * inset) ~ chunk|array2string }}
+{{ (2 * inset) ~ chunk|array2string(context.float_precision) }}
 {%- endfor %}
 {%- endfor %}
 {% elif how == "internal" %}
 {{ inset }}INTERNAL
 {% for chunk in value|array2chunks -%}
-{{ (2 * inset) ~ chunk|array2string }}
+{{ (2 * inset) ~ chunk|array2string(context.float_precision) }}
 {%- endfor %}
 {% elif how == "external" %}
 OPEN/CLOSE {{ value }}

flopy4/mf6/component.py

@@ -1,7 +1,7 @@
 from abc import ABC
 from collections.abc import MutableMapping
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Optional
 
 from attrs import fields
 from modflow_devtools.dfn import Dfn, Field
@@ -12,6 +12,7 @@
 from flopy4.mf6.constants import MF6
 from flopy4.mf6.spec import field, fields_dict, to_field
 from flopy4.mf6.utils.grid import update_maxbound
+from flopy4.mf6.write_context import WriteContext
 from flopy4.uio import IO, Loader, Writer
 
 COMPONENTS = {}
@@ -142,16 +143,31 @@ def load(self, format: str = MF6) -> None:
         for child in self.children.values():  # type: ignore
             child.load(format=format)
 
-    def write(self, format: str = MF6) -> None:
-        """Write the component and any children."""
+    def write(self, format: str = MF6, context: Optional[WriteContext] = None) -> None:
+        """
+        Write the component and any children.
+
+        Parameters
+        ----------
+        format : str, optional
+            Output format. Default is MF6.
+        context : WriteContext, optional
+            Configuration context for writing. If not provided,
+            uses the current context from the context manager stack,
+            or default settings.
+        """
         # TODO: setting filename is a temp hack to get the parent's
         # name as this component's filename stem, if it has one. an
         # actual solution is to auto-set the filename when children
         # are attached to parents.
         self.filename = self.filename or self.default_filename()
-        self._write(format=format)
+
+        # Determine active context: provided > current > default
+        active_context = context or WriteContext.current()
+
+        self._write(format=format, context=active_context)
         for child in self.children.values():  # type: ignore
-            child.write(format=format)
+            child.write(format=format, context=context)
 
     def to_dict(self, blocks: bool = False, strict: bool = False) -> dict[str, Any]:
         """

flopy4/mf6/context.py

@@ -52,9 +52,9 @@ def load(self, format=MF6):
         with cd(self.workspace):
             super().load(format=format)
 
-    def write(self, format=MF6):
+    def write(self, format=MF6, context=None):
         with cd(self.workspace):
-            super().write(format=format)
+            super().write(format=format, context=context)
 
     def to_xarray(self):
         return self.data  # type: ignore