Add GPU backend selection with load_cuda(), load_metal(), load_amdgpu(), load_oneapi() (#23)

* Add GPU backend selection functionality and example

* Fix linting errors

* Fix linting errors

* Simplify example, updated README to have gpu selection

* Fix linting errors

* Install GPU backends on first use

---------

Co-authored-by: Carlos Castillo Passi <cacp@stanford.edu>

Author: anvika-singhal

README.md

@@ -37,6 +37,8 @@ minutes. Subsequent runs are fast.
 import komamripy as km
 import numpy as np
 
+km.load_cuda()  # GPU backend selection, optional
+
 sys = km.Scanner()  # scanner hardware
 obj = km.brain_phantom2D()  # 2D brain phantom
 seq = km.PulseDesigner.EPI_example()  # example EPI sequence

examples/gpu_backend_selection.py

@@ -0,0 +1,27 @@
+"""Simulate MRI acquisition with optional GPU backend acceleration.
+
+This example demonstrates GPU backend selection. Uncomment one of the backend
+load functions to enable GPU acceleration (CUDA, Metal, AMDGPU, or oneAPI).
+Without a backend, simulation runs on CPU automatically.
+"""
+
+import numpy as np
+
+import komamripy as km
+
+# Optional: Load a GPU backend to enable GPU acceleration
+# Uncomment ONE of these (backend must be installed):
+# km.load_cuda()      # NVIDIA GPUs
+# km.load_metal()     # Apple Silicon
+# km.load_amdgpu()    # AMD GPUs
+# km.load_oneapi()    # Intel GPUs (experimental)
+
+# Define acquisition inputs
+sys = km.Scanner()
+obj = km.brain_phantom2D()
+seq = km.PulseDesigner.EPI_example()
+
+# Simulate with KomaMRI (gpu=true by default if backend loaded)
+sim_params = {"return_type": "mat"}
+signal = km.simulate(obj, seq, sys, sim_params=sim_params)
+signal = np.asarray(signal).reshape(-1)

src/komamripy/__init__.py

@@ -7,6 +7,8 @@
 Julia::
 
     using KomaMRI
+    using CUDA  # Load GPU backend
+
     sys = Scanner()
     obj = brain_phantom2D()
     seq = PulseDesigner.EPI_example()
@@ -18,6 +20,7 @@
     import komamripy as km
     import numpy as np
 
+    km.load_cuda()  # Load GPU backend
     sys = km.Scanner()
     obj = km.brain_phantom2D()
     seq = km.PulseDesigner.EPI_example()
@@ -26,10 +29,21 @@
 
 Simulation results are returned as Julia objects; use ``numpy.asarray`` to
 convert array-like results (such as a ``"mat"`` signal) into NumPy arrays.
+
+GPU backends can be loaded with:
+- km.load_cuda()      # NVIDIA GPUs
+- km.load_metal()     # Apple Silicon
+- km.load_amdgpu()    # AMD GPUs
+- km.load_oneapi()    # Intel GPUs (experimental)
+
+Once a backend is loaded, gpu=true is used by default in simulations.
 """
 
+from ._backends import load_amdgpu, load_cuda, load_metal, load_oneapi
 from ._session import get_julia
 
+__all__ = ["load_cuda", "load_metal", "load_amdgpu", "load_oneapi"]
+
 _JULIA_MODULE_ALIASES = {
     "base": "KomaMRIBase",
     "core": "KomaMRICore",

src/komamripy/_backends.py

@@ -0,0 +1,44 @@
+"""GPU backend selection for komamripy.
+
+Maps Python functions to Julia backend loading:
+- load_cuda() → using CUDA
+- load_metal() → using Metal
+- load_amdgpu() → using AMDGPU
+- load_oneapi() → using oneAPI
+
+Once a backend is loaded, KomaMRI.jl automatically uses it with gpu=true by default.
+"""
+
+from ._session import get_julia
+
+
+def load_cuda() -> None:
+    """Load CUDA backend: using CUDA"""
+    jl = get_julia()
+    if jl.Base.find_package("CUDA") is None:
+        jl.seval('import Pkg; Pkg.add("CUDA")')
+    jl.seval("using CUDA")
+
+
+def load_metal() -> None:
+    """Load Metal backend: using Metal"""
+    jl = get_julia()
+    if jl.Base.find_package("Metal") is None:
+        jl.seval('import Pkg; Pkg.add("Metal")')
+    jl.seval("using Metal")
+
+
+def load_amdgpu() -> None:
+    """Load AMDGPU backend: using AMDGPU"""
+    jl = get_julia()
+    if jl.Base.find_package("AMDGPU") is None:
+        jl.seval('import Pkg; Pkg.add("AMDGPU")')
+    jl.seval("using AMDGPU")
+
+
+def load_oneapi() -> None:
+    """Load oneAPI backend: using oneAPI (experimental)"""
+    jl = get_julia()
+    if jl.Base.find_package("oneAPI") is None:
+        jl.seval('import Pkg; Pkg.add("oneAPI")')
+    jl.seval("using oneAPI")