"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

[VeckoTheGecko]

Looks like the problem is in the C-conversion code in kernel.py (perhaps not being able to handle the dataset for some reason?)

[VeckoTheGecko]

After some investigation from @sruehs and @michaeldenes it looks to be a problem with the data and not Parcels

[michaeldenes]

From the periodic boundary conditions tutorial (https://docs.oceanparcels.org/en/latest/examples/tutorial_periodic_boundaries.html), the following lines are used to add the periodic halo:

fieldset.add_constant("halo_west", fieldset.U.grid.lon[0])
fieldset.add_constant("halo_east", fieldset.U.grid.lon[-1])

fieldset.add_periodic_halo(zonal=True)

The grid here is a RectilinearZGrid, and when we print the shape of the grid we get:

print(fieldset.U.grid.lon.shape, fieldset.U.grid.lat.shape)

(110,) (50,) i.e. 1D arrays.

So, when creating the halos, fieldset.halo_west and fieldset.halo_east are scalar values.

However, in @sruehs case, the shape of underlying grids are not 1D-arrays, rather 2D-arrays. When adding the two constants halo_west and halo_east, by providing fieldset.U.grid.lon[0] and fieldset.U.grid.lon[-1], you are providing an array of values, not a single scalar value. This is why the error pops up.

In the past, using an Arakawa B-grid (and MOM5 data), I've used the exact same code to add the halo, so I'm not sure why it's not working here, unless the grid is weird...

[VeckoTheGecko]

https://github.com/OceanParcels/Parcels/blob/93cc635b82397a826e4e0869f2db4ecc7b3bfb19/parcels/fieldset.py#L1352-L1372

Fieldset.add_constant() needs to have the value being a float, so in the case of a 2D array ...lon[0] would return an array instead of a single float. This isn't validated in the function hence the uninformative error message.

What is the usecase for having halo_west being an array instead of a constant?

---

I still am a bit concerned about the format of this data (ie., 'why is lon a 2D array?'), and how it compares against other datasets we've used. I can block out some time to look at this properly

EDIT: I see now. Since its a curvilinear grid lon and lat aren't dimensions (more the dimensions are xi and yi which unique combinations provide lon and lat)

[michaeldenes]

I don't think there is a use-case for halo_west or halo_east being an array, but I think the periodic BC tutorial should be updated to show how to implement periodic BC's for different models/grids.

I think parcels handles the periodicity in the nemo c-grids out of the box, but just to show as an example, if you load the data from the nemo tutorial (https://docs.oceanparcels.org/en/latest/examples/tutorial_nemo_3D.html), and print out the shapes of the grids:

print(fieldset.U.grid.lon.shape, fieldset.U.grid.lat.shape)

you get (201, 151) (201, 151).