Refactoring training and inference script to match source code changes by vratins · Pull Request #44 · diff-use/WaterFlow

vratins · 2026-03-03T22:29:33Z

Changes

scripts/train.py

Replace --use_slae flag with --encoder_type choices: gvp, slae, esm
Add dataset quality check args: --max_com_dist, --max_clash_fraction, --clash_dist, --interface_dist_threshold, --min_water_residue_ratio
Add per-water filtering args: --max_protein_dist, --min_edia, --max_bfactor_zscore with toggle flags
Add scheduler args: --scheduler, --warmup_steps, --eta_min_factor, --step_size, --step_gamma
Add model args: --slae_dim, --esm_dim, --grad_accum_steps, --drop_rate, --n_message_gvps, --n_update_gvps
New helpers: _extract_quality_config(), _build_dataset_config(), _required_embedding_field(), _resolve_embedding_dim(), resolve_encoder_config(), build_scheduler()
Training loop now supports gradient accumulation and dual schedulers (warmup + main)
Save resolved_encoder_config in config.json for inference compatibility

scripts/inference.py (+185 lines)

Import NUM_RBF from src.constants
Add --geometry_cache CLI arg to override model's geometry cache
Update build_model_from_config() to use resolved_encoder_config when available
Pass encoder_type, geometry_cache_name, include_mates to dataset

Note: A lot of the changes here are conditioned on changes in the source code that are part of other PRs

Copilot

Pull request overview

This PR refactors the training and inference entrypoint scripts to align CLI/config handling with recent model/dataset evolution (encoder selection via --encoder_type, expanded dataset filtering knobs, scheduler configuration, and saving a resolved encoder config for inference compatibility).

Changes:

Replaces --use_slae with --encoder_type {gvp,slae,esm} and adds helpers to resolve encoder config/dimensions from a sample batch.
Adds dataset quality + per-water filtering CLI args and logs the active/ignored filters into config.json.
Adds gradient accumulation and warmup+main scheduler support, and updates inference to prefer resolved_encoder_config plus configurable geometry cache selection.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 10 comments.

File	Description
`scripts/train.py`	Refactors CLI/config, adds dataset filter knobs, encoder resolution, gradient accumulation, warmup + main schedulers, and richer checkpoint/config saving.
`scripts/inference.py`	Updates config/model build logic to use resolved encoder config when present, adds geometry cache override, and passes encoder/cache settings into dataset creation.

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

Copilot · 2026-03-03T22:35:10Z

 def main():
+    """Run inference pipeline on a list of PDB structures."""
+    setup_logging_for_tqdm()
    args = parse_args()


setup_logging_for_tqdm() is already called at module import time, and it’s called again here in main(). This can lead to duplicate handlers or unexpected log formatting depending on implementation. Prefer initializing logging in one place (typically inside main() only) and removing the module-level side effect.

