eiko 0.8.0__py3-none-any.whl

eiko/__init__.py

@@ -0,0 +1,159 @@
+import os
+import sys
+
+# ---------------------------------------------------------
+# PATH RESOLUTION
+# ---------------------------------------------------------
+# Resolve paths once for the entire package.
+
+# Base directory where this file lives (Eiko/python/eiko/ or .../site-packages/eiko/).
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Option A: Path when the package is installed via pip (src is inside eiko/).
+SRC_DIR_INSTALLED = os.path.join(BASE_DIR, 'src')
+
+# Option B: Path when running locally during development (src is sibling to python/).
+SRC_DIR_DEV = os.path.abspath(os.path.join(BASE_DIR, '..', '..', 'src'))
+
+# Dynamically choose the folder that actually contains your files
+if os.path.exists(SRC_DIR_INSTALLED):
+    SRC_DIR = SRC_DIR_INSTALLED
+elif os.path.exists(SRC_DIR_DEV):
+    SRC_DIR = SRC_DIR_DEV
+else:
+    raise FileNotFoundError(
+        f"Could not find Eiko C++/CUDA source directory. Tried:\n"
+        f"1. {SRC_DIR_INSTALLED}\n"
+        f"2. {SRC_DIR_DEV}"
+    )
+
+# ---------------------------------------------------------
+# COMPILER FLAGS
+# ---------------------------------------------------------
+if sys.platform == "win32":
+    CXX_ARGS = ['/std:c++20', '/Zc:preprocessor', '/DNOMINMAX']
+    NVCC_ARGS = [
+        '-std=c++20', 
+        '-allow-unsupported-compiler', 
+        '-Xcompiler', '/Zc:preprocessor', 
+        '-Xcompiler', '/std:c++20', 
+        '-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH', 
+        '-DNOMINMAX', 
+        '--use_fast_math', 
+        '-arch=native'
+    ]
+else:
+    CXX_ARGS = ['-std=c++20', '-O3']
+    NVCC_ARGS = ['-std=c++20', '--use_fast_math', '-arch=native']
+
+
+EXTRA_INCLUDE_PATHS = [SRC_DIR]
+
+try:
+    from torch.utils.cpp_extension import CUDA_HOME
+    if CUDA_HOME:
+        cccl_base = os.path.join(CUDA_HOME, 'include', 'cccl')
+        if os.path.exists(cccl_base):
+            EXTRA_INCLUDE_PATHS.extend([
+                cccl_base,
+                os.path.join(cccl_base, 'thrust'),
+                os.path.join(cccl_base, 'libcudacxx', 'include'),
+                os.path.join(cccl_base, 'cub'),
+            ])
+except ImportError:
+    # Let individual wrappers handle the missing torch error cleanly
+    pass
+
+# ---------------------------------------------------------
+# PUBLIC API EXPORTS
+# ---------------------------------------------------------
+
+def _is_jax_array(obj):
+    """Safely checks if an object is a JAX array without importing JAX."""
+    return type(obj).__module__.startswith('jax')
+
+def eiko2d(u_init, f, v_init=None, dx=1.0, msfm=False, gated=False):
+    """
+    EIKO2D Computes the shortest time-of-flight in an arbitrary 2D medium.
+
+    Calculates the time-of-flight (u), given a slowness map (f = 1/c), and
+    initial conditions (u_init, initialized as infinity at unknown points).
+
+    EXAMPLE USAGES:
+        u = eiko2d(u_init, f)                                 # (Standard usage).
+        u = eiko2d(u_init, f, dx=0.5, msfm=True)              # (Named arguments).
+        u, v_out = eiko2d(u_init, f, v_init=advection_field)  # (Advection).
+
+    REQUIRED INPUTS:
+        u_init  - Initial conditions (known arrival times/delays).
+                  Shape: (H, W) for a single image, or (B, H, W) for a batch.
+        f       - Slowness. Can be (H, W) [broadcast to batch] or (B, H, W).
+
+    OPTIONAL INPUTS:
+        dx      - Input grid spacing. Default: 1.0.
+        v_init  - The initial advection field. Same size as u. Default: None (not used).
+        msfm    - Whether to enable Multi-Stencil Fast Marching (MSFM).
+                  Reduces bias along diagonal directions. Default: False.
+        gated   - Whether to enforce positive propagation along the first 
+                  data dimension. Speeds up computations. It is valid when 
+                  time only increases when moving axially. Default: False.
+
+    OUTPUTS:
+        u       - Computed arrival time (time-of-flight) map. Shape matches u_init.
+        v       - Output advection vectors (returned only if v_init was supplied).
+    """
+    if _is_jax_array(u_init):
+        from .eiko_jax import eiko2d as jax_eiko2d
+        return jax_eiko2d(u_init, f, v_init, dx, msfm, gated)
+    else:
+        from .eiko_torch import eiko2d as pt_eiko2d
+        return pt_eiko2d(u_init, f, v_init, dx, msfm, gated)
+
+def eiko3d(u_init, f, v_init=None, dx=1.0, msfm=False, gated=False):
+    """
+    EIKO3D Computes the shortest time-of-flight in an arbitrary 3D medium.
+
+    Calculates the time-of-flight (u), given a slowness map (f = 1/c), and
+    initial conditions (u_init, initialized as infinity at unknown points).
+
+    EXAMPLE USAGES:
+        u = eiko3d(u_init, f)                                 # (Standard usage).
+        u = eiko3d(u_init, f, dx=0.5, msfm=True)              # (Named arguments).
+        u, v_out = eiko3d(u_init, f, v_init=advection_field)  # (Advection).
+
+    REQUIRED INPUTS:
+        u_init  - Initial conditions (known arrival times/delays).
+                  Shape: (D, H, W) for a single volume, or (B, D, H, W) for a batch.
+        f       - Slowness. Can be (D, H, W) [broadcast to batch] or (B, D, H, W).
+
+    OPTIONAL INPUTS:
+        dx      - Input grid spacing. Default: 1.0.
+        v_init  - The initial advection field. Same size as u. Default: None (not used).
+        msfm    - Whether to enable Multi-Stencil Fast Marching (MSFM).
+                  Reduces bias along diagonal directions. Default: False.
+        gated   - Whether to enforce positive propagation along the first 
+                  data dimension. Speeds up computations. It is valid when 
+                  time only increases when moving axially. Default: False.
+
+    OUTPUTS:
+        u       - Computed arrival time (time-of-flight) map. Shape matches u_init.
+        v       - Output advection vectors (returned only if v_init was supplied).
+    """
+    if _is_jax_array(u_init):
+        from .eiko_jax import eiko3d as jax_eiko3d
+        return jax_eiko3d(u_init, f, v_init, dx, msfm, gated)
+    else:
+        from .eiko_torch import eiko3d as pt_eiko3d
+        return pt_eiko3d(u_init, f, v_init, dx, msfm, gated)
+
+eiko = eiko2d
+
+from .animate_eikonal import animate_eikonal
+
+# Define what is imported when a user runs `from eiko import *`.
+__all__ = ['eiko', 'eiko3d', 'animate_eikonal']
+
+# Define the version number, so it is easily accessible.
+from importlib.metadata import version
+__version__ = version("eiko")
+

