more robust prt -> mp7 adapter, deduplicate in pathline plot fns

modflowpy · Oct 6, 2023 · 5d9c497 · 5d9c497
1 parent 8d9cd0f
commit 5d9c497
Show file tree

Hide file tree

Showing 4 changed files with 131 additions and 132 deletions.
diff --git a/autotest/test_plotutil.py b/autotest/test_plotutil.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pandas as pd
-from modflow_devtools.markers import requires_exe
+import pytest
 
 from flopy.plot.plotutil import to_mp7_endpoints, to_mp7_pathlines
 
@@ -260,13 +260,26 @@
 )
 
 
-def test_to_mp7_pathlines():
-    mp7_pls = to_mp7_pathlines(pls)
+@pytest.mark.parametrize("dataframe", [True, False])
+def test_to_mp7_pathlines(dataframe):
+    inp_pls = pls if dataframe else pls.to_records(index=False)
+    mp7_pls = to_mp7_pathlines(inp_pls)
+    assert (
+        type(inp_pls)
+        == type(mp7_pls)
+        == (pd.DataFrame if dataframe else np.recarray)
+    )
     assert len(mp7_pls) == 10
-    assert set(dict(mp7_pls.dtypes).keys()) == set(mp7_pl_cols)
+    assert set(
+        dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names
+    ) == set(mp7_pl_cols)
 
 
-def test_to_mp7_endpoints():
-    mp7_eps = to_mp7_endpoints(pls)
+@pytest.mark.parametrize("dataframe", [True, False])
+def test_to_mp7_endpoints(dataframe):
+    inp_pls = pls if dataframe else pls.to_records(index=False)
+    mp7_eps = to_mp7_endpoints(inp_pls)
     assert len(mp7_eps) == 1
-    assert set(dict(mp7_eps.dtypes).keys()) == set(mp7_ep_cols)
+    assert set(
+        dict(mp7_eps.dtypes).keys() if dataframe else mp7_eps.dtype.names
+    ) == set(mp7_ep_cols)
diff --git a/flopy/plot/crosssection.py b/flopy/plot/crosssection.py
@@ -1114,41 +1114,14 @@ def plot_pathline(
                 )
             pl = [pl]
 
-        # convert dataframes to recarrays if needed
+        # convert prt to mp7 format
         pl = [
-            p.to_records(index=False) if isinstance(p, pd.DataFrame) else p
+            to_mp7_pathlines(
+                p.to_records(index=False) if isinstance(p, pd.DataFrame) else p
+            )
             for p in pl
         ]
 
-        def convert(p):
-            mp7_names = ["x", "y", "z", "time", "k", "particleid"]
-            prt_names = [
-                "x",
-                "y",
-                "z",
-                "t",
-                "trelease",
-                "imdl",
-                "iprp",
-                "irpt",
-                "ilay",
-            ]
-
-            # check format
-            if not (
-                all(n in p.dtype.names for n in mp7_names)
-                or all(n in p.dtype.names for n in prt_names)
-            ):
-                raise ValueError(
-                    "Pathline data must contain all of the following columns: "
-                    f"{mp7_names} for MODPATH 6-7, or {prt_names} for MODFLOW 6 PRT"
-                )
-
-            return to_mp7_pathlines(p) if "t" in p.dtype.names else p
-
-        # convert prt to mp7 format if needed
-        pl = [convert(p) for p in pl]
-
         # merge pathlines then split on particleid
         pls = stack_arrays(pl, asrecarray=True, usemask=False)
         pids = np.unique(pls["particleid"])

diff --git a/flopy/plot/map.py b/flopy/plot/map.py
@@ -748,41 +748,14 @@ def plot_pathline(self, pl, travel_time=None, **kwargs):
                 )
             pl = [pl]
 
-        # convert dataframes to recarrays if needed
+        # convert prt to mp7 format
         pl = [
-            p.to_records(index=False) if isinstance(p, pd.DataFrame) else p
+            to_mp7_pathlines(
+                p.to_records(index=False) if isinstance(p, pd.DataFrame) else p
+            )
             for p in pl
         ]
 
