diff --git a/flopy4/mf6/codec/__init__.py b/flopy4/mf6/codec/__init__.py index 34f14966..b9d5f05b 100644 --- a/flopy4/mf6/codec/__init__.py +++ b/flopy4/mf6/codec/__init__.py @@ -13,6 +13,7 @@ unstructure_array, unstructure_component, unstructure_oc, + unstructure_tdis, ) _JINJA_ENV = Environment( @@ -40,10 +41,12 @@ def _make_converter() -> Converter: from flopy4.mf6.component import Component from flopy4.mf6.gwf.oc import Oc + from flopy4.mf6.tdis import Tdis converter = Converter() converter.register_unstructure_hook_factory(xattree.has, lambda _: xattree.asdict) converter.register_unstructure_hook(Component, unstructure_component) + converter.register_unstructure_hook(Tdis, unstructure_tdis) converter.register_unstructure_hook(Oc, unstructure_oc) return converter diff --git a/flopy4/mf6/codec/converter.py b/flopy4/mf6/codec/converter.py index b614d713..ee0ad632 100644 --- a/flopy4/mf6/codec/converter.py +++ b/flopy4/mf6/codec/converter.py @@ -143,9 +143,39 @@ def unstructure_component(value: Component) -> dict[str, Any]: return data +def unstructure_tdis(value: Any) -> dict[str, Any]: + data = xattree.asdict(value) + blocks = get_blocks(value.dfn) + for block_name, block in blocks.items(): + if block_name == "perioddata": + array_fields = list(block.keys()) + + # Unstructure all arrays and collect all unique periods + arrays = {} + periods = set() # type: ignore + for field_name in array_fields: + arr = unstructure_array(data.get(field_name, {})) + arrays[field_name] = arr + periods.update(arr.keys()) + periods = sorted(periods) # type: ignore + + perioddata = {} # type: ignore + for kper in periods: + line = [] + for arr in arrays.values(): + if kper not in perioddata: + perioddata[kper] = [] # type: ignore + line.append(arr[kper]) + perioddata[kper] = tuple(line) + + data["perioddata"] = perioddata + return data + + def unstructure_oc(value: Any) -> dict[str, Any]: data = xattree.asdict(value) - for block_name, block in get_blocks(value.dfn).items(): + blocks = get_blocks(value.dfn) + for block_name, block in blocks.items(): if block_name == "period": # Dynamically collect all recarray fields in perioddata block array_fields = [] @@ -156,16 +186,15 @@ def unstructure_oc(value: Any) -> dict[str, Any]: # Unstructure all arrays and collect all unique periods arrays = {} - all_periods = set() # type: ignore + periods = set() # type: ignore for action, rtype, field_name in array_fields: arr = unstructure_array(data.get(field_name, {})) arrays[(action, rtype)] = arr - if isinstance(arr, dict): - all_periods.update(arr.keys()) - all_periods = sorted(all_periods) # type: ignore + periods.update(arr.keys()) + periods = sorted(periods) # type: ignore perioddata = {} # type: ignore - for kper in all_periods: + for kper in periods: for (action, rtype), arr in arrays.items(): if kper in arr: if kper not in perioddata: diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index 72a4b0de..78b658d0 100644 --- a/flopy4/mf6/component.py +++ b/flopy4/mf6/component.py @@ -100,14 +100,14 @@ def _preio(self, format: str) -> None: self.filename = self.default_filename() def load(self, format: str) -> None: - """Load the component from an input file.""" + """Load the component and any children.""" self._preio(format=format) self._load(format=format) for child in self.children.values(): # type: ignore child.load(format=format) def write(self, format: str) -> None: - """Write the component to an input file.""" + """Write the component and any children.""" self._preio(format=format) self._write(format=format) for child in self.children.values(): # type: ignore diff --git a/flopy4/mf6/filters.py b/flopy4/mf6/filters.py index 596a36f0..4b7d3e22 100644 --- a/flopy4/mf6/filters.py +++ b/flopy4/mf6/filters.py @@ -25,19 +25,17 @@ def dict_blocks(dfn: Dfn) -> dict: one or more fields, as opposed to a list block, which may only contain one recarray field, using list input. """ - x = { + return { block_name: block for block_name, block in get_blocks(dfn).items() if not _is_list_block(block) } - return x def list_blocks(dfn: Dfn) -> dict: - x = { + return { block_name: block for block_name, block in get_blocks(dfn).items() if _is_list_block(block) } - return x def field_type(field: Field) -> str: diff --git a/flopy4/mf6/simulation.py b/flopy4/mf6/simulation.py index bb1203f3..46484f01 100644 --- a/flopy4/mf6/simulation.py +++ b/flopy4/mf6/simulation.py @@ -54,11 +54,11 @@ def run(self, exe: str | PathLike = "mf6", verbose: bool = False) -> None: ) def load(self, format="ascii"): - """Load the simulation in the specified format.""" + """Load the simulation.""" with cd(self.workspace): - super().load(format) + super().load(format=format) def write(self, format="ascii"): - """Write the simulation in the specified format.""" + """Write the simulation.""" with cd(self.workspace): - super().write(format) + super().write(format=format) diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index 91f3bea5..79e8396b 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -33,18 +33,21 @@ class PeriodData: default=1.0, dims=("nper",), converter=Converter(structure_array, takes_self=True, takes_field=True), + reader="urword", ) nstp: NDArray[np.integer] = array( block="perioddata", default=1, dims=("nper",), converter=Converter(structure_array, takes_self=True, takes_field=True), + reader="urword", ) tsmult: NDArray[np.floating] = array( block="perioddata", default=1.0, dims=("nper",), converter=Converter(structure_array, takes_self=True, takes_field=True), + reader="urword", ) def to_time(self) -> ModelTime: diff --git a/flopy4/mf6/templates/blocks.jinja b/flopy4/mf6/templates/blocks.jinja index 1b0ba597..d15be65b 100644 --- a/flopy4/mf6/templates/blocks.jinja +++ b/flopy4/mf6/templates/blocks.jinja @@ -9,5 +9,5 @@ END {{ block_name.upper() }} {% endfor %} {% for block_name, block_ in (dfn|list_blocks).items() -%} -{{ macros.list(block_name, block_) }} +{{ macros.list(block_name, block_, stress=block_name == "period") }} {%- endfor%} diff --git a/flopy4/mf6/templates/macros.jinja b/flopy4/mf6/templates/macros.jinja index 63baf5b2..b1dc3f28 100644 --- a/flopy4/mf6/templates/macros.jinja +++ b/flopy4/mf6/templates/macros.jinja @@ -50,7 +50,7 @@ OPEN/CLOSE {{ value }} {% endif %} {% endmacro %} -{% macro list(block_name, block) %} +{% macro list(block_name, block, stress=False) %} {# from mf6's perspective, a list block (e.g. period data) always has just one variable, whose elements might be @@ -65,8 +65,9 @@ made sure exists in a sparse dict representation of an array. we need to expand this into a block for each stress period. #} -{% set dict = data[block_name] %} -{% for kper, value in dict.items() %} +{% set d = data[block_name] %} +{% if stress %} +{% for kper, value in d.items() %} BEGIN {{ block_name.upper() }} {{ kper }} {% for line in value %} {{ line|join(" ")|upper }} @@ -74,4 +75,11 @@ BEGIN {{ block_name.upper() }} {{ kper }} END {{ block_name.upper() }} {{ kper }} {% endfor %} +{% else %} +BEGIN {{ block_name.upper() }} +{% for line in d.values() %} +{{ line|join(" ")|upper }} +{% endfor %} +END {{ block_name.upper() }} +{% endif %} {% endmacro %}