eiko/animate_eikonal.py

@@ -0,0 +1,201 @@
+import sys
+
+def animate_eikonal(u, v=1.0, color_map='gray', video_filename='', 
+                    title='EIKONAL WAVEFRONT', pulse_width=80.0, 
+                    speed=0.5, overlay=None, outline=None, 
+                    style='real', render_mode='slice'):
+    """
+    ANIMATE_EIKONAL Visualizes Eikonal equation travel times as an animated wave.
+    """
+    
+    # Required dependencies to run the visualization.
+    try:
+        import numpy as np
+        import matplotlib.pyplot as plt
+        from matplotlib.animation import FuncAnimation
+    except ImportError:
+        raise ImportError(
+            "The animation utility requires 'numpy' and 'matplotlib'. "
+            "Please install it using: pip install \"eiko[examples]\""
+        )
+    
+    # Optional dependency for contours/isosurfaces
+    try:
+        from skimage import measure
+        HAS_SKIMAGE = True
+    except ImportError:
+        HAS_SKIMAGE = False
+        print("[Eiko] Warning: 'scikit-image' not found. Contours, overlays, and 3D isosurfaces will be disabled.", file=sys.stderr)
+    
+    # --- Data Type & Device Checks (e.g., PyTorch -> Numpy) ---
+    if hasattr(u, 'cpu'): u = u.detach().cpu().numpy()
+    if hasattr(v, 'cpu'): v = v.detach().cpu().numpy()
+    if hasattr(overlay, 'cpu'): overlay = overlay.detach().cpu().numpy()
+    if hasattr(outline, 'cpu'): outline = outline.detach().cpu().numpy()
+
+    is_3d = u.ndim == 3
+    
+    # --- Safe Data Initialization ---
+    valid_u = u[np.isfinite(u)]
+    if len(valid_u) == 0:
+        print("Warning: Travel time field u is completely Inf. Nothing to animate.", file=sys.stderr)
+        return
+
+    maxv = np.max(valid_u)
+    u_safe = np.copy(u).astype(np.float32)
+    u_safe[np.isinf(u_safe)] = maxv + pulse_width
+
+    # Handle Permutation and setup for 3D
+    if is_3d:
+        # In Numpy, typical 3D ordering is [Z, Y, X]. We keep it as is, 
+        # but map slices to dimensions.
+        Nz, Ny, Nx = u.shape
+        v_anim = v
+    else:
+        v_anim = v
+        Ny, Nx = u.shape
+
+    # --- Figure & Axis Setup ---
+    fig = plt.figure(figsize=(12, 8), facecolor='white')
+    
+    if style.lower() == 'real':
+        clim = [-1, 1]
+    elif style.lower() == 'abs':
+        clim = [0, 1]
+    elif style.lower() == 'db':
+        clim = [-30, 0]
+        db_ref = 20 * np.log10(1.0 + 1e-12) # +1e-12 prevents log10(0)
+
+    # Setup Render Objects
+    if is_3d:
+        # 3D Axes
+        ax = fig.add_subplot(111, projection='3d')
+        ax.set_title(title, fontsize=16, fontweight='bold')
+        ax.set_xlabel('X (Lateral)')
+        ax.set_ylabel('Y (Elevation)')
+        ax.set_zlabel('Z (Depth)')
+        ax.invert_zaxis() # Ultrasound convention: Depth increases downwards
+        
+        # We'll handle the 3D drawing inside the update loop, but prepare indices
+        sz, sy, sx = Nz//2, Ny//2, Nx//2
+        iso_val = -6 if style.lower() == 'db' else 0.5
+        
+    else:
+        # 2D Rendering
+        ax = fig.add_subplot(111)
+        ax.set_title(title, fontsize=16, fontweight='bold', fontname='serif')
+        ax.set_xlabel('X (Lateral)', fontsize=14, fontname='serif')
+        ax.set_ylabel('Z (Depth)', fontsize=14, fontname='serif')
+        
+        # Initialize empty image
+        im = ax.imshow(np.zeros((Ny, Nx)), vmin=clim[0], vmax=clim[1], cmap=color_map, aspect='equal')
+        
+        # In matplotlib, imshow puts (0,0) at top-left, which inherently matches 
+        # the ultrasound convention (Depth going down).
+        
+        if outline is not None and HAS_SKIMAGE:
+            ax.contour(outline, levels=[0.5], colors='k', linewidths=2)
+            
+        if overlay is not None and HAS_SKIMAGE:
+            contours = measure.find_contours(overlay, 0.5)
+            for contour in contours:
+                # contour is (row, col) -> (Y, X)
+                ax.fill(contour[:, 1], contour[:, 0], color='red', alpha=0.3, edgecolor='r', linewidth=2.5)
+
+    freq = 6 * np.pi
+    
+    # --- Animation Update Logic ---
+    def compute_frame(t):
+        """Calculates the analytical wave packet for a given time step."""
+        diff_t = u_safe - t
+        valid_mask = np.abs(diff_t) <= (pulse_width / 2.0)
+        
+        img = np.zeros_like(u_safe, dtype=np.complex64)
+        
+        if np.any(valid_mask):
+            d_valid = diff_t[valid_mask]
+            
+            # Analytical Hanning envelope
+            envelope = 0.5 * (1.0 + np.cos(2 * np.pi * d_valid / pulse_width))
+            
+            # Apply envelope and carrier phase
+            img[valid_mask] = envelope * np.exp(1j * freq * d_valid / pulse_width)
+            
+        img = img * v_anim
+        
+        # Apply Style
+        if style.lower() == 'real':
+            final_img = np.real(img)
+        elif style.lower() == 'abs':
+            final_img = np.abs(img)
+        elif style.lower() == 'db':
+            final_img = 20 * np.log10(np.abs(img) + 1e-12) - db_ref
+            
+        return final_img
+
+    # time_steps = np.arange(0, maxv + pulse_width, speed)
+    start_time = np.min(u_safe)
+    time_steps = np.arange(start_time, maxv + pulse_width, speed)
+
+    # Note: Global references for 3D plot collections so they can be removed/updated
+    frame_objs = [] 
+
+    
+    def update(frame_idx):
+        t = time_steps[frame_idx]
+        final_img = compute_frame(t)
+        
+        #if frame_idx % 20 == 0: # Check every 20 frames to avoid spamming
+        #    print(f"Frame {frame_idx} (t={t:.1f}): Data min/max = {final_img.min():.3f} / {final_img.max():.3f}")
+
+        if is_3d:
+            # Clear all collections (slices or isosurfaces) from the previous frame.
+            for coll in ax.collections:
+                coll.remove()
+            
+            # Define levels to force contouring (20 steps is usually enough for smooth waves)
+            levels = np.linspace(clim[0], clim[1], 20)
+            if render_mode == 'slice':
+                # Re-draw the slices.
+                X, Y = np.meshgrid(np.arange(Nx), np.arange(Ny))
+                ax.contourf(X, Y, final_img[sz, :, :], zdir='z', offset=sz, cmap=color_map, vmin=clim[0], vmax=clim[1], levels=levels)
+                
+                Y, Z = np.meshgrid(np.arange(Ny), np.arange(Nz))
+                ax.contourf(final_img[:, :, sx], Y, Z, zdir='x', offset=sx, cmap=color_map, vmin=clim[0], vmax=clim[1], levels=levels)
+                
+                X, Z = np.meshgrid(np.arange(Nx), np.arange(Nz))
+                ax.contourf(X, final_img[:, sy, :], Z, zdir='y', offset=sy, cmap=color_map, vmin=clim[0], vmax=clim[1], levels=levels)
+                
+            elif render_mode == 'isosurface' and HAS_SKIMAGE:
+                try:
+                    # Use marching cubes to find the isosurface.
+                    verts, faces, _, _ = measure.marching_cubes(final_img, level=iso_val)
+                    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+                    mesh = Poly3DCollection(verts[faces], alpha=0.7)
+                    mesh.set_facecolor([0.2, 0.6, 1.0])
+                    mesh.set_edgecolor('none')
+                    ax.add_collection3d(mesh)
+                except ValueError:
+                    pass # Fails cleanly if the volume doesn't contain the isovalue yet
+            
+            return ax.collections
+        else:
+            im.set_data(final_img)
+            return [im]
+
+    # --- Run / Export Animation ---
+    # Blit=True massively speeds up 2D rendering by only redrawing changed pixels.
+    # Blit is not well-supported for matplotlib 3D axes.
+    anim = FuncAnimation(fig, update, frames=len(time_steps), 
+                                   interval=1000/60, blit=not is_3d)
+
+    if video_filename:
+        print(f"Exporting video to {video_filename} (this may take a moment)...")
+        # Ensure you have ffmpeg installed on your Ubuntu machine (sudo apt install ffmpeg)
+        writer = animation.FFMpegWriter(fps=60, bitrate=2000)
+        anim.save(video_filename, writer=writer)
+        print("Video export complete.")
+    else:
+        plt.show()
+
+    return anim

