ic and oc looking better. several fixes. (#143)

Author: wpbonelli

flopy4/mf6/codec.py

@@ -13,7 +13,7 @@
     lstrip_blocks=True,
 )
 JINJA_ENV.filters["blocks"] = filters.blocks
-JINJA_ENV.filters["field_kind"] = filters.field_type
+JINJA_ENV.filters["field_type"] = filters.field_type
 JINJA_ENV.filters["field_value"] = filters.field_value
 JINJA_ENV.filters["array_delay"] = filters.array_delay
 JINJA_ENV.filters["array2string"] = filters.array2string

flopy4/mf6/filters.py

@@ -6,9 +6,20 @@
 from modflow_devtools.dfn import Dfn, Field
 from numpy.typing import NDArray
 
+from flopy4.mf6.spec import block_sort_key
+
 
 def blocks(dfn: Dfn) -> dict:
-    return {k: v for k, v in dfn.items() if k not in Dfn.__annotations__}
+    """
+    Get blocks from an MF6 input definition. Anything not an
+    explicitly defined key in the `Dfn` typed dict is a block.
+    """
+    return dict(
+        sorted(
+            {k: v for k, v in dfn.items() if k not in Dfn.__annotations__}.items(),
+            key=block_sort_key,
+        )
+    )
 
 
 def field_type(field: Field) -> str:
@@ -25,17 +36,26 @@ def field_value(ctx, field: Field):
     return getattr(ctx["data"], field["name"])
 
 
-def array_delay(value: xr.DataArray):
-    """Yield chunks (lines) from a Dask array."""
-    # TODO: Determine a good chunk size,
-    # because if the underlying array is only numpy, it will stay one block.
-    for chunk in value.chunk():
+def array_delay(value: xr.DataArray, chunks=None):
+    """
+    Yield chunks from an array. Each chunk becomes a line in the file.
+    If the array is not already chunked, it is chunked using the given
+    chunk size. If no chunk size is provided, the entire array becomes
+    a single chunk.
+    """
+    if value.chunks is None:
+        chunk_shape = chunks or {dim: size for dim, size in zip(value.dims, value.shape)}
+        value = value.chunk(chunk_shape)
+    for chunk in value.data.blocks:
         yield chunk.compute()
 
 
 def array2string(value: NDArray) -> str:
     """Convert an array to a string."""
-    return np.array2string(value, separator=" ")[1:-1]  # remove brackets
+    s = np.array2string(value, separator=" ")
+    if value.shape != ():
+        s = s[1:-1]  # remove brackets
+    return s.replace("'", "").replace('"', "")  # remove quotes
 
 
 def is_dict(value: Any) -> bool:

flopy4/mf6/spec.py

@@ -119,7 +119,7 @@ def array(
 Block = dict[str, Attribute]
 
 
-def _block_sort_key(item) -> int:
+def block_sort_key(item: tuple[str, dict]) -> int:
     k, _ = item
     if k == "options":
         return 0
@@ -153,7 +153,7 @@ def blocks_dict(cls) -> dict[str, Block]:
         if block not in blocks:
             blocks[block] = {}
         blocks[block][k] = v
-    return dict(sorted(blocks.items(), key=_block_sort_key))
+    return dict(sorted(blocks.items(), key=block_sort_key))
 
 
 def fields(cls) -> list[Attribute]:

flopy4/mf6/templates/blocks.jinja

@@ -1,8 +1,9 @@
 {% import 'macros.jinja' as macros with context %}
 {% for block_name, block_ in (dfn|blocks).items() %}
 BEGIN {{ block_name }}
-{% for field in block_.values() %}
+{% for field in block_.values() -%}
 {{ macros.field(field) }}
-{% endfor %}
+{%- endfor %}
 END {{ block_name }}
+
 {% endfor %}

flopy4/mf6/templates/macros.jinja

@@ -1,59 +1,54 @@
 {% macro field(field) %}
-{% set kind = field|field_kind %}
-{% if kind == 'attr' %}
-{% if field.type|is_dict %}{{ record(field) }}{% else %}{{ scalar(field) }}
-{% endif %}
-{% elif kind in ['array', 'coord'] %}
-{{ array(field) }}
-{% elif kind == 'dim' %}
+{% set type = field|field_type %}
+{% if type in ['keyword', 'integer', 'double precision', 'string'] %}
 {{ scalar(field) }}
-{% elif kind == 'child' %}
-{# TODO #}
-{% endif -%}
-{%- endmacro %}
+{% elif type == 'record' %}
+{{ record(field) }}
+{% elif type == 'keystring' %}
+{{ keystring(field) }}
+{% elif type == 'recarray' %}
+{{ recarray(field) }}
+{% endif %}
+{% endmacro %}
 
-{% macro scalar(field) -%}
+{% macro scalar(field) %}
 {% set value = field|field_value %}
 {% if value is not none %}{{ field.name }} {{ value }}{% endif %}
-{%- endmacro %}
+{% endmacro %}
 
-{% macro keystring(field) %} {# union #}
-{% for item in (field|field_value).items() %}
+{% macro keystring(field) %}
+{% for item in (field|field_value).values() -%}
 {{ field(item) }}
-{% endfor %}
-{%- endmacro %}
+{%- endfor %}
+{% endmacro %}
 
 {% macro record(field) %}
-{% for item in field|field_value %}
-{% if item.tagged %}{{ item.name }} {% endif %}{{ field(field) }}
-{% endfor %}
-{%- endmacro %}
+{% for item in (field|field_value).values() -%}
+{% if item.tagged %}{{ item.name }} {% endif %}{{ field(item) }}
+{%- endfor %}
+{% endmacro %}
 
-{% macro array(field, how="internal") %}
+{% macro recarray(field, how="internal") %}
 {% if how == "layered constant" %}
 {{ field.name }} LAYERED
 {% for val in field|field_value %}
 CONSTANT
 {% endfor %}
-{% elif how == "constant" %}
+{%- elif how == "constant" %}
 {{ field.name }} CONSTANT {{ field|field_value }}
-{% elif how == "layered" %}
+{%- elif how == "layered" %}
 {% if layered %}
 {{ field.name }}{% for val in field|field_value %} {{ val }}{% endfor %}
 {% endif %}
-{% elif how == "internal" %}
+{%- elif how == "internal" %}
 {{ field.name }} {{ internal_array(field) }}
-{% elif how == "external" %}
+{%- elif how == "external" %}
 {{ field.name}} OPEN/CLOSE {{ field|field_value }}
 {% endif %}
-{%- endmacro %}
+{% endmacro %}
 
 {% macro internal_array(field) %}
-{% for chunk in field|field_value|array_delay %}
+{% for chunk in field|field_value|array_delay -%}
 {{ chunk|array2string }}
-{% endfor %}
-{%- endmacro %}
-
-{% macro list(field) %}
-{# TODO #}
-{%- endmacro %}
+{%- endfor %}
+{% endmacro %}