From f5edfadf38d88869e2319d86a95d750c8f1e5333 Mon Sep 17 00:00:00 2001
From: wpbonelli <wbonelli@ucar.edu>
Date: Mon, 8 Apr 2024 17:45:06 -0400
Subject: [PATCH] fix(prt): fix timing-related issues (#1709)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* maximum tracking time was set incorrectly, particles were solved beyond the end of the time step for multi-step simulations — max time within the tracking loop should be end of time step unless it's the last time step in the simulation, in which case it should be particle stop time
* if, on recording a particle datum, the tracking time fell exactly on the boundary between time steps, the datum was assigned to the subsequent time step and stress period, while MODPATH 7 assigns it the previous — factor out a routine in the base MethodType to save particle data, with some extra logic for an inclusive upper bound on the time step like MP7
---
 autotest/test_prt_disv1.py                    | 42 +++++++++++-----
 src/Model/ModelUtilities/TrackData.f90        |  6 +--
 src/Model/ParticleTracking/prt.f90            | 21 ++++----
 src/Solution/ParticleTracker/Method.f90       | 49 ++++++++++++++-----
 .../ParticleTracker/MethodCellPassToBot.f90   |  7 +--
 .../ParticleTracker/MethodCellPollock.f90     |  4 +-
 .../ParticleTracker/MethodCellPollockQuad.f90 |  4 +-
 .../ParticleTracker/MethodCellTernary.f90     |  4 +-
 src/Solution/ParticleTracker/MethodDis.f90    |  5 +-
 src/Solution/ParticleTracker/MethodDisv.f90   |  4 +-
 .../ParticleTracker/MethodSubcellPollock.f90  |  9 ++--
 .../ParticleTracker/MethodSubcellTernary.f90  |  8 +--
 12 files changed, 92 insertions(+), 71 deletions(-)

diff --git a/autotest/test_prt_disv1.py b/autotest/test_prt_disv1.py
index d879ae12be0..d668d29cfbd 100644
--- a/autotest/test_prt_disv1.py
+++ b/autotest/test_prt_disv1.py
@@ -1,11 +1,26 @@
 """
 Tests particle tracking on a vertex (DISV) grid
-that reduces to a regular grid.
-
-Two cases are provided, one with valid release
-position and cell correspondences, and another
-with mismatching cell IDs; expect PRT to catch
-these and reject them.
+that reduces to a regular grid. This exercises
+PRT's ability to detect when a vertex grid can
+be solved via Pollock's method applied to quad-
+refined cells, instead of the new ternary method
+which applies more generally to polygonal cells.
+
+The simulation includes a single stress period
+with multiple time steps. This serves to test
+whether PRT properly solves trajectories over
+"internal" time steps, i.e. within the step's
+slice of simulation time, as well as extending
+tracking to termination or particle stop times
+during the simulation's final time step.
+
+Several cases are provided:
+    - default: No user-specified tracking times, MP7 in pathline mode.
+    - bprp: Mismatching cell IDs in PRP input, expect PRT to catch and reject these.
+    - trts: User-specified tracking times, some falling exactly on boundaries between
+            time steps. PRT and MP7 should both assign the datum to the prior time step.
+    - trtf: Same as trts, except tracking times are provided to PRT in a separate file,
+            rather than inline in the OC input file.
 """
 
 from pathlib import Path
@@ -52,7 +67,7 @@
 nouter, ninner = 100, 300
 hclose, rclose, relax = 1e-9, 1e-3, 0.97
 porosity = 0.1
-tracktimes = list(np.linspace(0, 11, 10))
+tracktimes = list(np.linspace(0, 19, 20))
 
 
 def tracktimes_file(path) -> Path:
@@ -225,6 +240,7 @@ def build_prt_sim(idx, gwf_ws, prt_ws, mf6):
             trackcsv_filerecord=[prt_track_csv_file],
             track_release=True,
             track_terminate=True,
+            track_transit=True,
             track_usertime=True,
             track_timesrecord=tracktimes if "trts" in name else None,
             track_timesfilerecord=(
@@ -266,6 +282,7 @@ def build_prt_sim(idx, gwf_ws, prt_ws, mf6):
 
 
 def build_mp7_sim(idx, ws, mp7, gwf):
+    name = cases[idx]
     partdata = get_partdata(gwf.modelgrid, releasepts_mp7)
     mp7_name = f"{cases[idx]}_mp7"
     pg = flopy.modpath.ParticleGroup(
@@ -287,10 +304,13 @@ def build_mp7_sim(idx, ws, mp7, gwf):
     )
     mpsim = flopy.modpath.Modpath7Sim(
         mp,
-        simulationtype="pathline",
+        simulationtype="combined"
+        if ("trts" in name or "trtf" in name)
+        else "pathline",
         trackingdirection="forward",
         budgetoutputoption="summary",
-        stoptimeoption="total",
+        stoptimeoption="extend",
+        timepointdata=[20, tracktimes],
         particlegroups=[pg],
     )
 
@@ -360,8 +380,6 @@ def check_output(idx, test):
 
     # load mf6 pathline results
     mf6_pls = pd.read_csv(prt_ws / prt_track_csv_file, na_filter=False)
-    if "trts" in name or "trtf" in name:
-        assert len(mf6_pls) == 100
 
     # make sure pathline df has "name" (boundname) column and default values
     assert "name" in mf6_pls
@@ -466,7 +484,7 @@ def check_output(idx, test):
     del mp7_pls["node"]
 
     # compare mf6 / mp7 pathline data
-    if "trts" in name or "trtf" in name:
+    if "bprp" in name:
         pass
     else:
         assert mf6_pls.shape == mp7_pls.shape
diff --git a/src/Model/ModelUtilities/TrackData.f90 b/src/Model/ModelUtilities/TrackData.f90
index b9b6006ebd6..bf2230d6b83 100644
--- a/src/Model/ModelUtilities/TrackData.f90
+++ b/src/Model/ModelUtilities/TrackData.f90
@@ -23,10 +23,8 @@ module TrackModule
 
   !> @brief Manages particle track (i.e. pathline) files.
   !!
-  !! Optionally filters events ("ireason" codes), e.g. release
-  !! or cell-to-cell transitions. Events ("ireason" codes) are:
+  !! Optionally filters events ("ireason" codes, selectable in the PRT-OC pkg):
   !!
-  !!   -1: ALL
   !!    0: RELEASE: particle is released
   !!    1: TRANSIT: particle moves from cell to cell
   !!    2: TIMESTEP: timestep ends
@@ -34,7 +32,7 @@ module TrackModule
   !!    4: WEAKSINK: particle exits a weak sink
   !!    5: USERTIME: user-specified tracking time
   !!
-  !! An arbitrary number of files can be managed. internal arrays
+  !! An arbitrary number of files can be managed. Internal arrays
   !! are resized as needed.
   !<
   type :: TrackFileControlType
diff --git a/src/Model/ParticleTracking/prt.f90 b/src/Model/ParticleTracking/prt.f90
index 94462f0b940..ac3af4f9f4c 100644
--- a/src/Model/ParticleTracking/prt.f90
+++ b/src/Model/ParticleTracking/prt.f90
@@ -928,7 +928,7 @@ end subroutine ftype_check
   !> @brief Solve the model
   subroutine prt_solve(this)
     ! -- modules
-    use TdisModule, only: kper, kstp, totimc, totim, nper, nstp
+    use TdisModule, only: kper, kstp, totimc, nper, nstp, delt
     use PrtPrpModule, only: PrtPrpType
     ! -- dummy variables
     class(PrtModelType) :: this
@@ -972,8 +972,7 @@ subroutine prt_solve(this)
 
           ! -- If particle is permanently unreleased, record its initial/terminal state
           if (particle%istatus == 8) &
-            call this%trackfilectl%save(particle, kper=kper, &
-                                        kstp=kstp, reason=3) ! reason=3: termination
+            call this%method%save(particle, reason=3) ! reason=3: termination
 
           ! -- If particle is inactive or not yet to be released, cycle
           if (particle%istatus > 1) cycle
@@ -981,14 +980,14 @@ subroutine prt_solve(this)
           ! -- If particle released this time step, record its initial state
           particle%istatus = 1
           if (particle%trelease >= totimc) &
-            call this%trackfilectl%save(particle, kper=kper, &
-                                        kstp=kstp, reason=0) ! reason=0: release
-
-          ! -- Unless in last stress period and it has only one time step,
-          ! -- limit max time to no later than end of time step
-          tmax = particle%tstop
-          if (kper == nper .and. nstp(kper) /= 1 .and. totim < particle%tstop) &
-            tmax = totim
+            call this%method%save(particle, reason=0) ! reason=0: release
+
+          ! -- Maximum time is end of time step unless this is the last
+          !    time step in the simulation, which case it's the particle
+          !    stop time.
+          tmax = totimc + delt
+          if (nper == kper .and. nstp(kper) == kstp) &
+            tmax = particle%tstop
 
           ! -- Get and apply the tracking method
           call this%method%apply(particle, tmax)
diff --git a/src/Solution/ParticleTracker/Method.f90 b/src/Solution/ParticleTracker/Method.f90
index 10ae60d5532..55ac806a35f 100644
--- a/src/Solution/ParticleTracker/Method.f90
+++ b/src/Solution/ParticleTracker/Method.f90
@@ -45,8 +45,9 @@ module MethodModule
     ! Overridden in subtypes that delegate
     procedure :: pass
     procedure :: load
-    ! Implemented in this class
+    ! Implemented here
     procedure :: init
+    procedure :: save
     procedure :: track
     procedure :: try_pass
     procedure :: update
@@ -151,6 +152,38 @@ subroutine pass(this, particle)
     call pstop(1, "pass must be overridden")
   end subroutine pass
 
+  !> @brief Save a particle's current state.
+  subroutine save(this, particle, reason)
+    use TdisModule, only: kper, kstp, totimc
+    ! dummy
+    class(MethodType), intent(inout) :: this
+    type(ParticleType), pointer, intent(inout) :: particle
+    integer(I4B), intent(in) :: reason
+    ! local
+    integer(I4B) :: per, stp
+
+    per = kper
+    stp = kstp
+
+    ! If tracking time falls exactly on a boundary between time steps,
+    ! report the previous time step for this datum. This is to follow
+    ! MP7's behavior, and because the particle will have been tracked
+    ! up to this instant under the previous time step's conditions, so
+    ! the time step we're about to start shouldn't get "credit" for it.
+    if (particle%ttrack == totimc .and. (per > 1 .or. stp > 1)) then
+      if (stp > 1) then
+        stp = stp - 1
+      else if (per > 1) then
+        per = per - 1
+        stp = 1
+      end if
+    end if
+
+    ! Save the particle's state to any registered tracking output files
+    call this%trackfilectl%save(particle, kper=per, &
+                                kstp=stp, reason=reason)
+  end subroutine save
+
   !> @brief Update particle state and check termination conditions
   !!
   !! Update the particle's properties (e.g. advancing flag, zone number,
@@ -159,8 +192,6 @@ end subroutine pass
   !! conditions apply, save particle state with the proper reason code.
   !<
   subroutine update(this, particle, cell_defn)
-    ! -- modules
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -171,28 +202,24 @@ subroutine update(this, particle, cell_defn)
       if (particle%istopzone .eq. cell_defn%izone) then
         particle%advancing = .false.
         particle%istatus = 6
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=3) ! reason=3: termination
+        call this%save(particle, reason=3) ! reason=3: termination
         return
       end if
     end if
     if (cell_defn%inoexitface .ne. 0) then
       particle%advancing = .false.
       particle%istatus = 5
-      call this%trackfilectl%save(particle, kper=kper, &
-                                  kstp=kstp, reason=3) ! reason=3: termination
+      call this%save(particle, reason=3) ! reason=3: termination
       return
     end if
     if (cell_defn%iweaksink .ne. 0) then
       if (particle%istopweaksink .ne. 0) then
         particle%advancing = .false.
         particle%istatus = 3
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=3) ! reason=3: termination
+        call this%save(particle, reason=3) ! reason=3: termination
         return
       else
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=4) ! reason=4: exited weak sink
+        call this%save(particle, reason=4) ! reason=4: exited weak sink
         return
       end if
     end if