eiko/eiko_jax.py

@@ -0,0 +1,326 @@
+import os
+import sys
+import struct
+from functools import partial
+from torch.utils.cpp_extension import load, _get_build_directory
+
+try:
+    import jax
+    import pybind11
+    import jaxlib
+except ImportError as e:
+    raise ImportError(
+        f"\n[Eiko] JAX bindings require 'jax', 'jaxlib', and 'pybind11' to be installed.\n"
+        f" Please install the required environment via: pip install \"eiko[jax]\"\n"
+    ) from e
+
+import jax.numpy as jnp
+from jax import core
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax.interpreters import mlir
+
+# Robust Primitive resolution for modern JAX.
+if not hasattr(core, "Primitive"):
+    from jax._src.core import Primitive
+    core.Primitive = Primitive
+
+# Version-agnostic xla_client import.
+try:
+    from jaxlib import xla_client
+except ImportError:
+    from jax.lib import xla_client
+
+# Version-agnostic custom_call import.
+#try:
+#    from jaxlib.hlo_helpers import custom_call
+#except ImportError:
+#    from jax.interpreters.mlir import custom_call
+
+from eiko import SRC_DIR, CXX_ARGS, NVCC_ARGS, EXTRA_INCLUDE_PATHS
+
+# ---------------------------------------------------------
+# JIT COMPILATION & LOADING
+# ---------------------------------------------------------
+build_dir = _get_build_directory('eiko_jax_impl', verbose=False)
+is_cached = os.path.exists(build_dir) and len(os.listdir(build_dir)) > 0
+
+if not is_cached:
+    print("[Eiko] First-time JAX initialization: JIT Compiling CUDA kernels for your GPU... (This may take a minute)")
+    sys.stdout.flush()
+
+jax_source = os.path.join(SRC_DIR, 'bindings', 'jax_bindings.cu')
+jax_includes = EXTRA_INCLUDE_PATHS + [pybind11.get_include()]
+
+_fim_jax_impl = load(
+    name="eiko_jax_impl",
+    sources=[jax_source],
+    extra_cflags=CXX_ARGS,
+    extra_cuda_cflags=NVCC_ARGS,
+    extra_include_paths=jax_includes,
+    verbose=False
+)
+
+for name, target in _fim_jax_impl.registrations().items():
+    xla_client.register_custom_call_target(name, target, platform="gpu")
+
+# =========================================================
+# 1. JAX PRIMITIVE DEFINITION & MLIR LOWERING
+# =========================================================
+_fim_prim = core.Primitive("jax_fim_solve")
+_fim_prim.multiple_results = False
+_fim_prim.def_impl(partial(xla.apply_primitive, _fim_prim))
+
+def _fim_abstract_eval(*args, opaque_data, out_shape, out_dtype):
+    return core.ShapedArray(out_shape, out_dtype)
+
+_fim_prim.def_abstract_eval(_fim_abstract_eval)
+
+def _build_custom_call_agnostic(call_target_name, result_types, operands, 
+                                operand_layouts, result_layouts, backend_config):
+    """
+    A custom call builder that checks for legacy wrappers 
+    and falls back to raw MLIR node generation for modern JAX.
+    """
+    # 1. Try mid-era JAX wrapper (JAX >= 0.4.15 and < 0.4.30)
+    try:
+        from jaxlib.hlo_helpers import custom_call
+        return custom_call(
+            call_target_name,
+            result_types=result_types,
+            operands=operands,
+            operand_layouts=operand_layouts,
+            result_layouts=result_layouts,
+            backend_config=backend_config
+        )
+    except ImportError:
+        pass
+        
+    # 2. Try legacy JAX wrapper (JAX < 0.4.15)
+    try:
+        from jax.interpreters.mlir import custom_call
+        return custom_call(
+            call_target_name,
+            result_types=result_types,
+            operands=operands,
+            operand_layouts=operand_layouts,
+            result_layouts=result_layouts,
+            backend_config=backend_config
+        )
+    except (ImportError, AttributeError):
+        pass
+        
+    # 3. Modern JAX (>= 0.4.30) where helpers are completely removed.
+    import jaxlib.mlir.ir as ir
+    try:
+        from jaxlib.mlir.dialects import mhlo as hlo
+    except ImportError:
+        from jaxlib.mlir.dialects import hlo
+
+    def _layout_attr(layouts):
+        if layouts is None: 
+            return None
+        import numpy as np
+        attr_list = []
+        for layout in layouts:
+            arr = np.array(layout, dtype=np.int64)
+            
+            # Define the element type as MLIR's 'index' type.
+            index_type = ir.IndexType.get()
+            
+            # Define the 1D tensor shape explicitly using the index type.
+            tensor_type = ir.RankedTensorType.get(arr.shape, index_type)
+            
+            # Create the attribute with the enforced tensor type.
+            attr = ir.DenseIntElementsAttr.get(arr, type=tensor_type)
+            attr_list.append(attr)
+            
+        return ir.ArrayAttr.get(attr_list)
+
+    kwargs = {
+        "call_target_name": ir.StringAttr.get(call_target_name),
+        "has_side_effect": ir.BoolAttr.get(False),
+        "api_version": ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 2),
+        "called_computations": ir.ArrayAttr.get([]),
+    }
+    
+    if backend_config is not None:
+        kwargs["backend_config"] = ir.StringAttr.get(backend_config)
+
+    if operand_layouts is not None:
+        kwargs["operand_layouts"] = _layout_attr(operand_layouts)
+        
+    if result_layouts is not None:
+        kwargs["result_layouts"] = _layout_attr(result_layouts)
+
+    return hlo.CustomCallOp(result_types, operands, **kwargs)
+
+# MLIR lowering rule: The bridge between JAX's Python graph and XLA C++.
+def _fim_lowering(ctx, *args, opaque_data, out_shape, out_dtype):
+    import jaxlib.mlir.ir as ir
+    
+    # Convert numpy dtype to MLIR IR type.
+    tensor_type = ir.RankedTensorType.get(
+        out_shape, mlir.dtype_to_ir_type(out_dtype)
+    )
+    
+    operand_layouts = [tuple(range(arg.type.rank)[::-1]) for arg in args]
+    result_layouts = [tuple(range(len(out_shape))[::-1])]
+        
+    # MLIR custom call builder.
+    call = _build_custom_call_agnostic(
+        "jax_fim_solve",
+        result_types=[tensor_type],
+        operands=args,
+        operand_layouts=operand_layouts,
+        result_layouts=result_layouts,
+        backend_config=opaque_data
+    )
+    return call.results
+
+
+mlir.register_lowering(_fim_prim, _fim_lowering, platform="gpu")
+
+# =========================================================
+# BATCHING RULE
+# =========================================================
+def _fim_batch_rule(batched_args, batch_dims, *, opaque_data, out_shape, out_dtype):
+    # Unpack original configuration.
+    unpacked = list(struct.unpack('=iiiifiiiiiii', opaque_data))
+    
+    # Find the new batch size. batched_args[0] is u_init.
+    bdim = batch_dims[0] if batch_dims[0] is not None else 0
+    new_batch_axis_size = batched_args[0].shape[bdim]
+    
+    # Update the batch_size (index 3 in our struct)
+    unpacked[3] = unpacked[3] * new_batch_axis_size
+    new_opaque_data = struct.pack('=iiiifiiiiiii', *unpacked)
+    
+    # Push the batch dimension to axis 0 for all batched arguments.
+    aligned_args = [
+        batching.moveaxis(arg, d, 0) if d is not None else arg 
+        for arg, d in zip(batched_args, batch_dims)
+    ]
+    
+    # Prepend the new batch dimension to the output shape.
+    new_out_shape = (new_batch_axis_size,) + tuple(out_shape)
+    
+    out = _fim_prim.bind(
+        *aligned_args, 
+        opaque_data=new_opaque_data, 
+        out_shape=new_out_shape, 
+        out_dtype=out_dtype
+    )
+    
+    # Tell JAX that the batched dimension of the output is at axis 0.
+    return out, 0
+
+batching.primitive_batchers[_fim_prim] = _fim_batch_rule
+
+# =========================================================
+# 2. BASE XLA CUSTOM CALL
+# =========================================================
+def _fim_custom_call(u_init, f, v, dx, msfm, is_3d, gated_x, is_backward, tof=None):
+    u_init = jnp.asarray(u_init)
+    f = jnp.asarray(f)
+    
+    has_v = v is not None
+    has_tof = tof is not None
+    operands = [u_init, f]
+
+    if has_v:
+        v = jnp.asarray(v)
+        operands.append(v)
+        
+    if is_backward and has_tof:
+        tof = jnp.asarray(tof)
+        operands.append(tof)
+
+    batch_size = u_init.shape[0]
+    if is_3d:
+        depth, height, width = u_init.shape[-3:]
+    else:
+        depth = 1
+        height, width = u_init.shape[-2:]
+    
+    if f.ndim == u_init.ndim - 1:
+        broadcast_f = True
+    else:
+        broadcast_f = (f.shape[0] == 1 and batch_size > 1)
+    
+    opaque_data = struct.pack(
+        '=iiiifiiiiiii', 
+        width, height, depth, batch_size, float(dx), 
+        int(is_3d), int(is_backward), int(msfm), int(has_v), 
+        int(broadcast_f), int(gated_x), int(has_tof)
+    )
+
+    return _fim_prim.bind(
+        *operands,
+        opaque_data=opaque_data,
+        out_shape=u_init.shape,
+        out_dtype=u_init.dtype
+    )
+
+# =========================================================
+# 3. VJP DEFINITION (Autograd support)
+# =========================================================
+@partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5, 6))
+def _solve_eikonal_base(u_init, f, v, dx, msfm, is_3d, gated_x):
+    if is_3d is None:
+        is_3d = u_init.ndim >= 4 and u_init.shape[-3] > 1
+        
+    return _fim_custom_call(u_init, f, v, dx, msfm, is_3d, gated_x, is_backward=False)
+
+def solve_eikonal_fwd(u_init, f, v, dx, msfm, is_3d, gated_x):
+    if is_3d is None:
+        is_3d = u_init.ndim >= 4 and u_init.shape[-3] > 1
+        
+    u_out = _fim_custom_call(u_init, f, v, dx, msfm, is_3d, gated_x, is_backward=False)
+    
+    res = (u_out, f)
+    return u_out, res
+
+def solve_eikonal_bwd(v, dx, msfm, is_3d, gated_x, res, grad_u):
+    u_out, f = res
+    lambda_init = jnp.zeros_like(u_out)
+    
+    lambda_adj = _fim_custom_call(
+        lambda_init, grad_u, None, dx, msfm, is_3d, gated_x, 
+        is_backward=True, tof=u_out
+    )
+    
+    grad_u_init = lambda_adj
+    grad_f = lambda_adj * f * dx * dx
+    
+    if f.ndim == lambda_adj.ndim - 1:
+        grad_f = jnp.sum(grad_f, axis=0)
+    elif f.ndim == lambda_adj.ndim and f.shape[0] == 1 and lambda_adj.shape[0] > 1:
+        grad_f = jnp.sum(grad_f, axis=0, keepdims=True)
+    
+    return grad_u_init, grad_f
+
+_solve_eikonal_base.defvjp(solve_eikonal_fwd, solve_eikonal_bwd)
+
+# =========================================================
+# 4. PUBLIC API
+# =========================================================
+def eiko2d(u_init, f, v_init=None, dx=1.0, msfm=False, gated=False):
+    if u_init.ndim > 3:
+        raise ValueError(f"eiko2d expects a 2D or batched 2D grid (max 3 dims), got {u_init.ndim} dims.")
+    
+    out = _solve_eikonal_base(u_init, f, v_init, dx, msfm, False, gated)
+    if v_init is not None:
+        return out, v_init
+    return out
+
+def eiko3d(u_init, f, v_init=None, dx=1.0, msfm=False, gated=False):
+    if u_init.ndim < 3:
+        raise ValueError(f"eiko3d expects a 3D or batched 3D grid (min 3 dims), got {u_init.ndim} dims.")
+        
+    out = _solve_eikonal_base(u_init, f, v_init, dx, msfm, True, gated)
+    if v_init is not None:
+        return out, v_init
+    return out
+
+eiko = eiko2d