Add bioheatExact solver

kwave/bioheatExact.py

@@ -0,0 +1,155 @@
+
+import numpy as np
+
+
+def bioheat_exact(
+    T0,
+    S,
+    material,
+    dx,
+    t,
+    device="cpu",
+    return_numpy=False,
+):
+    """
+    Fourier solution of Pennes' bioheat equation for a constant source.
+
+    Assumes periodic boundary conditions: opposite edges of each axis
+    are connected by the Fourier representation.
+
+    Uses NumPy on the CPU or CuPy on the GPU. Does not require PyTorch.
+
+    Equation
+    --------
+    dT/dt = D * Laplacian(T) - P * (T - Ta) + S
+
+    Parameters
+    ----------
+    T0 : array-like
+        Initial temperature field, 1D, 2D, or 3D.
+    S : array-like or scalar
+        Constant heat source [K/s], or scalar zero.
+    material : array-like
+        [D, P, Ta], where D is thermal diffusivity [m²/s],
+        P is the perfusion coefficient [1/s], and Ta is
+        the reference temperature [K].
+    dx : float or sequence
+        Grid spacing in metres.
+    t : float
+        Evaluation time in seconds.
+    device : {"cpu", "gpu"}
+        NumPy CPU execution or CuPy GPU execution.
+    return_numpy : bool
+        Return a NumPy array, including when device="gpu".
+
+    Returns
+    -------
+    numpy.ndarray or cupy.ndarray
+        Temperature field at time t. By default, the returned array
+        uses the selected backend.
+    """
+    if device not in ("cpu", "gpu"):
+        raise ValueError("device must be 'cpu' or 'gpu'.")
+
+    # Select the array backend.
+    if device == "gpu":
+        try:
+            import cupy as cp
+        except ImportError as exc:
+            raise ImportError(
+                "GPU execution requires CuPy and a compatible CUDA runtime."
+            ) from exc
+
+        try:
+            if cp.cuda.runtime.getDeviceCount() < 1:
+                raise RuntimeError("No usable CUDA GPU was found for CuPy.")
+        except cp.cuda.runtime.CUDARuntimeError as exc:
+            raise RuntimeError(
+                "No usable CUDA GPU was found for CuPy."
+            ) from exc
+
+        xp = cp
+    else:
+        xp = np
+
+    # Convert fields using the selected backend. In GPU mode, existing
+    # CuPy arrays stay on the GPU and can be reused as inputs.
+    T0_arr = xp.asarray(T0, dtype=xp.float64)
+
+    if T0_arr.ndim not in (1, 2, 3) or T0_arr.size <= 1:
+        raise ValueError("T0 must be a non-scalar 1D, 2D, or 3D array.")
+
+    S_arr = xp.asarray(S, dtype=xp.float64)
+    source_is_zero = S_arr.ndim == 0 and S_arr.item() == 0
+
+    if source_is_zero:
+        S_arr = None
+    elif S_arr.ndim == 0 or S_arr.shape != T0_arr.shape:
+        raise ValueError("S must be zero or have the same shape as T0.")
+
+    # Material parameters and grid spacing are small CPU-side metadata.
+    material_arr = np.asarray(material, dtype=float).ravel()
+
+    if material_arr.size != 3:
+        raise ValueError("material must contain [D, P, Ta].")
+
+    D, P, Ta = map(float, material_arr)
+
+    if not np.all(np.isfinite(material_arr)) or D < 0 or P < 0:
+        raise ValueError(
+            "Material parameters must be finite; D and P must be non-negative."
+        )
+
+    t = float(t)
+    if not np.isfinite(t) or t < 0:
+        raise ValueError("t must be finite and non-negative.")
+
+    if np.isscalar(dx):
+        spacing = [float(dx)] * T0_arr.ndim
+    else:
+        spacing = np.asarray(dx, dtype=float).ravel().tolist()
+
+    if len(spacing) != T0_arr.ndim:
+        raise ValueError(
+            f"dx must be a scalar or contain {T0_arr.ndim} values."
+        )
+
+    if not np.all(np.isfinite(spacing)) or any(d <= 0 for d in spacing):
+        raise ValueError("Grid spacing must be finite and positive.")
+
+    # Build squared angular wavenumber magnitude in FFT ordering.
+    # No fftshift/ifftshift is required with fftfreq.
+    k_axes = [
+        2.0 * np.pi * xp.fft.fftfreq(n, d=d)
+        for n, d in zip(T0_arr.shape, spacing)
+    ]
+
+    k_mesh = xp.meshgrid(*k_axes, indexing="ij")
+    k_squared = sum(ki**2 for ki in k_mesh)
+
+    denominator = D * k_squared + P
+    decay = xp.exp(-denominator * t)
+
+    # Stable evaluation of (1 - exp(-a*t)) / a, including a == 0.
+    nonzero = denominator != 0
+    safe_denominator = xp.where(nonzero, denominator, 1.0)
+    Q = xp.where(
+        nonzero,
+        -xp.expm1(-denominator * t) / safe_denominator,
+        t,
+    )
+
+    # Fourier solution for T - Ta.
+    T0_fft = xp.fft.fftn(T0_arr - Ta)
+
+    if S_arr is None:
+        T_fft = decay * T0_fft
+    else:
+        T_fft = decay * T0_fft + Q * xp.fft.fftn(S_arr)
+
+    T = xp.fft.ifftn(T_fft).real + Ta
+
+    if return_numpy:
+        return xp.asnumpy(T) if device == "gpu" else T
+
+    return T

