Fix alpha mode near unity

CLAUDE.md

@@ -0,0 +1,94 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What is k-wave-python?
+
+Python implementation of the [k-Wave](http://www.k-wave.org/) acoustics toolbox for time-domain acoustic and ultrasound simulations. Two backends: a pure Python/NumPy/CuPy solver and C++ binaries (OMP/CUDA).
+
+## Commands
+
+```bash
+# Setup
+uv sync --extra dev --extra test
+uv run pre-commit install
+
+# Run tests
+uv run pytest                              # all tests
+uv run pytest tests/test_kgrid.py          # single file
+uv run pytest tests/test_kgrid.py::test_name  # single test
+uv run pytest -m integration               # MATLAB-reference tests only
+
+# Lint/format
+uv run ruff check . --fix
+uv run ruff format .
+uv run pre-commit run --all-files
+
+# Run an example
+uv run examples/ivp_homogeneous_medium.py
+
+# Build
+uv build
+```
+
+Always use `uv run` (not `uv run python`) for pytest, examples, and scripts.
+
+## Code Style
+
+- Line length: 140 (configured in `pyproject.toml` under `[tool.ruff]`)
+- Pre-commit hooks: ruff (lint + format), codespell, nb-clean
+- Ruff ignores F821 (nested function refs) and F722 (jaxtyping annotations)
+- Type annotations use `beartype` + `jaxtyping`
+
+## Architecture
+
+### Entry point
+
+`kwave/kspaceFirstOrder.py` — `kspaceFirstOrder()` is the unified API. It accepts `kgrid`, `medium`, `source`, `sensor` and dispatches to the appropriate backend.
+
+```
+kspaceFirstOrder(kgrid, medium, source, sensor, backend="python"|"cpp", device="cpu"|"gpu")
+```
+
+### Two backends
+
+- **Python backend** (`kwave/solvers/kspace_solver.py`): `Simulation` class implementing k-space pseudospectral method in NumPy/CuPy. Supports 1D/2D/3D. Uses CuPy for GPU when `device="gpu"`.
+- **C++ backend** (`kwave/solvers/cpp_simulation.py` + `kwave/executor.py`): Serializes to HDF5, invokes compiled C++ binary, reads results back. `save_only=True` writes HDF5 without running (for cluster jobs).
+
+### Core data classes
+
+- `kWaveGrid` (`kwave/kgrid.py`) — domain discretization, spacing, time array
+- `kWaveMedium` (`kwave/kmedium.py`) — sound speed, density, absorption, nonlinearity
+- `kSource` (`kwave/ksource.py`) — pressure/velocity sources with masks and signals
+- `kSensor` (`kwave/ksensor.py`) — sensor mask and recording configuration
+
+### Legacy path
+
+`kspaceFirstOrder2D()` / `kspaceFirstOrder3D()` in their respective files route through `kWaveSimulation` (`kwave/kWaveSimulation.py`) — a large legacy dataclass used by the C++ backend path. New code should use `kspaceFirstOrder()` directly.
+
+### PML handling
+
+When `pml_inside=False` (default), `kspaceFirstOrder()` expands the grid by `2*pml_size` before simulation and strips PML from full-grid output fields afterward. `pml_size="auto"` selects optimal sizes via `get_optimal_pml_size()`.
+
+### Key utilities
+
+- `kwave/utils/pml.py` — PML sizing (`get_pml()`, `get_optimal_pml_size()`)
+- `kwave/utils/mapgen.py` — geometry generators (`make_disc`, `make_ball`, `make_cart_circle`, etc.)
+- `kwave/utils/signals.py` — signal generation (tone bursts, filtering)
+- `kwave/utils/filters.py` — spatial smoothing, Gaussian filters
+- `kwave/utils/io.py` — HDF5 read/write
+- `kwave/utils/conversion.py` — unit conversion, `cart2grid`
+
+## Testing
+
+- Tests in `tests/`, configured via `[tool.pytest.ini_options]` in `pyproject.toml`
+- Integration tests (`@pytest.mark.integration`) compare against MATLAB reference data
+- Test fixtures in `tests/integration/conftest.py`: `load_matlab_ref`, `assert_fields_close`
+- MATLAB reference data pointed to by `KWAVE_MATLAB_REF_DIR` env var
+
+## Naming Conventions
+
+- `backend="python"` or `"cpp"` (not "native")
+- `device="cpu"` or `"gpu"` (not `use_gpu` bool)
+- `quiet=True` to suppress output; `debug=True` for detailed output
+- `pml_size="auto"` for automatic PML sizing (not a separate boolean)

kwave/kWaveSimulation_helper/save_to_disk_func.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import warnings
 
 import numpy as np
 from scipy.io import savemat
@@ -229,6 +230,15 @@ def grab_medium_props(integer_variables, float_variables, medium, is_elastic_cod
         integer_variables.absorbing_flag = 0
 
     if medium.absorbing:
+        alpha_mode = getattr(medium, "alpha_mode", None)
+        if alpha_mode in ("no_absorption", "no_dispersion"):
+            warnings.warn(
+                f"medium.alpha_mode='{alpha_mode}' is not supported by the C++ backend "
+                f"and will be silently ignored. The C++ binary always computes both "
+                f"absorption and dispersion when absorbing_flag > 0. Use backend='python' "
+                f"to honor alpha_mode.",
+                stacklevel=4,
+            )
         if is_elastic_code:  # pragma: no cover
             # add to the variable list
             float_variables["chi"] = None

kwave/solvers/cpp_simulation.py

@@ -9,6 +9,7 @@
 import stat
 import subprocess
 import tempfile
+import warnings
 
 import numpy as np
 
@@ -36,6 +37,16 @@ def __init__(self, kgrid, medium, source, sensor, *, pml_size, pml_alpha, use_sg
         self.use_sg = use_sg
         self.ndim = kgrid.dim
 
+        alpha_mode = getattr(medium, "alpha_mode", None)
+        if alpha_mode in ("no_absorption", "no_dispersion"):
+            warnings.warn(
+                f"medium.alpha_mode='{alpha_mode}' is not supported by the C++ backend "
+                f"and will be silently ignored. The C++ binary always computes both "
+                f"absorption and dispersion when absorbing_flag > 0. Use backend='python' "
+                f"to honor alpha_mode.",
+                stacklevel=3,
+            )
+
     def prepare(self, data_path=None):
         """Write HDF5 input file. Returns (input_file, output_file)."""
         if data_path is None:
@@ -195,11 +206,14 @@ def _write_hdf5(self, filepath):
             "pml_x_size": pml_x_size,
             "pml_y_size": pml_y_size,
             "pml_z_size": pml_z_size,
-            "p_source_flag": int(has_p),
+            # NOTE: Despite the name "*_source_flag", the C++ binary expects
+            # the number of time points in the source signal (not a boolean).
+            # 0 = no source, >0 = number of time steps in the source signal.
+            "p_source_flag": np.asarray(source.p).shape[-1] if has_p else 0,
             "p0_source_flag": int(has_p0),
-            "ux_source_flag": int(has_ux),
-            "uy_source_flag": int(has_uy),
-            "uz_source_flag": int(has_uz),
+            "ux_source_flag": np.asarray(source.ux).shape[-1] if has_ux else 0,
+            "uy_source_flag": np.asarray(source.uy).shape[-1] if has_uy else 0,
+            "uz_source_flag": np.asarray(source.uz).shape[-1] if has_uz else 0,
             "sxx_source_flag": 0,
             "syy_source_flag": 0,
             "szz_source_flag": 0,

kwave/solvers/kspace_solver.py

@@ -411,6 +411,7 @@ def _setup_physics_operators(self):
 
         # Absorption/dispersion
         alpha_coeff_raw = getattr(self.medium, "alpha_coeff", 0)
+        alpha_mode = getattr(self.medium, "alpha_mode", None)
         if not _is_enabled(alpha_coeff_raw):
             self._absorption = lambda div_u: 0
             self._dispersion = lambda rho: 0
@@ -419,17 +420,23 @@ def _setup_physics_operators(self):
             alpha_power = float(xp.array(getattr(self.medium, "alpha_power", 1.5)).flatten()[0])
             alpha_np = 100 * alpha_coeff * (1e-6 / (2 * np.pi)) ** alpha_power / (20 * np.log10(np.e))
 
+            no_absorption = alpha_mode == "no_absorption"
+            no_dispersion = alpha_mode == "no_dispersion"
+
             if abs(alpha_power - 2.0) < 1e-10:  # Stokes
-                self._absorption = lambda div_u: -2 * alpha_np * self.c0 * self.rho0 * div_u
+                self._absorption = lambda div_u: 0 if no_absorption else -2 * alpha_np * self.c0 * self.rho0 * div_u
                 self._dispersion = lambda rho: 0
             else:  # Power-law with fractional Laplacian
                 tau = -2 * alpha_np * self.c0 ** (alpha_power - 1)
-                eta = 2 * alpha_np * self.c0**alpha_power * xp.tan(np.pi * alpha_power / 2)
                 nabla1 = self._fractional_laplacian(alpha_power - 2)
-                nabla2 = self._fractional_laplacian(alpha_power - 1)
-                # Fractional Laplacian already includes full k-space structure; no kappa needed
-                self._absorption = lambda div_u: tau * self._diff(self.rho0 * div_u, nabla1)
-                self._dispersion = lambda rho: eta * self._diff(rho, nabla2)
+                self._absorption = (lambda div_u: 0) if no_absorption else (lambda div_u: tau * self._diff(self.rho0 * div_u, nabla1))
+
+                if no_dispersion:
+                    self._dispersion = lambda rho: 0
+                else:
+                    eta = 2 * alpha_np * self.c0**alpha_power * xp.tan(np.pi * alpha_power / 2)
+                    nabla2 = self._fractional_laplacian(alpha_power - 1)
+                    self._dispersion = lambda rho: eta * self._diff(rho, nabla2)
 
         # Nonlinearity
         BonA_raw = getattr(self.medium, "BonA", 0)

kwave/solvers/validation.py

@@ -47,6 +47,14 @@ def validate_medium(medium, kgrid):
             power = float(np.asarray(medium.alpha_power).flat[0])
             if power < 0 or power > 3:
                 warnings.warn(f"medium.alpha_power={power} is outside typical range [0, 3].", stacklevel=3)
+            alpha_mode = getattr(medium, "alpha_mode", None)
+            if abs(power - 1.0) < 0.05 and alpha_mode != "no_dispersion":
+                raise ValueError(
+                    f"medium.alpha_power={power} is too close to 1.0. The dispersion term "
+                    f"contains tan(pi*alpha_power/2), which diverges at alpha_power=1. "
+                    f"Set medium.alpha_mode='no_dispersion' to disable the dispersion term, "
+                    f"or use an alpha_power value further from 1.0."
+                )
 
 
 def validate_pml(pml_size, kgrid):

tests/test_alpha_power_near_unity.py

@@ -0,0 +1,132 @@
+"""Tests for alpha_power near 1.0 (tan(pi*y/2) singularity)."""
+import numpy as np
+import pytest
+
+from kwave.kgrid import kWaveGrid
+from kwave.kmedium import kWaveMedium
+from kwave.ksensor import kSensor
+from kwave.ksource import kSource
+from kwave.kspaceFirstOrder import kspaceFirstOrder
+
+
+def _make_sim():
+    """Create a minimal 2D simulation with a point source."""
+    N = (64, 64)
+    dx = (1e-4, 1e-4)
+    kgrid = kWaveGrid(N, dx)
+    kgrid.makeTime(1500)
+
+    source = kSource()
+    source.p0 = np.zeros(N)
+    source.p0[32, 32] = 1.0
+
+    sensor = kSensor(mask=np.ones(N))
+    return kgrid, source, sensor
+
+
+def test_alpha_power_near_unity_raises_without_no_dispersion():
+    """alpha_power close to 1.0 must raise when dispersion is enabled."""
+    kgrid, source, sensor = _make_sim()
+    medium = kWaveMedium(sound_speed=1500, density=1000, alpha_coeff=0.5, alpha_power=1.01)
+
+    with pytest.raises(ValueError, match="too close to 1.0"):
+        kspaceFirstOrder(kgrid, medium, source, sensor, pml_inside=True, quiet=True)
+
+
+def test_alpha_power_near_unity_works_with_no_dispersion():
+    """alpha_power close to 1.0 must produce valid output when dispersion is disabled."""
+    kgrid, source, sensor = _make_sim()
+    medium = kWaveMedium(
+        sound_speed=1500,
+        density=1000,
+        alpha_coeff=0.5,
+        alpha_power=1.01,
+        alpha_mode="no_dispersion",
+    )
+
+    result = kspaceFirstOrder(kgrid, medium, source, sensor, pml_inside=True, quiet=True)
+    p = np.asarray(result["p"])
+    assert not np.any(np.isnan(p)), "Output contains NaN with alpha_power=1.01 and no_dispersion"
+    assert not np.all(p == 0), "Output is all zeros — absorption had no effect"
+
+
+def test_alpha_mode_no_absorption():
+    """alpha_mode='no_absorption' should disable absorption but keep dispersion."""
+    kgrid, source, sensor = _make_sim()
+    medium_lossy = kWaveMedium(sound_speed=1500, density=1000, alpha_coeff=0.5, alpha_power=1.5)
+    medium_no_abs = kWaveMedium(
+        sound_speed=1500,
+        density=1000,
+        alpha_coeff=0.5,
+        alpha_power=1.5,
+        alpha_mode="no_absorption",
+    )
+
+    result_lossy = kspaceFirstOrder(kgrid, medium_lossy, source, sensor, pml_inside=True, quiet=True)
+    result_no_abs = kspaceFirstOrder(kgrid, medium_no_abs, source, sensor, pml_inside=True, quiet=True)
+
+    p_lossy = np.asarray(result_lossy["p"])
+    p_no_abs = np.asarray(result_no_abs["p"])
+
+    assert not np.any(np.isnan(p_no_abs))
+    # Disabling absorption should change the output
+    assert not np.allclose(p_lossy, p_no_abs), "Disabling absorption should change the output"
+
+
+def test_alpha_mode_no_dispersion():
+    """alpha_mode='no_dispersion' should disable dispersion but keep absorption."""
+    kgrid, source, sensor = _make_sim()
+    medium_lossy = kWaveMedium(sound_speed=1500, density=1000, alpha_coeff=0.5, alpha_power=1.5)
+    medium_no_disp = kWaveMedium(
+        sound_speed=1500,
+        density=1000,
+        alpha_coeff=0.5,
+        alpha_power=1.5,
+        alpha_mode="no_dispersion",
+    )
+
+    result_lossy = kspaceFirstOrder(kgrid, medium_lossy, source, sensor, pml_inside=True, quiet=True)
+    result_no_disp = kspaceFirstOrder(kgrid, medium_no_disp, source, sensor, pml_inside=True, quiet=True)
+
+    p_lossy = np.asarray(result_lossy["p"])
+    p_no_disp = np.asarray(result_no_disp["p"])
+
+    assert not np.any(np.isnan(p_no_disp))
+    # Both should have absorption so similar amplitude, but waveforms differ
+    assert not np.allclose(p_lossy, p_no_disp), "Disabling dispersion should change the output"
+
+
+def test_alpha_power_normal_range_unaffected():
+    """alpha_power=1.5 (well away from singularity) should work as before."""
+    kgrid, source, sensor = _make_sim()
+    medium = kWaveMedium(sound_speed=1500, density=1000, alpha_coeff=0.5, alpha_power=1.5)
+
+    result = kspaceFirstOrder(kgrid, medium, source, sensor, pml_inside=True, quiet=True)
+    p = np.asarray(result["p"])
+    assert not np.any(np.isnan(p))
+    assert p.shape[0] > 0
+
+
+def test_cpp_backend_warns_on_alpha_mode(tmp_path):
+    """C++ backend should warn when alpha_mode is set (it cannot honor it)."""
+    kgrid, source, sensor = _make_sim()
+    medium = kWaveMedium(
+        sound_speed=1500,
+        density=1000,
+        alpha_coeff=0.5,
+        alpha_power=1.5,
+        alpha_mode="no_dispersion",
+    )
+
+    with pytest.warns(UserWarning, match="not supported by the C\\+\\+ backend"):
+        kspaceFirstOrder(
+            kgrid,
+            medium,
+            source,
+            sensor,
+            pml_inside=True,
+            quiet=True,
+            backend="cpp",
+            save_only=True,
+            data_path=str(tmp_path),
+        )