-        def convert(p):
-            mp7_names = ["x", "y", "z", "time", "k", "particleid"]
-            prt_names = [
-                "x",
-                "y",
-                "z",
-                "t",
-                "trelease",
-                "imdl",
-                "iprp",
-                "irpt",
-                "ilay",
-            ]
-
-            # check format
-            if not (
-                all(n in p.dtype.names for n in mp7_names)
-                or all(n in p.dtype.names for n in prt_names)
-            ):
-                raise ValueError(
-                    "Pathline data must contain all of the following columns: "
-                    f"{mp7_names} for MODPATH 6-7, or {prt_names} for MODFLOW 6 PRT"
-                )
-
-            return to_mp7_pathlines(p) if "t" in p.dtype.names else p
-
-        # convert prt to mp7 format if needed
-        pl = [convert(p) for p in pl]
-
         # merge pathlines then split on particleid
         pls = stack_arrays(pl, asrecarray=True, usemask=False)
         pids = np.unique(pls["particleid"])

diff --git a/flopy/plot/plotutil.py b/flopy/plot/plotutil.py
@@ -2594,7 +2594,9 @@ def parse_modpath_selection_options(
     return tep, istart, xp, yp
 
 
-def to_mp7_pathlines(data) -> pd.DataFrame:
+def to_mp7_pathlines(
+    data: Union[np.recarray, pd.DataFrame]
+) -> Union[np.recarray, pd.DataFrame]:
     """
     Convert MODFLOW 6 PRT pathline data to MODPATH 7 pathline format.
 
@@ -2605,13 +2607,46 @@ def to_mp7_pathlines(data) -> pd.DataFrame:
 
     Returns
     -------
-    pd.DataFrame
+    np.recarray or pd.DataFrame (consistent with input type)
     """
 
+    # define expected fields
+    mp7_names = ["x", "y", "z", "time", "k", "particleid"]
+    prt_names = [
+        "x",
+        "y",
+        "z",
+        "t",
+        "trelease",
+        "imdl",
+        "iprp",
+        "irpt",
+        "ilay",
+    ]
+
+    # determine return type
+    ret_type = type(data)
+
     # convert to dataframe if needed
-    if isinstance(data, np.recarray):
+    if not isinstance(data, pd.DataFrame):
         data = pd.DataFrame(data)
 
+    # check format
+    dt = data.dtypes
+    if not (
+        all(n in dt for n in mp7_names) or all(n in dt for n in prt_names)
+    ):
+        raise ValueError(
+            "Pathline data must contain all of the following columns: "
+            f"{mp7_names} for MODPATH 7, or {prt_names} for MODFLOW 6 PRT"
+        )
+
+    # return early if no need to convert
+    if "t" not in dt:
+        return (
+            data if ret_type == pd.DataFrame else data.to_records(index=False)
+        )
+
     # assign a unique particle index column incrementing an integer
     # for each unique combination of irpt, iprp, imdl, and trelease
     data = data.sort_values(["imdl", "iprp", "irpt", "trelease"])
@@ -2649,33 +2684,35 @@ def to_mp7_pathlines(data) -> pd.DataFrame:
         ]
     )
 
