Windowed Array

[fluidnumericsJoe]

Description

This PR adds a new WindowedArray class. For large datasets that do not fit entirely in memory, users often lean on dask-backed xarray/uxarray datasets . Particle advection on top of gridded data typically results in large out-of-order indexing requests which results in significant dask scheduler overhead. To mitigate the expense of scheduler overhead, the WindowedArray class overrides the .isel method to store 1 or more time slices in a numpy cache. If a time slice exists in the numpy cache, kernel execution sees zero overhead (there is no dask at this point); for time slices not in the cache, the dask scheduler overhead is minimal since large contiguous blocks of data (time slice) is loaded from dask into the numpy cache as needed.

A method is provided at the model level (.to_windowed_arrays()) to allow users to opt in to using the windowed arrays.

Supersedes #2671

Checklist

• Closes Out-of-core fields: move unstructured interpolators to .isel, and add a time-windowed array cache (WindowedArray) #2656
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): This is a claude-assisted port of Issue 2656 windowed array #2671 to the new main branch of parcels. After the initial port was completed, I made minor edits to docstrings and added support for caching that is independent of time integration direction.

[VeckoTheGecko]

#2719 is merged, so I think this is unblocked

[erikvansebille]

Looks good; a few minor comments below

[erikvansebille]

I don't understand what this means. Can you further clarify? I would say that the default should be 2 levels?

[fluidnumericsJoe]

I can change the default to 2 levels in the WindowedArray class. The thinking here in the default was to, by default, just retain in cache the span of time indices covered by a .isel() call. Admittedly this could lead to accidental OOM more readily than defaulting to 1 or 2 time levels.

[erikvansebille]

Ah, interesting idea to let the .isel decide the number of levels. But I didn't;t get that from the doctoring description, do perhaps that should be clarified then?

[erikvansebille]

Actually, I'm going to retract my Approval of this PR. I wanted to see how it would do on the tutorial_Argofloats notebook in terms of timing, so changed the ds.load() in that notebook for a fieldset.to_windowed_arrays() call (see eed1065)

