Edge type flags - #89
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe PR replaces fixed KNN edge construction with configurable radius and KNN policies. It adds fallback rescue, edge-type ablations, model-owned configuration, CLI and inference wiring, README documentation, and expanded regression tests. ChangesDynamic edge construction
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingCLI
participant FlowWaterGVP
participant ProteinWaterUpdate
participant build_dynamic_edges
TrainingCLI->>FlowWaterGVP: pass dynamic-edge configuration
FlowWaterGVP->>ProteinWaterUpdate: initialize active edge settings
ProteinWaterUpdate->>build_dynamic_edges: construct water interaction edges
build_dynamic_edges-->>ProteinWaterUpdate: return dynamic edges
ProteinWaterUpdate-->>FlowWaterGVP: return assembled graph
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
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
This PR makes edge construction in the flow model configurable (edge-type ablations and radius-vs-kNN dynamic graph building), and updates training/inference plumbing plus tests/docs to support replaying historical configs and new runtime edge policies.
Changes:
- Introduces dynamic edge construction controls (policy, cutoff, max_neighbors, fallback k, and WW/WP ablations) and threads them through
FlowWaterGVP/ProteinWaterUpdate. - Replaces the previous KNN-only edge builder with a unified radius/KNN builder plus optional “rescue” edges for isolated nodes.
- Updates CLI/config loading, tests, and README documentation to reflect the new edge configuration surface.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_train_config.py | Adds regression test to ensure recorded configs with dynamic_edge_policy: "auto" replay successfully and map to the intended runtime policy/etypes. |
| tests/test_forward.py | Updates forward-pass test to validate dynamic edges via the new build_edges() signature (no explicit k args). |
| tests/test_flow.py | Expands unit tests for radius/KNN edge building semantics, row conventions, fallback behavior, and policy resolution. |
| src/flow.py | Implements resolve_edge_policy, build_dynamic_edges, edge-type configuration/ablation, and fallback (“rescue”) logic in ProteinWaterUpdate. |
| src/constants.py | Adds get_active_edge_types() helper for WW/WP ablation while keeping PW/PP always enabled. |
| scripts/train.py | Replaces old k-only flags with a richer edge configuration CLI and wires args into model construction. |
| scripts/inference.py | Loads the same edge configuration keys from config.json when reconstructing models for inference. |
| README.md | Updates edge-type descriptions and adds an “Edge Construction” section documenting the new policy controls. |
Suppressed comments (2)
src/flow.py:259
- build_dynamic_edges currently treats any
policyother than "knn" as the radius path. That means typos (or accidentally passing "knn_if_isolated") silently change behavior instead of failing fast.
if policy == "knn":
# Asked for each destination's nearest sources, so sources come back second.
dst_idx, src_idx = knn(
x=src_pos, y=dst_pos, k=k, batch_x=batch_src, batch_y=batch_dst
)
README.md:310
- The CLI-flag summary table is out of sync with the implemented arguments: it lists
--dynamic_edge_policydefault asradiusand only mentionsradius/knn, but the CLI defaults toautoand also supportsknn_if_isolated; additionally,--max_neighborsand--k_wp(and the knn k's) are missing.
| `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](#edge-construction)) |
| `--cutoff` | `8.0` | Distance cutoff in Å for radius edges |
| `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query; `0` disables |
| `--disable_ww` | `false` | Ablate water→water edges |
| `--disable_wp` | `false` | Ablate water→protein edges |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| src_pos: (N_src, 3) source node positions. | ||
| dst_pos: (N_dst, 3) destination node positions. | ||
| k: Number of nearest source neighbors to find per destination node. | ||
| policy: One of DYNAMIC_EDGE_POLICIES. |
| if knn_fallback_k < 0: | ||
| raise ValueError(f"knn_fallback_k must be >= 0, got {knn_fallback_k}") | ||
|
|
||
| # build_edges only knows how to construct the four known relations, and | ||
| # HeteroConv would KeyError on any other, so reject it up front. | ||
| unknown = [et for et in (etypes or []) if et not in ALL_EDGE_TYPES] | ||
| if unknown: | ||
| raise ValueError( | ||
| f"etypes must be a subset of {ALL_EDGE_TYPES}, got unknown {unknown}" | ||
| ) | ||
|
|
||
| self.cutoff = cutoff | ||
| self.max_neighbors = max_neighbors | ||
| resolved = resolve_edge_policy(dynamic_edge_policy, sampling_strategy) | ||
| # Same edges as "radius"; the difference is the extra pass for nodes left with none. | ||
| self.rescue_isolated = resolved == "knn_if_isolated" and knn_fallback_k > 0 | ||
| self.dynamic_edge_policy = ( | ||
| "radius" if resolved == "knn_if_isolated" else resolved | ||
| ) | ||
| self.knn_fallback_k = knn_fallback_k | ||
| self.k_pw = k_pw | ||
| self.k_ww = k_ww | ||
| self.k_wp = k_wp |
| - Only PP edges are stored in the geometry cache; every water-touching edge is | ||
| rebuilt each forward pass, since water positions move during integration. See | ||
| [Edge Construction](#edge-construction) |
| | `radius` (default) | Connect every pair within `--cutoff`, capped at `--max_neighbors` per source | | ||
| | `knn` | Connect a fixed number of nearest neighbours (`--k_pw`, `--k_ww`, `--k_wp`) | |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
scripts/train.py (1)
1101-1108: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPass
sampling_strategytoFlowMatcher.
build_modelresolvesdynamic_edge_policyfrom the run's sampling strategy.FlowMatcherhere keeps its"uniform_ball"default. If a run selectsscaled_gaussian, the model resolves"auto"toknn_if_isolatedwhile the matcher still samples from the uniform ball prior. The two settings then describe different runs.🐛 Proposed fix
flow_matcher = FlowMatcher( model=model, p_self_cond=args.p_self_cond, use_distortion=args.use_distortion, p_distort=args.p_distort, t_distort=args.t_distort, sigma_distort=args.sigma_distort, + sampling_strategy=args.sampling_strategy, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/train.py` around lines 1101 - 1108, Pass the run’s sampling_strategy argument into the FlowMatcher construction alongside the existing distortion parameters, ensuring it uses the same resolved strategy as build_model. Preserve the current argument wiring and FlowMatcher behavior for callers whose strategy is unspecified.
🧹 Nitpick comments (3)
src/flow.py (2)
673-684: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument
sampling_strategyin the constructor docstring.The signature adds
sampling_strategy, and the value decides how"auto"resolves. The Args block does not list it. Readers cannot see that the prior affects edge construction.📝 Proposed docstring addition
dynamic_edge_policy: How water-touching edges are built, one of DYNAMIC_EDGE_POLICIES. Default: "radius" + sampling_strategy: Prior the run uses. Consulted only to resolve + a "auto" policy; see `resolve_edge_policy`. Default: "uniform_ball" knn_fallback_k: Nearest neighbours attached to waters the radius🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/flow.py` around lines 673 - 684, Update the constructor docstring’s Args section to document the sampling_strategy parameter, including that it controls how the “auto” strategy is resolved and affects edge construction. Place it alongside the other edge-construction configuration arguments and describe its default or accepted values using the existing symbols.
527-537: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a PP-specific neighbour count instead of
k_pwfor the dynamic protein-protein fallback.Under
policy="knn", this branch builds protein->protein edges withk=self.k_pw.k_pwis documented as the count for protein->water edges. The two relations have different densities, so the reuse couples an unrelated setting to PP. The branch runs only when the dataset carries no cached PP edges, but that is exactly the path where the value matters.Add a
k_ppparameter, or state the reuse in a comment so it is a deliberate choice.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/flow.py` around lines 527 - 537, Update the protein-protein fallback in the dynamic edge construction around build_dynamic_edges to use a dedicated k_pp neighbor-count parameter rather than k_pw. Add and propagate k_pp through the relevant flow configuration and call sites, preserving k_pw exclusively for protein-water edges.tests/test_flow.py (1)
422-461: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the water->protein rescue axis.
test_radius_strands_far_water_and_fallback_rescues_itcovers_add_knn_fallbackwithisolate_axis=1. Theisolate_axis=0branch remaps rows differently: it queries with the isolated sources and then swaps the returned rows. No test pins that remapping, so a row swap there would pass silently.Add a case that strands a protein atom far from every water and asserts that the rescued
EDGE_WPedges keep row 0 inside the water index range and row 1 inside the protein index range.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_flow.py` around lines 422 - 461, Add a test alongside test_radius_strands_far_water_and_fallback_rescues_it that places one protein atom far from all waters, enables radius isolation rescue with knn_fallback_k > 0, and inspects rescued EDGE_WP edges. Assert row 0 contains only valid water indices and row 1 contains only valid protein indices, pinning the isolate_axis=0 query-and-row-swap behavior in _add_knn_fallback.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@README.md`:
- Around line 211-231: Update the Edge Construction section to document all four
dynamic_edge_policy values: auto, radius, knn, and knn_if_isolated. Describe
that auto resolves to knn_if_isolated for scaled_gaussian sampling and otherwise
follows the standard default behavior, and state that stranded-water rescue via
knn_fallback_k requires the resolved knn_if_isolated policy; positive fallback
values do not rescue under plain radius.
- Around line 306-308: Update the README options table to match
scripts/train.py: document dynamic_edge_policy’s default as auto and include
auto, radius, knn, and knn_if_isolated as valid policies. Revise the
knn_fallback_k description to state that the rescue applies only to
knn_if_isolated runs, while preserving its default and disabled value.
In `@scripts/inference.py`:
- Around line 279-281: Update the FlowMatcher construction in
scripts/inference.py to pass the recorded sampling_strategy from config,
defaulting to "uniform_ball" consistently with the earlier configuration
handling. Preserve the existing model and p_self_cond arguments while ensuring
integration uses the training run’s selected prior.
In `@scripts/train.py`:
- Around line 628-638: The training CLI must define the sampling_strategy
argument before it is read. In scripts/train.py lines 628-638, retain
sampling_strategy=args.sampling_strategy in the FlowMatcher construction; in
scripts/train.py lines 1101-1108, add the parse_args option with uniform_ball
and scaled_gaussian choices; and in scripts/inference.py lines 279-281, make the
corresponding argument wiring consistent as required by the existing FlowMatcher
configuration.
In `@src/flow.py`:
- Around line 498-499: Update the comment above rescue_isolated in the relevant
flow method to accurately describe that the extra pass applies to both
protein→water and water→protein edges, while preserving the rescue assignment
itself.
- Around line 265-283: Update the homogeneous branch in build_dynamic_edges to
return source indices in row 0 by using the appropriate radius_graph flow or
flipping its output. Apply max_neighbors through the source-side neighbor limit
rather than max_num_neighbors, while preserving the existing candidate cap and
non-homogeneous radius behavior.
---
Outside diff comments:
In `@scripts/train.py`:
- Around line 1101-1108: Pass the run’s sampling_strategy argument into the
FlowMatcher construction alongside the existing distortion parameters, ensuring
it uses the same resolved strategy as build_model. Preserve the current argument
wiring and FlowMatcher behavior for callers whose strategy is unspecified.
---
Nitpick comments:
In `@src/flow.py`:
- Around line 673-684: Update the constructor docstring’s Args section to
document the sampling_strategy parameter, including that it controls how the
“auto” strategy is resolved and affects edge construction. Place it alongside
the other edge-construction configuration arguments and describe its default or
accepted values using the existing symbols.
- Around line 527-537: Update the protein-protein fallback in the dynamic edge
construction around build_dynamic_edges to use a dedicated k_pp neighbor-count
parameter rather than k_pw. Add and propagate k_pp through the relevant flow
configuration and call sites, preserving k_pw exclusively for protein-water
edges.
In `@tests/test_flow.py`:
- Around line 422-461: Add a test alongside
test_radius_strands_far_water_and_fallback_rescues_it that places one protein
atom far from all waters, enables radius isolation rescue with knn_fallback_k >
0, and inspects rescued EDGE_WP edges. Assert row 0 contains only valid water
indices and row 1 contains only valid protein indices, pinning the
isolate_axis=0 query-and-row-swap behavior in _add_knn_fallback.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ba98ee21-f527-4166-aab9-b78c54fed4fd
📒 Files selected for processing (8)
README.mdscripts/inference.pyscripts/train.pysrc/constants.pysrc/flow.pytests/test_flow.pytests/test_forward.pytests/test_train_config.py
| | `--dynamic_edge_policy` | Behaviour | | ||
| |-------------------------|-----------| | ||
| | `radius` (default) | Connect every pair within `--cutoff`, capped at `--max_neighbors` per source | | ||
| | `knn` | Connect a fixed number of nearest neighbours (`--k_pw`, `--k_ww`, `--k_wp`) | | ||
|
|
||
| The two differ in which side the neighbour budget applies to. KNN queries *per | ||
| destination*, so every destination is guaranteed edges but a source may have | ||
| none — coverage checks must read the destination row. Radius guarantees nothing: | ||
| a water with no protein atom inside `--cutoff` gets no PW edges at all. | ||
|
|
||
| `--knn_fallback_k` repairs that. Under `radius`, any water the query stranded is | ||
| reconnected to that many nearest protein atoms regardless of distance. Set it to | ||
| `0` to disable. It has no effect under `knn`, which cannot strand a node. | ||
|
|
||
| Set `--disable_ww` / `--disable_wp` to ablate those edge types; PW and PP are | ||
| always active. | ||
|
|
||
| > Configs written before the radius/KNN split recorded a three-valued | ||
| > `dynamic_edge_policy` (`auto`, `radius`, `knn_if_isolated`). All three built a | ||
| > radius graph, so they load and map to `radius`; whether stranded waters are | ||
| > rescued is now `--knn_fallback_k`'s job. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
The Edge Construction section does not match the shipped policy set.
scripts/train.py accepts four values: auto, radius, knn, and knn_if_isolated. The default is auto. This section lists only radius and knn, and it describes auto and knn_if_isolated as legacy values that map to radius. resolve_edge_policy maps auto to knn_if_isolated when the sampling strategy is scaled_gaussian.
Lines 221-223 also state that --knn_fallback_k repairs stranded waters "under radius". ProteinWaterUpdate.__init__ sets rescue_isolated only when the resolved policy is knn_if_isolated. A positive --knn_fallback_k under radius has no effect, which test_plain_radius_does_not_rescue pins.
Document all four values, the auto resolution rule, and the fact that the rescue requires knn_if_isolated.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@README.md` around lines 211 - 231, Update the Edge Construction section to
document all four dynamic_edge_policy values: auto, radius, knn, and
knn_if_isolated. Describe that auto resolves to knn_if_isolated for
scaled_gaussian sampling and otherwise follows the standard default behavior,
and state that stranded-water rescue via knn_fallback_k requires the resolved
knn_if_isolated policy; positive fallback values do not rescue under plain
radius.
| | `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](#edge-construction)) | | ||
| | `--cutoff` | `8.0` | Distance cutoff in Å for radius edges | | ||
| | `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query; `0` disables | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the documented defaults.
The table shows --dynamic_edge_policy default radius and lists only radius or knn. scripts/train.py sets default="auto" with choices=["auto", "radius", "knn", "knn_if_isolated"]. The --knn_fallback_k row implies that the rescue applies to radius runs; it applies only to knn_if_isolated.
📝 Proposed table fix
-| `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](`#edge-construction`)) |
+| `--dynamic_edge_policy` | `auto` | How water-touching edges are built: `auto`, `radius`, `knn`, or `knn_if_isolated` (see [Edge Construction](`#edge-construction`)) |
| `--cutoff` | `8.0` | Distance cutoff in Å for radius edges |
-| `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query; `0` disables |
+| `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query under `knn_if_isolated`; `0` disables |📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| | `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](#edge-construction)) | | |
| | `--cutoff` | `8.0` | Distance cutoff in Å for radius edges | | |
| | `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query; `0` disables | | |
| | `--dynamic_edge_policy` | `auto` | How water-touching edges are built: `auto`, `radius`, `knn`, or `knn_if_isolated` (see [Edge Construction](`#edge-construction`)) | | |
| | `--cutoff` | `8.0` | Distance cutoff in Å for radius edges | | |
| | `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query under `knn_if_isolated`; `0` disables | |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@README.md` around lines 306 - 308, Update the README options table to match
scripts/train.py: document dynamic_edge_policy’s default as auto and include
auto, radius, knn, and knn_if_isolated as valid policies. Revise the
knn_fallback_k description to state that the rescue applies only to
knn_if_isolated runs, while preserving its default and disabled value.
| # "auto" depends on which prior the run uses, so pass that through. | ||
| sampling_strategy=config.get("sampling_strategy", "uniform_ball"), | ||
| knn_fallback_k=config.get("knn_fallback_k", 8), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Also forward sampling_strategy to FlowMatcher.
This call restores the recorded sampling_strategy for the model. The FlowMatcher construction later in this file does not receive it, so integration samples the water prior with the "uniform_ball" default. A run trained with "scaled_gaussian" then integrates from a different prior than it was trained on.
🐛 Proposed fix at the FlowMatcher call site
FlowMatcher(
model=model,
p_self_cond=config.get("p_self_cond", 0.5),
sampling_strategy=config.get("sampling_strategy", "uniform_ball"),
)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/inference.py` around lines 279 - 281, Update the FlowMatcher
construction in scripts/inference.py to pass the recorded sampling_strategy from
config, defaulting to "uniform_ball" consistently with the earlier configuration
handling. Preserve the existing model and p_self_cond arguments while ensuring
integration uses the training run’s selected prior.
| cutoff=args.cutoff, | ||
| max_neighbors=args.max_neighbors, | ||
| dynamic_edge_policy=args.dynamic_edge_policy, | ||
| # "auto" depends on which prior the run uses, so pass that through. | ||
| sampling_strategy=args.sampling_strategy, | ||
| knn_fallback_k=args.knn_fallback_k, | ||
| disable_ww=args.disable_ww, | ||
| disable_wp=args.disable_wp, | ||
| k_pw=args.k_pw, | ||
| k_ww=args.k_ww, | ||
| k_wp=args.k_wp, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Trace every sampling_strategy producer and consumer.
rg -n -C3 'sampling_strategy' scripts/ src/ tests/Repository: diff-use/WaterFlow
Length of output: 6895
🏁 Script executed:
#!/bin/bash
# Inspect the argument parser and the FlowMatcher call sites to determine the exact wiring failure.
set -e
printf '\n## scripts/train.py parse_args sampling_strategy references\n'
rg -n -C4 'def parse_args|add_argument|sampling_strategy|args\.' scripts/train.py | sed -n '1,220p'
printf '\n## scripts/train.py FlowMatcher construction\n'
sed -n '1076,1112p' scripts/train.py | cat -n
printf '\n## scripts/inference.py FlowMatcher construction\n'
sed -n '260,290p' scripts/inference.py | cat -n
printf '\n## src/flow.py FlowMatcher signature and sampling implementation\n'
sed -n '860,955p' src/flow.py | cat -n
sed -n '850,898p' src/flow.py | cat -nRepository: diff-use/WaterFlow
Length of output: 17436
Wire sampling_strategy through training. scripts/train.py#L632 reads args.sampling_strategy, but scripts/train.py:parse_args() does not define --sampling_strategy, so training fails at startup. Add the CLI option with uniform_ball / scaled_gaussian choices, then pass sampling_strategy=args.sampling_strategy to the FlowMatcher construction.
📍 Affects 2 files
scripts/train.py#L628-L638(this comment)scripts/train.py#L1101-L1108scripts/inference.py#L279-L281
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/train.py` around lines 628 - 638, The training CLI must define the
sampling_strategy argument before it is read. In scripts/train.py lines 628-638,
retain sampling_strategy=args.sampling_strategy in the FlowMatcher construction;
in scripts/train.py lines 1101-1108, add the parse_args option with uniform_ball
and scaled_gaussian choices; and in scripts/inference.py lines 279-281, make the
corresponding argument wiring consistent as required by the existing FlowMatcher
configuration.
| # Cap against the number of reachable counterparts, minus the self-edge that | ||
| # a homogeneous query would otherwise spend a slot on. | ||
| num_candidates = dst_pos.size(0) - 1 if homogeneous else dst_pos.size(0) | ||
| cap = max(1, min(num_candidates, max_neighbors)) | ||
|
|
||
| if homogeneous: | ||
| return radius_graph( | ||
| src_pos, r=r, batch=batch_src, loop=False, max_num_neighbors=cap | ||
| ) | ||
|
|
||
| return idx.unique(dim=1) | ||
| # Asked for each source's neighbours within r, so sources already come back first. | ||
| return radius( | ||
| x=dst_pos, | ||
| y=src_pos, | ||
| r=r, | ||
| batch_x=batch_dst, | ||
| batch_y=batch_src, | ||
| max_num_neighbors=cap, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
torch_geometric.nn.radius_graph flow source_to_target row order max_num_neighbors semantics
💡 Result:
The torch_geometric.nn.radius_graph function computes a graph where edges connect all points within a specified radius r. Key parameters and their semantics are described below: max_num_neighbors: This parameter limits the maximum number of neighbors returned for each query point [1][2][3]. By default, it is set to 32 [1][4]. Note that the CPU implementation with this parameter enabled can be biased toward certain quadrants; it is recommended to set this to None or move inputs to the GPU for unbiased results [1][2][3]. flow: This parameter determines the direction of the edge index and is intended for use with message passing [1][2][3]. - 'source_to_target' (default): Produces an edge index in the format [source, target], where the source nodes point to the target nodes [1][2][3]. - 'target_to_source': Produces an edge index in the format [target, source] [2][3]. Row Order and Semantics: The function internally computes neighbors and returns an edge index tensor [2][3]. If flow is 'source_to_target', the function generates edges where the primary direction is from the source to the target [2][3]. Internally, the library computes the indices and then flips the rows (using.flip) if flow is 'source_to_target' to ensure the output matches the expected [source, target] orientation [2][3]. This ensures compatibility with downstream MessagePassing layers that expect specific source and target node orderings [1][2].
Citations:
- 1: https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.pool.radius_graph.html
- 2: https://pytorch-geometric.readthedocs.io/en/latest/%5Fmodules/torch%5Fgeometric/nn/pool.html
- 3: https://github.com/pyg-team/pytorch_geometric/blob/1f0661ce/torch_geometric/nn/pool/__init__.py
- 4: https://pytorch-geometric.readthedocs.io/en/stable/generated/torch_geometric.nn.pool.radius_graph.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the relevant function plus nearby calls/tests.
fd -a 'flow.py|test.*torch_geometric.*radius.*row_convention|test.*dynamic|test.*water' . | sed 's#^\./##' | head -100
echo '--- flow.py outline ---'
ast-grep outline src/flow.py --view compact || true
echo '--- relevant flow.py lines ---'
sed -n '180,300p' src/flow.py | nl -ba -v180
echo '--- heterogeneous test references ---'
rg -n "test_torch_geometric_radius_row_convention_unchanged|build_dynamic_edges|radius_graph|r_.*=|max_neighbors|max_num_neighbors|flow=" -S . | head -200Repository: diff-use/WaterFlow
Length of output: 472
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- src/flow.py lines 180-300 ---'
awk 'NR>=180 && NR<=300 {printf "%5d\t%s\n", NR, $0}' src/flow.py
echo '--- tests references ---'
rg -n "test_torch_geometric_radius_row_convention_unchanged|build_dynamic_edges|radius_graph|r_.*=|max_neighbors|max_num_neighbors|flow=" -S tests src | head -240
echo '--- tests around relevant names ---'
rg -n "torch_geometric|radius_row|radius_graph|neighbors|source_to_target|target_to_source|row0|row1" -S tests src | head -240
echo '--- deterministic probe: torch_geometric available? ---'
python3 - <<'PY'
import importlib.util
print("torch_geometric available:", importlib.util.find_spec("torch_geometric") is not None)
print("torch_cluster available:", importlib.util.find_spec("torch_cluster") is not None)
if importlib.util.find_spec("torch_geometric"):
import torch_geometric
print("torch_geometric version:", getattr(torch_geometric, "__version__", "unknown"))
PYRepository: diff-use/WaterFlow
Length of output: 26830
🌐 Web query:
torch_geometric nn radius graph source_to_target row order row 0 source row 1 target max_num_neighbors target
💡 Result:
In PyTorch Geometric (PyG), the radius_graph function computes a graph where edges connect nodes within a specified radius r [1][2]. Regarding your specific queries about flow and row structure: Flow and Row Order: The flow argument determines the direction of the edge index in the output Tensor [1][3]. When flow is set to 'source_to_target' (the default), the output edge_index follows a [target, source] row order [3][4]. Specifically, row 0 contains target indices, and row 1 contains source indices [3]. If set to 'target_to_source', the row order is [source, target] [3]. Max Number of Neighbors: The max_num_neighbors argument specifies the maximum number of neighbors to return for each node [1][2]. If the number of nodes within the radius r exceeds this limit, the function will return a subset of them (often chosen randomly) to respect this constraint [3][5]. Note that in some implementations, if this value is not provided, it defaults to 32 [1][2]. Additionally, users are advised that the CPU implementation may exhibit bias toward certain quadrants when using this parameter, and moving inputs to the GPU or setting max_num_neighbors to None is recommended if this bias is a concern [6].
Citations:
- 1: https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.pool.radius_graph.html
- 2: https://pytorch-geometric.readthedocs.io/en/2.7.0/generated/torch_geometric.nn.pool.radius_graph.html
- 3: https://github.com/rusty1s/pytorch_cluster/blob/master/torch_cluster/radius.py
- 4: https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/pool.html
- 5: https://github.com/rusty1s/pytorch_cluster/blob/master/README.md
- 6: https://pytorch-geometric.readthedocs.io/en/latest/generated/torch%5Fgeometric.nn.pool.radius%5Fgraph.html
Flip the homogeneous radius output row order and cap application.
radius_graph(..., flow="source_to_target") returns [target, source] as rows 0/1, but build_dynamic_edges() promises source indices in row 0. Either pass flow="target_to_source" or flip the result, and switch max_neighbors from max_num_neighbors to max_neighbors so the cap is applied to the source-side neighbours that the docstring exposes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/flow.py` around lines 265 - 283, Update the homogeneous branch in
build_dynamic_edges to return source indices in row 0 by using the appropriate
radius_graph flow or flipping its output. Apply max_neighbors through the
source-side neighbor limit rather than max_num_neighbors, while preserving the
existing candidate cap and non-homogeneous radius behavior.
| # Only protein-water edges get the extra pass; waters keep protein context anyway. | ||
| rescue = self.rescue_isolated |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the rescue comment.
The comment says that only protein-water edges get the extra pass. Line 564 applies the same rescue to water->protein edges.
📝 Proposed comment fix
- # Only protein-water edges get the extra pass; waters keep protein context anyway.
+ # Protein-water and water-protein edges get the extra pass; water-water does not.
rescue = self.rescue_isolated📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Only protein-water edges get the extra pass; waters keep protein context anyway. | |
| rescue = self.rescue_isolated | |
| # Protein-water and water-protein edges get the extra pass; water-water does not. | |
| rescue = self.rescue_isolated |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/flow.py` around lines 498 - 499, Update the comment above rescue_isolated
in the relevant flow method to accurately describe that the extra pass applies
to both protein→water and water→protein edges, while preserving the rescue
assignment itself.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (8)
src/flow.py:256
build_dynamic_edgestreats anypolicyvalue other than "knn" as the radius path, so typos or unexpected values silently change behavior. Since the docstring says the policy must be one ofDYNAMIC_EDGE_POLICIES, this should be validated and raise early.
# Same object
homogeneous = src_pos is dst_pos
if policy == "knn":
# Asked for each destination's nearest sources, so sources come back second.
README.md:90
- This says only PP edges are stored in the geometry cache and that every water-touching edge is rebuilt each forward pass, but
ProteinWaterUpdate.build_edgesexplicitly reuses cached PW edges when present (and there is a unit test asserting this behavior). The README should mention this exception so users understand when PW is rebuilt vs reused.
- Only PP edges are stored in the geometry cache; every water-touching edge is
rebuilt each forward pass, since water positions move during integration. See
[Edge Construction](#edge-construction)
README.md:215
- The Edge Construction table omits
autoandknn_if_isolated, and labelsradiusas the default, butscripts/train.pydefaults--dynamic_edge_policytoautoand the code supportsknn_if_isolated. Updating the table will prevent confusion when replaying configs or using CLI flags.
| `--dynamic_edge_policy` | Behaviour |
|-------------------------|-----------|
| `radius` (default) | Connect every pair within `--cutoff`, capped at `--max_neighbors` per source |
| `knn` | Connect a fixed number of nearest neighbours (`--k_pw`, `--k_ww`, `--k_wp`) |
README.md:231
- This blockquote implies that whether stranded waters are rescued is now solely
--knn_fallback_k’s job, but the code only enables rescue when the resolved policy isknn_if_isolated(includingautowithscaled_gaussian). Withdynamic_edge_policy=radius, rescue stays off even ifknn_fallback_k>0.
> Configs written before the radius/KNN split recorded a three-valued
> `dynamic_edge_policy` (`auto`, `radius`, `knn_if_isolated`). All three built a
> radius graph, so they load and map to `radius`; whether stranded waters are
> rescued is now `--knn_fallback_k`'s job.
README.md:306
- The documented default for
--dynamic_edge_policyisradius, but the CLI default inscripts/train.pyisauto. The README should match the actual flag default (and ideally list all supported values) to avoid confusion when starting new runs.
| `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](#edge-construction)) |
src/flow.py:499
- The comment says only protein→water edges get the fallback pass, but the code also applies
_add_knn_fallbackto water→protein edges whenrescueis enabled. Please update the comment so it matches the actual behavior (or restrict the rescue logic to PW only if that was the intent).
# Only protein-water edges get the extra pass; waters keep protein context anyway.
rescue = self.rescue_isolated
scripts/train.py:256
- The
--knn_fallback_khelp text implies the rescue runs under--dynamic_edge_policy radius, but the implementation only enables rescue when the resolved policy isknn_if_isolated(includingautowithscaled_gaussian). As written, the CLI help is misleading.
help=(
"Nearest neighbours attached to waters the radius query stranded; "
"0 disables the rescue. Ignored under --dynamic_edge_policy knn "
"(default: 8)"
),
README.md:223
- This paragraph says
--knn_fallback_krescues stranded waters “underradius”, but the implementation explicitly does not rescue under plainradius(only underknn_if_isolated/auto+scaled_gaussian). The README should reflect the actual gating so users don’t expectradius+fallback to work.
This issue also appears on line 228 of the same file.
`--knn_fallback_k` repairs that. Under `radius`, any water the query stranded is
reconnected to that many nearest protein atoms regardless of distance. Set it to
`0` to disable. It has no effect under `knn`, which cannot strand a node.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (3)
README.md:214
- The Edge Construction table says
radiusis the default and only documentsradius/knn, but the CLI default inscripts/train.pyisautoand the code also supportsknn_if_isolated. This mismatch will confuse users about what happens when they omit--dynamic_edge_policy.
| `--dynamic_edge_policy` | Behaviour |
|-------------------------|-----------|
| `radius` (default) | Connect every pair within `--cutoff`, capped at `--max_neighbors` per source |
| `knn` | Connect a fixed number of nearest neighbours (`--k_pw`, `--k_ww`, `--k_wp`) |
src/flow.py:242
build_dynamic_edgesdocumentspolicyas one ofDYNAMIC_EDGE_POLICIES, but the implementation only special-cases"knn"and otherwise falls back to radius behavior. That means typos (or passing"knn_if_isolated") silently change behavior instead of failing fast, and the current homogenous-graph detection usessrc_pos is dst_poswhich misses cases where the same tensor is passed via a different view/copy, leaving self-loops/cap logic inconsistent.
policy: One of DYNAMIC_EDGE_POLICIES.
k: Nearest neighbours per destination, used when policy is "knn".
r: Distance cutoff in Angstroms, used when policy is "radius".
max_neighbors: Per-source cap on radius results.
batch_src: (N_src,) batch assignment for source nodes, or None.
README.md:223
- The README says
--knn_fallback_krescues stranded waters underradius, but the implementation only enables the rescue underknn_if_isolated(orautoresolving to it) and explicitly disables rescue for plainradiuseven whenknn_fallback_k > 0. The docs should match the actual behavior so users can reason about isolation handling.
`--knn_fallback_k` repairs that. Under `radius`, any water the query stranded is
reconnected to that many nearest protein atoms regardless of distance. Set it to
`0` to disable. It has no effect under `knn`, which cannot strand a node.
etypes,cutoff,max_neighbors,dynamic_edge_policy,knn_fallback_konProteinWaterUpdate, which previously hardcodedALL_EDGE_TYPES. This is the dependencyConfidenceGVPneeds.get_active_edge_types(disable_ww, disable_wp)to ablate WW/WP edges; PW and PP always stay on.build_knn_edges→build_dynamic_edges, supporting both radius and kNN queries, plus a fallback that reconnects nodes a radius query left with no edges in the case of sampling with a scaled gaussian.--k_pw/--k_wwwith--dynamic_edge_policy,--cutoff,--max_neighbors,--knn_fallback_k,--disable_ww,--disable_wp;inference.pyreads the same keys back fromconfig.json.Summary by CodeRabbit
New Features
Documentation
Bug Fixes