Simsopt macromag: GPMOmr

[armulrich]

Summary

This PR adds and documents the GPMOmr workflow (GPMO with macromagnetic refinement) and significantly improves reproducibility and post-processing for the permanent-magnet MUSE study. The changes align the codebase with the methodology and results presented in:

A. Ulrich, M. Haberle, and A. A. Kaptanoglu,
Permanent magnet optimization of stellarators with coupling from finite permeability and demagnetization effects,
arXiv:2512.14997 (2025)
https://arxiv.org/abs/2512.14997

Key Improvements

• Macromagnetic refinement in the greedy loop
  Adds a device-scale MacroMag refinement path that accounts for finite-permeability and demagnetization effects. This includes a Python implementation of GPMO together with a Numba-accelerated implementation of the demagnetization tensors. Refinement is applied in a winner-only fashion via mm_refine_every, preserving the efficiency of the greedy loop while incorporating realistic magnetic coupling.

• Standardized run artifacts
  Introduces a consistent artifact format across GPMO and GPMOmr:

  • run_<id>.yaml for configuration and metadata
  • runhistory_<id>.csv for optimization histories
    This replaces legacy *.txt histories and enables uniform plotting and comparison across runs.

Overall, this PR makes the GPMO/GPMOmr workflows more realistic, reproducible, and directly traceable to the results reported in arXiv:2512.14997.

[codecov]

Codecov Report

❌ Patch coverage is 99.44994% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.20%. Comparing base (f8c9be3) to head (2a0d903).

Files with missing lines	Patch %	Lines
src/simsopt/solve/permanent_magnet_optimization.py	99.29%	4 Missing ⚠️
.../simsopt/util/permanent_magnet_helper_functions.py	99.58%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #587      +/-   ##
==========================================
+ Coverage   90.14%   90.20%   +0.05%     
==========================================
  Files          84       85       +1     
  Lines       17880    19757    +1877     
==========================================
+ Hits        16118    17821    +1703     
- Misses       1762     1936     +174     

Flag	Coverage Δ
unittests	90.20% <99.44%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[akaptano]

@armulrich

Made some improvements and refactoring changes. Sorry everyone but some significant PM refactoring got appended here, since it really helps and is past time. Most significantly:

1. Moved your finitemu calculations for MUSE into its own standalone example file, and removing your PM4Stell modifications, since these can go in your paper branch anyways. The original permanent magnet example files are now unchanged (except that they use the new vtk functionality to output cuboids in paraview and they use the new dataclasses -- see below).

2. The PM optimization, macromagnetics, etc. functions all took like 20 arguments, most of which are the same, so I have simplified all of this by using dataclasses like GPMOBacktrackParams. Check out the functionality in the code.

3. Updated a lot of the docstrings and type hinting with the power of Cursor.

4. Drastically improved the speed (by just vectorizing everything) of the matrix assembly -- all the code seems to running much faster than before. Try it out. Also, managed to circumvent the C++ memory errors by adding a XLA flag setting in init.py for the jax matrix assembly.

5. Tried to improve the test coverage to get codecov patch > 90%, including testing some functionality like the matrix_free solve and convergence w/ respect to Nadjacent if the direct_solve_sparse is used.

6. Combined the _toVTK functions in DipoleFields to one function, since there was a lot of duplicate code and really just need one function.

7. Drastically reduced duplicate code and refactored everything in macromag.py (now called macromagnetics.py). Total number of lines grew though after I added detailed docstrings with mathematical definitions. There was lots of duplicate code across functions (and same with in permanent_magnet_optimization.py).

Sorry this ballooned but I think its running much faster, has much more documentation, and removed lots of duplicate code. Let me know what you think and please check if I managed to destroy things. The finitemu analysis and optimizations in the example files looks about the same as before, with few percent differences, but I haven't taken the time to verify the Paraview plots and so on still look good. I have also not checked the updated docs -- please verify all the new docstrings (and their mathematical formulas) are building/visualizing well.

[armulrich]

@akaptano Thanks! Will take a loook.

[akaptano]

@armulrich Was there a reason to change _dipole_fields_from_symmetries logic here? I can't tell if this function actually changed at all

[armulrich]

@akaptano I think if I remember correctly the only real change here was replacing the old self.dipole_grid = dipole_grid with self.dipole_grid_unique = np.ascontiguousarray(dipole_grid), since self.dipole_grid gets overwritten later by the symmetry-expanded full grid anyway; we needed to keep the original half-period grid around for the _toVTK(..., dx, dy, dz) / _geometry_unique path. ALso needed to use ascontiguousarray here for the VTK export code. Other than that, I don’t think the symmetry logic changed here, and the long toroidal lines were just me reformatting them for readability.... Let me know if I should change that