Оптимизация огибающей пучка заряженных частиц с помощью генетического алгоритма

Задача, обратная расчёту огибающей по заданным полям, сложнее и встречается чаще. Известно, какую огибающую требуется получить, а подобрать необходимо поля соленоидов.

Перебор полей девяти соленоидов вручную занимает весь день, и с каждым дополнительным соленоидом задача становится всё менее выполнимой. Поскольку один расчёт огибающей занимает доли секунды, перебор целесообразно поручить машине.

Генетические алгоритмы работают по аналогии с естественным отбором в природе, а основу генетического алгоритма составляют особь, хромосомы (гены), популяция (набор особей), функции приспособленности, отбора, скрещивания и мутации. На практике процесс выглядит следующим образом: создание первого поколения, оценка, отбор, скрещивание, мутация, новое поколение, оценка, и так далее, пока не будет получен требуемый результат.

DEAP [41] представляет собой фреймворк для работы с генетическими алгоритмами, в котором уже есть множество готовых инструментов.

Создание пучка и ускорителя

Расчёт огибающей из предыдущей главы вынесен в библиотеку kenv, поэтому остаётся собрать установку из готовых частей.

import kenv as kv

Пучок задаётся четырьмя числами, которыми служат энергия 2 МэВ, ток 2 кА, начальный радиус 50 мм при такой же расходимости (rp) и нормализованный эмиттанс 1000 мм·мрад.

beam = kv.Beam(energy=2,
               current=2e3,
               radius=50e-3,
               rp=50e-3,
               normalized_emittance=1000e-6)

Ускоритель занимает отрезок от нуля до пяти метров с шагом сетки один сантиметр.

accelerator = kv.Accelerator(0, 5, 0.01)

На этом отрезке равномерно, через полметра, расставлены девять одинаковых соленоидов, пока включённых на одно и то же поле 0.02 Тл; это поле служит начальным приближением, которое алгоритму предстоит улучшить.

Solenoids = [ 
    [ 0.5000, 0.02, 'Bz.dat', 'Sol. 1'],
    [ 1.0000, 0.02, 'Bz.dat', 'Sol. 2'],
    [ 1.5000, 0.02, 'Bz.dat', 'Sol. 3'],
    [ 2.0000, 0.02, 'Bz.dat', 'Sol. 4'],
    [ 2.5000, 0.02, 'Bz.dat', 'Sol. 5'],
    [ 3.0000, 0.02, 'Bz.dat', 'Sol. 6'],
    [ 3.5000, 0.02, 'Bz.dat', 'Sol. 7'],
    [ 4.0000, 0.02, 'Bz.dat', 'Sol. 8'],
    [ 4.5000, 0.02, 'Bz.dat', 'Sol. 9'],
]

Перечисление преобразуется в элементы линейки.

for   z0, B0, filename, name in Solenoids:
    accelerator.Bz_beamline[name] = kv.Element(z0, B0, filename, name)

Вызов compile() сшивает профили всех соленоидов в одно поле вдоль тракта, как в предыдущей главе делалось вручную. После этого пучок и ускоритель связываются в расчёт, а track() интегрирует уравнение огибающей от начала до конца.

accelerator.compile()
simulation = kv.Simulation(beam, accelerator)
simulation.track()

График

Прежде чем приступать к оптимизации, необходимо увидеть исходное состояние.

matplotlib

Оформление настраивается заданием широких осей, полупрозрачных заливок и векторного вывода.

import holoviews as hv
hv.extension('matplotlib')

%opts Layout [tight=True]
%output size=150 backend='matplotlib' fig='svg'

%opts Area Curve [aspect=3 show_grid=True]
%opts Area  (alpha=0.25)
%opts Curve (alpha=0.5)
%opts Area.Beam [aspect=3 show_grid=True] (color='red' alpha=0.3)

import warnings
warnings.filterwarnings('ignore')

dim_z  = hv.Dimension('z',  unit='m', range=(accelerator.start, accelerator.stop))
dim_Bz = hv.Dimension('Bz', unit='T', label='Bz', range=(0, 0.1))
dim_Ez = hv.Dimension('Ez', unit='MV/m', label='Ez')

