diff --git a/README.md b/README.md
index 4724d45..bb6b100 100644
--- a/README.md
+++ b/README.md
@@ -72,14 +72,16 @@ WaterFlow processes structure files through several stages to create training-re
 - For atoms with alternate conformations, the highest-occupancy conformer is selected
 
 **Crystal Contact Detection**
-- Uses PyMOL's `symexp` to generate symmetry mates within 5.0Å cutoff
-- Symmetry mate atoms are included as additional protein context when `include_mates=True`
-- Mate atoms are stored separately for proper handling during training
+- Uses PyMOL's `symexp` to generate symmetry mates, keeping whole residues and whole ligand entities with any atom within the cutoff of the ASU. Runs only when `include_mates=True`; a no-mates cache never invokes PyMOL
+- Protein mates and ligand mates are selected separately by PyMOL's own classifiers, so `is_ligand` stays exact for mate nodes too
+- **Mate waters are never selected.** A mate water is a symmetry image of an ASU water, which is what the model predicts, so keeping it as context leaks the label
+- Symmetry also maps atoms onto themselves (special positions) and reaches one residue through two operators. Mate atoms within 0.3Å of an ASU atom, a target water, or an already-kept mate atom are dropped (`dedup_mate_atoms`); mate ligands are judged whole, so a ligand is never fragmented (`dedup_mate_ligands_by_residue`)
+- A mate keeps its source residue's `(chain, res_id, ins_code)`, so it inherits that residue's ESM row through `emb_res_idx` instead of a zero vector, and it joins the distance-filter reference so a water in a crystal contact — near a neighbour surface but far from the ASU — is not dropped as solvent-far
 
 **Graph Representation**
 - Node types: `protein` (ASU + symmetry mates + ligands), `water` (ground truth)
-- ASU ligand atoms are appended after ASU and mate atoms and carry the boolean `is_ligand` mask plus `residue_index = -1` (they have no residue embedding, so residue pooling masks them out)
-- `is_ligand` marks **ASU ligands only**. Symmetry-mate generation is currently unfiltered, so mate nodes can include HETATM and water atoms that `is_ligand` does not mark — see `TODO(mates)` in `ProteinWaterDataset._preprocess_one`. Don't treat `is_ligand` as an exhaustive ligand selector
+- Ligand atoms are appended after ASU and mate atoms and carry the boolean `is_ligand` mask plus `residue_index = -1` (they have no residue embedding, so residue pooling masks them out)
+- `is_mate` marks every non-ASU node, protein or ligand. The flow prior anchors on `~is_mate` so sampled waters start where the targets live
 - Edge types (defined in `src/constants.py`):
   - `('protein', 'pp', 'protein')`: protein-protein edges
   - `('protein', 'pw', 'water')`: protein to water
@@ -114,20 +116,22 @@ Preprocessed data is cached under `--processed_dir` in a three-layer architectur
 │       - protein_pos: centered protein coordinates (N, 3)
 │       - protein_x: element one-hot encoding (N, 16)
 │       - protein_res_idx: residue indices for grouping
-│       - is_ligand: bool mask marking the appended ASU ligand atoms (N,)
+│       - is_ligand: bool mask marking the ligand atoms (N,)
+│       - is_mate: bool mask marking the symmetry-mate atoms (N,)
+│       - emb_res_idx: embedding row per atom; -1 means no row (N,)
 │       - water_pos, water_x: water coordinates and features
 │       - num_asu_protein: ASU atom count (mate boundary metadata)
-│       # Note: When include_mates=True, mate atoms are concatenated into
-│       # protein_pos/protein_x, and ASU ligand atoms are appended after those.
-│       # Node order is [ASU protein | mates | ASU ligands]. Recover blocks via:
-│       #   ASU protein atoms = protein_pos[:num_asu_protein]
-│       #   ASU ligand atoms  = protein_pos[is_ligand]          # always last
-│       #   Mate atoms        = protein_pos[num_asu_protein:][~is_ligand[num_asu_protein:]]
+│       # Node order is [ASU protein | mate protein | ASU ligand | mate ligand],
+│       # so the two masks recover every block:
+│       #   ASU protein  = ~is_mate & ~is_ligand    (== the first num_asu_protein)
+│       #   mate protein =  is_mate & ~is_ligand
+│       #   ASU ligand   = ~is_mate &  is_ligand
+│       #   mate ligand  =  is_mate &  is_ligand
 │       #
-│       # is_ligand marks ASU ligands ONLY -- it is not an exhaustive ligand
-│       # selector. The mate block is unfiltered (see TODO(mates) in
-│       # _preprocess_one), so mate atoms may include HETATM/ligand/water atoms
-│       # that are NOT marked by is_ligand.
+│       # emb_res_idx indexes the ESM table: mate atoms carry the row of the ASU
+│       # residue they are a symmetry image of, and every ligand carries -1,
+│       # which reads as a zero row.
+├── <geometry_dir>/_filter_meta.json   # settings this directory was built with
 ├── esm/                   # ESM embeddings (per-residue)
 │   └── <pdb_id>_final.pt
 │       - residue_embeddings: ESM3 embeddings (N_res, embed_dim)
@@ -153,13 +157,29 @@ configs that produce different graphs never share a directory:
 
 The base name comes from `--geometry_cache_name` (default `geometry`).
 
+**Filter Provenance:**
+
+Filtering happens *before* the cache is written, so the thresholds are a property of the
+directory, not of the run reading it — and the `.pt` files record none of them. Each geometry
+directory therefore carries a `_filter_meta.json` sidecar holding the per-water filters and
+their toggles, the structure-level checks that decide which entries exist at all
+(`min_water_residue_ratio`, `max_com_dist`, `max_clash_fraction`, `clash_dist`,
+`interface_dist_threshold`), and the graph parameters behind the cached PP edges (`cutoff`,
+`max_neighbors`).
+
+The first run with `preprocess=True` writes it; every later run compares against it and
+**refuses to start** on a mismatch rather than appending differently filtered entries to the
+same directory. A disabled filter records `null` for its threshold, which cannot have changed
+the cached waters. Directories built before this existed have no sidecar: they load, and warn
+that their provenance is unverifiable, until a preprocessing run stamps them — so check your
+thresholds match the cache before that first run.
+
 **Cache Generation Notes:**
 - Geometry cache is generated automatically when `preprocess=True` (default)
 - ESM/SLAE caches require running the respective `generate_*_embeddings.py` scripts first
 - Preprocessing failures are logged to `<geometry_dir>/preprocessing_failures.log`
-- Geometry caches built before ligand support lack the `is_ligand` field and will fail to
-  load with a `KeyError`. Delete the geometry cache directory and let it regenerate — the
-  cached graphs are stale, not merely missing a field
+- A cache file missing any field the loader reads (`is_ligand`, `is_mate`, `emb_res_idx`, …)
+  raises `KeyError`. Delete the geometry cache directory and let it regenerate
 
 ## Environment Setup
 
@@ -303,7 +323,7 @@ These checks determine whether a structure is included in training:
 | `--max_com_dist` | `25.0` | Max protein-water center-of-mass distance (A) |
 | `--max_clash_fraction` | `0.05` | Max fraction of waters clashing with protein |
 | `--clash_dist` | `2.0` | Distance threshold for clash detection (A) |
-| `--min_water_residue_ratio` | `0.6` | Minimum waters per residue ratio |
+| `--min_water_residue_ratio` | `0.1` | Minimum waters per residue ratio |
 
 ### Per-Water Quality Filters
 
@@ -313,7 +333,7 @@ These filters remove individual low-quality waters (can be toggled):
 |-----------|---------|-------------|-------------|
 | `--max_protein_dist` | `5.0` | `--no_filter_by_distance` | Remove waters far from protein |
 | `--min_edia` | `0.4` | `--no_filter_by_edia` | Remove waters with low EDIA scores |
-| `--max_bfactor_zscore` | `1.5` | `--no_filter_by_bfactor` | Remove waters with high B-factor |
+| `--max_bfactor_zscore` | `2.0` | `--no_filter_by_bfactor` | Remove waters with high B-factor |
 
 <details>
 <summary><strong>About EDIA Scores</strong></summary>
