Open
Conversation
Model changes: - Add output_mode='torsion' to ProteinLigandFlowMatching - Translation/rotation heads: mean-pooled ligand features → EquivariantMLP → 3D - Torsion head: src/dst node scalar concat → MLP → 1 scalar per rotatable bond - Zero-initialized outputs with learnable scales for stable training Data changes: - Dataset returns torsion_data with translation, rotation, torsion_changes, rotatable_edges, mask_rotate (computed on-the-fly or from pre-computed data) Loss: - Component-wise weighted MSE (translation, rotation, torsion) - Circular loss for torsion angles (wrap to [-pi, pi]) - Optional coordinate reconstruction loss (end-to-end supervision) Sampling: - sample_trajectory_torsion: Euler ODE in decomposed space - Apply order: Torsion → Translation → Rotation (DiffDock convention) - Rodrigues rotation for both torsion and rigid body rotation Config: - configs/train_torsion.yaml with torsion-specific loss weights Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Training: - train_step dispatches to _train_step_torsion or _train_step_cartesian based on config output_mode - Torsion mode: samples timestep, interpolates coordinates, predicts decomposed velocities, computes component-wise loss - WandB logs torsion-specific metrics (loss_trans, loss_rot, loss_tor) - Progress bar shows component losses in torsion mode Validation: - Uses sample_trajectory_torsion for ODE integration in torsion mode - Falls back to cartesian Euler/RK4 for cartesian mode - Uses generate_timestep_schedule from sampling.py (DRY) Data: - collate_flowfix_batch now handles torsion_data collation - _collate_torsion_data: concatenates variable-length rotatable bonds, offsets atom indices per molecule, expands mask_rotate to batch size Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Revert all torsion changes from original files (flowmatching.py,
dataset.py, losses.py, sampling.py, train.py, model_builder.py)
and create clean separate files instead:
New files:
- src/models/flowmatching_torsion.py
ProteinLigandFlowMatchingTorsion inherits base model,
adds translation/rotation/torsion output heads
- src/data/dataset_torsion.py
FlowFixTorsionDataset inherits FlowFixDataset,
collate_torsion_batch handles variable-length rotatable bonds
- src/utils/losses_torsion.py
compute_se3_torsion_loss with circular torsion loss,
coord reconstruction via Rodrigues rotation
- src/utils/sampling_torsion.py
sample_trajectory_torsion for Euler ODE in decomposed space
- train_torsion.py
FlowFixTorsionTrainer - standalone torsion training script
Original files unchanged - cartesian mode works exactly as before.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ison - Remove 29 dead files: unused utils (train.py, loss_utils.py, metrics.py, plot.py), legacy configs (overfit_*, rectified_flow*), debug scripts, and SLURM scripts - Remove unused joint model files (config, SLURM script, test) - Reorganize inference results into output/cartesian_v4_baseline/ - Restructure reports/ into cartesian/ and torsion/ subdirectories - Add torsion architecture doc and results placeholder - Update all references from train_joint.yaml to train.yaml - Fix mermaid diagrams: use TB layout for complex diagrams Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.