Changes between Version 8 and Version 9 of misc/setup_for_ML/example_beamsize

Timestamp:

09/13/23 23:11:32 (10 months ago)

Author:

obina

Comment:

--

misc/setup_for_ML/example_beamsize

@@ -37 +37 @@
 from ocelot.gui.accelerator import *
 
-def func(x):
-    global cnt, QM1, QM2, tws, lat, tws0
-    x0=x[:,0][0]
-    x1=x[:,1][0]
-
+def calc_resp(x0, x1):
+    # Define Optics
+    ## Drift
+    L20  = Drift(l=0.2, eid='L20')
+    L50  = Drift(l=0.5, eid='L50')
+    ## Quadrupoles
+    QM1 = Quadrupole(l=0.2, k1=x0, eid='QM1')
+    QM2 = Quadrupole(l=0.2, k1=x1, eid='QM2')
+    ## Lattice
+    cell = (L50, QM1, L20, QM2, L50, L50)
+    lat = MagneticLattice(cell, stop=None)
+    # print("length of the cell: ", lat.totalLen, " m")
+    # Initial condition
+    tws0 = Twiss()
+    tws0.beta_x = 20.0
+    tws0.beta_y = 30.0
+    tws0.alpha_x = -10.0
+    tws0.alpha_y = -5.0
+    tws0.emit_x = 1.0
+    tws0.emit_y = 1.0
+    tws0.E = 1.0
     # update optics
-    QM1.k1 = x0
-    QM2.k1 = x1
     lat.update_transfer_maps()
     tws = twiss(lat, tws0, nPoints=1000)
 
     idx = -1  # last point of the lattice
-    sx = np.sqrt(tws[idx].beta_x)
-    sy = np.sqrt(tws[idx].beta_y)
+    sx = np.sqrt(tws0.emit_x * tws[idx].beta_x)
+    sy = np.sqrt(tws0.emit_y * tws[idx].beta_y)
+
+    return sx, sy, tws
+
+def show_beta(x0, x1):
+    sx, sy, tws = calc_resp(x0, x1)
+
+    s = [p.s for p in tws]
+    beta_x = [p.beta_x for p in tws]
+    beta_y = [p.beta_y for p in tws]
+
+    plt.plot(s, beta_x, 'b-');
+    plt.plot(s, beta_y, 'r-');
+    plt.legend(['Horiz', 'Vert'])
+    plt.xlabel('s [m]')
+    plt.ylabel('beta [m]')
+    plt.grid(True)
+    plt.show()
+
+def eval_func(x):
+    global cnt
+
+    x0 = x[0][0]
+    x1 = x[0][1]
+
+    sx, sy, tws = calc_resp(x0, x1) # calc optics for 2 QM
 
     # Evaluation Function Example
@@ -64 +103 @@
 
 # ==== Main =====
-
-# Define Optics
-## Drift
-L20  = Drift(l=0.2, eid='L20')
-L50  = Drift(l=0.5, eid='L50')
-## Quadrupoles
-QM1 = Quadrupole(l=0.2, k1=3.9, eid='QM1')
-QM2 = Quadrupole(l=0.2, k1=-3.25, eid='QM2')
-## Lattice
-cell = (L50, QM1, L20, QM2, L50, L50)
-lat = MagneticLattice(cell, stop=None)
-print("length of the cell: ", lat.totalLen, " m")
-tws0 = Twiss()
-tws0.beta_x = 20.0
-tws0.beta_y = 30.0
-tws0.alpha_x = -10.0
-tws0.alpha_y = -5.0
-tws0.emit_x = 1.0
-tws0.emit_y = 1.0
-tws0.E = 1.0
-
-tws = twiss(lat, tws0, nPoints=1000)
-
 # GP Optimization
 print("==== start ====")
 print("# cnt, x0, x1, sx, sy, eval_val")
-
 cnt = 0
 bounds = [{'name':'x0', 'type':'continuous', 'domain':(-20, 20)},
           {'name':'x1', 'type':'continuous', 'domain':(-20, 20)}]
 
-myBopt = GPyOpt.methods.BayesianOptimization(f=func,domain=bounds,initial_design_numdata=10,acquisition_type='LCB')
+myBopt = GPyOpt.methods.BayesianOptimization(f=eval_func,domain=bounds,initial_design_numdata=10,acquisition_type='LCB')
 
 myBopt.run_optimization(max_iter=50)
@@ -101 +116 @@
 myBopt.plot_convergence()
 print("Best     = ", myBopt.x_opt)
-print("eval_val = ", func(np.array([myBopt.x_opt])))
 
-# Plot results
-QM1.k1 = myBopt.x_opt[0]
-QM2.k1 = myBopt.x_opt[1]
-lat.update_transfer_maps()
-tws = twiss(lat, tws0, nPoints=1000)
+# plot betatron function
+x0, x1 = np.array(myBopt.x_opt)
+show_beta(x0, x1)
 
-s = [p.s for p in tws]
-beta_x = [p.beta_x for p in tws]
-beta_y = [p.beta_y for p in tws]
-
-plt.plot(s, beta_x, 'b-');
-plt.plot(s, beta_y, 'r-');
-
-plt.legend(['Horiz', 'Vert'])
-plt.xlabel('s [m]')
-plt.ylabel('beta [m]')
-plt.grid(True)
-plt.show()
 }}}