Now optimizing surfaces, near-axis and coils altogether

Author: rogeriojorge

essos/fields.py

@@ -3,7 +3,7 @@
 import jax.numpy as jnp
 from functools import partial
 from jax import jit, jacfwd, grad, vmap, tree_util, lax
-from essos.surfaces import SurfaceRZFourier
+from essos.surfaces import SurfaceRZFourier, BdotN_over_B
 from essos.plot import fix_matplotlib_3d
 from essos.util import newton
 
@@ -620,7 +620,7 @@ def compute_RZ(m, n):
     def B_mag(self, r, theta, phi):
         return self.B0*(1 + r * self.etabar * jnp.cos(theta - (self.iota - self.iotaN) * phi))
 
-    def plot(self, r=0.1, ntheta=80, nphi=150, ntheta_fourier=20, ax=None, show=True, close=False, axis_equal=True, **kwargs):
+    def plot(self, r=0.1, ntheta=40, nphi=120, ntheta_fourier=20, ax=None, show=True, close=False, axis_equal=True, **kwargs):
         kwargs.setdefault('alpha', 1)
         import matplotlib.pyplot as plt 
         from matplotlib import cm
@@ -657,6 +657,29 @@ def plot(self, r=0.1, ntheta=80, nphi=150, ntheta_fourier=20, ax=None, show=True
         if show:
             plt.show()
 
+    def to_vtk(self, filename, r=0.1, ntheta=40, nphi=120, ntheta_fourier=20, extra_data=None, field=None):
+        try: import numpy as np
+        except ImportError: raise ImportError("The 'numpy' library is required. Please install it using 'pip install numpy'.")
+        try: from pyevtk.hl import gridToVTK
+        except ImportError: raise ImportError("The 'pyevtk' library is required. Please install it using 'pip install pyevtk'.")
+        x, y, z, _ = self.get_boundary(r=r, ntheta=ntheta, nphi=nphi, ntheta_fourier=ntheta_fourier)
+        x = np.array(x.T.reshape((1, nphi, ntheta)).copy())
+        y = np.array(y.T.reshape((1, nphi, ntheta)).copy())
+        z = np.array(z.T.reshape((1, nphi, ntheta)).copy())
+        pointData = {}
+        if field is not None:
+            boundary = np.array([x, y, z]).transpose(1, 2, 3, 0)[0]
+            B_BiotSavart = np.array(vmap(lambda surf: vmap(lambda x: field.AbsB(x))(surf))(boundary)).reshape((1, nphi, ntheta)).copy()
+            pointData["B_BiotSavart"] = B_BiotSavart
+        theta1D = jnp.linspace(0, 2 * jnp.pi, ntheta)
+        phi1D = jnp.linspace(0, 2 * jnp.pi, nphi)
+        phi2D, theta2D = jnp.meshgrid(phi1D, theta1D)
+        Bmag = np.array(self.B_mag(r, theta2D, phi2D)).T.reshape((1, nphi, ntheta)).copy()
+        pointData["B_NearAxis"]=Bmag
+        if extra_data is not None:
+            pointData = {**pointData, **extra_data}
+        gridToVTK(str(filename), x, y, z, pointData=pointData)
+            
 tree_util.register_pytree_node(near_axis,
                                near_axis._tree_flatten,
                                near_axis._tree_unflatten)

essos/objective_functions.py

@@ -42,7 +42,7 @@ def loss_coils_for_nearaxis(x, field_nearaxis, dofs_curves, currents_scale, nfp,
 
     return B_difference_loss+gradB_difference_loss+coil_length_loss+coil_curvature_loss
 
-@partial(jit, static_argnums=(0, 1))
+# @partial(jit, static_argnums=(0, 1))
 def difference_B_gradB_onaxis(nearaxis_field, coils_field):
     Raxis = nearaxis_field.R0
     Zaxis = nearaxis_field.Z0
@@ -102,15 +102,15 @@ def loss_particle_drift(field, particles, maxtime=1e-5, num_steps=300, trace_tol
     angular_drift=(jnp.sum(jnp.diff(angular_drift,axis=1),axis=1))/num_steps
     return jnp.concatenate((jnp.max(radial_drift)*jnp.ravel(2./jnp.pi*jnp.abs(jnp.arctan(radial_drift/(angular_drift+1e-10)))), jnp.ravel(jnp.abs(radial_drift)), jnp.ravel(jnp.abs(vertical_factor))))
 
-@partial(jit, static_argnums=(0))
+# @partial(jit, static_argnums=(0))
 def loss_coil_length(field):
     return jnp.ravel(field.coils.length)
 
-@partial(jit, static_argnums=(0))
+# @partial(jit, static_argnums=(0))
 def loss_coil_curvature(field):
     return jnp.mean(field.coils.curvature, axis=1)
 
-@partial(jit, static_argnums=(0, 1))
+# @partial(jit, static_argnums=(0, 1))
 def loss_normB_axis(field, npoints=15):
     R_axis = jnp.mean(jnp.sqrt(vmap(lambda dofs: dofs[0, 0]**2 + dofs[1, 0]**2)(field.coils.dofs_curves)))
     phi_array = jnp.linspace(0, 2 * jnp.pi, npoints)

essos/optimization.py

@@ -34,7 +34,7 @@ def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e-
     loss_partial = partial(func, dofs_curves=coils.dofs_curves, currents_scale=currents_scale, nfp=nfp, n_segments=n_segments, stellsym=stellsym, **kwargs)
 
     ## Without JAX gradients, using finite differences
-    # result = least_squares(loss_partial, x0=initial_dofs, verbose=2, diff_step=1e-2,
+    # result = least_squares(loss_partial, x0=initial_dofs, verbose=2, diff_step=1e-4,
     #                        ftol=tolerance_optimization, gtol=tolerance_optimization,
     #                        xtol=1e-14, max_nfev=maximum_function_evaluations)
 
@@ -60,13 +60,23 @@ def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e-
             new_field_nearaxis = new_nearaxis_from_x_and_old_nearaxis(result.x[-len(kwargs['field_nearaxis'].x):], kwargs['field_nearaxis'])
             return new_coils, new_field_nearaxis
         elif 'surface_all' in kwargs and len(initial_dofs) == len(coils.x) + len(kwargs['surface_all'].x):
-            dofs_currents = result.x[len_dofs_curves:-len(kwargs['surface_all'].x)]
+            surface_all = kwargs['surface_all']
+            dofs_currents = result.x[len_dofs_curves:-len(surface_all.x)]
             curves = Curves(dofs_curves, n_segments, nfp, stellsym)
             new_coils = Coils(curves=curves, currents=dofs_currents * coils.currents_scale)
-            surface_all = kwargs['surface_all']
             new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta)
             new_surface.dofs = result.x[-len(surface_all.x):]
             return new_coils, new_surface
+        elif 'surface_all' in kwargs and 'field_nearaxis' in kwargs and len(initial_dofs) == len(coils.x) + len(kwargs['surface_all'].x) + len(kwargs['field_nearaxis'].x):
+            surface_all = kwargs['surface_all']
+            field_nearaxis = kwargs['field_nearaxis']
+            dofs_currents = result.x[len_dofs_curves:-len(surface_all.x)-len(field_nearaxis.x)]
+            curves = Curves(dofs_curves, n_segments, nfp, stellsym)
+            new_coils = Coils(curves=curves, currents=dofs_currents * coils.currents_scale)
+            new_surface = SurfaceRZFourier(rc=surface_all.rc, zs=surface_all.zs, nfp=nfp, range_torus=surface_all.range_torus, nphi=surface_all.nphi, ntheta=surface_all.ntheta)
+            new_surface.dofs = result.x[-len(surface_all.x)-len(field_nearaxis.x):-len(field_nearaxis.x)]
+            new_field_nearaxis = new_nearaxis_from_x_and_old_nearaxis(result.x[-len(field_nearaxis.x):], field_nearaxis)
+            return new_coils, new_surface, new_field_nearaxis
     except Exception as e:
         jax.debug.print("Error: {}", e)
         return None

examples/input_files/input.rotating_ellipse

@@ -4,11 +4,11 @@
   NFP = 0002
   MPOL = 002
   NTOR = 002
-!----- Boundary Parameters -----
-  RBC( 000,000) =    1.012658382028E+001    ZBS( 000,000) =    0.000000000000E+000
-  RBC( 001,000) =    1.656442246637E+000    ZBS( 001,000) =   -1.269527884010E+000
-  RBC(-001,001) =    3.742783485244E-001    ZBS(-001,001) =    3.400206123943E-001
-  RBC( 000,001) =    1.679774143201E+000    ZBS( 000,001) =    2.181203441748E+000
-  RBC( 001,001) =   -1.102571811326E+000    ZBS( 001,001) =    8.105522673578E-001
-  RBC(-002,002) =    6.572506820813E-003    ZBS(-002,002) =    4.002721688287E-003
+!----- Boundary Parameters (n,m) -----
+  RBC( 000,000) =    10    ZBS( 000,000) =    0
+  RBC( 001,000) =    1    ZBS( 001,000) =   -1
+  RBC(-001,001) =    0.1    ZBS(-001,001) =    0.1
+  RBC( 000,001) =    2.5    ZBS( 000,001) =    2.5
+  RBC( 001,001) =   -1    ZBS( 001,001) =    1
+  RBC(-002,002) =    1E-4    ZBS(-002,002) =    1E-4
 /