dim_r = hv.Dimension('r', label="Beam r", unit='mm', range=(0, 150))

Строятся две панели: сверху огибающая пучка в миллиметрах, снизу поле соленоидов вдоль тракта. По нижней панели видно, как алгоритм изменил установку, а по верхней — как это отразилось на пучке.

z_Bz= hv.Area((accelerator.parameter,accelerator.Bz(accelerator.parameter)), kdims=[dim_z], vdims=[dim_Bz])
z_r = hv.Area(((accelerator.parameter,simulation.envelope_x(accelerator.parameter)*1e3)), kdims=[dim_z], vdims=[dim_r], group='Beam')

(z_r+z_Bz).cols(1)

Огибающая далека от идеальной, поскольку одинаковые поля соленоидов дают периодическую перетяжку, и вместо плавного сжатия к выходу пучок проходит по тракту волной. Именно эту волну предстоит сгладить генетическим алгоритмом.

Генетический алгоритм

Модельную огибающую построим простейшим способом, проведя параболу через три точки: 150 мм на входе, 50 мм в середине тракта и 30 мм на выходе. Физического смысла эта кривая не имеет и представляет собой техническое задание для алгоритма.

import numpy as np

coefficients = np.polyfit([accelerator.start, (accelerator.start+accelerator.stop)/2, accelerator.stop], [0.15, 0.05, 0.03], 2) 
envelope_mod = coefficients[0]*accelerator.parameter**2 +  coefficients[1]*accelerator.parameter+ coefficients[2]

z_env = hv.Curve((accelerator.parameter, envelope_mod*1e3), kdims=[dim_z], vdims=[dim_r], label='Envelope_mod', group='Beam')

Полученная кривая приведена ниже.

z_env

Далее используется DEAP.

import random
from deap import creator, base, tools, algorithms

Первым шагом создаются необходимые типы, обычно функция приспособленности (fitness) и особь (individual).

creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
creator.create("Individual", list, fitness=creator.FitnessMin)

Далее необходимо собрать набор инструментов (toolbox), задающий основные параметры алгоритма. Сначала задаются границы, в пределах которых разрешено изменяться полям соленоидов, и размеры эволюции, то есть сто особей в популяции, сто поколений, вероятность скрещивания 0.4 и мутации 0.6.

Sol_B = 0.02     # [T] average Bz in the solenoids
Sol_B_max = 0.05  # [T] max Bz
Sol_B_min = 0.005 # [T] min Bz
Sol_Num = 9     # quantity

CXPB = 0.4       # cross chance
CX_INDPB = 0.6  # вероятность обмена одним геном
MUTPB = 0.6     # Mutation probability
NGEN = 100       # Number of generations
POP = 100       # Number of individuals
TOURN = 3       # Tournament size

В toolbox регистрируются три операции: порождение одного гена (случайное поле около 0.02 Тл), сборка особи из девяти генов и формирование популяции из особей. Особь здесь представляет собой список из девяти чисел, по одному на соленоид.

toolbox = base.Toolbox()
toolbox.register("attr_float", random.gauss, Sol_B, (Sol_B_max - Sol_B_min)/2)
toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_float, n=Sol_Num)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

Функция приспособленности является единственным местом, где алгоритм соприкасается с физикой. Она берёт девять чисел из особи, расставляет их по соленоидам, пересобирает ускоритель, заново рассчитывает огибающую и возвращает среднее по тракту расхождение с заданной кривой, в метрах. Чем меньше расхождение, тем лучше, поэтому вес приспособленности выше и задан отрицательным.

def evalution_envelope(individual):

    for x in range(Sol_Num):
        accelerator.Bz_beamline['Sol. %.d'%(x+1)].max_field  = individual[x]
    
    accelerator.compile()
    simulation = kv.Simulation(beam, accelerator)
    simulation.track()
    
    abs_errors = np.abs(envelope_mod - simulation.envelope_x(accelerator.parameter))
    
    return abs_errors.mean(),

