"only length-1 arrays can be converted to Python scalars" error when setting periodic boundaries

[sruehs]

Parcels version

3.0.3

Description

I encountered an error when running Parcels with b-SOSE (MITgcm) fields on Lorenz, which I have not seen before. It seems to be linked to some kind of conversion? Any help how to approach this is appreciated!

Code sample

import numpy as np
from glob import glob
from datetime import timedelta

from parcels import AdvectionRK4_3D, FieldSet, JITParticle, ParticleSet

filepath_u_5daily = sorted(
    glob("/nethome/ruhs0001/DATA/b-SOSE-iter154/5daily/Uvel_bsoseI154_*_5day.nc")
)
filepath_v_5daily = sorted(
    glob("/nethome/ruhs0001/DATA/b-SOSE-iter154/5daily/Vvel_bsoseI154_*_5day.nc")
)
filepath_w_5daily = sorted(
    glob("/nethome/ruhs0001/DATA/b-SOSE-iter154/5daily/Wvel_bsoseI154_*_5day.nc")
)
gridpath = "/nethome/ruhs0001/DATA/b-SOSE-iter154/grid.nc"


output_dir = "/nethome/ruhs0001/CLIMSATE_SouthAtlanticTransports/dev-lorenz/processed_data/traj_simulations/b-SOSE-iter154/Test/"


def get_fieldset(filepath_u, filepath_v, filepath_w):

    filenames = {
        "U": {
            "lon": gridpath,
            "lat": gridpath,
            "depth": filepath_w[0],
            "data": filepath_u,
        },
        "V": {
            "lon": gridpath,
            "lat": gridpath,
            "depth": filepath_w[0],
            "data": filepath_v,
        },
        "W": {
            "lon": gridpath,
            "lat": gridpath,
            "depth": filepath_w[0],
            "data": filepath_w,
        },
    }
    variables = {
        "U": "UVEL",
        "V": "VVEL",
        "W": "WVEL",
    }

    c_grid_dimensions = {"lon": "XG", "lat": "YG", "depth": "Zl", "time": "time"}
    dimensions = {
        "U": c_grid_dimensions,
        "V": c_grid_dimensions,
        "W": c_grid_dimensions,
    }
    fieldset = FieldSet.from_mitgcm(
        filenames,
        variables,
        dimensions,
        mesh="spherical",
        tracer_interp_method="cgrid_tracer",
        allow_time_extrapolation=False,
        time_periodic=False,
        deferred_load=True,
    )

    fieldset.add_constant("halo_west", fieldset.U.grid.lon[0])
    fieldset.add_constant("halo_east", fieldset.U.grid.lon[-1])
    fieldset.add_constant("halo_south", fieldset.U.grid.lat[0])
    fieldset.add_constant("halo_north", fieldset.U.grid.lat[-1])
    fieldset.add_periodic_halo(zonal=True, meridional=True)

    return fieldset


fieldset_5daily = get_fieldset(filepath_u_5daily, filepath_v_5daily, filepath_w_5daily)


init_lon_tmp = np.arange(-60, -56, 0.1)
init_lat_tmp = -44
init_depth_tmp = np.arange(10, 100, 10)
init_lon, init_lat, init_depth = np.meshgrid(init_lon_tmp, init_lat_tmp, init_depth_tmp)
init_lon = init_lon.flatten()
init_lat = init_lat.flatten()
init_depth = init_depth.flatten()


integration_dt = 20
runtime = 30
output_dt = 1


def periodicBC(particle, fieldset, time):
    if particle.lon < fieldset.halo_west:
        particle_dlon += fieldset.halo_east - fieldset.halo_west
    elif particle.lon > fieldset.halo_east:
        particle_dlon -= fieldset.halo_east - fieldset.halo_west


def create_run_particles(
    fieldset,
    output_filename,
    init_lon=init_lon,
    init_lat=init_lat,
    init_depth=init_depth,
    integration_dt=integration_dt,
    runtime=runtime,
    output_dt=output_dt,
):

    pset = ParticleSet.from_list(
        fieldset=fieldset,
        pclass=JITParticle,
        lon=init_lon,
        lat=init_lat,
        depth=init_depth,
    )

    output_file = pset.ParticleFile(
        name=output_dir + output_filename, outputdt=timedelta(days=output_dt)
    )

    pset.execute(
        pset.Kernel(AdvectionRK4_3D) + pset.Kernel(periodicBC),
        runtime=timedelta(days=runtime),
        dt=timedelta(minutes=integration_dt),
        output_file=output_file,
        verbose_progress=True,
    )


create_run_particles(
    fieldset=fieldset_5daily, output_filename="Trajectories5daily.zarr"
)

---------------------------------------------------------------------------
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[23], line 2
      1 # performing particle trajectory integration and saving trajectory data
----> 2 create_run_particles(fieldset=fieldset_5daily,
      3                      output_filename="Trajectories5daily.zarr")

Cell In[22], line 14
      7 pset = ParticleSet.from_list(fieldset=fieldset,
      8                              pclass=JITParticle,
      9                              lon=init_lon, lat=init_lat, depth=init_depth)
     11 output_file = pset.ParticleFile(name=output_dir + output_filename,
     12                                 outputdt=timedelta(days=output_dt))
---> 14 pset.execute(pset.Kernel(AdvectionRK4_3D) + pset.Kernel(periodicBC),
     15              runtime=timedelta(days=runtime),
     16              dt=timedelta(minutes=integration_dt),
     17              output_file=output_file,
     18              verbose_progress=True)

File /opt/apps/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/particleset.py:956, in ParticleSet.execute(self, pyfunc, pyfunc_inter, endtime, runtime, dt, output_file, verbose_progress, postIterationCallbacks, callbackdt, delete_cfiles)
    954 # If we don't perform interaction, only execute the normal kernel efficiently.
    955 if self.interaction_kernel is None:
--> 956     res = self.kernel.execute(self, endtime=next_time, dt=dt)
    957     if res == StatusCode.StopAllExecution:
    958         return StatusCode.StopAllExecution

File /opt/apps/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/kernel.py:578, in Kernel.execute(self, pset, endtime, dt)
    576 # Execute the kernel over the particle set
    577 if self.ptype.uses_jit:
--> 578     self.execute_jit(pset, endtime, dt)
    579 else:
    580     self.execute_python(pset, endtime, dt)

File /opt/apps/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/kernel.py:541, in Kernel.execute_jit(self, pset, endtime, dt)
    538 self.load_fieldset_jit(pset)
    540 fargs = [byref(f.ctypes_struct) for f in self.field_args.values()]
--> 541 fargs += [c_double(f) for f in self.const_args.values()]
    542 particle_data = byref(pset.ctypes_struct)
    543 return self._function(c_int(len(pset)), particle_data,
    544                       c_double(endtime), c_double(dt), *fargs)

TypeError: only length-1 arrays can be converted to Python scalars