Skip to content

Commit 7c8c04b

Browse files
author
wpbonelli
committed
refactor input file writing
1 parent 6060f9b commit 7c8c04b

File tree

19 files changed

+675
-380
lines changed

19 files changed

+675
-380
lines changed

‎flopy4/mf6/codec/__init__.py‎

Lines changed: 1 addition & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1,90 +1,6 @@
1-
import sys
2-
from os import PathLike
3-
from typing import Any
4-
5-
import numpy as np
6-
import xattree
7-
from cattrs import Converter
8-
from jinja2 import Environment, PackageLoader
9-
10-
from flopy4.mf6 import filters
11-
from flopy4.mf6.codec.converter import (
12-
unstructure_array,
13-
unstructure_chd,
14-
unstructure_component,
15-
unstructure_oc,
16-
unstructure_tdis,
17-
)
18-
19-
_JINJA_ENV = Environment(
20-
loader=PackageLoader("flopy4.mf6"),
21-
trim_blocks=True,
22-
lstrip_blocks=True,
23-
)
24-
_JINJA_ENV.filters["dict_blocks"] = filters.dict_blocks
25-
_JINJA_ENV.filters["list_blocks"] = filters.list_blocks
26-
_JINJA_ENV.filters["field_type"] = filters.field_type
27-
_JINJA_ENV.filters["field_value"] = filters.field_value
28-
_JINJA_ENV.filters["array_how"] = filters.array_how
29-
_JINJA_ENV.filters["array_chunks"] = filters.array_chunks
30-
_JINJA_ENV.filters["array2string"] = filters.array2string
31-
32-
_JINJA_TEMPLATE_NAME = "blocks.jinja"
33-
34-
_PRINT_OPTIONS = {
35-
"precision": 4,
36-
"linewidth": sys.maxsize,
37-
"threshold": sys.maxsize,
38-
}
39-
40-
41-
def _make_converter() -> Converter:
42-
# TODO: document what is converter's responsibility vs Jinja's
43-
# TODO: how can we make sure writing remains lazy for list input?
44-
# don't eagerly unstructure to dict, lazily access from the template?
45-
46-
from flopy4.mf6.component import Component
47-
from flopy4.mf6.gwf.chd import Chd
48-
from flopy4.mf6.gwf.oc import Oc
49-
from flopy4.mf6.tdis import Tdis
50-
51-
converter = Converter()
52-
converter.register_unstructure_hook_factory(xattree.has, lambda _: xattree.asdict)
53-
converter.register_unstructure_hook(Component, unstructure_component)
54-
converter.register_unstructure_hook(Tdis, unstructure_tdis)
55-
converter.register_unstructure_hook(Chd, unstructure_chd)
56-
converter.register_unstructure_hook(Oc, unstructure_oc)
57-
return converter
58-
59-
60-
_CONVERTER = _make_converter()
61-
62-
63-
def loads(data: str) -> Any:
64-
# TODO
65-
pass
66-
67-
68-
def load(path: str | PathLike) -> Any:
69-
# TODO
70-
pass
71-
72-
73-
def dumps(data) -> str:
74-
template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
75-
with np.printoptions(**_PRINT_OPTIONS): # type: ignore
76-
return template.render(dfn=type(data).dfn, data=_CONVERTER.unstructure(data))
77-
78-
79-
def dump(data, path: str | PathLike) -> None:
80-
template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
81-
iterator = template.generate(dfn=type(data).dfn, data=_CONVERTER.unstructure(data))
82-
with np.printoptions(**_PRINT_OPTIONS), open(path, "w") as f: # type: ignore
83-
f.writelines(iterator)
84-
1+
from flopy4.mf6.codec.writer import dump, dumps, load, loads
852

