Complex numbers in kernels when using JITParticle

[michaeldenes]

Question

Question

I'm wondering how to use complex numbers in kernels when using JIT mode. This may be quite a silly question, I simply need to extract the real part of a complex number, but x.real doesn't seem to work with cgen. See below a minimal working example, using the NEMO parcels tutorial example. There is an error when run in JIT mode, but no error (and the correct output) when run in Scipy mode! Scipy mode is just too slow for my application.

Parcels version:

Parcels version: 3.0.5.dev202

Supporting code/error messages

from datetime import timedelta
from glob import glob

import numpy as np
import math
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("real_from_complex", dtype=np.int32, initial=0.0),
         parcels.Variable("imag_from_complex", dtype=np.int32, initial=0.0)])

# Sample the grid indices of the particles  at each timestep
def computeRealFromComplex(particle, fieldset, time):
    someValue = 10 + 5j

    realPart = someValue.real
    imagPart = someValue.imag

    particle.real_from_complex = realPart
    particle.imag_from_complex = imagPart

kernels = [computeRealFromComplex, parcels.AdvectionRK4]

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

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

pset

When run in JIT mode:

Parcels version: 3.0.5.dev202
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
Using mode: jit
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[3], [line 70](vscode-notebook-cell:?execution_count=3&line=70)
     [63](vscode-notebook-cell:?execution_count=3&line=63) kernels = [computeRealFromComplex, parcels.AdvectionRK4]
     [65](vscode-notebook-cell:?execution_count=3&line=65) pset = parcels.ParticleSet.from_list(fieldset=fieldset,
     [66](vscode-notebook-cell:?execution_count=3&line=66)                                      pclass=SampleParticle,
     [67](vscode-notebook-cell:?execution_count=3&line=67)                                      lon=release_lons,
     [68](vscode-notebook-cell:?execution_count=3&line=68)                                      lat=release_lats)
---> [70](vscode-notebook-cell:?execution_count=3&line=70) pset.execute(kernels, runtime=timedelta(minutes=40), dt=timedelta(minutes=20))
     [72](vscode-notebook-cell:?execution_count=3&line=72) pset

File ~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:960, in ParticleSet.execute(self, pyfunc, pyfunc_inter, endtime, runtime, dt, output_file, verbose_progress, postIterationCallbacks, callbackdt, delete_cfiles)
    [958](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:958)     self.kernel = pyfunc
    [959](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:959) else:
--> [960](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:960)     self.kernel = self.Kernel(pyfunc, delete_cfiles=delete_cfiles)
    [961](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:961) # Prepare JIT kernel execution
    [962](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:962) if self.particledata.ptype.uses_jit:

File ~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:813, in ParticleSet.Kernel(self, pyfunc, c_include, delete_cfiles)
    [796](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:796) """Wrapper method to convert a `pyfunc` into a :class:`parcels.kernel.Kernel` object.
    [797](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:797) 
    [798](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:798) Conversion is based on `fieldset` and `ptype` of the ParticleSet.
   (...)
    [810](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:810)      (Default value = "")
    [811](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/particleset.py:811) """
...
    [269](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/compilation/codegenerator.py:269)     )
    [270](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/compilation/codegenerator.py:270) else:
--> [271](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/Projects/EddyPlasticEntrainment/simulations/plastic/~/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/compilation/codegenerator.py:271)     raise NotImplementedError(f"Cannot convert '{node.value.id}' used in kernel to C-code")

NotImplementedError: Cannot convert 'someValue' used in kernel to C-code

When run in Scipy mode, the output of pset is:

