Skip to content

Fix corner index ordering - #2801

Draft
willirath wants to merge 5 commits into
Parcels-code:mainfrom
willirath:fix-agrid-corner-index-ordering
Draft

Fix corner index ordering#2801
willirath wants to merge 5 commits into
Parcels-code:mainfrom
willirath:fix-agrid-corner-index-ordering

Conversation

@willirath

@willirath willirath commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Description

Fix corner index ordering in the A-grid and C-grid interpolators: the flattened gather currently in main silently produced wrong corners when particless time or depth indices were not aligned, so a particle could pick up a batch-mate's ti/zi.
Both interpolators now go through a single and tested _gather_corners helper that derives index arrays and output shape from the same layout.

Adds a regression test that a trajectory does not depend on its batch-mate release times, plus interpolation tests for mixed per-particle indices.

Also updates the expected z values in test_nemo_3D_curvilinear_fieldset[RK4_3D], which had recorded the buggy mixing.

Reproducer with U uniform in space, varying in time, so each particle's displacement depends only on its own clock:

# fieldset: flat A-grid, U = cos(2*pi*t/86400) m/s (uniform in space), V = 0, 3-hourly
t0 = np.timedelta64(0, "s")

def run(release_times):
    npart = len(release_times)
    pset = ParticleSet(
        fieldset, pclass=Particle,
        t=np.array(release_times),
        z=np.zeros(npart), y=np.zeros(npart), x=np.zeros(npart),
    )
    pset.execute(AdvectionRK4, dt=np.timedelta64(1, "h"), endtime=np.timedelta64(48, "h"))
    return pset.x[0]

alone     = run([t0])
staggered = run([t0] + [t0 + np.timedelta64(3, "h")] * 3)

assert staggered == pytest.approx(alone)  # fails before the fix

Particle 0 should be identical in both runs as only its batch-mates differ. Before the fix, staggered moves particle 0 onto the other particle time index and the assert fails. The depth analogue is the same shape: two particles at different z in one set get each other's zi.

Checklist

  • Tests added: Integration test and test for the new corner gather logic.
  • This PR targets the correct branch: Only relevant for v4. So targeted at main.

AI Disclosure

  • This PR contains AI-generated content.
    • I have tested any AI-generated content in my PR.
    • I take responsibility for any AI-generated content in my PR.
    • Describe how you used it (e.g., by pasting your prompt): Bug was found by Claude Opus when developing a downstream package. Diagnostics and implementation of the fix were done with Claude Opus. I guided the session and reviewed the code.

Draft PR because it's late and I want to have a fresh pair of eyes (my own or someone else's) tomorrow.

willirath and others added 2 commits August 4, 2026 23:36
Fails on main. Found while writing lcs_parcels.

Co-authored-by: Claude <noreply@anthropic.com>
The index arrays were built particle-major for T/Z but corner-major for
Y/X, while the reshape is corner-major. The orderings agree only when all
particles share a time and depth index, so batches with a non-uniform
clock or spread over depth levels were gathered from the wrong level.

Derive all four index arrays from one layout in _gather_corners, and route
the A-grid, C-grid velocity and C-grid tracer gathers through it. Also
fixes CGrid_Tracer's np.repeat(ti) missing its repeats argument.

test_nemo_3D_curvilinear_fieldset expectations recorded the old batch
behaviour; they now match what each particle gets advected on its own.

Co-authored-by: Claude <noreply@anthropic.com>
willirath and others added 2 commits August 5, 2026 00:47
The per-axis index arrays were laid out with manual stride arithmetic:
a running `before`, a `math.prod` of the trailing counts, and a comment
spelling out the flat position of each level.

Build them from the output shape instead. Each axis's levels are placed
in their slot of the (T, Z, Y, X, particle) grid, broadcast over the
other slots and flattened, which is exactly the inverse of the reshape
that ends the function. Same index arrays, no loop-carried state.

Co-authored-by: Claude <noreply@anthropic.com>
Drop the bold emphasis, the double-hyphen asides and the semicolons from
the docstrings and comments added with the corner gather rework, and give
the numpydoc parameter entries the `name : type` form used elsewhere.

Co-authored-by: Claude <noreply@anthropic.com>
@erikvansebille

Copy link
Copy Markdown
Member

Great that you spotted this bug @willirath, and thanks for the fix! I had a quick skim and I think your problem diagnosis and solution are sound. Let me know when you want me to do a full/thorough review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

2 participants