Fermat path-tracing with JAX

fpt-jax is a standalone library for differentiable path-tracing using the Fermat principle, implemented with JAX.

Installation

You can install this package from PyPI:

pip install fpt-jax

Usage

This library implements a single function, trace_rays, which traces rays undergoing specular reflections and diffractions on planar objects defined by origins and basis vectors:

> from fpt_jax import trace_rays; help(trace_rays)

trace_rays
   (tx: jax.Array, rx: jax.Array,
    object_origins: jax.Array, object_vectors: jax.Array,
    interaction_types: jax.Array | None = None, *,
    num_iters: int | None = None, max_num_iters: int | None = None,
    rtol: float | None = None, atol: float | None = None,
    unroll: int | bool = 1,
    num_iters_linesearch: int = 1, unroll_linesearch: int | bool = 1,
    implicit_diff: bool = True,
    use_image_method: bool = True) -> jax.Array:

Compute the points of interaction of rays with objects using Fermat's principle.

Each ray is obtained by minimizing the total travel distance from transmitter to receiver using a quasi-Newton optimization algorithm (BFGS). When use_image_method=True (default), intermediate specular reflections and refractions/transmissions are solved in closed form via the exact image method, drastically reducing the optimization dimension to only the diffraction edge parameters.

This function accepts batched inputs, where the leading dimensions must be broadcast-compatible.

Args:
    tx: Transmitter positions of shape (..., 3).
    rx: Receiver positions of shape (..., 3).
    object_origins: Origins of the objects of shape (..., num_interactions, 3).
    object_vectors: Vectors defining the objects of shape (..., num_interactions, num_dims, 3).
    interaction_types: Optional interaction types of shape (..., num_interactions)
        (0 = reflection, 1 = diffraction, 2 = transmission). If omitted, planar surfaces
        are treated as specular reflections.
    num_iters: Fixed number of iterations for the optimization algorithm.
        Mutually exclusive with max_num_iters.
    max_num_iters: Maximum number of iterations for adaptive optimization using a while loop
        with Cauchy termination. Mutually exclusive with num_iters. When specified, rtol
        and atol are required, and unroll must be 1 or False.
    rtol: Relative tolerance for the Cauchy termination criterion. Required when max_num_iters is set.
    atol: Absolute tolerance for the Cauchy termination criterion. Required when max_num_iters is set.
    unroll: If an integer, the number of optimization iterations to unroll in the JAX scan.
        If True, unroll all iterations. If False, do not unroll. Must be 1 or False if max_num_iters is set.
    num_iters_linesearch: Number of iterations for the line search fixed-point iteration.
    unroll_linesearch: If an integer, the number of fixed-point iterations to unroll in the JAX scan.
        If True, unroll all iterations. If False, do not unroll.
    implicit_diff: Whether to use implicit differentiation for computing the gradient.
        If True, assumes that the solution has converged and applies the implicit function theorem
        to differentiate the optimization problem with respect to the input parameters:
            tx, rx, object_origins, and object_vectors.
        If False, the gradient is computed by backpropagating through all iterations of the optimization algorithm.

        Using implicit differentiation is more memory- and computationally efficient,
        as it does not require storing intermediate values from all iterations,
        but it may be less accurate if the optimization has not fully converged.
        Moreover, implicit differentiation is not compatible with forward-mode autodiff in JAX.
    use_image_method: If True (default), specular planar interactions are solved exactly in closed
        form using the image method, reducing optimization to only diffracting edges. If False,
        runs standard BFGS over all coordinates simultaneously.

Returns:
    The points of interaction of shape (..., num_interactions, 3).
    To include the transmitter and receiver positions, concatenate tx and rx to the result.

This algorithm is also available within DiffeRT, our differentiable ray tracing library for radio propagation.

Getting help

For any question about the method or its implementation, make sure to first read the related paper.

If you want to report a bug in this library or the underlying algorithm, please open an issue on this GitHub repository. If you want to request a new feature, please consider opening an issue on DiffeRT's GitHub repository instead.

Citing

If you use this library in your research, please cite our paper:

@misc{eertmans2025fpt,
  title         = {Fast, Differentiable, GPU-Accelerated Ray Tracing for Multiple Diffraction and Reflection Paths},
  author        = {Jérome Eertmans and Sophie Lequeu and Benoît Legat and Laurent Jacques and Claude Oestges},
  year          = 2025,
  url           = {https://arxiv.org/abs/2510.16172},
  eprint        = {2510.16172},
  archiveprefix = {arXiv},
  primaryclass  = {eess.SP}
}