diff --git a/flopy4/mf6/codec.py b/flopy4/mf6/codec.py index d84f4efc..7ad9f3b2 100644 --- a/flopy4/mf6/codec.py +++ b/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 diff --git a/flopy4/mf6/filters.py b/flopy4/mf6/filters.py index 5d0c955b..d3fa1cf1 100644 --- a/flopy4/mf6/filters.py +++ b/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: diff --git a/flopy4/mf6/spec.py b/flopy4/mf6/spec.py index dfaecf09..87a31b39 100644 --- a/flopy4/mf6/spec.py +++ b/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]: diff --git a/flopy4/mf6/templates/blocks.jinja b/flopy4/mf6/templates/blocks.jinja index 4d6719e3..4081963a 100644 --- a/flopy4/mf6/templates/blocks.jinja +++ b/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 %} diff --git a/flopy4/mf6/templates/macros.jinja b/flopy4/mf6/templates/macros.jinja index 3603cd1b..a34ffcd2 100644 --- a/flopy4/mf6/templates/macros.jinja +++ b/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 %}