diff --git a/src/Solution/ParticleTracker/MethodCellPassToBot.f90 b/src/Solution/ParticleTracker/MethodCellPassToBot.f90
index ce28718113a..cb41dc4ca9f 100644
--- a/src/Solution/ParticleTracker/MethodCellPassToBot.f90
+++ b/src/Solution/ParticleTracker/MethodCellPassToBot.f90
@@ -5,7 +5,7 @@ module MethodCellPassToBotModule
   use CellDefnModule, only: CellDefnType, create_defn
   use PrtFmiModule, only: PrtFmiType
   use BaseDisModule, only: DisBaseType
-  use ParticleModule
+  use ParticleModule, only: ParticleType
   use CellModule, only: CellType
   use SubcellModule, only: SubcellType
   use TrackModule, only: TrackFileControlType
@@ -43,8 +43,6 @@ end subroutine destroy
 
   !> @brief Pass particle vertically and instantaneously to the cell bottom
   subroutine apply_ptb(this, particle, tmax)
-    ! -- modules
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodCellPassToBotType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -54,8 +52,7 @@ subroutine apply_ptb(this, particle, tmax)
     if (.not. particle%advancing) return
     particle%z = this%defn%bot
     particle%iboundary(2) = this%defn%npolyverts + 2
-    call this%trackfilectl%save(particle, kper=kper, &
-                                kstp=kstp, reason=1) ! reason=1: cell transition
+    call this%save(particle, reason=1) ! reason=1: cell transition
   end subroutine apply_ptb
 
 end module MethodCellPassToBotModule