863
__all__ = [
87-
"unstructure_array",
884
"loads",
895
"load",
906
"dumps",

‎flopy4/mf6/codec/converter.py‎

Lines changed: 0 additions & 154 deletions
This file was deleted.

‎flopy4/mf6/codec/reader/__init__.py‎

Whitespace-only changes.

‎flopy4/mf6/codec/writer/__init__.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
import sys
2+
from os import PathLike
3+
from typing import Any
4+
5+
import numpy as np
6+
from jinja2 import Environment, PackageLoader
7+
8+
from flopy4.mf6 import filters
9+
10+
_JINJA_ENV = Environment(
11+
loader=PackageLoader("flopy4.mf6.codec.writer"),
12+
trim_blocks=True,
13+
lstrip_blocks=True,
14+
)
15+
_JINJA_ENV.filters["dict_blocks"] = filters.dict_blocks
16+
_JINJA_ENV.filters["list_blocks"] = filters.list_blocks
17+
_JINJA_ENV.filters["array_how"] = filters.array_how
18+
_JINJA_ENV.filters["array_chunks"] = filters.array_chunks
19+
_JINJA_ENV.filters["array2string"] = filters.array2string
20+
_JINJA_ENV.filters["field_type"] = filters.field_type
21+
_JINJA_ENV.filters["array2list"] = filters.array2list
22+
_JINJA_ENV.filters["keystring2list"] = filters.keystring2list
23+
_JINJA_ENV.filters["keystring2list_multifield"] = filters.keystring2list_multifield
24+
_JINJA_TEMPLATE_NAME = "blocks.jinja"
25+
_PRINT_OPTIONS = {
26+
"precision": 4,
27+
"linewidth": sys.maxsize,
28+
"threshold": sys.maxsize,
29+
}
30+
31+
32+
def loads(data: str) -> Any:
33+
# TODO
34+
pass
35+
36+
37+
def load(path: str | PathLike) -> Any:
38+
# TODO
39+
pass
40+
41+
42+
def dumps(data) -> str:
43+
template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
44+
with np.printoptions(**_PRINT_OPTIONS): # type: ignore
45+
return template.render(data=data)
46+
47+
48+
def dump(data, path: str | PathLike) -> None:
49+
template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
50+
iterator = template.generate(data=data)
51+
with np.printoptions(**_PRINT_OPTIONS), open(path, "w") as f: # type: ignore
52+
f.writelines(iterator)

‎flopy4/mf6/codec/writer/templates/blocks.jinja‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
{% import 'macros.jinja' as macros with context %}
2+
{% for block_name, block_value in (data|dict_blocks).items() %}
3+
BEGIN {{ block_name.upper() }}
4+
{% for field_name, field_value in block_value.items() if (field_value) is not none -%}
5+
{{ macros.field(field_name, field_value) }}
6+
{%- endfor %}
7+
END {{ block_name.upper() }}
8+
9+
{% endfor %}
10+
11+
{% for block_name, block_value in (data|list_blocks).items() -%}
12+
{{ macros.list(block_name, block_value, multi=block_name in ["period"]) }}
13+
{%- endfor %}

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

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
{% macro field(name, value) %}
2+
{% set type = value|field_type %}
3+
{% if type in ['keyword', 'integer', 'double precision', 'string'] %}
4+
{{ scalar(name, value) }}
5+
{% elif type == 'record' %}
6+
{{ record(name, value) }}
7+
{% elif type == 'keystring' %}
8+
{{ keystring(name, value) }}
9+
{% elif type == 'recarray' %}
10+
{{ recarray(name, value, how=value|array_how) }}
11+
{% endif %}
12+
{% endmacro %}
13+
14+
{% macro scalar(name, value) %}
15+
{% set type = value|field_type %}
16+
{% if value is not none %}{{ name.upper() }}{% if type != 'keyword' %} {{ value }}{% endif %}{% endif %}
17+
{% endmacro %}
18+
19+
{% macro keystring(name, value) %}
20+
{% for item in value.values() -%}
21+
{{ field(item) }}
22+
{%- endfor %}
23+
{% endmacro %}
24+
25+
{% macro record(name, value) %}
26+
{% for itemitem in value.values() -%}
27+
{% if item.tagged %}{{ item.name.upper() }} {% endif %}{{ field(item) }}
28+
{%- endfor %}
29+
{% endmacro %}
30+
31+
{% macro recarray(name, value, how="internal") %}
32+
{{ name.upper() }}{% if "layered" in how %} LAYERED{% endif %}
33+
34+
{% if how == "constant" %}
35+
CONSTANT {{ value.item() }}
36+
{% elif how == "layered constant" %}
37+
{% for layer in value -%}
38+
CONSTANT {{ layer.item() }}
39+
{%- endfor %}
40+
{% elif how == "internal" %}
41+
INTERNAL
42+
{% for chunk in value|array_chunks -%}
43+
{{ chunk|array2string }}
44+
{%- endfor %}
45+
{% elif how == "external" %}
46+
OPEN/CLOSE {{ value }}
47+
{% endif %}
48+
{% endmacro %}
49+
50+
{% macro list(name, value, multi=False) %}
51+
{% if multi %}
52+
{# iterate through time periods and combine fields #}
53+
{% set field_arrays = {} %}
54+
{% for field_name, field_value in value.items() %}
55+
{% if field_value is not none %}
56+
{% set _ = field_arrays.update({field_name: field_value}) %}
57+
{% endif %}
58+
{% endfor %}
59+
60+
{% set first_array = field_arrays.values() | list | first %}
61+
{% if first_array is not none %}
62+
{% set nper = first_array.shape[0] %}
63+
{% for period_idx in range(nper) %}
64+
BEGIN {{ name.upper() }} {{ period_idx + 1 }}
65+
{% for row in (field_arrays|keystring2list_multifield(period_idx)) %}
66+
{{ row|join(" ") }}
67+
{% endfor %}
68+
END {{ name.upper() }} {{ period_idx + 1 }}
69+
70+
{% endfor %}
71+
{% endif %}
72+
{% else %}
73+
{% for field_name, field_value in value.items() %}
74+
{% if field_value is not none %}
75+
BEGIN {{ name.upper() }}
76+
{% for row in (field_value|array2list) %}
77+
{{ row|join(" ") }}
78+
{% endfor %}
79+
END {{ name.upper() }}
80+
{% endif %}
81+
{% endfor %}
82+
{% endif %}
83+
{% endmacro %}

0 commit comments

Comments
 (0)