more test cleanup

autotest/prt/prt_test_utils.py

@@ -1,20 +1,23 @@
 import os
 from types import SimpleNamespace
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import flopy
 import numpy as np
-import pandas as pd
 
 
-def all_equal(col, val):
-    a = col.to_numpy()
+def all_equal(series, val):
+    a = series.to_numpy()
     return a[0] == val and (a[0] == a).all()
 
 
 def get_gwf_sim(
     name, ws, mf6
 ) -> Tuple[flopy.mf6.MFSimulation, SimpleNamespace]:
+    """
+    Simple GWF simulation for use/modification by PRT tests
+    """
+
     # test case context
     ctx = SimpleNamespace(
         nlay=1,
@@ -30,15 +33,12 @@ def get_gwf_sim(
         # mp7 release points (cell-local coordinates)
         releasepts_mp7=[
             # node number, localx, localy, localz
-            # (0-based indexing converted to 1-based for mp7 by flopy)
             (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5)
             for i in range(9)
         ],
-        # expected release points in PRT format; use flopy
-        # to convert mp7 to prt format then check equality
+        # PRT release points (cell indices and global coordinates both required)
         releasepts_prt=[
             # particle index, k, i, j, x, y, z
-            # (0-based indexing converted to 1-based for mf6 by flopy)
             [i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5]
             for i in range(9)
         ],
@@ -119,7 +119,7 @@ def check_track_data(
     track_hdr: os.PathLike,
     track_csv: os.PathLike,
 ):
-    """Check that track data written to binary, CSV, and budget files are equal."""
+    """Check that binary and CSV track files are equal."""
 
     # get dtype from ascii header file
     dt = get_track_dtype(track_hdr)
@@ -135,8 +135,7 @@ def check_track_data(
         data_bin.shape == data_csv.shape
     ), f"Binary and CSV track data shapes do not match: {data_bin.shape} != {data_csv.shape}"
 
-    # check particle tracks written to all output files are equal
-    # check each column separately to avoid:
+    # check each column separately to avoid
     # TypeError: The DType <class 'numpy._FloatAbstractDType'> could not be promoted by <class 'numpy.dtype[void]'>
     for k in data_bin.dtype.names:
         if k == "name":
@@ -184,14 +183,20 @@ def get_model_name(name, mdl):
 
 
 def get_track_dtype(path: os.PathLike):
-    """Get the dtype of the track data recarray from the ascii header file."""
+    """Read a numpy dtype describing particle track
+    data format from the ascii track header file."""
 
     hdr_lns = open(path).readlines()
     hdr_lns_spl = [[ll.strip() for ll in l.split(",")] for l in hdr_lns]
     return np.dtype(list(zip(hdr_lns_spl[0], hdr_lns_spl[1])))
 
 
 def get_ireason_code(output_event):
+    """
+    Map output event to PRT ireason code specifing
+    the reason a particle track datum was recorded.
+    """
+
     return (
         0
         if output_event == "RELEASE"
@@ -208,6 +213,10 @@ def get_ireason_code(output_event):
 
 
 def get_partdata(grid, rpts):
+    """
+    Make a flopy.modpath.ParticleData from the given grid and release points.
+    """
+
     if grid.grid_type == "structured":
         return flopy.modpath.ParticleData(
             partlocs=[grid.get_lrc(p[0])[0] for p in rpts],
@@ -230,25 +239,6 @@ def get_partdata(grid, rpts):
         )
 
 
-def get_perioddata(name, periods=1, fraction=None) -> Optional[dict]:
-    if "reft" in name:
-        return None
-    opt = [
-        "FIRST"
-        if "frst" in name
-        else "ALL"
-        if "all" in name
-        else ("STEPS", 1)
-        if "stps" in name
-        else None
-    ]
-    if opt[0] is None:
-        raise ValueError(f"Invalid period option: {name}")
-    if fraction is not None:
-        opt.append(("FRACTION", fraction))
-    return {i: opt for i in range(periods)}
-
-
 def has_default_boundnames(data):
     name = [int(n.partition("0")[2]) for n in data["name"].to_numpy()]
     irpt = data["irpt"].to_numpy()

autotest/prt/test_prt_exg01.py

@@ -25,32 +25,14 @@
 from flopy.plot.plotutil import to_mp7_pathlines
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
-from prt_test_utils import (all_equal, check_budget_data, check_track_data,
-                            get_gwf_sim, has_default_boundnames)
+from prt_test_utils import check_budget_data, check_track_data, get_gwf_sim
 
 from framework import TestFramework
 from simulation import TestSimulation
 
-# simulation/model names
 simname = "prtexg01"
-
-# test cases
 ex = [simname, f"{simname}bnms"]
 
-# release points
-releasepts = [
-    # index, k, i, j, x, y, z
-    # (0-based indexing converted to 1-based for mf6 by flopy)
-    [i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5]
-    for i in range(9)
-]
-releasepts_mp7 = [
-    # node number, localx, localy, localz
-    # (0-based indexing converted to 1-based for mp7 by flopy)
-    (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5)
-    for i in range(9)
-]
-
 
 # model names
 def get_model_name(idx, mdl):
@@ -80,9 +62,9 @@ def build_sim(idx, ws, mf6):
 
     # create prp package
     rpts = (
-        [r + [str(r[0] + 1)] for r in releasepts]
+        [r + [str(r[0] + 1)] for r in ctx.releasepts_prt]
         if "bnms" in name
-        else releasepts
+        else ctx.releasepts_prt
     )
     flopy.mf6.ModflowPrtprp(
         prt,
@@ -137,10 +119,10 @@ def build_sim(idx, ws, mf6):
 
 def build_mp7_sim(ctx, idx, ws, mp7, gwf):
     partdata = flopy.modpath.ParticleData(
-        partlocs=[p[0] for p in releasepts_mp7],
-        localx=[p[1] for p in releasepts_mp7],
-        localy=[p[2] for p in releasepts_mp7],
-        localz=[p[3] for p in releasepts_mp7],
+        partlocs=[p[0] for p in ctx.releasepts_mp7],
+        localx=[p[1] for p in ctx.releasepts_mp7],
+        localy=[p[2] for p in ctx.releasepts_mp7],
+        localz=[p[3] for p in ctx.releasepts_mp7],
         timeoffset=0,
         drape=0,
     )

autotest/prt/test_prt_fmi01.py

@@ -36,14 +36,17 @@
 from flopy.plot.plotutil import to_mp7_pathlines
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
-from prt_test_utils import (all_equal, check_budget_data, check_track_data,
-                            get_gwf_sim, get_model_name, get_partdata,
-                            has_default_boundnames)
+from prt_test_utils import (
+    all_equal,
+    check_budget_data,
+    check_track_data,
+    get_gwf_sim,
+    get_model_name,
+    get_partdata,
+    has_default_boundnames,
+)
 
-# simulation name
 simname = "prtfmi01"
-
-# test cases
 ex = [simname, f"{simname}saws"]
 
 

autotest/prt/test_prt_fmi02.py

@@ -36,46 +36,41 @@
 from flopy.plot.plotutil import to_mp7_pathlines
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
-from prt_test_utils import (check_budget_data, check_track_data, get_gwf_sim,
-                            get_model_name, get_partdata)
+from prt_test_utils import (
+    check_budget_data,
+    check_track_data,
+    get_gwf_sim,
+    get_model_name,
+)
 
-# simulation/model names
 simname = "prtfmi02"
-
-# test cases
 ex = [
     f"{simname}all",
     f"{simname}rel",
     f"{simname}trst",
     f"{simname}tstp",
     f"{simname}wksk",
 ]
-
-# release points
 releasepts_prt = {
     "a": [
         # index, k, i, j, x, y, z
-        # (0-based indexing converted to 1-based for mf6 by flopy)
         [i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5]
         for i in range(4)
     ],
     "b": [
         # index, k, i, j, x, y, z
-        # (0-based indexing converted to 1-based for mf6 by flopy)
         [i, 0, 0, 0, float(f"0.{i + 5}"), float(f"9.{i + 5}"), 0.5]
         for i in range(5)
     ],
 }
 releasepts_mp7 = {
     "a": [
         # node number, localx, localy, localz
-        # (0-based indexing converted to 1-based for mf6 by flopy)
         (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5)
         for i in range(4)
     ],
     "b": [
         # node number, localx, localy, localz
-        # (0-based indexing converted to 1-based for mf6 by flopy)
         (0, float(f"0.{i + 5}"), float(f"0.{i + 5}"), 0.5)
         for i in range(5)
     ],

autotest/prt/test_prt_fmi03.py

@@ -37,16 +37,15 @@
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
 from matplotlib.collections import LineCollection
-from prt_test_utils import (check_budget_data, check_track_data, get_gwf_sim,
-                            get_model_name)
+from prt_test_utils import (
+    check_budget_data,
+    check_track_data,
+    get_gwf_sim,
+    get_model_name,
+)
 
-# simulation name
 simname = "prtfmi03"
-
-# test cases
 ex = [f"{simname}_l1", f"{simname}_l2"]
-
-# izone
 stopzone_cells = [(0, 1, 8), (0, 8, 1)]
 
 
@@ -57,22 +56,6 @@ def create_izone(nlay, nrow, ncol):
     return izone
 
 
-# release points
-# todo: define for mp7 first, then use flopy utils to convert to global coords for mf6 prt
-releasepts_mp7 = [
-    # node number, localx, localy, localz
-    # (0-based indexing converted to 1-based for mp7 by flopy)
-    (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5)
-    for i in range(9)
-]
-releasepts_prt = [
-    # particle index, k, i, j, x, y, z
-    # (0-based indexing converted to 1-based for mf6 by flopy)
-    (i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5)
-    for i in range(9)
-]
-
-
 def build_prt_sim(ctx, name, ws, mf6):
     # create simulation
     sim = flopy.mf6.MFSimulation(
@@ -122,8 +105,8 @@ def build_prt_sim(ctx, name, ws, mf6):
         prt,
         pname="prp1",
         filename=f"{prtname}_1.prp",
-        nreleasepts=len(releasepts_prt),
-        packagedata=releasepts_prt,
+        nreleasepts=len(ctx.releasepts_prt),
+        packagedata=ctx.releasepts_prt,
         perioddata={0: ["FIRST"]},
         istopzone=1,
     )
@@ -163,10 +146,10 @@ def build_prt_sim(ctx, name, ws, mf6):
 
 def build_mp7_sim(ctx, name, ws, mp7, gwf):
     partdata = flopy.modpath.ParticleData(
-        partlocs=[p[0] for p in releasepts_mp7],
-        localx=[p[1] for p in releasepts_mp7],
-        localy=[p[2] for p in releasepts_mp7],
-        localz=[p[3] for p in releasepts_mp7],
+        partlocs=[p[0] for p in ctx.releasepts_mp7],
+        localx=[p[1] for p in ctx.releasepts_mp7],
+        localy=[p[2] for p in ctx.releasepts_mp7],
+        localz=[p[3] for p in ctx.releasepts_mp7],
         timeoffset=0,
         drape=0,
     )

autotest/prt/test_prt_fmi04.py

@@ -35,29 +35,17 @@
 from flopy.plot.plotutil import to_mp7_pathlines
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
-from prt_test_utils import (check_budget_data, check_track_data, get_gwf_sim,
-                            get_ireason_code, get_model_name)
+from prt_test_utils import (
+    check_budget_data,
+    check_track_data,
+    get_gwf_sim,
+    get_ireason_code,
+    get_model_name,
+)
 
-# simulation name
 simname = "prtfmi04"
-
-# test cases
 ex = [simname, f"{simname}saws"]
 
-# release points
-releasepts_mp7 = [
-    # node number, localx, localy, localz
-    # (0-based indexing converted to 1-based for mp7 by flopy)
-    (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5)
-    for i in range(9)
-]
-releasepts_prt = [
-    # particle index, k, i, j, x, y, z
-    # (0-based indexing converted to 1-based for mf6 by flopy)
-    (i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5)
-    for i in range(9)
-]
-
 
 def build_prt_sim(ctx, name, ws, mf6):
     # output file names
@@ -106,8 +94,8 @@ def build_prt_sim(ctx, name, ws, mf6):
         prt,
         pname="prp1",
         filename=f"{prtname}_1.prp",
-        nreleasepts=len(releasepts_prt),
-        packagedata=releasepts_prt,
+        nreleasepts=len(ctx.releasepts_prt),
+        packagedata=ctx.releasepts_prt,
         perioddata={0: ["FIRST"]},
         stop_at_weak_sink="saws" in name,
     )
@@ -145,10 +133,10 @@ def build_mp7_sim(ctx, name, ws, mp7, gwf):
     mp7_pathline_file = f"{mp7name}.mppth"
 
     partdata = flopy.modpath.ParticleData(
-        partlocs=[p[0] for p in releasepts_mp7],
-        localx=[p[1] for p in releasepts_mp7],
-        localy=[p[2] for p in releasepts_mp7],
-        localz=[p[3] for p in releasepts_mp7],
+        partlocs=[p[0] for p in ctx.releasepts_mp7],
+        localx=[p[1] for p in ctx.releasepts_mp7],
+        localy=[p[2] for p in ctx.releasepts_mp7],
+        localz=[p[3] for p in ctx.releasepts_mp7],
         timeoffset=0,
         drape=0,
     )