Optimisation de l'espace des phases pour extraction lente¶
In [1]:
import matplotlib.pyplot as plt
import numpy as np
import xtrack as xt
from helpers import characterize_phase_space
%config InlineBackend.figure_format = "retina"
Chargement de la maille PIMMS¶
In [2]:
env = xt.load("inputs/pimms.json")
ring = env.pimms
ring.configure_bend_model(edge="full", core="adaptive", num_multipole_kicks=5)
# Alimentation des sextupoles d'extraction
env["kse1"] = 1
env["kse2"] = -6.5
Charactérisation de l'espace des phases¶
Les fonctionalités du notebook précédent ont été implémentées en une fonction, pour aisance par la suite.
In [3]:
characterize_phase_space(ring)
Out[3]:
{'stable_area': np.float64(6.488631516398322e-05),
'dpx_dx_at_septum': np.float64(-0.03428405181036467),
'x_fixed_points': array([-0.01049033, -0.00440707, 0.01401406]),
'px_fixed_points': array([ 0.00120107, -0.00162625, 0.00047853]),
'x_norm_fixed_points': array([-0.00357046, -0.00149996, 0.00476989]),
'px_norm_fixed_points': array([ 0.0035696 , -0.00475893, 0.00135925])}
In [4]:
# La chractérisation (sans visualisation) est rapide !
%time res = characterize_phase_space(ring, plot=False)
CPU times: user 555 ms, sys: 48 μs, total: 555 ms Wall time: 554 ms
Optimisation de la résonance¶
Utilisons les résultats de la function de charactérisation pour notre optimisation. Entre autres nous aimerions contrôler l'orientation de la séparatrice et la taille de la zone stable.
In [5]:
# Dans Xsuite, il est possible de définir une Action à fournir pour les optimiseurs
class ActionSeparatrix(xt.Action):
def __init__(self, line):
self.line = line
def run(self):
return characterize_phase_space(self.line, plot=False)
In [6]:
# Construisons et testons notre action
action = ActionSeparatrix(ring)
action.run()
Out[6]:
{'stable_area': np.float64(6.488631516398322e-05),
'dpx_dx_at_septum': np.float64(-0.03428405181036467),
'x_fixed_points': array([-0.01049033, -0.00440707, 0.01401406]),
'px_fixed_points': array([ 0.00120107, -0.00162625, 0.00047853]),
'x_norm_fixed_points': array([-0.00357046, -0.00149996, 0.00476989]),
'px_norm_fixed_points': array([ 0.0035696 , -0.00475893, 0.00135925])}
In [7]:
# Il est possible de créer des cibles d'optimization depuis l'action
# Nous agissons sur l'alimentation des sextupoles d'extraction
opt = ring.match(
solve=False,
method="4d",
vary=xt.VaryList(["kse1", "kse2"], step=0.5, limits=[-7, 7]),
targets=[
action.target("stable_area", 1.0e-4, tol=1e-5, weight=100),
action.target("dpx_dx_at_septum", 0.03, tol=5e-4),
],
)
opt.target_status()
Matching: model call n. 1 penalty = 6.4380e-02
Target status: id state tag tol_met residue current_val target_val description 0 ON stable_area False -3.51137e-05 6.48863e-05 0.0001 'stable_area', val=0.0001, tol=1e-05, we ... 1 ON dpx_dx_at_septum False -0.0642841 -0.0342841 0.03 'dpx_dx_at_septum', val=0.03, tol=0.0005 ...
Utilisation d'un optimiseur externe¶
On pourrait obtenir la solution simplement en appelant opt.solve(), qui utilise l’optimiseur interne d'Xsuite.
Cependant, pour certains problèmes l'on pourrait vouloir utiliser un optimiseur différent (par exemple sans dérivées) s’il est mieux adapté au problème à résoudre.
Ici, nous montrons comment appliquer l’optimiseur Py-BOBYQA à notre problème de matching non linéaire.
In [8]:
# Extraction d'une fonction de mérite pour l'optimiseur
merit_function = opt.get_merit_function(
return_scalar=True, # Py-BOBYQA veut une fonction retournant un scalaire
check_limits=False, # Py-BOBYQA aime explorer en-dehors des limites imposées
)
In [9]:
# Extraction des limites et du point de départ de la fonction de mérite
bounds = merit_function.get_x_limits()
x0 = merit_function.get_x()
In [10]:
# Recherche d'optimum en utilisant Py-BOBYQA
import pybobyqa
soln = pybobyqa.solve(
merit_function,
x0=x0,
bounds=bounds.T, # pybobyqa veut une transposée...
rhobeg=5,
rhoend=1e-4,
maxfun=30,
objfun_has_noise=True, # <-- utile pour notre cas
seek_global_minimum=True,
)
soln.x # correspond à kse1, kse2
Matching: model call n. 2 penalty = 6.4380e-02
Matching: model call n. 3 penalty = 4.1530e-02
Matching: model call n. 4 penalty = 5.2873e+01
Matching: model call n. 5 penalty = 9.8461e-02
Matching: model call n. 6 penalty = 6.5169e-02
Matching: model call n. 7 penalty = 4.4227e-02
Matching: model call n. 8 penalty = 4.1261e-02
Matching: model call n. 9 penalty = 4.5117e-02
Matching: model call n. 10 penalty = 3.7992e-02
Matching: model call n. 11 penalty = 3.5580e-02
Matching: model call n. 12 penalty = 2.8810e-02
Matching: model call n. 13 penalty = 4.1651e-03
Matching: model call n. 14 penalty = 2.4018e-02
Matching: model call n. 15 penalty = 9.9428e-03
Matching: model call n. 16 penalty = 4.1112e-03
Matching: model call n. 17 penalty = 5.2015e-03
Matching: model call n. 18 penalty = 3.0931e-03
Matching: model call n. 19 penalty = 2.6983e-03
Matching: model call n. 20 penalty = 4.2250e-03
Matching: model call n. 21 penalty = 3.2037e-03
Matching: model call n. 22 penalty = 3.5286e-03
Matching: model call n. 23 penalty = 1.9232e-03
Matching: model call n. 24 penalty = 1.8838e-03
Matching: model call n. 25 penalty = 1.9780e-03
Matching: model call n. 26 penalty = 1.4600e-03
Matching: model call n. 27 penalty = 1.4200e-03
Matching: model call n. 28 penalty = 6.2383e-04
Matching: model call n. 29 penalty = 1.8472e-03
Matching: model call n. 30 penalty = 1.1477e-03
Matching: model call n. 31 penalty = 2.0962e-03
Out[10]:
array([4.96356347, 0.4426206 ])
In [11]:
# Implémentation de ces valeurs
merit_function.set_x(soln.x)
Matching: model call n. 32 penalty = 6.2383e-04
In [12]:
# Status après optimisation
opt.target_status()
opt.vary_status()
Matching: model call n. 33 penalty = 6.2383e-04
Target status: id state tag tol_met residue current_val target_val description 0 ON stable_area True -4.88465e-06 9.51154e-05 0.0001 'stable_area', val=0.0001, tol=1e-05, we ... 1 ON dpx_dx_at_septum True 0.000388024 0.030388 0.03 'dpx_dx_at_septum', val=0.03, tol=0.0005 ... Vary status: id state tag met name lower_limit current_val upper_limit val_at_iter_0 step weight 0 ON OK kse1 -7 4.96356 7 1 0.5 1 1 ON OK kse2 -7 0.442621 7 -6.5 0.5 1
In [13]:
# Nous pouvons tagger la solution dans l'optimisateur d'Xsuite,
# ce qui correspond à un checkpoint ou l'on peut toujours revenir
opt.tag("bobyqa solution")
Matching: model call n. 34 penalty = 6.2383e-04
In [14]:
# Visualisons l'espace des phases après optimisation
_ = characterize_phase_space(ring)
Chargement d'un checkpoint de l'optimiseur¶
In [15]:
opt.log()
Out[15]:
Table: 2 rows, 16 cols iteration penalty alpha tag tol_met target_active hit_limits ... 0 0.0643799 -1 nn yy nn 1 0.000623827 -1 bobyqa solution yy yy nn
In [16]:
# Il est facile de charger un point spécifique
# Dans notre cas, le point de départ
opt.reload(0)
opt.vary_status()
Matching: model call n. 35 penalty = 6.4380e-02
Vary status: id state tag met name lower_limit current_val upper_limit val_at_iter_0 step weight 0 ON OK kse1 -7 1 7 1 0.5 1 1 ON OK kse2 -7 -6.5 7 -6.5 0.5 1
In [17]:
# Espace des phases correspondant au point chargé
_ = characterize_phase_space(ring)
In [18]:
# Retour à la solution de pybobyqa
opt.reload(tag="bobyqa solution")
Matching: model call n. 36 penalty = 6.2383e-04
In [19]:
# Espace des phases correspondant au point chargé
_ = characterize_phase_space(ring)
Export des alimentation optimisées¶
In [20]:
strengths: dict[str, float] = opt.get_knob_values()
xt.json.dump(strengths, "inputs/extraction_sextupoles.json")