Refactor optimization functions to use a unified loss function approach across examples

Author: rogeriojorge

essos/objective_functions.py

@@ -4,24 +4,10 @@
 from jax import jit, vmap
 from functools import partial
 from essos.dynamics import Tracing
-from essos.fields import BiotSavart, near_axis
+from essos.fields import BiotSavart
 from essos.surfaces import BdotN_over_B, BdotN
 from essos.coils import Curves, Coils
-
-def new_nearaxis_from_x_and_old_nearaxis(new_field_nearaxis_x, field_nearaxis):
-    len_rc = len(field_nearaxis.rc)
-    len_zs = len(field_nearaxis.zs)
-    # # keeping the first rc and zs the same
-    # new_field_nearaxis_rc = jnp.concatenate((jnp.array([field_nearaxis.rc[0]]),new_field_nearaxis_x[:len_rc][1:]))
-    # new_field_nearaxis_zs = jnp.concatenate((jnp.array([field_nearaxis.zs[0]]),new_field_nearaxis_x[len_rc:len_rc+len_zs][1:]))
-    new_field_nearaxis_rc = new_field_nearaxis_x[:len_rc]
-    new_field_nearaxis_zs = new_field_nearaxis_x[len_rc:len_rc+len_zs]
-    new_field_nearaxis_etabar = new_field_nearaxis_x[-1]
-    
-    new_field_nearaxis = near_axis(rc=new_field_nearaxis_rc, zs=new_field_nearaxis_zs, etabar=new_field_nearaxis_etabar,
-                                    B0=field_nearaxis.B0, sigma0=field_nearaxis.sigma0, I2=field_nearaxis.I2,
-                                    nphi=field_nearaxis.nphi, spsi=field_nearaxis.spsi, sG=field_nearaxis.sG, nfp=field_nearaxis.nfp)
-    return new_field_nearaxis
+from essos.optimization import new_nearaxis_from_x_and_old_nearaxis
 
 @partial(jit, static_argnums=(1, 4, 5, 6, 7, 8))
 def loss_coils_for_nearaxis(x, field_nearaxis, dofs_curves, currents_scale, nfp, max_coil_length=42,

essos/optimization.py

@@ -5,9 +5,22 @@
 from functools import partial
 from essos.coils import Curves, Coils
 from scipy.optimize import least_squares, minimize
-from essos.objective_functions import (loss_optimize_coils_for_particle_confinement,
-                                       loss_BdotN, loss_coils_for_nearaxis, loss_coils_and_nearaxis,
-                                       new_nearaxis_from_x_and_old_nearaxis)
+from essos.fields import near_axis
+
+def new_nearaxis_from_x_and_old_nearaxis(new_field_nearaxis_x, field_nearaxis):
+    len_rc = len(field_nearaxis.rc)
+    len_zs = len(field_nearaxis.zs)
+    # # keeping the first rc and zs the same
+    # new_field_nearaxis_rc = jnp.concatenate((jnp.array([field_nearaxis.rc[0]]),new_field_nearaxis_x[:len_rc][1:]))
+    # new_field_nearaxis_zs = jnp.concatenate((jnp.array([field_nearaxis.zs[0]]),new_field_nearaxis_x[len_rc:len_rc+len_zs][1:]))
+    new_field_nearaxis_rc = new_field_nearaxis_x[:len_rc]
+    new_field_nearaxis_zs = new_field_nearaxis_x[len_rc:len_rc+len_zs]
+    new_field_nearaxis_etabar = new_field_nearaxis_x[-1]
+    
+    new_field_nearaxis = near_axis(rc=new_field_nearaxis_rc, zs=new_field_nearaxis_zs, etabar=new_field_nearaxis_etabar,
+                                    B0=field_nearaxis.B0, sigma0=field_nearaxis.sigma0, I2=field_nearaxis.I2,
+                                    nphi=field_nearaxis.nphi, spsi=field_nearaxis.spsi, sG=field_nearaxis.sG, nfp=field_nearaxis.nfp)
+    return new_field_nearaxis
 
 def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e-4, maximum_function_evaluations=30, **kwargs):
     len_dofs_curves = len(jnp.ravel(coils.dofs_curves))
