Setting particles.state with boolean indexing in v4 when number of particles is changing

[reint-fischer]

Question

Question

I am running a simulation using v4 where particles are released at different times and I want to delete them when they run out of bounds. Setting particles.state with boolean indexing doesn’t work, since len(particles) != len(particles.state). I assume that has to do with particles not being released yet. Is this expected behavior, and if so, is there a suggested solution?

Supporting code/error messages

This works when the particles are all released at the same time, but breaks when using this delayed release:

def DeleteOutOfBounds(particles, fieldset):
    print(len(particles))
    print(len(particles.state))
    out_of_bounds = particles.state == parcels.StatusCode.ErrorOutOfBounds
    particles[out_of_bounds].state = parcels.StatusCode.Delete

INFO: Output files are stored in [/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/smagdiff.zarr](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/smagdiff.zarr)
Integration time: 2024-01-01T00:00:11.128799232:   0%|          | 0/604800.0 [00:00<?, ?it/s]
13
1

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[141], [line 15](vscode-notebook-cell:?execution_count=141&line=15)
      9 output_file = parcels.ParticleFile(
     10     store="smagdiff.zarr", outputdt=np.timedelta64(3,"h"), chunks=(1, 57)
     11 )
     13 random.seed(1636)  # Random seed for reproducibility
---> [15](vscode-notebook-cell:?execution_count=141&line=15) pset.execute(
     16     [smagdiff, DeleteOutOfBounds], # parcels.kernels.AdvectionRK4, 
     17     runtime=np.timedelta64(7,"D"),
     18     dt=np.timedelta64(5,"m"),
     19     output_file=output_file,
     20 )

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/particleset.py:549, in ParticleSet.execute(self, pyfunc, dt, endtime, runtime, output_file, verbose_progress)
    546 else:
    547     next_time = end_time
--> [549](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/particleset.py:549) self._kernel.execute(self, endtime=next_time, dt=dt)
    551 if next_output is not None:
    552     if np.abs(next_time - next_output) < np.timedelta64(1000, "ns"):

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/kernel.py:250, in Kernel.execute(self, pset, endtime, dt)
    248 # run kernels for all particles that need to be evaluated
    249 for f in self._pyfuncs:
--> [250](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/kernel.py:250)     f(pset[evaluate_particles], self._fieldset)
    252     # check for particles that have to be repeated
    253     repeat_particles = pset.state == StatusCode.Repeat

Cell In[135], [line 8](vscode-notebook-cell:?execution_count=135&line=8)
      6 print(f"state bool = {particles.state[out_of_bounds]}")
      7 print(f"bool = {out_of_bounds}")
----> [8](vscode-notebook-cell:?execution_count=135&line=8) particles[out_of_bounds].state = parcels.StatusCode.Delete

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/particle.py:137, in KernelParticle.__setattr__(self, name, value)
    135     object.__setattr__(self, name, value)
    136 else:
--> [137](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/particle.py:137)     self._data[name][self._index] = value

IndexError: boolean index did not match indexed array along axis 0; size of axis is 13 but size of corresponding boolean axis is 1

Alternatively, I tried to select just particles.state, but that does not allow me to set a new state:

def DeleteOutOfBounds_alt(particles, fieldset):
    print(len(particles))
    print(len(particles.state))
    out_of_bounds = particles.state == parcels.StatusCode.ErrorOutOfBounds
    particles.state[out_of_bounds] = parcels.StatusCode.Delete
    print(particles.state)

13
13
state = [60 10 10 10 10 10 10 10 10 10 10 10 10]

---------------------------------------------------------------------------
FieldOutOfBoundError                      Traceback (most recent call last)
Cell In[146], [line 15](vscode-notebook-cell:?execution_count=146&line=15)
      9 output_file = parcels.ParticleFile(
     10     store="smagdiff.zarr", outputdt=np.timedelta64(3,"h"), chunks=(1, 57)
     11 )
     13 random.seed(1636)  # Random seed for reproducibility
---> [15](vscode-notebook-cell:?execution_count=146&line=15) pset.execute(
     16     [parcels.kernels.AdvectionRK4, smagdiff, DeleteOutOfBounds],
     17     runtime=np.timedelta64(7,"D"),
     18     dt=np.timedelta64(5,"m"),
     19     output_file=output_file,
     20 )

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/particleset.py:549, in ParticleSet.execute(self, pyfunc, dt, endtime, runtime, output_file, verbose_progress)
    546 else:
    547     next_time = end_time
--> [549](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/particleset.py:549) self._kernel.execute(self, endtime=next_time, dt=dt)
    551 if next_output is not None:
    552     if np.abs(next_time - next_output) < np.timedelta64(1000, "ns"):

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/kernel.py:279, in Kernel.execute(self, pset, endtime, dt)
    277             error_func(pset[inds].time)
    278         else:
--> [279](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/kernel.py:279)             error_func(pset[inds].z, pset[inds].lat, pset[inds].lon)
    281 # Only add PositionUpdate kernel at the end of the first execute call to avoid adding dt to time too early
    282 if not self._positionupdate_kernel_added:

File ~/Documents/UU/parcels/Parcels/src/parcels/_core/statuscodes.py:54, in _raise_field_out_of_bound_error(z, y, x)
     53 def _raise_field_out_of_bound_error(z, y, x):
