diff --git a/flopy4/mf6/codec.py b/flopy4/mf6/codec.py deleted file mode 100644 index f0c67362..00000000 --- a/flopy4/mf6/codec.py +++ /dev/null @@ -1,35 +0,0 @@ -import sys -from os import PathLike - -import numpy as np -from jinja2 import Environment, PackageLoader - -from flopy4.mf6 import filters - -JINJA_ENV = Environment( - loader=PackageLoader("flopy4.mf6"), - trim_blocks=True, - lstrip_blocks=True, -) -JINJA_ENV.filters["blocks"] = filters.blocks -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 -JINJA_ENV.filters["is_dict"] = filters.is_dict -JINJA_TEMPLATE_NAME = "blocks.jinja" - - -def load(path: str | PathLike) -> None: - # TODO - pass - - -def dump(data, path: str | PathLike) -> None: - template = JINJA_ENV.get_template(JINJA_TEMPLATE_NAME) - iterator = template.generate(dfn=type(data).dfn, data=data) - # are these printoptions always applicable? - with np.printoptions(precision=4, linewidth=sys.maxsize, threshold=sys.maxsize): - # TODO don't hardcode the filename, maybe a filename attribute? - with open(path, "w") as f: # type: ignore - f.writelines(iterator) diff --git a/flopy4/mf6/codec/__init__.py b/flopy4/mf6/codec/__init__.py new file mode 100644 index 00000000..cbcc97f7 --- /dev/null +++ b/flopy4/mf6/codec/__init__.py @@ -0,0 +1,71 @@ +import sys +from os import PathLike +from typing import Any + +import numpy as np +from jinja2 import Environment, PackageLoader + +from flopy4.mf6 import filters +from flopy4.mf6.codec.converter import structure_array, unstructure_array + +_JINJA_ENV = Environment( + loader=PackageLoader("flopy4.mf6"), + trim_blocks=True, + lstrip_blocks=True, +) +_JINJA_ENV.filters["blocks"] = filters.blocks +_JINJA_ENV.filters["field_type"] = filters.field_type +_JINJA_ENV.filters["field_value"] = filters.field_value +_JINJA_ENV.filters["is_list"] = filters.is_list +_JINJA_ENV.filters["array_how"] = filters.array_how +_JINJA_ENV.filters["array_chunks"] = filters.array_chunks +_JINJA_ENV.filters["array2string"] = filters.array2string + +_JINJA_TEMPLATE_NAME = "blocks.jinja" + +_PRINT_OPTIONS = { + "precision": 4, + "linewidth": sys.maxsize, + "threshold": sys.maxsize, +} + + +def unstructure(data): + # TODO unstructure arrays into sparse dicts + # TODO combine OC fields into list input as defined in the MF6 dfn + # TODO return a dictionary instead of the component itself, then + # update filters to use dictinoary access instead of getattr() + return data + + +def loads(data: str) -> Any: + # TODO + pass + + +def load(path: str | PathLike) -> Any: + # TODO + pass + + +def dumps(data) -> str: + template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME) + with np.printoptions(**_PRINT_OPTIONS): # type: ignore + return template.render(dfn=type(data).dfn, data=unstructure(data)) + + +def dump(data, path: str | PathLike) -> None: + template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME) + iterator = template.generate(dfn=type(data).dfn, data=unstructure(data)) + with np.printoptions(**_PRINT_OPTIONS), open(path, "w") as f: # type: ignore + f.writelines(iterator) + + +__all__ = [ + "structure_array", + "unstructure_array", + "loads", + "load", + "dumps", + "dump", +] diff --git a/flopy4/mf6/converters.py b/flopy4/mf6/codec/converter.py similarity index 68% rename from flopy4/mf6/converters.py rename to flopy4/mf6/codec/converter.py index bd94b6ab..dee5b089 100644 --- a/flopy4/mf6/converters.py +++ b/flopy4/mf6/codec/converter.py @@ -3,13 +3,21 @@ import numpy as np import sparse from numpy.typing import NDArray +from xarray import DataArray from xattree import get_xatspec from flopy4.mf6.config import SPARSE_THRESHOLD from flopy4.mf6.constants import FILL_DNODATA -def convert_array(value, self_, field) -> NDArray: +# TODO: convert to a cattrs structuring hook so we don't have to +# apply separately to all array fields? +def structure_array(value, self_, field) -> NDArray: + """ + Convert a sparse dictionary representation of an array to a + dense numpy array or a sparse COO array. + """ + if not isinstance(value, dict): # if not a dict, assume it's a numpy array # and let xarray deal with it if it isn't @@ -91,3 +99,31 @@ def _get_nn(cellid): # a[(nn,)] = v return final(a) + + +def unstructure_array(value: DataArray) -> dict: + """ + Convert a dense numpy array or a sparse COO array to a sparse + dictionary representation suitable for serialization into the + MF6 list-based input format. + """ + # make sure dim 'kper' is present + if "kper" not in value.dims: + raise ValueError("array must have 'kper' dimension") + + if isinstance(value.data, sparse.COO): + coords = value.coords + data = value.data + else: + coords = np.array(np.nonzero(value)).T # type: ignore + data = value[tuple(coords.T)] # type: ignore + if not coords.size: # type: ignore + return {} + match value.ndim: + case 1: + return {k: v for k, v in zip(coords[:, 0], data)} # type: ignore + case 2: + return {(k, j): v for (k, j), v in zip(coords, data)} # type: ignore + case 3: + return {(k, i, j): v for (k, i, j), v in zip(coords, data)} # type: ignore + return {} diff --git a/flopy4/mf6/filters.py b/flopy4/mf6/filters.py index d3fa1cf1..ce6822eb 100644 --- a/flopy4/mf6/filters.py +++ b/flopy4/mf6/filters.py @@ -1,3 +1,5 @@ +from collections.abc import Hashable, Mapping +from io import StringIO from typing import Any import numpy as np @@ -36,28 +38,74 @@ def field_value(ctx, field: Field): return getattr(ctx["data"], field["name"]) -def array_delay(value: xr.DataArray, chunks=None): +def array_how(value: xr.DataArray) -> str: + return "internal" + + +def array_chunks(value: xr.DataArray, chunks: Mapping[Hashable, int] | None = 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. + Yield chunks from an array of up to 3 dimensions. If the + array is not already chunked, split it into chunks of the + specified sizes, given as a dictionary mapping dimension + names to chunk sizes. + + If chunk sizes are not specified, chunk the array with at + most 2 dimensions per chunk, where: + + - If the array is 3D, assume the first dimension is the + vertical (i.e. layers) and the others horizontal (rows and + columns, in that order), and yield a chunk per layer, such + that an array with indices (k, i, j) becomes k chunks, each + of shape (i, j). + + - If the array is 1D or 2D, yield it as 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) + if chunks is None: + match value.ndim: + case 1: + # 1D array, single chunk + chunks = {value.dims[0]: value.shape[0]} + case 2: + # 2D array, single chunk + chunks = {value.dims[0]: value.shape[0], value.dims[1]: value.shape[1]} + case 3: + # 3D array, chunk for each layer + chunks = { + value.dims[0]: 1, + value.dims[1]: value.shape[1], + value.dims[2]: value.shape[2], + } + value = value.chunk(chunks) for chunk in value.data.blocks: yield chunk.compute() def array2string(value: NDArray) -> str: - """Convert an array to a string.""" - s = np.array2string(value, separator=" ") - if value.shape != (): - s = s[1:-1] # remove brackets - return s.replace("'", "").replace('"', "") # remove quotes + """ + Convert an array to a string. The array can be 1D or 2D. + If the array is 1D, it is converted to a 1-line string, + with elements separated by whitespace. If the array is + 2D, each row becomes a line in the string. + """ + buffer = StringIO() + value = np.asarray(value) + if value.ndim > 2: + raise ValueError("Only 1D and 2D arrays are supported.") + if value.ndim == 1: + # add an axis to 1d arrays so np.savetxt writes elements on 1 line + value = value[None] + format = ( + "%d" + if np.issubdtype(value.dtype, np.integer) + else "%f" + if np.issubdtype(value.dtype, np.floating) + else "%s" + ) + np.savetxt(buffer, value, fmt=format, delimiter=" ") + return buffer.getvalue().strip() -def is_dict(value: Any) -> bool: - """Check if the value is a dictionary.""" - return isinstance(value, dict) +def is_list(value: Any) -> bool: + return isinstance(value, list) diff --git a/flopy4/mf6/gwf/chd.py b/flopy4/mf6/gwf/chd.py index 316cd37d..9a8c90f4 100644 --- a/flopy4/mf6/gwf/chd.py +++ b/flopy4/mf6/gwf/chd.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field @@ -40,7 +40,7 @@ class Steps: "nnodes", ), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) aux: Optional[NDArray[np.floating]] = array( block="period", @@ -49,7 +49,7 @@ class Steps: "nnodes", ), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) boundname: Optional[NDArray[np.str_]] = array( block="period", @@ -58,12 +58,12 @@ class Steps: "nnodes", ), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) steps: Optional[NDArray[np.object_]] = array( Steps, block="period", dims=("nper", "nnodes"), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/dis.py b/flopy4/mf6/gwf/dis.py index 6ab569b4..8529de4e 100644 --- a/flopy4/mf6/gwf/dis.py +++ b/flopy4/mf6/gwf/dis.py @@ -4,7 +4,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, dim, field @@ -42,31 +42,31 @@ class Dis(Package): block="griddata", default=1.0, dims=("ncol",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) delc: NDArray[np.floating] = array( block="griddata", default=1.0, dims=("nrow",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) top: NDArray[np.floating] = array( block="griddata", default=1.0, dims=("nrow", "ncol"), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) botm: NDArray[np.floating] = array( block="griddata", default=0.0, dims=("nlay", "nrow", "ncol"), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) idomain: NDArray[np.integer] = array( block="griddata", default=1, dims=("nlay", "nrow", "ncol"), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) nnodes: int = dim( coord="node", diff --git a/flopy4/mf6/gwf/ic.py b/flopy4/mf6/gwf/ic.py index e19d3b18..96cc58b8 100644 --- a/flopy4/mf6/gwf/ic.py +++ b/flopy4/mf6/gwf/ic.py @@ -3,7 +3,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field @@ -14,7 +14,7 @@ class Ic(Package): block="packagedata", dims=("nnodes",), default=1.0, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) export_array_ascii: bool = field(block="options", default=False) export_array_netcdf: bool = field(block="options", default=False) diff --git a/flopy4/mf6/gwf/npf.py b/flopy4/mf6/gwf/npf.py index b53e1342..4d8ea596 100644 --- a/flopy4/mf6/gwf/npf.py +++ b/flopy4/mf6/gwf/npf.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field @@ -51,47 +51,47 @@ class Xt3dOptions: block="griddata", dims=("nnodes",), default=0, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k: NDArray[np.floating] = array( block="griddata", dims=("nnodes",), default=1.0, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k22: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) k33: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle1: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle2: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) angle3: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) wetdry: Optional[NDArray[np.floating]] = array( block="griddata", dims=("nnodes",), default=None, - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/gwf/oc.py b/flopy4/mf6/gwf/oc.py index a6cb40ff..7e040b33 100644 --- a/flopy4/mf6/gwf/oc.py +++ b/flopy4/mf6/gwf/oc.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray from xattree import xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field from flopy4.utils import to_path @@ -55,26 +55,26 @@ class Period: block="perioddata", default="all", dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) save_budget: Optional[NDArray[np.object_]] = array( Steps, block="perioddata", default="all", dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) print_head: Optional[NDArray[np.object_]] = array( Steps, block="perioddata", default="all", dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) print_budget: Optional[NDArray[np.object_]] = array( Steps, block="perioddata", default="all", dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) diff --git a/flopy4/mf6/interface/flopy3.py b/flopy4/mf6/interface/flopy3.py index f0be4807..86019861 100644 --- a/flopy4/mf6/interface/flopy3.py +++ b/flopy4/mf6/interface/flopy3.py @@ -31,6 +31,8 @@ def __init__( self._time = modeltime self._ims = ims self._plist = list() + self.type = "Model" + self.name = model.name # type: ignore if self._model is None: raise Exception("Model Interface needs a model") diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index 069518c4..91f3bea5 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -7,7 +7,7 @@ from numpy.typing import NDArray from xattree import ROOT, xattree -from flopy4.mf6.converters import convert_array +from flopy4.mf6.codec import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, dim, field @@ -32,19 +32,19 @@ class PeriodData: block="perioddata", default=1.0, dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) nstp: NDArray[np.integer] = array( block="perioddata", default=1, dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) tsmult: NDArray[np.floating] = array( block="perioddata", default=1.0, dims=("nper",), - converter=Converter(convert_array, takes_self=True, takes_field=True), + converter=Converter(structure_array, takes_self=True, takes_field=True), ) def to_time(self) -> ModelTime: diff --git a/flopy4/mf6/templates/macros.jinja b/flopy4/mf6/templates/macros.jinja index a34ffcd2..9b4a2b8c 100644 --- a/flopy4/mf6/templates/macros.jinja +++ b/flopy4/mf6/templates/macros.jinja @@ -17,38 +17,47 @@ {% endmacro %} {% macro keystring(field) %} -{% for item in (field|field_value).values() -%} +{% for item in field.children.values() -%} {{ field(item) }} {%- endfor %} {% endmacro %} {% macro record(field) %} -{% for item in (field|field_value).values() -%} +{% for item in field.children.values() -%} {% if item.tagged %}{{ item.name }} {% endif %}{{ field(item) }} {%- endfor %} {% endmacro %} -{% macro recarray(field, how="internal") %} -{% if how == "layered constant" %} -{{ field.name }} LAYERED -{% for val in field|field_value %} -CONSTANT -{% endfor %} -{%- elif how == "constant" %} -{{ field.name }} CONSTANT {{ field|field_value }} -{%- elif how == "layered" %} -{% if layered %} -{{ field.name }}{% for val in field|field_value %} {{ val }}{% endfor %} -{% endif %} -{%- elif how == "internal" %} -{{ field.name }} {{ internal_array(field) }} -{%- elif how == "external" %} -{{ field.name}} OPEN/CLOSE {{ field|field_value }} +{% macro recarray(field) %} +{% set value = field|field_value %} +{% if value|is_list %} +{{ list(value) }} +{% else %} +{{ array(field.name, value, how=field|array_how) }} {% endif %} {% endmacro %} -{% macro internal_array(field) %} -{% for chunk in field|field_value|array_delay -%} +{% macro array(name, value, how="internal") %} +{{ name }}{% if "layered" in how %} LAYERED{% endif %} + +{% if how == "constant" %} +CONSTANT {{ value.item() }} +{% elif how == "layered constant" %} +{% for layer in value -%} +CONSTANT {{ layer.item() }} +{%- endfor %} +{% elif how == "internal" %} +INTERNAL +{% for chunk in value|array_chunks -%} {{ chunk|array2string }} {%- endfor %} +{% elif how == "external" %} +OPEN/CLOSE {{ value }} +{% endif %} +{% endmacro %} + +{% macro list(field) %} +{% for item in field.children.values() %} +{{ field(item) }} +{% endfor %} {% endmacro %} diff --git a/test/test_codec.py b/test/test_codec.py new file mode 100644 index 00000000..73b06d7d --- /dev/null +++ b/test/test_codec.py @@ -0,0 +1,64 @@ +from flopy4.mf6.codec import dumps +import pytest + + +def test_dumps_ic(): + from flopy4.mf6.gwf import Dis, Gwf, Ic + + dis = Dis() + gwf = Gwf(dis=dis) + ic = Ic( + parent=gwf, + export_array_ascii=True, + export_array_netcdf=True, + ) + + # TODO figure out how to adapt simulation_data + # and get this working, then compare results? + # old_gwf = Flopy3Model(model=gwf) + # old_ic = ModflowGwfic( + # old_gwf, + # save_flows=True, + # save_initial_conditions=True, + # export_array_ascii=True, + # export_array_netcdf=True, + # ) + + result = dumps(ic) + print(result) + assert result + + +def test_dumps_oc(): + from flopy4.mf6.gwf import Oc + + oc = Oc( + budget_file="test.bud", + head_file="test.hds", + save_head={"*": "all"}, + save_budget={"*": "all"}, + dims={"nper": 1}, + ) + + result = dumps(oc) + print(result) + assert result + + +@pytest.mark.skip(reason="TODO 3D arrays") +def test_dumps_dis(): + from flopy4.mf6.gwf import Dis + + dis = Dis( + nlay=1, + nrow=10, + ncol=10, + delr=100.0, + delc=100.0, + idomain=1, + length_units="feet", + ) + + result = dumps(dis) + print(result) + assert result diff --git a/test/test_jinja.py b/test/test_jinja.py deleted file mode 100644 index f3bde928..00000000 --- a/test/test_jinja.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest - -from flopy4.mf6.codec import JINJA_ENV, JINJA_TEMPLATE_NAME -from flopy4.mf6.gwf import Dis, Gwf, Ic, Oc - - -@pytest.fixture -def template(): - return JINJA_ENV.get_template(JINJA_TEMPLATE_NAME) - - -def test_render_ic(template): - dis = Dis() - gwf = Gwf(dis=dis) - ic = Ic( - parent=gwf, - export_array_ascii=True, - export_array_netcdf=True, - ) - - result = template.render(dfn=Ic.dfn, data=ic) - print(result) - assert result - - -def test_render_oc(template): - oc = Oc( - budget_file="test.bud", - head_file="test.hds", - save_head={"*": "all"}, - save_budget={"*": "all"}, - dims={"nper": 1}, - ) - - result = template.render(dfn=Oc.dfn, data=oc) - print(result) - assert result