Wavefront

Arrival time computation methods for track etching simulation.

This package provides multiple algorithms for computing arrival times on 2D non-uniform (e.g., log-spaced radial, linear depth) grids:

  • FMM (Fast Marching Method): For uniform grids

  • Dijkstra: For non-uniform grids, with Numba or C++ backends

Main algorithms

FMM (wavefront.fmm)

Fast Marching Method for smooth speed maps. Requires uniform grid (or interpolation from log-spaced grid).

Dijkstra (wavefront.dijkstra)

Dijkstra’s algorithm optimized for non-uniform grids using 8-connected neighborhood. Available backends:

  • dijkstra_numba (default): Numba JIT-compiled, ~5-20x speedup

  • dijkstra_cpp (optional): C++ with pybind11, ~50-100x speedup

Examples

Using Dijkstra with log-spaced grid (typical):

from tracketch.wavefront.dijkstra import arrival_time_dijkstra_fast

arrival_map = arrival_time_dijkstra_fast(
    r_log_um, z_um, etch_rate_map
)

Using FMM with interpolation to uniform grid:

from tracketch.wavefront.fmm import arrival_time_fmm

arrival_map = arrival_time_fmm(
    r_log_um, z_um, etch_rate_map, r_is_logscaled=True
)

Using C++ Dijkstra (if compiled):

from tracketch.wavefront.dijkstra import arrival_time_dijkstra_cpp

arrival_map = arrival_time_dijkstra_cpp(
    r_log_um, z_um, etch_rate_map
)
tracketch.wavefront.arrival_time_dijkstra_fast(r_um: ndarray, z_um: ndarray, etch_rate_map: ndarray, r_is_logscaled: bool = True, connectivity: Literal[8, 16, 32] = 8) → ndarray[source]

Compute arrival time using Dijkstra’s algorithm with Numba-accelerated graph construction.

Parameters:

  • r_um (np.ndarray) – Radial coordinates (can be non-uniform/log-spaced)

  • z_um (np.ndarray) – Depth coordinates (linear)

  • etch_rate_map (np.ndarray) – Etch rate (speed) at each (z, r) point, shape (n_z, n_r)

  • r_is_logscaled (bool) – Whether r grid is log-spaced (informational only)

  • connectivity ({8, 16, 32}) – Number of neighbor connections per node. - 8: Standard (0deg, 45deg, 90deg, …) - fastest, ~2% metrication error - 16: Knight’s move added (~1% error) - 32: Extended knight’s moves (~0.5% error)

Returns:

arrival_time_map – Arrival time at each grid point, shape (n_z, n_r)

Return type:

np.ndarray

tracketch.wavefront.arrival_time_dijkstra(r_um: ndarray, z_um: ndarray, etch_rate_map: ndarray, r_is_logscaled: bool = True, connectivity: Literal[8, 16, 32] = 8, backend: Literal['auto', 'cpp', 'numba'] = 'auto') → ndarray[source]

Compute arrival times using Dijkstra’s algorithm.

Auto-selects the fastest available backend (C++ > Numba).

Parameters:

  • r_um (np.ndarray) – Radial coordinates (can be non-uniform/log-spaced)

  • z_um (np.ndarray) – Depth coordinates (linear)

  • etch_rate_map (np.ndarray) – Etch rate (speed) at each (z, r) point, shape (n_z, n_r)

  • r_is_logscaled (bool) – Whether r grid is log-spaced (informational only)

  • connectivity ({8, 16, 32}) – Number of neighbor connections per node. - 8: Standard (fastest, ~2% metrication error) - 16: Knight’s moves added (~1% error) - 32: Extended knight’s moves (~0.5% error)

  • backend ({"auto", "cpp", "numba"}) – Which backend to use. “auto” selects C++ if available, else Numba.

Returns:

arrival_time_map – Arrival time at each grid point, shape (n_z, n_r)

Return type:

np.ndarray

tracketch.wavefront.arrival_time_fmm(r_um: ndarray, z_um: ndarray, etch_rate_map: ndarray, r_is_logscaled: bool = True, theta_deg: float = 0.0, n_uniform_multiplier: int = 3)[source]

Calculate arrival time map using fast marching method.

Parameters:

  • r_um (np.ndarray) – Radial grid (may be log-scaled)

  • z_um (np.ndarray) – Depth grid (linear)

  • etch_rate_map (np.ndarray) – Etch rate at each (z, r) point, shape (n_z, n_r)

  • r_is_logscaled (bool, optional) – Whether r grid is log-scaled (default: True)

  • theta_deg (float, optional) – Track angle relative to surface normal in degrees (default: 0.0). theta=0: track perpendicular to surface (wavefront starts at z=0) theta>0: tilted track, wavefront starts along z = r * tan(theta)

  • n_uniform_multiplier (int, optional) – Multiplier for number of uniform grid points (default: 3). Higher values better preserve variation from log grid, especially at small radii. Set to 1 to use same number of points as original grid.

Returns:

arrival_time_map – Arrival time at each (z, r) point

Return type:

np.ndarray

Arrival-time map dispatcher and iso-time contour extraction.

tracketch.wavefront.utils_march.get_arrival_time_map(r_um: ndarray, z_um: ndarray, etch_rate_map: ndarray, method: str = 'dijkstra', r_is_logscaled: bool = True, theta_deg: float = 0.0, n_uniform_multiplier: int = 3, connectivity: Literal[8, 16, 32] = 8) → ndarray[source]

Calculate the arrival-time map using the specified method.

Parameters:

  • r_um (ndarray) – Radial coordinates in um.

  • z_um (ndarray) – Axial coordinates in um.

  • etch_rate_map (ndarray) – Etch rates in um/hr, shape (n_z, n_r).

  • method (str) – "dijkstra_cpp" (fastest), "dijkstra_numba", or "fmm" (supports tilted wavefronts).

  • r_is_logscaled (bool) – Whether the radial coordinates are log-spaced.

  • theta_deg (float) – Track angle relative to the surface normal in degrees.

  • n_uniform_multiplier (int) – FMM only – multiplier for the uniform-grid resolution.

  • connectivity ({8, 16, 32}) – Dijkstra only – neighbour connectivity.

Returns:

Arrival-time map in hours, same shape as etch_rate_map.

Return type:

ndarray

tracketch.wavefront.utils_march.get_iso_time_contour(arrival_time_map: ndarray, r_um: ndarray, z_um: ndarray, etching_time_h: float) → tuple[ndarray, ndarray][source]

Extract the iso-time contour from the arrival-time map.

Parameters:

  • arrival_time_map (ndarray, shape (n_z, n_r)) – Arrival times in hours.

  • r_um (ndarray) – Radial coordinates (may be log-spaced).

  • z_um (ndarray) – Depth coordinates (linear).

  • etching_time_h (float) – Target arrival time in hours.

Returns:

  • r_coords (ndarray)

  • z_coords (ndarray) – Coordinates along the longest contour segment.