fix: correct type annotations in kWaveMedium class #636
Conversation
- Replace incorrect np.array type hints with proper Union types - Add support for scalar (float/int) and array (np.ndarray) inputs - Use Optional[] for parameters with None defaults - Improve IDE type checking and autocompletion support Addresses issue where IDE shows 'Expected type None, got float' warnings for alpha_coeff and alpha_power parameters.
✅ Actions performedReview triggered.
|
WalkthroughBroadened kWaveMedium type annotations to accept scalars, arrays, or None for several fields; added a kWaveMedium docstring; changed check_fields to accept Sequence[int]; updated alpha-related validations and axisymmetric/stokes handling; and added a comprehensive test suite for kWaveMedium. Changes
Sequence Diagram(s)sequenceDiagram
participant Test as Test / Caller
participant kMedium as kWaveMedium.__post_init__
participant Validator as check_fields / validators
participant Logger as logging
Note over kMedium,Validator: Initialization → normalization → validation
Test->>kMedium: instantiate with fields (scalar/array/None)
kMedium->>kMedium: normalize scalars to 1-element np.array
kMedium->>Validator: check_fields(kgrid_shape)
Validator->>Validator: validate alpha_filter shape == tuple(kgrid_shape)
Validator->>Validator: validate alpha_sign is numeric array of length 2
alt alpha_coeff provided
Validator->>Validator: validate alpha_coeff (non-negative, real)
end
alt axisymmetric & is_stokes
Validator->>Validator: coerce alpha_power, if not np.isclose(..., 2.0)
Validator->>Logger: logging.warning("alpha_power != 2, setting to 2")
Validator->>kMedium: set alpha_power = 2
else absorbing (non-stokes)
Validator->>Validator: assert alpha_power is scalar and within [0,3]
end
Validator-->>kMedium: validations complete
kMedium-->>Test: ready or raise AssertionError / ValueError
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Comment |
🧩 Analysis chainI'll examine the changes in this PR to provide a comprehensive review of the type annotation fixes in the kWaveMedium class. 🏁 Script executed: Length of output: 390 🏁 Script executed: Length of output: 9407 @waltsims I've reviewed the type annotation fixes in the kWaveMedium class. Here's my analysis: ✅ Excellent Improvements MadeThe type annotation fixes are well implemented:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
kwave/kmedium.py (1)
77-78: Potential null pointer exception with new Optional types.The validation logic directly accesses
self.alpha_coeffwithout checking if it'sNone, which could cause runtime errors sincealpha_coeffis nowOptional[Union[float, int, np.ndarray]].Apply this diff to add null checking:
- # check alpha_coeff is non-negative and real - if 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.") + # check alpha_coeff is non-negative and real + 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.")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
kwave/kmedium.py(2 hunks)
🧰 Additional context used
🪛 GitHub Actions: Ruff
kwave/kmedium.py
[error] 1-1: I001 Import block is un-sorted or un-formatted.
🔇 Additional comments (6)
kwave/kmedium.py (6)
14-14: LGTM! Improved type annotation for sound_speed.The change from
np.array(which is a function) toUnion[float, int, np.ndarray]correctly reflects that this field can accept scalar values or arrays, resolving IDE warnings.
16-16: LGTM! Consistent type improvements for numeric parameters.The updates to
sound_speed_ref,density,alpha_coeff,alpha_power, andBonAproperly useOptional[Union[float, int, np.ndarray]]withNonedefaults, allowing both scalar and array inputs while maintaining optional behavior.Also applies to: 18-18, 20-20, 22-22, 31-31
25-25: LGTM! Correct string type for alpha_mode.The change from
np.array = NonetoOptional[str] = Noneis correct sincealpha_modeaccepts string values like 'no_absorption' or 'no_dispersion' as seen in the validation logic.
27-27: LGTM! Appropriate array-only types for filter parameters.Using
Optional[np.ndarray]foralpha_filterandalpha_signis correct since these parameters specifically require array structures based on the validation logic (shape checking and size requirements).Also applies to: 29-29
35-35: Document and confirm newstokesboolean (Stokes absorption / axisymmetric)Verified: medium.stokes is set in kwave/kmedium.py (set_absorbing) and referenced in kwave/kWaveSimulation.py (selects "stokes" vs "absorbing") and kwave/kWaveSimulation_helper/save_to_disk_func.py (sets absorbing_flag=2). Update the PR description to explain this field, why the default is False, when it should be enabled (axisymmetric or alpha_mode == "stokes"), and any backward-compat/serialization implications.
46-46: Verify sound_speed handling in post_initnp.atleast_1d accepts scalar int/float and returns a 1D ndarray; sandbox verification failed because numpy is not installed. Run locally to confirm:
python3 - <<'PY'
import numpy as np
for v in (1500.0, 1500, [1500,1600], np.array([1500,1600])):
print(repr(v), '->', np.atleast_1d(v))
PY
- Consolidate typing imports to fix Ruff linter errors - Add None check for alpha_coeff validation to prevent runtime errors - Resolves CodeRabbit review comments
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
kwave/kmedium.py (3)
65-67: Bug: alpha_filter shape check can raise AttributeError and mis-compare.Comparing tuple (alpha_filter.shape) to kgrid_shape (often a tuple) and then calling .all() on a bool will fail; also tuple==ndarray semantics are wrong.
- if self.alpha_filter is not None and not (self.alpha_filter.shape == kgrid_shape).all(): - raise ValueError("medium.alpha_filter must be the same size as the computational grid.") + 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.")
70-73: Bug: wrong module name and fragile numeric check for alpha_sign.
kwave.utils.checkutilsdoesn’t exist here (import is checks). Also, validating arrays via is_number isn’t appropriate.- 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." + )
170-175: Bug: attribute access on scalar and deprecated logging level.
self.alpha_power.sizebreaks for Python scalars; use np.size. Also prefer logging.warning and tolerant equality for floats.- 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(float(ap), 2.0): + logging.warning("the axisymmetric code and stokes absorption assume alpha_power = 2, user value ignored.")
♻️ Duplicate comments (1)
kwave/kmedium.py (1)
3-3: Import consolidation fixed; prior pipeline issue resolved.Single, sorted typing import looks good.
🧹 Nitpick comments (3)
kwave/kmedium.py (3)
3-3: Prefer Sequence[int] for kgrid_shape to reflect actual usage (tuple).This matches callers passing kgrid_k.shape and avoids confusion.
-from typing import List, Optional, Union +from typing import List, Optional, Sequence, Union @@ - def check_fields(self, kgrid_shape: np.ndarray) -> None: + def check_fields(self, kgrid_shape: Sequence[int]) -> None:Also applies to: 47-47
11-12: Promote the commented guidance to a real class docstring (and remove the commented block).Improves IDE/help introspection and keeps the class self-documented.
class kWaveMedium(object): + """ + 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. + """ @@ - # """ - # 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. - # """Also applies to: 36-42
105-105: Nit: fix grammar in assertion message.- 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"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
kwave/kmedium.py(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
kwave/kmedium.py (4)
kwave/kWaveSimulation.py (2)
check_medium(576-605)kWaveSimulation(34-1552)kwave/kWaveSimulation_helper/expand_grid_matrices.py (1)
expand_medium(106-127)kwave/kWaveSimulation_helper/create_absorption_variables.py (1)
create_stokes_medium_variables(33-39)kwave/kWaveSimulation_helper/save_to_disk_func.py (1)
grab_medium_props(214-241)
🔇 Additional comments (2)
kwave/kmedium.py (2)
13-31: Broadened type hints are correct and align with runtime conversions.Accepting float|int|ndarray for these fields matches downstream usage.
76-77: Nice guard added for optional alpha_coeff.The None-check avoids exceptions when alpha_coeff is unspecified.
- Fix alpha_filter shape comparison bug - Fix alpha_sign validation with proper numpy checks - Fix alpha_power scalar access using np.asarray() - Fix grammar error in assertion message - Improve type hints and add class docstring
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (8)
kwave/kmedium.py (8)
12-20: Docstring OK; consider expanding later for parameters.High-level class docstring is fine. A future pass could add per-field notes (e.g., units, homogeneous vs heterogeneous expectations) for IDE tooltips.
31-33: Fix alpha_mode docs: 'stokes' is also a valid value.The comment says valid inputs are 'no_absorption' or 'no_dispersion', but check_fields also accepts 'stokes'. Update to avoid misleading users.
Apply this diff:
- # valid inputs are 'no_absorption' or 'no_dispersion' + # valid inputs are 'no_absorption', 'no_dispersion', or 'stokes'
45-47: Add explicit return type to dunder for clarity.Minor typing polish: annotate post_init return as None.
- def __post_init__(self): + def __post_init__(self) -> None:
71-76: Tighten alpha_sign validation to enforce 1-D, exactly two elements.Current check allows any shape with size 2 (e.g., (2,1)). If you truly require a flat 2-vector, enforce 1-D shape.
- 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): + if self.alpha_sign is not None: + alpha_sign_arr = np.asarray(self.alpha_sign) + if alpha_sign_arr.ndim != 1 or alpha_sign_arr.shape != (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." )
97-109: Prefer explicit exceptions over assert for user input validation.Asserts can be stripped with Python -O, disabling critical validation. Consider raising ValueError/TypeError for public API checks like ensure_defined.
- for f in fields: - assert getattr(self, f) is not None, f"The field {f} must not be None" + for f in fields: + if getattr(self, f) is None: + raise ValueError(f"The field {f} must not be None")
147-154: Separate validations and fix misleading error message.This assert mixes alpha_coeff “real” with alpha_power range but the message mentions only alpha_power. Split checks and raise explicit exceptions.
- # check y is scalar - assert np.isscalar(self.alpha_power), "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." + # alpha_power must be a real scalar in [0, 3) + if not np.isscalar(self.alpha_power): + raise ValueError("medium.alpha_power must be a real scalar.") + if not (0 <= float(self.alpha_power) < 3): + raise ValueError("medium.alpha_power must be in the range [0, 3).") + # alpha_coeff should be real (non-negativity is checked elsewhere) + if not np.all(np.isreal(self.alpha_coeff)): + raise ValueError("medium.alpha_coeff must be real.")
182-185: Clean up message string concatenation.Join the error message as a single literal for readability.
- raise NotImplementedError( - "Input option medium.alpha_mode is not supported with the axisymmetric code " "or medium.alpha_mode = " "stokes" "." - ) + raise NotImplementedError( + "Input option medium.alpha_mode is not supported with the axisymmetric code or with medium.alpha_mode = 'stokes'." + )
188-190: Same here: simplify concatenated assert message.- assert self.alpha_filter is None, ( - "Input option medium.alpha_filter is not supported with the axisymmetric code " "or medium.alpha_mode = 'stokes'. " - ) + assert self.alpha_filter is None, ( + "Input option medium.alpha_filter is not supported with the axisymmetric code or with medium.alpha_mode = 'stokes'." + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
kwave/kmedium.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
kwave/kmedium.py (3)
kwave/kgrid.py (2)
size(297-301)kWaveGrid(14-701)kwave/kWaveSimulation_helper/expand_grid_matrices.py (1)
expand_medium(106-127)kwave/kWaveSimulation.py (2)
check_medium(576-605)kWaveSimulation(34-1552)
🔇 Additional comments (4)
kwave/kmedium.py (4)
22-40: Type-hint broadening looks correct and consistent.Allowing scalars, arrays, and None via Union/Optional for the public fields aligns with usage in the rest of the codebase and resolves prior type-checker friction for scalar inputs.
48-57: check_fields now accepts Sequence[int]: LGTM.Accepting Sequence[int] makes call sites with tuple shapes natural and keeps the function flexible.
67-68: alpha_filter shape validation: LGTM.Comparing to tuple(kgrid_shape) is robust and clear.
79-81: Good: alpha_coeff guard handles None before numeric checks.Prevents exceptions when alpha_coeff is omitted.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #636 +/- ##
==========================================
+ Coverage 73.95% 74.15% +0.20%
==========================================
Files 50 50
Lines 7000 7004 +4
Branches 1338 1340 +2
==========================================
+ Hits 5177 5194 +17
+ Misses 1280 1269 -11
+ Partials 543 541 -2
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- Fix NumPy deprecation warning in _check_absorbing_with_stokes by using ap.item() instead of float(ap) - Add comprehensive test suite for kWaveMedium class to improve code coverage - Test all edge cases including elastic properties, validation methods, and Stokes mode restrictions - All tests pass without warnings or errors
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (6)
tests/test_kmedium.py (6)
16-37: Parametrize elastic-property access checks to reduce repetition.Same assertion pattern across six properties; a small parametrize cleans it up.
Apply this refactor:
-def test_elastic_properties_access(): - """Test access to elastic code related properties (should raise NotImplementedError)""" - medium = kWaveMedium(sound_speed=1500) - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.sound_speed_shear - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.sound_speed_ref_shear - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.sound_speed_compression - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.sound_speed_ref_compression - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.alpha_coeff_compression - - with pytest.raises(NotImplementedError, match="Elastic simulation"): - _ = medium.alpha_coeff_shear +@pytest.mark.parametrize( + "attr", + [ + "sound_speed_shear", + "sound_speed_ref_shear", + "sound_speed_compression", + "sound_speed_ref_compression", + "alpha_coeff_compression", + "alpha_coeff_shear", + ], +) +def test_elastic_properties_access(attr): + """Elastic properties should raise NotImplementedError.""" + medium = kWaveMedium(sound_speed=1500) + with pytest.raises(NotImplementedError, match="Elastic simulation"): + getattr(medium, attr)
59-72: Tighten error-message matching (optional).Matching the full message guards against regressions.
-with pytest.raises(AssertionError, match="alpha_coeff must not be None"): +with pytest.raises(AssertionError, match=r"^The field alpha_coeff must not be None$") medium.ensure_defined('alpha_coeff') -with pytest.raises(AssertionError, match="alpha_power must not be None"): +with pytest.raises(AssertionError, match=r"^The field alpha_power must not be None$") medium.ensure_defined('alpha_power')
127-139: Also assert valid alpha_mode values pass.Covers the allowed enum surface.
def test_alpha_mode_validation_edge_cases(): @@ with pytest.raises(AssertionError): medium2.check_fields(kgrid.N) + +@pytest.mark.parametrize("mode", ["no_absorption", "no_dispersion", "stokes"]) +def test_alpha_mode_valid_values(mode): + kgrid = kWaveGrid(Vector([64, 64]), Vector([0.1e-3, 0.1e-3])) + medium = kWaveMedium(sound_speed=1500, alpha_mode=mode) + # Should not raise + medium.check_fields(kgrid.N)
210-236: Avoid direct mutation of flags; use the API.Keeps tests aligned with public surface and avoids redundant state changes.
- # Reset - medium.absorbing = False - medium.stokes = False - - # Set to Stokes absorbing + # Set to Stokes absorbing (no need to reset via direct assignment) medium.set_absorbing(is_absorbing=True, is_stokes=True) @@ - # Set to non-absorbing - medium.set_absorbing(is_absorbing=False, is_stokes=False) + # Set to non-absorbing + medium.set_absorbing(is_absorbing=False, is_stokes=False)
197-208: Consider testing that alpha_filter is forbidden in Stokes mode.Covers the assertion path in _check_absorbing_with_stokes.
def test_stokes_mode_alpha_mode_restrictions(): @@ with pytest.raises(NotImplementedError, match="not supported with the axisymmetric code"): medium2.set_absorbing(is_absorbing=True, is_stokes=True) + +def test_stokes_mode_rejects_alpha_filter(): + kgrid = kWaveGrid(Vector([64, 64]), Vector([0.1e-3, 0.1e-3])) + filt = np.zeros(tuple(kgrid.N)) + medium = kWaveMedium(sound_speed=1500, alpha_coeff=0.75, alpha_filter=filt) + with pytest.raises(AssertionError, match=r"alpha_filter.*not supported"): + medium.set_absorbing(is_absorbing=True, is_stokes=True)
238-240: Drop the pytest entry point.Tests are discovered by runners; this block isn’t needed.
-if __name__ == '__main__': - pytest.main([__file__, '-v']) +# No module entry point needed; use `pytest -q`.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
kwave/kmedium.py(5 hunks)tests/test_kmedium.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- kwave/kmedium.py
🧰 Additional context used
🧬 Code graph analysis (1)
tests/test_kmedium.py (2)
kwave/data.py (1)
Vector(7-49)kwave/kmedium.py (12)
kWaveMedium(11-237)sound_speed_shear(198-202)sound_speed_ref_shear(205-209)sound_speed_compression(212-216)sound_speed_ref_compression(219-223)alpha_coeff_compression(226-230)alpha_coeff_shear(233-237)is_defined(82-95)ensure_defined(97-108)is_nonlinear(110-117)set_absorbing(119-135)check_fields(48-80)
🔇 Additional comments (10)
tests/test_kmedium.py (10)
1-5: Nice, clear module docstring.Good context-setting for the suite.
7-14: Imports look tidy and minimal.No unnecessary deps; explicit module paths are good.
39-57: is_defined scenarios are thorough.Covers single and multi-field cases; matches the current API returning a list.
74-83: is_nonlinear coverage is crisp.Linear vs nonlinear paths validated.
85-92: Stokes + alpha_power=None path validated.Confirms coercion to 2; good.
157-163: LGTM.None passes as intended.
165-182: Alpha-coeff array validation tests are solid.Covers happy path, negatives, and complex values.
184-195: Post-init normalization assertions are good.Catches scalar-to-array promotion and shape preservation.
197-208: Restrictions in Stokes mode covered.Correctly targets NotImplementedError path.
94-105: Verified — patch('logging.warning') is correct.
kwave/kmedium.py imports logging and calls logging.warning(...) (line 176), so patch('logging.warning') will capture the call.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
- Reorder imports according to PEP 8 standards (stdlib, third-party, local) - Remove unnecessary __main__ block from test file - Fix I001 import block formatting issue
…r/k-wave-python into fix/type-annotations-kmedium
1 participant
Problem
The current type annotations in
kWaveMediumusenp.arraywhich is a function, not a type. This causes IDE warnings like "Expected type 'None', got 'float'" when using scalar values foralpha_coeffandalpha_power.Solution
np.arraytype hints with properUnion[float, int, np.ndarray]Optional[]for parameters that acceptNoneExample Usage