Skip to content

Add CPU unit tests and tighten type-checking (foundry rigid.py + rf3 helpers)#317

Closed
lyskov-ai wants to merge 8 commits into
RosettaCommons:productionfrom
lyskov-ai:0034-foundry-rigid-mypy-strict-and-tests
Closed

Add CPU unit tests and tighten type-checking (foundry rigid.py + rf3 helpers)#317
lyskov-ai wants to merge 8 commits into
RosettaCommons:productionfrom
lyskov-ai:0034-foundry-rigid-mypy-strict-and-tests

Conversation

@lyskov-ai

Copy link
Copy Markdown
Contributor

Adds CPU unit tests and tightens type-checking across the shared foundry layer and the rf3 model. No runtime behaviour changes — the source edits are type annotations plus a relocation of one type-check suppression; everything else is new tests.

foundry/utils/rigid.py — type annotations + tests

  • Annotated the previously-untyped Rotation/Rigid methods and the module-level quaternion/matrix helpers, and enabled strict type-checking for the module (disallow_untyped_defs + check_untyped_defs). Annotation-only. Also widened the identity_*/identity shape parameters from Tuple[int] to Tuple[int, ...] (they accept arbitrary-rank shapes).
  • New tests/test_rigid.py (33 tests) pinning the geometry: quaternion↔rotation-matrix round-trips, composition/inversion, point application, the 4×4 and 7-vector encodings, and from_3_points frame construction — all against independent references (Rodrigues, einsum, QR-built rotations).

rf3 — CPU unit tests

New fixture-backed CPU tests (no cluster data / GPU / checkpoints) for pure rf3 logic: loss gradients and utilities, model-layer shape helpers (layer_utils, mlff, attention, outer-product), the diffusion sampler's noise schedule and initial structure, the af3_losses symmetry-resolution geometry, and the chiral / iPTM metric orchestration.

rf3.data.paired_msa — type-check hygiene

Moved the module's mypy suppression out of the central pyproject.toml reject-list into a file-level # mypy: ignore-errors directive, so the suppression is visible at the broken code and deleting one line re-enables checking. (The module is broken against the installed atomworks and needs a data-loader refactor validated on cluster data — tracked separately.)

lyskov and others added 8 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>
Replace the central pyproject ignore_errors override for rf3.data.paired_msa with a file-level '# mypy: ignore-errors' directive in the module itself. The module is broken against the installed atomworks (subclasses a now-function) and needs a PandasDataset-API refactor to clear honestly; keeping it in mypy's files scope with an in-file directive makes the suppression visible at the point of breakage and re-enables checking the moment the directive is removed.

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

Bring src/foundry/utils/rigid.py under a per-module disallow_untyped_defs +
check_untyped_defs override (the first Track 1 direction-(b) strictness slice) by
annotating its ~47 previously-untyped Rotation/Rigid methods and quaternion/matrix
helpers. Annotation-only, no behaviour change; self-references use string forward-refs
and the identity helpers' shape params are widened to Tuple[int, ...].

Add tests/test_rigid.py: 33 fixture-backed CPU unit tests pinning the pure geometry
(quat<->matrix round-trips, compose/invert, apply, the 4x4/7-vector encodings, and
from_3_points) against independent references.

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
@lyskov-ai lyskov-ai requested a review from woodsh17 June 10, 2026 01:25
@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 0034-foundry-rigid-mypy-strict-and-tests 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