Add MLX backend for Apple Silicon processors

kwave/__init__.py

@@ -7,6 +7,7 @@
 import warnings
 from pathlib import Path
 from typing import List
+from urllib.error import URLError
 from urllib.request import urlretrieve
 
 # Test installation with:
@@ -78,8 +79,15 @@
     "vcruntime140_threads.dll",
 ]
 
+# Metal kernels of the MLX GPU binary (Apple Silicon). The darwin MLX binary links MLX statically and loads
+# this file from its own directory, so it is downloaded next to it.
+MLX_METALLIB = "mlx.metallib"
+
 EXECUTABLE_PREFIX = "kspaceFirstOrder-"
-ARCHITECTURES = ["omp", "cuda"]
+# GPU backend per platform: CUDA on Linux/Windows, MLX (Metal) on macOS/Apple Silicon.
+ARCHITECTURES = ["omp", "cuda", "mlx"]
+# Backends whose binaries are optional: a failed download only disables that backend instead of aborting `import kwave`.
+OPTIONAL_ARCHITECTURES = ["mlx"]
 
 
 def _platform_binary_url(platform_name: str, architecture: str) -> list:
@@ -92,11 +100,15 @@ def _platform_binary_url(platform_name: str, architecture: str) -> list:
     """
     if architecture == "cuda" and platform_name == "darwin":
         return []
+    if architecture == "mlx" and platform_name != "darwin":
+        return []
     if platform_name == "darwin" and platform_name == PLATFORM and _darwin_unsupported:
         return []
-    backend = "CUDA" if architecture == "cuda" else "OMP"
+    backend = {"cuda": "CUDA", "mlx": "MLX"}.get(architecture, "OMP")
     ext = ".exe" if platform_name == "windows" else ""
     exe_url = f"{_UNIFIED_RELEASE_URL}{EXECUTABLE_PREFIX}{backend}-{platform_name}{ext}"
+    if architecture == "mlx":
+        return [exe_url, _UNIFIED_RELEASE_URL + MLX_METALLIB]
     # Shared runtime DLLs ship with the OMP entry only — both backends find
     # them at the same BINARY_PATH, so attaching to both would download the
     # full 21-DLL bundle twice on first install.
@@ -238,6 +250,19 @@ def _record_binary_metadata(binary_version: str, binary_filepath: str, binary_ur
         json.dump(metadata, outfile, indent=4)
 
 
+def _remove_installed_binaries(system_os: str, bin_type: str) -> None:
+    """Remove every local file (and its metadata) of a backend, e.g. after an incomplete download of a bundle."""
+    for url in URL_DICT[system_os][bin_type]:
+        filename = _local_filename(url.split("/")[-1])
+        for path in (BINARY_PATH / filename, BINARY_PATH / f"{filename}_metadata.json"):
+            try:
+                os.remove(path)
+            except FileNotFoundError:
+                pass
+            except OSError:  # pragma: no cover - defensive
+                logging.warning(f"Could not remove {path}; please delete it manually.")
+
+
 def download_binaries(system_os: str, bin_type: str):
     """
     Download binary from release url
@@ -255,6 +280,10 @@ def download_binaries(system_os: str, bin_type: str):
         binary_version, asset_name = url.split("/")[-2:]
         filename = _local_filename(asset_name)
 
+        # Keep valid cached files: a missing optional backend must not force (network) re-downloads of the others.
+        if system_os == PLATFORM and _is_binary_present(filename, bin_type):
+            continue
+
         logging.log(logging.INFO, f"Downloading {asset_name} -> {filename} in {BINARY_PATH}...")
 
         # Create the directory if it does not yet exist
@@ -267,6 +296,15 @@ def download_binaries(system_os: str, bin_type: str):
             _ensure_executable(binary_filepath)
             _record_binary_metadata(binary_version=binary_version, binary_filepath=binary_filepath, binary_url=url, filename=filename)
 
