JaxCurrent actually not quite working yet. Made some little fixes but…

… need to fully fix.

examples/2_Intermediate/stage_two_optimization_dipole_arrays_QA.py

@@ -29,7 +29,7 @@
 # filename = TEST_DIR / input_name
 
 # Directory for output
-OUT_DIR = "./dipole_array_optimization_QA/"
+OUT_DIR = "./dipole_array_optimization_QA_debug/"
 if os.path.exists(OUT_DIR):
     shutil.rmtree(OUT_DIR)
 os.makedirs(OUT_DIR, exist_ok=True)
@@ -43,7 +43,7 @@
 nphi = 64
 ntheta = 64
 poff = 2.0
-coff = 1.0
+coff = 2.0
 s = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi, ntheta=ntheta)
 s_inner = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi * 4, ntheta=ntheta * 4)
 s_outer = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi * 4, ntheta=ntheta * 4)

src/simsopt/field/coil.py

@@ -1,10 +1,13 @@
 from math import pi
 import numpy as np
+from jax import vjp, jacfwd, jvp
+import jax.numpy as jnp
 
 from simsopt._core.optimizable import Optimizable
 from simsopt._core.derivative import Derivative
 from simsopt.geo.curvexyzfourier import CurveXYZFourier
 from simsopt.geo.curve import RotatedCurve
+from simsopt.geo.jit import jit
 import simsoptpp as sopp
 
 
@@ -142,6 +145,7 @@ def current(self):
         self.current_jax = jit(lambda dofs: self.current_pure(dofs))
         self.dcurrent_by_dcoeff_jax = jit(jacfwd(self.current_jax))
         self.dcurrent_by_dcoeff_vjp_jax = jit(lambda x, v: vjp(self.current_jax, x)[1](v)[0])
+        self.vjp = dcurrent_by_dcoeff_vjp_jax
 
     def set_dofs(self, dofs):
         self.local_x = dofs