Copilot · 2026-03-03T22:35:11Z

    train_loader = get_dataloader(
-        args.train_list, args.processed_dir,
-        batch_size=args.batch_size, shuffle=True,
+        pdb_list_file=args.train_list,
+        processed_dir=args.processed_dir,
+        batch_size=args.batch_size,
+        shuffle=True,
        num_workers=args.num_workers,
-        base_pdb_dir=args.base_pdb_dir,
-        include_mates=args.include_mates,
+        pin_memory=args.pin_memory,
+        prefetch_factor=args.prefetch_factor,
+        persistent_workers=args.persistent_workers,
        duplicate_single_sample=args.duplicate_single_sample,
+        **dataset_kwargs,


get_dataloader() in src/dataset.py does not accept prefetch_factor/persistent_workers and will forward them into ProteinWaterDataset(**dataset_kwargs), where they are not valid constructor args either. This will raise a TypeError at runtime. Either update src.dataset.get_dataloader to accept these DataLoader-only kwargs and pass them into torch.utils.data.DataLoader, or stop passing them here until the dataset/loader APIs support them.

Copilot · 2026-03-03T22:35:11Z

+    # Handle remaining gradients at end of epoch
+    if (step + 1) % args.grad_accum_steps != 0:
+        if args.grad_clip > 0:


train_epoch() uses step after the dataloader loop to decide whether to flush remaining gradients. If train_loader is empty, step is never set and this will raise UnboundLocalError. Consider handling the empty-loader case explicitly (e.g., initialize step = -1 before the loop or early-return when len(train_loader) == 0).

Copilot · 2026-03-03T22:35:11Z

+        encoder_type=encoder_type,
        include_mates=include_mates,
+        geometry_cache_name=geometry_cache_name,


ProteinWaterDataset.__init__ in src/dataset.py does not currently accept encoder_type or geometry_cache_name. Passing these here will raise a TypeError. Either extend ProteinWaterDataset to support these options, or remove them from the constructor call and rely on the dataset’s existing parameters until the underlying dataset API is updated.

Suggested change

encoder_type=encoder_type,

include_mates=include_mates,

geometry_cache_name=geometry_cache_name,

include_mates=include_mates,

Copilot · 2026-03-03T22:35:11Z

    val_loader = get_dataloader(
-        args.val_list, args.processed_dir,
-        batch_size=args.batch_size, shuffle=False,
+        pdb_list_file=args.val_list,
+        processed_dir=args.processed_dir,
+        batch_size=args.batch_size,
+        shuffle=False,
        num_workers=args.num_workers,
-        base_pdb_dir=args.base_pdb_dir,
-        include_mates=args.include_mates,
-        duplicate_single_sample=1,  # Don't duplicate for validation
+        pin_memory=args.pin_memory,
+        prefetch_factor=args.prefetch_factor,
+        persistent_workers=args.persistent_workers,
+        duplicate_single_sample=args.duplicate_single_sample,
+        **dataset_kwargs,
    )


dataset_kwargs includes encoder_type and geometry_cache_name, but ProteinWaterDataset.__init__ (in src/dataset.py) currently does not accept those parameters. As written, get_dataloader(..., **dataset_kwargs) will error when constructing the dataset. Either extend ProteinWaterDataset to accept/use these fields, or remove them from dataset_kwargs until the core dataset supports encoder/geometry cache selection.

Copilot · 2026-03-03T22:35:12Z

+        pin_memory=args.pin_memory,
+        prefetch_factor=args.prefetch_factor,
+        persistent_workers=args.persistent_workers,
+        duplicate_single_sample=args.duplicate_single_sample,


Validation loader is now using duplicate_single_sample=args.duplicate_single_sample, which will duplicate the validation set when running single-sample overfitting experiments and can artificially inflate/alter validation metrics. Consider forcing duplicate_single_sample=1 for validation (as was done previously) so eval reflects the true dataset size.

Suggested change

duplicate_single_sample=args.duplicate_single_sample,

duplicate_single_sample=1,

Copilot · 2026-03-03T22:35:12Z

+        n_message_gvps=args.n_message_gvps,
+        n_update_gvps=args.n_update_gvps,


FlowWaterGVP (in src/flow.py) does not currently accept n_message_gvps or n_update_gvps keyword arguments, so this constructor call will raise TypeError: __init__() got an unexpected keyword argument .... Either plumb these through FlowWaterGVP/ProteinWaterUpdate (and into GVPMultiEdgeConv) or remove these arguments from the script until the model supports them.

Suggested change

n_message_gvps=args.n_message_gvps,

n_update_gvps=args.n_update_gvps,

Copilot · 2026-03-03T22:35:12Z

        metrics = flow_matcher.training_step(
-            batch, optimizer,
-            grad_clip=args.grad_clip,
+            batch,
            use_self_conditioning=args.use_self_cond,
+            accumulation_steps=args.grad_accum_steps,
        )


FlowMatcher.training_step() in src/flow.py currently requires an optimizer argument and performs optimizer.zero_grad() / backward() / step() internally. This call now omits the optimizer and passes accumulation_steps, which is not a supported parameter, so training will fail with a TypeError (and gradient accumulation logic would also double-step if both places manage stepping). To fix: either update FlowMatcher.training_step to support gradient accumulation (accept accumulation_steps, do scaled loss/backward only, and let the outer loop step), or revert this call to the existing API and keep optimizer stepping inside FlowMatcher.

Copilot · 2026-03-03T22:35:12Z

 from loguru import logger
 from tqdm import tqdm

+from src.constants import NUM_RBF


src/constants.py does not define NUM_RBF, so from src.constants import NUM_RBF will fail at import time. Either add NUM_RBF to src/constants.py (and keep it consistent with the model’s expected RBF dimension) or stop importing it here and use the correct existing constant/config key.

Suggested change

from src.constants import NUM_RBF

Copilot · 2026-03-03T22:35:13Z

+        n_message_gvps=config.get("n_message_gvps", 2),
+        n_update_gvps=config.get("n_update_gvps", 2),


FlowWaterGVP.__init__ (in src/flow.py) does not currently accept n_message_gvps / n_update_gvps keyword args, so this call will raise TypeError. If you intend these to be configurable, the parameters need to be added/plumbed through the model implementation; otherwise remove them from inference model construction to match the actual FlowWaterGVP signature.

Suggested change

n_message_gvps=config.get("n_message_gvps", 2),

n_update_gvps=config.get("n_update_gvps", 2),

DorisMai · 2026-03-04T18:26:02Z

        # build result dicts similar to rk4
        results = []
        for graph, water_pred in zip(graphs, water_preds):
-            results.append({
-                "protein_pos": graph["protein"].pos.numpy(),
-                "water_true": graph["water"].pos.numpy(),
-                "water_pred": water_pred,
-                "trajectory": None,
-                "pdb_id": getattr(graph, 'pdb_id', None),
-            })
+            results.append(
+                {
+                    "protein_pos": graph["protein"].pos.numpy(),
+                    "water_true": graph["water"].pos.numpy(),
+                    "water_pred": water_pred,
+                    "trajectory": None,
+                    "pdb_id": getattr(graph, "pdb_id", None),
+                }
+            )


why is this dict building here rather than inside euler_integration function?

Yep, fixed this

DorisMai · 2026-03-04T18:33:58Z

A more detailed/updated explanation would be helpful. I had to go to PR #43 to see the difference between processed_dir and geometry_cache

Added more detail to the help message

DorisMai · 2026-03-04T19:04:00Z

I would assume all the metrics to compute that rely on ground truth comparison are technically not part of the inference logic. If there is no ground truth or one thinks the deposited info is wrong, your script should still be able to make a prediction of water placement?

Added a skip_metrics flag for the case one does not have/does not want to use ground truth data. In the future I would assume something cleaner with config files.

DorisMai · 2026-03-09T22:11:45Z

+    # TODO: Remove hardcoded default paths. These should be required arguments
+    # or loaded from environment variables / config files for portability.
+    # Current hardcoded paths:
+    #   - processed_dir: /home/srivasv/flow_cache/
+    #   - base_pdb_dir: /sb/wankowicz_lab/data/srivasv/pdb_redo_data
+    #   - edia_dir: /sb/wankowicz_lab/data/srivasv/edia_results
+    #   - save_dir: /home/srivasv/flow_checkpoints
+    #   - wandb_dir: /home/srivasv/wandb_logs


DorisMai · 2026-03-09T22:12:12Z

+    # TODO: Add support for loading configuration from YAML/JSON config files.
+    # This would allow users to save and share training configurations easily.
+    # Example: --config config.yaml would load all arguments from the file,
+    # with CLI args taking precedence for overrides.


DorisMai · 2026-03-09T22:51:08Z

-            pred=out['water_pred'],
-            true=out['water_true'],
-            threshold=1.0
+            pred=out["water_pred"], true=out["water_true"], threshold=1.0


should probably expose threshold as an argument, same as in inference.py?

Yep, addressing this

DorisMai · 2026-03-09T22:57:43Z

+    p.add_argument(
+        "--no_persistent_workers",
+        dest="persistent_workers",
+        action="store_false",
+        help="Disable persistent_workers",
+    )


I don't think you need these kind of paired flags anymore. You can just do p.add_argument("--persistent_workers", action="store_true") and don't set a default value and it should work fine. similarly for pin_memory.

Yep, simplifying this

DorisMai · 2026-03-09T23:01:09Z

            device=device,
            water_ratio=water_ratio,
        )
