wip

Author: wpbonelli

flopy4/mf6/component.py

@@ -2,13 +2,14 @@
 from collections.abc import MutableMapping
 
 from xattree import xattree
+from flopy4.io import Writer
 
 COMPONENTS = {}
 """MF6 component registry."""
 
 
 @xattree
-class Component(ABC, MutableMapping):
+class Component(ABC, MutableMapping, Writer):
     @classmethod
     def __attrs_init_subclass__(cls):
         COMPONENTS[cls.__name__.lower()] = cls
@@ -27,8 +28,3 @@ def __iter__(self):
 
     def __len__(self):
         return len(self.children)  # type: ignore
-
-    def write(self) -> None:
-        # TODO: write with jinja to file
-        for child in self.children.values():  # type: ignore
-            child.write()

flopy4/mf6/filters.py

@@ -0,0 +1,27 @@
+from collections.abc import Iterable, Mapping
+from inspect import isclass
+import types
+from typing import Union, get_args, get_origin
+import numpy as np
+import xarray as xr
+from jinja2 import pass_context
+
+from attrs import Attribute
+from xattree import Xattribute
+
+
+@pass_context
+def value(ctx, field: Xattribute):
+    """Return the kind of the field."""
+    # TODO
+    pass
+
+
+def dask_expand(data: xr.DataArray):
+    for block in data.data.to_delayed():
+        block_data = block.compute()
+        yield block_data
+
+
+def nparray2string(data: np.ndarray):
+    return np.array2string(data, separator=" ")[1:-1]  # remove brackets

flopy4/mf6/io.py

@@ -0,0 +1,44 @@
+import sys
+import numpy as np
+from jinja2 import Environment, PackageLoader
+from flopy4.mf6 import filters
+from flopy4.mf6.spec import blocks_dict
+
+
+env = Environment(
+loader=PackageLoader("flopy4.mf6"),
+    trim_blocks=True,
+    lstrip_blocks=True,
+)
+env.filters["dask_expand"] = filters.dask_expand
+env.filters["nparray2string"] = filters.nparray2string
+
+
+class Writer:
+    def _write_ascii(self, path) -> None:
+        # TODO: factor out an ascii writer separately
+
+        block_spec = blocks_dict(type(self))
+        blocks = {}
+        for block_name, block in block_spec.items():
+            blocks[block_name] = {}
+            for field_name, field in block.items():
+                if field_name == "data":
+                    continue
+                blocks[block_name][field_name] = {
+                    "spec": field,
+                    "value": getattr(self, field_name),
+                }
+
+        template = env.get_template("blocks.jinja")
+        iterator = template.generate(blocks=blocks)
+        with np.printoptions(
+            precision=4, linewidth=sys.maxsize, threshold=sys.maxsize
+        ):
+            with open(path, "w") as f:
+                f.writelines(iterator)
+
+    def write(self) -> None:
+        self._write_ascii()
+        for child in self.children.values():  # type: ignore
+            child.write()

flopy4/mf6/templates/blocks.jinja

@@ -1,17 +1,8 @@
 {% for block in blocks %}
-begin {{ block.name }} {{ block.index }}
-{% for var in block %}
-{% if var.kind in ['keyword', 'integer', 'double precision'] %}
-{{ macros.scalar(var) }}
-{% elif var.kind == 'keystring' %}
-{{ macros.union(var) }}
-{% elif var.kind == 'record' %}
-{{ macros.record(var) }}
-{% elif var.kind == 'array' %}
-{{ macros.array(var) }}
-{% elif var.kind == 'list' %}
-{{ macros.list(var) }}
-{% endif %}
+BEGIN {{ block.name }} {% if block.index is not none %}{{ block.index }}{% endif %}
+{% for field in block %}
+{{ macros.field(field) }}
 {% endfor %}
-end {{ block.name }}
+END {{ block.name }}
+
 {% endfor %}

flopy4/mf6/templates/macros.jinja

@@ -1,36 +1,58 @@
-{% macro scalar(var) %}
-{{ var.name }} {% endif %}{{ var.value }}
+{% macro field(field) %}
+{% if field|kind == 'scalar' %}
+{{ macros.scalar(field) }}
+{% elif field|kind == 'union' %}
+{{ macros.union(field) }}
+{% elif field|kind == 'record' %}
+{{ macros.record(field) }}
+{% elif field|kind == 'array' %}
+{{ macros.array(field) }}
+{% elif field|kind == 'list' %}
+{{ macros.list(field) }}
+{% endif %}
 {% endmacro %}
 
-{% macro union(var) %}
+{% macro scalar(field) %}
+{{ field.name }} {% endif %}{{ field.value }}
 {% endmacro %}
 
-{% macro record(var) %}
-{% if var.tagged %}{{ var.name }} {% endif %}{{ var.value|join(' ') }}
+{% macro union(field) %}
+{% for item in field.value.items() %}
+{{ macros.field(item) }}
+{% endfor %}
 {% endmacro %}
 
-{% macro array(var) %}
-{% if constant %}{% if layered %}
-{# layered constant #}
-{{ var.name }} LAYERED
-{% for val in var.value %}
+{% macro record(field) %}
+{% for item in field.values() %}{% if item.tagged %}{{ item.name }} {% endif %}{{ macros.field(field) }}{%
+endfor
+%}
+{% endmacro %}
+
+{% macro array(field, how="internal") %}
+{% if how == "layered constant" %}
+{{ field.name }} LAYERED
+{% for val in field.value %}
 CONSTANT
 {% endfor %}
-{% else %}
-{# constant #}
-{{ var.name }} CONSTANT {{ var.value }}
-{% endif %}
-{% else %}
-{# layered #}
+{% elif how == "constant" %}
+{{ field.name }} CONSTANT {{ field.value }}
+{% elif how == "layered" %}
 {% if layered %}
-{{ var.name }}
-{% for val in var.value %}{{ val }}{% endfor %}
-{% else %}
-{# grid-shaped #}
-{{ var.name }} {{ var.value }}
-{% endif %}
+{{ field.name }}
+{% for val in field.value %}{{ val }}{% endfor %}
+{% elif how == "internal" %}
+{{ field.name }} {{ macro.internal_array(field) }}
+{% elif how == "external" %}
+{{ field.name}} OPEN/CLOSE {{ field.value }}
 {% endif %}
 {% endmacro %}
 
-{% macro list(var) %}
+{% macro internal_array(field) %}
+{% for chunk in field.value|dask_expand %}
+{{ chunk|nparray2string }}
+{% endfor %}
+{% endmacro %}
+
+{% macro list(field) %}
+{# TODO #}
 {% endmacro %}

pyproject.toml

@@ -34,8 +34,8 @@ classifiers = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "attrs",                                                                            # todo: bounds?
-    "cattrs",                                                                           # todo: bounds?
+    "attrs",
+    "cattrs",
     "flopy",
     "Jinja2>=3.0",
     "numpy>=1.20.3",