Confidence Model - #90
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Adds a new “confidence model” subsystem to WaterFlow to score generated candidate waters, generate soft supervision targets from GT distances, and post-process candidates via SuperWater-style vdW clustering/NMS. This complements the existing flow generator by enabling a two-stage pipeline (generate → score/cluster) while keeping the scorer backbone warm-start compatible with flow checkpoints.
Changes:
- Introduces
src/confidence.pywith smootherstep-based targets, vdW clustering, and aConfidenceGVPscorer built onProteinWaterUpdatewith PW+PP edges. - Adds comprehensive unit tests covering target behavior, clustering invariants, and cached-vs-dynamic edge behavior for
ConfidenceGVP. - Updates README tree to document the new module and tests.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/confidence.py |
Implements smootherstep target/confidence, vdW clustering, and ConfidenceGVP scorer using the existing GVP backbone components. |
tests/test_confidence.py |
Adds unit tests for targets, clustering, ConfidenceGVP forward behavior, gradients, and cached edge usage. |
README.md |
Documents the new confidence module and its test file in the repository layout. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| s_all, v_all, pp_edge_attr = self.encoder(data) | ||
| encoder_input = (s_all, v_all) if self.encoder.output_dims[1] > 0 else s_all | ||
| s_p_latent, v_p_latent = self.encoder_to_flow(encoder_input) | ||
|
|
||
| if "water" not in data.node_types or data["water"].num_nodes == 0: | ||
| return torch.zeros(0, device=device) | ||
|
|
Implementing the confidence model.
smootherstep_target: soft cutoff on nearest-GT distance, 1 inside 0.5 Å, 0 outside 1.5 Å.cluster_waters_vdw: absorb into confidence-weighted centroids at 1.52 Å (oxygen vdW), then NMS over those centroids. Centroids carry the cluster max.ConfidenceGVP—FlowWaterGVP's backbone minus time and self-conditioning, one scalar per candidate, PW+PP edges only. Backbone stays structurally identical so the scorer warm-starts from a flow checkpoint. This needs the dynamic edge type policy and is hence built off PR Edge type flags #89