Остаётся указать DEAP операторы скрещивания, мутации и отбора. Скрещивание равномерное, потомок берёт каждый ген у одного из двух родителей; мутация гауссова, к полю соленоида добавляется случайная добавка; отбор турнирный.

toolbox.register("evaluate", evalution_envelope)
toolbox.register("mate", tools.cxUniform, indpb=CX_INDPB)
toolbox.register("mutate", tools.mutGaussian, mu=0, sigma=(Sol_B_max - Sol_B_min)/2, indpb=MUTPB)
toolbox.register("select", tools.selTournament, tournsize=TOURN)


def clip(func):
    """Загоняет поля обратно в коридор после скрещивания и мутации."""
    def wrapper(*args, **kwargs):
        offspring = func(*args, **kwargs)
        for child in offspring:
            for i, b in enumerate(child):
                child[i] = min(max(b, Sol_B_min), Sol_B_max)
        return offspring
    return wrapper

Без последней функции объявленные Sol_B_min и Sol_B_max остались бы комментарием, поскольку мутация гауссовым шумом ничего не знает о коридоре и выведет поле за паспортный предел источника питания. Само по себе объявление clip ничего не делает: чтобы ограничение срабатывало после каждой операции над особью, функция навешивается декоратором на скрещивание и мутацию.

toolbox.decorate("mate", clip)
toolbox.decorate("mutate", clip)

Размер турнира здесь является отдельной константой, а не долей от числа поколений. Турнир задаёт давление отбора: чем он больше, тем увереннее побеждают текущие лидеры и тем быстрее исчезает разнообразие. Три-пять особей представляют собой обычный рабочий выбор, а турнир в треть популяции превращает эволюцию в жадный поиск.

Для наблюдения за эволюцией DEAP ведёт журнал (logbook). Ему поручается вычислять по каждому поколению среднее, разброс, лучшее и худшее значения приспособленности.

stats_fit = tools.Statistics(lambda ind: ind.fitness.values)
mstats = tools.MultiStatistics(fitness=stats_fit)
mstats.register("avg", np.mean)
mstats.register("std", np.std)
mstats.register("min", np.min)
mstats.register("max", np.max)

Запускается eaSimple, реализующий простейшую эволюционную схему: популяция оценивается, отбирается, скрещивается, мутирует, и всё это повторяется сто раз.

population = toolbox.population(n=POP)
pop, logbook = algorithms.eaSimple(population, toolbox, cxpb=CXPB, mutpb=MUTPB, ngen=NGEN, stats=mstats, verbose=True)
   	      	                                 fitness                                  
   	      	--------------------------------------------------------------------------
