perf(uzf): allocate unsaturated zone work arrays once and refactor the package - #2909
Draft
jdhughes-dev wants to merge 12 commits into
Draft
perf(uzf): allocate unsaturated zone work arrays once and refactor the package#2909jdhughes-dev wants to merge 12 commits into
jdhughes-dev wants to merge 12 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors and optimizes the UZF package by moving per-cell/per-iteration temporary allocations into one-time setup allocations, and by splitting UZF implementation across submodules while also renaming memory-manager-registered variables to clearer names (affecting BMI/XMI variable names).
Changes:
- Allocate UZF unsaturated-zone work arrays once during package setup to reduce runtime overhead.
- Split UZF routines into new submodules (waves/ET and UZF package input/obs/budget) and refactor call paths accordingly.
- Rename memory-manager-visible UZF variables (wave arrays and cell properties) to clearer, modern names and document the interface impact.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| src/Model/ModelUtilities/submodules/uzf-waves.f90 | New UZF wave-routing submodule (routing, wave storage helpers, leadspeed). |
| src/Model/ModelUtilities/submodules/uzf-et.f90 | New UZF evapotranspiration submodule and ET-driven wave manipulation. |
| src/Model/GroundWaterFlow/submodules/gwf-uzf-obs.f90 | New UZF observations submodule extracted from gwf-uzf.f90. |
| src/Model/GroundWaterFlow/submodules/gwf-uzf-input.f90 | New UZF input/static property parsing submodule extracted from gwf-uzf.f90. |
| src/Model/GroundWaterFlow/submodules/gwf-uzf-budget.f90 | New UZF budget assembly submodule extracted from gwf-uzf.f90. |
| src/Model/GroundWaterFlow/gwf-uzf.f90 | Integrates submodules, removes per-call work object, renames TS shadow arrays, updates wave/ET interactions. |
| src/meson.build | Adds new UZF submodule sources to meson build. |
| msvs/mf6core.vfproj | Adds new UZF submodule sources to Visual Studio project. |
| doc/ReleaseNotes/develop.toml | Adds release notes for the performance change and the BMI/XMI variable renames. |
| autotest/TestUzfCellGroup.f90 | New unit tests for wave snapshotting, shifting, storage, WC queries, and caph. |
| autotest/tester.f90 | Registers the new UzfCellGroup unit test suite. |
| autotest/meson.build | Enables building/running the new UzfCellGroup test suite under meson. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+66
to
+72
| do j = 1, 6 | ||
| this%wave_theta(j, icell) = this%theta_res(icell) | ||
| this%wave_depth(j, icell) = DZERO | ||
| this%wave_speed(j, icell) = DZERO | ||
| this%wave_flux(j, icell) = DZERO | ||
| this%nwaves(icell) = 1 | ||
| end do |
Comment on lines
+703
to
+706
| if (abs(comp1) < DEM30) then | ||
| if (comp3 > DEM30) fhold = (comp3 * thsrinv)**eps | ||
| if (fhold < DEM30) fhold = DEM30 | ||
| leadspeed = epsfksths * (fhold**eps_m1) |
Comment on lines
+305
to
+310
| this%wave_depth(this%nwaves, icell) = DZERO | ||
| this%wave_flux(this%nwaves, icell) = & | ||
| this%vks(icell) * (((this%wave_theta(this%nwaves, icell) - & | ||
| this%theta_res(icell)) * dtheta_inv)** & | ||
| this%bc_eps(icell)) | ||
| this%wave_theta(this%nwaves, icell) = theta_surf |
Comment on lines
+175
to
+184
| if (this%wave_theta(this%nwaves(icell), icell) - thetaout > & | ||
| theta_min) then | ||
| this%wave_theta(this%nwaves(icell) + 1, icell) = & | ||
| this%wave_theta(this%nwaves(icell), icell) - thetaout | ||
| numadd = 1 | ||
| else if (this%wave_theta(this%nwaves(icell), icell) > & | ||
| theta_min) then | ||
| this%wave_theta(this%nwaves(icell) + 1, icell) = theta_min | ||
| numadd = 1 | ||
| end if |
| this%nwaves(icell) = this%nwaves(icell) + 1 | ||
| if (this%nwaves(icell) > this%nwaves_max) then | ||
| ! | ||
| ! -- too many waves error, deallocate temp arrays and return |
| if (vks <= DZERO) then | ||
| write (errmsg, '(a,1x,i0,1x,a,1x,g0,a)') & | ||
| 'VKS for uzf cell', n, & | ||
| 'must be greater than 0 (specified value ia', vks, ').' |
| this%nsets = this%parser%GetInteger() | ||
| write (this%iout, '(4x,a,i0)') 'NTRAILSETS = ', this%nsets | ||
| this%nwavesets = this%parser%GetInteger() | ||
| write (this%iout, '(4x,a,i0)') 'NTRAILSETS = ', this%nwavesets |
6 tasks
UzfCellGroupType%init took an optional memory path: with it, all 41 arrays went to the memory manager; without it, plain allocatables. The second mode existed only to build two throwaway objects that hold a snapshot of four wave arrays, so the type carried an imem_manager flag and duplicated 41-line allocate and deallocate blocks to support a case that needs four arrays. The larger cost was that one of those objects, uzfktemp in uzet, was built and torn down on every call, which is once per UZF cell per outer iteration per time step. leadwav separately allocated checktime and more on every call. That is roughly 84 heap operations per cell per outer iteration in routines whose real work is a handful of flops. Both snapshots are now UzfWaveStoreType members allocated once in init, as are the leadwav work arrays, and init has a single memory-manager path. wave_shift was overloaded to copy between objects and to shift within a cell; it is split into store_waves, load_waves, and shift_waves, which is what the call sites actually meant. nwav and ntrail hold the same value in every cell and become scalars. Wall-clock time drops 19 to 34 percent per model, 28 percent over the twelve slowest UZF autotest models, repeatable to within half a point over three runs. Memory grows by nwav*44 - 8*ncells bytes per package, about 12 to 20 KB, and does not scale with cell count. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces.
The cell group carried names inherited from MODFLOW-2005 UZF1 that say nothing about what they hold: uzthst, uzdpst, uzflst, uzspst, nwavst, thtr, thts, thti, extwc, eps, ha, hroot, rootact, etact, totflux, sinf, surflux, watab, extdpuz, ivertcon. Since the cell group is registered with the memory manager, those names are also what a BMI or XMI consumer sees. The four wave arrays and the wave count become wave_depth, wave_theta, wave_flux, wave_speed and nwaves, with WAVE_DEPTH, WAVE_THETA, WAVE_FLUX, WAVE_SPEED and NWAVES in the memory manager. Water contents become theta_res, theta_sat, theta_init and theta_ext; the remaining opaque members are renamed on the same principle. Names that already read clearly, such as celtop, celbot, landflag, uzfarea, vks, finf and pet, are left alone so the diff stays reviewable. Local variables in the wave routines get the same treatment: itester, ffcheck, feps1, feps2, more, jshort, shortest, fluxhld2 and iflx were the worst of them. A module header now records what a wave is, the Brooks-Corey flux relation, how the train is ordered and packed, and what creates and removes waves, with the Smith (1983) and Niswonger and others (2006) references. Every member carries a one-line doxygen comment, and solve() states the reset_state protocol, which is the least obvious thing about the package. findcellabove is deleted; it was declared, defined, and never called. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces.
The UZF package had no unit test coverage below the model level; the only test-drive suite was UzfEtUtil, which covers the groundwater ET functions and nothing in the cell group. The wave-train primitives are worth testing directly because they are small, pure functions of the wave state, and because the kinematic wave routines that use them are the hardest part of the package to reason about from a model-level test. Nine tests cover the store, load and shift primitives, the mobile water integral over a full and a partial wave train, the cell water content, the water content at a depth, and the Brooks-Corey capillary head. Expected values are derived from the physical definitions in the test rather than copied from the implementation, so they check behavior and not the current arithmetic. The routines under test are ones the outstanding UZF defects do not touch, so these should survive the fixes unchanged.
UzfCellGroupModule was a 2000-line module holding the type, the property setters, the kinematic wave solution and the evapotranspiration solution, with no boundary between them. It is now an interface module of 937 lines plus two submodules, following the pattern of the gwf-sfr-* and gwf-lak-* submodules: uzf-waves.f90 has the wave train and its routing, uzf-et.f90 has everything that removes water by evapotranspiration. Each moved procedure is declared once, in the parent's separate-module-procedure interface, and the body carries no duplicate dummy declarations. Wave indices are made consistent. The wave arrays were indexed by j, jk, jj, k, kk, kj, iwav and i, seven names for one concept. A plain loop or scan over waves now uses j, and a wave position that is remembered and read later gets a name that says what it is: jabove for the deepest wave above a depth, jext for the wave at the extinction depth, jwet for the next wave still above the extinction water content, jmerge for the wave being merged away. Two declaration mistakes are corrected on the way through, both harmless but both misleading: setwaves declared its icell dummy under the local marker, and factors declared its two locals under the dummy marker. caph declared its result variable on the same line as its locals. The new files are registered in src/meson.build and msvs/mf6core.vfproj; check-vfproj passes. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces, and the unit tests pass.
gwf-uzf.f90 used four conventions at once. Most routines indexed uzf cells with i and held the host groundwater node in n, which is the reverse of every other advanced package: SFR, LAK and MAW all use n for the feature number. uzf_cf, uzf_cc, uzf_setup_budobj and uzf_fill_budobj already used n for the uzf cell, so n meant a uzf cell in some routines and a groundwater node in others. uzf_rp and print_cell_properties used a third spelling, node, for the groundwater node. read_cell_properties was the worst of it: i was a uzf cell in three loops and a groundwater node in a fourth, n was a uzf cell in one loop and a groundwater node two hundred lines later, and ic was a third name for a groundwater node. Throughout the package n is now the uzf cell number and node is the groundwater node number, matching the advanced packages. Where a second uzf cell is involved it is n2, as in the budget routines, so the connected-cell terms read update_term(n, n2) for a uzf-to-uzf flow and update_term(n, node) for a uzf-to-groundwater flow; both were previously written with the same variable. Loop counters that index something other than a cell, such as observations and auxiliary variables, are left alone. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces.
…ubmodules gwf-uzf.f90 was 3050 lines holding the package type, the stress-period and solution routines, the packagedata reader and its validation, the observation machinery, and the budget object construction. Reading the solve path meant scrolling past 900 lines of input checking and budget bookkeeping that has nothing to do with it. Following the gwf-sfr-* and gwf-lak-* pattern, the reader and its checks move to gwf-uzf-input.f90, the observations to gwf-uzf-obs.f90, and the budget object to gwf-uzf-budget.f90, leaving 1987 lines in the parent. uzf_process_obsID moves whole rather than through the interface, since only uzf_df_obs refers to it and a private module procedure with no caller left in the module can be discarded by the compiler. uzf_solve now documents the two things about this package that cannot be worked out by reading it. First, that it is called three times per time step and only the uzf_bd_obs call, with reset_state false, commits the wave state; the uzf_fc and uzf_cq calls roll it back. Second, that the cell loop cascades water downward by writing the infiltration of the cell below, so it is correct only where UZF cells are numbered from the top down within a vertical stack. The new files are registered in src/meson.build and msvs/mf6core.vfproj; check-vfproj passes. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces, and every procedure in the seven UZF source files now carries a doxygen brief.
The seven _pvar arrays look like pure duplication of uzfobj arrays of the same meaning, and the comment above them said only that they have uzfobj counterparts, which invites removing them by pointing the time series manager straight at the uzfobj arrays. That would be wrong. uzf_ad walks the cells in order, and setdataet, setdataetwc and setdataetha each write the current cell and the cell below it. A cell that is the cell_below of another therefore has its uzfobj value overwritten by its neighbor before the loop reaches it. The _pvar arrays hold what the user supplied for each cell, untouched by that propagation, and are what uzf_ad reads from every time step. Without them the input for any vertically connected cell would be lost after the first stress period.
_pvar stood for package variable, which every member of the type is, and it appears nowhere else in MODFLOW 6. It was also attached to three unrelated kinds of thing: six time series input targets, one output array, and two DIMENSIONS scalars. What it was really doing is avoiding name collisions, because the cell group is initialized on the package memory path and the two share one namespace. Renaming the cell group in the previous commits removed most of those collisions, so the suffix is now mostly vestigial. The time series targets take an _input suffix, which says the thing that matters about them: they hold what the user supplied and are never written by the solution. Their doxygen names the input variable each one carries. This also corrects the memory name of extdp, which was registered as EXDP_PVAR, missing a letter. nwav_pvar is deleted. It was assigned once and handed to the cell group constructor, which keeps it as nwaves_max. gwet_pvar is deleted too; it was a copy of the cell group gwet array made in uzf_cq, and the budget, listing and observation code now reads the cell group directly. nsets becomes nwavesets, after the input keyword it holds. No input file, keyword, or definition file changes. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces.
unsat_stor and update_wav both located a wave by counting down from the top of the train and reassigning the result on every match, so the answer was the last assignment rather than an obvious one, and every wave was visited whatever the answer. Both now scan up and exit at the first match, which is the same wave. uzet evaluated theta_res(icell) + extwc1 fifteen times. That expression is the lowest water content evapotranspiration can draw a wave down to, so it is computed once into theta_min. It is formed from extwc1 rather than from theta_ext directly because extwc1 is clamped, and the two are not the same number in floating point. Neither change is measurable: a UZF-dominated benchmark moves from -53.0 to -53.5 percent at 2500 cells and -51.0 to -51.5 percent at 10000 cells against the pre-refactor baseline, which is inside the run-to-run spread. They are kept for readability, not for speed. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces.
… notes check-format fails on any fprettify output, and fprettify warns rather than reformats when a statement cannot be indented within the 82 column limit. The longer variable names introduced by this branch pushed 34 statements past that point, all of them reported but none of them fixed by fix-style. Most are split at an argument or operator boundary. The Brooks-Corey flux expressions could not be split usefully, being four levels of parentheses deep, so the saturation term they share is assigned to a named local first. That is the same arithmetic in the same order, and it reads better than the nested form it replaces. The one flux expression that is left nested is the one indexed by the whole nwaves array rather than a single element, where a scalar temporary would not compile. develop.toml gains two items: the run time reduction, and the memory manager variable names, which are visible through the BMI and XMI interfaces and so are a change users may need to act on. Output is bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces, and check-format, check-spelling, check-python-lint, check-python-format and check-vfproj all pass.
The too-many-waves error path in uzet said it deallocated temporary arrays. It has not done so since the work arrays were moved to one-time allocation earlier in this branch; the comment survived the change and misleads anyone tracing the error handling.
The rebase onto develop moved the five UZF defect fixes from MODFLOW-ORG#2911 into files this branch had already split into submodules, so git resolved the non-moving hunks and silently dropped four of the fixes with the procedures that carried them. The vector subscript in the trailing-wave routine, the uninitialized wave speed working variable, the six wave reset bound, and the three evapotranspiration capacity guards are re-applied to uzf-waves.f90 and uzf-et.f90. The early-return restore in solve and the NTRAILWAVES check survived and are unchanged. Verified against a binary built from develop at 61b19eb: all 45 UZF, UZT and UZE autotest workspaces produce byte-identical output, so this branch remains a pure refactor of the fixed code.
jdhughes-dev
force-pushed
the
uzf-refactor
branch
from
August 7, 2026 13:47
45bf05e to
5d5d329
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Unsaturated Zone Flow (UZF) Package allocated and deallocated a complete set of unsaturated-zone work arrays on every evapotranspiration calculation, and a second set on every wave-routing calculation, both of which occur once per UZF cell per outer iteration per time step. The work arrays are now allocated once when the package is set up, which reduces simulation run times by 20 to 36 percent for the UZF test problems and by about 52 percent for problems in which UZF accounts for most of the work; the reduction does not diminish as the number of UZF cells increases.
The kinematic wave variables are registered with the memory manager, and are therefore visible through the BMI and XMI interfaces, under names inherited from MODFLOW-2005 that did not describe their contents.
UZTHST,UZDPST,UZFLST,UZSPSTandNWAVSTare nowWAVE_THETA,WAVE_DEPTH,WAVE_FLUX,WAVE_SPEEDandNWAVES, and the cell properties are renamed on the same principle.UzfCellGroup.f90andgwf-uzf.f90are each split into submodules following thegwf-sfr-*pattern, and the wave and cell indices are made consistent.Results are bit-for-bit identical over all 45 UZF, UZT and UZE autotest workspaces. Input files, keywords, and definition files are unchanged.
Checklist of items for pull request
fprettifyFor additional information see instructions for contributing and instructions for developing.