keep supporting 3-tuples with dummy values for backcompat

Author: wpbonelli

autotest/test_cellbudgetfile.py

@@ -1010,3 +1010,88 @@ def test_cellbudgetfile_get_ts_aux_vars_mf6_disu(dis_sim):
     chd = cbc.get_data(text="CHD")
     for field in ["node", "node2", "q"]:
         assert field in chd[0].dtype.names
+
+
+@pytest.mark.requires_exe("mf6")
+def test_cellbudgetfile_get_ts_backwards_compatible_idx_format(
+    dis_sim, function_tmpdir
+):
+    from flopy.mf6 import ModflowGwfchd, ModflowGwfdisu, ModflowGwfdisv
+
+    sim = dis_sim
+    gwf = sim.get_model()
+    dis_grid = gwf.modelgrid
+
+    # DISV
+    disv_kwargs = get_disv_kwargs(
+        nlay=dis_grid.nlay,
+        nrow=dis_grid.nrow,
+        ncol=dis_grid.ncol,
+        delr=dis_grid.delr,
+        delc=dis_grid.delc,
+        tp=dis_grid.top,
+        botm=dis_grid.botm,
+    )
+    gwf.remove_package("dis")
+    gwf.remove_package("chd")
+    disv = ModflowGwfdisv(gwf, **disv_kwargs)
+    chd_spd = [[0, 0, 10.0], [0, 24, 0.0]]
+    chd = ModflowGwfchd(gwf, stress_period_data=chd_spd)
+
+    sim.set_sim_path(function_tmpdir / "disv_backcompat")
+    sim.write_simulation()
+    success, _ = sim.run_simulation(silent=False)
+    assert success
+
+    cbc = gwf.output.budget()
+
+    ts_new = cbc.get_ts(idx=(0, 4), text="DATA-SPDIS", variable="qx")
+    ts_old = cbc.get_ts(idx=(0, 0, 4), text="DATA-SPDIS", variable="qx")
+
+    np.testing.assert_array_equal(
+        ts_new,
+        ts_old,
+    )
+
+    # DISU
+    disu_kwargs = get_disu_kwargs(
+        nlay=dis_grid.nlay,
+        nrow=dis_grid.nrow,
+        ncol=dis_grid.ncol,
+        delr=dis_grid.delr,
+        delc=dis_grid.delc,
+        tp=dis_grid.top,
+        botm=dis_grid.botm,
+        return_vertices=True,
+    )
+    gwf.remove_package("disv")
+    gwf.remove_package("chd")
+    disu = ModflowGwfdisu(gwf, **disu_kwargs)
+    chd_spd = [[0, 10.0], [24, 0.0]]
+    chd = ModflowGwfchd(gwf, stress_period_data=chd_spd)
+
+    sim.set_sim_path(function_tmpdir / "disu_backcompat")
+    sim.write_simulation()
+    success, _ = sim.run_simulation(silent=False)
+    assert success
+
+    cbc = gwf.output.budget()
+
+    ts_new = cbc.get_ts(idx=4, text="DATA-SPDIS", variable="qx")
+    ts_old = cbc.get_ts(idx=(0, 0, 4), text="DATA-SPDIS", variable="qx")
+
+    np.testing.assert_array_equal(
+        ts_new,
+        ts_old,
+    )
+
+    # test lists too
+    ts_new_list = cbc.get_ts(idx=[4, 10], text="DATA-SPDIS", variable="qx")
+    ts_old_list = cbc.get_ts(
+        idx=[(0, 0, 4), (0, 0, 10)], text="DATA-SPDIS", variable="qx"
+    )
+
+    np.testing.assert_array_equal(
+        ts_new_list,
+        ts_old_list,
+    )

flopy/utils/binaryfile/__init__.py

@@ -1615,9 +1615,13 @@ def get_ts(self, idx, text=None, times=None, variable="q"):
 
             - Structured grids (DIS): (layer, row, column) or list of such
             - Vertex grids (DISV): (layer, cellid) or list of such
-            - Unstructured grids (DISU): node number, (node number,) or list of such
+            - Unstructured grids (DISU): node number or list of such
 
             All indices must be zero-based.
+
+            For backwards compatibility, DISV and DISU grids also accept the old
+            3-tuple format with dummy values: (layer, dummy, cellid) for DISV and
+            (dummy, dummy, node) for DISU.
         text : str
             The text identifier for the record.  Examples include
             'RIVER LEAKAGE', 'STORAGE', 'FLOW RIGHT FACE', etc.
@@ -1760,11 +1764,23 @@ def _cellids(self, idx):
                     raise ValueError(f"Column index {j} out of range [0, {self.ncol})")
                 cellid.append((k, i, j))
             elif grid_type == "vertex":
-                if not isinstance(item, (list, tuple)) or len(item) != 2:
+                if isinstance(item, (list, tuple)):
+                    if len(item) == 2:
+                        # proper format: (layer, cellid)
+                        k, cell = item
+                    elif len(item) == 3:
+                        # old format: (layer, dummy, cellid)
+                        k, cell = item[0], item[2]
+                    else:
+                        raise ValueError(
+                            f"Expected DISV cell index (layer, cellid) "
+                            f"or (layer, dummy, cellid), got: {item}"
+                        )
+                else:
                     raise ValueError(
-                        f"Expected DISV cell index (layer, cellid), got: {item}"
+                        f"Expected DISV cell index (layer, cellid) or "
+                        f"(layer, dummy, cellid), got: {item}"
                     )
-                k, cell = item
                 if k < 0 or k >= self.nlay:
                     raise Exception(f"Layer index {k} out of range [0, {self.nlay})")
                 if cell < 0 or cell >= self.modelgrid.ncpl:
@@ -1773,9 +1789,26 @@ def _cellids(self, idx):
                     )
                 cellid.append((k, cell))
             else:
-                if not isinstance(item, (int, np.integer)):
-                    raise Exception(f"Expected DISU node number , got: {item}")
-                node = int(item)
+                if isinstance(item, (int, np.integer)):
+                    # proper format: just node number
+                    node = int(item)
+                elif isinstance(item, (list, tuple)):
+                    if len(item) == 3:
+                        # old format: (dummy, dummy, node)
+                        node = int(item[2])
+                    elif len(item) == 1:
+                        # Also support single-element tuple
+                        node = int(item[0])
+                    else:
+                        raise ValueError(
+                            f"Expected DISU node number or (dummy, dummy, node), "
+                            f"got: {item}"
+                        )
+                else:
+                    raise ValueError(
+                        f"Expected DISU node number or (dummy, dummy, node), "
+                        f"got: {item}"
+                    )
                 if node < 0 or node >= self.modelgrid.nnodes:
                     raise Exception(
                         f"Node index {node} out of range [0, {self.modelgrid.nnodes})"