diff --git a/scripts/generate_slae_embeddings.py b/scripts/generate_slae_embeddings.py
index 2da2926..8e42f61 100644
--- a/scripts/generate_slae_embeddings.py
+++ b/scripts/generate_slae_embeddings.py
@@ -1,9 +1,8 @@
 """
 Precompute SLAE embeddings for protein structures and save to separate cache files.
 
-NOTE: This SLAE encoder is legacy and is NOT currently used. We primarily use the
-ESM encoder (see scripts/generate_esm_embeddings.py). This script is retained for
-reference/reproducibility only.
+NOTE: The SLAE encoder is NOT currently used. We primarily use the ESM encoder
+(see scripts/generate_esm_embeddings.py); this script is kept for reproducibility.
 
 This script:
 1. Reads a split file containing PDB entries
diff --git a/scripts/inference.py b/scripts/inference.py
index e73b919..d263a6d 100644
--- a/scripts/inference.py
+++ b/scripts/inference.py
@@ -215,10 +215,10 @@ def _extract_dataset_filter_config(config: dict) -> dict:
         "max_clash_fraction": config.get("max_clash_fraction", 0.05),
         "clash_dist": config.get("clash_dist", 2.0),
         "interface_dist_threshold": config.get("interface_dist_threshold", 4.0),
-        "min_water_residue_ratio": config.get("min_water_residue_ratio", 0.6),
+        "min_water_residue_ratio": config.get("min_water_residue_ratio", 0.1),
         "max_protein_dist": config.get("max_protein_dist", 5.0),
         "min_edia": config.get("min_edia", 0.4),
-        "max_bfactor_zscore": config.get("max_bfactor_zscore", 1.5),
+        "max_bfactor_zscore": config.get("max_bfactor_zscore", 2.0),
         "filter_by_distance": config.get("filter_by_distance", True),
         "filter_by_edia": config.get("filter_by_edia", True),
         "filter_by_bfactor": config.get("filter_by_bfactor", True),
diff --git a/scripts/train.py b/scripts/train.py
index f5bc0be..7cbad8f 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -149,8 +149,11 @@ def parse_args():
     p.add_argument(
         "--min_water_residue_ratio",
         type=float,
-        default=0.6,
-        help="Quality: minimum waters/residue ratio required per structure.",
+        default=0.1,
+        help=(
+            "Quality: minimum waters/residue ratio required per structure. Applied "
+            "at cache-write time, so it decides which structures the cache holds."
+        ),
     )
 
     # per-water filtering (toggleable)
@@ -169,8 +172,12 @@ def parse_args():
     p.add_argument(
         "--max_bfactor_zscore",
         type=float,
-        default=1.5,
-        help="Water filter: remove waters with normalized B-factor above this threshold.",
+        default=2.0,
+        help=(
+            "Water filter: remove waters with normalized B-factor above this "
+            "threshold. Baked in at cache-write time, so a warm cache built at a "
+            "different value is refused rather than extended."
+        ),
     )
     p.add_argument(
         "--no_filter_by_distance",
diff --git a/src/constants.py b/src/constants.py
index 5e2604f..9a06350 100644
--- a/src/constants.py
+++ b/src/constants.py
@@ -7,9 +7,6 @@
 NODE_FEATURE_DIM = 16  # Default node scalar feature dimension
 
 # Native widths of the cached embeddings produced by scripts/generate_*_embeddings.py.
-# These are fixed by the upstream models, not tunable: ESM3-open emits 1536-wide
-# per-residue vectors, SLAE emits 128-wide per-atom vectors. Cached encoders take the
-# width as a required config key (embedding_dim); these are the values to pass.
 ESM_EMBEDDING_DIM = 1536
 SLAE_EMBEDDING_DIM = 128
 
diff --git a/src/dataset.py b/src/dataset.py
index 4a6ba0b..609c5e3 100644
--- a/src/dataset.py
+++ b/src/dataset.py
@@ -13,6 +13,8 @@
 
 import itertools
 import json
+import os
+import re
 from collections import OrderedDict
 from pathlib import Path
 
@@ -24,6 +26,7 @@
 from biotite.structure.io.pdb import get_structure, PDBFile
 from biotite.structure.io.pdbx import CIFFile, get_structure as get_structure_cif
 from loguru import logger
+from scipy.spatial import cKDTree
 from scipy.spatial.distance import cdist
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
@@ -44,6 +47,10 @@
 )
 
 
+# Per-directory record of the settings a geometry cache was built with.
+FILTER_META_FILENAME = "_filter_meta.json"
+
+
 def element_onehot(symbols: list[str]) -> Tensor:
     """One-hot encoding with 'other' bucket at end."""
     other_idx = len(ELEMENT_VOCAB)
@@ -112,24 +119,38 @@ def parse_asu_with_biotite(
 
 
 def get_crystal_contacts_pymol(
-    struc_path: str, cutoff: float = 5.0
+    struc_path: str,
+    cutoff: float = 5.0,
+    include_ligands: bool = False,
 ) -> dict[str, np.ndarray | list]:
     """
-    Extract ASU and symmetry mate atoms within crystal contact distance.
+    Extract ASU and symmetry-mate atoms within crystal contact distance.
 
-    Uses PyMOL's symexp command to generate symmetry mates and selects
-    interface atoms within the specified cutoff distance.
+    PyMOL's symexp generates the mates; `byres` then keeps whole residues and
+    whole ligand entities with any atom within `cutoff` of the ASU. Protein and
+    ligand mates come back under separate keys, classified by PyMOL itself
+    (`polymer.protein` vs the non-protein, non-solvent remainder), so nothing
+    downstream needs residue-name heuristics.
+
+    Mate waters are never selected: a mate water is a symmetry image of an ASU
+    water, which is a prediction target, so keeping it as context is a label
+    leak. Protein and ligand contact surfaces are genuine context.
 
     Args:
-        struc_path: Path to structure file (PDB/CIF) with crystal symmetry information
-        cutoff: Distance cutoff in Angstroms for interface detection
+        struc_path: Structure file (PDB/CIF) carrying crystal symmetry.
+        cutoff: Interface distance cutoff in Angstroms.
+        include_ligands: Also collect whole ligand/ion/cofactor/nucleic-acid mate
+            entities (still never waters). Off by default: protein mates alone.
 
     Returns:
         Dict with keys:
             - 'asu_coords': (N_asu, 3) ASU atom coordinates