diff --git a/src/Solution/ParticleTracker/MethodCellPollock.f90 b/src/Solution/ParticleTracker/MethodCellPollock.f90
index 357fde67c7c..ad2bdbe5642 100644
--- a/src/Solution/ParticleTracker/MethodCellPollock.f90
+++ b/src/Solution/ParticleTracker/MethodCellPollock.f90
@@ -117,7 +117,6 @@ end subroutine pass_mcp
 
   !> @brief Apply Pollock's method to a rectangular cell
   subroutine apply_mcp(this, particle, tmax)
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodCellPollockType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -144,8 +143,7 @@ subroutine apply_mcp(this, particle, tmax)
       ! -- the particle state to output file(s).
       if (particle%z > cell%defn%top) then
         particle%z = cell%defn%top
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=1) ! reason=1: cell transition
+        call this%save(particle, reason=1) ! reason=1: cell transition
       end if
 
       !  Transform particle location into local cell coordinates
diff --git a/src/Solution/ParticleTracker/MethodCellPollockQuad.f90 b/src/Solution/ParticleTracker/MethodCellPollockQuad.f90
index 9335731fe9c..fd033cdcb5d 100644
--- a/src/Solution/ParticleTracker/MethodCellPollockQuad.f90
+++ b/src/Solution/ParticleTracker/MethodCellPollockQuad.f90
@@ -197,7 +197,6 @@ end subroutine pass_mcpq
 
   !> @brief Solve the quad-rectangular cell via Pollock's method
   subroutine apply_mcpq(this, particle, tmax)
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodCellPollockQuadType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -217,8 +216,7 @@ subroutine apply_mcpq(this, particle, tmax)
       ! -- the particle state to output file(s).
       if (particle%z > cell%defn%top) then
         particle%z = cell%defn%top
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=1) ! reason=1: cell transition
+        call this%save(particle, reason=1) ! reason=1: cell transition
       end if
 
       ! -- Transform particle location into local cell coordinates,