+        except URLError as e:
+            if bin_type not in OPTIONAL_ARCHITECTURES:
+                raise
+            # e.g. the pinned release does not ship this backend yet: keep the other backends usable. A partial
+            # bundle (e.g. the MLX executable without mlx.metallib) cannot run, so remove what is already there.
+            _remove_installed_binaries(system_os, bin_type)
+            logging.warning(f"Could not download {asset_name} ({e}); the {bin_type.upper()} backend will not be available.")
+            return
+
         except TimeoutError:
             logging.log(
                 logging.WARN,

kwave/compat.py

@@ -54,7 +54,7 @@ def options_to_kwargs(simulation_options=None, execution_options=None):
         if opts.backend == "python":
             kwargs["backend"] = "python"
             kwargs["device"] = "gpu" if opts.is_gpu_simulation else "cpu"
-        elif opts.backend == "CUDA":
+        elif opts.backend in ("CUDA", "MLX"):
             kwargs["backend"] = "cpp"
             kwargs["device"] = "gpu"
         elif opts.backend == "OMP":

kwave/executor.py

@@ -26,7 +26,8 @@ def _make_binary_executable(self):
         if not binary_path.exists():
             if kwave.PLATFORM == "darwin" and self.execution_options.is_gpu_simulation:
                 raise ValueError(
-                    "GPU simulations are currently not supported on MacOS. "
+                    f"GPU binary not found at {binary_path}. GPU simulations on macOS require Apple Silicon "
+                    "and the MLX binary (kspaceFirstOrder-MLX). "
                     "Try running the simulation on CPU by setting is_gpu_simulation=False."
                 )
             raise FileNotFoundError(f"Binary not found at {binary_path}")

kwave/kspaceFirstOrder.py

@@ -177,7 +177,8 @@ def kspaceFirstOrder(
             the compiled C++ binary.  Default ``"python"``.
         device: ``"cpu"`` or ``"gpu"``.  For ``backend="python"`` this
             selects NumPy (cpu) vs CuPy (gpu).  For ``backend="cpp"`` it
-            selects the OMP vs CUDA binary.  Default ``"cpu"``.
+            selects the OMP vs GPU binary (CUDA on Linux/Windows,
+            MLX on macOS with Apple Silicon).  Default ``"cpu"``.
         dtype: Numerical precision for state arrays in the Python backend.
             Accepts dtype-like input — a numpy dtype (``np.float32``,
             ``np.float64``; cupy aliases like ``cp.float32`` work since cupy
@@ -207,7 +208,7 @@ def kspaceFirstOrder(
         debug: Print detailed diagnostic output.  Default ``False``.
         num_threads: Thread count for the C++ OMP binary.  ``None`` uses all
             available cores.  Default ``None``.
-        device_num: GPU device index for CUDA execution.  Default ``None``.
+        device_num: GPU device index for CUDA/MLX execution.  Default ``None``.
         binary_path: Path to a custom C++ binary.  When ``None`` (default),
             the binary bundled with ``k-wave-data`` is used.  Only applies
             when ``backend="cpp"``.

kwave/kspaceFirstOrder2D.py

@@ -33,7 +33,8 @@ def kspace_first_order_2d_gpu(
     2D time-domain simulation of wave propagation on a GPU using C++ CUDA code.
 
     This function provides a blind interface to the C++/CUDA version of kspaceFirstOrder2D
-    (called kspaceFirstOrder-CUDA). The function works by saving the input files to disk,
+    (called kspaceFirstOrder-CUDA; kspaceFirstOrder-MLX on macOS with Apple Silicon).
+    The function works by saving the input files to disk,
     running the C++/CUDA simulation, and loading the output files back into Python.
 
     Parameters:

kwave/kspaceFirstOrder3D.py

@@ -34,6 +34,8 @@ def kspaceFirstOrder3DG(
     This function provides a blind interface to the C++/CUDA version of kspaceFirstOrder3D
     (called kspaceFirstOrder-CUDA). The function works by saving the input files to disk,
     running the C++/CUDA simulation, and loading the output files back into Python.
+    On macOS with Apple Silicon, the MLX port of the same code (kspaceFirstOrder-MLX) is used
+    instead, running on the Apple GPU through Metal.
 
     Parameters:
     -----------

kwave/options/simulation_execution_options.py

@@ -17,7 +17,8 @@ class SimulationExecutionOptions:
     Attributes:
         backend: Execution backend to use. Options:
             - "OMP": C++ OpenMP binary (default for CPU)
-            - "CUDA": C++ CUDA binary (default for GPU)
+            - "CUDA": C++ CUDA binary (default for GPU on Linux/Windows)
+            - "MLX": C++ MLX binary for the Apple Silicon GPU (default for GPU on macOS)
             - "python": Pure Python/CuPy solver (no external binaries required)
     """
 
@@ -39,7 +40,7 @@ def __init__(
         checkpoint_interval: Optional[int] = None,  # [seconds]
         checkpoint_timesteps: Optional[int] = None,  # [timestep integer]
         checkpoint_file: Optional[Path | str] = None,  # [path to hdf5 file]
-        backend: Optional[str] = None,  # "OMP", "CUDA", or "python"
+        backend: Optional[str] = None,  # "OMP", "CUDA", "MLX", or "python"
     ):
         self._num_threads_explicit = num_threads is not None
         self.backend = backend
@@ -113,7 +114,7 @@ def backend(self) -> Optional[str]:
 
     @backend.setter
     def backend(self, value: Optional[str]):
-        valid_backends = [None, "OMP", "CUDA", "python"]
+        valid_backends = [None, "OMP", "CUDA", "MLX", "python"]
         if value not in valid_backends:
             raise ValueError(f"Backend must be one of {valid_backends}, got '{value}'")
         self._backend = value
@@ -136,14 +137,18 @@ def is_gpu_simulation(self, value: Optional[bool]):
 
     @property
     def binary_name(self) -> str:
-        valid_binary_names = ["kspaceFirstOrder-OMP", "kspaceFirstOrder-CUDA"]
+        valid_binary_names = ["kspaceFirstOrder-OMP", "kspaceFirstOrder-CUDA", "kspaceFirstOrder-MLX"]
         if self._binary_name is None:
             # set default binary name based on GPU simulation value
             if self.is_gpu_simulation is None:
                 raise ValueError("`is_gpu_simulation` must be set to either True or False before determining the binary name.")
 
-            if self.is_gpu_simulation:
-                self._binary_name = "kspaceFirstOrder-CUDA"
+            if self.backend == "MLX":
+                # explicit choice of the Apple Silicon GPU binary
+                self._binary_name = "kspaceFirstOrder-MLX"
+            elif self.is_gpu_simulation:
+                # macOS has no CUDA: the GPU binary runs on the Apple Silicon GPU through MLX (Metal)
+                self._binary_name = "kspaceFirstOrder-MLX" if PLATFORM == "darwin" else "kspaceFirstOrder-CUDA"
             else:
                 self._binary_name = "kspaceFirstOrder-OMP"
 

kwave/solvers/cpp_simulation.py

@@ -339,7 +339,8 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path:
 
         When *binary_path* is provided the caller's custom binary is used and
         verified to exist.  Otherwise the bundled binary is selected based on
-        *device* (``"gpu"`` → CUDA, anything else → OMP) and verified to exist.
+        *device* (``"gpu"`` → CUDA, or MLX on macOS; anything else → OMP) and
+        verified to exist.
 
         Args:
             device: ``"cpu"`` or ``"gpu"``.
@@ -352,8 +353,8 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path:
         Raises:
             FileNotFoundError: When the binary does not exist at the resolved
                 path.
-            ValueError: When ``device="gpu"`` is requested on macOS where no
-                CUDA binary is available.
+            ValueError: When ``device="gpu"`` is requested on macOS and the MLX
+                binary is not available (e.g. on an Intel Mac).
         """
         if binary_path is not None:
             resolved = Path(binary_path)
@@ -363,11 +364,18 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path:
 
         import kwave
 
-        binary_name = "kspaceFirstOrder-CUDA" if device == "gpu" else "kspaceFirstOrder-OMP"
+        if device == "gpu":
+            # macOS has no CUDA: use the MLX binary, which runs on the Apple Silicon GPU (Metal)
+            binary_name = "kspaceFirstOrder-MLX" if kwave.PLATFORM == "darwin" else "kspaceFirstOrder-CUDA"
+        else:
+            binary_name = "kspaceFirstOrder-OMP"
         resolved = kwave.BINARY_PATH / binary_name
         if not resolved.exists():
             if kwave.PLATFORM == "darwin" and device == "gpu":
-                raise ValueError("GPU simulations are not supported on macOS. Use device='cpu'.")
+                raise ValueError(
+                    f"GPU binary {binary_name} not found at {resolved}. GPU simulations on macOS require Apple Silicon "
+                    "and the MLX binary from the k-wave-python binary release. Use device='cpu' otherwise."
+                )
             raise FileNotFoundError(f"C++ binary not found at {resolved}. Install with: pip install k-wave-data")
         return resolved
 
@@ -402,7 +410,7 @@ def _execute(self, input_file, output_file, *, device, num_threads, device_num,
         resolved = self._resolve_binary_path(device, binary_path)
         resolved.chmod(resolved.stat().st_mode | stat.S_IEXEC)
 
-        if device == "gpu":
+        if device == "gpu" and not resolved.name.startswith("kspaceFirstOrder-MLX"):
             self._warn_if_unsupported_gpu()
 
         # Build command-line options

tests/test__init__.py

@@ -102,5 +102,101 @@ def test_existing_non_executable_binary_is_healed(tmp_path, monkeypatch):
     assert post_mode & stat.S_IXOTH, "other exec bit not healed on cache hit"
 
 
+def test_darwin_gpu_urls_point_to_mlx_binary_and_metallib():
+    """macOS has no CUDA: its GPU backend is the MLX binary plus the Metal kernel library it loads at run time."""
+    import kwave
+
+    mlx_urls = kwave._platform_binary_url("darwin", "mlx")
+    if kwave.PLATFORM == "darwin" and kwave._darwin_unsupported:
+        assert mlx_urls == []
+        return
+    assert [url.rsplit("/", 1)[-1] for url in mlx_urls] == ["kspaceFirstOrder-MLX-darwin", "mlx.metallib"]
+    assert all(url.startswith(kwave._UNIFIED_RELEASE_URL) for url in mlx_urls)
+    assert kwave._platform_binary_url("darwin", "cuda") == []
+    assert kwave._platform_binary_url("linux", "mlx") == []
+    assert kwave._platform_binary_url("windows", "mlx") == []
+
+
+def test_local_filename_strips_platform_suffix_for_mlx():
+    import kwave
+
+    assert kwave._local_filename("kspaceFirstOrder-MLX-darwin") == "kspaceFirstOrder-MLX"
+    assert kwave._local_filename("mlx.metallib") == "mlx.metallib"
+
+
+def test_failed_optional_mlx_download_does_not_abort(tmp_path, monkeypatch):
+    """A release without the MLX assets must not break `import kwave` (the OMP backend stays usable)."""
+    from urllib.error import HTTPError
+
+    import kwave
+
+    monkeypatch.setattr(kwave, "BINARY_PATH", tmp_path)
+
+    def fake_urlretrieve(url, dest):
+        raise HTTPError(url, 404, "Not Found", None, None)
+
+    monkeypatch.setattr(kwave, "urlretrieve", fake_urlretrieve)
+    monkeypatch.setattr(kwave, "URL_DICT", {kwave.PLATFORM: {"mlx": ["https://example.invalid/kspaceFirstOrder-MLX-darwin"]}})
+    kwave.download_binaries(kwave.PLATFORM, "mlx")  # must not raise
+
+    monkeypatch.setattr(kwave, "URL_DICT", {kwave.PLATFORM: {"omp": ["https://example.invalid/kspaceFirstOrder-OMP-darwin"]}})
+    with pytest.raises(HTTPError):
+        kwave.download_binaries(kwave.PLATFORM, "omp")
+
+
+def _fake_release(tmp_path, monkeypatch, fail_assets):
+    """Point kwave at a fake release with an OMP binary and a two-file MLX bundle; assets in fail_assets 404."""
+    from urllib.error import HTTPError
+
+    import kwave
+
+    base = "https://example.invalid/releases/download/v9.9.9/"
+    url_dict = {
+        kwave.PLATFORM: {
+            "omp": [base + "kspaceFirstOrder-OMP-darwin"],
+            "mlx": [base + "kspaceFirstOrder-MLX-darwin", base + "mlx.metallib"],
+        }
+    }
+    downloads = []
+
+    def fake_urlretrieve(url, dest):
+        downloads.append(url.rsplit("/", 1)[-1])
+        if url.rsplit("/", 1)[-1] in fail_assets:
+            raise HTTPError(url, 404, "Not Found", None, None)
+        Path(dest).write_bytes(b"payload-" + url.encode())
+
+    monkeypatch.setattr(kwave, "BINARY_PATH", tmp_path)
+    monkeypatch.setattr(kwave, "URL_DICT", url_dict)
+    monkeypatch.setattr(kwave, "ARCHITECTURES", ["omp", "mlx"])
+    monkeypatch.setattr(kwave, "urlretrieve", fake_urlretrieve)
+    return downloads
+
+
+def test_cached_omp_is_not_downloaded_again_when_mlx_is_missing(tmp_path, monkeypatch):
+    """A missing optional MLX bundle must not re-download (or require the network for) a valid cached OMP binary."""
+    import kwave
+
+    downloads = _fake_release(tmp_path, monkeypatch, fail_assets={"kspaceFirstOrder-MLX-darwin", "mlx.metallib"})
+    kwave.install_binaries()
+    assert (tmp_path / "kspaceFirstOrder-OMP").exists()
+    assert not kwave.binaries_present()  # MLX still missing
+
+    downloads.clear()
+    kwave.install_binaries()  # second import: OMP is cached and valid
+    assert "kspaceFirstOrder-OMP-darwin" not in downloads
+    assert downloads == ["kspaceFirstOrder-MLX-darwin"]
+
+
+def test_incomplete_mlx_bundle_is_removed(tmp_path, monkeypatch):
+    """If mlx.metallib fails after the MLX executable was downloaded, the unusable executable must not stay installed."""
+    import kwave
+
+    _fake_release(tmp_path, monkeypatch, fail_assets={"mlx.metallib"})
+    kwave.download_binaries(kwave.PLATFORM, "mlx")
+    assert not (tmp_path / "kspaceFirstOrder-MLX").exists()
+    assert not (tmp_path / "kspaceFirstOrder-MLX_metadata.json").exists()
+    assert not (tmp_path / "mlx.metallib").exists()
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

tests/test_cpp_simulation.py

@@ -70,10 +70,17 @@ def test_default_missing_cuda_on_linux_raises_file_not_found(self, tmp_path):
             with pytest.raises(FileNotFoundError, match="pip install k-wave-data"):
                 CppSimulation._resolve_binary_path("gpu")
 
-    def test_default_gpu_on_macos_raises_value_error(self, tmp_path):
-        # macOS + gpu must raise before the binary-existence check.
+    def test_default_gpu_on_macos_selects_mlx_binary(self, tmp_path):
+        mlx_binary = tmp_path / "kspaceFirstOrder-MLX"
+        mlx_binary.write_text("#!/bin/sh\n")
         with patch("kwave.BINARY_PATH", tmp_path), patch("kwave.PLATFORM", "darwin"):
-            with pytest.raises(ValueError, match="not supported on macOS"):
+            resolved = CppSimulation._resolve_binary_path("gpu")
+        assert resolved.name == "kspaceFirstOrder-MLX"
+
+    def test_default_gpu_on_macos_missing_mlx_raises_value_error(self, tmp_path):
+        # macOS + gpu without the MLX binary (e.g. Intel Mac) must raise a ValueError pointing to device='cpu'.
+        with patch("kwave.BINARY_PATH", tmp_path), patch("kwave.PLATFORM", "darwin"):
+            with pytest.raises(ValueError, match="kspaceFirstOrder-MLX not found"):
                 CppSimulation._resolve_binary_path("gpu")
 
     @pytest.mark.skipif(sys.platform == "win32", reason="Windows does not honor Unix executable bits")

tests/test_executor.py

@@ -255,9 +255,7 @@ def test_cpu_environment_variable(self):
         del os.environ["KWAVE_FORCE_CPU"]
 
     def test_executor_gpu_cuda_failure_darwin(self):
-        expected_error_msg = (
-            "GPU simulations are currently not supported on MacOS. Try running the simulation on CPU by setting is_gpu_simulation=False."
-        )
+        expected_error_msg = "GPU simulations on macOS require Apple Silicon and the MLX binary"
         # Configure the mock path object
         mock_binary_path = MagicMock(spec=Path)
         mock_binary_path.chmod.side_effect = FileNotFoundError