diff --git a/kwave/bioheatExact.py b/kwave/bioheatExact.py new file mode 100644 index 00000000..8908e7e8 --- /dev/null +++ b/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 diff --git a/tests/test_bioheat_exact.py b/tests/test_bioheat_exact.py new file mode 100644 index 00000000..4aa5696e --- /dev/null +++ b/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)