[erikvansebille]

Description

This PR adds tests for XGrid and UxGrid to test that the particle element index (particles.ei) are updated correctly when a particle is moved.

It also fixes a small bug in ParticleSet.populate_indices(), which was not compatible with UxGrid

Checklist

• Closes Grid search (ei/igrid) not caching between successive grid search calls #2693
• Closes Check that there are unit tests that check for correct particles._ei setting and updating #2669
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

AI Disclosure

None

[wyatt-fluidnumerics]

This looks good to me, but I don't think it should close #2693.

I'd been looking into #2963 already and I think the igrid issue original raised by @VeckoTheGecko is still a problem

This was already in main.

AFAICT the caching of ei isn't even working (searching for .igrid = gives no meaningful results). I already suspected this was the case, hence why I was mentioning before that we should have explicit tests for this as part of our release roadmap.

For another issue/pr

For example, the following code queries the hash table on every particle search despite stationary particles:

import numpy as np

from parcels import FieldSet, ParticleSet
from parcels._core.mesh import FlatMesh
from parcels._core.spatialhash import SpatialHash
from parcels._datasets.structured.generic import datasets

hashed = []
original_query = SpatialHash.query

def mock_query(self, y, x):
    hashed.append(len(y))
    return original_query(self, y, x)

SpatialHash.query = mock_query

ds_curvilinear = datasets["2d_left_rotated"][["data_g", "grid"]].rename({"data_g": "curvilinear"})
ds_rectilinear = datasets["ds_2d_left"][["data_g", "grid"]].rename({"data_g": "rectilinear"})
fieldset = FieldSet.from_sgrid_conventions(ds_curvilinear, mesh=FlatMesh()) + FieldSet.from_sgrid_conventions(
    ds_rectilinear, mesh=FlatMesh()
)

rng = np.random.default_rng(0)
n_particles = 50
x, y, z = rng.uniform(1, 5, n_particles), rng.uniform(1, 5, n_particles), rng.uniform(1, 80, n_particles)


def SampleBothGrids(particles, fieldset):
    _ = fieldset.curvilinear[particles]
    _ = fieldset.rectilinear[particles]


def hashed_per_step(n_steps: int) -> list[int]:
    pset = ParticleSet(fieldset, x=x, y=y, z=z)
    dt = np.timedelta64(1, "D")
    counts = []
    for _ in range(n_steps):
        hashed.clear()
        pset.execute(SampleBothGrids, runtime=dt, dt=dt, verbose_progress=False)
        counts.append(sum(hashed))
    return counts

print(f"igrid as Parcels leaves it: curvilinear {fieldset.curvilinear.igrid}, rectilinear {fieldset.rectilinear.igrid}")
print(f"particles sent to the spatial hash per step: {hashed_per_step(5)} (of {n_particles})")

igrid as Parcels leaves it: curvilinear -1, rectilinear -1
particles sent to the spatial hash per step: [50, 50, 50, 50, 50] (of 50)