FastSIMUS 0.0.1__py3-none-any.whl

fast_simus/__init__.py

@@ -0,0 +1,33 @@
+"""FastSIMUS - Fast Simulator for Medical Ultrasound based on SIMUS/MUST."""
+
+from fast_simus.medium_params import MediumParams
+from fast_simus.pfield import PfieldPlan, PfieldStrategy, pfield, pfield_compute, pfield_precompute
+from fast_simus.simus import SimusPlan, SimusResult, SimusStrategy, simus, simus_compute, simus_precompute
+from fast_simus.transducer_params import BaffleType, TransducerParams
+from fast_simus.tx_delay import (
+    diverging_wave,
+    focused,
+    plane_wave,
+)
+from fast_simus.utils.geometry import element_positions
+
+__all__ = [
+    "BaffleType",
+    "MediumParams",
+    "PfieldPlan",
+    "PfieldStrategy",
+    "SimusPlan",
+    "SimusResult",
+    "SimusStrategy",
+    "TransducerParams",
+    "diverging_wave",
+    "element_positions",
+    "focused",
+    "pfield",
+    "pfield_compute",
+    "pfield_precompute",
+    "plane_wave",
+    "simus",
+    "simus_compute",
+    "simus_precompute",
+]

fast_simus/_pfield_math.py