-            - 'mate_coords': (N_mate, 3) symmetry mate atom coordinates
-            - 'asu_atoms': List of PyMOL atom objects for ASU
-            - 'mate_atoms': List of PyMOL atom objects for mates
+            - 'asu_atoms': List of PyMOL atom objects for the ASU
+            - 'mate_coords': (N_mate, 3) whole protein-mate residues
+            - 'mate_atoms': List of PyMOL atom objects for protein mates
+            - 'mate_ligand_coords': (M, 3) whole ligand-mate entities, empty
+              unless include_ligands
+            - 'mate_ligand_atoms': List of PyMOL atom objects for ligand mates
     """
     with pymol2.PyMOL() as pm:
         cmd = pm.cmd
@@ -138,24 +159,39 @@ def get_crystal_contacts_pymol(
         obj = "struct"
         cmd.load(struc_path, obj)
         cmd.symexp("sym", obj, obj, cutoff)
-        cmd.select("interface", f"byres (sym* within {cutoff} of {obj})")
 
-        asu_coords = cmd.get_coords(obj, state=1)
-        mate_coords = cmd.get_coords("sym* and interface", state=1)
-        asu_atoms = cmd.get_model(obj, state=1).atom
-        mate_atoms = cmd.get_model("sym* and interface", state=1).atom
+        def _coords(selection: str) -> np.ndarray:
+            coords = cmd.get_coords(selection, state=1)
+            return coords if coords is not None else np.zeros((0, 3), dtype=float)
 
-        asu_coords = (
-            asu_coords if asu_coords is not None else np.zeros((0, 3), dtype=float)
-        )
-        mate_coords = (
-            mate_coords if mate_coords is not None else np.zeros((0, 3), dtype=float)
+        # Whole protein-mate residues with any atom within cutoff of the ASU.
+        cmd.select(
+            "iface_prot", f"byres ((sym* and polymer.protein) within {cutoff} of {obj})"
         )
+        mate_coords = _coords("iface_prot")
+        mate_atoms = cmd.get_model("iface_prot", state=1).atom
+
+        # Whole ligand-mate entities (non-protein, non-water het atoms: ligands,
+        # ions, cofactors, nucleic acids).
+        if include_ligands:
+            cmd.select(
+                "iface_lig",
+                f"byres ((sym* and (not polymer.protein) and (not solvent)) "
+                f"within {cutoff} of {obj})",
+            )
+            mate_ligand_coords = _coords("iface_lig")
+            mate_ligand_atoms = cmd.get_model("iface_lig", state=1).atom
+        else:
+            mate_ligand_coords = np.zeros((0, 3), dtype=float)
+            mate_ligand_atoms = []
+
         return {
-            "asu_coords": asu_coords,
+            "asu_coords": _coords(obj),
+            "asu_atoms": cmd.get_model(obj, state=1).atom,
             "mate_coords": mate_coords,
-            "asu_atoms": asu_atoms,
             "mate_atoms": mate_atoms,
+            "mate_ligand_coords": mate_ligand_coords,
+            "mate_ligand_atoms": mate_ligand_atoms,
         }
 
 
@@ -182,8 +218,6 @@ def match_atoms_to_coords(
     if target_coords.shape[0] == 0 or len(atoms) == 0:
         return []
 
-    from scipy.spatial import cKDTree
-
     tree = cKDTree(atoms.coord)
     dists, nearest = tree.query(target_coords, k=1, distance_upper_bound=tolerance)
     within = np.isfinite(dists) & (nearest < len(atoms))
@@ -199,6 +233,114 @@ def match_atoms_to_coords(
     return matched
 
 
+def dedup_mate_atoms(
+    mate_coords: np.ndarray,
+    mate_atoms: list,
+    reference_coords: np.ndarray,
+    tol: float = 0.3,
+) -> tuple[np.ndarray, list]:
+    """
+    Drop mate atoms coincident with a reference atom or an already-kept mate atom.
+
+    Crystal symmetry creates coincidences: an atom on a rotation or screw axis
+    maps onto itself, and one residue can be reached through two operators. Left
+    alone each becomes an independent node, giving duplicates joined by ~0 A
+    edges -- and, for a target water on a special position, a label leak.
+
+    Args:
+        mate_coords: (N, 3) mate atom coordinates, uncentered.
+        mate_atoms: Parallel list of mate atom objects, kept in lockstep.
+        reference_coords: (M, 3) uncentered ASU coordinates.
+        tol: Coincidence radius in Angstroms.
+
+    Returns:
+        (kept_coords, kept_atoms). The first atom of a coincident group is the
+        one kept, so the result depends on input order.
+    """
+    n = mate_coords.shape[0]
+    if n == 0:
+        return mate_coords, mate_atoms
+
+    # An empty reference tree answers inf, so no guard is needed here.
+    drop = cKDTree(reference_coords).query(mate_coords, k=1)[0] < tol
+
+    # Self-dedup is a first-win sweep. One tree answers every lookup, so the
+    # sweep only walks each atom's coincident neighbors.
+    neighbors = cKDTree(mate_coords).query_ball_point(mate_coords, r=tol)
+    kept = np.zeros(n, dtype=bool)
+    for i in range(n):
+        if drop[i]:
+            continue
+        earlier = [j for j in neighbors[i] if j < i and kept[j]]
+        # query_ball_point includes r, this sweep is strict.
+        dists = np.linalg.norm(mate_coords[earlier] - mate_coords[i], axis=1)
+        kept[i] = not (dists < tol).any()
+
+    keep_idx = np.flatnonzero(kept)
+    return mate_coords[keep_idx], [mate_atoms[i] for i in keep_idx]
+
+
+def dedup_mate_ligands_by_residue(
+    lig_coords: np.ndarray,
+    lig_atoms: list,
+    reference_coords: np.ndarray,
+    tol: float = 0.3,
+    image_frac: float = 0.5,
+) -> tuple[np.ndarray, list]:
+    """
+    Drop whole mate-ligand entities that are symmetry images of ASU atoms.
+
+    Unlike `dedup_mate_atoms`, which works per atom, this works per entity so a
+    ligand is never fragmented: it goes only when the whole ligand is a redundant
+    copy. Genuine neighbor-cell ligands are kept whole.
+
+    Args:
+        lig_coords: (M, 3) mate-ligand atom coordinates, uncentered.
+        lig_atoms: Parallel list of mate-ligand atom objects.
+        reference_coords: Uncentered ASU coordinates.
+        tol: Coincidence radius in Angstroms.
+        image_frac: Drop a ligand when more than this fraction of its atoms are
+            coincident with the reference.
+
+    Returns:
+        (kept_coords, kept_atoms) with whole symmetry-image ligands removed.
+    """
+    if len(lig_atoms) == 0:
+        return lig_coords, lig_atoms
+
+    ref_tree = cKDTree(reference_coords)
+    # Group atom indices by ligand entity (chain, residue id, segment).
+    groups = {}
+    for i, atom in enumerate(lig_atoms):
+        key = (atom.chain, atom.resi, getattr(atom, "segi", ""))
+        groups.setdefault(key, []).append(i)
+
+    keep_idx: list[int] = []
+    for idxs in groups.values():
+        if np.mean(ref_tree.query(lig_coords[idxs], k=1)[0] < tol) <= image_frac:
+            keep_idx.extend(idxs)  # genuine neighbor ligand: keep whole
+    keep_idx.sort()
+
+    return lig_coords[keep_idx], [lig_atoms[i] for i in keep_idx]
+
+
+def _parse_pdb_resi(resi) -> tuple[int, str] | None:
+    """
+    Parse a PyMOL residue identifier, which may carry an insertion code.
+
+    Args:
+        resi: Residue id as PyMOL exposes it, e.g. "52", "-3", "52A".
+
+    Returns:
+        (res_id, ins_code), or None when there is no integer part, which the
+        caller scores with a zero embedding rather than crashing on.
+    """
+    match = re.match(r"^\s*(-?\d+)\s*([A-Za-z]?)\s*$", str(resi))
+    if match is None:
+        return None
+    return int(match.group(1)), match.group(2).strip()
+
+
 def _make_undirected(edge_index: torch.Tensor) -> torch.Tensor:
     """
     Convert directed edges to undirected by adding reverse edges.