tests/test_bioheat_exact.py

@@ -0,0 +1,99 @@
+
+import numpy as np
+import pytest
+
+from kwave.bioheatExact import bioheat_exact
+
+
+def test_zero_time_preserves_initial_temperature():
+    T0 = np.arange(12, dtype=float).reshape(3, 4) + 20.0
+
+    result = bioheat_exact(
+        T0=T0,
+        S=0,
+        material=[1.2e-7, 0.02, 37.0],
+        dx=(1e-3, 2e-3),
+        t=0,
+    )
+
+    np.testing.assert_allclose(result, T0, rtol=1e-12, atol=1e-12)
+
+
+def test_zero_diffusion_and_perfusion():
+    T0 = np.arange(12, dtype=float).reshape(3, 4)
+    S = np.full_like(T0, 2.5)
+    t = 3.0
+
+    result = bioheat_exact(
+        T0=T0,
+        S=S,
+        material=[0.0, 0.0, 37.0],
+        dx=(1e-3, 2e-3),
+        t=t,
+    )
+
+    np.testing.assert_allclose(result, T0 + t * S, rtol=1e-12, atol=1e-12)
+
+
+def test_single_fourier_mode_on_uneven_grid():
+    nx, ny = 16, 12
+    dx = (0.2, 0.07)
+    D = 0.3
+    P = 0.15
+    Ta = 37.0
+    t = 0.4
+
+    x = np.arange(nx)[:, None]
+    y = np.arange(ny)[None, :]
+
+    # A single periodic Fourier mode with nonzero wavenumbers
+    # along both axes.
+    mode = np.cos(2 * np.pi * x / nx + 2 * np.pi * y / ny)
+    T0 = Ta + mode
+
+    kx = 2 * np.pi / (nx * dx[0])
+    ky = 2 * np.pi / (ny * dx[1])
+    decay = np.exp(-(D * (kx**2 + ky**2) + P) * t)
+
+    expected = Ta + decay * mode
+
+    result = bioheat_exact(
+        T0=T0,
+        S=0,
+        material=[D, P, Ta],
+        dx=dx,
+        t=t,
+    )
+
+    np.testing.assert_allclose(result, expected, rtol=1e-11, atol=1e-11)
+
+
+def test_gpu_result_can_be_reused():
+    cp = pytest.importorskip("cupy")
+
+    try:
+        if cp.cuda.runtime.getDeviceCount() < 1:
+            pytest.skip("No CUDA GPU available")
+        cp.zeros(1).sum().item()  # Check that GPU execution works.
+    except cp.cuda.runtime.CUDARuntimeError:
+        pytest.skip("No usable CUDA GPU available")
+
+    T0 = cp.full((8, 8), 37.0)
+    S = cp.ones((8, 8), dtype=cp.float64)
+
+    kwargs = {
+        "S": S,
+        "material": [0.0, 0.0, 37.0],
+        "dx": 1e-3,
+        "t": 0.1,
+        "device": "gpu",
+    }
+
+    T1 = bioheat_exact(T0=T0, **kwargs)
+    T2 = bioheat_exact(T0=T1, **kwargs)
+
+    assert isinstance(T1, cp.ndarray)
+    assert isinstance(T2, cp.ndarray)
+
+    cp.testing.assert_allclose(T1, 37.1)
+    cp.testing.assert_allclose(T2, 37.2)