@@ -48,34 +61,3 @@ def optimize_loss_function(func, initial_dofs, coils, tolerance_optimization=1e-
     except Exception as e:
         jax.debug.print("Error: {}", e)
         return None
-
-def optimize_coils_for_particle_confinement(coils, particles, target_B_on_axis=5.7, max_coil_length=22, model='GuidingCenter',
-                                            maxtime=5e-6, num_steps=500, trace_tolerance=1e-5, tolerance_optimization=1e-4,
-                                            maximum_function_evaluations=30, max_coil_curvature=0.1):
-    return optimize_loss_function(loss_optimize_coils_for_particle_confinement, initial_dofs=coils.x, coils=coils,
-                           tolerance_optimization=tolerance_optimization, particles=particles,
-                           maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature,
-                           target_B_on_axis=target_B_on_axis, max_coil_length=max_coil_length, model=model,
-                           maxtime=maxtime, num_steps=num_steps, trace_tolerance=trace_tolerance)
-
-def optimize_coils_for_vmec_surface(vmec, coils, tolerance_optimization=1e-10,
-                                    maximum_function_evaluations=30,
-                                    max_coil_length=42, max_coil_curvature=0.1):
-    return optimize_loss_function(loss_BdotN, initial_dofs=coils.x, coils=coils, tolerance_optimization=tolerance_optimization,
-                                  maximum_function_evaluations=maximum_function_evaluations, vmec=vmec,
-                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
-    
-def optimize_coils_for_nearaxis(field_nearaxis, coils, tolerance_optimization=1e-10,
-                                    maximum_function_evaluations=30,
-                                    max_coil_length=42, max_coil_curvature=0.1):
-    return optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils.x, coils=coils, tolerance_optimization=tolerance_optimization,
-                                  maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis,
-                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
-
-def optimize_coils_and_nearaxis(field_nearaxis, coils, tolerance_optimization=1e-10,
-                                    maximum_function_evaluations=30,
-                                    max_coil_length=42, max_coil_curvature=0.1):
-    initial_dofs = jnp.concatenate((coils.x, field_nearaxis.x))
-    return optimize_loss_function(loss_coils_and_nearaxis, initial_dofs=initial_dofs, coils=coils, tolerance_optimization=tolerance_optimization,
-                                  maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis,
-                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)

examples/optimize_coils_and_nearaxis.py

@@ -4,8 +4,9 @@
 from essos.coils import Coils, CreateEquallySpacedCurves
 from essos.fields import near_axis, BiotSavart
 from essos.dynamics import Tracing
-from essos.optimization import optimize_coils_and_nearaxis, optimize_coils_for_nearaxis
-from essos.objective_functions import difference_B_gradB_onaxis
+from essos.optimization import optimize_loss_function
+from essos.objective_functions import (difference_B_gradB_onaxis,
+                    loss_coils_and_nearaxis, loss_coils_for_nearaxis)
 
 # Optimization parameters
 max_coil_length = 5.0
@@ -38,17 +39,19 @@
 # Optimize coils
 print(f'Optimizing coils for initial near=axis with {maximum_function_evaluations} function evaluations.')
 time0 = time()
-coils_optimized_initial_nearaxis = optimize_coils_for_nearaxis(field_nearaxis_initial, coils_initial, maximum_function_evaluations=maximum_function_evaluations,
-                                                    max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,
-                                                    tolerance_optimization=tolerance_optimization)
+initial_dofs = coils_initial.x
+coils_optimized_initial_nearaxis = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils_initial.x, coils=coils_initial, tolerance_optimization=tolerance_optimization,
+                                  maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis_initial,
+                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
 print(f"Optimization took {time()-time0:.2f} seconds")
 
 # Optimize coils
 print(f'Optimizing coils and near-axis with {maximum_function_evaluations} function evaluations.')
 time0 = time()
-coils_optimized, field_nearaxis_optimized = optimize_coils_and_nearaxis(field_nearaxis_initial, coils_initial, maximum_function_evaluations=maximum_function_evaluations,
-                                                    max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,
-                                                    tolerance_optimization=tolerance_optimization)
+initial_dofs = jnp.concatenate((coils_initial.x, field_nearaxis_initial.x))
+coils_optimized, field_nearaxis_optimized = optimize_loss_function(loss_coils_and_nearaxis, initial_dofs=initial_dofs, coils=coils_initial, tolerance_optimization=tolerance_optimization,
+                                  maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field_nearaxis_initial,
+                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
 print(f"Optimization took {time()-time0:.2f} seconds")
 
 B_difference_initial, gradB_difference_initial = difference_B_gradB_onaxis(field_nearaxis_initial, BiotSavart(coils_optimized_initial_nearaxis))

examples/optimize_coils_for_nearaxis.py

@@ -4,7 +4,8 @@
 from essos.coils import Coils, CreateEquallySpacedCurves
 from essos.fields import near_axis, BiotSavart
 from essos.dynamics import Tracing
-from essos.optimization import optimize_coils_for_nearaxis
+from essos.optimization import optimize_loss_function
+from essos.objective_functions import loss_coils_for_nearaxis
 
 # Optimization parameters
 max_coil_length = 5.0
@@ -37,9 +38,11 @@
 # Optimize coils
 print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.')
 time0 = time()
-coils_optimized = optimize_coils_for_nearaxis(field, coils_initial, maximum_function_evaluations=maximum_function_evaluations,
-                                                    max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,
-                                                    tolerance_optimization=tolerance_optimization)
+initial_dofs = coils_initial.x
+coils_optimized = optimize_loss_function(loss_coils_for_nearaxis, initial_dofs=coils_initial.x,
+                                  coils=coils_initial, tolerance_optimization=tolerance_optimization,
+                                  maximum_function_evaluations=maximum_function_evaluations, field_nearaxis=field,
+                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
 print(f"Optimization took {time()-time0:.2f} seconds")
 
 # Trace fieldlines

examples/optimize_coils_particle_confinement_fullorbit.py

@@ -8,7 +8,8 @@
 from essos.dynamics import Particles, Tracing
 from essos.fields import BiotSavart
 from essos.coils import Coils, CreateEquallySpacedCurves
-from essos.optimization import optimize_coils_for_particle_confinement
+from essos.optimization import optimize_loss_function
+from essos.objective_functions import loss_optimize_coils_for_particle_confinement
 
 # Optimization parameters
 target_B_on_axis = 5.7
@@ -17,7 +18,7 @@
 nparticles = 8
 order_Fourier_series_coils = 4
 number_coil_points = 80
-maximum_function_evaluations = 100
+maximum_function_evaluations = 60
 maxtime_tracing = 1e-5
 number_coils_per_half_field_period = 3
 number_of_field_periods = 2
@@ -45,8 +46,11 @@
 # Optimize coils
 print(f'Optimizing coils with {maximum_function_evaluations} function evaluations and maxtime_tracing={maxtime_tracing}')
 time0 = time()
-coils_optimized = optimize_coils_for_particle_confinement(coils_initial, particles, target_B_on_axis=target_B_on_axis, maxtime=maxtime_tracing, model=model,
-                                        max_coil_length=max_coil_length, maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature)
+coils_optimized = optimize_loss_function(loss_optimize_coils_for_particle_confinement, initial_dofs=coils_initial.x,
+                           coils=coils_initial, tolerance_optimization=1e-4, particles=particles,
+                           maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature,
+                           target_B_on_axis=target_B_on_axis, max_coil_length=max_coil_length, model=model,
+                           maxtime=maxtime_tracing, num_steps=timesteps, trace_tolerance=1e-5)
 print(f"  Optimization took {time()-time0:.2f} seconds")
 particles.to_full_orbit(BiotSavart(coils_optimized))
 tracing_optimized = Tracing(field=coils_optimized, particles=particles, maxtime=maxtime_tracing, model=model, timesteps=timesteps)

examples/optimize_coils_particle_confinement_guidingcenter.py

@@ -7,7 +7,8 @@
 import matplotlib.pyplot as plt
 from essos.dynamics import Particles, Tracing
 from essos.coils import Coils, CreateEquallySpacedCurves
-from essos.optimization import optimize_coils_for_particle_confinement
+from essos.optimization import optimize_loss_function
+from essos.objective_functions import loss_optimize_coils_for_particle_confinement
 
 # Optimization parameters
 target_B_on_axis = 5.7
@@ -42,8 +43,13 @@
 # Optimize coils
 print(f'Optimizing coils with {maximum_function_evaluations} function evaluations and maxtime_tracing={maxtime_tracing}')
 time0 = time()
-coils_optimized = optimize_coils_for_particle_confinement(coils_initial, particles, target_B_on_axis=target_B_on_axis, maxtime=maxtime_tracing, model=model,
-                                        max_coil_length=max_coil_length, maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature)
+coils_optimized = optimize_loss_function(loss_optimize_coils_for_particle_confinement, initial_dofs=coils_initial.x, coils=coils_initial,
+                           tolerance_optimization=1e-4, particles=particles,
+                           maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature,
+                           target_B_on_axis=target_B_on_axis, max_coil_length=max_coil_length, model=model,
+                           maxtime=maxtime_tracing, num_steps=500, trace_tolerance=1e-5)
+# coils_optimized = optimize_coils_for_particle_confinement(coils_initial, particles, target_B_on_axis=target_B_on_axis, maxtime=maxtime_tracing, model=model,
+#                                         max_coil_length=max_coil_length, maximum_function_evaluations=maximum_function_evaluations, max_coil_curvature=max_coil_curvature)
 print(f"  Optimization took {time()-time0:.2f} seconds")
 tracing_optimized = Tracing(field=coils_optimized, particles=particles, maxtime=maxtime_tracing, model=model)
 

examples/optimize_coils_vmec_surface.py

@@ -2,10 +2,13 @@
 number_of_processors_to_use = 8 # Parallelization, this should divide ntheta*nphi
 os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}'
 from time import time
+import jax.numpy as jnp
 import matplotlib.pyplot as plt
+from essos.surfaces import BdotN_over_B
 from essos.coils import Coils, CreateEquallySpacedCurves
-from essos.fields import Vmec
-from essos.optimization import optimize_coils_for_vmec_surface
+from essos.fields import Vmec, BiotSavart
+from essos.objective_functions import loss_BdotN
+from essos.optimization import optimize_loss_function
 
 # Optimization parameters
 max_coil_length = 40
@@ -38,11 +41,16 @@
 # Optimize coils
 print(f'Optimizing coils with {maximum_function_evaluations} function evaluations.')
 time0 = time()
-coils_optimized = optimize_coils_for_vmec_surface(vmec, coils_initial, maximum_function_evaluations=maximum_function_evaluations,
-                                                    max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,
-                                                    tolerance_optimization=tolerance_optimization)
+coils_optimized = optimize_loss_function(loss_BdotN, initial_dofs=coils_initial.x, coils=coils_initial, tolerance_optimization=tolerance_optimization,
+                                  maximum_function_evaluations=maximum_function_evaluations, vmec=vmec,
+                                  max_coil_length=max_coil_length, max_coil_curvature=max_coil_curvature,)
 print(f"Optimization took {time()-time0:.2f} seconds")
 
+BdotN_over_B_initial = BdotN_over_B(vmec.surface, BiotSavart(coils_initial))
+BdotN_over_B_optimized = BdotN_over_B(vmec.surface, BiotSavart(coils_optimized))
+print(f"Maximum BdotN/B before optimization: {jnp.max(BdotN_over_B_initial):.2e}")
+print(f"Maximum BdotN/B after optimization: {jnp.max(BdotN_over_B_optimized):.2e}")
+
 # Plot coils, before and after optimization
 fig = plt.figure(figsize=(8, 4))
 ax1 = fig.add_subplot(121, projection='3d')