@@ -741,10 +883,10 @@ def __init__(
         max_clash_fraction: float = 0.05,
         clash_dist: float = 2.0,
         interface_dist_threshold: float = 4.0,
-        min_water_residue_ratio: float = 0.6,
+        min_water_residue_ratio: float = 0.1,
         max_protein_dist: float = 5.0,
         min_edia: float = 0.4,
-        max_bfactor_zscore: float = 1.5,
+        max_bfactor_zscore: float = 2.0,
         filter_by_distance: bool = True,
         filter_by_edia: bool = True,
         filter_by_bfactor: bool = True,
@@ -850,6 +992,8 @@ def __init__(
 
         self.entries = self._parse_pdb_list(pdb_list_file)
 
+        self._sync_filter_meta(write=preprocess)
+
         if preprocess:
             self._preprocess_all()
 
@@ -914,6 +1058,76 @@ def _parse_pdb_list(self, pdb_list_file: str) -> list[dict]:
         logger.info(f"Loaded {len(entries)} entries from {pdb_list_file}")
         return entries
 
+    def _sync_filter_meta(self, write: bool) -> None:
+        """
+        Refuse to read or extend a cache built under different settings.
+
+        Filtering happens before the cache is written, so these are properties of
+        the directory rather than of the run reading it -- and the .pt files
+        record none of them. Writing entries under different settings would leave
+        one directory holding two populations no later reader can tell apart.
+
+        Args:
+            write: Create when it is absent. Only runs that may add
+                entries (preprocess=True) claim a directory this way.
+
+        Raises:
+            ValueError: If the recorded settings differ from this run's.
+        """
+        meta_path = self.geometry_dir / FILTER_META_FILENAME
+        # A disabled filter's threshold is None: it never touched the cached
+        # waters, so it must not make two identical caches look incompatible.
+        current = {
+            "filter_by_distance": self.filter_by_distance,
+            "filter_by_edia": self.filter_by_edia,
+            "filter_by_bfactor": self.filter_by_bfactor,
+            "max_protein_dist": self.max_protein_dist
+            if self.filter_by_distance
+            else None,
+            "min_edia": self.min_edia if self.filter_by_edia else None,
+            "max_bfactor_zscore": self.max_bfactor_zscore
+            if self.filter_by_bfactor
+            else None,
+            "min_water_residue_ratio": self.min_water_residue_ratio,
+            "max_com_dist": self.max_com_dist,
+            "max_clash_fraction": self.max_clash_fraction,
+            "clash_dist": self.clash_dist,
+            "interface_dist_threshold": self.interface_dist_threshold,
+            "cutoff": self.cutoff,
+            "max_neighbors": self.max_neighbors,
+        }
+
+        if meta_path.is_file():
+            with open(meta_path) as f:
+                recorded = json.load(f)
+            differing = [
+                f"{name}: cache={recorded.get(name)!r} run={value!r}"
+                for name, value in current.items()
+                if recorded.get(name) != value
+            ]
+            if differing:
+                raise ValueError(
+                    f"Settings disagree with {meta_path}: {', '.join(differing)}. "
+                    "The cache was filtered at write time, so one directory cannot "
+                    "hold both. Match the recorded values or point "
+                    "geometry_cache_name at a different directory."
+                )
+            return
+
+        if write:
+            self.geometry_dir.mkdir(parents=True, exist_ok=True)
+            # Written through a temp file: cache builds fan out over processes,
+            # and a reader must never catch a half-written sidecar.
+            tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp")
+            with open(tmp_path, "w") as f:
+                json.dump(current, f, indent=2)
+            tmp_path.replace(meta_path)
+        elif any(self.geometry_dir.glob("*.pt")):
+            logger.warning(
+                f"{self.geometry_dir} has no {FILTER_META_FILENAME}; the settings "
+                "its entries were built with cannot be verified."
+            )
+
     def _preprocess_all(self):
         """
         Preprocess all PDB files that don't have cached geometry results.
@@ -986,20 +1200,31 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
         if not chain_valid:
             raise ValueError(f"Quality filter failed: {chain_reason}")
 
-        crystal_data = get_crystal_contacts_pymol(struc_path, self.cutoff)
+        # PyMOL is only needed for symmetry expansion, so a no-mates cache skips
+        # it (and with it the water cross-check below) entirely.
+        if self.include_mates:
+            crystal_data = get_crystal_contacts_pymol(
+                struc_path, self.cutoff, include_ligands=self.include_ligands
+            )
 
-        # Keep only the waters PyMOL also saw. PyMOL's ASU is a superset of
-        # biotite's (it keeps every altloc conformer), so a water missing from it
-        # means the two parses disagree rather than that the water is unwanted.
-        asu_water_indices = match_atoms_to_coords(
-            water_atoms, crystal_data["asu_coords"]
-        )
-        if asu_water_indices:
-            asu_water_mask = np.zeros(len(water_atoms), dtype=bool)
-            asu_water_mask[asu_water_indices] = True
-            water_atoms = water_atoms[asu_water_mask]
-        else:
-            water_atoms = water_atoms[:0]
+            # Keep only the waters PyMOL also saw. PyMOL's ASU is a superset of
+            # biotite's (it keeps every altloc conformer), so a water missing from
+            # it means the two parses disagree rather than that the water is
+            # unwanted.
+            asu_water_indices = match_atoms_to_coords(
+                water_atoms, crystal_data["asu_coords"]
+            )
+            if asu_water_indices:
+                asu_water_mask = np.zeros(len(water_atoms), dtype=bool)
+                asu_water_mask[asu_water_indices] = True
+                water_atoms = water_atoms[asu_water_mask]
+            else:
+                if len(water_atoms) > 0:
+                    logger.warning(
+                        f"{entry['pdb_id']}: no waters survived the biotite/PyMOL "
+                        f"cross-check (had {len(water_atoms)})"
+                    )
+                water_atoms = water_atoms[:0]
 
         # Per-water filtering is optional; structure-level quality checks below always run.
         use_distance_filter = self.filter_by_distance
@@ -1039,11 +1264,19 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
                 )
             )
 
-            # apply quality filters
+            # Apply quality filters. Mate protein atoms join the distance
+            # reference so a genuine crystal-contact water is not dropped as solvent-far.
+            if self.include_mates and crystal_data["mate_coords"].shape[0] > 0:
+                filter_protein_coords = np.concatenate(
+                    [protein_atoms.coord, crystal_data["mate_coords"]], axis=0
+                )
+            else:
+                filter_protein_coords = protein_atoms.coord
+
             keep_mask = filter_waters_by_quality(
                 water_atoms.coord,
                 water_keys,
-                protein_atoms.coord if use_distance_filter else None,
+                filter_protein_coords if use_distance_filter else None,
                 edia_lookup,
                 bfactor_lookup,
                 max_protein_dist=self.max_protein_dist,
@@ -1100,6 +1333,20 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
         protein_res_idx = torch.from_numpy(
             bts.spread_residue_wise(sanitized_for_idx, np.arange(num_residues))
         ).long()
+
+        # (chain, res_id, ins_code) -> residue index, so a symmetry mate can
+        # inherit the embedding row of the ASU residue it is an image of. Keyed
+        # off the same sanitized parse that defines protein_res_idx, so the index
+        # lines up with the stored ESM rows.
+        asu_reskey_to_residx: dict[tuple[str, int, str], int] = {}
+        for res_i, start in enumerate(bts.get_residue_starts(sanitized_for_idx)):
+            key = (
+                str(sanitized_for_idx.chain_id[start]).strip(),
+                int(sanitized_for_idx.res_id[start]),
+                normalize_ins_code(sanitized_for_idx.ins_code[start]),
+            )
+            asu_reskey_to_residx.setdefault(key, res_i)
+
         num_waters = len(water_atoms)
         ratio_valid, ratio_reason = check_water_residue_ratio(
             num_waters,
@@ -1118,33 +1365,69 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
             water_pos = torch.zeros((0, 3), dtype=torch.float32)
             water_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32)
 
-        # process symmetry mate atoms
-        #
-        # TODO(mates): the mate atom set is unfiltered and inconsistent with the ASU
-        # path. get_crystal_contacts_pymol runs symexp over the whole object and
-        # selects "sym* and interface" with no polymer filter, so mate_atoms carries
-        # het atoms and waters as well as protein, and every one of them becomes a
-        # protein-type node below. Consequences: mate ligands are already included
-        # but never marked in is_ligand (unlike ASU ligands) and are not gated by
-        # include_ligands; mate waters -- symmetry images of the prediction target --
-        # enter as protein context. Fix in dev_crystal_mates.
-        mate_coords = crystal_data["mate_coords"]
-        if mate_coords.shape[0] > 0:
-            mate_pos = torch.tensor(mate_coords, dtype=torch.float32) - center
-            mate_elements = [a.symbol.upper() for a in crystal_data["mate_atoms"]]
-            mate_x = element_onehot(mate_elements)
-
-            # compute mate residue indices (group atoms by actual residue)
-            mate_residue_keys = [(a.chain, a.resi) for a in crystal_data["mate_atoms"]]
-            unique_mate_res = list(dict.fromkeys(mate_residue_keys))  # preserves order
-            mate_res_map = {k: i for i, k in enumerate(unique_mate_res)}
-            mate_res_idx = torch.tensor(
-                [mate_res_map[k] for k in mate_residue_keys], dtype=torch.long
+        # Mate blocks stay empty unless include_mates (and, for ligands,
+        # include_ligands) filled them in below.
+        mate_pos = torch.zeros((0, 3), dtype=torch.float32)
+        mate_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32)
+        mate_res_idx = torch.empty(0, dtype=torch.long)
+        mate_emb_res_idx = torch.empty(0, dtype=torch.long)
+        mate_lig_coords = np.zeros((0, 3), dtype=float)
+        mate_lig_atoms: list = []
+
+        if self.include_mates:
+            # Drop mate atoms coincident with an ASU atom, a target water, or an
+            # already-kept mate atom: special positions and redundant symmetry
+            # images. Uncentered coords; mates are centered below.
+            ref_parts = [protein_atoms.coord]
+            if len(water_atoms):
+                ref_parts.append(water_atoms.coord)
+            # ASU ligands join the reference so a mate ligand that is only their
+            # symmetry image goes too; neighbor-cell ligands stay.
+            if self.include_ligands and len(ligand_atoms) > 0:
+                ref_parts.append(ligand_atoms.coord)
+            reference = np.concatenate(ref_parts, axis=0)
+            mate_coords, mate_atoms = dedup_mate_atoms(
+                crystal_data["mate_coords"], crystal_data["mate_atoms"], reference
             )
-        else:
-            mate_pos = torch.zeros((0, 3), dtype=torch.float32)
-            mate_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32)
-            mate_res_idx = torch.empty(0, dtype=torch.long)
+            # Ligand mates dedup at entity granularity, so a ligand is never
+            # fragmented: whole symmetry images go, genuine neighbors stay.
+            if self.include_ligands:
+                mate_lig_coords, mate_lig_atoms = dedup_mate_ligands_by_residue(
+                    crystal_data["mate_ligand_coords"],
+                    crystal_data["mate_ligand_atoms"],
+                    reference,
+                )
+
+            if mate_coords.shape[0] > 0:
+                mate_pos = torch.tensor(mate_coords, dtype=torch.float32) - center
+                mate_x = element_onehot([a.symbol.upper() for a in mate_atoms])
+
+                # compute mate residue indices (group atoms by actual residue)
+                mate_residue_keys = [(a.chain, a.resi) for a in mate_atoms]
+                unique_mate_res = list(dict.fromkeys(mate_residue_keys))  # keeps order
+                mate_res_map = {k: i for i, k in enumerate(unique_mate_res)}
+                mate_res_idx = torch.tensor(
+                    [mate_res_map[k] for k in mate_residue_keys], dtype=torch.long
+                )
+
+                # A mate keeps its source residue's (chain, resi) and the embedding
+                # is coordinate-free, so the two share one row. -1 means no match,
+                # which reads as a zero embedding.
+                mate_emb_idx = []
+                for atom in mate_atoms:
+                    parsed = _parse_pdb_resi(atom.resi)
+                    mate_emb_idx.append(
+                        asu_reskey_to_residx.get((str(atom.chain).strip(), *parsed), -1)
+                        if parsed is not None
+                        else -1
+                    )
+                mate_emb_res_idx = torch.tensor(mate_emb_idx, dtype=torch.long)
+                unmatched = int((mate_emb_res_idx < 0).sum())
+                if unmatched:
+                    logger.warning(
+                        f"{entry['cache_key']}: {unmatched}/{len(mate_emb_idx)} mate "
+                        "atoms unmatched to an ASU residue (zero embedding for those)"
+                    )
 
         # Compute final protein data based on include_mates flag
         num_asu_protein = protein_pos.size(0)
@@ -1164,29 +1447,48 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
             final_protein_x = protein_x
             final_protein_res_idx = protein_res_idx
 
-        # Append ASU ligand atoms after protein (and mate) atoms when enabled.
-        # is_ligand mask marks which protein-type nodes are ligand atoms.
-        # Ligands always go last so num_asu_protein and mate counts are unaffected,
-        # preserving ESM/SLAE embedding alignment via _pad_atom_embeddings_for_mates.
-        # Only ASU ligands are handled here -- mate het atoms come in unfiltered via
-        # the mate block above, see TODO(mates) there.
+        # Append ligand atoms last, giving the node order ASU protein -> mate
+        # protein -> ASU ligand -> mate ligand: num_asu_protein and the mate count
+        # stay meaningful, which is what keeps ESM/SLAE aligned. Ligands get
+        # residue_index = emb_res_idx = -1 (no residue embedding); residue pooling
+        # masks out those negatives before any scatter (GVPEncoder._pool_by_residue).
+        ligand_blocks = []
         if self.include_ligands and len(ligand_atoms) > 0:
-            ligand_pos = torch.tensor(ligand_atoms.coord, dtype=torch.float32) - center
-            ligand_elements = [str(e).upper() for e in ligand_atoms.element]
-            ligand_x = element_onehot(ligand_elements)
-            final_protein_pos = torch.cat([final_protein_pos, ligand_pos], dim=0)
-            final_protein_x = torch.cat([final_protein_x, ligand_x], dim=0)
-            # Ligand atoms get residue_index = -1 (sentinel; no residue embedding).
-            # The is_ligand mask identifies them; residue-pooling masks out these
-            # negative indices before any scatter (see GVPEncoder._pool_by_residue).
-            ligand_res_idx = torch.full((len(ligand_atoms),), -1, dtype=torch.long)
-            final_protein_res_idx = torch.cat(
-                [final_protein_res_idx, ligand_res_idx], dim=0
+            ligand_blocks.append(
+                (
+                    ligand_atoms.coord,
+                    [str(e).upper() for e in ligand_atoms.element],
+                    False,
+                )
+            )
+        if len(mate_lig_atoms) > 0:
+            ligand_blocks.append(
+                (mate_lig_coords, [a.symbol.upper() for a in mate_lig_atoms], True)
             )
-            is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool)
-            is_ligand[-len(ligand_atoms) :] = True
-        else:
-            is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool)
+
+        # Mate proteins inherit their source ASU residue's embedding row; ligands
+        # get -1 whichever cell they came from.
+        n_protein = final_protein_pos.size(0)
+        emb_res_idx = torch.cat([protein_res_idx, mate_emb_res_idx], dim=0)
+        is_mate = torch.zeros(n_protein, dtype=torch.bool)
+        is_mate[num_asu_protein:] = True
+
+        for coords, elements, from_mate in ligand_blocks:
+            n_lig = len(elements)
+            pos = torch.tensor(coords, dtype=torch.float32) - center
+            final_protein_pos = torch.cat([final_protein_pos, pos], dim=0)
+            final_protein_x = torch.cat(
+                [final_protein_x, element_onehot(elements)], dim=0
+            )
+            sentinel = torch.full((n_lig,), -1, dtype=torch.long)
+            final_protein_res_idx = torch.cat([final_protein_res_idx, sentinel], dim=0)
+            emb_res_idx = torch.cat([emb_res_idx, sentinel], dim=0)
+            is_mate = torch.cat(
+                [is_mate, torch.full((n_lig,), from_mate, dtype=torch.bool)], dim=0
+            )
+
+        is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool)
+        is_ligand[n_protein:] = True
 
         # Compute PP edges and features
         if final_protein_pos.size(0) > 0:
@@ -1219,6 +1521,8 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
                 "protein_x": final_protein_x,
                 "protein_res_idx": final_protein_res_idx,
                 "is_ligand": is_ligand,
+                "is_mate": is_mate,
+                "emb_res_idx": emb_res_idx,
                 "water_pos": water_pos,
                 "water_x": water_x,
                 # PP topology and features (precomputed)
@@ -1240,9 +1544,9 @@ def _annotate_data_with_embeddings(
         self,
         data: HeteroData,
         cache_key: str,
-        asu_protein_res_idx: torch.Tensor,
         num_asu_protein: int,
         num_protein_residues: int,
+        emb_res_idx: torch.Tensor,
     ) -> None:
         """
         Load encoder-specific embeddings and attach to data object.
@@ -1255,9 +1559,11 @@ def _annotate_data_with_embeddings(
         Args:
             data: HeteroData object to attach embeddings to (modified in-place)
             cache_key: Identifier for cached embedding files
-            asu_protein_res_idx: (N_asu,) residue index per ASU atom
             num_asu_protein: Number of ASU protein atoms
             num_protein_residues: Number of unique protein residues
+            emb_res_idx: (N_total,) embedding row per atom -- mates inherit their
+                source ASU residue's row; ligands and unmatched atoms are -1 and
+                get a zero row.
         """
         if self.encoder_type == "slae":
             data["protein"].embedding = load_slae_embedding(
@@ -1276,10 +1582,15 @@ def _annotate_data_with_embeddings(
                 num_protein_residues=num_protein_residues,
                 cache_load_mmap=self.cache_load_mmap,
             )
-            esm_atom_emb = residue_embeddings[asu_protein_res_idx]
-            data["protein"].embedding = _pad_atom_embeddings_for_mates(
-                esm_atom_emb, data["protein"].num_nodes
+            # Per-atom inheritance: a mate atom takes the row of the ASU residue
+            # it images; ligands and unmatched atoms (-1) stay zero.
+            atom_emb = residue_embeddings.new_zeros(
+                data["protein"].num_nodes, residue_embeddings.size(1)
             )
+            valid = emb_res_idx >= 0
+            if valid.any():
+                atom_emb[valid] = residue_embeddings[emb_res_idx[valid]]
+            data["protein"].embedding = atom_emb
             data["protein"].embedding_type = "esm"
 
     def __getitem__(self, idx: int) -> HeteroData:
@@ -1323,6 +1634,8 @@ def __getitem__(self, idx: int) -> HeteroData:
         protein_x = cached["protein_x"]
         protein_res_idx = cached["protein_res_idx"]
         is_ligand = cached["is_ligand"]
+        is_mate = cached["is_mate"]
+        emb_res_idx = cached["emb_res_idx"]
         pp_edge_index = cached["pp_edge_index"]
         pp_edge_unit_vectors = cached["pp_edge_unit_vectors"]
         pp_edge_rbf = cached["pp_edge_rbf"]
@@ -1331,9 +1644,6 @@ def __getitem__(self, idx: int) -> HeteroData:
         water_pos = cached["water_pos"]
         water_x = cached["water_x"]
 
-        # extract ASU protein residue indices for embedding loading
-        asu_protein_res_idx = protein_res_idx[:num_asu_protein]
-
         data = HeteroData()
 
         # compute total num_residues (protein + mates)
@@ -1345,6 +1655,7 @@ def __getitem__(self, idx: int) -> HeteroData:
         data["protein"].pos = protein_pos
         data["protein"].residue_index = protein_res_idx
         data["protein"].is_ligand = is_ligand
+        data["protein"].is_mate = is_mate
         data["protein"].num_nodes = protein_pos.size(0)
         data["protein"].num_residues = num_residues
         data["protein"].num_protein_residues = num_protein_residues
@@ -1352,9 +1663,9 @@ def __getitem__(self, idx: int) -> HeteroData:
         self._annotate_data_with_embeddings(
             data=data,
             cache_key=entry["embedding_key"],  # use base key for embeddings
-            asu_protein_res_idx=asu_protein_res_idx,
             num_asu_protein=num_asu_protein,
             num_protein_residues=num_protein_residues,
+            emb_res_idx=emb_res_idx,
         )
 
         data["water"].x = water_x
diff --git a/src/flow.py b/src/flow.py
index 6a4f151..a138378 100644
--- a/src/flow.py
+++ b/src/flow.py
@@ -57,6 +57,7 @@ def sample_waters_uniform_ball(
     batch_w: Tensor,
     cutoff: float = 8.0,
     device: torch.device | None = None,
+    anchor_mask: Tensor | None = None,
 ) -> Tensor:
     """
     Sample water positions uniformly inside balls of radius *cutoff* centred
@@ -73,6 +74,10 @@ def sample_waters_uniform_ball(
             batch vector and get samples aligned to it.
         cutoff: Ball radius in Angstroms
         device: Optional output device (defaults to protein_pos.device)
+        anchor_mask: Optional (N_protein,) bool selecting eligible anchors. Used to
+            anchor on ASU atoms only, so the prior spawns where the targets live
+            instead of dispersing onto symmetry mates that OT must then transport
+            back. Ignored if it would leave a water-requesting graph with no anchor.
 
     Returns:
         water_pos: (N_water, 3) sampled positions, one per entry of batch_w
@@ -97,6 +102,14 @@ def sample_waters_uniform_ball(
     if batch_p.numel() > 0:
         num_graphs = max(num_graphs, int(batch_p.max().item()) + 1)
 
+    # Drop to the eligible anchors
+    if anchor_mask is not None:
+        eligible = anchor_mask.to(device).bool()
+        counts = torch.bincount(batch_p[eligible], minlength=num_graphs)
+        if not (counts[batch_w] == 0).any():
+            protein_pos = protein_pos.to(device)[eligible]
+            batch_p = batch_p[eligible]
+
     # per-graph protein atom counts and cumulative offsets
     num_p_per_graph = scatter(
         torch.ones(batch_p.size(0), device=device, dtype=torch.long),
@@ -405,15 +418,14 @@ def forward(
         if EDGE_PP in data.edge_types:
             pp_edge = data[EDGE_PP]
 
-            # V_edge fallback is for backward compatibility with datasets
-            # that don't have cached edge features. A given model only sees one or the other.
+            # A given model sees one source or the other, never both.
             if pp_edge_attr is not None:
                 # Use encoder-learned scalar features (s_edge) with unit vectors
                 s_edge, V_edge = pp_edge_attr
                 if hasattr(pp_edge, "edge_unit_vectors"):
                     cached_edge_attr_dict[EDGE_PP] = (s_edge, pp_edge.edge_unit_vectors)
                 else:
-                    # Fallback for datasets without cached unit vectors
+                    # Graphs built outside the dataset carry vectors on the encoder side
                     cached_edge_attr_dict[EDGE_PP] = (s_edge, V_edge.squeeze(1))
             elif hasattr(pp_edge, "edge_rbf") and hasattr(pp_edge, "edge_unit_vectors"):
                 # No encoder edge features (e.g., SLAE/ESM) - use cached geometric features
@@ -716,12 +728,20 @@ def _sample_waters(
         """Dispatch to the configured sampling strategy, sampling one water per
         entry of batch_w and returning them in that order."""
         if self.sampling_strategy == "uniform_ball":
+            # With crystal mates in the graph, anchor on ASU atoms only: the
+            # targets are ASU-only, so mate anchors just disperse the prior.
+            # Runs without mates pass no mask and are unchanged.
+            is_mate = getattr(batch_data["protein"], "is_mate", None)
+            anchor_mask = (
+                ~is_mate.bool() if is_mate is not None and bool(is_mate.any()) else None
+            )
             return sample_waters_uniform_ball(
                 protein_pos=batch_data["protein"].pos,
                 batch_p=batch_data["protein"].batch,
                 batch_w=batch_w,
                 cutoff=self.graph_cutoff,
                 device=device,
+                anchor_mask=anchor_mask,
             )
         # scaled_gaussian
         sigma_per_graph = self.compute_sigma_per_graph(batch_data, device)
diff --git a/tests/conftest.py b/tests/conftest.py
index 7fa5117..9be3d92 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -83,7 +83,8 @@ def pdb_1deu():
 
 @pytest.fixture
 def pdb_4h0b():
-    """4h0b - has non-water ligand HETATMs for ligand support tests."""
+    """4h0b - has non-water ligand HETATMs for ligand support tests. P6 space group,
+    so a water on the 6-fold axis has a symmetry copy ~0A away (special position)."""
     return _resolve_test_path("4h0b", ".pdb")
 
 
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index b58f2f2..1120517 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -36,13 +36,17 @@
 from src.dataset import (
     _make_undirected,
     _pad_atom_embeddings_for_mates,
+    _parse_pdb_resi,
     apply_threshold_filter,
     check_chain_interactions,
     check_com_distance,
     check_water_clashes,
     check_water_residue_ratio,
     compute_normalized_bfactors,
+    dedup_mate_atoms,
+    dedup_mate_ligands_by_residue,
     element_onehot,
+    FILTER_META_FILENAME,
     filter_waters_by_quality,
     get_crystal_contacts_pymol,
     get_dataloader,
@@ -305,6 +309,172 @@ def test_warns_on_odd_counts_below_half(self, warning_log, n_atoms, n_matched):
         assert f"{n_matched}/{n_atoms} atoms matched" in warning_log[0]
 
 
+@pytest.mark.unit
+class TestDedupMateAtoms:
+    """Tests for symmetry-mate coordinate deduplication."""
+
+    @staticmethod
+    def _atoms(n):
+        return [object() for _ in range(n)]
+
+    def test_empty_passthrough(self):
+        coords = np.zeros((0, 3))
+        out_coords, out_atoms = dedup_mate_atoms(coords, [], np.zeros((0, 3)))
+
+        assert out_coords.shape == (0, 3)
+        assert out_atoms == []
+
+    def test_drops_atoms_coincident_with_reference(self):
+        """A mate atom sitting on an ASU/target atom is a leak and is removed."""
+        mate_coords = np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]])
+        reference = np.array([[0.0, 0.0, 0.0]])
+
+        out_coords, out_atoms = dedup_mate_atoms(
+            mate_coords, self._atoms(2), reference, tol=0.3
+        )
+
+        assert out_coords.shape[0] == 1
+        assert len(out_atoms) == 1
+        np.testing.assert_allclose(out_coords[0], [10.0, 0.0, 0.0])
+
+    def test_keeps_atoms_aligned(self):
+        """Returned coords and atom objects stay in lockstep."""
+        mate_coords = np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]])
+        first, second = object(), object()
+        reference = np.array([[0.0, 0.0, 0.0]])
+
+        out_coords, out_atoms = dedup_mate_atoms(
+            mate_coords, [first, second], reference, tol=0.3
+        )
+
+        assert out_atoms == [second]
+        np.testing.assert_allclose(out_coords[0], [10.0, 0.0, 0.0])
+
+    def test_keeps_atoms_beyond_tolerance(self):
+        """Separations at or past tol are distinct atoms, not duplicates."""
+        mate_coords = np.array([[0.0, 0.0, 0.0], [0.3, 0.0, 0.0], [0.6, 0.0, 0.0]])
+
+        out_coords, _ = dedup_mate_atoms(
+            mate_coords, self._atoms(3), np.zeros((0, 3)), tol=0.3
+        )
+
+        assert out_coords.shape[0] == 3
+
+    def test_self_dedup_is_first_wins(self):
+        """A chain of near-coincident mate atoms collapses onto the earliest."""
+        mate_coords = np.array([[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.2, 0.0, 0.0]])
+        first, second, third = object(), object(), object()
+
+        out_coords, out_atoms = dedup_mate_atoms(
+            mate_coords, [first, second, third], np.zeros((0, 3)), tol=0.3
+        )
+
+        assert out_atoms == [first]
+        np.testing.assert_allclose(out_coords[0], [0.0, 0.0, 0.0])
+
+
+class _FakeLigandAtom:
+    """Stand-in for a PyMOL atom object with the fields the dedup reads."""
+
+    def __init__(self, chain, resi, segi=""):
+        self.chain = chain
+        self.resi = resi
+        self.segi = segi
+
+
+@pytest.mark.unit
+class TestDedupMateLigandsByResidue:
+    """Tests for whole-entity symmetry-image ligand removal."""
+
+    def test_empty_passthrough(self):
+        coords = np.zeros((0, 3))
+        out_coords, out_atoms = dedup_mate_ligands_by_residue(
+            coords, [], np.zeros((0, 3))
+        )
+
+        assert out_coords.shape == (0, 3)
+        assert out_atoms == []
+
+    def test_drops_whole_symmetry_image_ligand(self):
+        """A ligand whose atoms mostly land on ASU atoms is dropped entirely."""
+        lig_coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]])
+        lig_atoms = [_FakeLigandAtom("A", "1") for _ in range(3)]
+        reference = lig_coords.copy()
+
+        out_coords, out_atoms = dedup_mate_ligands_by_residue(
+            lig_coords, lig_atoms, reference, tol=0.3
+        )
+
+        assert out_coords.shape[0] == 0
+        assert out_atoms == []
+
+    def test_keeps_neighbour_ligand_whole(self):
+        """A genuine neighbour-cell ligand keeps every atom, including any that
+        happen to coincide with the ASU."""
+        lig_coords = np.array([[0.0, 0.0, 0.0], [9.0, 0.0, 0.0], [10.0, 0.0, 0.0]])
+        lig_atoms = [_FakeLigandAtom("B", "7") for _ in range(3)]
+        reference = np.array([[0.0, 0.0, 0.0]])  # only one atom coincides
+
+        out_coords, out_atoms = dedup_mate_ligands_by_residue(
+            lig_coords, lig_atoms, reference, tol=0.3
+        )
+
+        assert out_coords.shape[0] == 3
+        assert len(out_atoms) == 3
+
+    def test_entities_are_judged_independently(self):
+        """One ligand being an image does not remove its neighbours."""
+        lig_coords = np.array([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]])
+        lig_atoms = [_FakeLigandAtom("A", "1"), _FakeLigandAtom("A", "2")]
+        reference = np.array([[0.0, 0.0, 0.0]])
+
+        out_coords, out_atoms = dedup_mate_ligands_by_residue(
+            lig_coords, lig_atoms, reference, tol=0.3
+        )
+
+        assert out_coords.shape[0] == 1
+        np.testing.assert_allclose(out_coords[0], [20.0, 0.0, 0.0])
+        assert out_atoms[0].resi == "2"
+
+    def test_segment_separates_entities(self):
+        """Two ligands sharing (chain, resi) but not segi stay independent."""
+        lig_coords = np.array([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]])
+        lig_atoms = [
+            _FakeLigandAtom("A", "1", segi="X"),
+            _FakeLigandAtom("A", "1", segi="Y"),
+        ]
+        reference = np.array([[0.0, 0.0, 0.0]])
+
+        out_coords, _ = dedup_mate_ligands_by_residue(
+            lig_coords, lig_atoms, reference, tol=0.3
+        )
+
+        assert out_coords.shape[0] == 1
+        np.testing.assert_allclose(out_coords[0], [20.0, 0.0, 0.0])
+
+
+@pytest.mark.unit
+class TestParsePdbResi:
+    """Tests for PyMOL residue-identifier parsing."""
+
+    @pytest.mark.parametrize(
+        "resi,expected",
+        [
+            ("52", (52, "")),
+            ("-3", (-3, "")),
+            ("52A", (52, "A")),
+            (" 52 ", (52, "")),
+            (7, (7, "")),
+        ],
+    )
+    def test_parses(self, resi, expected):
+        assert _parse_pdb_resi(resi) == expected
+
+    @pytest.mark.parametrize("resi", ["", "A", "52AB", "5.2"])
+    def test_returns_none_when_unparseable(self, resi):
+        assert _parse_pdb_resi(resi) is None
+
+
 @pytest.mark.unit
 class TestCheckComDistance:
     """Tests for center of mass distance quality filter."""
