Fix issue #23

CHANGELOG.md

@@ -10,6 +10,7 @@ Changelog
 * Add support for Beutler soft-core Lennard-Jones form.
 * Fixed type check for ``water_template``.
 * Add support for simulations in the osmotic ensemble.
+* Fixed non-uniform bulk insertion positions caused by use of normal rather than uniform random numbers.
 
 [2025.2.0](https://github.com/openbiosim/loch/compare/2025.1.0...2025.2.0) - Feb 2026
 -------------------------------------------------------------------------------------

src/loch/_kernels.py

@@ -272,7 +272,8 @@
             GLOBAL float* water_position,
             int is_target,
             GLOBAL float* randoms_rotation,
-            GLOBAL float* randoms_position,
+            GLOBAL float* randoms_position_sphere,
+            GLOBAL float* randoms_position_bulk,
             GLOBAL float* randoms_radius,
             GLOBAL const float* cell_matrix)
         {
@@ -319,9 +320,9 @@
                 if (is_target == 1)
                 {
                     // Generate a random position around the target using pre-generated normals.
-                    xyz[0] = randoms_position[tidx * 3];
-                    xyz[1] = randoms_position[tidx * 3 + 1];
-                    xyz[2] = randoms_position[tidx * 3 + 2];
+                    xyz[0] = randoms_position_sphere[tidx * 3];
+                    xyz[1] = randoms_position_sphere[tidx * 3 + 1];
+                    xyz[2] = randoms_position_sphere[tidx * 3 + 2];
 
                     float norm = sqrtf(xyz[0] * xyz[0] + xyz[1] * xyz[1] + xyz[2] * xyz[2]);
                     xyz[0] /= norm;
@@ -337,9 +338,9 @@
                 {
                     // Use pre-generated uniform randoms for bulk sampling.
                     float r[3];
-                    r[0] = randoms_position[tidx * 3];
-                    r[1] = randoms_position[tidx * 3 + 1];
-                    r[2] = randoms_position[tidx * 3 + 2];
+                    r[0] = randoms_position_bulk[tidx * 3];
+                    r[1] = randoms_position_bulk[tidx * 3 + 1];
+                    r[2] = randoms_position_bulk[tidx * 3 + 2];
 
                     for (int i = 0; i < 3; i++)
                     {

src/loch/_platforms/_rng.py

@@ -42,8 +42,11 @@ class BatchRandoms:
     rotation : numpy.ndarray
         Shape (batch_size, 3) array of uniform [0,1) randoms for water rotation.
 
-    position : numpy.ndarray
-        Shape (batch_size, 3) array of normal randoms for position direction.
+    position_sphere : numpy.ndarray
+        Shape (batch_size, 3) array of normal randoms for sphere position direction.
+
+    position_bulk : numpy.ndarray
+        Shape (batch_size, 3) array of uniform [0,1) randoms for bulk box sampling.
 
     radius : numpy.ndarray
         Shape (batch_size,) array of uniform [0,1) randoms for radial distance.
@@ -53,7 +56,8 @@ class BatchRandoms:
     """
 
     rotation: _np.ndarray
-    position: _np.ndarray
+    position_sphere: _np.ndarray
+    position_bulk: _np.ndarray
     radius: _np.ndarray
     acceptance: _np.ndarray
 
@@ -101,7 +105,10 @@ def _generate_batch(self) -> BatchRandoms:
             rotation=self._rng.uniform(0, 1, size=(self._batch_size, 3)).astype(
                 _np.float32
             ),
-            position=self._rng.normal(0, 1, size=(self._batch_size, 3)).astype(
+            position_sphere=self._rng.normal(0, 1, size=(self._batch_size, 3)).astype(
+                _np.float32
+            ),
+            position_bulk=self._rng.uniform(0, 1, size=(self._batch_size, 3)).astype(
                 _np.float32
             ),
             radius=self._rng.uniform(0, 1, size=self._batch_size).astype(_np.float32),

src/loch/_sampler.py

@@ -1479,7 +1479,10 @@ def move(self, context: _openmm.Context) -> list[int]:
             # Get pre-computed random numbers for this batch.
             batch_randoms = self._rng_manager.get_batch_randoms()
             randoms_rotation = self._backend.to_gpu(batch_randoms.rotation)
-            randoms_position = self._backend.to_gpu(batch_randoms.position)
+            randoms_position_sphere = self._backend.to_gpu(
+                batch_randoms.position_sphere
+            )
+            randoms_position_bulk = self._backend.to_gpu(batch_randoms.position_bulk)
             randoms_radius = self._backend.to_gpu(batch_randoms.radius)
 
             # Generate the random water positions and orientations.
@@ -1492,7 +1495,8 @@ def move(self, context: _openmm.Context) -> list[int]:
                 self._water_positions,
                 is_target,
                 randoms_rotation,
-                randoms_position,
+                randoms_position_sphere,
+                randoms_position_bulk,
                 randoms_radius,
                 self._gpu_cell_matrix,
                 block=(self._num_threads, 1, 1),

tests/test_energy.py

@@ -285,17 +285,17 @@ def test_platform_consistency(fixture, request):
     )
 
 
-# Reference energy values captured with seed=42 on the original kernel implementation.
+# Reference energy values captured with seed=42.
 # These anchor the kernel output to exact values so that refactors (e.g. moving from
 # __device__ static arrays to buffer arguments) can be validated.
 _REFERENCE_ENERGIES = {
     "water_box": {
-        "energy_coul": -9.45853172201302,
-        "energy_lj": 3.2191088,
+        "energy_coul": -4.674884,
+        "energy_lj": 0.82380486,
     },
     "bpti": {
-        "energy_coul": -15.377882774621897,
-        "energy_lj": -0.58867246,
+        "energy_coul": -13.205343,
+        "energy_lj": 5.061536,
     },
 }
 

tests/test_rng.py

@@ -11,12 +11,14 @@ def test_dataclass_fields(self):
         """Test that BatchRandoms has the expected fields."""
         batch = BatchRandoms(
             rotation=np.zeros((10, 3)),
-            position=np.zeros((10, 3)),
+            position_sphere=np.zeros((10, 3)),
+            position_bulk=np.zeros((10, 3)),
             radius=np.zeros(10),
             acceptance=np.zeros(10),
         )
         assert hasattr(batch, "rotation")
-        assert hasattr(batch, "position")
+        assert hasattr(batch, "position_sphere")
+        assert hasattr(batch, "position_bulk")
         assert hasattr(batch, "radius")
         assert hasattr(batch, "acceptance")
 
@@ -32,7 +34,8 @@ def test_batch_shapes(self):
         batch = rng.get_batch_randoms()
 
         assert batch.rotation.shape == (batch_size, 3)
-        assert batch.position.shape == (batch_size, 3)
+        assert batch.position_sphere.shape == (batch_size, 3)
+        assert batch.position_bulk.shape == (batch_size, 3)
         assert batch.radius.shape == (batch_size,)
         assert batch.acceptance.shape == (batch_size,)
 
@@ -45,7 +48,8 @@ def test_batch_dtypes(self):
         batch = rng.get_batch_randoms()
 
         assert batch.rotation.dtype == np.float32
-        assert batch.position.dtype == np.float32
+        assert batch.position_sphere.dtype == np.float32
+        assert batch.position_bulk.dtype == np.float32
         assert batch.radius.dtype == np.float32
         assert batch.acceptance.dtype == np.float32
 
@@ -60,6 +64,7 @@ def test_uniform_range(self):
             batch = rng.get_batch_randoms()
 
             assert np.all(batch.rotation >= 0) and np.all(batch.rotation < 1)
+            assert np.all(batch.position_bulk >= 0) and np.all(batch.position_bulk < 1)
             assert np.all(batch.radius >= 0) and np.all(batch.radius < 1)
             assert np.all(batch.acceptance >= 0) and np.all(batch.acceptance < 1)
 
@@ -73,7 +78,7 @@ def test_normal_distribution(self):
         samples = []
         for _ in range(10):
             batch = rng.get_batch_randoms()
-            samples.append(batch.position.flatten())
+            samples.append(batch.position_sphere.flatten())
 
         all_samples = np.concatenate(samples)
 
@@ -92,7 +97,8 @@ def test_reproducibility_with_seed(self):
         batch2 = rng2.get_batch_randoms()
 
         np.testing.assert_array_equal(batch1.rotation, batch2.rotation)
-        np.testing.assert_array_equal(batch1.position, batch2.position)
+        np.testing.assert_array_equal(batch1.position_sphere, batch2.position_sphere)
+        np.testing.assert_array_equal(batch1.position_bulk, batch2.position_bulk)
         np.testing.assert_array_equal(batch1.radius, batch2.radius)
         np.testing.assert_array_equal(batch1.acceptance, batch2.acceptance)