Skip to content

Commit

Permalink
Compare to MAML environment
Browse files Browse the repository at this point in the history
  • Loading branch information
@jank324
jank324 committed Nov 18, 2023
1 parent 3589e32 commit dcd6319
Show file tree
Hide file tree
Showing 10 changed files with 1,289 additions and 13 deletions.
125 changes: 125 additions & 0 deletions compare/simon_maml/electron_design.mad
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
!********************************************
! TT43-TT41 (AWAKE e-line) model
!
! F.M. Velotti, C. Bracco
! MADX version MAD-X 5.04.01
!********************************************

title, 'AWAKE e-line';

option, echo;
option, RBARC=FALSE;


!system, "ln -fns /opt/vkain/PycharmProjects/awakett43/2018 tt43_repo";

!*********************************************
! Import sequence, elements and strengths
! N.B. FINT values calulated empirically from
! beam-based measurements
!*********************************************

call, file = "tt43_repo/elements/electron_elements.ele";
call, file = "tt43_repo/sequence/electron_sequence.seq";
call, file = "tt43_repo/strength/electron_strength.str";

call, file = "tt43_repo/beams/electron_beam.beamx";

fint_h_var = 0.25;
fint_v_var = 0.09;
!***********************************************
! Select strengths for needed optics
!***********************************************

//Focus at BTV50
//kqf430205 := -7.74315 ;
//kqd412344 := -5.08493 ;
//kqf412346 :=21.0534;
//kqd412347 := -25.2612 ;


//Focus at BTV53
//kqf430205 := -7.83841 ;
//kqd412344 := - 3.71018;
//kqf412346 := 18.396;
//kqd412347 := -16.3205;


//Focus at BTV54
//kqf430205 := -7.70116 ;
//kqd412344 := -3.90342 ;
//kqf412346 := 17.9845;
//kqd412347 := -15.1748;

//Focus at Iris +1 m
//kqf430205 := -7.64757 ;
//kqd412344 := -5.91917 ;
//kqf412346 := 17.6078;
//kqd412347 := -13.3947;

//Focus at Iris +2 m
//kqf430205 := -7.69869 ;
//kqd412344 := -6.32961;
//kqf412346 := 17.3987 ;
//kqd412347 := -12.7635;


//Focus at Iris +3 m
//kqf430205 := -7.6056 ;
//kqd412344 := -6.02714;
//kqf412346 := 17.1487;
//kqd412347 := -12.5029;

//Focus at Iris +4 m
//kqf430205 := -7.75767;
//kqd412344 := -3.12061;
//kqf412346 := 16.2752;
//kqd412347 := -12.1149;

//Focus at Iris +5 m
//kqf430205 := -7.61516;
//kqd412344 := -4.42116;
//kqf412346 := 16.4406;
//kqd412347 := -11.8455;


!***************************************
! Default optics
!***************************************

//Focus at Iris
kqf430205 := -7.92377;
kqd412344 := -1.07907;
kqf412346 := 16.7789;
kqd412347 := -14.0290;


use, sequence= TT43;

set, format="15.9f";

! Normalised emittance about 7 um in both planes
emittx:=1.93e-7;
emitty:=1.93e-7;

! RMS momentum spread (estimated from measurements)
dpp = 8e-3;

!Spot Size [m]
sigma_x := sqrt((table(twiss, betx)*emittx) + ((table(twiss,dx))*dpp)^2);
sigma_y := sqrt((table(twiss, bety)*emitty) + ((table(twiss,dy))*dpp)^2);


! Measured ini conditions: beta_x = 3.1 m, alfa_x = 0.8, beta_y = 3.5 m, alfa_y = 0.8

Select,flag=twiss, clear;
Select,flag=twiss,column=name,keyword,s,l,sigma_x,sigma_y,betx,bety,dx, dpx, dy, dpy, alfx,alfy,x,y,k1l,mux,muy;
//Select,flag=twiss,column=name,keyword,s,l,aper_1,aper_2;
TWISS, file="tt43_repo/output/electron_tt43.out",

! Ideal initial conditions
initbeta0:beta0,DELTAP=0.0,BETX=5,ALFX=0,DX=0,DPX=0,
BETY=5,ALFY=0,DY=0.0,DPY=0.0,
x=0,px=0,y=0,py=0;

!stop;
34 changes: 21 additions & 13 deletions compare/simon_maml/environment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,22 @@
import random
from enum import Enum

import gym
import gymnasium as gym
import numpy as np
import pandas as pd
from cpymad.madx import Madx
from gym import spaces
from gymnasium import spaces


# Standard environment for the AWAKE environment, adjusted so it can be used for the MAML therefore containing functions
# for creating and sampling tasks
# Standard environment for the AWAKE environment, adjusted so it can be used for the
# MAML therefore containing functions for creating and sampling tasks
def generate_optics():
OPTIONS = ["WARN"] # ['ECHO', 'WARN', 'INFO', 'DEBUG', 'TWISS_PRINT']
MADX_OUT = [f"option, -{ele};" for ele in OPTIONS]

madx = Madx(stdout=False)
madx.input("\n".join(MADX_OUT))

tt43_ini = "maml_rl/envs/electron_design.mad"
tt43_ini = "../compare/simon_maml/electron_design.mad"

madx.call(file=tt43_ini, chdir=True)

Expand All @@ -34,7 +33,7 @@ def generate_optics():
madx.globals.update(quads)

madx.input(
"initbeta0:beta0,BETX=5,ALFX=0,DX=0,DPX=0,BETY=5,ALFY=0,DY=0.0,DPY=0.0,x=0,px=0,y=0,py=0;"
"initbeta0:beta0,BETX=5,ALFX=0,DX=0,DPX=0,BETY=5,ALFY=0,DY=0.0,DPY=0.0,x=0,px=0,y=0,py=0;" # noqa: E501
)
twiss_cpymad = madx.twiss(beta0="initbeta0").dframe()

Expand Down Expand Up @@ -117,7 +116,8 @@ def recalculate_response(self):
self.twiss_bpms, self.twiss_correctors, self.plane
)

# MAML specific function, while training samples fresh new tasks and for testing it uses previously saved tasks
# MAML specific function, while training samples fresh new tasks and for testing it
# uses previously saved tasks
def sample_tasks(self, num_tasks):
goals = []
print("Number of tasks: " + str(num_tasks))
Expand All @@ -141,7 +141,9 @@ def sample_tasks(self, num_tasks):
goals = pickle.load(input_file)

# Code used to create sample Tasks
# with open("../AWAKE_MB_META_RL/meta_environment/Tasks_data/Tasks", "wb") as fp: # Pickling
# with open(
# "../AWAKE_MB_META_RL/meta_environment/Tasks_data/Tasks", "wb"
# ) as fp: # Pickling
# pickle.dump(goals, fp)
# print("Saved Tasks")

Expand Down Expand Up @@ -177,7 +179,13 @@ def step(self, action):
if (return_reward > self.threshold) or self.current_steps > self.MAX_TIME:
self.is_finalized = True

return return_state, return_reward, self.is_finalized, {"task": self._id}
return (
return_state.astype(np.float32),
return_reward,
self.is_finalized,
False,
{"task": self._id},
)

def step_opt(self, action):
state, reward = self._take_action(action, is_optimisation=True)
Expand Down Expand Up @@ -261,11 +269,11 @@ def reset(self, **kwargs):
self.kicks_0 = self.settingsV * self.action_scale

if self.plane == Plane.horizontal:
init_positions = np.zeros(len(self.positionsH)) # self.positionsH
# init_positions = np.zeros(len(self.positionsH)) # self.positionsH
rmatrix = self.responseH

if self.plane == Plane.vertical:
init_positions = np.zeros(len(self.positionsV)) # self.positionsV
# init_positions = np.zeros(len(self.positionsV)) # self.positionsV
rmatrix = self.responseV