@@ -768,6 +938,13 @@ def test_is_ligand_mask_shape(self, pdb_4h0b, tmp_path):
         assert data["protein"].is_ligand.shape == (data["protein"].num_nodes,)
         assert data["protein"].is_ligand.dtype == torch.bool
 
+        # No residue owns a ligand, so it carries the -1 embedding sentinel.
+        cached = torch.load(
+            tmp_path / "processed" / "geometry" / "4h0b_final.pt", weights_only=False
+        )
+        assert cached["is_ligand"].any()
+        assert (cached["emb_res_idx"][cached["is_ligand"]] == -1).all()
+
     def test_protein_x_dim_unchanged(self, pdb_4h0b, tmp_path):
         """protein.x should still be 16-dim one-hot for both protein and ligand atoms."""
         list_file = tmp_path / "list.txt"
@@ -843,6 +1020,45 @@ def test_different_cutoffs(self, pdb_6eey):
             result_large["mate_coords"].shape[0] >= result_small["mate_coords"].shape[0]
         )
 
+    def test_mates_never_include_solvent(self, pdb_8dzt):
+        """No mate atom, protein or ligand, is a water: a mate water is a symmetry
+        image of an ASU water, which is a prediction target."""
+        result = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0, include_ligands=True)
+
+        protein_resns = {str(a.resn).upper() for a in result["mate_atoms"]}
+        ligand_resns = {str(a.resn).upper() for a in result["mate_ligand_atoms"]}
+        assert not {"HOH", "WAT"} & protein_resns
+        assert not {"HOH", "WAT"} & ligand_resns
+
+    def test_ligand_mates_are_gated_and_separate(self, pdb_8dzt):
+        """include_ligands=False suppresses ligand mates and leaves protein mates
+        alone -- the two sets come back under separate keys."""
+        full = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0, include_ligands=True)
+        protein_only = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0)
+
+        assert len(protein_only["mate_ligand_atoms"]) == 0
+        assert protein_only["mate_ligand_coords"].shape[0] == 0
+        assert len(full["mate_ligand_atoms"]) > 0
+        assert protein_only["mate_coords"].shape[0] == full["mate_coords"].shape[0]
+
+    def test_special_position_water_never_selected(self, pdb_4h0b):
+        """4h0b has a target water on the 6-fold axis whose symmetry copy lands
+        ~0 A away. Since mate waters are never selected, no mate of any kind may
+        coincide with a target water -- the special-position leak cannot happen."""
+        from scipy.spatial import cKDTree
+
+        _, water_atoms, _ = parse_asu_with_biotite(pdb_4h0b)
+        result = get_crystal_contacts_pymol(pdb_4h0b, cutoff=5.0, include_ligands=True)
+
+        mate_coords = result["mate_coords"]
+        if result["mate_ligand_coords"].shape[0]:
+            mate_coords = np.concatenate(
+                [mate_coords, result["mate_ligand_coords"]], axis=0
+            )
+        assert mate_coords.shape[0] > 0
+        nearest = cKDTree(mate_coords).query(water_atoms.coord, k=1)[0]
+        assert (nearest < 0.3).sum() == 0
+
 
 @pytest.mark.integration
 class TestProteinWaterDataset:
