Add type hints + unit tests across foundry.utils and rf3#321
Merged
woodsh17 merged 12 commits intoJun 10, 2026
Merged
Conversation
…_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>
…d unit tests Add foundry.utils.alignment and foundry.utils.weights to the per-module strict-mypy override (disallow_untyped_defs + check_untyped_defs) and annotate their untyped definitions: - alignment.py: weighted_rigid_align / get_rmsd / the dead superimpose stub. get_rmsd's torch.sum(..., axis=) -> dim= (torch numpy-compat alias; the stubs only declare dim, behaviour-identical). - weights.py: -> None on the two __post_init__ methods. Annotating weighted_rigid_align's params honestly as torch.Tensor surfaced a real arg-type mismatch in rf3's dump_trajectories (param loosely typed list[torch.Tensor | np.ndarray], body tensor-only); resolved with documented cast(torch.Tensor, ...) at the alignment call site. Extend tests/test_weight_loading.py with 3 tests for documented but untested behaviour: the '?' and '[...]' glob wildcards plus literal-dot escaping, and the COPY-policy missing-parameter fallback chain. Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
Add foundry.utils.rotation_augmentation to the per-module strict mypy override (disallow_untyped_defs + check_untyped_defs) and annotate its four tensor-geometry functions. Annotation-only; already covered by tests/test_rotation_augmentation.py. Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>
woodsh17
approved these changes
Jun 10, 2026
This was referenced Jun 16, 2026
Closed
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
3 participants
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.
Brings the shared
foundry.utilslayer and severalrf3modules under stricter type checking and fills unit-test gaps. No behaviour changes — the source edits are type annotations and test additions only, plus two small documented exceptions noted below.Type hints (shared layer)
Annotated and opted into strict mypy checking (untyped defs disallowed, function bodies checked) for six
foundry.utilsmodules:rigid,components,torch,alignment,weights, androtation_augmentation. Any new untyped function in these modules now fails the type-check gate.Unit tests
tests/test_rigid.pyandtests/test_components.py, plus extensions totests/test_torch_utils.pyandtests/test_weight_loading.py, pinning the geometry, contig-parsing, tensor, and weight-loading helpers.Other
rf3.data.paired_msa: moved its type-check suppression from the centralpyproject.tomlreject-list into a file-level directive in the module itself. The module is currently broken against the installed atomworks (it subclasses a symbol that became a factory function) and needs a separate refactor; the in-file directive keeps the suppression visible next to the broken code.rf3/utils/io.py: one documentedcastwhere annotatingweighted_rigid_alignsurfaced a loosely-typed caller (dump_trajectories, whose parameter is wider than its tensor-only body). Behaviour is preserved; narrowing the caller is filed as follow-up.