Parcels version: 3.0.5.dev202
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
[/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.conda/envs/parcels-v304-bugfix/lib/python3.12/site-packages/parcels/field.py:511): FileWarning: File [/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc](https://vscode-remote+ssh-002dremote-002blorenz-002dshort.vscode-resource.vscode-cdn.net/nethome/denes001/.cache/parcels/NemoNorthSeaORCA025-N006_data/coordinates.nc) could not be decoded properly by xarray (version 2024.7.0). It will be opened with no decoding. Filling values might be wrongly parsed.
  with _grid_fb_class(lonlat_filename, dimensions, indices, netcdf_engine) as filebuffer:
Using mode: scipy
100%|██████████| 2400.0/2400.0 [00:00<00:00, 4749.10it/s]
P[0](lon=-9.996915, lat=59.999753, depth=0.000000, real_from_complex=10.000000, imag_from_complex=5.000000, time=1200.000000)
P[1](lon=-7.497746, lat=61.249643, depth=0.000000, real_from_complex=10.000000, imag_from_complex=5.000000, time=1200.000000)
P[2](lon=-4.993494, lat=62.499670, depth=0.000000, real_from_complex=10.000000, imag_from_complex=5.000000, time=1200.000000)
P[3](lon=-2.498075, lat=63.750454, depth=0.000000, real_from_complex=10.000000, imag_from_complex=5.000000, time=1200.000000)
P[4](lon=0.001060, lat=64.999866, depth=0.000000, real_from_complex=10.000000, imag_from_complex=5.000000, time=1200.000000)

[michaeldenes]

To add to this, as the kernel I'm developing doesn't just define some complex number and then compute the real/imaginary parts of it (the example above), I'd also like to make sure any computation with complex numbers is correct also! For example:

import math
def SomeComplexNumberFunction(particle, fieldset, time):
    rdelta = 1./1.15 # Some relation between densities of two things
    
    psi = (1.j * math.sqrt(1. - (rdelta - 1.)**2) + rdelta - 1.)**(1./3.)
    Phi = (1.j * math.sqrt(3.) / 2.) * (1. / psi - psi) - 1. / (2. * psi) - psi / 2. + 1.

    Phi = Phi.real # This not only needs to not throw an error, but be computed correctly too!

    # some other computations
    ...

[erikvansebille]

Thanks for the question, @michaeldenes. Parcels currently doesn't support Complex numbers in its python-to-C-codeconverter. It would be a major undertaking to implement support, requiring a lot of extra code. Is there no way around using complex numbers?

One alternative, if you really need to use complex number arithmetic, is to use "custom headers". This is a little-used feature in Parcels, that allows you to directly inject C-code into your python-Kernels. See e.g. below for how this is used in the unit-tests

Parcels/tests/test_kernel_language.py

Lines 370 to 397 in e3c4615

	def test_c_kernel(fieldset_unit_mesh, mode, c_inc):
	coord_type = np.float32 if c_inc == "str" else np.float64
	pset = ParticleSet(fieldset_unit_mesh, pclass=ptype[mode], lon=[0.5], lat=[0], lonlatdepth_dtype=coord_type)
	def func(U, lon, dt):
	u = U.data[0, 2, 1]
	return lon + u * dt
	if c_inc == "str":
	c_include = """
	static inline StatusCode func(CField *f, double *particle_dlon, double *dt)
	{
	float data2D[2][2][2];
	StatusCode status = getCell2D(f, 1, 2, 0, data2D, 1); CHECKSTATUS(status);
	float u = data2D[0][0][0];
	*particle_dlon = +u * *dt;
	return SUCCESS;
	}
	"""
	else:
	c_include = os.path.join(os.path.dirname(__file__), "customed_header.h")
	def ckernel(particle, fieldset, time):
	func("parcels_customed_Cfunc_pointer_args", fieldset.U, particle_dlon, particle.dt) # noqa
	kernel = pset.Kernel(ckernel, c_include=c_include)
	pset.execute(kernel, endtime=4.0, dt=1.0)
	assert np.allclose(pset.lon[0], 0.81578948)

If you can write a function that uses complex arithmetic in C and put it in such a header file, you could(?) then put it into a kernel and call it that way. But be warned; getting this method to work may be very finicky...