Add opencl_platform_index option to match OpenMM's OpenCL device

CHANGELOG.md

@@ -9,6 +9,7 @@ Changelog
 * Let the OpenCL runtime choose the work-group size rather than passing the CUDA block size, which exceeds the limit on AMD GPUs [#47](https://github.com/OpenBioSim/loch/pull/47).
 * Keep an existing ghost residue file on restart rather than clearing it, so that it stays aligned with the trajectory frames written before the restart [#50](https://github.com/OpenBioSim/loch/pull/50).
 * Flag virtual sites on perturbable molecules as ghosts when their parent atom is a ghost, so that they use soft-core rather than full Coulomb interactions with the waters [#53](https://github.com/OpenBioSim/loch/pull/53).
+* Add an `opencl_platform_index` option so that the OpenCL device is chosen from the same platform as the OpenMM context [#55](https://github.com/OpenBioSim/loch/pull/55).
 
 [2026.2.0](https://github.com/openbiosim/loch/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------

README.md

@@ -158,7 +158,10 @@ rather than just within the GCMC sphere.
 The GPU platform is controlled via the `platform` argument, which can be set to
 `"cuda"`, `"opencl"` or `"auto"` (default). When set to `"auto"`, `loch` will
 attempt to use the CUDA platform first, falling back to OpenCL if CUDA is not
-available.
+available. When using OpenCL alongside OpenMM, set `opencl_platform_index` to
+the same value as OpenMM's `OpenCLPlatformIndex` property so that `device`
+refers to the same device as OpenMM's `DeviceIndex`. Otherwise, `device`
+indexes the GPUs from all OpenCL platforms.
 
 1) Get the GCMC system:
 

src/loch/_platforms/__init__.py

@@ -86,6 +86,7 @@ def create_backend(
     num_threads: int,
     nvcc: _Optional[str] = None,
     compiler_optimisations: bool = True,
+    opencl_platform_index: _Optional[int] = None,
 ) -> _PlatformBackend:
     """
     Create a GPU platform backend instance.
@@ -122,6 +123,10 @@ def create_backend(
         Enable compiler optimisations for faster math operations.
         Default: True (matches OpenMM defaults).
 
+    opencl_platform_index : int, optional
+        OpenCL platform index (OpenCL only). If None, GPU devices from all
+        OpenCL platforms are used.
+
     Returns
     -------
     PlatformBackend
@@ -158,6 +163,7 @@ def create_backend(
             num_threads,
             nvcc,
             compiler_optimisations,
+            opencl_platform_index,
         )
     else:
         raise ValueError(

src/loch/_platforms/_opencl.py

@@ -35,8 +35,8 @@
 from .._kernels import code as _kernel_code
 from ._base import PlatformBackend as _PlatformBackend
 
-# Module-level kernel compilation cache. Keyed on
-# (device_index, compiler_optimisations). Stores compiled program binaries.
+# Module-level kernel compilation cache. Keyed on (device,
+# compiler_optimisations). Stores compiled program binaries.
 _kernel_cache = {}
 
 
@@ -58,6 +58,7 @@ def __init__(
         num_threads,
         nvcc=None,
         compiler_optimisations=True,
+        platform_index=None,
     ):
         """
         Initialize the OpenCL platform backend.
@@ -66,7 +67,9 @@ def __init__(
         ----------
 
         device : int
-            The OpenCL device index to use.
+            The index of the device on the chosen OpenCL platform, or in the
+            list of GPU devices from all platforms if 'platform_index' is
+            None.
 
         num_points : int
             Number of atoms per water molecule (typically 3).
@@ -91,15 +94,37 @@ def __init__(
             Enable compiler optimisations for faster math operations.
             When True, passes -cl-mad-enable and -cl-no-signed-zeros to the compiler.
             Default: True (matches OpenMM defaults).
+
+        platform_index : int, optional
+            The index of the OpenCL platform to use, matching OpenMM's
+            OpenCLPlatformIndex. Devices are then indexed as by OpenMM's
+            DeviceIndex, i.e. GPU and CPU devices on that platform. If None,
+            GPU devices from all platforms are used.
         """
         # Get platforms and devices
         platforms = _cl.get_platforms()
-        devices = []
-        for platform in platforms:
-            devices.extend(platform.get_devices(device_type=_cl.device_type.GPU))
 
-        if not devices:
-            raise RuntimeError("No OpenCL GPU devices found")
+        if platform_index is not None:
+            if not isinstance(platform_index, int):
+                raise ValueError("'platform_index' must be of type 'int'")
+            if platform_index < 0 or platform_index >= len(platforms):
+                raise ValueError(
+                    f"'platform_index' must be between 0 and {len(platforms) - 1}"
+                )
+            devices = platforms[platform_index].get_devices(
+                device_type=_cl.device_type.GPU | _cl.device_type.CPU
+            )
+            if not devices:
+                raise RuntimeError(
+                    f"No OpenCL devices found on platform {platform_index}"
+                )
+        else:
+            devices = []
+            for platform in platforms:
+                devices.extend(platform.get_devices(device_type=_cl.device_type.GPU))
+
+            if not devices:
+                raise RuntimeError("No OpenCL GPU devices found")
 
         # Validate device index
         if device is not None:
@@ -137,7 +162,7 @@ def compile_kernels(self) -> _Dict[str, _Callable]:
         dict
             Dictionary mapping kernel names to callable kernel functions.
         """
-        cache_key = (self._device_index, self._compiler_optimisations)
+        cache_key = (self._device, self._compiler_optimisations)
 
         # Build compiler options
         build_options = []

src/loch/_sampler.py

@@ -85,6 +85,7 @@ def __init__(
         water_template: _Optional[_Any] = None,
         device: _Optional[int] = None,
         platform: str = "auto",
+        opencl_platform_index: _Optional[int] = None,
         tolerance: float = 0.0,
         lambda_schedule: _Optional[_Any] = None,
         lambda_value: float = 0.0,
@@ -190,6 +191,12 @@ def __init__(
             or 'opencl'. When 'auto', CUDA will be preferred if available,
             falling back to OpenCL.
 
+        opencl_platform_index: int
+            The OpenCL platform index, matching OpenMM's OpenCLPlatformIndex,
+            so that 'device' refers to the same device as the OpenMM context.
+            If None, devices from all OpenCL platforms are used. Ignored for
+            CUDA.
+
         tolerance: float
             The tolerance for the acceptance probability, i.e. the minimum
             probability of acceptance for a move. This can be used to exclude
@@ -703,6 +710,11 @@ def __init__(
                 raise ValueError("'device' must be of type 'int'")
         self._device = device
 
+        if opencl_platform_index is not None:
+            if not isinstance(opencl_platform_index, int):
+                raise ValueError("'opencl_platform_index' must be of type 'int'")
+        self._opencl_platform_index = opencl_platform_index
+
         if not isinstance(compiler_optimisations, bool):
             raise ValueError("'compiler_optimisations' must be of type 'bool'")
         self._compiler_optimisations = compiler_optimisations
@@ -718,6 +730,7 @@ def __init__(
             num_threads=self._num_threads,
             nvcc=self._nvcc,
             compiler_optimisations=self._compiler_optimisations,
+            opencl_platform_index=self._opencl_platform_index,
         )
 
         # Compile kernels
@@ -852,6 +865,7 @@ def __str__(self) -> str:
             f"water_template={self._water_template}, "
             f"platform={self._platform}, "
             f"device={self._device}, "
+            f"opencl_platform_index={self._opencl_platform_index}, "
             f"tolerance={self._tolerance}, "
             f"nvcc={self._nvcc}, "
             f"compiler_optimisations={self._compiler_optimisations}, "

tests/test_opencl_platform.py

@@ -0,0 +1,101 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+# Skip the entire module if PyOpenCL is not installed.
+pytest.importorskip("pyopencl")
+
+_CPU = 2
+_GPU = 4
+
+
+def _mock_platform(devices):
+    """A mock OpenCL platform with (name, type) devices, filtered by type."""
+    platform = MagicMock()
+    platform.get_devices.side_effect = lambda device_type: [
+        name for name, kind in devices if kind & device_type
+    ]
+    return platform
+
+
+def _make_backend(mock_cl, device, platform_index, platforms=None):
+    """Instantiate OpenCLPlatform with a mocked PyOpenCL module."""
+    from loch._platforms._opencl import OpenCLPlatform
+
+    if platforms is None:
+        platforms = [
+            _mock_platform([("a0", _GPU), ("a1", _GPU)]),
+            _mock_platform([("b_cpu", _CPU), ("b0", _GPU)]),
+        ]
+
+    mock_cl.device_type.CPU = _CPU
+    mock_cl.device_type.GPU = _GPU
+    mock_cl.get_platforms.return_value = platforms
+
+    return OpenCLPlatform(
+        device=device,
+        num_points=3,
+        num_batch=10,
+        num_waters=5,
+        num_atoms=100,
+        num_threads=32,
+        platform_index=platform_index,
+    )
+
+
+class TestOpenCLPlatformIndex:
+    """Tests for OpenCLPlatform device selection by platform index."""
+
+    def test_all_platforms_when_none(self):
+        """GPU devices from all platforms are indexed when platform_index is None."""
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            backend = _make_backend(mock_cl, device=2, platform_index=None)
+            assert backend._device == "b0"
+
+    def test_selected_platform(self):
+        """Devices are indexed within the chosen platform."""
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            backend = _make_backend(mock_cl, device=1, platform_index=0)
+            assert backend._device == "a1"
+
+    def test_cpu_devices_indexed_as_openmm(self):
+        """CPU devices count towards the index, matching OpenMM's DeviceIndex."""
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            backend = _make_backend(mock_cl, device=1, platform_index=1)
+            assert backend._device == "b0"
+
+    def test_device_out_of_range_for_platform(self):
+        """A device index beyond the chosen platform's devices raises."""
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            with pytest.raises(ValueError, match="'device'"):
+                _make_backend(mock_cl, device=2, platform_index=1)
+
+    def test_platform_index_out_of_range(self):
+        """An invalid platform index raises."""
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            with pytest.raises(ValueError, match="'platform_index'"):
+                _make_backend(mock_cl, device=0, platform_index=2)
+
+    def test_platform_without_devices(self):
+        """A platform with no devices raises a clear error."""
+        platforms = [_mock_platform([])]
+        with patch("loch._platforms._opencl._cl") as mock_cl:
+            with pytest.raises(RuntimeError, match="platform 0"):
+                _make_backend(mock_cl, device=0, platform_index=0, platforms=platforms)
+
+    def test_sampler_forwards_platform_index(self):
+        """create_backend passes opencl_platform_index to OpenCLPlatform."""
+        from loch._platforms import create_backend
+
+        with patch("loch._platforms._opencl.OpenCLPlatform") as mock_backend:
+            create_backend(
+                platform="opencl",
+                device=0,
+                num_points=3,
+                num_batch=10,
+                num_waters=5,
+                num_atoms=100,
+                num_threads=32,
+                opencl_platform_index=1,
+            )
+            assert mock_backend.call_args.args[-1] == 1