diff --git a/CHANGELOG.md b/CHANGELOG.md index b3e4c90..2b17936 100644 --- a/CHANGELOG.md +++ b/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 ------------------------------------------------------------------------------------- diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index 146accc..8959791 100644 --- a/src/loch/_kernels.py +++ b/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++) { diff --git a/src/loch/_platforms/_rng.py b/src/loch/_platforms/_rng.py index 5b80e66..26ae59c 100644 --- a/src/loch/_platforms/_rng.py +++ b/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), diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 79166f0..0fec54b 100644 --- a/src/loch/_sampler.py +++ b/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), diff --git a/tests/test_energy.py b/tests/test_energy.py index 5493069..1dc878b 100644 --- a/tests/test_energy.py +++ b/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, }, } diff --git a/tests/test_rng.py b/tests/test_rng.py index cbb514f..6dda0a8 100644 --- a/tests/test_rng.py +++ b/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)