gen	nevals	avg      	gen	max      	min      	nevals	std       
0  	100   	0.0373577	0  	0.0408331	0.0337452	100   	0.00113532
1  	82    	0.0348016	1  	0.0379858	0.0299872	82    	0.00134257
2  	69    	0.0317311	2  	0.0355503	0.0268419	69    	0.00158446
3  	70    	0.0284586	3  	0.0310798	0.0258454	70    	0.00116153
4  	77    	0.0262835	4  	0.028602 	0.0245999	77    	0.00073413
5  	72    	0.0248689	5  	0.0268379	0.0234427	72    	0.000541174
6  	84    	0.0237668	6  	0.0256165	0.0221611	84    	0.00067109 
7  	73    	0.0225147	7  	0.0241554	0.0203067	73    	0.000610466
8  	88    	0.0210856	8  	0.0235622	0.0194378	88    	0.000755119
9  	80    	0.0197334	9  	0.0209437	0.0190342	80    	0.000337098
10 	81    	0.0191692	10 	0.0204253	0.0183268	81    	0.00035722 
11 	75    	0.0185557	11 	0.019661 	0.0178986	75    	0.000283725
12 	78    	0.0181075	12 	0.0190688	0.0175483	78    	0.00029574 
13 	72    	0.0177022	13 	0.018703 	0.0172446	72    	0.000248748
14 	75    	0.0174452	14 	0.0183427	0.0169873	75    	0.000190804
15 	70    	0.017232 	15 	0.0178082	0.0169012	70    	0.000181106
16 	79    	0.0170371	16 	0.0180127	0.0167734	79    	0.000208848
17 	80    	0.0168772	17 	0.0176019	0.0165999	80    	0.000167511
18 	80    	0.0167085	18 	0.0175836	0.0164198	80    	0.000174468
19 	77    	0.0165316	19 	0.0171106	0.0162831	77    	0.000152431
20 	76    	0.0163782	20 	0.0168837	0.0160568	76    	0.000171506
21 	69    	0.0161959	21 	0.0170411	0.015828 	69    	0.000194345
22 	77    	0.0159477	22 	0.0166562	0.0156399	77    	0.000157436
23 	79    	0.0158229	23 	0.0184114	0.0155106	79    	0.000314105
24 	80    	0.0156802	24 	0.0167603	0.0153466	80    	0.000234712
25 	72    	0.0154841	25 	0.0165787	0.0152393	72    	0.000184083
26 	74    	0.0153536	26 	0.0160787	0.0150713	74    	0.000201229
27 	83    	0.0152175	27 	0.0160113	0.0147709	83    	0.000237473
28 	79    	0.0149943	28 	0.0159486	0.014535 	79    	0.00027873 
29 	80    	0.0147364	29 	0.0155165	0.0143068	80    	0.000219007
30 	78    	0.0145096	30 	0.0155232	0.01421  	78    	0.000219074
31 	76    	0.0143674	31 	0.0162176	0.0141153	76    	0.000292944
32 	77    	0.0142465	32 	0.0154183	0.0139755	77    	0.000228094
33 	74    	0.0141755	33 	0.016872 	0.0137429	74    	0.000370584
34 	65    	0.0140263	34 	0.0165405	0.0135535	65    	0.000369079
35 	79    	0.0138775	35 	0.0165023	0.0134757	79    	0.000469816
36 	78    	0.0136499	36 	0.0165048	0.0132913	78    	0.00036574 
37 	77    	0.0135281	37 	0.0162187	0.0131399	77    	0.000370037
38 	72    	0.0133894	38 	0.0161935	0.013018 	72    	0.000407104
39 	74    	0.0132855	39 	0.0151958	0.0129085	74    	0.000356725
40 	73    	0.0132202	40 	0.0158132	0.0127831	73    	0.000502361
41 	78    	0.0131469	41 	0.0154469	0.0126621	78    	0.000549534
42 	72    	0.0128525	42 	0.0153473	0.0125449	72    	0.000371909
43 	81    	0.0128408	43 	0.0143709	0.0124426	81    	0.00037498 
44 	80    	0.0128018	44 	0.0143546	0.0122802	80    	0.00042712 
45 	76    	0.0126884	45 	0.0149749	0.0122105	76    	0.000484442
46 	74    	0.0124913	46 	0.0141759	0.0121807	74    	0.000374844
47 	84    	0.0124814	47 	0.0150409	0.012149 	84    	0.000478495
48 	65    	0.012378 	48 	0.0143345	0.0120466	65    	0.000396209
49 	77    	0.0123797	49 	0.0141054	0.0119723	77    	0.000465145
50 	66    	0.0122275	50 	0.0136186	0.0119299	66    	0.000324548
51 	77    	0.0122757	51 	0.0142704	0.0118622	77    	0.000472376
52 	74    	0.0121928	52 	0.0138265	0.0118082	74    	0.000460488
53 	73    	0.0121423	53 	0.0145869	0.0117309	73    	0.000487374
54 	70    	0.0120654	54 	0.0132641	0.0116092	70    	0.000405956
55 	77    	0.012012 	55 	0.0137579	0.0115895	77    	0.00050723 
56 	73    	0.0118692	56 	0.0141939	0.0115563	73    	0.000437811
57 	83    	0.0118798	57 	0.0139667	0.0114851	83    	0.000426248
58 	71    	0.0117967	58 	0.01344  	0.0114796	71    	0.000419319
59 	72    	0.0117301	59 	0.0145933	0.0114334	72    	0.000445674
60 	78    	0.0116524	60 	0.0135045	0.0113808	78    	0.000333071
61 	78    	0.0116611	61 	0.0132638	0.0113299	78    	0.000437407
62 	78    	0.0116445	62 	0.0140102	0.0113001	78    	0.000467812
63 	80    	0.0116662	63 	0.0138093	0.0112541	80    	0.000496356
64 	78    	0.0115682	64 	0.0126411	0.0112403	78    	0.000360796
65 	73    	0.0115044	65 	0.0133941	0.0112293	73    	0.000428872
66 	81    	0.0114277	66 	0.0134346	0.0112152	81    	0.000315043
67 	72    	0.0114574	67 	0.0132544	0.0111495	72    	0.000390321
68 	81    	0.0114764	68 	0.0135241	0.0110932	81    	0.000486222
69 	75    	0.011434 	69 	0.0132443	0.0110874	75    	0.000469529
70 	74    	0.0113903	70 	0.0129446	0.0109944	74    	0.000420948
71 	80    	0.0112841	71 	0.0132937	0.0109944	80    	0.000408004
72 	75    	0.0112891	72 	0.0133998	0.010978 	75    	0.000465794
73 	79    	0.0112265	73 	0.0122116	0.0109464	79    	0.00029769 
74 	80    	0.0112526	74 	0.0129213	0.0108979	80    	0.000417258
75 	71    	0.0112418	75 	0.0128844	0.0108979	71    	0.000447559
76 	73    	0.0111925	76 	0.0138475	0.0108402	73    	0.000519714
77 	78    	0.0111465	77 	0.0132176	0.010835 	78    	0.000422602
78 	79    	0.0112091	78 	0.0133698	0.0107964	79    	0.000558782
79 	73    	0.0110862	79 	0.0131931	0.0107964	73    	0.000465272
80 	67    	0.011036 	80 	0.0127483	0.0107964	67    	0.000386788

