fix(mfdataarray): support time array series (#2703)

Handle time-indexed and period-indexed transient array data independently. TAS data uses float keys (0.0, 6.0, 12.0) representing actual simulation times. Standard transient data uses integer keys (0, 1, 2) representing stress period indices. Only the latter were supported before.

Also allow setting array data of unknown shape (like TAS) external, this was not supported before.

Fix #2702.

Author: wpbonelli

autotest/test_mf6_tas.py

@@ -0,0 +1,139 @@
+"""
+Test TIMEARRAYSERIES (TAS) package functionality.
+
+Tests for GitHub issue #2702:
+- Problem 1: Array entries repeated when accessing TAS data via .array property
+- Problem 2: set_all_data_external() not externalizing TAS arrays
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+import flopy
+
+pytestmark = pytest.mark.mf6
+
+
+def test_tas_array_property(function_tmpdir):
+    sim_ws = Path(function_tmpdir) / "tas_array_test"
+    sim_ws.mkdir(exist_ok=True)
+
+    sim = flopy.mf6.MFSimulation(sim_name="tas_test", version="mf6", sim_ws=str(sim_ws))
+    tdis = flopy.mf6.ModflowTdis(
+        sim,
+        time_units="DAYS",
+        nper=3,
+        perioddata=[(1.0, 1, 1.0), (10.0, 10, 1.0), (10.0, 10, 1.0)],
+    )
+    model = flopy.mf6.ModflowGwf(sim, modelname="model")
+    ims = flopy.mf6.ModflowIms(sim, print_option="SUMMARY")
+    sim.register_ims_package(ims, [model.name])
+    dis = flopy.mf6.ModflowGwfdis(
+        model,
+        nlay=1,
+        nrow=10,
+        ncol=10,
+        delr=100.0,
+        delc=100.0,
+        top=0.0,
+        botm=-10.0,
+    )
+    ic = flopy.mf6.ModflowGwfic(model, strt=0.0)
+    npf = flopy.mf6.ModflowGwfnpf(model, k=10.0)
+    oc = flopy.mf6.ModflowGwfoc(model)
+    rch = flopy.mf6.ModflowGwfrcha(model, recharge="TIMEARRAYSERIES rcharray")
+
+    rch_array_0 = 0.001 * np.ones((10, 10))
+    rch_array_0[0, 0] = 0.01
+
+    rch_array_6 = 0.002 * np.ones((10, 10))
+    rch_array_6[5, 5] = 0.02
+
+    rch_array_12 = 0.003 * np.ones((10, 10))
+    rch_array_12[9, 9] = 0.03
+
+    tas = flopy.mf6.ModflowUtltas(
+        rch,
+        filename="model.rch.tas",
+        tas_array={0.0: rch_array_0, 6.0: rch_array_6, 12.0: rch_array_12},
+        time_series_namerecord="rcharray",
+        interpolation_methodrecord="stepwise",
+    )
+
+    sim.write_simulation()
+    sim2 = flopy.mf6.MFSimulation.load(sim_ws=str(sim_ws))
+    model2 = sim2.get_model("model")
+    rch2 = model2.get_package("rcha")
+    tas2 = rch2.get_package("tas")
+
+    active_keys = tas2.tas_array.get_active_key_list()
+    keys = [key for key, _ in active_keys]
+    assert len(active_keys) == 3
+    assert keys == [0.0, 6.0, 12.0]
+
+    data_0 = tas2.tas_array.get_data(0.0)
+    data_6 = tas2.tas_array.get_data(6.0)
+    data_12 = tas2.tas_array.get_data(12.0)
+
+    assert np.allclose(data_0.max(), 0.01)
+    assert np.allclose(data_6.max(), 0.02)
+    assert np.allclose(data_12.max(), 0.03)
+
+    array_data = tas2.tas_array.array
+    assert array_data.shape[0] == 3
+    assert np.allclose(array_data[0].max(), 0.01)
+    assert np.allclose(array_data[1].max(), 0.02)
+    assert np.allclose(array_data[2].max(), 0.03)
+
+
+def test_tas_set_all_data_external(function_tmpdir):
+    sim_ws = Path(function_tmpdir) / "tas_external_test"
+    sim_ws.mkdir(exist_ok=True)
+
+    sim = flopy.mf6.MFSimulation(sim_name="tas_ext", version="mf6", sim_ws=str(sim_ws))
+    tdis = flopy.mf6.ModflowTdis(
+        sim,
+        time_units="DAYS",
+        nper=2,
+        perioddata=[(10.0, 10, 1.0), (10.0, 10, 1.0)],
+    )
+    model = flopy.mf6.ModflowGwf(sim, modelname="model")
+    ims = flopy.mf6.ModflowIms(sim)
+    sim.register_ims_package(ims, [model.name])
+    dis = flopy.mf6.ModflowGwfdis(
+        model,
+        nlay=1,
+        nrow=5,
+        ncol=5,
+        delr=100.0,
+        delc=100.0,
+        top=0.0,
+        botm=-10.0,
+    )
+    ic = flopy.mf6.ModflowGwfic(model, strt=0.0)
+    npf = flopy.mf6.ModflowGwfnpf(model, k=10.0)
+    oc = flopy.mf6.ModflowGwfoc(model)
+    rch = flopy.mf6.ModflowGwfrcha(model, recharge="TIMEARRAYSERIES rcharray")
+    tas = flopy.mf6.ModflowUtltas(
+        rch,
+        filename="model.rch.tas",
+        tas_array={0.0: 0.001 * np.ones((5, 5)), 10.0: 0.002 * np.ones((5, 5))},
+        time_series_namerecord="rcharray",
+        interpolation_methodrecord="stepwise",
+    )
+
+    sim.set_all_data_external()
+    sim.write_simulation()
+
+    tas_file = sim_ws / "model.rch.tas"
+    assert tas_file.exists()
+
+    with open(tas_file, "r") as f:
+        content = f.read()
+        assert "OPEN/CLOSE" in content
+        assert "INTERNAL" not in content
+
+    external_files = list(sim_ws.glob("*.txt"))
+    assert len(external_files) >= 2

flopy/mf6/data/mfdataarray.py

@@ -616,9 +616,12 @@ def store_as_external_file(
             if isinstance(data, str) and self._tas_info(data)[0] is not None:
                 # data must not be time array series information
                 continue
+            # Check if dimensions are well-defined. For arrays with unknown
+            # shape (e.g., TAS arrays with shape="(unknown)"), allow external
+            # storage if we have actual numpy array data at runtime
             if storage.get_data_dimensions(current_layer)[0] == -9999:
-                # data must have well defined dimensions to make external
-                continue
+                if not isinstance(data, np.ndarray):
+                    continue
             try:
                 # store layer's data in external file
                 if (
@@ -1852,33 +1855,61 @@ def _get_array(self, num_sp, apply_mult, **kwargs):
         """
         data = None
         output = None
-        for sp in range(0, num_sp):
-            if sp in self._data_storage:
-                self.get_data_prep(sp)
+
+        # Check if we have float-type keys (e.g., TAS data with time values)
+        # vs. integer-type stress period keys
+        if self._data_storage:
+            has_float_keys = any(
+                isinstance(key, float)
+                for key in self._data_storage.keys()
+            )
+        else:
+            has_float_keys = False
+
+        if has_float_keys:
+            # For TAS-style data with floating-point time keys,
+            # iterate through actual keys in sorted order
+            sorted_keys = sorted(self._data_storage.keys())
+            for key in sorted_keys:
+                self.get_data_prep(key)
                 data = super().get_data(apply_mult=apply_mult, **kwargs)
                 data = np.expand_dims(data, 0)
-            else:
-                # if there is no previous data provide array of
-                # zeros, otherwise provide last array of data found
-                if data is None:
-                    # get any data
-                    data_dimensions = None
-                    for key in self._data_storage.keys():
-                        self.get_data_prep(key)
-                        data = super().get_data(
-                            apply_mult=apply_mult, **kwargs
-                        )
-                        break
-                    if data is not None:
-                        if self.structure.type == DatumType.integer:
-                            data = np.full_like(data, 1)
-                        else:
-                            data = np.full_like(data, 0.0)
-                        data = np.expand_dims(data, 0)
-            if output is None or data is None:
-                output = data
-            else:
-                output = np.concatenate((output, data))
+
+                if output is None:
+                    output = data
+                else:
+                    output = np.concatenate((output, data))
+        else:
+            # For standard stress period data with integer keys,
+            # iterate through stress periods and fill missing ones
+            for sp in range(0, num_sp):
+                if sp in self._data_storage:
+                    self.get_data_prep(sp)
+                    data = super().get_data(apply_mult=apply_mult, **kwargs)
+                    data = np.expand_dims(data, 0)
+                else:
+                    # if there is no previous data provide array of
+                    # zeros, otherwise provide last array of data found
+                    if data is None:
+                        # get any data
+                        data_dimensions = None
+                        for key in self._data_storage.keys():
+                            self.get_data_prep(key)
+                            data = super().get_data(
+                                apply_mult=apply_mult, **kwargs
+                            )
+                            break
+                        if data is not None:
+                            if self.structure.type == DatumType.integer:
+                                data = np.full_like(data, 1)
+                            else:
+                                data = np.full_like(data, 0.0)
+                            data = np.expand_dims(data, 0)
+                if output is None or data is None:
+                    output = data
+                else:
+                    output = np.concatenate((output, data))
+
         return output
 
     def has_data(self, layer=None):