fix(space): correct centre offset and unbias SphericalShell sampling (#38)

## Summary

Two related bugs in `maws/space.py` that silently distort sampling for
`Sphere` and `SphericalShell`:

- **Centre offset ignored in `generator()`**
([maws/space.py:270](maws/space.py#L270),
[maws/space.py:339](maws/space.py#L339)) — both classes returned
coordinates around the origin instead of `self.centre`. Any caller
passing a non-zero centre (e.g. the ligand COM in `maws2023.py`) was
silently sampling in the wrong region of space. `Box.generator()`
already does this correctly; matched the same pattern.
- **Biased radial + polar-angle sampling in
`SphericalShell.generator()`** — `r ~ Uniform(R_in, R_out)` over-weights
the inner shell (no r² Jacobian) and `psi ~ Uniform(0, π)` over-weights
the poles. `Sphere.generator()` was already corrected in a prior change;
brought `SphericalShell` to the same volume-correct form:
  - `r = (u·(R_out³ − R_in³) + R_in³)^(1/3)` for u ∈ [0, 1]
  - direction via `cos(psi)` uniform in [-1, 1]

These matter because MAWS energy/entropy averages over the sampled poses
(`S(energies, beta=BETA)` in the first MAWS step) — biased sampling
biases the result.

## Tests

Existing `tests/test_space.py` only used `centre=[0,0,0]`, which masked
the centre-offset bug entirely. Added tests that:
- Pin samples within the radius/shell measured from a non-zero centre
(deterministic, would fail before fix).
- Statistically pin the radial mean (E[r] ≈ 8.04 vs. biased 7.5) and
direction (E[(z/r)²] ≈ 1/3 vs. biased 0.50) at N=10_000 with a tolerance
band that cleanly separates correct from biased.

`pytest tests/test_space.py` — 17/17 pass after the fix.

## Test plan

- [ ] CI passes on `main`
- [ ] `pytest tests/test_space.py` locally — 17 passed
- [ ] No callers regressed (Sphere/SphericalShell are not currently
invoked in `maws2023.py`; `Box`/`Cube` paths unchanged)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Author: siddharth7113

maws/space.py

@@ -269,9 +269,9 @@ def generator(self):
         rotation_angle = np.random.uniform(0, 2 * np.pi)
         result = np.array(
             [
-                r * np.cos(phi) * np.sin(psi),
-                r * np.sin(phi) * np.sin(psi),
-                r * np.cos(psi),
+                self.centre[0] + r * np.cos(phi) * np.sin(psi),
+                self.centre[1] + r * np.sin(phi) * np.sin(psi),
+                self.centre[2] + r * np.cos(psi),
                 x,
                 y,
                 z,
@@ -332,15 +332,21 @@ def generator(self):
         """
         axis = np.random.uniform(-1, 1, 3)
         x, y, z = axis / np.linalg.norm(axis)
-        r = np.random.uniform(self.innerRadius, self.outerRadius)
+        # Uniform-in-volume radial sampling: r ~ (u·(R_out³ − R_in³) + R_in³)^(1/3)
+        u = np.random.uniform(0, 1)
+        r = (u * (self.outerRadius**3 - self.innerRadius**3) + self.innerRadius**3) ** (
+            1 / 3
+        )
         phi = np.random.uniform(0, 2 * np.pi)
-        psi = np.random.uniform(0, np.pi)
+        # Uniform direction on the sphere: cos(psi) uniform in [-1, 1]
+        cos_psi = np.random.uniform(-1, 1)
+        psi = np.arccos(cos_psi)
         rotation_angle = np.random.uniform(0, 2 * np.pi)
         result = np.array(
             [
-                r * np.cos(phi) * np.sin(psi),
-                r * np.sin(phi) * np.sin(psi),
-                r * np.cos(psi),
+                self.centre[0] + r * np.cos(phi) * np.sin(psi),
+                self.centre[1] + r * np.sin(phi) * np.sin(psi),
+                self.centre[2] + r * np.cos(psi),
                 x,
                 y,
                 z,

tests/test_space.py

@@ -93,6 +93,26 @@ def test_sphere_sample_within_radius(self):
             distance = np.linalg.norm(pos)
             assert distance <= 10.0, f"Sample {pos} outside sphere radius"
 
+    def test_sphere_sample_offset_by_centre(self):
+        """Sphere samples are within radius of the (non-zero) centre, not the origin."""
+        centre = np.array([50.0, -30.0, 12.0])
+        sphere = Sphere(radius=10.0, centre=centre)
+        for _ in range(50):
+            sample = sphere.generator()
+            pos = np.array(sample[:3])
+            distance = np.linalg.norm(pos - centre)
+            assert distance <= 10.0 + 1e-9, (
+                f"Sample {pos} not within radius of centre {centre}"
+            )
+
+    def test_sphere_is_in_membership(self):
+        """Sphere.is_in() correctly accepts inside points and rejects outside ones."""
+        centre = np.array([10.0, 10.0, 10.0])
+        sphere = Sphere(radius=5.0, centre=centre)
+        assert bool(sphere.is_in(np.array([10.0, 10.0, 10.0])))
+        assert bool(sphere.is_in(np.array([13.0, 10.0, 10.0])))
+        assert not bool(sphere.is_in(np.array([100.0, 0.0, 0.0])))
+
 
 class TestSphericalShell:
     """Tests for SphericalShell sampling space."""
@@ -112,6 +132,60 @@ def test_shell_sample_within_bounds(self):
             distance = np.linalg.norm(pos)
             assert 5.0 <= distance <= 10.0, f"Sample {pos} outside shell bounds"
 
+    def test_shell_sample_offset_by_centre(self):
+        """SphericalShell samples sit in the shell measured from centre, not origin."""
+        centre = np.array([50.0, -30.0, 12.0])
+        shell = SphericalShell(innerRadius=5.0, outerRadius=10.0, centre=centre)
+        for _ in range(50):
+            sample = shell.generator()
+            pos = np.array(sample[:3])
+            distance = np.linalg.norm(pos - centre)
+            assert 5.0 - 1e-9 <= distance <= 10.0 + 1e-9, (
+                f"Sample {pos} outside shell of centre {centre}"
+            )
+
+    def test_shell_radial_distribution_is_uniform_in_volume(self):
+        """Shell sampling must be uniform per unit volume (∝ r²), not uniform in r.
+
+        For R_in=5, R_out=10 with the volume-correct CDF
+            r = (u·(R_out³ − R_in³) + R_in³)^(1/3),
+        E[r] = (3/4)·(R_out⁴ − R_in⁴)/(R_out³ − R_in³) ≈ 8.036.
+        A naive `r ~ Uniform(R_in, R_out)` gives E[r] = 7.5 — well outside
+        the statistical band of the correct sampler at N=10_000.
+        """
+        rng = np.random.default_rng(0)
+        np.random.seed(0)  # generator() uses np.random
+        R_in, R_out = 5.0, 10.0
+        shell = SphericalShell(innerRadius=R_in, outerRadius=R_out, centre=[0, 0, 0])
+        radii = np.array([np.linalg.norm(shell.generator()[:3]) for _ in range(10_000)])
+        expected = (3 / 4) * (R_out**4 - R_in**4) / (R_out**3 - R_in**3)
+        # Tight band: correct sampler lands within ~0.05 of 8.036; biased
+        # sampler is at 7.5, well outside this band.
+        assert abs(radii.mean() - expected) < 0.1, (
+            f"E[r] = {radii.mean():.3f}, expected ~{expected:.3f} (uniform-in-r bias?)"
+        )
+        _ = rng  # silence unused
+
+    def test_shell_direction_is_uniform_on_sphere(self):
+        """Shell directions must be uniform on the sphere (E[z²/r²] = 1/3).
+
+        With `psi ~ Uniform(0, π)` the polar angle is biased toward the
+        poles, giving E[cos²(psi)] = 1/2 instead of 1/3.
+        """
+        np.random.seed(1)
+        R_in, R_out = 5.0, 10.0
+        shell = SphericalShell(innerRadius=R_in, outerRadius=R_out, centre=[0, 0, 0])
+        samples = np.array([shell.generator()[:3] for _ in range(10_000)])
+        # Direction vectors (unit), then mean of z²
+        radii = np.linalg.norm(samples, axis=1, keepdims=True)
+        unit_z = (samples / radii)[:, 2]
+        mean_z2 = float((unit_z**2).mean())
+        # Uniform-on-sphere: 1/3 ≈ 0.333. Biased: 1/2 = 0.500. Tolerance
+        # 0.04 cleanly separates the two at N=10_000.
+        assert abs(mean_z2 - 1 / 3) < 0.04, (
+            f"E[(z/r)²] = {mean_z2:.3f}, expected ~0.333 (polar-angle bias?)"
+        )
+
 
 class TestNAngles:
     """Tests for NAngles (torsion angle) sampling space."""