diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py index b629ea624..2ab059b28 100644 --- a/src/parcels/_core/field.py +++ b/src/parcels/_core/field.py @@ -98,8 +98,6 @@ def __init__( self.name = name self.model = model - self.igrid = -1 # Default the grid index to -1 - @property def data(self): return self.model.field_data(self.name) @@ -108,6 +106,14 @@ def data(self): def grid(self): # TODO PR: Remove in favour of referencing model grid directly return self.model.grid + @property + def igrid(self): + return self._igrid + + @igrid.setter + def igrid(self, value): + self._igrid = value + @property def time_interval(self): # TODO PR: Remove in favour of referencing model time_interval directly if "time" not in self.model.data[self.name].dims: @@ -215,7 +221,6 @@ def __init__( self.V = V self.W = W self.grid = U.grid - self.igrid = U.igrid if W is None: _assert_same_time_interval((U, V)) @@ -247,6 +252,10 @@ def interp_method(self, method: VectorInterpolator): raise ValueError(f"method must be a `VectorInterpolator` object. Got {type(method)=!r}") self._interp_method = method + @property + def igrid(self): + return self.U.igrid + def eval(self, t: datetime, z, y, x, particles=None): """Interpolate vectorfield values in space and time. diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index ca6f3aea7..c87575ae5 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -106,6 +106,9 @@ def reconstruct_fields(self): if self.constant_model is not None: fields += self.constant_model.construct_fields() self._fields = {f.name: f for f in fields} + for field in self._fields.values(): + if isinstance(field, Field): + field.igrid = self.gridset.index(field.grid) def __getattr__(self, name): """Get the field by name. If the field is not found, check if it's a context variable.""" diff --git a/tests/test_particleset_execute.py b/tests/test_particleset_execute.py index 2b37e325a..1dc70f6ed 100644 --- a/tests/test_particleset_execute.py +++ b/tests/test_particleset_execute.py @@ -16,6 +16,7 @@ Variable, ) from parcels._core.mesh import FlatMesh +from parcels._core.spatialhash import SpatialHash from parcels._core.statuscodes import GridSearchingError from parcels._core.utils.time import timedelta_to_float from parcels._datasets.structured.generated import simple_UV_dataset @@ -193,6 +194,39 @@ def SampleU(particles, fieldset): # pragma: no cover pset.execute(SampleU, runtime=np.timedelta64(1, "D"), dt=np.timedelta64(1, "D")) +def test_ei_cache_skips_spatial_hash_with_two_grids(monkeypatch): + """Each grid keeps its own cached ei, so stationary particles only go through the spatial hash on the first step.""" + n_hashed = [] + query = SpatialHash.query + + def counting_query(self, y, x): + n_hashed.append(len(y)) + return query(self, y, x) + + monkeypatch.setattr(SpatialHash, "query", counting_query) + + ds_curvilinear = datasets_structured["2d_left_rotated"][["data_g", "grid"]].rename({"data_g": "curvilinear"}) + ds_rectilinear = datasets_structured["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() + ) + + def SampleBothGrids(particles, fieldset): # pragma: no cover + _ = fieldset.curvilinear[particles] + _ = fieldset.rectilinear[particles] + + npart = 50 + rng = np.random.default_rng(0) + pset = ParticleSet(fieldset, x=rng.uniform(1, 5, npart), y=rng.uniform(1, 5, npart), z=rng.uniform(1, 80, npart)) + dt = np.timedelta64(1, "D") + pset.execute(SampleBothGrids, runtime=dt, dt=dt) + assert sum(n_hashed) == npart + + n_hashed.clear() + pset.execute(SampleBothGrids, runtime=4 * dt, dt=dt) + assert sum(n_hashed) == 0 + + def test_particleset_run_to_endtime(fieldset): starttime = fieldset.time_interval.left endtime = fieldset.time_interval.right