Bug in Morton encoding query

[erikvansebille]

Parcels version

v4-dev

Description

The code below leads to the following error. @fluidnumerics-joe, do you have any idea what's going on?

Code sample

import xarray as xr
import numpy as np
import parcels
from glob import glob
DATA_ROOT = "/Users/erik/Desktop/Parcelsv4_test/MOi"

fileU = f"{DATA_ROOT}/GLO12/psy4v3r1-daily_U_2010-01-0[1-3].nc"
filenames = {"U": glob(fileU), "V": glob(fileU.replace("_U_", "_V_")), "W": glob(fileU.replace("_U_", "_W_"))}
mesh_mask = f"{DATA_ROOT}/domain_ORCA0083-N006/PSY4V3R1_mesh_hgr.nc"

fileargs = {"concat_dim": "time_counter",
    "combine": "nested",
    "data_vars": 'minimal',
    "coords": 'minimal',
    "compat": 'override',
    "chunks": {"time_counter": 1, "depth":2, "y": 64, "x": 64}
}

ds_u = xr.open_mfdataset(filenames["U"], **fileargs)[["vozocrtx"]].drop_vars(["nav_lon", "nav_lat"])
ds_v = xr.open_mfdataset(filenames["V"], **fileargs)[["vomecrty"]].drop_vars(["nav_lon", "nav_lat"])
ds_depth = xr.open_mfdataset(filenames["W"], **fileargs)[["depthw"]]
ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0)

ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"})

coords={
    "X": {"left": "x"},
    "Y": {"left": "y"},
    "T": {"center": "time"},
}

ds = ds.isel(depth=0, deptht=0)
print(ds)

xgcm_grid = parcels.xgcm.Grid(ds, coords=coords, periodic=False)
grid = parcels.xgrid.XGrid(xgcm_grid)

U = parcels.Field("U", ds["U"], grid, interp_method=parcels.XLinear)
V = parcels.Field("V", ds["V"], grid, interp_method=parcels.XLinear)
U.units = parcels.GeographicPolar()
V.units = parcels.Geographic()
UV = parcels.VectorField("UV", U, V)

fieldset = parcels.FieldSet([U, V, UV])

lon = -70
lat = -70

pset = parcels.ParticleSet(fieldset=fieldset, lon=lon, lat=lat)

runtime = np.timedelta64(2, "D")
dt = np.timedelta64(15, "m")

pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False)

Traceback (most recent call last):
  File "/Users/erik/Library/CloudStorage/OneDrive-UniversiteitUtrecht/ParcelsSnippets/test_mortonhashing_bug.py", line 55, in <module>
    pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False)
  File "/Users/erik/Codes/ParcelsCode/parcels/particleset.py", line 583, in execute
    self._kernel.execute(self, endtime=next_time, dt=dt)
  File "/Users/erik/Codes/ParcelsCode/parcels/kernel.py", line 250, in execute
    f(pset[evaluate_particles], self._fieldset, None)
  File "/Users/erik/Codes/ParcelsCode/parcels/application_kernels/advection.py", line 103, in AdvectionEE
    (u1, v1) = fieldset.UV[particle]
               ~~~~~~~~~~~^^^^^^^^^^
  File "/Users/erik/Codes/ParcelsCode/parcels/field.py", line 340, in __getitem__
    return self.eval(key.time, key.depth, key.lat, key.lon, key)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Codes/ParcelsCode/parcels/field.py", line 315, in eval
    position = self.grid.search(z, y, x, ei=_ei)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Codes/ParcelsCode/parcels/xgrid.py", line 298, in search
    yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Codes/ParcelsCode/parcels/_index_search.py", line 258, in _search_indices_curvilinear_2d
    yi, xi = grid.get_spatial_hash().query(y, x)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/erik/Codes/ParcelsCode/parcels/spatialhash.py", line 258, in query
    j_best[has_hits] = j[src_best[has_hits]]
                       ~^^^^^^^^^^^^^^^^^^^^