However, I then got the following error, which seems a problem in this implementation?

    ---------------------------------------------------------------------------
    ValueError                                Traceback (most recent call last)
    Cell In[2], line 56
         52     mode="w",
         53 )
         54 
         55 # Now execute the Kernels for 30 days, saving data every 30 minutes
    ---> 56 pset.execute(
         57     kernels,
         58     runtime=timedelta(days=30),
         59     dt=timedelta(minutes=15),
    
    File ~/Codes/parcels/src/parcels/_core/particleset.py:446, in ParticleSet.execute(self, kernels, dt, endtime, runtime, output_file, verbose_progress)
        443 else:
        444     next_time = end_time
    --> 446 self._kernel.execute(self, endtime=next_time, dt=dt)
        448 if next_output is not None:
        449     if np.abs(next_time - next_output) < 0.001:
    
    File ~/Codes/parcels/src/parcels/_core/kernel.py:212, in Kernel.execute(self, pset, endtime, dt)
        209 with warnings.catch_warnings():
        210     warnings.simplefilter("ignore", FieldEvalWarning)
    --> 212     f(pset[evaluate_particles], self._fieldset)
        214     # check for particles that have to be repeated
        215     repeat_particles = pset.state == StatusCode.Repeat
    
    Cell In[1], line 40, in ArgoVerticalMovement(particles, fieldset)
         36     ptcls2.dz[next_phase] = maxdepth - ptcls2.z[next_phase]  # avoid overshoot
         37 
         38     # Phase 3: Rising with vertical_speed until at surface
         39     ptcls3.dz -= vertical_speed * ptcls3.dt
    ---> 40     ptcls3.temp = fieldset.thetao[ptcls3.t, ptcls3.z, ptcls3.y, ptcls3.x]
         41     next_phase = ptcls3.z + ptcls3.dz <= mindepth
         42     ptcls3.cycle_phase[next_phase] = 4
         43     ptcls3.dz[next_phase] = mindepth - ptcls3.z[next_phase]  # avoid overshoot
    
    File ~/Codes/parcels/src/parcels/_core/field.py:189, in Field.__getitem__(self, key)
        187         return self.eval(key.t, key.z, key.y, key.x, key)
        188     else:
    --> 189         return self.eval(*key)
        190 except tuple(AllParcelsErrorCodes.keys()) as error:
        191     return _deal_with_errors(error, key, vector_type=None)
    
    File ~/Codes/parcels/src/parcels/_core/field.py:176, in Field.eval(self, t, z, y, x, particles)
        172 x = np.atleast_1d(x)
        174 particle_positions, grid_positions = _get_positions(self, t, z, y, x, particles, _ei)
    --> 176 value = self.interp_method.interp(particle_positions, grid_positions, self)
        178 _update_particle_states_interp_value(particles, value)
        179 _mask_outofbounds_values(grid_positions, value)
    
    File ~/Codes/parcels/src/parcels/interpolators/_xinterpolators.py:102, in XLinear.interp(self, particle_positions, grid_positions, field)
         99 lenT = 2 if np.any(tau > 0) else 1
        100 lenZ = 2 if np.any(zeta > 0) else 1
    --> 102 corner_data = _get_corner_data_Agrid(data, ti, zi, yi, xi, lenT, lenZ, len(xsi), axis_dim)
        104 if lenT == 2:
        105     tau = tau[np.newaxis, :]
    
    File ~/Codes/parcels/src/parcels/interpolators/_xinterpolators.py:65, in _get_corner_data_Agrid(data, ti, zi, yi, xi, lenT, lenZ, npart, axis_dim)
         62 if "time" in data.dims:
         63     selection_dict["time"] = xr.DataArray(ti, dims=("points"))
    ---> 65 return data.isel(selection_dict).data.reshape(lenT, lenZ, 2, 2, npart)
    
    File ~/Codes/parcels/src/parcels/_core/_windowed_array.py:86, in WindowedArray.isel(self, indexers, **kwargs)
         84 t_vals = np.asarray(t_ind.values if isinstance(t_ind, xr.DataArray) else t_ind)
         85 levels = np.unique(t_vals)
    ---> 86 self._ensure(levels)
         88 # stack the resident levels into one small NumPy block; remap to local indices
         89 block = np.stack([self._cache[int(lvl)] for lvl in levels])  # (nlevels, *rest)
    
    File ~/Codes/parcels/src/parcels/_core/_windowed_array.py:67, in WindowedArray._ensure(self, levels)
         62         self.bytes_read += self._slab_bytes
         63 # retire cached levels outside the span this call requested. Direction never
         64 # enters here: a forward (dt > 0) or backward (dt < 0) clock both shed their
         65 # trailing edge. Consecutive brackets overlap on one endpoint (inside [lo, hi]),
         66 # so it is retained and each level is still read at most once per pass.
    ---> 67 lo, hi = int(np.min(levels)), int(np.max(levels))
         68 for old in [k for k in self._cache if k < lo or k > hi]:
         69     del self._cache[old]
    
    File ~/Codes/parcels/.pixi/envs/docs/lib/python3.14/site-packages/numpy/_core/fromnumeric.py:3261, in min(a, axis, out, keepdims, initial, where)
       3149 @array_function_dispatch(_min_dispatcher)
       3150 def min(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue,
       3151         where=np._NoValue):
       3152     """
       3153     Return the minimum of an array or minimum along an axis.
       3154 
       (...)   3259     6
       3260     """
    -> 3261     return _wrapreduction(a, np.minimum, 'min', axis, None, out,
       3262                           keepdims=keepdims, initial=initial, where=where)
    
    File ~/Codes/parcels/.pixi/envs/docs/lib/python3.14/site-packages/numpy/_core/fromnumeric.py:83, in _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs)
         80         else:
         81             return reduction(axis=axis, out=out, **passkwargs)
    ---> 83 return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
    
    ValueError: zero-size array to reduction operation minimum which has no identity

Can you fix this error, @fluidnumericsJoe?

[fluidnumericsJoe]

This error message is curious

    ValueError: zero-size array to reduction operation minimum which has no identity

it suggests that levels in the _ensure(levels) function is length 0 . This suggests that the time indices requested in _xinterpolators.py:65 is empty .

Looking at potential solution here.

[fluidnumericsJoe]

@erikvansebille - cb588b1 resolves the issue in the modified argo float tutorial

[erikvansebille]

Looks good. Happy to be merged when CI gives green ticks