@@ -952,6 +1168,8 @@ def test_getitem_passes_mmap_flag_to_geometry_loader(self, tmp_path, monkeypatch
             "protein_x": torch.zeros((1, len(ELEMENT_VOCAB) + 1), dtype=torch.float32),
             "protein_res_idx": torch.zeros(1, dtype=torch.long),
             "is_ligand": torch.zeros(1, dtype=torch.bool),
+            "is_mate": torch.zeros(1, dtype=torch.bool),
+            "emb_res_idx": torch.zeros(1, dtype=torch.long),
             "pp_edge_index": torch.empty((2, 0), dtype=torch.long),
             "pp_edge_unit_vectors": torch.empty((0, 3), dtype=torch.float32),
             "pp_edge_rbf": torch.empty((0, 16), dtype=torch.float32),
@@ -2138,9 +2356,9 @@ def test_gvp_encoder_no_embeddings(self, tmp_path, pdb_base_dir):
         dataset._annotate_data_with_embeddings(
             data=data,
             cache_key="test",
-            asu_protein_res_idx=torch.tensor([0]),
             num_asu_protein=100,
             num_protein_residues=50,
+            emb_res_idx=torch.zeros(100, dtype=torch.long),
         )
 
         # Should not have added any embedding attributes
@@ -2176,9 +2394,9 @@ def test_slae_encoder_loads_slae(self, tmp_path, pdb_base_dir):
         dataset._annotate_data_with_embeddings(
             data=data,
             cache_key="test_final",
-            asu_protein_res_idx=torch.tensor([0]),
             num_asu_protein=100,
             num_protein_residues=50,
+            emb_res_idx=torch.zeros(100, dtype=torch.long),
         )
 
         assert hasattr(data["protein"], "embedding")
@@ -2217,9 +2435,9 @@ def test_esm_encoder_loads_esm(self, tmp_path, pdb_base_dir):
         dataset._annotate_data_with_embeddings(
             data=data,
             cache_key="test_final",
-            asu_protein_res_idx=asu_res_idx,
             num_asu_protein=50,
             num_protein_residues=10,
+            emb_res_idx=asu_res_idx,
         )
 
         assert hasattr(data["protein"], "embedding")
@@ -2254,9 +2472,9 @@ def test_slae_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir):
         dataset._annotate_data_with_embeddings(
             data=data,
             cache_key="test_final",
-            asu_protein_res_idx=torch.zeros(num_asu, dtype=torch.long),
             num_asu_protein=num_asu,
             num_protein_residues=1,
+            emb_res_idx=torch.zeros(num_asu + num_mate + num_ligand, dtype=torch.long),
         )
 
         emb = data["protein"].embedding
