Skip to content

Fix kWaveMedium type annotations (rebased #636 + Greptile fixes) #736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
waltsims merged 2 commits into from
May 16, 2026
Merged

Fix kWaveMedium type annotations (rebased #636 + Greptile fixes) #736

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8e33fc7
Fix kWaveMedium type annotations and broaden test coverage
waltsims May 16, 2026
2c9b66a
Merge branch 'master' into chore/636-kmedium-type-annotations-rebased
waltsims May 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 40 additions & 36 deletions kwave/kmedium.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import logging
from dataclasses import dataclass
from typing import List, Optional, Union
from typing import List, Optional, Sequence, Union

import numpy as np

import kwave.utils.checks
from kwave.enums import AlphaMode


Expand All @@ -16,49 +15,51 @@ def _to_alpha_mode(value):
return AlphaMode(value)
except (ValueError, TypeError):
raise ValueError(
f"medium.alpha_mode must be an AlphaMode enum value or one of " f"'no_absorption', 'no_dispersion', 'stokes', got {value!r}"
f"medium.alpha_mode must be an AlphaMode enum value or one of 'no_absorption', 'no_dispersion', 'stokes', got {value!r}"
) from None


@dataclass
class kWaveMedium(object):
"""
Medium properties for k-Wave simulations.

Note: For heterogeneous medium parameters, medium.sound_speed and medium.density
must be given in matrix form with the same dimensions as kgrid. For homogeneous
medium parameters, these can be given as single numeric values. If the medium is
homogeneous and velocity inputs or outputs are not required, it is not necessary
to specify medium.density.
"""

# sound speed distribution within the acoustic medium [m/s] | required to be defined
sound_speed: np.array
sound_speed: Union[float, int, np.ndarray]
# reference sound speed used within the k-space operator (phase correction term) [m/s]
sound_speed_ref: np.array = None
sound_speed_ref: Optional[Union[float, int, np.ndarray]] = None
# density distribution within the acoustic medium [kg/m^3]
density: np.array = None
density: Optional[Union[float, int, np.ndarray]] = None
# power law absorption coefficient [dB/(MHz^y cm)]
alpha_coeff: np.array = None
alpha_coeff: Optional[Union[float, int, np.ndarray]] = None
# power law absorption exponent
alpha_power: np.array = None
alpha_power: Optional[Union[float, int, np.ndarray]] = None
# optional input to force either the absorption or dispersion terms in the equation of state to be excluded;
# valid inputs are AlphaMode.NO_ABSORPTION, AlphaMode.NO_DISPERSION, or the equivalent strings
alpha_mode: Optional[Union[AlphaMode, str]] = None
# frequency domain filter applied to the absorption and dispersion terms in the equation of state
alpha_filter: np.array = None
alpha_filter: Optional[np.ndarray] = None
# two element array used to control the sign of absorption and dispersion terms in the equation of state
alpha_sign: np.array = None
alpha_sign: Optional[np.ndarray] = None
# parameter of nonlinearity
BonA: np.array = None
BonA: Optional[Union[float, int, np.ndarray]] = None
# is the medium absorbing?
absorbing: bool = False
# is the medium absorbing stokes?
stokes: bool = False

# """
# Note: For heterogeneous medium parameters, medium.sound_speed and
# medium.density must be given in matrix form with the same dimensions as
# kgrid. For homogeneous medium parameters, these can be given as single
# numeric values. If the medium is homogeneous and velocity inputs or
# outputs are not required, it is not necessary to specify medium.density.
# """

def __post_init__(self):
self.sound_speed = np.atleast_1d(self.sound_speed)
self.alpha_mode = _to_alpha_mode(self.alpha_mode)

def check_fields(self, kgrid_shape: np.ndarray) -> None:
def check_fields(self, kgrid_shape: Sequence[int]) -> None:
"""
Check whether the given properties are valid

Expand All @@ -72,17 +73,17 @@ def check_fields(self, kgrid_shape: np.ndarray) -> None:
self.alpha_mode = _to_alpha_mode(self.alpha_mode)

# check the absorption filter input is valid
if self.alpha_filter is not None and not (self.alpha_filter.shape == kgrid_shape).all():
if self.alpha_filter is not None and self.alpha_filter.shape != tuple(kgrid_shape):
raise ValueError("medium.alpha_filter must be the same size as the computational grid.")

# check the absorption sign input is valid
if self.alpha_sign is not None and (not kwave.utils.checkutils.is_number(self.alpha_sign) or (self.alpha_sign.size != 2)):
raise ValueError(
"medium.alpha_sign must be given as a " "2 element numerical array controlling absorption and dispersion, respectively."
)
if self.alpha_sign is not None:
alpha_sign_arr = np.atleast_1d(self.alpha_sign)
if alpha_sign_arr.size != 2 or not np.issubdtype(alpha_sign_arr.dtype, np.number):
raise ValueError("medium.alpha_sign must be a 2 element numeric array controlling absorption and dispersion, respectively.")

# check alpha_coeff is non-negative and real
if not np.all(np.isreal(self.alpha_coeff)) or np.any(self.alpha_coeff < 0):
if self.alpha_coeff is not None and (not np.all(np.isreal(self.alpha_coeff)) or np.any(self.alpha_coeff < 0)):
raise ValueError("medium.alpha_coeff must be non-negative and real.")

def is_defined(self, *fields) -> List[bool]:
Expand Down Expand Up @@ -111,7 +112,7 @@ def ensure_defined(self, *fields) -> None:
None
"""
for f in fields:
assert getattr(self, f) is not None, f"The field {f} must be not be None"
assert getattr(self, f) is not None, f"The field {f} must not be None"

def is_nonlinear(self) -> bool:
"""
Expand Down Expand Up @@ -150,13 +151,14 @@ def _check_absorbing_without_stokes(self) -> None:
# enforce both absorption parameters
self.ensure_defined("alpha_coeff", "alpha_power")

# check y is a scalar
assert np.isscalar(self.alpha_power), "medium.alpha_power must be scalar."
# check y is a scalar (np.isscalar rejects 0-d ndarrays — accept np.array(1.5) too)
is_scalar = np.isscalar(self.alpha_power) or (isinstance(self.alpha_power, np.ndarray) and self.alpha_power.size == 1)
Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps Bot May 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 is_scalar silently accepts 1-D single-element arrays

The condition self.alpha_power.size == 1 admits np.array([1.5]) (a 1-D ndarray), not just a 0-d array. A 1-D array passes the check, then propagates through the remainder of _check_absorbing_without_stokes without being reduced to a true scalar — 0 <= np.array([1.5]) < 3 works at the assertion level due to Python/NumPy single-element coercion, but downstream simulation code that expects a bare scalar for the power-law exponent may see unexpected shapes. Consider extracting the scalar value with .item() once is_scalar has been confirmed, so self.alpha_power is always a Python float at that point.

assert is_scalar, "medium.alpha_power must be scalar."

# check y is real and within 0 to 3
assert (
np.all(np.isreal(self.alpha_coeff)) and 0 <= self.alpha_power < 3
), "medium.alpha_power must be a real number between 0 and 3."
assert np.all(np.isreal(self.alpha_coeff)) and 0 <= self.alpha_power < 3, (
"medium.alpha_power must be a real number between 0 and 3."
)

# display warning if y is close to 1 and the dispersion term has not been set to zero
if self.alpha_mode != "no_dispersion":
Expand All @@ -176,21 +178,23 @@ def _check_absorbing_with_stokes(self):
self.ensure_defined("alpha_coeff")

# give warning if y is specified
if self.alpha_power is not None and (self.alpha_power.size != 1 or self.alpha_power != 2):
logging.log(logging.WARN, "the axisymmetric code and stokes absorption assume alpha_power = 2, user value ignored.")
if self.alpha_power is not None:
ap = np.asarray(self.alpha_power)
if ap.size != 1 or not np.isclose(ap.item(), 2.0):
logging.warning("the axisymmetric code and stokes absorption assume alpha_power = 2, user value ignored.")

# overwrite y value
self.alpha_power = 2

# don't allow medium.alpha_mode with the axisymmetric code
if self.alpha_mode is not None and (self.alpha_mode in ["no_absorption", "no_dispersion"]):
raise NotImplementedError(
"Input option medium.alpha_mode is not supported with the axisymmetric code " "or medium.alpha_mode = " "stokes" "."
"Input option medium.alpha_mode is not supported with the axisymmetric code or medium.alpha_mode = stokes."
)

# don't allow alpha_filter with stokes absorption (no variables are applied in k-space)
assert self.alpha_filter is None, (
"Input option medium.alpha_filter is not supported with the axisymmetric code " "or medium.alpha_mode = 'stokes'. "
"Input option medium.alpha_filter is not supported with the axisymmetric code or medium.alpha_mode = 'stokes'. "
)

##########################################
Expand Down
Loading
Loading