From e2ecec054636474cbefd79285c55fb24ab928cfe Mon Sep 17 00:00:00 2001 From: rhh006 Date: Wed, 4 Nov 2026 11:02:04 -0800 Subject: [PATCH 1/5] Add files via upload This PR adds a Python port of the kWaveDiffusion class from the MATLAB k-Wave toolbox to k-Wave-Python, supporting 1D, 2D, and 3D thermal diffusion and perfusion simulations in homogeneous and heterogeneous media, with optional GPU acceleration. --- kwave/kwaveDiffusion.py | 1663 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 1663 insertions(+) create mode 100644 kwave/kwaveDiffusion.py diff --git a/kwave/kwaveDiffusion.py b/kwave/kwaveDiffusion.py new file mode 100644 index 00000000..36bc6600 --- /dev/null +++ b/kwave/kwaveDiffusion.py @@ -0,0 +1,1663 @@ +""" +kWaveDiffusion +============== + +Python/NumPy port of the MATLAB kWaveDiffusion class from k-Wave. + +Solves: + + A * dT/dt = div(K * grad(T)) - B * (T - Ta) + Q + +or, when diffusion_coeff is supplied: + + dT/dt = div(D * grad(T)) - P * (T - Ta) + S + +where + + D = K / A + P = B / A + S = Q / A + +The implementation below currently supports periodic boundaries. + +Supported dimensions: + 1D + 2D + 3D + +Supported media: + homogeneous + heterogeneous + +Supported properties: + diffusion_coeff + density + specific_heat + thermal_conductivity + + perfusion_coeff + blood_density + blood_specific_heat + blood_perfusion_rate + blood_ambient_temperature + +Supported source: + T0 + Q + +References +---------- +Treeby & Saratoon, k-Wave MATLAB kWaveDiffusion. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +import numpy as np + +try: + import cupy as cp +except ImportError: + cp = None + + +ArrayLike = Union[np.ndarray, float, int] + + +@dataclass +class DiffusionMedium: + """ + Medium parameters. + + Either provide: + + diffusion_coeff + + or: + + density + specific_heat + thermal_conductivity + + For perfusion, either provide: + + perfusion_coeff + + or: + + blood_density + blood_specific_heat + blood_perfusion_rate + + blood_ambient_temperature is required when perfusion is used. + """ + + # Diffusion representation + diffusion_coeff: Optional[ArrayLike] = None + + density: Optional[ArrayLike] = None + specific_heat: Optional[ArrayLike] = None + thermal_conductivity: Optional[ArrayLike] = None + + # Perfusion representation + perfusion_coeff: Optional[ArrayLike] = None + + blood_density: Optional[ArrayLike] = None + blood_specific_heat: Optional[ArrayLike] = None + blood_perfusion_rate: Optional[ArrayLike] = None + + blood_ambient_temperature: Optional[ArrayLike] = None + + # k-space reference coefficients + diffusion_coeff_ref: Union[str, float] = "max" + perfusion_coeff_ref: Union[str, float] = "max" + + # Boundary condition + boundary_condition: str = "periodic" + + +@dataclass +class DiffusionSource: + """ + Thermal source. + + T0: + Initial temperature [degC] + + Q: + Volumetric heat deposition [W/m^3] + """ + + T0: ArrayLike + Q: ArrayLike = 0.0 + + +class kWaveDiffusion: + """ + Time-domain k-space solver for the diffusion / Pennes bioheat equation. + + This is a Python port of the MATLAB kWaveDiffusion class. + + Parameters + ---------- + kgrid: + kWaveGrid instance. + + medium: + DiffusionMedium or a dictionary containing the medium parameters. + + source: + DiffusionSource or dictionary containing T0 and optionally Q. + + sensor: + Optional binary sensor mask. + + use_kspace: + If True, use the k-space temporal correction. + + dtype: + np.float64 or np.float32. + + display_updates: + Print simulation information. + + device: + "cpu" (NumPy) or "gpu"/"cuda" (CuPy CUDA backend). + GPU execution requires a CUDA-compatible CuPy installation. + + Notes + ----- + Current implementation supports periodic boundaries. + + Insulating/conducting DCT/DST boundaries are intentionally not + implemented here because those require the k-Wave DTT operators. + """ + + def __init__( + self, + kgrid, + medium, + source, + sensor=None, + *, + use_kspace: bool = True, + dtype=np.float64, + display_updates: bool = True, + device: str = "cpu", + ): + + self.kgrid = kgrid + self.dtype = np.dtype(dtype) + + device = str(device).lower() + if device in ("cuda", "gpu"): + if cp is None: + raise ImportError( + "CUDA execution requires CuPy. Install a CuPy build " + "compatible with your CUDA version." + ) + if cp.cuda.runtime.getDeviceCount() < 1: + raise RuntimeError("No CUDA device is available.") + self.device = "gpu" + self.xp = cp + elif device == "cpu": + self.device = "cpu" + self.xp = np + else: + raise ValueError("device must be 'cpu', 'gpu', or 'cuda'.") + + self.use_kspace = use_kspace + self.display_updates = display_updates + + # ------------------------------------------------------------ + # Determine grid dimensions + # ------------------------------------------------------------ + + self._read_grid(kgrid) + + # ------------------------------------------------------------ + # Read medium/source + # ------------------------------------------------------------ + + if isinstance(medium, dict): + medium = DiffusionMedium(**medium) + + if isinstance(source, dict): + source = DiffusionSource(**source) + + self.medium = medium + self.source = source + + # ------------------------------------------------------------ + # Validate boundary condition + # ------------------------------------------------------------ + + boundary = medium.boundary_condition + + if boundary not in ( + "periodic", + "insulating", + "conducting", + ): + raise ValueError( + f"Unknown boundary condition: {boundary}" + ) + + if boundary != "periodic": + raise NotImplementedError( + "This Python implementation currently supports " + "periodic boundaries only. Insulating/conducting " + "DCT/DST support can be added separately." + ) + + self.boundary_condition = boundary + + # ------------------------------------------------------------ + # Read material properties + # ------------------------------------------------------------ + + self._setup_diffusion(medium) + + self._setup_perfusion(medium) + + self._setup_source(source) + + # ------------------------------------------------------------ + # Determine whether medium is heterogeneous + # ------------------------------------------------------------ + + self.flag_homogeneous = ( + self.diffusion_p1.size == 1 + and self.diffusion_p2.size == 1 + and self.perfusion_coeff.size == 1 + and self.blood_ambient_temperature.size == 1 + ) + + # ------------------------------------------------------------ + # Reference coefficients + # ------------------------------------------------------------ + + self.diffusion_coeff_ref = self._reference_value( + medium.diffusion_coeff_ref, + self.diffusion_p1 * self.diffusion_p2, + "diffusion_coeff_ref", + ) + + if bool(self.xp.any(self.perfusion_coeff != 0)): + self.perfusion_coeff_ref = self._reference_value( + medium.perfusion_coeff_ref, + self.perfusion_coeff, + "perfusion_coeff_ref", + ) + else: + self.perfusion_coeff_ref = 0.0 + + # ------------------------------------------------------------ + # Initial temperature + # ------------------------------------------------------------ + + T0 = self.xp.asarray(source.T0, dtype=self.dtype) + + if T0.ndim == 0: + self.T = self.xp.full( + self.shape, + T0.item(), + dtype=self.dtype, + ) + else: + if T0.shape != self.shape: + raise ValueError( + f"T0 has shape {T0.shape}, expected {self.shape}" + ) + + self.T = T0.copy() + + # ------------------------------------------------------------ + # Heat source + # ------------------------------------------------------------ + + Q = self.xp.asarray(source.Q, dtype=self.dtype) + + if Q.ndim == 0: + self.Q = Q.item() + else: + if Q.shape != self.shape: + raise ValueError( + f"Q has shape {Q.shape}, expected {self.shape}" + ) + + self.Q = Q.copy() + + # ------------------------------------------------------------ + # Thermal damage + # ------------------------------------------------------------ + + self.cem43 = self.xp.zeros( + self.shape, + dtype=self.dtype, + ) + + # ------------------------------------------------------------ + # Time/sensor bookkeeping + # ------------------------------------------------------------ + + self.time_steps_taken = 0 + + if sensor is not None: + if isinstance(sensor, dict): + sensor = sensor["mask"] + + self.sensor_mask = self.xp.asarray(sensor) + + if self.sensor_mask.shape != self.shape: + raise ValueError( + f"Sensor mask has shape {self.sensor_mask.shape}, " + f"expected {self.shape}" + ) + + # MATLAB requires binary mask + if not bool(self.xp.all( + (self.sensor_mask == 0) + | (self.sensor_mask == 1) + )): + raise ValueError( + "sensor.mask must contain only 0 and 1." + ) + + self.sensor_mask_index = self.xp.flatnonzero( + self.sensor_mask.ravel(order="C") + ) + + self.num_sensor_points = len( + self.sensor_mask_index + ) + + self.sensor_data = None + + else: + self.sensor_mask = None + self.sensor_mask_index = None + self.num_sensor_points = 0 + self.sensor_data = None + + # ------------------------------------------------------------ + # Build Fourier operators + # ------------------------------------------------------------ + + self._build_wavenumbers() + + # ================================================================== + # GRID + # ================================================================== + + def _read_grid(self, kgrid): + + """ + Read grid dimensions and spacings. + + k-Wave-Python's kWaveGrid stores the grid dimensions in N + and the spacings in spacing. This function also supports + common legacy attribute layouts. + """ + + # New k-Wave-Python API + if hasattr(kgrid, "N"): + N = np.asarray(kgrid.N).astype(int).ravel() + + if hasattr(kgrid, "spacing"): + spacing = np.asarray( + kgrid.spacing + ).astype(float).ravel() + else: + spacing = None + + else: + N = [] + + for name in ("Nx", "Ny", "Nz"): + if hasattr(kgrid, name): + N.append(int(getattr(kgrid, name))) + + N = np.asarray(N, dtype=int) + + spacing = [] + + for name in ("dx", "dy", "dz"): + if hasattr(kgrid, name): + spacing.append(float(getattr(kgrid, name))) + + spacing = np.asarray( + spacing, + dtype=float, + ) + + if N.size == 0: + raise ValueError( + "Could not determine grid dimensions from kgrid." + ) + + self.dim = int(N.size) + + self.Nx = int(N[0]) + + self.Ny = int(N[1]) if self.dim >= 2 else 1 + self.Nz = int(N[2]) if self.dim >= 3 else 1 + + self.shape = tuple( + int(v) for v in N + ) + + if spacing is None or spacing.size < self.dim: + spacing = np.ones( + self.dim, + dtype=float, + ) + + for i, name in enumerate( + ("dx", "dy", "dz") + ): + if hasattr(kgrid, name): + spacing[i] = float( + getattr(kgrid, name) + ) + + self.dx = float(spacing[0]) + + self.dy = ( + float(spacing[1]) + if self.dim >= 2 + else self.dx + ) + + self.dz = ( + float(spacing[2]) + if self.dim >= 3 + else self.dx + ) + + # ================================================================== + # MATERIAL + # ================================================================== + + def _as_field( + self, + value, + name, + allow_none=False, + ): + + if value is None: + if allow_none: + return self.xp.asarray( + 0.0, + dtype=self.dtype, + ) + + raise ValueError( + f"{name} must be specified." + ) + + arr = self.xp.asarray( + value, + dtype=self.dtype, + ) + + if arr.ndim == 0: + return arr + + if arr.shape != self.shape: + raise ValueError( + f"{name} has shape {arr.shape}; " + f"expected scalar or {self.shape}." + ) + + return arr + + def _setup_diffusion(self, medium): + + # ------------------------------------------------------------ + # Either D or rho/c/K + # ------------------------------------------------------------ + + if medium.diffusion_coeff is not None: + + self.diffusion_p1 = self.xp.asarray( + 1.0, + dtype=self.dtype, + ) + + self.diffusion_p2 = self._as_field( + medium.diffusion_coeff, + "diffusion_coeff", + ) + + else: + + if ( + medium.density is None + or medium.specific_heat is None + or medium.thermal_conductivity is None + ): + raise ValueError( + "Specify either diffusion_coeff or " + "density, specific_heat, and " + "thermal_conductivity." + ) + + density = self._as_field( + medium.density, + "density", + ) + + specific_heat = self._as_field( + medium.specific_heat, + "specific_heat", + ) + + conductivity = self._as_field( + medium.thermal_conductivity, + "thermal_conductivity", + ) + + self.diffusion_p1 = ( + 1.0 / (density * specific_heat) + ) + + self.diffusion_p2 = conductivity + + # ================================================================== + + def _setup_perfusion(self, medium): + + # No perfusion parameters at all + if ( + medium.perfusion_coeff is None + and medium.blood_density is None + and medium.blood_specific_heat is None + and medium.blood_perfusion_rate is None + and medium.blood_ambient_temperature is None + ): + self.perfusion_coeff = self.xp.asarray( + 0.0, + dtype=self.dtype, + ) + + self.blood_ambient_temperature = self.xp.asarray( + 0.0, + dtype=self.dtype, + ) + + return + + # Direct P + if medium.perfusion_coeff is not None: + + if medium.blood_ambient_temperature is None: + raise ValueError( + "blood_ambient_temperature is required " + "when perfusion_coeff is specified." + ) + + self.perfusion_coeff = self._as_field( + medium.perfusion_coeff, + "perfusion_coeff", + ) + + self.blood_ambient_temperature = self._as_field( + medium.blood_ambient_temperature, + "blood_ambient_temperature", + ) + + return + + # Full blood parameters + required = ( + medium.blood_density, + medium.blood_specific_heat, + medium.blood_perfusion_rate, + medium.blood_ambient_temperature, + medium.density, + medium.specific_heat, + ) + + if any(x is None for x in required): + + raise ValueError( + "For blood-property perfusion specify:\n" + "blood_density\n" + "blood_specific_heat\n" + "blood_perfusion_rate\n" + "blood_ambient_temperature\n" + "density\n" + "specific_heat" + ) + + blood_density = self._as_field( + medium.blood_density, + "blood_density", + ) + + blood_specific_heat = self._as_field( + medium.blood_specific_heat, + "blood_specific_heat", + ) + + blood_perfusion_rate = self._as_field( + medium.blood_perfusion_rate, + "blood_perfusion_rate", + ) + + density = self._as_field( + medium.density, + "density", + ) + + specific_heat = self._as_field( + medium.specific_heat, + "specific_heat", + ) + + self.perfusion_coeff = ( + blood_density + * blood_perfusion_rate + * blood_specific_heat + / (density * specific_heat) + ) + + self.blood_ambient_temperature = self._as_field( + medium.blood_ambient_temperature, + "blood_ambient_temperature", + ) + + # ================================================================== + + def _setup_source(self, source): + """ + Q is retained in W/m^3. + + If Q is supplied, the update term is: + + Q / (rho * c) + + when rho*c is known. + + When diffusion_coeff was supplied, density and specific_heat + are therefore required for physical Q scaling. + """ + + if source.Q is None: + source.Q = 0.0 + + if bool(self.xp.all( + self.xp.asarray(source.Q) == 0 + )): + return + + # If density and specific heat are available, + # calculate 1/A. + if ( + self.medium.density is not None + and self.medium.specific_heat is not None + ): + + density = self._as_field( + self.medium.density, + "density", + ) + + specific_heat = self._as_field( + self.medium.specific_heat, + "specific_heat", + ) + + self.q_scale_factor = ( + 1.0 / (density * specific_heat) + ) + + else: + + # MATLAB requires these properties whenever Q + # is defined. + raise ValueError( + "density and specific_heat must be supplied " + "when source.Q is nonzero." + ) + + # ================================================================== + # REFERENCE COEFFICIENT + # ================================================================== + + def _reference_value( + self, + ref, + values, + name, + ): + + values = self.xp.asarray(values) + + if isinstance(ref, str): + + ref_lower = ref.lower() + + if ref_lower == "min": + return float(float(self.xp.min(values).item())) + + if ref_lower == "mean": + return float(float(self.xp.mean(values).item())) + + if ref_lower == "max": + return float(float(self.xp.max(values).item())) + + raise ValueError( + f"Invalid {name}: {ref}" + ) + + return float(ref) + + # ================================================================== + # WAVENUMBERS + # ================================================================== + + def _build_wavenumbers(self): + + """ + Construct FFT wavenumbers. + + Ordering is shifted so that the resulting arrays correspond + to the ordering used by MATLAB's: + + ifftshift(kgrid.k) + """ + + axes = [] + + for n, d in zip( + self.shape, + (self.dx, self.dy, self.dz), + ): + + if n <= 1: + k = self.xp.zeros( + 1, + dtype=self.dtype, + ) + + else: + + # Equivalent Fourier frequencies in rad/m + k = ( + 2.0 + * np.pi + * self.xp.fft.fftfreq( + n, + d=d, + ) + ) + + axes.append(self.xp.asarray(k, dtype=self.dtype)) + + # ------------------------------------------------------------ + # MATLAB kWave ordering: + # + # ifftshift(kgrid.k) + # + # The normal numpy FFT frequencies are already in FFT order, + # so we use those directly for multiplication with fftn(). + # ------------------------------------------------------------ + + if self.dim == 1: + + self.kx = axes[0] + self.k = self.xp.abs(self.kx) + + elif self.dim == 2: + + self.kx = axes[0][:, None] + self.ky = axes[1][None, :] + + self.k = self.xp.sqrt( + self.kx**2 + + self.ky**2 + ) + + elif self.dim == 3: + + self.kx = axes[0][:, None, None] + self.ky = axes[1][None, :, None] + self.kz = axes[2][None, None, :] + + self.k = self.xp.sqrt( + self.kx**2 + + self.ky**2 + + self.kz**2 + ) + + else: + + raise ValueError( + "Only 1D, 2D, and 3D are supported." + ) + + self.k_squared = self.k**2 + + # ================================================================== + # FOURIER TRANSFORMS + # ================================================================== + + def FT(self, x): + """ + Forward Fourier transform. + """ + + return self.xp.fft.fftn(x) + + # ------------------------------------------------------------------ + + def IT(self, x): + """ + Inverse Fourier transform. + + The MATLAB implementation explicitly takes real(). + """ + + return self.xp.fft.ifftn(x).real + + # ================================================================== + # K-SPACE CORRECTION + # ================================================================== + + def _calculate_kappa(self, dt): + + if not self.use_kspace: + return self.xp.asarray( + 1.0, + dtype=self.dtype, + ) + + kappa_argument = ( + dt + * ( + self.diffusion_coeff_ref + * self.k_squared + + self.perfusion_coeff_ref + ) + ) + + kappa = self.xp.empty_like( + kappa_argument, + dtype=self.dtype, + ) + + # MATLAB: + # + # kappa = (1 - exp(-x)) / x + # + # Need explicit limit at x=0: + # + # lim x->0 = 1 + + zero = ( + self.xp.abs(kappa_argument) == 0 + ) + + nonzero = ~zero + + kappa[zero] = 1.0 + + x = kappa_argument[nonzero] + + kappa[nonzero] = ( + -self.xp.expm1(-x) / x + ) + + return kappa + + # ================================================================== + # DERIVATIVE OPERATORS + # ================================================================== + + def _derivative_operators(self, kappa): + + """ + Construct: + + i*kx*sqrt(kappa) + i*ky*sqrt(kappa) + i*kz*sqrt(kappa) + + exactly corresponding to the periodic-boundary branch of + MATLAB getDerivMatrix(). + """ + + sqrt_kappa = self.xp.sqrt(kappa) + + deriv_x = ( + 1j + * self.kx + * sqrt_kappa + ) + + if self.dim >= 2: + deriv_y = ( + 1j + * self.ky + * sqrt_kappa + ) + else: + deriv_y = None + + if self.dim >= 3: + deriv_z = ( + 1j + * self.kz + * sqrt_kappa + ) + else: + deriv_z = None + + return ( + deriv_x, + deriv_y, + deriv_z, + ) + + # ================================================================== + # SOURCE TERM + # ================================================================== + + def _source_update(self, kappa): + + if np.isscalar(self.Q): + + if self.Q == 0: + return 0.0 + + else: + + if bool(self.xp.all(self.Q == 0)): + return 0.0 + + Q_ft = self.FT(self.Q) + + result = self.IT( + kappa * Q_ft + ) + + # Physical source: + # + # Q / A + # + # where A = rho*c + + result = ( + self.q_scale_factor + * result + ) + + return result + + # ================================================================== + # PERFUSION TERM + # ================================================================== + + def _perfusion_update( + self, + kappa, + ): + + if bool(self.xp.all( + self.perfusion_coeff == 0 + )): + return 0.0 + + temperature_difference = ( + self.T + - self.blood_ambient_temperature + ) + + result = self.IT( + kappa + * self.FT( + temperature_difference + ) + ) + + result = ( + -self.perfusion_coeff + * result + ) + + return result + + # ================================================================== + # HOMOGENEOUS DIFFUSION + # ================================================================== + + def _homogeneous_diffusion_update( + self, + kappa, + ): + + """ + MATLAB: + + d_term = + diffusion_p1 .* diffusion_p2 .* ... + IT(-k.^2 .* kappa .* FT(T)) + """ + + coefficient = ( + self.diffusion_p1 + * self.diffusion_p2 + ) + + T_ft = self.FT(self.T) + + result = self.IT( + -self.k_squared + * kappa + * T_ft + ) + + return coefficient * result + + # ================================================================== + # HETEROGENEOUS DIFFUSION + # ================================================================== + + def _heterogeneous_diffusion_update( + self, + kappa, + deriv_x, + deriv_y, + deriv_z, + ): + + """ + MATLAB heterogeneous formulation: + + div(K grad(T)) / A + + implemented spectrally as: + + p1 * [ + D_x^-1 K D_x T + + + D_y^-1 K D_y T + + + D_z^-1 K D_z T + ] + + where the derivative operators contain sqrt(kappa). + """ + + T_ft = self.FT(self.T) + + d_term = 0.0 + + # ------------------------------------------------------------ + # X + # ------------------------------------------------------------ + + temp_x = self.IT( + deriv_x * T_ft + ) + + flux_x = ( + self.diffusion_p2 + * temp_x + ) + + d_term = d_term + self.IT( + deriv_x + * self.FT(flux_x) + ) + + # ------------------------------------------------------------ + # Y + # ------------------------------------------------------------ + + if self.dim >= 2: + + temp_y = self.IT( + deriv_y * T_ft + ) + + flux_y = ( + self.diffusion_p2 + * temp_y + ) + + d_term = d_term + self.IT( + deriv_y + * self.FT(flux_y) + ) + + # ------------------------------------------------------------ + # Z + # ------------------------------------------------------------ + + if self.dim >= 3: + + temp_z = self.IT( + deriv_z * T_ft + ) + + flux_z = ( + self.diffusion_p2 + * temp_z + ) + + d_term = d_term + self.IT( + deriv_z + * self.FT(flux_z) + ) + + return ( + self.diffusion_p1 + * d_term + ) + + # ================================================================== + # TIME STEP + # ================================================================== + + def take_time_step( + self, + Nt: int, + dt: float, + ): + """ + Advance the temperature field Nt time steps. + + Parameters + ---------- + Nt: + Number of time steps. + + dt: + Time step [s]. + + Returns + ------- + np.ndarray + Final temperature field. + """ + + Nt = int(Nt) + + if Nt <= 0: + raise ValueError( + "Nt must be positive." + ) + + dt = float(dt) + + if dt <= 0: + raise ValueError( + "dt must be positive." + ) + + if self.display_updates: + + print( + "Running k-Wave thermal simulation..." + ) + + print( + f" dimensions: {self.shape}" + ) + + print( + f" dt: {dt:.6g} s" + ) + + print( + f" Nt: {Nt}" + ) + + print( + f" t_end: {Nt * dt:.6g} s" + ) + + print( + f" heterogeneous: " + f"{not self.flag_homogeneous}" + ) + + # ------------------------------------------------------------ + # k-space correction + # ------------------------------------------------------------ + + kappa = self._calculate_kappa(dt) + + # ------------------------------------------------------------ + # derivative operators + # ------------------------------------------------------------ + + if not self.flag_homogeneous: + + ( + deriv_x, + deriv_y, + deriv_z, + ) = self._derivative_operators( + kappa + ) + + else: + + deriv_x = None + deriv_y = None + deriv_z = None + + # ------------------------------------------------------------ + # source update + # + # Q does not depend on temperature and therefore can be + # calculated once. + # ------------------------------------------------------------ + + q_term = self._source_update( + kappa + ) + + # ------------------------------------------------------------ + # Time loop + # ------------------------------------------------------------ + + for t_index in range(Nt): + + # -------------------------------------------------------- + # Perfusion + # -------------------------------------------------------- + + p_term = self._perfusion_update( + kappa + ) + + # -------------------------------------------------------- + # Diffusion + # -------------------------------------------------------- + + if self.flag_homogeneous: + + d_term = ( + self._homogeneous_diffusion_update( + kappa + ) + ) + + else: + + d_term = ( + self._heterogeneous_diffusion_update( + kappa, + deriv_x, + deriv_y, + deriv_z, + ) + ) + + # -------------------------------------------------------- + # Temperature update + # + # MATLAB: + # + # obj.T = obj.T + dt .* ... + # (d_term + p_term + q_term) + # -------------------------------------------------------- + + self.T = ( + self.T + + dt + * ( + d_term + + p_term + + q_term + ) + ) + + # -------------------------------------------------------- + # CEM43 + # -------------------------------------------------------- + + self._update_cem43( + dt + ) + + # -------------------------------------------------------- + # Sensor + # -------------------------------------------------------- + + if self.num_sensor_points > 0: + + values = self.T.ravel( + order="C" + )[ + self.sensor_mask_index + ] + + if self.sensor_data is None: + + self.sensor_data = self.xp.empty( + ( + self.num_sensor_points, + Nt, + ), + dtype=self.dtype, + ) + + self.sensor_data[ + :, + t_index, + ] = values + + self.time_steps_taken += Nt + + if self.display_updates: + + print( + "Simulation completed." + ) + + return self.T + + # ================================================================== + # CEM43 + # ================================================================== + + def _update_cem43(self, dt): + + """ + Port of: + + obj.cem43 = obj.cem43 + dt ./ 60 .* ... + (0.5 .* (obj.T >= 43) + ... + 0.25 .* (obj.T >= 37 & obj.T < 43)).^(43 - obj.T); + """ + + factor = self.xp.where( + self.T >= 43.0, + 0.5, + self.xp.where( + (self.T >= 37.0) + & (self.T < 43.0), + 0.25, + 0.0, + ), + ) + + # Avoid 0**positive -> 0 naturally. + # For temperatures below 37 this corresponds to + # zero thermal dose in the MATLAB expression. + + self.cem43 += ( + dt + / 60.0 + * self.xp.where( + factor > 0, + factor ** ( + 43.0 - self.T + ), + 0.0, + ) + ) + + # ================================================================== + # LESION MAP + # ================================================================== + + @property + def lesion_map(self): + + return self.cem43 >= 240.0 + + # ================================================================== + # LESION SIZE + # ================================================================== + + @property + def lesion_size(self): + + if self.dim == 1: + + cell_volume = self.dx + + elif self.dim == 2: + + cell_volume = ( + self.dx + * self.dy + ) + + else: + + cell_volume = ( + self.dx + * self.dy + * self.dz + ) + + value = self.xp.sum(self.lesion_map) * cell_volume + return float(value.item()) if self.device == "gpu" else value + + # ================================================================== + # DT LIMIT + # ================================================================== + + @property + def dt_limit(self): + + """ + Port of the MATLAB dt_limit property. + + This is primarily useful for determining whether the chosen + timestep is inside the conditionally-stable region. + + The homogeneous k-space formulation is unconditionally stable + when the reference coefficient is sufficiently large. + """ + + diffusion_coeff = ( + self.diffusion_p1 + * self.diffusion_p2 + ) + + D_max = float( + float(self.xp.max(diffusion_coeff).item()) + ) + + k_max = float( + float(self.xp.max(self.k).item()) + ) + + k_min = float( + float(self.xp.min(self.k).item()) + ) + + P = self.perfusion_coeff + + # ------------------------------------------------------------ + # No perfusion + # ------------------------------------------------------------ + + if bool(self.xp.all(P == 0)): + + if ( + self.diffusion_coeff_ref + >= D_max / 2.0 + ): + + return np.inf + + denominator = ( + self.diffusion_coeff_ref + * k_max**2 + ) + + if denominator == 0: + return np.inf + + argument = ( + 1.0 + - 2.0 + * self.diffusion_coeff_ref + / D_max + ) + + if argument <= 0: + return np.inf + + return ( + -self.xp.log(argument) + / denominator + ) + + # ------------------------------------------------------------ + # Perfusion + # ------------------------------------------------------------ + + reference = ( + self.diffusion_coeff_ref + * k_max**2 + + self.perfusion_coeff_ref + ) + + condition = ( + 0.5 + * ( + diffusion_coeff + * k_max**2 + + P + ) + ) + + if reference >= float(self.xp.max(condition).item()): + + return np.inf + + denominator_max = self.xp.max( + diffusion_coeff + * k_max**2 + + P + ) + + arg_max = ( + 1.0 + - 2.0 + * ( + self.diffusion_coeff_ref + * k_max**2 + + self.perfusion_coeff_ref + ) + / denominator_max + ) + + if arg_max <= 0: + kmax_val = np.inf + + else: + + kmax_val = ( + -self.xp.log(arg_max) + / ( + self.diffusion_coeff_ref + * k_max**2 + + self.perfusion_coeff_ref + ) + ) + + denominator_min = self.xp.max( + diffusion_coeff + * k_min**2 + + P + ) + + arg_min = ( + 1.0 + - 2.0 + * ( + self.diffusion_coeff_ref + * k_min**2 + + self.perfusion_coeff_ref + ) + / denominator_min + ) + + if arg_min <= 0: + + kmin_val = np.inf + + else: + + denom = ( + self.diffusion_coeff_ref + * k_min**2 + + self.perfusion_coeff_ref + ) + + if denom == 0: + + kmin_val = np.inf + + else: + + kmin_val = ( + -self.xp.log(arg_min) + / denom + ) + + return min( + kmax_val, + kmin_val, + ) + + # ================================================================== + # RESET + # ================================================================== + + def reset( + self, + T0=None, + ): + """ + Reset temperature and thermal dose. + + If T0 is omitted, the original source.T0 is used. + """ + + if T0 is None: + T0 = self.source.T0 + + T0 = self.xp.asarray( + T0, + dtype=self.dtype, + ) + + if T0.ndim == 0: + + self.T = self.xp.full( + self.shape, + T0.item(), + dtype=self.dtype, + ) + + else: + + if T0.shape != self.shape: + raise ValueError( + f"T0 shape {T0.shape} " + f"does not match {self.shape}" + ) + + self.T = T0.copy() + + self.cem43.fill(0) + + self.time_steps_taken = 0 + + self.sensor_data = None From 7e8261b1b5962ad1698c643bf6873b4a62770b5e Mon Sep 17 00:00:00 2001 From: rhh006 Date: Wed, 4 Nov 2026 11:07:40 -0800 Subject: [PATCH 2/5] Add files via upload Add a 3D thermal diffusion test for kWaveDiffusion function demonstrating localized heating, perfusion, and CPU/CUDA device selection. Reports temperature, CEM43 thermal dose, and lesion volume. --- tests/kwaveDiffusion_test.py | 120 +++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 tests/kwaveDiffusion_test.py diff --git a/tests/kwaveDiffusion_test.py b/tests/kwaveDiffusion_test.py new file mode 100644 index 00000000..2c9fa485 --- /dev/null +++ b/tests/kwaveDiffusion_test.py @@ -0,0 +1,120 @@ +import numpy as np + +from kwave.kgrid import kWaveGrid +from kwaveDiffusion import DiffusionMedium, DiffusionSource, kWaveDiffusion + + +# ---------------------------------------------------------------------- +# Grid +# ---------------------------------------------------------------------- + +Nx, Ny, Nz = 64, 64, 64 +dx = 0.5e-3 + +kgrid = kWaveGrid([Nx, Ny, Nz], [dx, dx, dx]) + + +# ---------------------------------------------------------------------- +# Initial temperature and heat source +# ---------------------------------------------------------------------- + +T0 = np.full((Nx, Ny, Nz), 37.0, dtype=np.float64) + +S = np.zeros_like(T0) + +# Small spherical-ish heating region +cx, cy, cz = Nx // 2, Ny // 2, Nz // 2 +radius = 5 + +x, y, z = np.indices(T0.shape) + +mask = ( + (x - cx)**2 + + (y - cy)**2 + + (z - cz)**2 + <= radius**2 +) + +S[mask] = 1.0 # K/s + +# The solver accepts volumetric heat deposition Q [W/m^3]. +# Convert the desired temperature source S [K/s] using rho * cp. +rho = 1000.0 # kg/m^3 +cp = 4000.0 # J/(kg K) +Q = S * rho * cp + + +# ---------------------------------------------------------------------- +# Material properties +# ---------------------------------------------------------------------- + +D = 1.4e-7 # m^2/s +P = 0.01 # 1/s +Ta = 37.0 # degC + +medium = DiffusionMedium( + diffusion_coeff=D, + density=rho, + specific_heat=cp, + perfusion_coeff=P, + blood_ambient_temperature=Ta, +) + +source = DiffusionSource( + T0=T0, + Q=Q, +) + + +# ---------------------------------------------------------------------- +# CPU / CUDA selection +# ---------------------------------------------------------------------- + +# Uses CUDA when CuPy and a CUDA device are available; otherwise CPU. +try: + import cupy as cp + + device = "gpu" if cp.cuda.runtime.getDeviceCount() > 0 else "cpu" +except (ImportError, RuntimeError): + device = "cpu" + +print(f"Using device: {device}") + + +# ---------------------------------------------------------------------- +# Run simulation +# ---------------------------------------------------------------------- + +t = 20.0 # total simulation time [s] + +solver = kWaveDiffusion( + kgrid, + medium, + source, + use_kspace=True, + dtype=np.float64, + display_updates=True, + device=device, +) + +# A single k-space-corrected time step advances the homogeneous +# solution over the full interval. +T = solver.take_time_step( + Nt=1, + dt=t, +) + + +# ---------------------------------------------------------------------- +# Results +# ---------------------------------------------------------------------- + +# Convert to NumPy for reporting/plotting when running on CUDA. +T = cp.asnumpy(T) + +print("Temperature shape:", T.shape) +print("Minimum temperature:", T.min(), "degC") +print("Maximum temperature:", T.max(), "degC") +print("Maximum temperature rise:", T.max() - T0.max(), "degC") +print("Maximum CEM43:", cp.asnumpy(solver.cem43).max()) +print("Lesion volume:", cp.asnumpy(solver.lesion_size), "m^3") From 08e0b2483eaa0b824e55666bfa083573bab1623e Mon Sep 17 00:00:00 2001 From: rhh006 Date: Wed, 4 Nov 2026 11:11:23 -0800 Subject: [PATCH 3/5] Rename kwaveDiffusion_test.py to test_kwaveDiffusion.py --- tests/{kwaveDiffusion_test.py => test_kwaveDiffusion.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{kwaveDiffusion_test.py => test_kwaveDiffusion.py} (100%) diff --git a/tests/kwaveDiffusion_test.py b/tests/test_kwaveDiffusion.py similarity index 100% rename from tests/kwaveDiffusion_test.py rename to tests/test_kwaveDiffusion.py From 66aa1957340e5021e69bdc717818bd87b02478fa Mon Sep 17 00:00:00 2001 From: rhh006 Date: Sun, 8 Nov 2026 21:18:58 -0800 Subject: [PATCH 4/5] Refactor test_kwaveDiffusion for improved structure --- tests/test_kwaveDiffusion.py | 397 +++++++++++++++++++++++++++-------- 1 file changed, 312 insertions(+), 85 deletions(-) diff --git a/tests/test_kwaveDiffusion.py b/tests/test_kwaveDiffusion.py index 2c9fa485..24f688ca 100644 --- a/tests/test_kwaveDiffusion.py +++ b/tests/test_kwaveDiffusion.py @@ -1,120 +1,347 @@ + import numpy as np +import pytest from kwave.kgrid import kWaveGrid -from kwaveDiffusion import DiffusionMedium, DiffusionSource, kWaveDiffusion +from kwave.kwaveDiffusion import ( + DiffusionMedium, + DiffusionSource, + kWaveDiffusion, +) -# ---------------------------------------------------------------------- -# Grid -# ---------------------------------------------------------------------- -Nx, Ny, Nz = 64, 64, 64 -dx = 0.5e-3 +@pytest.fixture +def grid(): + N = 32 + dx = 0.5e-3 + return kWaveGrid([N], [dx]) -kgrid = kWaveGrid([Nx, Ny, Nz], [dx, dx, dx]) +def to_numpy(value): + if hasattr(value, "get"): + return value.get() -# ---------------------------------------------------------------------- -# Initial temperature and heat source -# ---------------------------------------------------------------------- + return np.asarray(value) -T0 = np.full((Nx, Ny, Nz), 37.0, dtype=np.float64) +def test_uniform_heating(grid): -S = np.zeros_like(T0) + rho = 1000.0 + specific_heat = 4000.0 -# Small spherical-ish heating region -cx, cy, cz = Nx // 2, Ny // 2, Nz // 2 -radius = 5 + heating_rate = 1.0 + dt = 2.0 -x, y, z = np.indices(T0.shape) + medium = DiffusionMedium( + diffusion_coeff=1.4e-7, + density=rho, + specific_heat=specific_heat, + ) -mask = ( - (x - cx)**2 - + (y - cy)**2 - + (z - cz)**2 - <= radius**2 -) + source = DiffusionSource( + T0=37.0, + Q=heating_rate * rho * specific_heat, + ) -S[mask] = 1.0 # K/s + solver = kWaveDiffusion( + grid, + medium, + source, + use_kspace=True, + display_updates=False, + ) -# The solver accepts volumetric heat deposition Q [W/m^3]. -# Convert the desired temperature source S [K/s] using rho * cp. -rho = 1000.0 # kg/m^3 -cp = 4000.0 # J/(kg K) -Q = S * rho * cp + T = to_numpy(solver.take_time_step(1, dt)) + expected = 37.0 + heating_rate * dt -# ---------------------------------------------------------------------- -# Material properties -# ---------------------------------------------------------------------- + np.testing.assert_allclose( + T, + expected, + rtol=1e-10, + atol=1e-10, + ) -D = 1.4e-7 # m^2/s -P = 0.01 # 1/s -Ta = 37.0 # degC -medium = DiffusionMedium( - diffusion_coeff=D, - density=rho, - specific_heat=cp, - perfusion_coeff=P, - blood_ambient_temperature=Ta, -) +def test_fourier_mode_decay(grid): -source = DiffusionSource( - T0=T0, - Q=Q, -) + N = 32 + dx = 0.5e-3 + D = 1.4e-7 -# ---------------------------------------------------------------------- -# CPU / CUDA selection -# ---------------------------------------------------------------------- + x = np.arange(N) -# Uses CUDA when CuPy and a CUDA device are available; otherwise CPU. -try: - import cupy as cp + mode = 3 - device = "gpu" if cp.cuda.runtime.getDeviceCount() > 0 else "cpu" -except (ImportError, RuntimeError): - device = "cpu" + T0 = ( + 37.0 + + np.cos(2.0 * np.pi * mode * x / N) + ) -print(f"Using device: {device}") + medium = DiffusionMedium( + diffusion_coeff=D, + ) + source = DiffusionSource(T0=T0) -# ---------------------------------------------------------------------- -# Run simulation -# ---------------------------------------------------------------------- + solver = kWaveDiffusion( + grid, + medium, + source, + use_kspace=True, + display_updates=False, + ) -t = 20.0 # total simulation time [s] + dt = 2.0 -solver = kWaveDiffusion( - kgrid, - medium, - source, - use_kspace=True, - dtype=np.float64, - display_updates=True, - device=device, -) + T = to_numpy(solver.take_time_step(1, dt)) -# A single k-space-corrected time step advances the homogeneous -# solution over the full interval. -T = solver.take_time_step( - Nt=1, - dt=t, -) + k = 2.0 * np.pi * mode / (N * dx) + + expected = ( + 37.0 + + np.cos(2.0 * np.pi * mode * x / N) + * np.exp(-D * k**2 * dt) + ) + + np.testing.assert_allclose( + T, + expected, + rtol=1e-9, + atol=1e-9, + ) + + +def test_scalar_and_constant_array_agree(grid): + + N = 32 + + D = 1.4e-7 + + T0 = 37.0 + np.cos( + 2.0 * np.pi * 3 * np.arange(N) / N + ) + + source = DiffusionSource(T0=T0) + + solver_scalar = kWaveDiffusion( + grid, + DiffusionMedium(diffusion_coeff=D), + source, + display_updates=False, + ) + + solver_array = kWaveDiffusion( + grid, + DiffusionMedium( + diffusion_coeff=np.full(N, D) + ), + source, + display_updates=False, + ) + + T_scalar = solver_scalar.take_time_step(1, 2.0) + T_array = solver_array.take_time_step(1, 2.0) + + np.testing.assert_allclose( + T_scalar, + T_array, + rtol=1e-10, + atol=1e-10, + ) + + +def test_sensor_recording_multiple_calls(grid): + + N = 32 + + mask = np.zeros(N) + mask[[5, 10, 15]] = 1 + + medium = DiffusionMedium( + diffusion_coeff=1.4e-7, + ) + + source = DiffusionSource(T0=37.0) + + solver = kWaveDiffusion( + grid, + medium, + source, + sensor=mask, + display_updates=False, + ) + + to_numpy(solver.take_time_step(1, 0.1)) + to_numpy(solver.take_time_step(2, 0.1)) + + assert solver.sensor_data.shape == (3, 3) + assert solver.time_steps_taken == 3 + + +def test_zero_source_can_be_activated(grid): + + rho = 1000.0 + specific_heat = 4000.0 + + medium = DiffusionMedium( + diffusion_coeff=1.4e-7, + density=rho, + specific_heat=specific_heat, + ) + + source = DiffusionSource( + T0=37.0, + Q=0.0, + ) + + solver = kWaveDiffusion( + grid, + medium, + source, + display_updates=False, + ) + + solver.Q = 1.0e6 + + T = to_numpy(solver.take_time_step(1, 1.0)) + + assert np.max(T) > 37.0 + + +def test_dt_limit_uncorrected(grid): + + D = 1.4e-7 + + solver = kWaveDiffusion( + grid, + DiffusionMedium(diffusion_coeff=D), + DiffusionSource(T0=37.0), + use_kspace=False, + display_updates=False, + ) + + assert np.isfinite(solver.dt_limit) + assert solver.dt_limit > 0 + +def test_nyquist_mode_decay(grid): + + N = 32 + D = 1.4e-7 + + x = np.arange(N) + + T0 = 37.0 + (-1.0) ** x + + solver = kWaveDiffusion( + grid, + DiffusionMedium( + diffusion_coeff=np.full(N, D) + ), + DiffusionSource(T0=T0), + display_updates=False, + ) + + T = to_numpy(solver.take_time_step(1, 2.0)) + + k_nyquist = np.pi / grid.dx + + expected = ( + 37.0 + + (-1.0) ** x + * np.exp(-D * k_nyquist**2 * 2.0) + ) + + np.testing.assert_allclose( + T, + expected, + rtol=1e-8, + atol=1e-8, + ) + +def test_dt_limit_zero_reference(grid): + + D = 1.4e-7 + + solver = kWaveDiffusion( + grid, + DiffusionMedium( + diffusion_coeff=D, + diffusion_coeff_ref=0.0, + ), + DiffusionSource(T0=37.0), + use_kspace=True, + display_updates=False, + ) + + expected = 2.0 / ( + D * np.max(solver.k) ** 2 + ) + + np.testing.assert_allclose( + solver.dt_limit, + expected, + rtol=1e-10, + ) + +def test_cpu_backend(grid): + + solver = kWaveDiffusion( + grid, + DiffusionMedium( + diffusion_coeff=1.4e-7, + ), + DiffusionSource(T0=37.0), + device="cpu", + display_updates=False, + ) + + T = to_numpy(solver.take_time_step(1, 1.0)) + + assert isinstance(T, np.ndarray) + + assert np.all(np.isfinite(T)) + +def test_temperature_remains_constant(grid): + solver = kWaveDiffusion( + grid, + DiffusionMedium( + diffusion_coeff=1.4e-7, + ), + DiffusionSource(T0=37.0), + display_updates=False, + ) + + T = to_numpy(solver.take_time_step(10, 0.1)) + + np.testing.assert_allclose( + T, + 37.0, + atol=1e-12, + ) -# ---------------------------------------------------------------------- -# Results -# ---------------------------------------------------------------------- +def test_heterogeneous_diffusion_changes_temperature(grid): -# Convert to NumPy for reporting/plotting when running on CUDA. -T = cp.asnumpy(T) + N = 32 + D = np.full(N, 1.4e-7) + D[8:16] = 2.8e-7 + + x = np.arange(N) + + T0 = 37.0 + np.cos( + 2.0 * np.pi * 3 * x / N + ) + + solver = kWaveDiffusion( + grid, + DiffusionMedium(diffusion_coeff=D), + DiffusionSource(T0=T0), + display_updates=False, + ) + + T = to_numpy(solver.take_time_step(1, 2.0)) -print("Temperature shape:", T.shape) -print("Minimum temperature:", T.min(), "degC") -print("Maximum temperature:", T.max(), "degC") -print("Maximum temperature rise:", T.max() - T0.max(), "degC") -print("Maximum CEM43:", cp.asnumpy(solver.cem43).max()) -print("Lesion volume:", cp.asnumpy(solver.lesion_size), "m^3") + assert np.all(np.isfinite(T)) + assert not np.allclose(T, T0) From 9c7ed5477ee1697fb600133418783a6178616661 Mon Sep 17 00:00:00 2001 From: rhh006 Date: Sun, 8 Nov 2026 21:19:16 -0800 Subject: [PATCH 5/5] Implement uniformity checks for diffusion and perfusion Added methods to check uniformity of diffusion and perfusion coefficients, and updated the source setup to handle material properties more robustly. --- kwave/kwaveDiffusion.py | 304 +++++++++++++++++++++++----------------- 1 file changed, 173 insertions(+), 131 deletions(-) diff --git a/kwave/kwaveDiffusion.py b/kwave/kwaveDiffusion.py index 36bc6600..17c7852f 100644 --- a/kwave/kwaveDiffusion.py +++ b/kwave/kwaveDiffusion.py @@ -175,6 +175,17 @@ class kWaveDiffusion: implemented here because those require the k-Wave DTT operators. """ + def _is_uniform(self, value): + + value = self.xp.asarray(value) + + if value.size == 0: + return False + + return bool( + self.xp.all(value == value.flat[0]) + ) + def __init__( self, kgrid, @@ -268,11 +279,18 @@ def __init__( # Determine whether medium is heterogeneous # ------------------------------------------------------------ + self.flag_homogeneous_diffusion = ( + self._is_uniform(self.diffusion_p1) + and self._is_uniform(self.diffusion_p2) + ) + + self.flag_homogeneous_perfusion = ( + self._is_uniform(self.perfusion_coeff) + ) + self.flag_homogeneous = ( - self.diffusion_p1.size == 1 - and self.diffusion_p2.size == 1 - and self.perfusion_coeff.size == 1 - and self.blood_ambient_temperature.size == 1 + self.flag_homogeneous_diffusion + and self.flag_homogeneous_perfusion ) # ------------------------------------------------------------ @@ -516,7 +534,7 @@ def _as_field( return arr def _setup_diffusion(self, medium): - + # ------------------------------------------------------------ # Either D or rho/c/K # ------------------------------------------------------------ @@ -533,6 +551,11 @@ def _setup_diffusion(self, medium): "diffusion_coeff", ) + if bool(self.xp.any(self.diffusion_p2 < 0)): + raise ValueError( + "diffusion_coeff must be non-negative." + ) + else: if ( @@ -561,6 +584,17 @@ def _setup_diffusion(self, medium): "thermal_conductivity", ) + if bool(self.xp.any(density <= 0)): + raise ValueError("density must be positive.") + + if bool(self.xp.any(specific_heat <= 0)): + raise ValueError("specific_heat must be positive.") + + if bool(self.xp.any(conductivity <= 0)): + raise ValueError( + "thermal_conductivity must be positive." + ) + self.diffusion_p1 = ( 1.0 / (density * specific_heat) ) @@ -673,57 +707,52 @@ def _setup_perfusion(self, medium): # ================================================================== + def _setup_source(self, source): - """ - Q is retained in W/m^3. - - If Q is supplied, the update term is: - - Q / (rho * c) - - when rho*c is known. - - When diffusion_coeff was supplied, density and specific_heat - are therefore required for physical Q scaling. - """ - + if source.Q is None: source.Q = 0.0 - - if bool(self.xp.all( - self.xp.asarray(source.Q) == 0 - )): - return - - # If density and specific heat are available, - # calculate 1/A. - if ( + + has_material_properties = ( self.medium.density is not None and self.medium.specific_heat is not None - ): - + ) + + if has_material_properties: + density = self._as_field( self.medium.density, "density", ) - + specific_heat = self._as_field( self.medium.specific_heat, "specific_heat", ) - + + if bool(self.xp.any(density <= 0)): + raise ValueError("density must be positive.") + + if bool(self.xp.any(specific_heat <= 0)): + raise ValueError( + "specific_heat must be positive." + ) + self.q_scale_factor = ( 1.0 / (density * specific_heat) ) - + else: - - # MATLAB requires these properties whenever Q - # is defined. - raise ValueError( - "density and specific_heat must be supplied " - "when source.Q is nonzero." - ) + + self.q_scale_factor = None + + if bool(self.xp.any( + self.xp.asarray(source.Q) != 0 + )): + raise ValueError( + "density and specific_heat must be " + "supplied when source.Q is nonzero." + ) # ================================================================== # REFERENCE COEFFICIENT @@ -970,34 +999,37 @@ def _derivative_operators(self, kappa): def _source_update(self, kappa): - if np.isscalar(self.Q): - - if self.Q == 0: - return 0.0 - - else: - - if bool(self.xp.all(self.Q == 0)): + Q = self.xp.asarray( + self.Q, + dtype=self.dtype, + ) + + if Q.ndim == 0: + if Q.item() == 0: return 0.0 - - Q_ft = self.FT(self.Q) - + + Q = self.xp.full( + self.shape, + Q.item(), + dtype=self.dtype, + ) + + elif not bool(self.xp.any(Q != 0)): + return 0.0 + + if self.q_scale_factor is None: + raise ValueError( + "A nonzero heat source requires density " + "and specific_heat." + ) + + Q_ft = self.FT(Q) + result = self.IT( kappa * Q_ft ) - - # Physical source: - # - # Q / A - # - # where A = rho*c - - result = ( - self.q_scale_factor - * result - ) - - return result + + return self.q_scale_factor * result # ================================================================== # PERFUSION TERM @@ -1093,73 +1125,50 @@ def _heterogeneous_diffusion_update( where the derivative operators contain sqrt(kappa). """ - + T_ft = self.FT(self.T) - + d_term = 0.0 - - # ------------------------------------------------------------ - # X - # ------------------------------------------------------------ - - temp_x = self.IT( + + # X direction + temp_x = self.xp.fft.ifftn( deriv_x * T_ft ) - - flux_x = ( - self.diffusion_p2 - * temp_x - ) - - d_term = d_term + self.IT( - deriv_x - * self.FT(flux_x) - ) - - # ------------------------------------------------------------ - # Y - # ------------------------------------------------------------ - + + flux_x = self.diffusion_p2 * temp_x + + d_term += self.xp.fft.ifftn( + deriv_x * self.FT(flux_x) + ).real + + # Y direction if self.dim >= 2: - - temp_y = self.IT( + + temp_y = self.xp.fft.ifftn( deriv_y * T_ft ) - - flux_y = ( - self.diffusion_p2 - * temp_y - ) - - d_term = d_term + self.IT( - deriv_y - * self.FT(flux_y) - ) - - # ------------------------------------------------------------ - # Z - # ------------------------------------------------------------ - + + flux_y = self.diffusion_p2 * temp_y + + d_term += self.xp.fft.ifftn( + deriv_y * self.FT(flux_y) + ).real + + # Z direction if self.dim >= 3: - - temp_z = self.IT( + + temp_z = self.xp.fft.ifftn( deriv_z * T_ft ) + + flux_z = self.diffusion_p2 * temp_z + + d_term += self.xp.fft.ifftn( + deriv_z * self.FT(flux_z) + ).real + + return self.diffusion_p1 * d_term - flux_z = ( - self.diffusion_p2 - * temp_z - ) - - d_term = d_term + self.IT( - deriv_z - * self.FT(flux_z) - ) - - return ( - self.diffusion_p1 - * d_term - ) # ================================================================== # TIME STEP @@ -1201,6 +1210,8 @@ def take_time_step( "dt must be positive." ) + sensor_start = self.time_steps_taken + if self.display_updates: print( @@ -1335,34 +1346,39 @@ def take_time_step( if self.num_sensor_points > 0: - values = self.T.ravel( - order="C" - )[ - self.sensor_mask_index - ] + required_steps = sensor_start + Nt - if self.sensor_data is None: + if ( + self.sensor_data is None + or self.sensor_data.shape[1] < required_steps + ): + + old_data = self.sensor_data self.sensor_data = self.xp.empty( ( self.num_sensor_points, - Nt, + required_steps, ), dtype=self.dtype, ) + if old_data is not None: + self.sensor_data[:, :old_data.shape[1]] = old_data + + values = self.T.ravel(order="C")[ + self.sensor_mask_index + ] + self.sensor_data[ :, - t_index, + sensor_start + t_index, ] = values self.time_steps_taken += Nt if self.display_updates: - - print( - "Simulation completed." - ) + print("Simulation completed.") return self.T @@ -1480,6 +1496,32 @@ def dt_limit(self): ) P = self.perfusion_coeff + + # Uncorrected explicit diffusion/perfusion update. + # + # Applies when k-space correction is disabled or the + # reference coefficients are zero. + + if ( + not self.use_kspace + or ( + self.diffusion_coeff_ref == 0 + and self.perfusion_coeff_ref == 0 + ) + ): + + max_rate = float( + self.xp.max( + diffusion_coeff * self.k_squared + + P + ).item() + ) + + if max_rate <= 0: + return np.inf + + return 2.0 / max_rate + # ------------------------------------------------------------ # No perfusion