sedtrails.particle_tracer.position_calculator_numba

Numba-optimized ParticlePositionCalculator

This module provides a Numba-accelerated implementation of particle position calculation using a 4-stage Runge-Kutta integration scheme on unstructured grids. It replaces the matplotlib triangulation and multiprocessing approaches with Numba-compatible alternatives.

Dependencies: numpy, numba, scipy (only for initial triangulation)

Module Contents

Functions

create_numba_particle_calculator	Factory function to create a Numba-optimized particle calculator.
find_triangle	Find which triangle contains the point (x, y).
interpolate_field	Interpolate field values at given points using barycentric coordinates.
update_particles_rk4	Update particle positions using a 4-stage Runge-Kutta integration scheme.
update_particles_rk4_parallel	Parallel version of update_particles_rk4 using Numba’s prange.

API

sedtrails.particle_tracer.position_calculator_numba.create_numba_particle_calculator(grid_x, grid_y, triangles=None)

Factory function to create a Numba-optimized particle calculator.

This function preprocesses the grid data and returns optimized Numba functions for particle position calculation.

grid_x, grid_yarray_like

Coordinates of grid nodes

trianglesarray_like, optional

Triangle connectivity (node indices). If None, Delaunay triangulation will be computed once (not using Numba).

dict

Dictionary containing optimized Numba functions

sedtrails.particle_tracer.position_calculator_numba.find_triangle(x, y, grid_x, grid_y, triangles)

Find which triangle contains the point (x, y).

x, yfloat

Coordinates of the point

grid_x, grid_yarray_like

Coordinates of grid nodes

trianglesarray_like

Triangle connectivity (node indices)

int

Triangle index or -1 if outside all triangles

sedtrails.particle_tracer.position_calculator_numba.interpolate_field(field, x_points, y_points, grid_x, grid_y, triangles)

Interpolate field values at given points using barycentric coordinates.

fieldarray_like

Field values at grid nodes

x_points, y_pointsarray_like

Coordinates of points where to interpolate

grid_x, grid_yarray_like

Coordinates of grid nodes

trianglesarray_like

Triangle connectivity (node indices)

array_like

Interpolated field values

sedtrails.particle_tracer.position_calculator_numba.update_particles_rk4(x0, y0, grid_u, grid_v, grid_x, grid_y, triangles, dt, igeo=0)

Update particle positions using a 4-stage Runge-Kutta integration scheme.

x0, y0array_like

Initial particle positions

grid_u, grid_varray_like

Velocity components at grid nodes

grid_x, grid_yarray_like

Coordinates of grid nodes

trianglesarray_like

Triangle connectivity (node indices)

dtfloat

Time step

igeoint, optional

Flag for geographic coordinates adjustment (default: 0)

tuple

Updated particle positions (x_new, y_new)

sedtrails.particle_tracer.position_calculator_numba.update_particles_rk4_parallel(x0, y0, grid_u, grid_v, grid_x, grid_y, triangles, dt, igeo=0)

Parallel version of update_particles_rk4 using Numba’s prange.

This function is optimized for large particle sets.

Parameters are the same as update_particles_rk4.