@@ -0,0 +1,261 @@
+"""Physics helpers for pfield computation.
+
+Pure Array API functions for geometry, phase initialization, frequency
+selection, and obliquity. No loop structure or backend-specific code.
+"""
+
+from __future__ import annotations
+
+from math import ceil, log, pi
+from typing import NamedTuple
+
+from beartype import beartype as typechecker
+from jaxtyping import Complex, Float, jaxtyped
+
+from fast_simus.spectrum import probe_spectrum, pulse_spectrum
+from fast_simus.transducer_params import BaffleType
+from fast_simus.utils._array_api import Array, _ArrayNamespace
+
+# Conversion factor: Nepers to dB -- 20/log(10) ≈ 8.6859
+# Shared by Array API path and Metal kernel wrapper.
+NEPER_TO_DB = 20.0 / log(10.0)
+
+
+class _FrequencyPlan(NamedTuple):
+    """Frequency sampling plan for pfield computation.
+
+    Attributes:
+        selected_freqs: Selected frequencies, shape
+        pulse_spectrum: Pulse spectrum at selected frequencies, shape
+        probe_spectrum: Probe spectrum at selected frequencies, shape
+        freq_step: Frequency step in Hz
+    """
+
+    selected_freqs: Float[Array, " n_frequencies"]
+    pulse_spectrum: Complex[Array, " n_frequencies"]
+    probe_spectrum: Float[Array, " n_frequencies"]
+    freq_step: float
+
+
+@jaxtyped(typechecker=typechecker)
+def _subelement_centroids(
+    element_width: float,
+    n_sub: int,
+    theta_e: Float[Array, " n_elements"],
+    xp: _ArrayNamespace,
+) -> Float[Array, "n_elements n_sub 2"]:
+    """Compute sub-element centroid positions relative to element centers.
+
+    Args:
+        element_width: Element width in meters.
+        n_sub: Number of sub-elements per element.
+        theta_e: Element angular positions in radians.
+        xp: Array namespace.
+
+    Returns:
+        Sub-element offsets with shape (n_elements, n_sub, 2) where [..., 0]
+        is lateral (x) and [..., 1] is axial (z).
+    """
+    seg_length = element_width / n_sub
+    seg_offsets = xp.asarray([-element_width / 2.0 + seg_length / 2.0 + i * seg_length for i in range(n_sub)])
+    # Broadcasting: (n_sub,) -> (1, n_sub), (n_elements,) -> (n_elements, 1)
+    seg_offsets_2d = xp.reshape(seg_offsets, (1, n_sub))
+    cos_theta = xp.cos(theta_e)[:, None]
+    sin_neg_theta = xp.sin(-theta_e)[:, None]
+    subelement_dx = seg_offsets_2d * cos_theta
+    subelement_dz = seg_offsets_2d * sin_neg_theta
+    return xp.stack([subelement_dx, subelement_dz], axis=-1)
+
+
+@jaxtyped(typechecker=typechecker)
+def _distances_and_angles(
+    points: Float[Array, "*batch 2"],
+    subelement_offsets: Float[Array, "n_elements n_sub 2"],
+    element_pos: Float[Array, "n_elements 2"],
+    theta_e: Float[Array, " n_elements"],
+    speed_of_sound: float,
+    freq_center: float,
+    xp: _ArrayNamespace,
+) -> tuple[
+    Float[Array, "*batch n_elements n_sub"],
+    Float[Array, "*batch n_elements n_sub"],
+    Float[Array, "*batch n_elements n_sub"],
+]:
+    """Compute distances and angles from grid points to sub-elements.
+
+    Args:
+        points: Grid point positions.
+        subelement_offsets: Sub-element offsets.
+        element_pos: Element positions.
+        theta_e: Element angular positions.
+        speed_of_sound: Speed of sound in m/s.
+        freq_center: Center frequency in Hz.
+        xp: Array namespace.
+
+    Returns:
+        Tuple of (distances, sin_theta, theta_arr):
+        - distances: Distances from grid points to sub-elements.
+        - sin_theta: Sine of angles relative to element normal.
+        - theta_arr: Angles relative to element normal.
+    """
+    delta: Float[Array, "*batch n_elements n_sub xz=2"] = (
+        points[..., None, None, :] - subelement_offsets - element_pos[:, None, :]
+    )
+    delta_x = delta[..., 0]
+    delta_z = delta[..., 1]
+    dist_squared = delta_x**2 + delta_z**2
+    distances = xp.sqrt(dist_squared)
+
+    # Distances with clipping (use unclipped sqrt for angle computation)
+    min_distance = xp.asarray(speed_of_sound / freq_center / 2.0)
+    distances_clipped = xp.where(distances < min_distance, min_distance, distances)
+
+    # Angle relative to element normal
+    _div_eps = xp.asarray(1e-16)  # Numerical stability for division
+    theta_arr = xp.asin((delta_x + _div_eps) / (distances + _div_eps)) - theta_e[:, None]
+    sin_theta = xp.sin(theta_arr)
+
+    return distances_clipped, sin_theta, theta_arr
+
+
+def _select_frequencies(
+    fc: float,
+    bandwidth: float,
+    tx_n_wavelengths: float,
+    db_thresh: float,
+    max_freq_step: float,
+    xp: _ArrayNamespace,
+) -> _FrequencyPlan:
+    """Select frequency samples for pfield computation.
+
+    Args:
+        fc: Center frequency in Hz.
+        bandwidth: Fractional bandwidth.
+        tx_n_wavelengths: Number of wavelengths in TX pulse.
+        db_thresh: Threshold in dB for frequency component selection.
+        max_freq_step: Upper bound for frequency step.
+        xp: Array namespace.
+
+    Returns:
+        FrequencyPlan with selected frequencies and spectra.
+    """
+    # Frequency samples
+    n_freq = int(2 * ceil(fc / max_freq_step) + 1)
+    frequencies = xp.linspace(0, 2 * fc, n_freq)
+    freq_step = float(frequencies[1])
+
+    # Keep only significant components (dB threshold)
+    angular_freqs_all = xp.asarray(2.0 * pi) * frequencies
+    spectrum_magnitude = xp.abs(
+        pulse_spectrum(angular_freqs_all, fc, tx_n_wavelengths) * probe_spectrum(angular_freqs_all, fc, bandwidth)
+    )
+    gain_db = 20.0 * xp.log10(xp.asarray(1e-200) + spectrum_magnitude / xp.max(spectrum_magnitude))
+    above_threshold = gain_db > db_thresh
+    idx_first, idx_last = _first_last_true(xp, above_threshold)
+
+    selected_freqs = frequencies[idx_first : idx_last + 1]
+
+    angular_freqs_sel = xp.asarray(2.0 * pi) * selected_freqs
+    pulse_spect = pulse_spectrum(angular_freqs_sel, fc, tx_n_wavelengths)
+    probe_spect = probe_spectrum(angular_freqs_sel, fc, bandwidth)
+
+    return _FrequencyPlan(selected_freqs, pulse_spect, probe_spect, freq_step)
+
+
+@jaxtyped(typechecker=typechecker)
+def _obliquity_factor(
+    theta_arr: Float[Array, "*batch n_elements n_sub"],
+    baffle: BaffleType | float,
+    xp: _ArrayNamespace,
+) -> Float[Array, "*batch n_elements n_sub"]:
+    """Compute obliquity factor based on baffle type.
+
+    Args:
+        theta_arr: Angles relative to element normal.
+        baffle: Baffle type or impedance ratio.
+        xp: Array namespace.
+
+    Returns:
+        Obliquity factor.
+    """
+    non_rigid_baffle = baffle != BaffleType.RIGID
+    _horizon_floor = xp.asarray(1e-16)  # Near-zero for beyond-hemisphere angles
+
+    if non_rigid_baffle:
+        if baffle == BaffleType.SOFT:
+            obliquity_factor = xp.cos(theta_arr)
+        else:
+            cos_th = xp.cos(theta_arr)
+            obliquity_factor = cos_th / (cos_th + float(baffle))
+    else:
+        obliquity_factor = xp.ones(theta_arr.shape)
+
+    obliquity_factor = xp.where(
+        xp.abs(theta_arr) >= xp.asarray(pi / 2),
+        _horizon_floor,
+        obliquity_factor,
+    )
+
+    return obliquity_factor
+
+
+@jaxtyped(typechecker=typechecker)
+def _init_exponentials(
+    freq_start: float | Float[Array, ""],
+    speed_of_sound: float,
+    attenuation: float,
+    distances: Float[Array, "*batch n_elements n_sub"],
+    obliquity_factor: Float[Array, "*batch n_elements n_sub"],
+    freq_step: float | Float[Array, ""],
+    xp: _ArrayNamespace,
+) -> tuple[
+    Complex[Array, "*batch n_elements n_sub"],
+    Complex[Array, "*batch n_elements n_sub"],
+]:
+    """Initialize exponential arrays for frequency loop.
+
+    Args:
+        freq_start: Initial frequency in Hz (scalar or 0-d array).
+        speed_of_sound: Speed of sound in m/s.
+        attenuation: Attenuation coefficient in dB/cm/MHz.
+        distances: Distances.
+        obliquity_factor: Obliquity factor.
+        freq_step: Frequency step in Hz (scalar or 0-d array).
+        xp: Array namespace.
+
+    Returns:
+        Tuple of (phase_decay, phase_decay_step):
+        - phase_decay: Initial complex exponential array.
+        - phase_decay_step: Complex exponential increment per frequency step.
+    """
+    wavenumber_init = 2.0 * pi * freq_start / speed_of_sound
+    attenuation_wavenum = attenuation / NEPER_TO_DB * freq_start / 1e6 * 1e2
+
+    # exp(-kwa*distances + 1j*mod(kw*distances, 2pi))
+    kw0_r = xp.asarray(wavenumber_init) * distances
+    two_pi = xp.asarray(2.0 * pi)
+    phase_mod = kw0_r - two_pi * xp.floor(kw0_r / two_pi)
+    phase_decay = xp.exp(xp.asarray(-attenuation_wavenum) * distances + xp.asarray(1j) * phase_mod)
+
+    wavenumber_step = 2.0 * pi * freq_step / speed_of_sound
+    attenuation_step = attenuation / NEPER_TO_DB * freq_step / 1e6 * 1e2
+    phase_decay_step = xp.exp(xp.asarray(-attenuation_step + 1j * wavenumber_step) * distances)
+
+    # Incorporate obliquity / sqrt(distances) (2D, no elevation)
+    phase_decay = phase_decay * obliquity_factor / xp.sqrt(distances)
+
+    return phase_decay, phase_decay_step
+
+
+def _first_last_true(xp: _ArrayNamespace, mask: Array) -> tuple[int, int]:
+    """Find first and last True index in 1D boolean array. JAX-compatible (no nonzero)."""
+    n = mask.shape[0]
+    if n == 0:
+        return 0, 0
+    # Cast to int: argmax on bool not allowed by array_api_strict
+    mask_int = xp.asarray(mask, dtype=xp.int32)
+    first = int(xp.argmax(mask_int))
+    if int(xp.max(mask_int)) == 0:
+        return 0, 0
+    last = n - 1 - int(xp.argmax(mask_int[::-1]))
+    return first, last

