Accelerated ts.optimize by batching Frechet Cell Filter

[falletta]

Update Summary

---

1. torch_sim/math.py

Removed unbatched/legacy code:

• expm_frechet_block_enlarge (helper function for block enlargement method)
• _diff_pade3, _diff_pade5, _diff_pade7, _diff_pade9 (Padé approximation helpers)
• expm_frechet_algo_64 (original algorithm implementation)
• matrix_exp (custom matrix exponential function)
• vec, expm_frechet_kronform (Kronecker form helpers)
• expm_cond (condition number estimation)
• class expm (autograd Function class)
• _is_valid_matrix, _determine_eigenvalue_case (unbatched helpers)

Refactored expm_frechet:

• Now optimized specifically for batched 3x3 matrices (common case for cell operations)
• Handles both (B, 3, 3) batched input and (3, 3) single matrix input (auto-adds batch dim)
• Removed method parameter (was SPS"or blockEnlarge)
• Inlined the algorithm directly instead of calling helper functions

Refactored matrix_log_33:

• Added _ensure_batched, _determine_matrix_log_cases, _process_matrix_log_case helpers
• Made the matrix log computation work in batched mode

---

2. torch_sim/optimizers/cell_filters.py

Vectorized compute_cell_forces:

• Before: used nested loops over systems and directions (9 iterations per system)
• After: uses batched matrix operations across all systems and all 9 directions simultaneously
• Key optimization: expm_frechet(A_batch, E_batch) is now called once with all n_systems * 9 matrices batched together

---

3. tests/test_math.py

Refactored tests:

• TestExpmFrechet: test_expm_frechet, test_small_norm_expm_frechet, test_fuzz
• TestExpmFrechetTorch: test_expm_frechet, test_fuzz

All updated to use 3x3 matrices and simplified by removing method parameter testing. Fuzz tests streamlined with fewer iterations.

Removed tests:

• test_problematic_matrix, test_medium_matrix (both numpy and torch versions)
• TestExpmFrechetTorchGrad class

Tests for comparing computation methods and large matrix performance no longer apply to the 3x3-specialized implementation.

Added tests:

• TestExpmFrechet.test_large_norm_matrices - Tests scaling behavior for larger norm matrices
• TestLogM33.test_batched_positive_definite - Tests batched matrix logarithm with round-trip verification
• TestFrechetCellFilterIntegration - Integration tests for the cell filter pipeline
• test_wrap_positions_* - Tests for the new wrap_positions property

---

Results

The figure below shows the speedup achieved for 10-step atomic relaxation. The test is performed for a 8-atom cubic supercell of MgO using the mace-mpa model. Prior results are shown in blue, while new results are shown in red. The speedup is calculated as speedup (%) = (baseline_time / current_time − 1) × 100. We observe a speedup up to 564% for large batches.

[orionarcher]

Could we get some tests verifying identical numerical behavior of the old and new versions? Can be deleted before merging when we get rid of the unbatched version.

[falletta]

@orionarcher I added test_math_frechet.py to compare the batched code against the SciPy implementations. Please have a look—happy to revise it as needed, and if it looks good, we can include it directly in test_math.py.

[orionarcher]

Why support both batched and unbatched versions?

[falletta]

I now removed all unbatched code

[orionarcher]

Are we testing the batched or unbatched versions here?

[falletta]

I incorporated the batched tests only in test_math.py

[thomasloux]

Haven't looked carefully on the implementation but I would potentially support to have 2 separate functions for batched (B, 3, 3) and unbatched (3,3) algorithms. This would also prevent graph breaks in the future, be easier to read, and in practice a state.cell is always (B, 3, 3), potentially with B=1. So we would always use the batched version anyway.

[falletta]

@orionarcher I removed all unbatched and unused code while preserving the new performance speedups. Please see the PR description for a detailed list of changes.

@thomasloux It’s indeed a good point, but for now it’s probably better to keep things clean and stick to the batched implementation only. By keeping only the batched implementation, we can remove quite a few lines of dead code.

[orionarcher]

Again not checking the integration math, but noting that TestExpmFrechetTorch.test_expm_frechet and test_fuzz both discard the Frechet derivative result (observed_expm, _ = ...). Since the Frechet derivative is the core output used by the cell filter, these should verify it too. Can you add the Frechet assertions back?