-    # build mp7 format dataframe
-    return pd.DataFrame(
-        np.core.records.fromarrays(
-            [
-                data[seqn_key],
-                data["iprp"],
-                data[seqn_key],
-                data["irpt"],
-                data["t"],
-                data["x"],
-                data["y"],
-                data["z"],
-                data["ilay"],
-                data["icell"],
-                # todo local coords (xloc, yloc, zloc)
-                np.zeros(data.shape[0]),
-                np.zeros(data.shape[0]),
-                np.zeros(data.shape[0]),
-                data["kper"],
-                data["kstp"],
-            ],
-            dtype=mp7_dtypes,
-        )
+    # build mp7 format recarray
+    ret = np.core.records.fromarrays(
+        [
+            data[seqn_key],
+            data["iprp"],
+            data[seqn_key],
+            data["irpt"],
+            data["t"],
+            data["x"],
+            data["y"],
+            data["z"],
+            data["ilay"],
+            data["icell"],
+            # todo local coords (xloc, yloc, zloc)
+            np.zeros(data.shape[0]),
+            np.zeros(data.shape[0]),
+            np.zeros(data.shape[0]),
+            data["kper"],
+            data["kstp"],
+        ],
+        dtype=mp7_dtypes,
     )
 
+    return pd.DataFrame(ret) if ret_type == pd.DataFrame else ret
 
-def to_mp7_endpoints(data) -> pd.DataFrame:
+
+def to_mp7_endpoints(
+    data: Union[np.recarray, pd.DataFrame]
+) -> Union[np.recarray, pd.DataFrame]:
     """
     Convert MODFLOW 6 PRT pathline data to MODPATH 7 endpoint format.
 
@@ -2686,9 +2723,12 @@ def to_mp7_endpoints(data) -> pd.DataFrame:
 
     Returns
     -------
-    pd.DataFrame
+    np.recarray or pd.DataFrame (consistent with input type)
     """
 
+    # determine return type
+    ret_type = type(data)
+
     # convert to dataframe if needed
     if isinstance(data, np.recarray):
         data = pd.DataFrame(data)
@@ -2771,39 +2811,39 @@ def to_mp7_endpoints(data) -> pd.DataFrame:
         ]
     )
 
-    # build mp7 format dataframe
-    return pd.DataFrame(
-        np.core.records.fromarrays(
-            [
-                endpts["sequencenumber"],
-                endpts["iprp"],
-                endpts["irpt"],
-                endpts["istatus"],
-                endpts["trelease"],
-                endpts["t"],
-                endpts["node0"],
-                endpts["k0"],
-                # todo initial local coords (xloc0, yloc0, zloc0)
-                np.zeros(endpts.shape[0]),
-                np.zeros(endpts.shape[0]),
-                np.zeros(endpts.shape[0]),
-                endpts["x0"],
-                endpts["y0"],
-                endpts["z0"],
-                endpts["zone0"],
-                np.zeros(endpts.shape[0]),  # todo initial cell face?
-                endpts["icell"],
-                endpts["ilay"],
-                # todo local coords (xloc, yloc, zloc)
-                np.zeros(endpts.shape[0]),
-                np.zeros(endpts.shape[0]),
-                np.zeros(endpts.shape[0]),
-                endpts["x"],
-                endpts["y"],
-                endpts["z"],
-                endpts["izone"],
-                np.zeros(endpts.shape[0]),  # todo cell face?
-            ],
-            dtype=mp7_dtype,
-        )
+    # build mp7 format recarray
+    ret = np.core.records.fromarrays(
+        [
+            endpts["sequencenumber"],
+            endpts["iprp"],
+            endpts["irpt"],
+            endpts["istatus"],
+            endpts["trelease"],
+            endpts["t"],
+            endpts["node0"],
+            endpts["k0"],
+            # todo initial local coords (xloc0, yloc0, zloc0)
+            np.zeros(endpts.shape[0]),
+            np.zeros(endpts.shape[0]),
+            np.zeros(endpts.shape[0]),
+            endpts["x0"],
+            endpts["y0"],
+            endpts["z0"],
+            endpts["zone0"],
+            np.zeros(endpts.shape[0]),  # todo initial cell face?
+            endpts["icell"],
+            endpts["ilay"],
+            # todo local coords (xloc, yloc, zloc)
+            np.zeros(endpts.shape[0]),
+            np.zeros(endpts.shape[0]),
+            np.zeros(endpts.shape[0]),
+            endpts["x"],
+            endpts["y"],
+            endpts["z"],
+            endpts["izone"],
+            np.zeros(endpts.shape[0]),  # todo cell face?
+        ],
+        dtype=mp7_dtype,
     )
+
+    return pd.DataFrame(ret) if ret_type == pd.DataFrame else ret