Flow + GVP updates

[vratins]

• Add support for encoder-provided edge features (pp_edge_attr) in flow model, enabling GVP encoder to pass learned edge representations through the message passing layers
• Add configurable n_message_gvps and n_update_gvps parameters to control GVP layer depth
• Refactor training_step to separate gradient accumulation from optimizer step (caller now handles optimizer)
• Use constants (NUM_RBF, RBF_CUTOFF) instead of hardcoded values throughout GVP module
• Add comprehensive docstrings and standardize formatting

[Copilot]

Pull request overview

This PR extends the flow model to optionally consume encoder-produced protein–protein edge features through the GVP message passing stack, while also refactoring training to support external optimizer stepping / gradient accumulation and improving GVP module configurability + documentation.

Changes:

• Add support for passing pp_edge_attr from the encoder into ProteinWaterUpdate / GVPMultiEdgeConv and make GVP depth configurable via n_message_gvps / n_update_gvps.
• Refactor FlowMatcher.training_step to perform forward+backward only (caller performs optimizer.step()), with an accumulation_steps parameter.
• Introduce constants (NUM_RBF, RBF_CUTOFF) usage in GVP/flow modules and expand docstrings/formatting; update tests accordingly.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

File	Description
tests/test_flow.py	Updates mocks for the new encoder return shape and adds a forward test intended for an ESM-based encoder.
src/gvp.py	Adds documentation and adds optional cached edge attribute plumbing to heterogeneous GVP message passing; switches some defaults to constants.
src/flow.py	Threads encoder-provided PP edge features into the updater, adds configurable GVP depths, changes the vector field head dims, and refactors training_step for accumulation/caller-managed optimizer.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[marcuscollins]

I think these PRs need to be reorganized so that they don't break main. There are too many things here where it is hard to evaluate what's happening because they require changes in other places not included in this PR.

I've put comments in here, so have a look, but try to isolate each functional change, like adding an argument to a method's signature, and make separate PRs for those. Make a separate PR for changes that are just formatting.

[marcuscollins]

You get rid of some code in this PR (gradient clipping, stepping the optimizer) that worries me. I don't see where the optimizer.step() method is called, which can't be right, unless it was previously called in multiple places and you're removing one. Otherwise mostly requests for clarification.

[marcuscollins]

why get rid of gradient clipping? And optimizer.step()? This feels like something isn't right, since I don't see that step taken anywhere else.

[vratins]

I've kept only the forward pass and loss computation in the training step here, and moved the grad clipping to the training script -- this is primarily because it makes gradient accumulation easier to control.

[marcuscollins]

LGTM