Skip to content

Add CPU unit tests for models/rf3 (losses, layers, metrics, sampler)#314

Closed
lyskov-ai wants to merge 6 commits into
RosettaCommons:productionfrom
lyskov-ai:0032-rf3-test-metric-orchestration
Closed

Add CPU unit tests for models/rf3 (losses, layers, metrics, sampler)#314
lyskov-ai wants to merge 6 commits into
RosettaCommons:productionfrom
lyskov-ai:0032-rf3-test-metric-orchestration

Conversation

@lyskov-ai

Copy link
Copy Markdown
Contributor

Adds fixture-backed CPU unit tests for models/rf3, with no source changes. The tests pin the model's pure numeric and behavioural contracts using tiny committed fixtures (no GPU, downloaded checkpoints, or cluster data), so they run in the standard CI environment.

What's covered:

  • Loss functions & gradients — distogram and smoothed-LDDT losses; the closed-form dihedral-MSE gradient checked against autograd; chiral-gradient routing/accumulation.
  • Symmetry resolution geometry — the Kabsch alignment (recovers a known rotation+translation, corrects reflections to a proper rotation) and the residue-automorphism selection.
  • Model layers — the layer_utils shape helpers, the MLFF conformer embedding, triangle attention/multiplication (vanilla CPU path), and the outer-product mean.
  • Metrics — chiral metric orchestration (polymer/non-polymer routing, unresolved-atom handling) and the predicted-TM / interface-TM scoring.
  • Diffusion sampler & utilities — the AF3 noise schedule and initial-structure construction, plus assorted pure helpers (RBF encoding, LeCun init, loss-dict reshaping).

Why: building out test coverage of models/rf3 ahead of refactoring work, as a regression net that pins current behaviour.

lyskov and others added 6 commits June 9, 2026 17:49
…_module

Cover the pure numeric helpers in rf3.loss.loss and rf3.util_module:

- calc_ddihedralmse_dxyz: the hand-derived closed-form dihedral-loss
  gradient, pinned against torch.autograd of a mirrored forward, plus
  zero-at-truth and leading-dim shape preservation.
- calc_chiral_grads_flat_impl: empty-centres, scatter routing onto only
  the centre atoms, index_add_ accumulation for shared atoms, and the
  no_grad_on_chiral_center flag.
- rbf: Gaussian distance-encoding (feature dim, exact values, [0,1]).
- init_lecun_normal: module returned, Parameter weight, +/-2*stddev
  bound, and Lecun std sqrt(scale/fan_in).

Test-only; both modules are already mypy-clean. Tests run in float32
(the production coordinate dtype).

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
…duct layer

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
…helpers

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
@lyskov-ai lyskov-ai requested a review from woodsh17 June 9, 2026 21:45
@lyskov-ai

Copy link
Copy Markdown
Contributor Author

Closing as redundant: this PR's cumulative diff is fully contained in #321, which was squash-merged into production (commit e412591). All of its changes are now in production.

@lyskov-ai lyskov-ai closed this Jun 16, 2026
@lyskov-ai lyskov-ai deleted the 0032-rf3-test-metric-orchestration branch June 16, 2026 19:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants