diff --git a/examples/main.py b/examples/main.py index 4a24c71..b6dbdfb 100644 --- a/examples/main.py +++ b/examples/main.py @@ -3,33 +3,36 @@ import pickle import warnings +import numpy as np +import scipy.constants as const + from astracore import * from plotplugins import * from gaoptimizer import * -beamline = '/home/tangcx/projects/twofreqrfgun/beamline' # the beamline folder -simroot = '/home/tangcx/WORK/data' # the root simulation folder +beamline = '/home/tangcx/projects/ttxmod/beamline' # the beamline folder +simroot = '/home/tangcx/WORK/zhangzhe/data' # the root simulation folder # Constants -eta = 0.48 # sigma ratio of the truncated gaussian dist -norm_emit = 0.9 # mm.mrad/mm, normalized emittance -z_stop = 11.0 # m, end position -ntot = 10000 # number of macro particles -e_gun = 120 # MV/m, gun gradient +z_drift = 0.1 # m, end position relative to the linac exit +z_stop = 5.1 # m, end position of the simulation +qtot = 0.2 # nC, total charge +ntot = 50000 # number of macro particles phi_l1 = 0 # deg, linac 1 phase -phi_l2 = 0 # deg, linac 2 phase +shift_lsol = 0.1 # m, position shift of the linac solenoid relative to the linac +e_photon = 4.66 # eV, photon energy +phi_w = 4.31 # eV, work function # Variables LT = [1e-3, 20e-3] # ns, laser pulse length SIGX = [0.1, 1.0] # mm, laser radius (gaussian cut at 1 sigma) PSOL = [0.21, 0.3] # m, gun solenoid position BSOL = [0.1, 0.3] # T, gun solenoid strength +EGUN = [80, 150] # MV/m, gun gradient PHIGUN = [-20, 20] # deg, gun launch phase -PHICAV = [0, 360] # deg, high order cavity phase -ECAV = [-10, 100] # MV/m, high order cavity gradient EL1 = [0, 35] # MV/m, linac 1 gradient -EL2 = [0, 35] # MV/m, linac 2 gradient -PL1 = [1, 3] # m, linac 1 position +PL1 = [1, 2] # m, linac 1 position +BLSOL = [0, 0.3] # T, linac solenoid strength def recover(x, bound): @@ -41,29 +44,33 @@ def gen_patch(x): sig_x = recover(x[1], SIGX) pos_sol = recover(x[2], PSOL) b_sol = recover(x[3], BSOL) - phi_gun = recover(x[4], PHIGUN) - phi_cav = recover(x[5], PHICAV) - e_cav = recover(x[6], ECAV) - e_l1 = recover(x[7], EL1) - e_l2 = recover(x[8], EL2) - pos_l1 = recover(x[9], PL1) - pos_l2 = pos_l1 + 4 + e_gun = recover(x[4], EGUN) + phi_gun = recover(x[5], PHIGUN) + e_l1 = recover(x[6], EL1) + pos_l1 = recover(x[7], PL1) + b_lsol = recover(x[8], BLSOL) + pos_lsol = pos_l1 + shift_lsol + + phi_sch = float(np.sqrt(const.e * e_gun / (4 * const.pi * const.epsilon_0)) * 1e3) + phi_eff = phi_w - phi_sch patch = {'input': {'lt': lt, 'sig_x': sig_x, - 'nemit_x': sig_x * eta * norm_emit, - 'ipart': ntot}, + 'ipart': ntot, + 'q_total': qtot, + 'phi_eff': phi_eff, + 'e_photon': e_photon}, 'newrun': {'auto_phase': True, 'zstop': z_stop}, 'cavity': {'lefield': True, - 'maxe': [e_gun, e_cav, e_l1, e_l2], - 'phi': [phi_gun, phi_cav, phi_l1, phi_l2], - 'c_pos': [0, 0, pos_l1, pos_l2]}, + 'maxe': [e_gun, e_l1], + 'phi': [phi_gun, phi_l1], + 'c_pos': [0, pos_l1]}, 'charge': {'lspch': True, 'lmirror': True}, 'solenoid': {'lbfield': True, - 'maxb': [b_sol], - 's_pos': [pos_sol]}} + 'maxb': [b_sol, b_lsol], + 's_pos': [pos_sol, pos_lsol]}} return patch @@ -73,13 +80,20 @@ def evaluate(x, sim): try: core.run(patch, sim) data = core.get_data(sim, -1, 'g') - fitness = [nemit_t(data)[0], current_r(data), skewness(data)] - if pnum(data) < 0.9 * ntot: # a small penalty! - fitness[0] += 1 - fitness[1] -= 100 - fitness[2] += 1 + emitx = nemit_t(data)[0] + sig = sig_z(data) + fitness = [emitx, sig] + # Constrains + if emitx > 1: # emittance too large + fitness[0] += (emitx - 1) ** 2 + if sig > 1.2: # current too low + fitness[1] += (sig - 1.2) ** 2 + elif sig < 0.2: # current too high + fitness[0] += (1 / sig - 5) ** 2 + if pnum(data) < 0.9 * ntot: # particle loss + fitness = [100, 100] # almost death penalty except: - fitness = [100, 0, 0] # almost death penalty + fitness = [100, 100] # almost death penalty return fitness @@ -103,14 +117,10 @@ def parse_ga_input(arg1, arg2): print('Evolving based on the previous {0} generations from ghist file {1}...'.format(_ngen, arg1)) except FileNotFoundError as exc: raise OptimizerError('error in reading ghist file {}!'.format(arg1)) from exc - except: - raise try: ngen = int(arg2) except ValueError as exc: raise OptimizerError('error in parsing number of generations!') from exc - except: - raise return npop, ngen @@ -120,8 +130,8 @@ def parse_ga_input(arg1, arg2): # Optimizer setup opt = NSGAII(evaluate) -opt.NDIM = 10 -opt.OBJ = (-1, 1, -1) +opt.NDIM = 9 +opt.OBJ = (-1, -1) opt.setup() if __name__ == '__main__': diff --git a/examples/test.py b/examples/test.py new file mode 100644 index 0000000..b6dbdfb --- /dev/null +++ b/examples/test.py @@ -0,0 +1,168 @@ +import os +import sys +import pickle +import warnings + +import numpy as np +import scipy.constants as const + +from astracore import * +from plotplugins import * +from gaoptimizer import * + +beamline = '/home/tangcx/projects/ttxmod/beamline' # the beamline folder +simroot = '/home/tangcx/WORK/zhangzhe/data' # the root simulation folder + +# Constants +z_drift = 0.1 # m, end position relative to the linac exit +z_stop = 5.1 # m, end position of the simulation +qtot = 0.2 # nC, total charge +ntot = 50000 # number of macro particles +phi_l1 = 0 # deg, linac 1 phase +shift_lsol = 0.1 # m, position shift of the linac solenoid relative to the linac +e_photon = 4.66 # eV, photon energy +phi_w = 4.31 # eV, work function + +# Variables +LT = [1e-3, 20e-3] # ns, laser pulse length +SIGX = [0.1, 1.0] # mm, laser radius (gaussian cut at 1 sigma) +PSOL = [0.21, 0.3] # m, gun solenoid position +BSOL = [0.1, 0.3] # T, gun solenoid strength +EGUN = [80, 150] # MV/m, gun gradient +PHIGUN = [-20, 20] # deg, gun launch phase +EL1 = [0, 35] # MV/m, linac 1 gradient +PL1 = [1, 2] # m, linac 1 position +BLSOL = [0, 0.3] # T, linac solenoid strength + + +def recover(x, bound): + return bound[0] + (bound[1] - bound[0]) * x + + +def gen_patch(x): + lt = recover(x[0], LT) + sig_x = recover(x[1], SIGX) + pos_sol = recover(x[2], PSOL) + b_sol = recover(x[3], BSOL) + e_gun = recover(x[4], EGUN) + phi_gun = recover(x[5], PHIGUN) + e_l1 = recover(x[6], EL1) + pos_l1 = recover(x[7], PL1) + b_lsol = recover(x[8], BLSOL) + pos_lsol = pos_l1 + shift_lsol + + phi_sch = float(np.sqrt(const.e * e_gun / (4 * const.pi * const.epsilon_0)) * 1e3) + phi_eff = phi_w - phi_sch + + patch = {'input': {'lt': lt, + 'sig_x': sig_x, + 'ipart': ntot, + 'q_total': qtot, + 'phi_eff': phi_eff, + 'e_photon': e_photon}, + 'newrun': {'auto_phase': True, + 'zstop': z_stop}, + 'cavity': {'lefield': True, + 'maxe': [e_gun, e_l1], + 'phi': [phi_gun, phi_l1], + 'c_pos': [0, pos_l1]}, + 'charge': {'lspch': True, + 'lmirror': True}, + 'solenoid': {'lbfield': True, + 'maxb': [b_sol, b_lsol], + 's_pos': [pos_sol, pos_lsol]}} + + return patch + + +def evaluate(x, sim): + patch = gen_patch(x) + try: + core.run(patch, sim) + data = core.get_data(sim, -1, 'g') + emitx = nemit_t(data)[0] + sig = sig_z(data) + fitness = [emitx, sig] + # Constrains + if emitx > 1: # emittance too large + fitness[0] += (emitx - 1) ** 2 + if sig > 1.2: # current too low + fitness[1] += (sig - 1.2) ** 2 + elif sig < 0.2: # current too high + fitness[0] += (1 / sig - 5) ** 2 + if pnum(data) < 0.9 * ntot: # particle loss + fitness = [100, 100] # almost death penalty + except: + fitness = [100, 100] # almost death penalty + return fitness + + +def parse_ga_input(arg1, arg2): + try: + npop = int(arg1) + except: + try: + with open(arg1, 'rb') as f: + _ngen = 0 + while True: + try: + npop = pickle.load(f) + _ngen += 1 + except EOFError: + break + except: + raise + if not _ngen: + raise OptimizerError('ghist file {} is empty!'.format(arg1)) + print('Evolving based on the previous {0} generations from ghist file {1}...'.format(_ngen, arg1)) + except FileNotFoundError as exc: + raise OptimizerError('error in reading ghist file {}!'.format(arg1)) from exc + try: + ngen = int(arg2) + except ValueError as exc: + raise OptimizerError('error in parsing number of generations!') from exc + + return npop, ngen + + +# Core initialization +core = AstraCore(beamline) + +# Optimizer setup +opt = NSGAII(evaluate) +opt.NDIM = 9 +opt.OBJ = (-1, -1) +opt.setup() + +if __name__ == '__main__': + # Parse the input arguments + nargs = len(sys.argv[1:]) + if not nargs: + raise OptimizerError('initial population (number) and number of generations are needed!') + elif nargs == 1: + fcount = 1 + while os.path.exists(HISTFNAME + (' {:d}'.format(fcount) if fcount > 1 else '')): + fcount += 1 + fcount -= 1 + ghist = HISTFNAME + (' {:d}'.format(fcount) if fcount > 1 else '') + npop, ngen = parse_ga_input(ghist, sys.argv[1]) + else: + npop, ngen = parse_ga_input(sys.argv[1], sys.argv[2]) + if nargs > 2: + _simroot = os.path.join(core.root, sys.argv[3]) + if os.path.exists(_simroot): + simroot = _simroot + else: + warnings.warn('customized sim root folder not found, fallback to the default one!') + + # Let's rock! + opt.evolve(npop, ngen, pre=simroot) + + # Record + with open('gapop', 'w') as f: + f.write(str(opt.pop)) + with open('galog', 'w') as f: + f.write(str(opt.log)) + with open('gafit', 'w') as f: + fit = [ind.fitness.values for ind in opt.pop] + f.write(str(fit))