diff --git a/src/Solution/ParticleTracker/MethodCellTernary.f90 b/src/Solution/ParticleTracker/MethodCellTernary.f90
index bee0f1edac1..e7deaadb7cd 100644
--- a/src/Solution/ParticleTracker/MethodCellTernary.f90
+++ b/src/Solution/ParticleTracker/MethodCellTernary.f90
@@ -146,7 +146,6 @@ end subroutine pass_mct
   !> @brief Apply the ternary method to a polygonal cell
   subroutine apply_mct(this, particle, tmax)
     use ConstantsModule, only: DZERO, DONE, DHALF
-    use TdisModule, only: kper, kstp
     ! dummy
     class(MethodCellTernaryType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -193,8 +192,7 @@ subroutine apply_mct(this, particle, tmax)
       ! -- particle state to file
       if (particle%z > cell%defn%top) then
         particle%z = cell%defn%top
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=1) ! reason=1: cell transition
+        call this%save(particle, reason=1) ! reason=1: cell transition
       end if
 
       npolyverts = cell%defn%npolyverts
diff --git a/src/Solution/ParticleTracker/MethodDis.f90 b/src/Solution/ParticleTracker/MethodDis.f90
index dc8a35d0ffd..b80c715492e 100644
--- a/src/Solution/ParticleTracker/MethodDis.f90
+++ b/src/Solution/ParticleTracker/MethodDis.f90
@@ -135,8 +135,6 @@ end subroutine load_dis
 
   !> @brief Pass a particle to the next cell, if there is one
   subroutine pass_dis(this, particle)
-    ! -- modules
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodDisType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -174,8 +172,7 @@ subroutine pass_dis(this, particle)
           ! particle%idomain(2) = -abs(particle%idomain(2))   ! kluge???
           particle%istatus = 2 ! kluge note: use -2 to allow check for transfer to another model???
           particle%advancing = .false.
-          call this%trackfilectl%save(particle, kper=kper, &
-                                      kstp=kstp, reason=3) ! reason=3: termination
+          call this%save(particle, reason=3) ! reason=3: termination
           ! particle%iboundary(2) = -1
         else
           idiag = dis%con%ia(cell%defn%icell)
