Sampling and storing `particle.xi` as a variable gives a different result to calling `particle.xi` (in JIT mode)

[michaeldenes]

Parcels version

3.0.2

Description

In JIT mode, it looks like sampling the particle grid indices doesn't work as expected. All particles have the same grid indices, and these indices do not change in time with advection. In Scipy mode, things seem to work as expected (below, just change from JITParticle to ScipyParticle). Unfortunately my code takes too long to run in Scipy mode, so ideally we can get this working in JIT mode!

I've provided a minimum working example based on the NEMO tutorial below.

Code sample

from datetime import timedelta
from glob import glob

import numpy as np
import parcels

print('Parcels version:', parcels.__version__)

example_dataset_folder = parcels.download_example_dataset(
    "NemoNorthSeaORCA025-N006_data"
)
ufiles = sorted(glob(f"{example_dataset_folder}/ORCA*U.nc"))
vfiles = sorted(glob(f"{example_dataset_folder}/ORCA*V.nc"))
wfiles = sorted(glob(f"{example_dataset_folder}/ORCA*W.nc"))
mesh_mask = f"{example_dataset_folder}/coordinates.nc"

filenames = {
    "U": {"lon": mesh_mask, "lat": mesh_mask, "depth": wfiles[0], "data": ufiles},
    "V": {"lon": mesh_mask, "lat": mesh_mask, "depth": wfiles[0], "data": vfiles},
    "W": {"lon": mesh_mask, "lat": mesh_mask, "depth": wfiles[0], "data": wfiles},
}

variables = {"U": "uo", "V": "vo", "W": "wo"}

# Note that all variables need the same dimensions in a C-Grid
c_grid_dimensions = {
    "lon": "glamf",
    "lat": "gphif",
    "depth": "depthw",
    "time": "time_counter",
}
dimensions = {
    "U": c_grid_dimensions,
    "V": c_grid_dimensions,
    "W": c_grid_dimensions,
}

fieldset = parcels.FieldSet.from_nemo(filenames, variables, dimensions)

n_particles = 5
release_lons = np.linspace(-10,0,n_particles)
release_lats = np.linspace(60,65,n_particles)

# Create a custom particle type to store the grid indices
mode = 'jit'
print('Using mode:', mode)
ParticleType = parcels.ScipyParticle if mode == 'scipy' else parcels.JITParticle
SampleParticle = ParticleType.add_variables(
        [parcels.Variable("grid0_xi", dtype=np.int32, initial=0.0),
         parcels.Variable("grid0_yi", dtype=np.int32, initial=0.0),
         parcels.Variable("grid0_zi", dtype=np.int32, initial=0.0)])

# Sample the grid indices of the particles  at each timestep
def SampleGridIndices(particle, fieldset, time):
    particle.grid0_xi = particle.xi[0]
    particle.grid0_yi = particle.yi[0]
    particle.grid0_zi = particle.zi[0]

kernels = [SampleGridIndices, parcels.AdvectionRK4]

pset = parcels.ParticleSet.from_list(fieldset=fieldset,
                                     pclass=SampleParticle,
                                     lon=release_lons,
                                     lat=release_lats)

pset.execute(kernels, runtime=timedelta(days=2), dt=timedelta(minutes=20))

for particle in pset:
    print('Particle id:', particle._index)
    if particle.grid0_xi != particle.xi[0] or particle.grid0_yi != particle.yi[0] or particle.grid0_zi != particle.zi[0]:
        print('ERROR: The sampled indices do not match the particle indices')
    print('Sampled indices:', particle.grid0_xi, particle.grid0_yi, particle.grid0_zi)
    print('particle.xi,yi,zi:', particle.xi, particle.yi, particle.zi, '\n')

pset.execute(kernels, runtime=timedelta(days=2), dt=timedelta(minutes=20))

for particle in pset:
    print('Particle id:', particle._index)
    if particle.grid0_xi != particle.xi[0] or particle.grid0_yi != particle.yi[0] or particle.grid0_zi != particle.zi[0]:
        print('ERROR: The sampled indices do not match the particle indices')
    print('Sampled indices:', particle.grid0_xi, particle.grid0_yi, particle.grid0_zi)
    print('particle.xi,yi,zi:', particle.xi, particle.yi, particle.zi, '\n')
    

Parcels version: 3.0.2
WARNING: File [/Users/denes001/Library/Caches/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://untitled+.vscode-resource.vscode-cdn.net/Users/denes001/Library/Caches/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2023.6.0). It will be opened with no decoding. Filling values might be wrongly parsed.
Using mode: jit
WARNING: Be careful when sampling particle.xi, as this is updated in the kernel loop. Best to place the sampling statement before advection.
WARNING: Be careful when sampling particle.yi, as this is updated in the kernel loop. Best to place the sampling statement before advection.
WARNING: Be careful when sampling particle.zi, as this is updated in the kernel loop. Best to place the sampling statement before advection.
100%|██████████| 172800.0/172800.0 [00:00<00:00, 37744804.25it/s]
Particle id: 0
Sampled indices: 50 90 0
particle.xi,yi,zi: [50] [90] [0] 

Particle id: 1
ERROR: The sampled indices do not match the particle indices
Sampled indices: 50 90 0
particle.xi,yi,zi: [57] [100] [0] 

Particle id: 2
ERROR: The sampled indices do not match the particle indices
Sampled indices: 50 90 0
particle.xi,yi,zi: [66] [110] [0] 

Particle id: 3
ERROR: The sampled indices do not match the particle indices
Sampled indices: 50 90 0
particle.xi,yi,zi: [70] [120] [0] 

Particle id: 4
ERROR: The sampled indices do not match the particle indices
Sampled indices: 50 90 0
particle.xi,yi,zi: [75] [130] [0] 

100%|██████████| 172800.0/172800.0 [00:00<00:00, 22788106.62it/s]
Particle id: 0
Sampled indices: 51 90 0
particle.xi,yi,zi: [51] [90] [0] 

Particle id: 1
ERROR: The sampled indices do not match the particle indices
Sampled indices: 51 90 0
particle.xi,yi,zi: [59] [99] [0] 

Particle id: 2
ERROR: The sampled indices do not match the particle indices
Sampled indices: 51 90 0
particle.xi,yi,zi: [67] [108] [0] 

Particle id: 3
ERROR: The sampled indices do not match the particle indices
Sampled indices: 51 90 0
particle.xi,yi,zi: [71] [121] [0] 

Particle id: 4
ERROR: The sampled indices do not match the particle indices
Sampled indices: 51 90 0
particle.xi,yi,zi: [76] [130] [0]

[michaeldenes]

As an (obvious) comment, it looks like, in JIT mode, when sampling particle.xi[0] etc. in the SampleGridIndices kernel, it is only taking values from the very first particle. It seems to match with the pset[0].xi[0] fine, and changes temporally, but all the other particles sample the wrong particle grid indices.

[michaeldenes]

This looks to already be fixed in #1614. I've tested the above code (and my separate, more complicated script where I need this) using the current master branch, and all is working as expected! Going to close this as the bug will be fixed in the next release.