perf: speed up unit tests 63% by removing JAX from triangle tests#252
Merged
perf: speed up unit tests 63% by removing JAX from triangle tests#252
Conversation
…ucing test sizes Switched all triangle unit tests from JAX classes (ArrayTriangles, CoordinateArrayTriangles) to numpy equivalents (ArrayTrianglesNp, CoordinateArrayTrianglesNp), eliminating ~10s of JAX JIT compilation overhead. Removed pure-JAX test files covered by integration tests. Reduced PSF weighted noise test mask from 51x51 to 21x21. Total suite time: 33.9s → 12.5s (63% faster, 722 tests). Co-Authored-By: Claude Opus 4.6 (1M context) <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.
Summary
Profiled the full PyAutoArray test suite and applied targeted optimizations to reduce total test time from 33.9s to 12.5s (63% faster). The main bottleneck was JAX JIT compilation overhead in triangle unit tests — each unique array shape triggered a fresh compilation costing 0.5-1s.
Key changes:
ArrayTrianglesNp,CoordinateArrayTrianglesNp)test_vertex_coordinates.py(contained only print statements, no assertions)API Changes
None — internal test changes only.
Test Plan
Full API Changes (for automation & release notes)
No public API changes. All modifications are confined to test files:
test_autoarray/structures/triangles/conftest.pytest_autoarray/structures/triangles/test_coordinate.pytest_autoarray/structures/triangles/test_area.pytest_autoarray/structures/triangles/test_extended_source.pytest_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.pyDeleted files:
test_autoarray/structures/triangles/test_jax.py(covered by integration tests)test_autoarray/structures/triangles/test_nan_triangles.py(JAX-specific NaN handling)test_autoarray/structures/triangles/test_vertex_coordinates.py(no assertions)🤖 Generated with Claude Code