Fixing all Examples

debug_from_simsopt_import.py

@@ -0,0 +1,96 @@
+import os
+import jax
+import jax.numpy as jnp
+
+from essos.coils import Coils
+from essos.fields import BiotSavart, Vmec
+from essos.constants import (
+    ALPHA_PARTICLE_MASS,
+    ALPHA_PARTICLE_CHARGE,
+    FUSION_ALPHA_PARTICLE_ENERGY,
+)
+from essos.dynamics import GuidingCenter, Particles
+from essos.objective_functions import normB_axis
+
+
+def describe_case(label, coils, vmec_surface, initial_xyz, initial_vpar):
+    print(f"\n=== {label} ===")
+    print("n_base_curves:", coils.curves.n_base_curves)
+    print("n_segments:", coils.n_segments)
+    print("nfp:", coils.nfp, "stellsym:", coils.stellsym)
+    print("dofs_curves shape:", coils.dofs_curves.shape)
+    print("dofs_currents_raw shape:", coils.dofs_currents_raw.shape)
+    print("currents raw min/max:", float(jnp.min(coils.dofs_currents_raw)), float(jnp.max(coils.dofs_currents_raw)))
+    print("currents scale:", float(coils.currents_scale))
+    print("currents normalized min/max:", float(jnp.min(coils.dofs_currents)), float(jnp.max(coils.dofs_currents)))
+    print("gamma shape:", coils.gamma.shape)
+    print("gamma finite:", bool(jnp.all(jnp.isfinite(coils.gamma))))
+    print("gamma_dash finite:", bool(jnp.all(jnp.isfinite(coils.gamma_dash))))
+    print("gamma_dashdash finite:", bool(jnp.all(jnp.isfinite(coils.gamma_dashdash))))
+
+    field = BiotSavart(coils)
+    B_axis = normB_axis(field, npoints=200)
+    print("axis B mean before renorm:", float(jnp.mean(B_axis)))
+    coils.dofs_currents = coils.dofs_currents * 5.7 / jnp.mean(B_axis)
+    field = BiotSavart(coils)
+    print("axis B mean after renorm:", float(jnp.mean(normB_axis(field, npoints=200))))
+
+    print("surface gamma finite:", bool(jnp.all(jnp.isfinite(vmec_surface.gamma))))
+    print("initial xyz:", initial_xyz)
+    print("field.B(initial):", field.B(initial_xyz))
+    print("field.AbsB(initial):", float(field.AbsB(initial_xyz)))
+    print("field.dAbsB_by_dX(initial):", field.dAbsB_by_dX(initial_xyz))
+    print("field.curl_b(initial):", field.curl_b(initial_xyz))
+    print("field.kappa(initial):", field.kappa(initial_xyz))
+
+    finite_checks = {
+        "B": jnp.all(jnp.isfinite(field.B(initial_xyz))),
+        "AbsB": jnp.isfinite(field.AbsB(initial_xyz)),
+        "dAbsB_by_dX": jnp.all(jnp.isfinite(field.dAbsB_by_dX(initial_xyz))),
+        "curl_b": jnp.all(jnp.isfinite(field.curl_b(initial_xyz))),
+        "kappa": jnp.all(jnp.isfinite(field.kappa(initial_xyz))),
+    }
+    print("finite checks:", {k: bool(v) for k, v in finite_checks.items()})
+
+    particles = Particles(
+        initial_xyz=jnp.expand_dims(initial_xyz, 0),
+        mass=ALPHA_PARTICLE_MASS,
+        charge=ALPHA_PARTICLE_CHARGE,
+        energy=FUSION_ALPHA_PARTICLE_ENERGY,
+    )
+    initial_condition = jnp.array([initial_xyz[0], initial_xyz[1], initial_xyz[2], initial_vpar])
+    rhs = GuidingCenter(0.0, initial_condition, (field, particles, type("E", (), {"E_covariant": lambda self, x: jnp.zeros(3)})()))
+    print("GuidingCenter rhs:", rhs)
+    print("GuidingCenter rhs finite:", bool(jnp.all(jnp.isfinite(rhs))))
+
+
+def main():
+    print(jax.devices())
+
+    simsopt_json = os.path.join("examples", "input_files", "QH_simple_scaled.json")
+    wout_file = os.path.join("examples", "input_files", "wout_QH_simple_scaled.nc")
+    vmec = Vmec(wout_file)
+
+    R0 = 17.0
+    initial_xyz = jnp.array([R0, 0.0, 0.0])
+    initial_vpar = 0.0
+
+    coils_true = Coils.from_simsopt(simsopt_json, nfp=4, stellsym=True)
+    describe_case("from_simsopt stellsym=True", coils_true, vmec.surface, initial_xyz, initial_vpar)
+
+    coils_false = Coils.from_simsopt(simsopt_json, nfp=4, stellsym=False)
+    describe_case("from_simsopt stellsym=False", coils_false, vmec.surface, initial_xyz, initial_vpar)
+
+    from simsopt import load
+
+    bs = load(simsopt_json)
+    base_coils = bs.coils[:6]
+    print("\n=== direct simsopt base coils ===")
+    print("n base coils:", len(base_coils))
+    print("current values:", [float(c.current.get_value()) for c in base_coils])
+    print("curve gamma shape:", jnp.asarray(base_coils[0].curve.gamma()).shape)
+    print("curve gamma first point:", jnp.asarray(base_coils[0].curve.gamma())[0])
+
+
+if __name__ == "__main__":
+    main()