self.current_episode += 1
Expand All @@ -285,7 +293,7 @@ def reset(self, **kwargs):
return_initial_state = np.array(state * self.state_scale)
self.initial_conditions.append([return_initial_state])
return_value = return_initial_state
return return_value
return return_value.astype(np.float32), {"task": self._id}

def seed(self, seed=None):
random.seed(seed)
Expand Down
8 changes: 8 additions & 0 deletions compare/simon_maml/tt43_repo/beams/electron_beam.beamx
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
!****************************
! Beam specification
!****************************

!option,debug,-echo;
!sigt=TBUNCH*CLIGHT*(BEAM->PC/BEAM->ENERGY)
BEAM, particle=electron, PC=19e-3, exn=10e-6, eyn=10e-6, npart=3e9, bunched, sigt=TBUNCH*CLIGHT*(BEAM->PC/BEAM->ENERGY);

76 changes: 76 additions & 0 deletions compare/simon_maml/tt43_repo/elements/electron_elements.ele
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
!-------------------------------------------------------
! Elements
!-------------------------------------------------------

//-------------- Bending Magnets
MBH: RBEND, L := 0.1755 , apertype = ellipse , aperture = {0.02,0.02}, FINT := fint_h_var, hgap = 0.035, k1 := mbh_k1 ; ! Bending magnet, horizontal

MBV: RBEND, L := 0.1755 , apertype = ellipse , aperture = {0.02,0.02}, FINT := fint_v_var, hgap = 0.035 ; ! Bending Magnet, vertical

!specification magnetic lenght 0.1755
!specification mechanical lenght 0.294


//-------------- Quadrupoles
QUAD: QUADRUPOLE, L := 0.0708, apertype = ellipse, aperture = {0.02,0.02};
!specification magnetic length 0.0718, but the measured gradient correspond to
!a magnetic lenght slightly smaller, i.e. 0.0708 m => this is very important to take
!into account



//-------------- Correctors
COR: KICKER, L := 0.0406, apertype = ellipse, aperture = {0.02,0.02};
!COR : KICKER , L := 0.1651; !magnetic length
!specification magnetic length 0.1654
!specification mechanical lenght 0.073


//-------------- Beam position monitors
BPM : MONITOR , L := 0, apertype = ellipse, aperture = {0.02,0.02};
BPMShort : MONITOR , L := 0.207, apertype = ellipse, aperture = {0.02,0.02};
BPMLong : MONITOR , L := 0.207, apertype = ellipse, aperture = {0.02,0.02};


//-------------- Beam observation TV monitors based on screens

BTV : INSTRUMENT , L := 0.273, apertype = ellipse, aperture = {0.02,0.02};
BTV2 : INSTRUMENT , L := 0.354, apertype = ellipse, aperture = {0.02,0.02};



//-------------- Markers
Start_line : MARKER , L := 0;
End_line : MARKER , L := 0; ! apertype = rectangle, aperture = {0.00025,0.00025};
Mon : INSTRUMENT , L := 0;

//-------------- BeamStopper
Beam_Stopper : INSTRUMENT , L := 0.15; !0.069775;



//-------------- Elements proton and secondary beam line
bpg412445 : INSTRUMENT , L := 0.248 , aperture = ellipse, aperture = {0.060,0.060}; !old 0.248
otr2 : INSTRUMENT , L := 0.35 , aperture = ellipse, aperture = {0.06,0.06};

plasmacell: INSTRUMENT , L = 10 , aperture=ellipse ,aperture={0.020 ,0.020} ;
iris: INSTRUMENT , l = 0.001 , aperture=ellipse ,aperture={0.005 ,0.005} ;

table: INSTRUMENT , L = 2 ;

MQNBR : QUADRUPOLE , L := 0.285;


MBXFB : SBEND , L := 1;


BTVSPS : INSTRUMENT , L = 0.448;

HEATER : SOLENOID , L = 0.025 , KS = 0.00017/0.033;







Loading

0 comments on commit dcd6319

Please sign in to comment.