@@ -2264,10 +2482,10 @@ def test_slae_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir):
         assert torch.equal(emb[:num_asu], asu_emb), "ASU rows must be left untouched"
         assert (emb[num_asu:] == 0).all(), "mate and ligand rows must be zero-padded"
 
-    def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir):
-        """ESM residue embeddings broadcast to ASU atoms only. Mate and ligand atoms
-        are zero-padded -- ligands carry residue_index=-1 and must never be used to
-        index the residue embedding table."""
+    def test_esm_mates_inherit_and_ligands_zero(self, tmp_path, pdb_base_dir):
+        """ESM rows broadcast to ASU atoms, and mate atoms inherit the row of the
+        ASU residue they image. Ligands carry -1 and must never index the residue
+        table, so they stay zero."""
         from torch_geometric.data import HeteroData
 
         num_residues = 4
@@ -2293,16 +2511,19 @@ def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir):
         data = HeteroData()
         data["protein"].num_nodes = num_asu + num_mate + num_ligand
 
-        # ASU res idx only -- ligand sentinels (-1) live past num_asu_protein and are
-        # sliced off by __getitem__ before this call.
         asu_res_idx = torch.arange(num_residues).repeat_interleave(atoms_per_residue)
+        # Mates image residue 0; ligands get the -1 sentinel.
+        mate_res_idx = torch.zeros(num_mate, dtype=torch.long)
+        emb_res_idx = torch.cat(
+            [asu_res_idx, mate_res_idx, torch.full((num_ligand,), -1)]
+        )
 
         dataset._annotate_data_with_embeddings(
             data=data,
             cache_key="test_final",
-            asu_protein_res_idx=asu_res_idx,
             num_asu_protein=num_asu,
             num_protein_residues=num_residues,
+            emb_res_idx=emb_res_idx,
         )
 
         emb = data["protein"].embedding
@@ -2310,7 +2531,10 @@ def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir):
         assert torch.equal(emb[:num_asu], residue_emb[asu_res_idx]), (
             "each ASU atom must carry its own residue's embedding"
         )
-        assert (emb[num_asu:] == 0).all(), "mate and ligand rows must be zero-padded"
+        assert torch.equal(
+            emb[num_asu : num_asu + num_mate], residue_emb[mate_res_idx]
+        ), "mate atoms must inherit their source residue's embedding"
+        assert (emb[num_asu + num_mate :] == 0).all(), "ligand rows must stay zero"
 
 
 # ============== Tests for caching behavior ==============
