From e500ac47bb8d371a5d745f20ec60ae4e3ae81d83 Mon Sep 17 00:00:00 2001 From: mjreno Date: Fri, 19 Jul 2024 15:04:45 -0400 Subject: [PATCH] add basic MFArray tests Co-authored-by: mjreno --- docs/examples/array_example.py | 12 +++---- flopy4/array.py | 61 ++++++++++++++-------------------- flopy4/block.py | 28 +++++++++++++--- flopy4/constants.py | 1 + flopy4/package.py | 5 ++- flopy4/parameter.py | 13 ++++++-- flopy4/scalar.py | 12 +++++++ test/test_array.py | 58 ++++++++++++++++++++++++++++++++ test/test_block.py | 17 +++++++--- test/test_package.py | 37 +++++++++++++-------- 10 files changed, 177 insertions(+), 67 deletions(-) create mode 100644 test/test_array.py diff --git a/docs/examples/array_example.py b/docs/examples/array_example.py index d75f2a4..afc641f 100644 --- a/docs/examples/array_example.py +++ b/docs/examples/array_example.py @@ -49,7 +49,7 @@ # Open and load a NumPy array representation fhandle = open(internal) -imfa = MFArray.load(fhandle, data_path, shape) +imfa = MFArray.load(fhandle, data_path, shape, header=False) # Get values @@ -70,7 +70,7 @@ plt.colorbar() fhandle = open(constant) -cmfa = MFArray.load(fhandle, data_path, shape) +cmfa = MFArray.load(fhandle, data_path, shape, header=False) cvals = cmfa.value plt.imshow(cvals[0:100]) plt.colorbar() @@ -93,7 +93,7 @@ # External fhandle = open(external) -emfa = MFArray.load(fhandle, data_path, shape) +emfa = MFArray.load(fhandle, data_path, shape, header=False) evals = emfa.value evals @@ -118,7 +118,7 @@ fhandle = open(ilayered) shape = (3, 1000, 100) -ilmfa = MFArray.load(fhandle, data_path, shape, layered=True) +ilmfa = MFArray.load(fhandle, data_path, shape, header=False, layered=True) vals = ilmfa.value ilmfa._value # internal storage @@ -165,7 +165,7 @@ fhandle = open(clayered) shape = (3, 1000, 100) -clmfa = MFArray.load(fhandle, data_path, shape, layered=True) +clmfa = MFArray.load(fhandle, data_path, shape, header=False, layered=True) clmfa._value @@ -218,7 +218,7 @@ fhandle = open(mlayered) shape = (3, 1000, 100) -mlmfa = MFArray.load(fhandle, data_path, shape, layered=True) +mlmfa = MFArray.load(fhandle, data_path, shape, header=False, layered=True) mlmfa.how diff --git a/flopy4/array.py b/flopy4/array.py index 5b9cf10..e07cc8a 100644 --- a/flopy4/array.py +++ b/flopy4/array.py @@ -183,8 +183,8 @@ class MFArray(MFParameter, NumPyArrayMixin): def __init__( self, - array, shape, + array=None, how=MFArrayType.internal, factor=None, block=None, @@ -328,53 +328,42 @@ def write(self, f): pass @classmethod - def load(cls, f, cwd, shape, layered=False): - """ - - Parameters - ---------- - f + def load(cls, f, cwd, shape, header=True, **kwargs): + layered = kwargs.pop("layered", False) + + if header: + tokens = multi_line_strip(f).split() + name = tokens[0] + kwargs.pop("name", None) + if len(tokens) > 1 and tokens[1] == "layered": + layered = True + else: + name = kwargs.pop("name", None) - Returns - ------- - MFArray - """ if layered: nlay = shape[0] - lay_shape = shape[1:] + lshp = shape[1:] objs = [] for _ in range(nlay): - mfa = cls._load(f, cwd, lay_shape) + mfa = cls._load(f, cwd, lshp, name) objs.append(mfa) - mfa = MFArray( - np.array(objs, dtype=object), + return MFArray( shape, + array=np.array(objs, dtype=object), how=None, factor=None, + name=name, layered=True, ) - else: - mfa = cls._load(f, cwd, shape, layered=layered) - - return mfa + kwargs.pop("layered", None) + return cls._load( + f, cwd, shape, layered=layered, name=name, **kwargs + ) @classmethod - def _load(cls, f, cwd, shape, layered=False): - """ - - Parameters - ---------- - f - cwd - shape - layered - - Returns - ------- - - """ + def _load(cls, f, cwd, shape, layered=False, **kwargs): control_line = multi_line_strip(f).split() if CommonNames.iprn.lower() in control_line: @@ -392,7 +381,7 @@ def _load(cls, f, cwd, shape, layered=False): array = float(control_line[clpos]) clpos += 1 - elif how == how.external: + elif how == MFArrayType.external: ext_path = Path(control_line[clpos]) fpath = cwd / ext_path with open(fpath) as foo: @@ -406,8 +395,7 @@ def _load(cls, f, cwd, shape, layered=False): if len(control_line) > 2: factor = float(control_line[clpos + 1]) - mfa = cls(array, shape, how, factor=factor) - return mfa + return cls(shape, array=array, how=how, factor=factor, **kwargs) @staticmethod def read_array(f): @@ -433,6 +421,7 @@ def read_array(f): CommonNames.internal in line or CommonNames.external in line or CommonNames.constant in line + or CommonNames.end in line.upper() ): f.seek(pos, 0) break diff --git a/flopy4/block.py b/flopy4/block.py index aaad808..5668fa9 100644 --- a/flopy4/block.py +++ b/flopy4/block.py @@ -3,6 +3,7 @@ from dataclasses import asdict from typing import Any +from flopy4.array import MFArray from flopy4.parameter import MFParameter, MFParameters from flopy4.utils import strip @@ -30,6 +31,21 @@ class MFBlockMappingMeta(MFBlockMeta, ABCMeta): class MFBlock(MFParameters, metaclass=MFBlockMappingMeta): + """ + MF6 input block. Maps parameter names to parameters. + + Notes + ----- + This class is dynamically subclassed by `MFPackage` + to match each block within a package parameter set. + + Supports dictionary and attribute access. The class + attributes specify the block's parameters. Instance + attributes contain both the specification and value. + + The block's name and index are discovered upon load. + """ + def __init__(self, name=None, index=None, params=None): self.name = name self.index = index @@ -49,7 +65,7 @@ def params(self) -> MFParameters: return self.data @classmethod - def load(cls, f): + def load(cls, f, **kwargs): name = None index = None found = False @@ -72,9 +88,13 @@ def load(cls, f): param = members.get(key) if param is not None: f.seek(pos) - params[key] = type(param).load( - f, **asdict(param.with_name(key).with_block(name)) - ) + spec = asdict(param.with_name(key).with_block(name)) + kwargs = {**kwargs, **spec} + if type(param) is MFArray: + # TODO: inject from model somehow? + # and remove special handling here + kwargs["cwd"] = "" + params[key] = type(param).load(f, **kwargs) return cls(name, index, params) diff --git a/flopy4/constants.py b/flopy4/constants.py index d07b9c8..f1fc31b 100644 --- a/flopy4/constants.py +++ b/flopy4/constants.py @@ -8,3 +8,4 @@ class CommonNames: vertex = "vertex" unstructured = "unstructured" empty = "" + end = "END" diff --git a/flopy4/package.py b/flopy4/package.py index 0ebce30..af73721 100644 --- a/flopy4/package.py +++ b/flopy4/package.py @@ -51,7 +51,8 @@ class MFPackageMappingMeta(MFPackageMeta, ABCMeta): class MFPackage(UserDict, metaclass=MFPackageMappingMeta): """ - MF6 package base class. + MF6 model or simulation component package. + TODO: reimplement with `ChainMap`? """ @@ -93,6 +94,8 @@ def load(cls, f): line = f.readline() if line == "": break + if line == "\n": + continue line = strip(line).lower() words = line.split() key = words[0] diff --git a/flopy4/parameter.py b/flopy4/parameter.py index 9f48960..eb7f94b 100644 --- a/flopy4/parameter.py +++ b/flopy4/parameter.py @@ -1,8 +1,9 @@ from abc import abstractmethod +from ast import literal_eval from collections import UserDict from dataclasses import dataclass, fields from enum import Enum -from typing import Any, Optional +from typing import Any, Optional, Tuple class MFReader(Enum): @@ -40,6 +41,7 @@ class MFParamSpec: repeating: bool = False tagged: bool = True reader: MFReader = MFReader.urword + shape: Optional[Tuple[int]] = None default_value: Optional[Any] = None @classmethod @@ -74,6 +76,8 @@ def load(cls, f) -> "MFParamSpec": spec[key] = val == "true" elif key == "reader": spec[key] = MFReader.from_str(val) + elif key == "shape": + spec[key] = literal_eval(val) else: spec[key] = val return cls(**spec) @@ -134,6 +138,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): super().__init__( @@ -150,6 +155,7 @@ def __init__( repeating=repeating, tagged=tagged, reader=reader, + shape=shape, default_value=default_value, ) @@ -161,7 +167,10 @@ def value(self) -> Optional[Any]: class MFParameters(UserDict): - """Mapping of parameter names to parameters.""" + """ + Mapping of parameter names to parameters. + Supports dictionary and attribute access. + """ def __init__(self, params=None): super().__init__(params) diff --git a/flopy4/scalar.py b/flopy4/scalar.py index c511703..b99d98b 100644 --- a/flopy4/scalar.py +++ b/flopy4/scalar.py @@ -26,6 +26,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): self._value = value @@ -43,6 +44,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) @@ -68,6 +70,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): super().__init__( @@ -85,6 +88,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) @@ -122,6 +126,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): super().__init__( @@ -139,6 +144,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) @@ -174,6 +180,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): super().__init__( @@ -191,6 +198,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) @@ -226,6 +234,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): super().__init__( @@ -243,6 +252,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) @@ -290,6 +300,7 @@ def __init__( repeating=False, tagged=False, reader=MFReader.urword, + shape=None, default_value=None, ): self.inout = inout @@ -308,6 +319,7 @@ def __init__( repeating, tagged, reader, + shape, default_value, ) diff --git a/test/test_array.py b/test/test_array.py new file mode 100644 index 0000000..61d4220 --- /dev/null +++ b/test/test_array.py @@ -0,0 +1,58 @@ +import numpy as np + +from flopy4.array import MFArray + + +def test_array_load_1d(tmp_path): + name = "array" + fpth = tmp_path / f"{name}.txt" + how = "INTERNAL" + v = [1.0, 2.0, 3.0] + value = " ".join(str(x) for x in v) + with open(fpth, "w") as f: + f.write(f"{name.upper()}\n{how}\n{value}\n") + + with open(fpth, "r") as f: + array = MFArray.load(f, cwd=tmp_path, shape=(3)) + assert array.name == name + assert np.allclose(array.value, np.array(v)) + + +def test_array_load_3d(tmp_path): + name = "array" + fpth = tmp_path / f"{name}.txt" + how = "INTERNAL" + v = [[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]], [[7.0, 8.0, 9.0]]] + value = "" + for a in v: + for b in a: + value += " ".join(str(x) for x in b) + value += " " + with open(fpth, "w") as f: + f.write(f"{name.upper()}\n{how}\n{value}\n") + + with open(fpth, "r") as f: + array = MFArray.load(f, cwd=tmp_path, shape=(3, 1, 3)) + assert array.name == name + assert np.allclose(array.value, np.array(v)) + + +def test_array_load_layered(tmp_path): + name = "array" + fpth = tmp_path / f"{name}.txt" + how = "INTERNAL" + v = [[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]], [[7.0, 8.0, 9.0]]] + value = "" + for a in v: + for b in a: + # TODO: MFArray expects this on separate line + value += f"{how}\n" + value += " ".join(str(x) for x in b) + value += "\n" + with open(fpth, "w") as f: + f.write(f"{name.upper()} LAYERED\n{value}") + + with open(fpth, "r") as f: + array = MFArray.load(f, cwd=tmp_path, shape=(3, 1, 3)) + assert array.name == name + assert np.allclose(array.value, np.array(v)) diff --git a/test/test_block.py b/test/test_block.py index c7b37c3..d367b42 100644 --- a/test/test_block.py +++ b/test/test_block.py @@ -1,3 +1,6 @@ +import numpy as np + +from flopy4.array import MFArray from flopy4.block import MFBlock from flopy4.scalar import MFDouble, MFFilename, MFInteger, MFKeyword, MFString @@ -10,12 +13,12 @@ class TestBlock(MFBlock): d = MFDouble(description="double") s = MFString(description="string", optional=False) f = MFFilename(description="filename", optional=False) - # a = MFArray(description="array") + a = MFArray(description="array", shape=(3)) def test_members(): params = TestBlock.params - assert len(params) == 5 + assert len(params) == 6 k = params["k"] assert isinstance(k, MFKeyword) @@ -42,9 +45,10 @@ def test_members(): assert f.description == "filename" assert not f.optional - # a = params["a"] - # assert isinstance(f, MFArray) - # assert a.description == "array" + a = params["a"] + assert isinstance(a, MFArray) + assert a.description == "array" + assert a.optional def test_load_write(tmp_path): @@ -57,6 +61,7 @@ def test_load_write(tmp_path): f.write(" D 1.0\n") f.write(" S value\n") f.write(f" F FILEIN {fpth}\n") + f.write(" A\n INTERNAL\n 1.0 2.0 3.0\n") f.write(f"END {name.upper()}\n") # test block load @@ -76,6 +81,7 @@ def test_load_write(tmp_path): assert block.d == 1.0 assert block.s == "value" assert block.f == fpth + assert np.allclose(block.a, np.array([1.0, 2.0, 3.0])) # test block write fpth2 = tmp_path / f"{name}2.txt" @@ -88,3 +94,4 @@ def test_load_write(tmp_path): assert " D 1.0\n" in lines assert " S value\n" in lines assert f" F FILEIN {fpth}\n" in lines + # assert " A\n INTERNAL\n 1.0 2.0 3.0\n" in lines diff --git a/test/test_package.py b/test/test_package.py index 529e278..6e32255 100644 --- a/test/test_package.py +++ b/test/test_package.py @@ -1,3 +1,4 @@ +from flopy4.array import MFArray from flopy4.block import MFBlock from flopy4.package import MFPackage from flopy4.scalar import MFDouble, MFFilename, MFInteger, MFKeyword, MFString @@ -20,12 +21,12 @@ class TestPackage(MFPackage): ) s = MFString(block="options", description="string", optional=False) f = MFFilename(block="options", description="filename", optional=False) - # a = MFArray(block="packagedata", description="array") + a = MFArray(block="packagedata", description="array", shape=(3)) def test_member_params(): params = TestPackage.params - assert len(params) == 5 + assert len(params) == 6 k = params["k"] assert isinstance(k, MFKeyword) @@ -57,15 +58,15 @@ def test_member_params(): assert f.description == "filename" assert not f.optional - # a = params["a"] - # assert isinstance(f, MFArray) - # assert a.block == "packagedata" - # assert a.description == "array" + a = params["a"] + assert isinstance(a, MFArray) + assert a.block == "packagedata" + assert a.description == "array" def test_member_blocks(): blocks = TestPackage.blocks - assert len(blocks) == 1 + assert len(blocks) == 2 block = blocks["options"] assert isinstance(block, MFBlock) @@ -96,14 +97,19 @@ def test_member_blocks(): assert f.description == "filename" assert not f.optional - # a = block["a"] - # assert isinstance(f, MFArray) - # assert a.description == "array" + block = blocks["packagedata"] + assert isinstance(block, MFBlock) + assert len(block.params) == 1 + + a = block["a"] + assert isinstance(a, MFArray) + assert a.description == "array" def test_load_write(tmp_path): name = "test" fpth = tmp_path / f"{name}.txt" + array = " ".join(str(x) for x in [1.0, 2.0, 3.0]) with open(fpth, "w") as f: f.write("BEGIN OPTIONS\n") f.write(" K\n") @@ -111,15 +117,19 @@ def test_load_write(tmp_path): f.write(" D 1.0\n") f.write(" S value\n") f.write(f" F FILEIN {fpth}\n") + f.write(f" A\nINTERNAL\n{array}\n") f.write("END OPTIONS\n") - # todo another block + f.write("\n") + f.write("BEGIN PACKAGEDATA\n") + f.write(f" A\nINTERNAL\n{array}\n") + f.write("END PACKAGEDATA\n") # test package load with open(fpth, "r") as f: package = TestPackage.load(f) - assert len(package.blocks) == 1 - assert len(package.params) == 5 + assert len(package.blocks) == 2 + assert len(package.params) == 6 # class attributes: param specifications assert isinstance(TestPackage.k, MFKeyword) @@ -147,4 +157,5 @@ def test_load_write(tmp_path): assert " D 1.0\n" in lines assert " S value\n" in lines assert f" F FILEIN {fpth}\n" in lines + # todo check array assert "END OPTIONS\n" in lines