Средняя приспособленность падает почти монотонно, а разброс std сначала падает, к пятнадцатому поколению достигает минимума и далее немного растёт. Это не признак того, что популяция нашла несколько разных хороших решений, а хвост неудачных мутантов. Вероятность мутации здесь высокая, и часть потомков каждого поколения заведомо оказывается непригодной.

Извлечём десять лучших особей и сохраним историю эволюции в массивы.

top = tools.selBest(pop, k=10)
gen = np.array(logbook.select("gen"))
fit_min = np.array(logbook.chapters["fitness"].select("min"))
fit_max = np.array(logbook.chapters["fitness"].select("max"))
fit_avg = np.array(logbook.chapters["fitness"].select("avg"))

Лучшая особь берётся из популяции, её поля расставляются по соленоидам, и огибающая рассчитывается ещё раз, уже для графика. Заданная кривая накладывается поверх расчётной, чтобы была видна степень их совпадения.

for x in range(Sol_Num):
            accelerator.Bz_beamline['Sol. %.d'%(x+1)].max_field  = top[0][x]
        
accelerator.compile()
simulation = kv.Simulation(beam, accelerator)
simulation.track()

z_Bz_gen= hv.Area((accelerator.parameter,accelerator.Bz(accelerator.parameter)), kdims=[dim_z], vdims=[dim_Bz])
z_r_gen = hv.Area(((accelerator.parameter,simulation.envelope_x(accelerator.parameter)*1e3)), kdims=[dim_z], vdims=[dim_r], group='Beam')

(z_r_gen*z_env+z_Bz_gen).cols(1)

Результаты

Девять соленоидов, сто особей в популяции, сто поколений, и алгоритм подогнал огибающую под заданную, ни разу не потребовав понимания того, как поля соленоидов влияют на пучок. Журнал выше напечатан по восьмидесятое поколение, и к нему приспособленность лучшей особи упала с 0.034 до 0.0108, средняя по популяции с 0.037 до 0.011. Обе вышли на плато задолго до конца: последние двадцать напечатанных поколений не дали практически ничего. Дальнейшее улучшение невозможно, поскольку заданная огибающая физически недостижима с большей точностью.