IndexError: index 9223372036854775807 is out of bounds for axis 0 with size 13213618

[fluidnumericsJoe]

I will work on reproducing and get back to you..

[fluidnumericsJoe]

Just dug into this a little bit with pdb. On entry to the query call, the y and x values are NaN at this failure point. We probably need to trace back to see where that started.

  File "/home/joe/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/spatialhash.py", line 258, in query
    j_best[has_hits] = j[src_best[has_hits]]
                       ~^^^^^^^^^^^^^^^^^^^^
IndexError: index 9223372036854775807 is out of bounds for axis 0 with size 13213618
Uncaught exception. Entering post mortem debugging
Running 'cont' or 'step' will restart the program
> /home/joe/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/spatialhash.py(258)query()
-> j_best[has_hits] = j[src_best[has_hits]]
(Pdb) query_codes
array([0], dtype=uint32)
(Pdb) pos
array([0])
(Pdb) qx
array([nan], dtype=float32)
(Pdb) qy
array([nan], dtype=float32)
(Pdb) x
array([nan], dtype=float32)
(Pdb) y
array([nan], dtype=float32)

[fluidnumericsJoe]

I've traced the source of the NaN's back to the interpolation of u and v within parcels.field.VectorField.eval on the first time step. All inputs to this time level look good. I did confirm that the yi,xi index found by the morton indexing (on the first evaluation) is the correct one, meaning that the subsequent search in _search_indices_curvilinear_2d stops in the first iteration with the computational coordinates shown below. Perhaps there's something wrong in the interpolator here ?

(Pdb) l
318  
319             if self._vector_interp_method is None:
320                 u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x)
321                 v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x)
322                 breakpoint()
323  ->             if "3D" in self.vector_type:
324                     w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
325             else:
326                 (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x)
327  
328             if applyConversion:
(Pdb) print(u,v)
[nan] [nan]
(Pdb) print(position)
{'Z': (0, 0.0), 'Y': (array([300]), array([0.31629649])), 'X': (array([2604]), array([0.5]))}
(Pdb) print(tau,ti,time)
[0.] [0] ['2025-01-01T12:00:00.000000000']

[erikvansebille]

Thanks for tracking! You're right that indeed the value in the array is np.nan so that the XLinear interpolator returns np.nan and everything breaks. In v3, we change all np.nan to 0 in all Fields; I'll think about whether we want to keep that behaviour or give users more flexibility about how to deal with NaNs.

In the meantime, the error message here in spatialhash.query() might also be a bit more graceful/informative? Could we return a (-1,-1) if this error occurs (as the doctoring for query suggests)?

[erikvansebille]

I'll think about whether we want to keep that behaviour or give users more flexibility about how to deal with NaNs.

Note that #2152 already (partly) fixed this issue in cad34e0, by doing a call to _update_particle_states_velocity that filters out NaNs. Running the example code above with #2152 results in the (slightly) more informative error

Traceback (most recent call last):
  File "/Users/erik/Library/CloudStorage/OneDrive-UniversiteitUtrecht/ParcelsSnippets/test_mortonhashing_bug.py", line 54, in <module>
    pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False)
  File "/Users/erik/Codes/ParcelsCode/parcels/particleset.py", line 585, in execute
    self._kernel.execute(self, endtime=next_time, dt=dt)
  File "/Users/erik/Codes/ParcelsCode/parcels/kernel.py", line 286, in execute
    error_func(pset[inds].depth, pset[inds].lat, pset[inds].lon)
  File "/Users/erik/Codes/ParcelsCode/parcels/tools/statuscodes.py", line 52, in _raise_field_sampling_error
    raise FieldSamplingError(f"Field sampled at (depth={z}, lat={y}, lon={x})")
parcels.tools.statuscodes.FieldSamplingError: Field sampled at (depth=[0.], lat=[-70.], lon=[-70.])