fast_simus/_pfield_strategies.py

@@ -0,0 +1,203 @@
+"""Loop drivers for pfield frequency sweep (Layer 3).
+
+Each driver iterates the same _freq_step_body() using a different mechanism:
+
+- _freq_outer_python: Python for-loop (NumPy/CuPy, constant memory)
+- _pfield_freq_vectorized: tensor broadcast (small grids only)
+- _freq_outer_scan: JAX lax.scan for O(1) compilation cost
+"""
+
+from __future__ import annotations
+
+from math import pi
+
+import array_api_extra as xpx
+from beartype import beartype as typechecker
+from jaxtyping import Bool, Complex, Float, jaxtyped
+
+from fast_simus.utils._array_api import Array, _ArrayNamespace
+
+
+def _freq_step_body(
+    phase: Complex[Array, " *grid n_sources"],
+    phase_step: Complex[Array, " *grid n_sources"],
+    spectrum_k: complex | Array,
+    xp: _ArrayNamespace,
+    *,
+    directivity_k: Float[Array, " *grid n_sources"] | None = None,
+) -> tuple[Complex[Array, " *grid n_sources"], Float[Array, " *grid"]]:
+    """One frequency step: geometric update, source contraction, spectrum weight.
+
+    Single source of truth for per-frequency math. Source points are
+    already flattened (n_elements * n_sub) with 1/n_sub absorbed.
+
+    Args:
+        phase: Current phase state (geometric progression).
+        phase_step: Per-step multiplier for geometric progression.
+        spectrum_k: Combined pulse*probe spectrum weight for this frequency.
+        xp: Array namespace.
+        directivity_k: Per-source directivity for this frequency (optional).
+
+    Returns:
+        Tuple of (updated_phase, rp_k) where rp_k = |P_k|^2 at this frequency.
+    """
+    if directivity_k is not None:
+        phase_weighted = phase * directivity_k
+    else:
+        phase_weighted = phase
+    pressure_k = spectrum_k * xp.sum(phase_weighted, axis=-1)
+    phase = phase * phase_step
+    return phase, xp.real(pressure_k * xp.conj(pressure_k))
+
+
+def _freq_outer_python(
+    phase_decay_init: Complex[Array, " *grid n_sources"],
+    phase_decay_step: Complex[Array, " *grid n_sources"],
+    is_out: Bool[Array, " *grid"],
+    wavenumbers: Float[Array, " n_freq"],
+    pulse_spect: Complex[Array, " n_freq"],
+    probe_spect: Float[Array, " n_freq"],
+    seg_length: float,
+    sin_theta: Float[Array, " *grid n_sources"],
+    full_frequency_directivity: bool,
+    xp: _ArrayNamespace,
+) -> Float[Array, " *grid"]:
+    """Python for-loop driver for NumPy/CuPy: constant O(grid * sources) memory.
+
+    Iterates one frequency at a time using _freq_step_body, accumulating
+    |P_k|^2 into the result. Peak memory is independent of n_freq.
+    """
+    spectra = pulse_spect * probe_spect
+    n_freq = int(wavenumbers.shape[0])
+    zero = xp.asarray(0.0)
+
+    phase = phase_decay_init
+    rp = xp.zeros(phase.shape[:-1])
+
+    if full_frequency_directivity:
+        for k in range(n_freq):
+            sinc_arg = wavenumbers[k] * seg_length / 2.0 * sin_theta / pi
+            directivity_k = xpx.sinc(sinc_arg, xp=xp)
+            phase, rp_k = _freq_step_body(phase, phase_decay_step, spectra[k], xp, directivity_k=directivity_k)
+            rp = rp + xp.where(is_out, zero, rp_k)
+    else:
+        for k in range(n_freq):
+            phase, rp_k = _freq_step_body(phase, phase_decay_step, spectra[k], xp)
+            rp = rp + xp.where(is_out, zero, rp_k)
+
+    return rp
+
+
+@jaxtyped(typechecker=typechecker)
+def _pfield_freq_vectorized(
+    phase_decay_init: Complex[Array, " *grid n_sources"],
+    phase_decay_step: Complex[Array, " *grid n_sources"],
+    is_out: Bool[Array, " *grid"],
+    wavenumbers: Float[Array, " n_freq"],
+    pulse_spect: Complex[Array, " n_freq"],
+    probe_spect: Float[Array, " n_freq"],
+    seg_length: float,
+    sin_theta: Float[Array, " *grid n_sources"],
+    full_frequency_directivity: bool,
+    xp: _ArrayNamespace,
+) -> Float[Array, " *grid"]:
+    """Vectorized frequency sweep: broadcast all frequencies at once.
+
+    Reference implementation kept for testing and small-grid use cases.
+    Production code uses _freq_outer_python (constant memory) instead.
+
+    Uses the geometric progression: phase_decay[k] = init * step^k.
+    Source points are pre-flattened with 1/n_sub absorbed.
+
+    Best for small grids where the (*grid, n_sources, n_freq) tensor fits
+    in memory. For large grids, use an iterative driver instead.
+    """
+    n_freq = wavenumbers.shape[0]
+    exponents = xp.arange(n_freq, dtype=wavenumbers.dtype)
+
+    # (*grid, n_sources, n_freq) via geometric progression
+    phase_k = phase_decay_init[..., None] * phase_decay_step[..., None] ** exponents
+
+    if full_frequency_directivity:
+        sinc_arg = wavenumbers * seg_length / 2.0 * sin_theta[..., None] / pi
+        phase_k = xpx.sinc(sinc_arg, xp=xp) * phase_k
+
+    # Contract over sources: (*grid, n_sources, n_freq) -> (*grid, n_freq)
+    pressure_all = xp.sum(phase_k, axis=-2)
+
+    pressure_all = pulse_spect * probe_spect * pressure_all
+
+    pressure_all = xp.where(is_out[..., None], xp.asarray(0.0 + 0j), pressure_all)
+
+    return xp.sum(xp.abs(pressure_all) ** 2, axis=-1)
+
+
+def _freq_outer_scan(
+    phase_decay_init: Complex[Array, " *grid n_sources"],
+    phase_decay_step: Complex[Array, " *grid n_sources"],
+    is_out: Bool[Array, " *grid"],
+    wavenumbers: Float[Array, " n_freq"],
+    pulse_spect: Complex[Array, " n_freq"],
+    probe_spect: Float[Array, " n_freq"],
+    seg_length: float,
+    sin_theta: Float[Array, " *grid n_sources"],
+    full_frequency_directivity: bool,
+    xp: _ArrayNamespace,
+) -> Float[Array, " *grid"]:
+    """JAX vmap+scan driver: per-grid-point kernel vmapped over the grid.
+
+    The scan carry is only (n_sources,) per grid point, matching the Metal
+    kernel's per-thread model. vmap handles parallelism across grid points.
+    """
+    import jax
+    import jax.numpy as jnp
+
+    grid_shape = phase_decay_init.shape[:-1]
+    n_sources = phase_decay_init.shape[-1]
+
+    # Flatten grid to (n_grid, n_sources) and (n_grid,)
+    init_flat = xp.reshape(phase_decay_init, (-1, n_sources))
+    step_flat = xp.reshape(phase_decay_step, (-1, n_sources))
+    is_out_flat = xp.reshape(is_out, (-1,))
+    sin_theta_flat = xp.reshape(sin_theta, (-1, n_sources))
+
+    spectra = pulse_spect * probe_spect
+
+    if full_frequency_directivity:
+
+        def _single_point_scan(
+            phase_init_g: jax.Array, phase_step_g: jax.Array, is_out_g: jax.Array, sin_theta_g: jax.Array
+        ) -> jax.Array:
+            def scan_fn(carry, k):
+                phase, rp = carry
+                sinc_arg = wavenumbers[k] * seg_length / 2.0 * sin_theta_g / pi
+                directivity_k = xpx.sinc(sinc_arg, xp=xp)
+                p_k = spectra[k] * xp.sum(phase * directivity_k)
+                rp_k = xp.real(p_k * xp.conj(p_k))
+                rp = rp + xp.where(is_out_g, xp.asarray(0.0), rp_k)
+                phase = phase * phase_step_g
+                return (phase, rp), None
+
+            (_, rp), _ = jax.lax.scan(scan_fn, (phase_init_g, jnp.float32(0.0)), jnp.arange(spectra.shape[0]))
+            return rp
+
+        rp_flat = jax.vmap(_single_point_scan)(init_flat, step_flat, is_out_flat, sin_theta_flat)
+    else:
+
+        def _single_point_scan_no_dir(
+            phase_init_g: jax.Array, phase_step_g: jax.Array, is_out_g: jax.Array
+        ) -> jax.Array:
+            def scan_fn(carry, spectrum_k):
+                phase, rp = carry
+                p_k = spectrum_k * xp.sum(phase)
+                rp_k = xp.real(p_k * xp.conj(p_k))
+                rp = rp + xp.where(is_out_g, xp.asarray(0.0), rp_k)
+                phase = phase * phase_step_g
+                return (phase, rp), None
+
+            (_, rp), _ = jax.lax.scan(scan_fn, (phase_init_g, jnp.float32(0.0)), spectra)
+            return rp
+
+        rp_flat = jax.vmap(_single_point_scan_no_dir)(init_flat, step_flat, is_out_flat)
+
+    return xp.reshape(rp_flat, grid_shape)