@@ -2665,6 +2889,156 @@ def test_num_asu_protein_metadata_correct(
         assert data.num_asu_protein_atoms <= data["protein"].num_nodes
         assert data.num_asu_protein_atoms > 0
 
+    def _mate_dataset(self, single_pdb_list_file, tmp_path, pdb_base_dir, **kwargs):
+        return ProteinWaterDataset(
+            pdb_list_file=single_pdb_list_file,
+            processed_dir=str(tmp_path),
+            base_pdb_dir=str(pdb_base_dir),
+            include_mates=True,
+            preprocess=True,
+            **kwargs,
+        )
+
+    def test_mate_provenance_fields(self, single_pdb_list_file, tmp_path, pdb_base_dir):
+        """is_mate splits ASU from mate at num_asu_protein, and the mate atoms
+        behind it carry the embedding row of the ASU residue they image, not the
+        -1 that reads as a zero row."""
+        dataset = self._mate_dataset(single_pdb_list_file, tmp_path, pdb_base_dir)
+        data = dataset[0]
+        cached = torch.load(
+            tmp_path / "geometry_mates" / "6eey_final.pt", weights_only=False
+        )
+        is_mate, emb_res_idx = data["protein"].is_mate, cached["emb_res_idx"]
+        num_asu = data.num_asu_protein_atoms
+
+        assert is_mate.shape == emb_res_idx.shape == (data["protein"].num_nodes,)
+        assert not is_mate[:num_asu].any()
+        assert is_mate.sum().item() > 0
+        # Ligands ride behind the mates, so the mate block is contiguous but need
+        # not run to the end; every marked node is past the ASU.
+        assert is_mate[num_asu:][: is_mate.sum().item()].all()
+
+        mate_protein = is_mate & ~cached["is_ligand"]
+        assert (emb_res_idx[mate_protein] >= 0).all()
+        assert (emb_res_idx[mate_protein] < data["protein"].num_protein_residues).all()
+
+    def test_no_mates_run_marks_nothing_as_mate(
+        self, single_pdb_list_file, tmp_path, pdb_base_dir
+    ):
+        """Without mates every node is ASU, so the prior's anchor mask is inert."""
+        dataset = ProteinWaterDataset(
+            pdb_list_file=single_pdb_list_file,
+            processed_dir=str(tmp_path),
+            base_pdb_dir=str(pdb_base_dir),
+            include_mates=False,
+            preprocess=True,
+        )
+
+        data = dataset[0]
+        assert not data["protein"].is_mate.any()
+
+    def test_mate_waters_never_enter_the_graph(
+        self, single_pdb_list_file, tmp_path, pdb_base_dir
+    ):
+        """No protein node may sit on a target water: that is the label leak the
+        mate selection and the dedup pass exist to prevent."""
+        dataset = self._mate_dataset(single_pdb_list_file, tmp_path, pdb_base_dir)
+        data = dataset[0]
+
+        waters = data["water"].pos
+        if waters.size(0) == 0:
+            pytest.skip("structure has no waters after filtering")
+        nearest = torch.cdist(waters, data["protein"].pos).min(dim=1).values
+        assert nearest.min().item() > 0.3
+
+
+@pytest.mark.unit
+class TestFilterMetaSidecar:
+    """A geometry directory records the settings its entries were built with."""
+
+    def _dataset(self, tmp_path, *, preprocess=True, **kwargs):
+        """Dataset over an empty list: claims the directory, preprocesses nothing."""
+        list_file = tmp_path / "empty.txt"
+        list_file.write_text("")
+        return ProteinWaterDataset(
+            pdb_list_file=str(list_file),
+            processed_dir=str(tmp_path / "processed"),
+            base_pdb_dir=str(tmp_path),
+            preprocess=preprocess,
+            **kwargs,
+        )
+
+    def _meta_path(self, tmp_path):
+        return tmp_path / "processed" / "geometry_mates" / FILTER_META_FILENAME
+
+    def test_written_on_preprocess(self, tmp_path):
+        self._dataset(tmp_path, max_bfactor_zscore=2.0)
+
+        recorded = json.loads(self._meta_path(tmp_path).read_text())
+        assert recorded["max_bfactor_zscore"] == 2.0
+        assert recorded["min_edia"] == 0.4
+        assert recorded["filter_by_bfactor"] is True
+        # Structure-level checks decide which entries exist at all, and the graph
+        # parameters decide the cached edges: both belong to the directory too.
+        assert recorded["min_water_residue_ratio"] == 0.1
+        assert recorded["cutoff"] == 8.0
+        assert recorded["max_neighbors"] == 256
+
+    def test_matching_settings_accepted(self, tmp_path):
+        self._dataset(tmp_path, max_bfactor_zscore=2.0)
+        self._dataset(tmp_path, max_bfactor_zscore=2.0)  # must not raise
+
+    @pytest.mark.parametrize(
+        "changed,preprocess",
+        [
+            ({"max_bfactor_zscore": 1.5}, True),  # water threshold
+            ({"filter_by_edia": False}, True),  # water filter toggle
+            ({"min_water_residue_ratio": 0.6}, True),  # which entries exist
+            ({"cutoff": 6.0}, True),  # which PP edges were cached
+            # A read-only run is refused too: it would report metrics over waters
+            # filtered differently than it asked for.
+            ({"min_edia": 0.6}, False),
+        ],
+    )
+    def test_mismatch_refused(self, tmp_path, changed, preprocess):
+        self._dataset(tmp_path)
+
+        with pytest.raises(ValueError, match=next(iter(changed))):
+            self._dataset(tmp_path, preprocess=preprocess, **changed)
+
+    def test_disabled_filter_ignores_its_threshold(self, tmp_path):
+        """A disabled filter never touched the cached waters, so its threshold
+        must not make two identical caches look incompatible."""
+        self._dataset(tmp_path, filter_by_bfactor=False, max_bfactor_zscore=2.0)
+
+        assert (
+            json.loads(self._meta_path(tmp_path).read_text())["max_bfactor_zscore"]
+            is None
+        )
+        self._dataset(tmp_path, filter_by_bfactor=False, max_bfactor_zscore=1.5)
+
+    def test_directories_are_claimed_independently(self, tmp_path):
+        """Mates and no-mates are separate directories, so they may disagree."""
+        self._dataset(tmp_path, include_mates=True, max_bfactor_zscore=2.0)
+        self._dataset(tmp_path, include_mates=False, max_bfactor_zscore=1.5)
+
+    def test_unlabelled_cache_warns(self, tmp_path, warning_log):
+        """An existing directory with no sidecar is usable but unverifiable."""
+        geometry_dir = tmp_path / "processed" / "geometry_mates"
+        geometry_dir.mkdir(parents=True)
+        (geometry_dir / "6eey_final.pt").write_bytes(b"cache")
+
+        self._dataset(tmp_path, preprocess=False)
+
+        assert any(FILTER_META_FILENAME in message for message in warning_log)
+
+    def test_empty_directory_does_not_warn(self, tmp_path, warning_log):
+        (tmp_path / "processed" / "geometry_mates").mkdir(parents=True)
+
+        self._dataset(tmp_path, preprocess=False)
+
+        assert not any(FILTER_META_FILENAME in message for message in warning_log)
+
 
 # ============== Tests for residue index assignment ==============
 
diff --git a/tests/test_flow.py b/tests/test_flow.py
index 9d2f62a..1ae9b39 100644
--- a/tests/test_flow.py
+++ b/tests/test_flow.py
@@ -633,6 +633,83 @@ def test_zero_protein_graph_raises(self, device):
                 device=device,
             )
 
+    def test_anchor_mask_is_per_graph(self, device):
+        """Masked-out atoms are never ball centres, and each graph stays within
+        its own eligible atoms rather than its neighbour's."""
+        torch.manual_seed(0)
+        protein_pos = torch.tensor(
+            [[0.0, 0.0, 0.0], [500.0, 0.0, 0.0], [100.0, 0.0, 0.0], [900.0, 0.0, 0.0]],
+            device=device,
+        )
+        batch_p = torch.tensor([0, 0, 1, 1], dtype=torch.long, device=device)
+        batch_w = _batch_from_counts(
+            torch.tensor([50, 50], dtype=torch.long, device=device), device
+        )
+        anchor_mask = torch.tensor([True, False, True, False], device=device)
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=2.0,
+            device=device,
+            anchor_mask=anchor_mask,
+        )
+
+        assert pos[batch_w == 0][:, 0].abs().max().item() < 5.0
+        assert (pos[batch_w == 1][:, 0] - 100.0).abs().max().item() < 5.0
+
+    def test_anchor_mask_ignored_when_it_starves_a_graph(self, device):
+        """A mask leaving a water-requesting graph with no anchor is dropped, not
+        allowed to shift that graph's waters onto another graph's atoms."""
+        torch.manual_seed(0)
+        protein_pos = torch.tensor([[0.0, 0.0, 0.0], [1000.0, 0.0, 0.0]], device=device)
+        batch_p = torch.tensor([0, 1], dtype=torch.long, device=device)
+        batch_w = _batch_from_counts(
+            torch.tensor([10, 10], dtype=torch.long, device=device), device
+        )
+        # graph 1 has no eligible atom
+        anchor_mask = torch.tensor([True, False], device=device)
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=2.0,
+            device=device,
+            anchor_mask=anchor_mask,
+        )
+
+        assert (pos[batch_w == 1][:, 0] - 1000.0).abs().max().item() < 5.0
+
+    def test_anchor_mask_none_matches_all_true_mask(self, device):
+        """An all-True mask changes neither the draws nor their order."""
+        protein_pos = torch.randn(12, 3, device=device) * 10
+        batch_p = torch.cat([torch.zeros(6), torch.ones(6)]).long().to(device)
+        batch_w = _batch_from_counts(
+            torch.tensor([20, 15], dtype=torch.long, device=device), device
+        )
+
+        torch.manual_seed(7)
+        without = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=8.0,
+            device=device,
+        )
+        torch.manual_seed(7)
+        with_mask = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=8.0,
+            device=device,
+            anchor_mask=torch.ones(12, dtype=torch.bool, device=device),
+        )
+
+        assert torch.equal(without, with_mask)
+
     def test_large_spread_protein_succeeds(self, device):
         """The scenario that crashes truncated Gaussian (sigma~50) works here."""
         torch.manual_seed(0)