Задача решается генетическим алгоритмом, а не градиентным спуском, по двум причинам, типичным для вычислительной физики.

  • Градиента нет. Целевая функция получается численным интегрированием уравнения огибающей; аналитической производной по полям соленоидов у неё не существует, а численная потребовала бы девяти дополнительных расчётов на каждый шаг.
  • Один расчёт дёшев. KENV решает уравнение огибающей за миллисекунды, на порядки быстрее PIC-кода. Поэтому можно позволить себе десять тысяч расчётов (100 особей × 100 поколений), и поэтому эволюционные методы, расточительные по числу вызовов целевой функции, здесь уместны.

Таким образом, метод оптимизации выбирается не только по виду целевой функции, но и по цене одного её вычисления.

Обратная задача. Восстановление огибающей по эксперименту

Подгонка под заданную огибающую остаётся учебной постановкой, тогда как реальная задача является обратной. Имеются измерения, и требуется определить, какой пучок их породил.

На линейном индукционном ускорителе средств наблюдения внутри тракта нет, поэтому измеряется не сама огибающая, а диаметр пучка на выходном экране при разных значениях поля одного из соленоидов. Получается кривая «размер пучка от силы линзы», и по ней требуется восстановить начальный радиус пучка \(r\), угловой разлёт \(dr/dz\) и нормализованный эмиттанс \(\epsilon_n\).

Постановка для генетического алгоритма меняется незначительно. Особью становится вектор из двенадцати сил соленоидальных линз и трёх начальных условий пучка, а функцией приспособленности — среднеквадратичное отклонение расчётной кривой «диаметр от поля соленоида» от измеренной. Стартовая популяция разыгрывается гауссовыми случайными величинами вокруг ожидаемых значений.

ПараметрСреднее \(\mu\)СКО \(\sigma\)
Поле соленоидов \(B_z\)0.04 Тл0.02 Тл
Начальный радиус \(r\)48 мм24 мм
Угловой разлёт \(dr/dz\)38 мрад20 мрад
Нормализованный эмиттанс \(\epsilon_n\)1150 мм·мрад500 мм·мрад

Размер популяции здесь другой, 1500 особей против 100 в прямой задаче, поскольку размерность пространства поиска выросла с 9 до 15 и популяцию необходимо наращивать соответственно, иначе алгоритм останавливается в локальном минимуме. Задача решается за те же 100 поколений.

Для пучка с током 1.5 кА восстановленные значения приведены ниже.

$$ r = 42 \div 49\ \text{мм}, \qquad \frac{dr}{dz} = 35 \div 38\ \text{мрад}, \qquad \epsilon_n = 1125 \div 1215\ \text{мм}\cdot\text{мрад}. $$

Это не погрешности измерения. Это область эквивалентности, множество разных наборов начальных условий, дающих одинаково хорошее согласие с экспериментом в пределах его точности. Единственного решения у обратной задачи здесь нет, и генетический алгоритм это показывает, выдавая не точку, а облако.

При решении обратной задачи необходимо проверять единственность решения и, если она отсутствует, приводить в качестве результата область, а не точку. Алгоритм оптимизации всегда возвращает одно число; насколько оно определено данными, обязан понимать исследователь. Запуск оптимизации из нескольких разных начальных приближений и оценка разброса ответов являются проверкой, занимающей минуты.

Проверка результата

Восстановленные начальные условия были подставлены обратно в модель, а полученная огибающая сравнена с результатами двух других программ, PIC-кода ASTRA и программы UltraSAM. На пятнадцатиметровом ускорительном тракте огибающие совпали.

KENV выполняет расчёт настолько быстрее ASTRA и UltraSAM, что поверх него удалось построить интерактивный интерфейс для настройки проводки пучка в реальном времени. На установке это означает меньшее число тестовых импульсов на настройку, больший ресурс установки и больше времени на физику.

Исходный код KENV опубликован как библиотека для Python и зарегистрирован в Реестре программ для ЭВМ (№ 2024611244 от 18.01.2024).