Skip to content

Add CPU unit tests and tighten type-checking (foundry utils + rf3 helpers)#319

Closed
lyskov-ai wants to merge 10 commits into
RosettaCommons:productionfrom
lyskov-ai:0036-foundry-torch-utils-mypy-strict-and-tests
Closed

Add CPU unit tests and tighten type-checking (foundry utils + rf3 helpers)#319
lyskov-ai wants to merge 10 commits into
RosettaCommons:productionfrom
lyskov-ai:0036-foundry-torch-utils-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 two behaviour-preserving refactors and one relocated type-check suppression); everything else is new tests.

foundry shared layer — type annotations + tests

  • utils/rigid.py — annotated the Rotation/Rigid methods and quaternion/matrix helpers and enabled strict type-checking for the module. New tests/test_rigid.py (33 tests): quaternion↔rotation-matrix round-trips, composition/inversion, point application, the 4×4 and 7-vector encodings, and from_3_points.
  • utils/components.py — annotated the contig/component parsers and atom-name mask getters and enabled strict type-checking (plus a strint collapse in split_contig, a tuple-typed fixed_parts, and widening get_name_mask's query_names to str | list[str] to match its isinstance branch). New tests/test_components.py (25 tests) pinning the parsing grammar and get_name_mask.
  • utils/torch.py — annotated the remaining untyped helpers (map_to, the tracer-warning contextmanager, assert_shape, Timer/Timers) and enabled strict type-checking; the two warnings.filters mutations carry a documented # type: ignore[attr-defined] (typeshed types it immutable, but it is a list at runtime). Extended tests/test_torch_utils.py (+12 tests): scatter_mean, assert_shape/assert_same_shape, device_of, and Timer/Timers.

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.

lyskov and others added 10 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>
…g tests

Bring src/foundry/utils/components.py under the per-module disallow_untyped_defs +
check_untyped_defs override by annotating the contig parsers and mask getters.
Annotation-only apart from two behaviour-preserving refactors (collapse a str->int
variable reuse in split_contig; store fixed_parts as tuples to avoid object-typed
heterogeneous lists). Also widen get_name_mask's query_names to str | list[str] to
match the isinstance branch and docstring.

Add tests/test_components.py: 25 unit tests pinning the contig/component parsing
grammar (split_contig, extract_pn_unit_info, get_motif_components_and_breaks,
get_design_pattern_with_constraints) and the get_name_mask atom selector.

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
Bring src/foundry/utils/torch.py under the per-module disallow_untyped_defs +
check_untyped_defs override by annotating map_to's kwargs, the tracer-warning
contextmanager, assert_shape, and the Timer/Timers methods. Annotation-only; the
two warnings.filters mutations carry a documented type: ignore[attr-defined]
(typeshed types it immutable, but it is a list at runtime).

Extend tests/test_torch_utils.py with 12 tests for the untested pure helpers:
scatter_mean, assert_shape/assert_same_shape, device_of, and Timer/Timers.

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
@lyskov-ai lyskov-ai requested a review from woodsh17 June 10, 2026 02:57
@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 0036-foundry-torch-utils-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