diff --git a/README.md b/README.md index 9ecad5c..2604349 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ WaterFlow/ ├── scripts/ # Executable scripts │ ├── train.py # Training pipeline │ ├── inference.py # Run inference on trained models +│ ├── cache_candidates.py # Sample candidate waters for confidence training │ ├── generate_esm_embeddings.py # Precompute ESM embeddings │ └── generate_slae_embeddings.py # Precompute SLAE embeddings ├── tests/ # Test suite @@ -172,7 +173,7 @@ The base name comes from `--geometry_cache_name` (default `geometry`). 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 +directory therefore carries a `_filter_meta.json` file 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`, @@ -181,7 +182,7 @@ their toggles, the structure-level checks that decide which entries exist at all 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 +the cached waters. Directories built before this existed have no such file: 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. diff --git a/scripts/cache_candidates.py b/scripts/cache_candidates.py new file mode 100644 index 0000000..808902b --- /dev/null +++ b/scripts/cache_candidates.py @@ -0,0 +1,358 @@ +""" +Generate the candidate cache for confidence-model training. + +Samples candidate waters from a trained flow checkpoint over the flow dataset +cache layout and writes one `.pt = {"candidate_pos": (Nc, 3)}` per +structure, plus a `generation.json` record of how it was made. Train the confidence +model on the result with `scripts/train_confidence.py --candidate_dir `. + +Everything the confidence model needs (protein graph, embeddings, GT waters, PP edges) +is loaded from the flow caches at train time. Model loading and integration reuse the flow inference +machinery verbatim, so candidates are sampled exactly as `scripts/inference.py` +would sample them. + +Example: + python -m scripts.cache_candidates \\ + --flow_run_dir \\ + --pdb_list splits/conf_train.txt \\ + --processed_dir \\ + --base_pdb_dir \\ + --water_ratio 3.0 --seed 0 --method euler --num_steps 100 +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from loguru import logger +from tqdm import tqdm + +from scripts.inference import ( + _extract_dataset_filter_config, + build_model_from_config, + load_checkpoint, + load_config, + run_inference_batch, +) +from src.dataset import ProteinWaterDataset +from src.flow import FlowMatcher +from src.utils import setup_logging_for_tqdm + + +def default_candidate_dir( + processed_dir: str | Path, + flow_run_dir: str | Path, + water_ratio: float, + seed: int, +) -> Path: + """Namespaced candidate dir so caches from different checkpoints, ratios, or + seeds never collide: `{processed_dir}/candidate_cache/_r{ratio}_s{seed}`. + """ + run_name = Path(flow_run_dir).name + return ( + Path(processed_dir) / "candidate_cache" / f"{run_name}_r{water_ratio:g}_s{seed}" + ) + + +def _write_candidate_cache( + dataset, + sample_batch, + out: Path, + *, + run_info: dict, + batch_size: int = 8, + overwrite: bool = False, +) -> dict: + """Sample and write one thin candidate file per uncached structure. + + Split from `generate_candidate_cache` so the loop and skip logic runs against a + plain dataset and sampler instead of a live flow checkpoint. + + Args: + dataset: Indexable structures. `dataset.entries[i]["cache_key"]` names the + output file, which equals the sampled graph's `pdb_id` and the key + `ConfidenceDataset` reads back. + sample_batch: Callable `(graphs) -> [{"pdb_id", "water_pred"}, ...]`. + out: Directory for the `.pt` files and `generation.json`. + run_info: Generation parameters recorded, with the counts, in + `generation.json`. + batch_size: Structures sampled per `sample_batch` call. + overwrite: Re-generate even if a `.pt` already exists. + + Returns: + dict stats: {"out_dir", "n_written", "n_skipped", "n_total"}. + """ + out.mkdir(parents=True, exist_ok=True) + entries = dataset.entries + n_total = len(dataset) + n_written = 0 + n_skipped = 0 + for start in tqdm(range(0, n_total, batch_size), desc="candidate-gen"): + chunk = range(start, min(start + batch_size, n_total)) + # Decide what to (re)generate from the cache key alone, so an already + # cached structure never pays for a graph build. + todo = [ + i + for i in chunk + if overwrite or not (out / f"{entries[i]['cache_key']}.pt").exists() + ] + n_skipped += len(chunk) - len(todo) + if not todo: + continue + for result in sample_batch([dataset[i] for i in todo]): + candidate_pos = torch.as_tensor(result["water_pred"], dtype=torch.float32) + torch.save({"candidate_pos": candidate_pos}, out / f"{result['pdb_id']}.pt") + n_written += 1 + + (out / "generation.json").write_text( + json.dumps( + { + **run_info, + "n_written": n_written, + "n_skipped": n_skipped, + "n_total": n_total, + }, + indent=2, + ) + ) + + logger.info( + f"Candidate cache written to {out}: {n_written} written, " + f"{n_skipped} skipped (already cached)." + ) + return { + "out_dir": str(out), + "n_written": n_written, + "n_skipped": n_skipped, + "n_total": n_total, + } + + +def generate_candidate_cache( + flow_run_dir: str | Path, + pdb_list: str | Path, + processed_dir: str | Path, + base_pdb_dir: str | Path, + out_dir: str | Path | None = None, + *, + checkpoint: str = "best.pt", + water_ratio: float = 3.0, + seed: int = 0, + num_steps: int = 100, + method: str = "euler", + batch_size: int = 8, + geometry_cache_name: str | None = None, + include_mates: bool | None = None, + device: str = "cuda", + overwrite: bool = False, +) -> dict: + """ + Sample candidate waters from a trained flow checkpoint into candidate files. + + Args: + flow_run_dir: Flow training run dir (contains config.json + checkpoints/). + pdb_list: Text file of `_final` keys, one per line. + processed_dir: Cache root shared with flow training (geometry + esm). + base_pdb_dir: Base PDB dir, as used by flow training. + out_dir: Output dir. Defaults to `default_candidate_dir(...)`. + checkpoint: Checkpoint filename under `{flow_run_dir}/checkpoints`. + water_ratio: Oversampling ratio — sample `num_residues * water_ratio` waters. + seed: RNG seed for the sampling prior (reproducible candidates). + num_steps, method: Integration settings, as in inference. + batch_size: Graphs per integration batch. + geometry_cache_name / include_mates: Optional overrides; default to the + flow config's values so the graph matches what the flow model saw. + device: 'cuda' or 'cpu'. + overwrite: Re-generate even if a `.pt` already exists. + + Returns: + dict stats: {"out_dir", "n_written", "n_skipped", "n_total"}. + """ + run_dir = Path(flow_run_dir) + config = load_config(run_dir) + device_t = torch.device(device if torch.cuda.is_available() else "cpu") + + # Frozen flow model, loaded exactly as scripts/inference.py loads it. + model = build_model_from_config(config, device_t) + checkpoint_path = run_dir / "checkpoints" / checkpoint + epoch = load_checkpoint(model, checkpoint_path, device_t) + logger.info(f"Loaded flow checkpoint {checkpoint_path} (epoch {epoch})") + + flow_matcher = FlowMatcher( + model=model, + sampling_strategy=config.get("sampling_strategy", "uniform_ball"), + ) + + if include_mates is None: + include_mates = config.get("include_mates", False) + if geometry_cache_name is None: + geometry_cache_name = config.get("geometry_cache_name", "geometry") + encoder_type = config.get("encoder_type", "gvp") + + dataset = ProteinWaterDataset( + pdb_list_file=str(pdb_list), + processed_dir=str(processed_dir), + base_pdb_dir=str(base_pdb_dir), + encoder_type=encoder_type, + include_mates=include_mates, + # Also picks the cache directory, so it has to track the flow run. + include_ligands=config.get("include_ligands", True), + geometry_cache_name=geometry_cache_name, + preprocess=True, + **_extract_dataset_filter_config(config), + ) + logger.info( + f"Generating candidates for {len(dataset)} structures " + f"(encoder={encoder_type}, geometry={geometry_cache_name}, " + f"mates={include_mates}, ratio={water_ratio}, seed={seed})" + ) + + out = ( + Path(out_dir) + if out_dir is not None + else default_candidate_dir(processed_dir, run_dir, water_ratio, seed) + ) + torch.manual_seed(seed) # reproducible sampling prior + + def sample_batch(graphs): + return run_inference_batch( + flow_matcher, + graphs, + method=method, + num_steps=num_steps, + device=str(device_t), + water_ratio=water_ratio, + ) + + run_info = { + "flow_run_dir": str(run_dir), + "pdb_list": str(pdb_list), + "checkpoint": checkpoint, + "epoch": epoch, + "water_ratio": water_ratio, + "seed": seed, + "num_steps": num_steps, + "method": method, + "encoder_type": encoder_type, + "geometry_cache_name": geometry_cache_name, + "include_mates": include_mates, + } + return _write_candidate_cache( + dataset, + sample_batch, + out, + run_info=run_info, + batch_size=batch_size, + overwrite=overwrite, + ) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Generate the candidate cache for confidence-model training." + ) + p.add_argument( + "--flow_run_dir", + required=True, + help="Flow training run dir (config.json + checkpoints/).", + ) + p.add_argument( + "--pdb_list", + required=True, + help="Text file of _final keys, one per line.", + ) + p.add_argument( + "--processed_dir", + required=True, + help="Cache root shared with flow training (geometry + esm).", + ) + p.add_argument( + "--base_pdb_dir", + required=True, + help="Base PDB dir, as used by flow training.", + ) + p.add_argument( + "--out_dir", + default=None, + help="Output dir. Default: " + "{processed_dir}/candidate_cache/_r{ratio}_s{seed}.", + ) + p.add_argument( + "--checkpoint", + default="best.pt", + help="Checkpoint filename under {flow_run_dir}/checkpoints.", + ) + p.add_argument( + "--water_ratio", + type=float, + default=3.0, + help="Oversampling: num_residues * water_ratio waters per structure.", + ) + p.add_argument( + "--seed", + type=int, + default=0, + help="RNG seed for the sampling prior (reproducible candidates).", + ) + p.add_argument("--num_steps", type=int, default=100, help="Integration steps.") + p.add_argument( + "--method", + choices=["euler", "rk4"], + default="euler", + help="Integration method.", + ) + p.add_argument( + "--batch_size", type=int, default=8, help="Graphs per integration batch." + ) + p.add_argument( + "--geometry_cache_name", + default=None, + help="Override geometry cache base name (default: flow config).", + ) + p.add_argument( + "--include_mates", + action="store_true", + default=None, + help="Force-include symmetry mates (default: flow config).", + ) + p.add_argument("--device", default="cuda") + p.add_argument( + "--overwrite", + action="store_true", + help="Re-generate even if a .pt already exists.", + ) + p.add_argument("--log_level", default="INFO") + return p.parse_args() + + +def main() -> None: + args = parse_args() + setup_logging_for_tqdm(level=args.log_level) + stats = generate_candidate_cache( + flow_run_dir=args.flow_run_dir, + pdb_list=args.pdb_list, + processed_dir=args.processed_dir, + base_pdb_dir=args.base_pdb_dir, + out_dir=args.out_dir, + checkpoint=args.checkpoint, + water_ratio=args.water_ratio, + seed=args.seed, + num_steps=args.num_steps, + method=args.method, + batch_size=args.batch_size, + geometry_cache_name=args.geometry_cache_name, + include_mates=args.include_mates, + device=args.device, + overwrite=args.overwrite, + ) + print( + f"Done: {stats['n_written']} written, {stats['n_skipped']} skipped " + f"-> {stats['out_dir']}" + ) + + +if __name__ == "__main__": + main() diff --git a/src/dataset.py b/src/dataset.py index 5928314..567e5d4 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -1142,7 +1142,7 @@ def _sync_filter_meta(self, write: bool) -> None: 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. + # and a reader must never catch a half-written file. tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp") with open(tmp_path, "w") as f: json.dump(current, f, indent=2) diff --git a/tests/test_cache_candidates.py b/tests/test_cache_candidates.py new file mode 100644 index 0000000..b148029 --- /dev/null +++ b/tests/test_cache_candidates.py @@ -0,0 +1,115 @@ +"""Unit tests for scripts/cache_candidates.py -- candidate cache generation.""" + +import json +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import scripts.cache_candidates as cc + + +class _FakeDataset: + """Flow-dataset stand-in: `.entries` with cache keys and indexable graphs whose + `pdb_id` matches, exactly as ProteinWaterDataset exposes them.""" + + def __init__(self, keys): + self.entries = [{"cache_key": k} for k in keys] + self._graphs = [SimpleNamespace(pdb_id=k) for k in keys] + + def __len__(self): + return len(self._graphs) + + def __getitem__(self, idx): + return self._graphs[idx] + + +def _sampler(n_cand): + """A `sample_batch` stand-in: a fixed-size candidate set per structure.""" + + def sample_batch(graphs): + return [ + {"pdb_id": g.pdb_id, "water_pred": np.zeros((n_cand, 3), dtype=np.float32)} + for g in graphs + ] + + return sample_batch + + +# Static run-info params; their exact values are opaque to _write_candidate_cache. +_STATIC = {"flow_run_dir": "run", "checkpoint": "best.pt", "epoch": 7} + + +@pytest.mark.unit +class TestWriteCandidateCache: + def test_writes_thin_candidate_files(self, tmp_path): + out = tmp_path / "cand" + stats = cc._write_candidate_cache( + _FakeDataset(["a_final", "b_final"]), + _sampler(4), + out, + run_info=_STATIC, + batch_size=1, + ) + + assert stats["n_written"] == 2 + assert (out / "a_final.pt").exists() and (out / "b_final.pt").exists() + payload = torch.load(out / "a_final.pt", weights_only=True) + assert payload["candidate_pos"].shape == (4, 3) + assert payload["candidate_pos"].dtype == torch.float32 + # generation.json records run_info plus this run's counts + info = json.loads((out / "generation.json").read_text()) + assert info["n_written"] == 2 and info["n_total"] == 2 + assert info["checkpoint"] == "best.pt" + + def test_skips_existing_without_overwrite(self, tmp_path): + out = tmp_path / "cand" + out.mkdir() + torch.save({"candidate_pos": torch.ones(9, 3)}, out / "a_final.pt") + + stats = cc._write_candidate_cache( + _FakeDataset(["a_final"]), + _sampler(2), + out, + run_info=_STATIC, + ) + + assert stats["n_written"] == 0 and stats["n_skipped"] == 1 + untouched = torch.load(out / "a_final.pt", weights_only=True) + assert untouched["candidate_pos"].shape == (9, 3) + + def test_overwrite_regenerates(self, tmp_path): + out = tmp_path / "cand" + out.mkdir() + torch.save({"candidate_pos": torch.ones(9, 3)}, out / "a_final.pt") + + stats = cc._write_candidate_cache( + _FakeDataset(["a_final"]), + _sampler(2), + out, + run_info=_STATIC, + overwrite=True, + ) + + assert stats["n_written"] == 1 + regenerated = torch.load(out / "a_final.pt", weights_only=True) + assert regenerated["candidate_pos"].shape == (2, 3) + + +@pytest.mark.unit +class TestDefaultCandidateDir: + def test_namespaced_by_run_ratio_and_seed(self, tmp_path): + out = cc.default_candidate_dir(tmp_path, "/runs/my_run", 3.0, 1) + + assert out == tmp_path / "candidate_cache" / "my_run_r3_s1" + + def test_distinct_configs_never_collide(self, tmp_path): + dirs = { + cc.default_candidate_dir(tmp_path, run, ratio, seed) + for run in ("/runs/a", "/runs/b") + for ratio in (2.0, 3.0) + for seed in (0, 1) + } + + assert len(dirs) == 8 diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 13fa53b..46dc88f 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -3168,7 +3168,7 @@ def _dataset(self, tmp_path, pdb_base_dir, **kwargs): include_mates=True, include_ligands=True, preprocess=True, - # 4h0b ships no EDIA sidecar; this class is about the node layout + # 4h0b ships no EDIA file; this class is about the node layout filter_by_distance=False, filter_by_edia=False, filter_by_bfactor=False, @@ -3281,7 +3281,7 @@ def test_esm_encoder_over_the_real_preprocessing_path(self, tmp_path, pdb_base_d @pytest.mark.unit -class TestFilterMetaSidecar: +class TestFilterMetaFile: """A geometry directory records the settings its entries were built with.""" def _dataset(self, tmp_path, *, preprocess=True, **kwargs): @@ -3352,7 +3352,7 @@ def test_directories_are_claimed_independently(self, tmp_path): @pytest.mark.parametrize("preprocess", [False, True]) def test_unlabelled_cache_warns(self, tmp_path, warning_log, preprocess): - """An existing directory with no sidecar is unverifiable whether or not this + """An existing directory with no such file is unverifiable whether or not this run also claims it.""" geometry_dir = tmp_path / "processed" / "geometry_mates" geometry_dir.mkdir(parents=True)