---> [54](https://file+.vscode-resource.vscode-cdn.net/Users/Gebruiker/Documents/UU/parcels/Parcels/docs/user_guide/examples_v3/~/Documents/UU/parcels/Parcels/src/parcels/_core/statuscodes.py:54)     raise FieldOutOfBoundError(f"Field sampled out-of-bound, at (z={z}, lat={y}, lon={x})")

FieldOutOfBoundError: Field sampled out-of-bound, at (z=[0.494025], lat=[-57.846943], lon=[30.899977])

Simulation setup

# Load the CopernicusMarine data in the Agulhas region from the example_datasets
example_dataset_folder = parcels.download_example_dataset(
    "CopernicusMarine_data_for_Argo_tutorial"
)

ds_fields = xr.open_mfdataset(f"{example_dataset_folder}/*.nc", combine="by_coords")
ds_fields.load()  # load the dataset into memory

x = fieldset.U.grid.lon
y = fieldset.U.grid.lat

def smagdiff(particles, fieldset):
    dt = particles.dt / np.timedelta64(1, "s")
    dx = 0.01
    # gradients are computed by using a local central difference.
    updx, vpdx = fieldset.UV[particles.time, particles.z, particles.lat, particles.lon + dx, particles]
    umdx, vmdx = fieldset.UV[particles.time, particles.z, particles.lat, particles.lon - dx, particles]
    updy, vpdy = fieldset.UV[particles.time, particles.z, particles.lat + dx, particles.lon, particles]
    umdy, vmdy = fieldset.UV[particles.time, particles.z, particles.lat - dx, particles.lon, particles]

    dudx = (updx - umdx) / (2 * dx)
    dudy = (updy - umdy) / (2 * dx)

    dvdx = (vpdx - vmdx) / (2 * dx)
    dvdy = (vpdy - vmdy) / (2 * dx)

    A = fieldset.cell_areas[particles.time, particles.z, particles.lat, particles.lon]
    sq_deg_to_sq_m = (1852 * 60) ** 2 * np.cos(particles.lat * np.pi / 180)
    A = A / sq_deg_to_sq_m
    Kh = fieldset.Cs * A * np.sqrt(dudx**2 + 0.5 * (dudy + dvdx) ** 2 + dvdy**2)

    dlat = random.normalvariate(0.0, 1.0) * np.sqrt(2 * np.fabs(dt) * Kh)
    dlon = random.normalvariate(0.0, 1.0) * np.sqrt(2 * np.fabs(dt) * Kh)

    particles.dlat += dlat
    particles.dlon += dlon

def calc_cell_edge_sizes(grid):
    """Calculate cell sizes based on numpy.gradient method.

    Currently only works for Rectilinear Grids. Operates in place adding a `cell_edge_sizes`
    attribute to the grid.
    """

    cell_edge_sizes_x = np.zeros((grid.ydim+1, grid.xdim+1), dtype=np.float32)
    cell_edge_sizes_y = np.zeros((grid.ydim+1, grid.xdim+1), dtype=np.float32)

    x_conv = parcels.GeographicPolar() if grid._mesh == "spherical" else parcels.UnitConverter()
    y_conv = parcels.Geographic() if grid._mesh == "spherical" else parcels.UnitConverter()
    for y, (lat, dy) in enumerate(zip(grid.lat, np.gradient(grid.lat), strict=False)):
        for x, (lon, dx) in enumerate(
            zip(grid.lon, np.gradient(grid.lon), strict=False)
        ):
            cell_edge_sizes_x[y, x] = x_conv.to_source(dx, grid.depth[0], lat, lon)
            cell_edge_sizes_y[y, x] = y_conv.to_source(dy, grid.depth[0], lat, lon)
    return cell_edge_sizes_x, cell_edge_sizes_y


def calc_cell_areas(grid):
    cell_edge_sizes_x, cell_edge_sizes_y = calc_cell_edge_sizes(grid)
    return cell_edge_sizes_x * cell_edge_sizes_y

ds_fields["cell_areas"] = ds_fields["uo"] * 0 + calc_cell_areas(fieldset.U.grid)

fieldset = parcels.FieldSet.from_copernicusmarine(ds_fields)

fieldset.add_constant("Cs", 0.1)

time = np.arange(fieldset.time_interval.left, fieldset.time_interval.left+np.timedelta64(13,"h"), np.timedelta64(1,'h')) # np.repeat(fieldset.time_interval.left,13)
lon = [31] * len(time)
lat = [-33] * len(time)
z = [ds_fields.depth[0]]  * len(time)
pset = parcels.ParticleSet(
    fieldset=fieldset, pclass=parcels.Particle, time=time, z=z, lat=lat, lon=lon,
)

output_file = parcels.ParticleFile(
    store="smagdiff.zarr", outputdt=np.timedelta64(3,"h"), chunks=(1, 57)
)

random.seed(1636)  # Random seed for reproducibility

pset.execute(
    [smagdiff, DeleteOutOfBounds], # parcels.kernels.AdvectionRK4, 
    runtime=np.timedelta64(7,"D"),
    dt=np.timedelta64(5,"m"),
    output_file=output_file,
)

[erikvansebille]

I think this is directly related to #2143, that @VeckoTheGecko is working on. Good for him to have this background info, I guess