essos/augmented_lagrangian.py

@@ -21,6 +21,11 @@ class LagrangeMultiplier(NamedTuple):
 
 
 def _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad):
+    if out is None:
+        raise ValueError(
+            "Constraint function returned None during initialization. "
+            "Constraints used with eq()/ineq() must return a scalar or array."
+        )
     z = jnp.zeros_like(out)
     return LagrangeMultiplier(
         value=multiplier + z,

essos/coil_perturbation.py

@@ -238,7 +238,7 @@ def perturb_curves(curves, sampler:GaussianSampler, key=None, perturbation_type=
             or "statistical"/"statistic" to perturb every expanded coil independently.
 
     Returns:
-        A new perturbed object of the same family as the input.
+        A new perturbed object of the same type as the input.
     """
     if perturbation_type == "systematic":
         if isinstance(curves, DiscretizedCoils):
@@ -248,19 +248,21 @@ def perturb_curves(curves, sampler:GaussianSampler, key=None, perturbation_type=
                 currents=curves.dofs_currents_raw,
                 nfp=curves.nfp,
                 stellsym=curves.stellsym,
+                currents_scale=curves.currents_scale,
+                scale_fixed=curves.scale_fixed,
             )
 
         if isinstance(curves, Coils):
             perturbation = _draw_curve_perturbation(sampler, key, curves.curves.n_base_curves)
-            base_curves = _make_curves_like(curves.curves, curves.dofs_curves, nfp=1, stellsym=False)
+            base_curves = _make_curves_like(curves.curves, curves.curves._dofs, nfp=1, stellsym=False)
             perturbed_base_gamma = base_curves.gamma + perturbation
             dofs_new, _ = fit_dofs_from_coils(perturbed_base_gamma, curves.order, curves.n_segments, assume_uniform=True)
             new_curves = _make_curves_like(curves.curves, dofs_new, nfp=curves.nfp, stellsym=curves.stellsym)
-            return Coils(curves=new_curves, currents=curves.dofs_currents_raw)
+            return Coils(curves=new_curves, currents=curves.dofs_currents_raw, currents_scale=curves.currents_scale)
 
         if isinstance(curves, Curves):
             perturbation = _draw_curve_perturbation(sampler, key, curves.n_base_curves)
-            base_curves = _make_curves_like(curves, curves.dofs, nfp=1, stellsym=False)
+            base_curves = _make_curves_like(curves, curves._dofs, nfp=1, stellsym=False)
             perturbed_base_gamma = base_curves.gamma + perturbation
             dofs_new, _ = fit_dofs_from_coils(perturbed_base_gamma, curves.order, curves.n_segments, assume_uniform=True)
             return _make_curves_like(curves, dofs_new, nfp=curves.nfp, stellsym=curves.stellsym)
@@ -270,12 +272,19 @@ def perturb_curves(curves, sampler:GaussianSampler, key=None, perturbation_type=
         gamma_perturbed = curves.gamma + perturbation
 
         if isinstance(curves, DiscretizedCoils):
-            return DiscretizedCoils(gamma_perturbed, currents=curves.currents, nfp=1, stellsym=False)
+            return DiscretizedCoils(
+                gamma_perturbed,
+                currents=curves.currents,
+                nfp=1,
+                stellsym=False,
+                currents_scale=curves.currents_scale,
+                scale_fixed=curves.scale_fixed,
+            )
 
         if isinstance(curves, Coils):
             dofs_new, _ = fit_dofs_from_coils(gamma_perturbed, curves.order, curves.n_segments, assume_uniform=True)
             new_curves = _make_curves_like(curves.curves, dofs_new, nfp=1, stellsym=False)
-            return Coils(curves=new_curves, currents=curves.currents)
+            return Coils(curves=new_curves, currents=curves.currents, currents_scale=curves.currents_scale)
 
         if isinstance(curves, Curves):
             dofs_new, _ = fit_dofs_from_coils(gamma_perturbed, curves.order, curves.n_segments, assume_uniform=True)