-        # build result dicts similar to rk4


in which PR did you @vratins update src/flow.py to have euler_integrate return dicts? If it hasn't happened make an issue?

That would be PR #38

DorisMai

Remaining issues to address in future PRs: issue #13 #29 #49 #50

refactoring training and inference script to match source code changes

708458c

vratins requested review from DorisMai, Copilot and marcuscollins March 3, 2026 22:29

Copilot started reviewing on behalf of vratins March 3, 2026 22:30 View session

Copilot AI reviewed Mar 3, 2026

View reviewed changes

DorisMai reviewed Mar 4, 2026

View reviewed changes

addressing changes in inference script

a1ae936

DorisMai reviewed Mar 9, 2026

View reviewed changes

addressing PR comments for train.py

2907952

vratins linked an issue Mar 10, 2026 that may be closed by this pull request

use loguru for logging instead of print statements #18

Closed

This was referenced Mar 10, 2026

Clarify usage of duplicate_single_sample flag #49

Open

Rename CLI flags to be consistent between train.py and inference.py #50

Open

DorisMai approved these changes Mar 10, 2026

View reviewed changes

vratins merged commit 63c646e into main Mar 10, 2026

vratins mentioned this pull request Mar 17, 2026

Flow + GVP updates #38

Merged

vratins deleted the dev_training_refactor branch March 19, 2026 15:43

	duplicate_single_sample=args.duplicate_single_sample,
	duplicate_single_sample=1,

		n_message_gvps=args.n_message_gvps,
		n_update_gvps=args.n_update_gvps,

		n_message_gvps=config.get("n_message_gvps", 2),
		n_update_gvps=config.get("n_update_gvps", 2),

Uh oh!

Conversation

vratins commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DorisMai Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DorisMai left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

vratins commented Mar 3, 2026 •

edited

Loading

DorisMai Mar 9, 2026 •

edited

Loading