Feat: Freeze DOFs Capability

[eduardolneto]

Summary

Adds JAX-compatible degree-of-freedom freezing to ESSOS optimizations without changing the existing optimizer-vector size or pytree structure.

Frozen DOFs remain at their initial optimizer-space values during loss/constraint evaluation, and their gradient entries are zero. This works for both normal custom_loss / composite_loss optimization and augmented-Lagrangian constraints.

New capabilities

# Generic escape hatch: freeze any flat Boolean mask or matching pytree.
loss.freeze_dofs(mask)

# Remove all frozen selections.
loss.unfreeze_dofs()

Freeze base-coil currents

L_total.freeze_current("field", coil=0)
C_total_constraint.freeze_current("field", coil=0)

Freezes the normalized current DOF for base coil 0. Since currents_scale remains fixed, the corresponding physical current remains fixed too.

Freeze complete coil geometry

L_total.freeze_coil("field", coil=2)

Freezes every curve coefficient of base coil 2, while leaving its current and other coils free.

Freeze selected curve coefficients

L_total.freeze_curve_dofs(
    "field",
    coil=1,
    coordinates="z",
    modes=1,
)

Freezes the stored z curve coefficient at mode index 1 for base coil 1.

Freeze surface Fourier modes

L_total.freeze_surface_modes(
    "surface",
    rc=[(1, 0)],
    zs=[(1, 0)],
)

Freezes the SurfaceRZFourier radial and vertical coefficients for physical (m, n) mode (1, 0)—equivalent to RBC(0,1) and ZBS(0,1) in VMEC notation.

Examples added

• examples/coil_optimization/optimize_coils_frozen_current.py
• examples/coil_optimization/optimize_coils_frozen_current_augmented_lagrangian.py
• examples/coil_optimization/optimize_coils_vmec_surface_frozen_current.py
• examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_frozen_current.py
• examples/simple_examples/freeze_surface_modes.py

The VMEC-surface examples use the new user-facing form:

L_total.freeze_current("field", coil=FROZEN_CURRENT_INDEX)

and:

C_total_constraint.freeze_current("field", coil=FROZEN_CURRENT_INDEX)

Files changed

• essos/frozen_dofs.py — shared DOF-freezing and selector implementation.
• essos/losses.py — integrates freezing with base_loss, custom_loss, and composite_loss.
• essos/augmented_lagrangian.py — integrates freezing with selective and composite ALM constraints.
• tests/test_multiobjectives.py — tests generic, current, coil, curve, and surface selection.
• tests/test_augmented_lagrangian.py — tests ALM frozen-DOF projection.
• Added example files listed above.

Validation

• Python syntax compilation and whitespace/diff checks pass.
• Focused runtime tests should be run in a JAX-enabled environment:

pytest tests/test_multiobjectives.py tests/test_augmented_lagrangian.py -q

[codecov]

Codecov Report

❌ Patch coverage is 76.00000% with 36 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
essos/frozen_dofs.py	76.66%	11 Missing and 10 partials ⚠️
essos/losses.py	68.57%	9 Missing and 2 partials ⚠️
essos/augmented_lagrangian.py	84.00%	2 Missing and 2 partials ⚠️

Files with missing lines	Coverage Δ
essos/augmented_lagrangian.py	63.31% <84.00%> (+4.16%)	⬆️
essos/losses.py	83.33% <68.57%> (+10.79%)	⬆️
essos/frozen_dofs.py	76.66% <76.66%> (ø)

... and 3 files with indirect coverage changes