diff --git a/src/Solution/ParticleTracker/MethodDisv.f90 b/src/Solution/ParticleTracker/MethodDisv.f90
index f7da0e7fe7d..7eb45c76cfe 100644
--- a/src/Solution/ParticleTracker/MethodDisv.f90
+++ b/src/Solution/ParticleTracker/MethodDisv.f90
@@ -127,7 +127,6 @@ end subroutine load_disv
   subroutine pass_disv(this, particle)
     ! -- modules
     use DisvModule, only: DisvType
-    use TdisModule, only: kper, kstp
     ! -- dummy
     class(MethodDisvType), intent(inout) :: this
     type(ParticleType), pointer, intent(inout) :: particle
@@ -158,8 +157,7 @@ subroutine pass_disv(this, particle)
           ! particle%idomain(2) = -abs(particle%idomain(2))   ! kluge???
           particle%istatus = 2 ! kluge note, todo: use -2 to check for transfer to another model???
           particle%advancing = .false.
-          call this%trackfilectl%save(particle, kper=kper, &
-                                      kstp=kstp, reason=3) ! reason=3: termination
+          call this%save(particle, reason=3) ! reason=3: termination
           ! particle%iboundary(2) = -1
         else
           idiag = dis%con%ia(cell%defn%icell)
diff --git a/src/Solution/ParticleTracker/MethodSubcellPollock.f90 b/src/Solution/ParticleTracker/MethodSubcellPollock.f90
index f4d54e37b99..3b122739c61 100644
--- a/src/Solution/ParticleTracker/MethodSubcellPollock.f90
+++ b/src/Solution/ParticleTracker/MethodSubcellPollock.f90
@@ -86,7 +86,6 @@ end subroutine apply_msp
   subroutine track_subcell(this, subcell, particle, tmax)
     ! modules
     use ParticleModule, only: get_particle_id
-    use TdisModule, only: kper, kstp
     ! dummy
     class(MethodSubcellPollockType), intent(inout) :: this
     class(SubcellRectType), intent(in) :: subcell
@@ -186,6 +185,7 @@ subroutine track_subcell(this, subcell, particle, tmax)
     !    within the current period and time step only.
     call this%tracktimes%try_advance()
     tslice = this%tracktimes%selection
+
     if (all(tslice > 0)) then
       do i = tslice(1), tslice(2)
         t = this%tracktimes%times(i)
@@ -202,8 +202,7 @@ subroutine track_subcell(this, subcell, particle, tmax)
         particle%z = z * subcell%dz
         particle%ttrack = t
         particle%istatus = 1
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=5)
+        call this%save(particle, reason=5)
       end do
     end if
 
@@ -266,9 +265,7 @@ subroutine track_subcell(this, subcell, particle, tmax)
     particle%iboundary(3) = exitFace
 
     ! -- Save particle track record
-    if (reason > -1) &
-      call this%trackfilectl%save(particle, kper=kper, &
-                                  kstp=kstp, reason=reason)
+    if (reason >= 0) call this%save(particle, reason=reason)
 
   end subroutine track_subcell
 
diff --git a/src/Solution/ParticleTracker/MethodSubcellTernary.f90 b/src/Solution/ParticleTracker/MethodSubcellTernary.f90
index 33721911b10..edf31873aa7 100644
--- a/src/Solution/ParticleTracker/MethodSubcellTernary.f90
+++ b/src/Solution/ParticleTracker/MethodSubcellTernary.f90
@@ -64,8 +64,6 @@ end subroutine apply_mst
 
   !> @brief Track a particle across a triangular subcell using the ternary method
   subroutine track_subcell(this, subcell, particle, tmax)
-    ! modules
-    use TdisModule, only: kper, kstp
     ! dummy
     class(MethodSubcellTernaryType), intent(inout) :: this
     class(SubcellTriType), intent(in) :: subcell
@@ -283,8 +281,7 @@ subroutine track_subcell(this, subcell, particle, tmax)
         particle%z = z
         particle%ttrack = t
         particle%istatus = 1
-        call this%trackfilectl%save(particle, kper=kper, &
-                                    kstp=kstp, reason=5)
+        call this%save(particle, reason=5)
       end do
     end if
 
@@ -349,8 +346,7 @@ subroutine track_subcell(this, subcell, particle, tmax)
 
     ! -- Save particle track record
     if (reason > -1) &
-      call this%trackfilectl%save(particle, kper=kper, &
-                                  kstp=kstp, reason=reason) ! reason=2: timestep
+      call this%save(particle, reason=reason) ! reason=2: timestep
   end subroutine track_subcell
 
   !> @brief Do calculations related to analytical z solution