Merge pull request #18 from mjreno/array_write

Update array write

flopy4/array.py

@@ -203,6 +203,7 @@ def __init__(
         tagged=False,
         reader=MFReader.urword,
         default_value=None,
+        path: Optional[Path] = None,
     ):
         MFParam.__init__(
             self,
@@ -227,6 +228,7 @@ def __init__(
         self._shape = shape
         self._how = how
         self._factor = factor
+        self._path = path
 
     def __getitem__(self, item):
         return self.raw[item]
@@ -343,8 +345,77 @@ def how(self):
         return self._how
 
     def write(self, f, **kwargs):
-        # todo
-        pass
+        PAD = "  "
+        if self.layered:
+            f.write(f"{PAD}" + f"{self.name.upper()} LAYERED\n")
+            for mfa in self._value:
+                values = mfa.raw
+                if mfa._how == MFArrayType.internal:
+                    if len(values.shape) == 1:
+                        v = f"{PAD*3}"
+                        v += " ".join([str(x) for x in values])
+                        v += "\n"
+                    elif len(values.shape) == 2:
+                        v = f"\n{PAD*3}"
+                        v = v.join(" ".join(str(x) for x in y) for y in values)
+                    elif len(values.shape) == 3:
+                        v = f"{PAD*3}"
+                        for i in range(len(values)):
+                            v += " ".join(
+                                f"\n{PAD*3}" + " ".join(str(x) for x in y)
+                                for y in values[i]
+                            )
+                    lines = f"{PAD*2}" + f"{MFArrayType.to_string(mfa._how)}"
+                    if mfa._factor:
+                        lines += f" FACTOR {mfa._factor}"
+                    lines += f"\n{PAD*3}" + f"{v}\n"
+                elif mfa._how == MFArrayType.external:
+                    lines = (
+                        f"{PAD*2}" + f"{MFArrayType.to_string(mfa._how)} "
+                        f"{mfa._path}\n"
+                    )
+                elif mfa._how == MFArrayType.constant:
+                    lines = (
+                        f"{PAD*2}" + f"{MFArrayType.to_string(mfa._how)} "
+                        f"{str(mfa._value)}\n"
+                    )
+                f.write(lines)
+        else:
+            values = self.raw
+            if self._how == MFArrayType.internal:
+                if len(values.shape) == 1:
+                    v = f"\n{PAD*3}"
+                    v += " ".join([str(x) for x in values])
+                elif len(values.shape) == 2:
+                    v = f"\n{PAD*3}"
+                    v = v.join(" ".join(str(x) for x in y) for y in values)
+                elif len(values.shape) == 3:
+                    v = f"{PAD*3}"
+                    for i in range(len(values)):
+                        v += " ".join(
+                            f"\n{PAD*3}" + " ".join(str(x) for x in y)
+                            for y in values[i]
+                        )
+                lines = (
+                    f"{PAD}" + f"{self.name.upper()}\n"
+                    f"{PAD*2}" + f"{MFArrayType.to_string(self._how)}"
+                )
+                if self._factor:
+                    lines += f" FACTOR {self._factor}"
+                lines += f"{v}\n"
+            elif self._how == MFArrayType.external:
+                lines = (
+                    f"{PAD}" + f"{self.name.upper()}\n"
+                    f"{PAD*2}"
+                    + f"{MFArrayType.to_string(self._how)} {self._path}\n"
+                )
+            elif self._how == MFArrayType.constant:
+                lines = (
+                    f"{PAD}" + f"{self.name.upper()}\n"
+                    f"{PAD*2}" + f"{MFArrayType.to_string(self._how)} "
+                    f"{str(self._value)}\n"
+                )
+            f.write(lines)
 
     @classmethod
     def load(cls, f, cwd, shape, header=True, **kwargs):
@@ -391,6 +462,7 @@ def _load(cls, f, cwd, shape, layered=False, **kwargs):
             control_line.pop(idx)
 
         how = MFArrayType.from_string(control_line[0])
+        extpath = None
         clpos = 1
 
         if how == MFArrayType.internal:
@@ -401,8 +473,8 @@ def _load(cls, f, cwd, shape, layered=False, **kwargs):
             clpos += 1
 
         elif how == MFArrayType.external:
-            ext_path = Path(control_line[clpos])
-            fpath = cwd / ext_path
+            extpath = Path(control_line[clpos])
+            fpath = cwd / extpath
             with open(fpath) as foo:
                 array = cls.read_array(foo)
             clpos += 1
@@ -414,7 +486,14 @@ def _load(cls, f, cwd, shape, layered=False, **kwargs):
         if len(control_line) > 2:
             factor = float(control_line[clpos + 1])
 
-        return cls(shape, array=array, how=how, factor=factor, **kwargs)
+        return cls(
+            shape,
+            array=array,
+            how=how,
+            factor=factor,
+            path=extpath,
+            **kwargs,
+        )
 
     @staticmethod
     def read_array(f):

spec/make_ipkg.py

@@ -35,7 +35,7 @@ def __init__(
         with open(f"{PROJ_ROOT}/spec/mf6pkg.template") as f:
             tout = Template(f.read())
 
-        if not (tin or tout):
+        if not (tin and tout):
             raise ValueError("FileSystem NO-OPT")
 
         jinja_blocks = []

spec/make_toml.py

@@ -176,7 +176,6 @@ def _set_var_d(self):
         self._var_d = vardict
 
     def _substitute(self, blockname, component, subcomponent):
-        block_d = dict()
         block_d = {}
         for k in self._var_d:
             varname, block = k

test/test_array.py

@@ -9,9 +9,9 @@ def test_array_load_1d(tmp_path):
     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
@@ -28,9 +28,31 @@ def test_array_load_3d(tmp_path):
         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_3d_external(tmp_path):
+    name = "array"
+    fpth = tmp_path / f"{name}.txt"
+    extfpth = f"external_{name}.txt"
+    how = "OPEN/CLOSE"
+    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(tmp_path / extfpth, "w") as f:
+        f.write(f"{value}")
+    with open(fpth, "w") as f:
+        f.write(f"{name.upper()}\n{how} {extfpth}\n")
     with open(fpth, "r") as f:
         array = MFArray.load(f, cwd=tmp_path, shape=(3, 1, 3))
         assert array.name == name
@@ -45,13 +67,12 @@ def test_array_load_layered(tmp_path):
     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 += "  " + 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

test/test_block.py

@@ -122,7 +122,9 @@ 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
+        assert "  A\n" in lines
+        assert "    INTERNAL\n" in lines
+        assert "      1.0 2.0 3.0\n" in lines
         assert "  R  RK  RI 2  RD 2.0\n" in lines
         assert "END OPTIONS\n" in lines
 

test/test_package.py

@@ -277,7 +277,9 @@ 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 "  A\n" in lines
+        assert "    INTERNAL\n" in lines
+        assert f"      {array}\n" in lines
         assert "END OPTIONS\n" in lines