From 017197e8271b60ac5ab6f17aef87b1da76a20d0f Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 26 Oct 2026 15:24:14 +0000
Subject: [PATCH] Flag virtual sites on ghost atoms as FEP ghosts.

---
 CHANGELOG.md         |  1 +
 src/loch/_sampler.py | 11 +++++++
 tests/test_vsites.py | 73 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8d0b1dd..d1cc2c6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ Changelog
 * Fix OpenCL kernel address-space errors on AMD GPUs by placing program-scope constants in `__constant` memory and copying the GCMC target into private memory before calling `distance2` [#45](https://github.com/OpenBioSim/loch/pull/45).
 * 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).
 
 [2026.2.0](https://github.com/openbiosim/loch/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index f7a43a1..39e833b 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -2813,6 +2813,17 @@ def _initialise_gpu_memory(self):
                             omm_idx = sire_idx + int(self._vsite_atom_offsets[sire_idx])
                             is_ghost_fep[omm_idx] = 1
 
+                # Virtual sites follow the molecule's atoms and inherit the
+                # ghost status of their parent atom, matching Sire.
+                if mol.number() in self._mol_vsite_charges:
+                    sire_start = atoms.find(mol.atoms()[0])
+                    omm_start = sire_start + int(self._vsite_atom_offsets[sire_start])
+                    vsites = mol.property("virtual_sites")
+                    for k in range(len(self._mol_vsite_charges[mol.number()])):
+                        parent = int(vsites[str(k)]["vs_indices"][0])
+                        if is_ghost_fep[omm_start + parent] == 1:
+                            is_ghost_fep[omm_start + mol.num_atoms() + k] = 1
+
             # Convert to GPU array.
             is_ghost_fep = self._backend.to_gpu(is_ghost_fep.astype(_np.int32))
 
diff --git a/tests/test_vsites.py b/tests/test_vsites.py
index 225b491..c9e8e4e 100644
--- a/tests/test_vsites.py
+++ b/tests/test_vsites.py
@@ -208,3 +208,76 @@ def test_flag_ghost_waters_sire_indices(bpti, platform):
     # Every initially-ghost water molecule should be flagged.
     n_flagged = len(flagged["property is_ghost_water"].molecules())
     assert n_flagged == sampler._num_ghost_waters
+
+
+@pytest.mark.skipif(
+    "CUDA_VISIBLE_DEVICES" not in os.environ,
+    reason="Requires CUDA enabled GPU.",
+)
+@pytest.mark.parametrize("platform", ["cuda", "opencl"])
+def test_fep_vsite_ghost_flags(sd12, platform):
+    """
+    Virtual sites on a perturbable molecule inherit the ghost status of their
+    parent atom, so that they use soft-core interactions with the waters.
+    """
+    import numpy as np
+
+    mols, reference = sd12
+
+    lig = mols["property is_perturbable"].molecules()[0]
+    n_atoms = lig.num_atoms()
+
+    ghost = [
+        i
+        for i, atom in enumerate(lig.atoms())
+        if np.isclose(atom.property("charge0").value(), 0.0)
+        and np.isclose(atom.property("LJ0").epsilon().value(), 0.0)
+    ]
+    real = [i for i in range(n_atoms) if i not in ghost]
+    assert len(ghost) > 0
+
+    # Virtual site 0 is parented to a ghost atom, virtual site 1 to a real atom.
+    vsite_dict = {}
+    for k, parent in enumerate([ghost[0], real[0]]):
+        vsite_dict[str(k)] = {
+            "vs_indices": [parent] + [i for i in real if i != parent][:2],
+            "vs_ows": [1, 0, 0],
+            "vs_xs": [1, -1, 0],
+            "vs_ys": [0, 1, -1],
+            "vs_local": [0.03, 0, 0],
+        }
+    parents = {str(i): [] for i in range(n_atoms)}
+    parents[str(ghost[0])] = [0]
+    parents[str(real[0])] = [1]
+
+    cursor = lig.cursor()
+    cursor.set("n_virtual_sites", 2)
+    for name in ("vs_charges", "vs_charges0", "vs_charges1"):
+        cursor.set(name, [0.1, -0.1])
+    cursor.set("virtual_sites", vsite_dict)
+    cursor.set("parents", parents)
+
+    mols_with_vs = mols.clone()
+    mols_with_vs.update(cursor.commit())
+
+    sampler = GCMCSampler(
+        mols_with_vs,
+        reference=reference,
+        cutoff_type="rf",
+        cutoff="10 A",
+        lambda_schedule=sr.cas.LambdaSchedule.standard_morph(),
+        lambda_value=0.5,
+        ghost_file=None,
+        log_file=None,
+        test=True,
+        platform=platform,
+        seed=42,
+    )
+
+    is_ghost_fep = sampler._backend.from_gpu(sampler._gpu_is_ghost_fep)
+
+    lig_start = mols_with_vs.atoms().find(lig.atoms()[0])
+    vsite_start = lig_start + sampler._vsite_atom_offsets[lig_start] + n_atoms
+
+    assert is_ghost_fep[vsite_start] == 1
+    assert is_ghost_fep[vsite_start + 1] == 0