def setUp(self): self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> ')) here = os.path.dirname(__file__) there = os.path.join(here, 'testseq.madx') self.doc = open(there).read() for line in self.doc.splitlines(): self.mad._libmadx.input(line)
def test_zero_errors(): # check that zero-errors are loaded without error cpymad_spec = util.find_spec("cpymad") if cpymad_spec is None: print("cpymad is not available - abort test") sys.exit(0) from cpymad.madx import Madx madx = Madx() madx.input(''' qd: multipole, knl={0,-0.3}; qf: multipole, knl={0, 0.3}; testseq: sequence, l = 1; qd, at = 0.3; qf, at = 0.6; endsequence; ''') madx.beam() madx.use('testseq') madx.select(flag='error', pattern='qf') madx.command.efcomp(dkn=[0, 0, 0, 0, 0.0, 0.0, 0.0], dks=[0.0, 0.0, 0, 0]) madx.command.ealign(dx=0.0, dy=0.0, ds=0.0, DPHI=0.0, DTHETA=0.0, DPSI=0.0, MREX=0.0, MREY=0.0, MSCALX=0.0, MSCALY=0.0, AREX=0.0, AREY=0.0)
def make_lhc_beams(madx: Madx, energy: float = 6500, emittance: float = 3.75e-6, **kwargs) -> None: """ Define beams with default configuratons for `LHCB1` and `LHCB2` sequences. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. energy (float): beam energy in GeV. Defaults to 6500. emittance (float): emittance in meters, which will be used to calculate geometric emittance, then fed to the BEAM command. Keyword Args: Any keyword argument that can be given to the MAD-X BEAM command. """ logger.info("Making default beams for 'lhcb1' and 'lhbc2' sequences") madx.globals["NRJ"] = energy madx.globals["brho"] = energy * 1e9 / madx.globals.clight geometric_emit = madx.globals["geometric_emit"] = emittance / (energy / 0.938) for beam in (1, 2): logger.trace(f"Defining beam for sequence 'lhcb{beam:d}'") madx.command.beam( sequence=f"lhcb{beam:d}", particle="proton", bv=1 if beam == 1 else -1, energy=energy, npart=1.0e10, ex=geometric_emit, ey=geometric_emit, **kwargs, )
def apply_lhc_colinearity_knob(madx: Madx, colinearity_knob_value: float = 0, ir: int = None) -> None: """ Applies the LHC colinearity knob. If you don't know what this is, you should not be using this function. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. colinearity_knob_value (float): Units of the colinearity knob to apply. Defaults to 0 so users don't mess up local coupling by mistake. This should be a positive integer, normally between 1 and 10. ir (int): The Interaction Region to apply the knob to, should be one of [1, 2, 5, 8]. Classically 1 or 5. """ logger.info( f"Applying Colinearity knob with a unit setting of {colinearity_knob_value}" ) logger.warning("You should re-match tunes & chromaticities after this") knob_variables = (f"KQSX3.R{ir:d}", f"KQSX3.L{ir:d}" ) # MQSX IP coupling correctors right_knob, left_knob = knob_variables madx.globals[right_knob] = colinearity_knob_value * 1e-4 logger.trace(f"Set '{right_knob}' to {madx.globals[right_knob]}") madx.globals[left_knob] = -1 * colinearity_knob_value * 1e-4 logger.trace(f"Set '{left_knob}' to {madx.globals[left_knob]}")
def reset(self): """ If MAD-X fails, re-spawn process """ self.madx = Madx(stdout=False) self.madx.option(echo=False, warn=False, info=False, debug=False, verbose=False) self.madx.input('BEAM, PARTICLE=PROTON, PC = 2.14') self.madx.input('BRHO := BEAM->PC * 3.3356;') self.madx.call(file='ps_mu.seq') self.madx.call(file='ps_ss_mod.seq') self.madx.call(file='ps_50LeQ.str') self.madx.call(file='ps_pro_bare_machine.str') self.madx.call(file='remove_elements.seq') self.madx.input('seqedit, sequence = PS;') self.madx.input('select, flag=seqedit, class = MQNAAIAP;') self.madx.input('select, flag=seqedit, class = MQNABIAP;') self.madx.input('select, flag=seqedit, class = MQSAAIAP;') self.madx.input('select, flag=seqedit, class = QNSD;') self.madx.input('select, flag=seqedit, class = QNSF;') self.madx.input('use, sequence = PS;') self.madx.input('seqedit,sequence = PS;flatten;endedit;') self.madx.input( 'seqedit,sequence = PS;remove, element=SELECTED;endedit;') self.madx.input('endedit;') self.madx.input('use, sequence = PS;')
def madxp(inputFile, pythonData=None, outputDF='output.pkl', command_log_file='log.madx', stdout_file='stdout.madx', verbose=False): ''' It runs the MADX dataframe using the MADX extended syntax. inputFile: the MADX input file. outputDF: the file to dump the output DF. The MADX variable and the pythonData will be available for all code sections. command_log_file: the filename of the logging file. Use the None variable not to log. stdout_file: the filename of the file to redirect the stdout. Use the None variable not to log. verbose: boolean flag to have verbose output during the execution. ''' if command_log_file==None: if stdout_file==None: madx = Madx() else: with open(stdout_file, 'w') as f: madx = Madx(stdout=f) else: if stdout_file==None: madx = Madx(command_log=command_log_file) else: with open(stdout_file, 'w') as f: madx = Madx(stdout=f,command_log=command_log_file) aux=df2run(madx, madx2df(inputFile), pythonData=pythonData, verbose=verbose) if outputDF!=None: aux.to_pickle(outputDF) print('Profiling DF saved.')
def test_independent_instances(): # Check independence by defining a variable differently in each # instance: with Madx(prompt='X1:> ') as mad1, Madx(prompt='X2:> ') as mad2: mad1.input('ANSWER=42;') mad2.input('ANSWER=43;') assert mad1.eval('ANSWER') == 42 assert mad2.eval('ANSWER') == 43
def test_independent_instances(self): # create a second Madx instance (1st one is created in setUp) madxness = Madx() # Check independence by defining a variable differently in each # instance: self.mad.input('ANSWER=42;') madxness.input('ANSWER=43;') self.assertEqual(self.mad.evaluate('ANSWER'), 42); self.assertEqual(madxness.evaluate('ANSWER'), 43);
def test_independent_instances(self): # create a second Madx instance (1st one is created in setUp) madxness = Madx() # Check independence by defining a variable differently in each # instance: self.mad.input('ANSWER=42;') madxness.input('ANSWER=43;') self.assertEqual(self.mad.evaluate('ANSWER'), 42) self.assertEqual(madxness.evaluate('ANSWER'), 43)
def read_tfsfile(filename): """ Read TFS table. """ from cpymad.madx import Madx global TFS_READER if not TFS_READER: TFS_READER = Madx() TFS_READER.command.readmytable(file=abspath(filename), table='user') return TFS_READER.get_table('user')
def setUp(self): self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> ')) here = os.path.dirname(__file__) there = os.path.join(here, 'testseq.madx') with open(there) as f: self.doc = f.read() for line in self.doc.splitlines(): line = line.split('!')[0].strip() if line: self.mad._libmadx.input(line)
def apply_lhc_rigidity_waist_shift_knob(madx: Madx, rigidty_waist_shift_value: float = 0, ir: int = None, side: str = "left") -> None: """ Applies the LHC rigidity waist shift knob, moving the waist left or right of IP. If you don't know what this is, you should not be using this function. The waist shift is done by unbalancing the triplet powering knob of the left and right-hand sides of the IP. Warning: Applying the shift will modify your tunes and most likely flip them, making a subsequent matching impossible if your lattice has coupling. To avoid this, match to tunes split further apart before applying the waist shift knob, and then match to the desired working point. For instance for the LHC, matching to (62.27, 60.36) before applying and afterwards rematching to (62.31, 60.32) usually works well. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. rigidty_waist_shift_value (float): Units of the rigidity waist shift knob (positive values only). ir (int): The Interaction Region to apply the knob to, should be one of [1, 2, 5, 8]. Classically 1 or 5. side (str): Which side of the IP to move the waist to, determines a sign in the calculation. Defaults to 'left', which means s_waist < s_ip (and setting it to 'right' would move the waist to s_waist > s_ip). """ logger.info( f"Applying Rigidity Waist Shift knob with a unit setting of {rigidty_waist_shift_value}" ) logger.warning("You should re-match tunes & chromaticities after this") right_knob, left_knob = f"kqx.r{ir:d}", f"kqx.l{ir:d}" # IP triplet default knob (no trims) current_right_knob = madx.globals[right_knob] current_left_knob = madx.globals[left_knob] if side.lower() == "left": madx.globals[right_knob] = ( 1 - rigidty_waist_shift_value * 0.005) * current_right_knob madx.globals[left_knob] = ( 1 + rigidty_waist_shift_value * 0.005) * current_left_knob elif side.lower() == "right": madx.globals[right_knob] = ( 1 + rigidty_waist_shift_value * 0.005) * current_right_knob madx.globals[left_knob] = ( 1 - rigidty_waist_shift_value * 0.005) * current_left_knob else: logger.error( f"Given side '{side}' invalid, only 'left' and 'right' are accepted values." ) raise ValueError("Invalid value for parameter 'side'.") logger.trace(f"Set '{right_knob}' to {madx.globals[right_knob]}") logger.trace(f"Set '{left_knob}' to {madx.globals[left_knob]}")
def test_command_log(self): """Check that the command log contains all input commands.""" # create a new Madx instance that uses the history feature: history_filename = '_test_madx.madx.tmp' mad = Madx(command_log=history_filename) # feed some input and compare with history file: for line in self.doc.splitlines(): mad.input(line) with open(history_filename) as history_file: history = history_file.read() self.assertEqual(history.strip(), self.doc.strip()) # remove history file del mad os.remove(history_filename)
def main(): np.set_printoptions( **{ 'precision': 5, 'suppress': True, # no scientific notation 'linewidth': 120, }) # load gantry model m = Madx(stdout=False) m.call('../hit_models/hht3/run.madx', True) d_kick = 0.1e-3 # define elements els = m.sequence.hht3.expanded_elements H = els.index('h1ms2h') V = els.index('h1ms1v') M = [ el.index for el in els if el.index > max(H, V) and el.base_name == 'monitor' ] orm_tab_1 = calc_orms(m, H, V, M, d_kick) m.globals.kl_b3qd12 += 1e-7 orm_tab_2 = calc_orms(m, H, V, M, d_kick) m.globals.kl_b3qd12 -= 1e-7 orm_tab = (orm_tab_2 - orm_tab_1) / 1e-7 orm_tab = orm_tab_2 xlabel = [els[m].name for m in M] plt.plot(xlabel, orm_tab[:, 0], 'o', label="sin x") plt.plot(xlabel, orm_tab[:, 2], label="var x") plt.plot(xlabel, orm_tab[:, 4], label="sec x") plt.legend() plt.setp(plt.xticks()[1], rotation=50) plt.show() plt.clf() plt.plot(xlabel, orm_tab[:, 1], 'o', label="sin y") plt.plot(xlabel, orm_tab[:, 3], label="var y") plt.plot(xlabel, orm_tab[:, 5], label="sec y") plt.legend() plt.setp(plt.xticks()[1], rotation=50) plt.show()
def main(): filenames = glob.glob( '../data/orm/2019-01-20_quadscan/M8-E108-F1-I9-G1/*/*_X.yml') raw_data = [yaml.safe_load(read_file(filename)) for filename in filenames] m = Madx(stdout=False) m.verbose() m.call('../hit_models/hht3/sequence.madx', chdir=True) m.command.beam() m.use('hht3') #m.call('../hit_models/hht3/strengths0.madx', chdir=True) for data in raw_data: mean = mean_values(data) err = stddevs(data) kl = np.array([ sum(optics.values()) for optics, group in itertools.groupby(data['records'], key=lambda r: r['optics']) ]) mon = data['monitors'][0] knob, = data['optics'][0].keys() quad = [ elem.name for elem in m.sequence.hht3.expanded_elements if elem.base_name == 'quadrupole' and knob in m.expr_vars(elem.defs.k1) ][0] m.globals.update(data['base_optics']) for i, x in enumerate('xy'): #plt.subplot(2, 1, i+1) plt.title(f"pos{x}_{mon}({knob})") plt.ylabel("{x} [m]") plt.xlabel(f"{knob} [$m^{{-1}}$]") plt.errorbar(kl, mean[:, 0, i], err[:, 0, i], label=x) for pos in np.linspace(-0.002, 0.002, 5): modelled = np.array([ track(m, optics, range=f'{quad}/{mon}', **{x: pos})[i] for optics, group in itertools.groupby( data['records'], key=lambda r: r['optics']) ]) i0 = np.argmin(kl) modelled += mean[i0, 0, i] - modelled[i0] plt.plot(kl, modelled, label=f'${x}_0$={pos}') plt.legend() plt.show()
def switch_magnetic_errors(madx: Madx, **kwargs) -> None: """ INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD). Applies magnetic field orders. This will only work for LHC and HLLHC machines. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. Keyword Args: default: sets global default to this value. Defaults to `False`. AB#: sets the default for all of that order, the order being the # number. A# or B#: sets the default for systematic and random of this id. A#s, B#r etc.: sets the specific value. In all kwargs, the order # should be in the range (1...15), where 1 == dipolar field. """ logger.debug("Setting magnetic errors") global_default = kwargs.get("default", False) for order in range(1, 16): logger.trace(f"Setting up for order {order}") order_default = kwargs.get(f"AB{order:d}", global_default) for ab in "AB": ab_default = kwargs.get(f"{ab}{order:d}", order_default) for sr in "sr": name = f"{ab}{order:d}{sr}" error_value = int(kwargs.get(name, ab_default)) logger.trace( f"Setting global for 'ON_{name}' to {error_value}") madx.globals[f"ON_{name}"] = error_value
def test_command_log(): """Check that the command log contains all input commands.""" # create a new Madx instance that uses the history feature: history_filename = '_test_madx.madx.tmp' try: # feed some input lines and compare with history file: lines = """ l = 5; f = 200; fodo: sequence, refer=entry, l=100; QF: quadrupole, l=5, at= 0, k1= 1/(f*l); QD: quadrupole, l=5, at=50, k1=-1/(f*l); endsequence; beam, particle=proton, energy=2; use, sequence=fodo; """.splitlines() lines = [line.strip() for line in lines if line.strip()] with Madx(command_log=history_filename) as mad: for line in lines: mad.input(line) with open(history_filename) as history_file: history = history_file.read() assert history.strip() == '\n'.join(lines).strip() finally: # remove history file os.remove(history_filename)
def _test_transfer_map(seq, range_, doc, rtol=1e-7, atol=1e-15): with Madx() as mad: mad.input(doc) mad.use(seq) par = ['x', 'px', 'y', 'py', 't', 'pt'] val = [+0.0010, -0.0015, -0.0020, +0.0025, +0.0000, +0.0000] twiss = {'betx': 0.0012, 'alfx': 0.0018, 'bety': 0.0023, 'alfy': 0.0027} twiss.update(zip(par, val)) elems = range_.split('/') smap = mad.sectormap(elems, sequence=seq, **twiss)[-1] tw = mad.twiss(sequence=seq, range=range_, **twiss) # transport of coordinate vector: x_init = np.array(val) x_final_tw = np.array([tw[p][-1] for p in par]) x_final_sm = np.dot(smap, np.hstack((x_init, 1))) assert_allclose(x_final_tw[:4], x_final_sm[:4], rtol=rtol, atol=atol) # transport of beam matrix: tm = smap[0:6, 0:6] tab_len = len(tw['sig11']) sig_init = tw.sigmat(0) sig_final_tw = tw.sigmat(tab_len-1) sig_final_sm = np.dot(tm, np.dot(sig_init, tm.T)) assert_allclose(sig_final_tw[0:4, 0:4], sig_final_sm[0:4, 0:4], rtol=rtol, atol=atol)
def setUp(self): self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> ')) here = os.path.dirname(__file__) there = os.path.join(here, 'data', 'lebt', 'init.madx') self.doc = open(there).read() for line in self.doc.splitlines(): self.mad._libmadx.input(line)
def test_format_command(self): m = self.madx = Madx() self.assertEqual(util.format_command('twiss', sequence='lhc'), 'twiss, sequence="lhc";') self.assertEqual(util.format_command('option', echo=True), 'option, echo=true;') self.assertEqual( util.format_command('constraint', betx=Constraint(max=3.13)), 'constraint, betx<3.13;') self.assertEqual(util.format_command(m.command.quadrupole, k1='pi/2'), 'quadrupole, k1:=pi/2;') self.assertEqual(util.format_command('multipole', knl=[0.0, 1.0, 2.0]), 'multipole, knl={0.0,1.0,2.0};') self.assertEqual(util.format_command('twiss', range=Range('#s', '#e')), 'twiss, range=#s/#e;') self.assertEqual(util.format_command( 'select', class_='quadrupole', ), 'select, class="quadrupole";') self.assertEqual( util.format_command(m.elements.quadrupole, k1="hello + world"), 'quadrupole, k1:=hello + world;') self.assertEqual( util.format_command( # match->sequence parameter is list in MAD-X! m.command.match, sequence="foo"), "match, sequence=foo;") self.assertEqual( util.format_command( m.command.select, class_='quadrupole', ), 'select, class=quadrupole;')
def load_file(cls, filename, madx=None, *, undo_stack=None, interpolate=0, **madx_kwargs): """Load model from .madx or .yml file and pass additional arguments to the Model constructor.""" madx = madx or Madx(**madx_kwargs) madx.option(echo=False) filename = os.path.abspath(filename) path, name = os.path.split(filename) ext = os.path.splitext(name)[1].lower() if ext in ('.yml', '.yaml'): data = yaml.load_file(filename) path = os.path.join(path, data.get('path', '.')) _load_params(data, 'beam', path) _load_params(data, 'twiss', path) for fname in data.get('init-files', []): _call(madx, path, fname) else: _call(madx, path, filename) seqname = _guess_main_sequence(madx) data = _get_seq_model(madx, seqname) data['init-files'] = [filename] return cls(madx, data, undo_stack=undo_stack, filename=filename, interpolate=interpolate)
def apply_lhc_coupling_knob(madx: Madx, coupling_knob: float = 0, beam: int = 1, telescopic_squeeze: bool = False) -> None: """ Applies the LHC coupling knob to reach the desired C- value. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. coupling_knob (float): Desired value for the Cminus, typically a few units of 1E-3. Defaults to 0 so users don't mess up coupling by mistake beam (int): beam to apply the knob to, defaults to beam 1. telescopic_squeeze (bool): if set to True, uses the knobs for Telescopic Squeeze configuration. Defaults to False. """ logger.info("Applying coupling knob") logger.warning("You should re-match tunes & chromaticities after this") suffix = "_sq" if telescopic_squeeze else "" knob_name = f"CMRS.b{beam:d}{suffix}" logger.trace( f"Knob '{knob_name}' is {madx.globals[knob_name]} before implementation" ) madx.globals[knob_name] = coupling_knob logger.trace(f"Set '{knob_name}' to {madx.globals[knob_name]}")
def __init__(self, dataFile, madxModelFile, profilePath): self.madxModelFile = madxModelFile self.dataFile = dataFile self.profilePath = profilePath self.data = self.readData(dataFile) self.monitor = self.getMonitor() self.kickers, self.kicks = self.getKicks() self.sequence = self.getSequence() self.madx = Madx(stdout=False) self.madx.call(file=self.madxModelFile, chdir=True) # This are the initial conditions for the Twiss Module of MAD-X # there doesn't seem to be a strong dependence on them self.dx = 1.0e-4 self.dpx = 1.0e-6 self.dy = 1.0e-4 self.dpy = 1.0e-6
def test_context_manager(): output = [] with Madx(stdout=output.append) as m: m.input('foo = 3;') assert m.globals.foo == 3 assert b'+ MAD-X finished normally ' in output[-1] assert not bool(m) with raises(RuntimeError): m.input(';')
def test_context_manager(self): output = [] with Madx(stdout=output.append) as m: m.input('foo = 3;') self.assertEqual(m.globals.foo, 3) self.assertIn(b'+ MAD-X finished normally ', output[-1]) self.assertFalse(bool(m)) with self.assertRaises(RuntimeError): m.input(';')
def test_append_semicolon(): """Check that semicolon is automatically appended to input() text.""" # Regression test for #73 log = [] with Madx(command_log=log.append) as mad: mad.input('a = 0') mad.input('b = 1') assert log == ['a = 0;', 'b = 1;'] assert mad.globals.a == 0 assert mad.globals.b == 1
def match_no_coupling_through_ripkens( madx: Madx, sequence: str = None, location: str = None, vary_knobs: Sequence[str] = None) -> None: """ Matching routine to get cross-term Ripken parameters beta_12 and beta_21 to be 0 at a given location. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. sequence (str): name of the sequence to activate for the matching. location (str): the name of the element at which one wants the cross-term Ripkens to be 0. vary_knobs (Sequence[str]): the variables names to 'vary' in the MADX routine. """ logger.info( f"Matching Ripken parameters for no coupling at location {location}") logger.debug("Creating macro tu update Ripkens") madx.input("do_ripken: macro = {twiss, ripken=True;}" ) # cpymad needs .input for macros logger.debug("Matching Parameters") madx.command.match(sequence=sequence, use_macro=True, chrom=True) for knob in vary_knobs: madx.command.vary(name=knob) madx.command.use_macro(name="do_ripken") madx.input(f"constraint, expr=table(twiss, {location}, beta12)=0" ) # need input else includes " and fails madx.input(f"constraint, expr=table(twiss, {location}, beta21)=0" ) # need input else includes " and fails madx.command.lmdif(calls=500, tolerance=1e-21) madx.command.endmatch()
def correct_lhc_orbit(madx: Madx, orbit_tolerance: float = 1e-14, iterations: int = 3, mode: str = "svd", **kwargs) -> None: """ Routine for orbit correction using 'MCB.*' elements in the LHC. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. orbit_tolerance (float): the tolerance for the correction. Defaults to 1e-14. iterations (int): the number of iterations of the correction to perform. Defaults to 3. mode (str): the method to use for the correction. Defaults to 'svd'. Keyword Args: Any keyword argument that can be given to the MAD-X CORRECT command, such as `mode`, `ncorr`, etc. """ logger.info("Starting orbit correction") for default_kicker in ("kicker", "hkicker", "vkicker", "virtualcorrector"): logger.trace(f"Disabling default corrector class '{default_kicker}'") madx.command.usekick(status="off", class_=default_kicker) logger.debug("Selecting '^MCB.*' correctors") madx.command.usekick(status="on", pattern="^MCB.*") madx.command.usemonitor(status="on", class_="monitor") for _ in range(iterations): logger.trace("Doing orbit correction for Y then X plane") madx.twiss(chrom=True) madx.command.correct(plane="y", flag="ring", error=orbit_tolerance, mode=mode, **kwargs) madx.command.correct(plane="x", flag="ring", error=orbit_tolerance, mode=mode, **kwargs)
def __init__(self, dataFile, madxModelFile, profilePath): self.madxModelFile = madxModelFile self.dataFile = dataFile self.profilePath = profilePath self.data = self.readData(dataFile) self.monitor = self.getMonitor() self.kickers, self.kicks = self.getKicks() self.sequence = self.getSequence() self.madx = Madx(stdout=False) self.madx.call(file=self.madxModelFile, chdir=True) # This are the initial conditions for the Twiss Module of MAD-X # there doesn't seem to be a strong dependence on them self.dx = 1.0e-4 self.dpx = 1.0e-6 self.dy = 1.0e-4 self.dpy = 1.0e-6 # Initial twiss parameters self.alfax = -3.036572962956109 self.alfay = 0.24718095605022355 self.betax = 20.549681146151855 self.betay = 2.671346837756801
def get_rdts(madx: Madx, order: int = 4, file: Union[Path, str] = None) -> tfs.TfsDataFrame: """ INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD). Calculate the RDTs via PTC_TWISS. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. order (int): maximum derivative order (only 0, 1 or 2 implemented in PTC). Defaults to `2`. file (Union[Path, str]): path to output file. Default `None` Returns: A TfsDataframe with results. """ logger.info(f"Entering PTC to calculate RDTs up to order {order}") madx.ptc_create_universe() logger.trace("Creating PTC layout") madx.ptc_create_layout(model=3, method=4, nst=3, exact=True) # madx.ptc_create_layout(model=3, method=6, nst=1) # from Michi logger.trace("Incorporating MAD-X alignment errors") madx.ptc_align() # use madx alignment errors # madx.ptc_setswitch(fringe=True) # include fringe effects logger.debug("Executing PTC Twiss") madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True) madx.ptc_end() logger.debug("Extracting results to TfsDataFrame") dframe = tfs.TfsDataFrame(madx.table.twissrdt.dframe()) dframe.columns = dframe.columns.str.upper() dframe.NAME = dframe.NAME.str.upper() if file: logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'") tfs.write(file, dframe) return dframe
def make_sixtrack_output(madx: Madx, energy: int) -> None: """ INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD). Prepare output for sixtrack run. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. energy (float): beam energy in GeV. """ logger.info("Preparing outputs for SixTrack") logger.debug("Powering RF cavities") madx.globals[ "VRF400"] = 8 if energy < 5000 else 16 # is 6 at injection for protons iirc? madx.globals[ "LAGRF400.B1"] = 0.5 # cavity phase difference in units of 2pi madx.globals["LAGRF400.B2"] = 0.0 logger.debug("Executing TWISS and SIXTRACK commands") madx.twiss() # used by sixtrack madx.sixtrack( cavall=True, radius=0.017) # this value is only ok for HL(LHC) magnet radius
def __init__(self, nturns=20000, npart=1000): self.nturns = nturns self.npart = npart mass = self.A * nmass * 1e9 * e / c**2 # in kg charge = self.Q * e # in Coul Ekin = self.Ekin_per_nucleon * self.A p0c = np.sqrt(Ekin**2 + 2 * Ekin * mass / e * c**2) # in eV Etot = np.sqrt(p0c**2 + (mass / e)**2 * c**4) * 1e-9 # in GeV p0 = p0c / c * e # in SI units gamma = np.sqrt(1 + (p0 / (mass * c))**2) beta = np.sqrt(1 - gamma**-2) self.beta = beta self.gamma = gamma self.p0 = p0 self.Etot = Etot self.p0c = p0c self.charge = charge self.mass = mass epsx_gauss = self.epsx_rms_fin * 1.43 epsy_gauss = self.epsy_rms_fin * 1.41 self.epsn_x = epsx_gauss * beta * gamma self.epsn_y = epsy_gauss * beta * gamma self.sig_z = self.sig_z * 1.22 self.sig_dp = self.sig_dp * 1.22 self.beta_z = self.sig_z / self.sig_dp self.madx = Madx() self.madx.options.echo = False self.madx.options.warn = False self.madx.options.info = False
def test_streamreader(self): output = [] m = Madx(stdout=output.append) self.assertEqual(len(output), 1) self.assertIn(b'++++++++++++++++++++++++++++++++++++++++++++', output[0]) self.assertIn(b'+ Support: [email protected],', output[0]) self.assertIn(b'+ Release date: ', output[0]) self.assertIn(b'+ Execution date: ', output[0]) # self.assertIn(b'+ Support: [email protected], ', output[1]) m.input('foo = 3;') self.assertEqual(len(output), 1) m.input('foo = 3;') self.assertEqual(len(output), 2) self.assertEqual(output[1], b'++++++ info: foo redefined\n') m.quit() self.assertEqual(len(output), 3) self.assertIn(b'+ MAD-X finished normally ', output[2])
from cpymad.madx import Madx import pysixtracklib as pyst import time from scipy.constants import e, m_p, c import numpy as np p0c = 6 * 1e9 # in eV Etot = np.sqrt(p0c**2 + (m_p/e)**2 * c**4) * 1e-9 # in GeV mad = Madx() mad.options.echo = False mad.call(file="fodo.madx") mad.command.beam(particle='proton', energy=str(Etot)) mad.use(sequence="FODO") mad.twiss() mad.command.select(flag="makethin", class_="quadrupole", slice='8') mad.command.select(flag="makethin", class_="sbend", slice='8') mad.command.makethin(makedipedge=False, style="teapot", sequence="fodo") mad.twiss() sis18 = mad.sequence.FODO nturns = 1 elements = pyst.Elements.from_mad(sis18) def prepare(npart=int(1e6), p0c=p0c, elements=elements, device='cpu'):
class TestMadx(unittest.TestCase, _TestCaseCompat): """ Test methods for the Madx class. The tests are directly based on the specifics of the sequence in test/testseq.madx Please compare this file for reference. """ def setUp(self): self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> ')) here = os.path.dirname(__file__) there = os.path.join(here, 'testseq.madx') with open(there) as f: self.doc = f.read() for line in self.doc.splitlines(): line = line.split('!')[0].strip() if line: self.mad._libmadx.input(line) def tearDown(self): self.mad.quit() del self.mad def test_copyright(self): import cpymad notice = cpymad.get_copyright_notice() self.assertIsInstance(notice, type(u"")) def test_version(self): """Check that the Madx.version attribute can be used as expected.""" version = self.mad.version # check format: major, minor, micro = map(int, version.release.split('.')) # We need at least MAD-X 5.04.02: self.assertGreaterEqual((major, minor, micro), (5, 4, 2)) # check format: year, month, day = map(int, version.date.split('.')) self.assertGreaterEqual((year, month, day), (2018, 10, 3)) self.assertLessEqual(month, 12) self.assertLessEqual(day, 31) self.assertTrue(str(version).startswith( 'MAD-X {}'.format(version.release))) def test_metadata(self): version = self.mad.version self.assertEqual(metadata.__version__, version.release) self.assertIsInstance(metadata.get_copyright_notice(), type(u"")) def test_independent_instances(self): # create a second Madx instance (1st one is created in setUp) madxness = Madx() # Check independence by defining a variable differently in each # instance: self.mad.input('ANSWER=42;') madxness.input('ANSWER=43;') self.assertEqual(self.mad.eval('ANSWER'), 42) self.assertEqual(madxness.eval('ANSWER'), 43) madxness.quit() def test_streamreader(self): output = [] m = Madx(stdout=output.append) self.assertEqual(len(output), 1) self.assertIn(b'++++++++++++++++++++++++++++++++++++++++++++', output[0]) self.assertIn(b'+ Support: [email protected],', output[0]) self.assertIn(b'+ Release date: ', output[0]) self.assertIn(b'+ Execution date: ', output[0]) # self.assertIn(b'+ Support: [email protected], ', output[1]) m.input('foo = 3;') self.assertEqual(len(output), 1) m.input('foo = 3;') self.assertEqual(len(output), 2) self.assertEqual(output[1], b'++++++ info: foo redefined\n') m.quit() self.assertEqual(len(output), 3) self.assertIn(b'+ MAD-X finished normally ', output[2]) def test_quit(self): self.mad.quit() self.assertIsNot(self.mad._process.returncode, None) self.assertFalse(bool(self.mad)) with self.assertRaises(RuntimeError): self.mad.input(';') def test_context_manager(self): output = [] with Madx(stdout=output.append) as m: m.input('foo = 3;') self.assertEqual(m.globals.foo, 3) self.assertIn(b'+ MAD-X finished normally ', output[-1]) self.assertFalse(bool(m)) with self.assertRaises(RuntimeError): m.input(';') def test_command_log(self): """Check that the command log contains all input commands.""" # create a new Madx instance that uses the history feature: history_filename = '_test_madx.madx.tmp' mad = Madx(command_log=history_filename) # feed some input and compare with history file: for line in self.doc.splitlines(): mad.input(line) with open(history_filename) as history_file: history = history_file.read() self.assertEqual(history.strip(), self.doc.strip()) # remove history file mad.quit() del mad os.remove(history_filename) def test_call_and_chdir(self): folder = os.path.abspath(os.path.dirname(__file__)) parent = os.path.dirname(folder) getcwd = self.mad._libmadx.getcwd g = self.mad.globals self.mad.chdir(folder) self.assertEqual(getcwd(), folder) self.mad.call('answer_42.madx') self.assertEqual(g.answer, 42) with self.mad.chdir('..'): self.assertEqual(getcwd(), parent) self.mad.call('test/answer_43.madx') self.assertEqual(g.answer, 43) self.mad.call('test/answer_call42.madx', True) self.assertEqual(g.answer, 42) self.assertEqual(getcwd(), folder) self.mad.call('answer_43.madx') self.assertEqual(g.answer, 43) self.mad.chdir('..') self.assertEqual(getcwd(), parent) def _check_twiss(self, seq_name): beam = 'ex=1, ey=2, particle=electron, sequence={0};'.format(seq_name) self.mad.command.beam(beam) self.mad.use(seq_name) initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5) twiss = self.mad.twiss(sequence=seq_name, **initial) # Check initial values: self.assertAlmostEqual(twiss['alfx'][0], initial['alfx']) self.assertAlmostEqual(twiss['alfy'][0], initial['alfy']) self.assertAlmostEqual(twiss['betx'][0], initial['betx']) self.assertAlmostEqual(twiss['bety'][0], initial['bety']) self.assertAlmostEqual(twiss.summary['ex'], 1) self.assertAlmostEqual(twiss.summary['ey'], 2) # Check that keys are all lowercase: for k in twiss: self.assertEqual(k, k.lower()) for k in twiss.summary: self.assertEqual(k, k.lower()) def test_error(self): self.mad.input(""" seq: sequence, l=1; endsequence; beam; use, sequence=seq; """) # Errors in MAD-X must not crash, but return False instead: self.assertFalse(self.mad.input('twiss;')) self.assertTrue(self.mad.input('twiss, betx=1, bety=1;')) def test_twiss_1(self): self._check_twiss('s1') # s1 can be computed at start self._check_twiss('s1') # s1 can be computed multiple times self._check_twiss('s2') # s2 can be computed after s1 def test_twiss_2(self): self._check_twiss('s2') # s2 can be computed at start self._check_twiss('s1') # s1 can be computed after s2 def test_twiss_with_range(self): beam = 'ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command.beam(beam) self.mad.use('s1') params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1') # Compute TWISS on full sequence, then on a sub-range, then again on # the full sequence. This checks that none of the range selections # have side-effects on each other: betx_full1 = self.mad.twiss(**params)['betx'] betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx'] betx_full2 = self.mad.twiss(**params)['betx'] # Check that the results have the expected lengths: self.assertEqual(len(betx_full1), 9) self.assertEqual(len(betx_range), 4) self.assertEqual(len(betx_full2), 9) # Check numeric results. Since the first 3 elements of range and full # sequence are identical, equal results are expected. And non-equal # results afterwards. self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1 self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1 self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2 self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2 def test_range_row_api(self): beam = 'ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command.beam(beam) self.mad.use('s1') params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1') tab = self.mad.twiss(range=('dr[2]', 'sb'), **params) self.assertEqual(tab.range, ('dr[2]', 'sb')) self.assertIn('betx', tab) def test_survey(self): self.mad.beam() self.mad.use('s1') tab = self.mad.survey() self.assertEqual(tab._name, 'survey') self.assertIn('x', tab) self.assertIn('y', tab) self.assertIn('z', tab) self.assertIn('theta', tab) self.assertIn('phi', tab) self.assertIn('psi', tab) self.assertLess(tab.x[-1], -1) assert_allclose(tab.y, 0) self.assertGreater(tab.z[-1], 7) def test_match(self): beam = 'ex=1, ey=2, particle=electron, sequence=s2;' self.mad.command.beam(beam) self.mad.use('s2') params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s2') self.mad.match(constraints=[dict(range='s1$end', betx=2.0)], weight={'betx': 2}, vary=['qp2->k1'], **params) twiss = self.mad.twiss(**params) val = twiss.betx[-1] self.assertAlmostEqual(val, 2.0, places=2) def test_verbose(self): self.mad.verbose(False) self.assertEqual(self.mad.options.echo, False) self.assertEqual(self.mad.options.info, False) self.mad.verbose(True) self.assertEqual(self.mad.options.echo, True) self.assertEqual(self.mad.options.info, True) def test_active_sequence(self): self.mad.command.beam('ex=1, ey=2, particle=electron, sequence=s1;') self.mad.use('s1') self.assertEqual(self.mad.sequence(), 's1') self.mad.beam() self.mad.use('s2') self.assertEqual(self.mad.sequence().name, 's2') def test_get_sequence(self): with self.assertRaises(KeyError): self.mad.sequence['sN'] s1 = self.mad.sequence['s1'] self.assertEqual(s1.name, 's1') seqs = self.mad.sequence self.assertItemsEqual(seqs, ['s1', 's2']) def test_eval(self): self.assertEqual(self.mad.eval(True), True) self.assertEqual(self.mad.eval(13), 13) self.assertEqual(self.mad.eval(1.3), 1.3) self.assertEqual(self.mad.eval([2, True, 'QP_K1']), [2, True, 2.0]) self.assertAlmostEqual(self.mad.eval("1/QP_K1"), 0.5) def test_globals(self): g = self.mad.globals # Membership: self.assertNotIn('FOO', g) # Setting values: g['FOO'] = 2 self.assertIn('FOO', g) self.assertEqual(g['FOO'], 2) self.assertEqual(self.mad.eval('FOO'), 2) # Re-setting values: g['FOO'] = 3 self.assertEqual(self.mad.eval('FOO'), 3) # Setting expressions: g['BAR'] = '3*foo' self.assertEqual(self.mad.eval('BAR'), 9) g['FOO'] = 4 self.assertEqual(self.mad.eval('BAR'), 12) self.assertEqual(g.defs.bar, "3*foo") self.assertEqual(g.cmdpar.bar.definition, "3*foo") # attribute access: g.bar = 42 self.assertEqual(g.defs.bar, 42) self.assertEqual(g.cmdpar.bar.definition, 42) self.assertEqual(g.BAR, 42) # repr self.assertIn("'bar': 42.0", str(g)) with self.assertRaises(NotImplementedError): del g['bar'] with self.assertRaises(NotImplementedError): del g.bar self.assertEqual(g.bar, 42) # still there self.assertIn('bar', list(g)) self.assertIn('foo', list(g)) # self.assertEqual(list(g), list(g.defs)) # self.assertEqual(list(g), list(g.cmdpar)) self.assertEqual(len(g), len(list(g))) self.assertEqual(len(g.defs), len(list(g.defs))) self.assertEqual(len(g.cmdpar), len(list(g.cmdpar))) def test_elements(self): self.assertIn('sb', self.mad.elements) self.assertIn('sb', list(self.mad.elements)) self.assertNotIn('foobar', self.mad.elements) self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14/4) idx = self.mad.elements.index('qp1') elem = self.mad.elements[idx] self.assertEqual(elem['k1'], 3) def test_sequence_map(self): seq = self.mad.sequence self.assertEqual(len(seq), 2) self.assertEqual(set(seq), {'s1', 's2'}) self.assertIn('s1', seq) self.assertNotIn('s3', seq) self.assertTrue(hasattr(seq, 's1')) self.assertFalse(hasattr(seq, 's3')) self.assertEqual(seq.s1.name, 's1') self.assertEqual(seq.s2.name, 's2') with self.assertRaises(AttributeError): seq.s3 def test_table_map(self): self.mad.beam() self.mad.use('s2') self.mad.survey(sequence='s2') tab = self.mad.table self.assertIn('survey', list(tab)) self.assertIn('survey', tab) self.assertNotIn('foobar', tab) self.assertEqual(len(tab), len(list(tab))) with self.assertRaises(AttributeError): tab.foobar def test_sequence(self): s1 = self.mad.sequence.s1 self.assertEqual(str(s1), '<Sequence: s1>') self.assertEqual(s1, self.mad.sequence.s1) self.assertEqual(s1, 's1') self.assertNotEqual(s1, self.mad.sequence.s2) self.assertNotEqual(s1, 's2') with self.assertRaises(RuntimeError): s1.beam with self.assertRaises(RuntimeError): s1.twiss_table with self.assertRaises(RuntimeError): s1.twiss_table_name self.assertFalse(s1.has_beam) self.assertFalse(s1.is_expanded) s1.expand() self.assertTrue(s1.has_beam) self.assertTrue(s1.is_expanded) s1.expand() # idempotent self.assertTrue(s1.has_beam) self.assertTrue(s1.is_expanded) initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5) self.mad.twiss(sequence='s1', sectormap=True, table='my_twiss', **initial) # Now works: self.assertEqual(s1.beam.particle, 'positron') self.assertEqual(s1.twiss_table_name, 'my_twiss') self.assertEqual(s1.twiss_table.betx[0], 2.5) self.assertEqual(s1.element_names(), [ 's1$start', 'dr', 'qp', 'dr[2]', 'qp[2]', 'dr[3]', 'sb', 'dr[4]', 's1$end', ]) self.assertEqual(s1.expanded_element_names(), s1.element_names()) self.assertEqual(len(s1.element_names()), len(s1.element_positions())) self.assertEqual(s1.element_positions(), [ 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0]) self.assertEqual(s1.expanded_element_positions(), [ 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0]) self.assertEqual(s1.elements[0].name, 's1$start') self.assertEqual(s1.elements[-1].name, 's1$end') self.assertEqual(s1.elements[-1].index, len(s1.elements)-1) self.assertEqual(s1.elements[3].index, 3) self.assertEqual(s1.elements.index('#s'), 0) self.assertEqual(s1.elements.index('#e'), len(s1.elements)-1) self.assertEqual(s1.elements.index('sb'), 6) def _get_elems(self, seq_name): elems = self.mad.sequence[seq_name].elements elem_idx = dict((el.node_name, i) for i, el in enumerate(elems)) return elems, elem_idx def test_sequence_get_elements_s1(self): s1, idx = self._get_elems('s1') qp1 = s1['qp[1]'] qp2 = s1['qp[2]'] sb1 = s1['sb[1]'] self.assertLess(idx['qp'], idx['qp[2]']) self.assertLess(idx['qp[2]'], idx['sb']) self.assertAlmostEqual(qp1['at'], 1.5) self.assertAlmostEqual(qp2['at'], 3.5) self.assertAlmostEqual(sb1['at'], 6) self.assertAlmostEqual(qp1.position, 1) self.assertAlmostEqual(qp2.position, 3) self.assertAlmostEqual(sb1.position, 5) self.assertAlmostEqual(qp1['l'], 1) self.assertAlmostEqual(qp2['l'], 1) self.assertAlmostEqual(sb1['l'], 2) self.assertAlmostEqual(float(qp1['k1']), 2) self.assertAlmostEqual(float(qp2['k1']), 2) self.assertAlmostEqual(float(sb1['angle']), 3.14/4) self.assertEqual(qp1.cmdpar.k1.expr.lower(), "qp_k1") def test_sequence_get_elements_s2(self): s2, idx = self._get_elems('s2') qp1 = s2['qp1[1]'] qp2 = s2['qp2[1]'] self.assertLess(idx['qp1'], idx['qp2']) self.assertAlmostEqual(qp1['at'], 0) self.assertAlmostEqual(qp2['at'], 1) self.assertAlmostEqual(qp1['l'], 1) self.assertAlmostEqual(qp2['l'], 2) self.assertAlmostEqual(float(qp1['k1']), 3) self.assertAlmostEqual(float(qp2['k1']), 2) # def test_sequence_get_expanded_elements(self): def test_crash(self): """Check that a RuntimeError is raised in case MAD-X crashes.""" self.assertTrue(self.mad) # a.t.m. MAD-X crashes on this input, because the L (length) # parametere is missing: self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;') self.assertFalse(self.mad) def test_sequence_elements(self): elements = self.mad.sequence['s1'].elements iqp2 = elements.index('qp[2]') qp1 = elements['qp[1]'] qp2 = elements[iqp2] self.assertAlmostEqual(qp1['at'], 1.5) self.assertAlmostEqual(qp2['at'], 3.5) self.assertAlmostEqual(qp1.position, 1) self.assertAlmostEqual(qp2.position, 3) self.assertEqual(iqp2, elements.at(3.1)) def test_sequence_expanded_elements(self): beam = 'ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command.beam(beam) self.mad.use('s1') elements = self.mad.sequence['s1'].expanded_elements iqp2 = elements.index('qp[2]') qp1 = elements['qp[1]'] qp2 = elements[iqp2] self.assertAlmostEqual(qp1['at'], 1.5) self.assertAlmostEqual(qp2['at'], 3.5) self.assertAlmostEqual(qp1.position, 1) self.assertAlmostEqual(qp2.position, 3) self.assertEqual(iqp2, elements.at(3.1)) def test_element_inform(self): beam = 'ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command.beam(beam) self.mad.use('s1') elem = self.mad.sequence.s1.expanded_elements['qp'] self.assertSetEqual({'k1', 'l', 'at'}, { name for name in elem if elem.cmdpar[name].inform }) def test_table(self): beam = 'ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command.beam(beam) self.mad.use('s1') initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5) twiss = self.mad.twiss(sequence='s1', sectormap=True, **initial) sector = self.mad.table.sectortable self.assertTrue(str(twiss).startswith("<Table 'twiss': ")) self.assertTrue(str(sector).startswith("<Table 'sectortable': ")) self.assertIn('betx', twiss) self.assertIn('t111', sector) self.assertNotIn('t111', twiss) self.assertNotIn('betx', sector) self.assertEqual(len(twiss), len(list(twiss))) self.assertEqual(set(twiss), set(twiss[0])) self.assertEqual(twiss.s[5], twiss[5].s) self.assertEqual(twiss.s[-1], twiss[-1].s) copy = twiss.copy() assert_allclose(copy['betx'], twiss.betx) self.assertEqual(set(copy), set(twiss)) copy = twiss.copy(['betx']) self.assertEqual(set(copy), {'betx'}) ALL = slice(None) self.assertEqual(sector.tmat(0).shape, (6, 6, 6)) assert_allclose(sector.tmat(ALL)[0, 0, 0, :], sector.t111) assert_allclose(sector.tmat(ALL)[1, 5, 3, :], sector.t264) assert_allclose(sector.tmat(ALL)[3, 0, 3, :], sector.t414) assert_allclose(sector.tmat(ALL)[4, 4, 4, :], sector.t555) assert_allclose(sector.rmat(ALL)[0, 0, :], sector.r11) assert_allclose(sector.rmat(ALL)[1, 5, :], sector.r26) assert_allclose(sector.rmat(ALL)[3, 0, :], sector.r41) assert_allclose(sector.rmat(ALL)[4, 4, :], sector.r55) assert_allclose(sector.kvec(ALL)[0, :], sector.k1) assert_allclose(sector.kvec(ALL)[1, :], sector.k2) assert_allclose(sector.kvec(ALL)[3, :], sector.k4) assert_allclose(sector.kvec(ALL)[4, :], sector.k5) r = self.mad.sectortable()[:, :6, :6] k = self.mad.sectortable()[:, 6, :6] t = self.mad.sectortable2() num_elems = len(self.mad.sequence.s1.elements) self.assertEqual(t.shape, (num_elems, 6, 6, 6)) self.assertEqual(r.shape, (num_elems, 6, 6)) self.assertEqual(k.shape, (num_elems, 6)) assert_allclose(t[:, 0, 0, 0], sector.t111) assert_allclose(t[:, 1, 5, 3], sector.t264) assert_allclose(t[:, 3, 0, 3], sector.t414) assert_allclose(t[:, 4, 4, 4], sector.t555) assert_allclose(r[:, 0, 0], sector.r11) assert_allclose(r[:, 1, 5], sector.r26) assert_allclose(r[:, 3, 0], sector.r41) assert_allclose(r[:, 4, 4], sector.r55) assert_allclose(k[:, 0], sector.k1) assert_allclose(k[:, 1], sector.k2) assert_allclose(k[:, 3], sector.k4) assert_allclose(k[:, 4], sector.k5) def test_attr(self): self.assertTrue(hasattr(self.mad, 'constraint')) self.assertTrue(hasattr(self.mad, 'constraint_')) self.assertTrue(hasattr(self.mad, 'global_')) self.assertFalse(hasattr(self.mad, 'foobar')) self.assertFalse(hasattr(self.mad, '_constraint')) def test_expr(self): g = self.mad.globals vars = self.mad.expr_vars g.foo = 1 g.bar = 2 self.assertEqual(set(vars('foo')), {'foo'}) self.assertEqual(set(vars('(foo) * sin(2*pi*bar)')), {'foo', 'bar'}) def test_command(self): twiss = self.mad.command.twiss sbend = self.mad.elements.sb clone = sbend.clone('foobar', angle="pi/5", l=1) self.assertIn('betx=0', str(twiss)) self.assertIn('angle=', str(sbend)) self.assertIn('tilt', sbend) self.assertEqual(sbend.tilt, 0) self.assertEqual(len(sbend), len(list(sbend))) self.assertIn('tilt', list(sbend)) self.assertEqual(clone.name, 'foobar') self.assertEqual(clone.base_type.name, 'sbend') self.assertEqual(clone.parent.name, 'sb') self.assertEqual(clone.defs.angle, 'pi / 5') self.assertAlmostEqual(clone.angle, 0.6283185307179586) self.assertEqual(len(clone), len(sbend)) self.assertIn('angle=0.628', str(clone)) self.assertNotIn('tilt', str(clone)) clone.angle = 0.125 clone = self.mad.elements.foobar # need to update cache self.assertEqual(clone.angle, 0.125) self.assertEqual(len(twiss), len(list(twiss))) self.assertIn('betx', list(twiss)) self.assertNotEqual(clone.angle, clone.parent.angle) del clone.angle clone = self.mad.elements.foobar # need to update cache self.assertEqual(clone.angle, clone.parent.angle) with self.assertRaises(AttributeError): clone.missing_attribute with self.assertRaises(NotImplementedError): del twiss['betx'] with self.assertRaises(NotImplementedError): del clone.base_type.angle def test_array_attribute(self): self.mad.globals.nine = 9 clone = self.mad.elements.multipole.clone('foo', knl=[0, 'nine/3', 4]) knl = clone.knl self.assertEqual(knl[0], 0) self.assertEqual(knl[1], 3) self.assertEqual(knl[2], 4) self.assertEqual(len(knl), 3) self.assertEqual(list(knl), [0.0, 3.0, 4.0]) self.assertEqual(str(knl), '[0.0, 3.0, 4.0]') knl[1] = '3*nine' self.assertEqual(self.mad.elements.foo.defs.knl[1], '3 * nine') self.assertEqual(self.mad.elements.foo.knl[1], 27) def test_command_map(self): command = self.mad.command self.assertIn('match', command) self.assertIn('sbend', command) self.assertNotIn('foooo', command) self.assertIn('match', list(command)) self.assertEqual(len(command), len(list(command))) self.assertIn('match', str(command)) self.assertIn('sbend', str(command)) self.assertIn('sbend', self.mad.base_types) self.assertNotIn('match', self.mad.base_types) def test_comments(self): mad = self.mad var = mad.globals mad('x = 1; ! x = 2;') self.assertEqual(var.x, 1) mad('x = 2; // x = 3;') self.assertEqual(var.x, 2) mad('x = 3; /* x = 4; */') self.assertEqual(var.x, 3) mad('/* x = 3; */ x = 4;') self.assertEqual(var.x, 4) mad('x = 5; ! /* */ x = 6;') self.assertEqual(var.x, 5) mad('x = 5; /* ! */ x = 6;') self.assertEqual(var.x, 6) def test_multiline_input(self): mad = self.mad var = mad.globals mad(''' x = 1; y = 2; ''') self.assertEqual(var.x, 1) self.assertEqual(var.y, 2) mad(''' x = /* 3; y =*/ 4; ''') self.assertEqual(var.x, 4) self.assertEqual(var.y, 2) mad(''' x = 1; /* */ x = 2; */ if (x == 1) { x = 3; } ''') self.assertEqual(var.x, 2) mad(''' x = 1; /* x = 2; */ if (x == 1) { x = 3; } ''') self.assertEqual(var.x, 3)
class TestMadx(unittest.TestCase, _compat.TestCase): """ Test methods for the Madx class. The tests are directly based on the specifics of the sequence in test/data/lebt/init.madx Please compare this file for reference. """ def setUp(self): self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> ')) here = os.path.dirname(__file__) there = os.path.join(here, 'data', 'lebt', 'init.madx') self.doc = open(there).read() for line in self.doc.splitlines(): self.mad._libmadx.input(line) def tearDown(self): del self.mad def test_version(self): """Check that the Madx.version attribute can be used as expected.""" version = self.mad.version # check format: major, minor, mini = map(int, version.release.split('.')) # We need at least MAD-X 5.02.03: self.assertGreaterEqual(major, 5) self.assertGreaterEqual(minor, 2) self.assertGreaterEqual(mini, 3) # check format: year, month, day = map(int, version.date.split('.')) self.assertGreaterEqual(year, 2014) self.assertGreaterEqual(month, 1) self.assertGreaterEqual(day, 1) self.assertLessEqual(month, 12) self.assertLessEqual(day, 31) def test_independent_instances(self): # create a second Madx instance (1st one is created in setUp) madxness = Madx() # Check independence by defining a variable differently in each # instance: self.mad.input('ANSWER=42;') madxness.input('ANSWER=43;') self.assertEqual(self.mad.evaluate('ANSWER'), 42); self.assertEqual(madxness.evaluate('ANSWER'), 43); def test_command_log(self): """Check that the command log contains all input commands.""" # create a new Madx instance that uses the history feature: history_filename = '_test_madx.madx.tmp' mad = Madx(command_log=history_filename) # feed some input and compare with history file: for line in self.doc.splitlines(): mad.input(line) with open(history_filename) as history_file: history = history_file.read() self.assertEqual(history.strip(), self.doc.strip()) # remove history file del mad os.remove(history_filename) # TODO: # def test_command(self): # def test_help(self): # def test_call(self): def _check_twiss(self, seq_name): beam = 'beam, ex=1, ey=2, particle=electron, sequence={0};'.format(seq_name) self.mad.command(beam) initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5) # by explicitly specifying the 'columns' parameter a persistent copy # is returned. We check that this copy contains the data we want and # that it has a 'summary' attribute: twiss = self.mad.twiss(sequence=seq_name, columns=['betx', 'bety', 'alfx', 'alfy'], **initial) betx, bety = twiss['betx'], twiss['bety'] alfx, alfy = twiss['alfx'], twiss['alfy'] # Check initial values: self.assertAlmostEqual(twiss['alfx'][0], initial['alfx']) self.assertAlmostEqual(twiss['alfy'][0], initial['alfy']) self.assertAlmostEqual(twiss['betx'][0], initial['betx']) self.assertAlmostEqual(twiss['bety'][0], initial['bety']) self.assertAlmostEqual(twiss.summary['ex'], 1) self.assertAlmostEqual(twiss.summary['ey'], 2) # Check that keys are all lowercase: for k in twiss: self.assertEqual(k, k.lower()) for k in twiss.summary: self.assertEqual(k, k.lower()) def test_twiss_1(self): self._check_twiss('s1') # s1 can be computed at start self._check_twiss('s1') # s1 can be computed multiple times self._check_twiss('s2') # s2 can be computed after s1 def test_twiss_2(self): self._check_twiss('s2') # s2 can be computed at start self._check_twiss('s1') # s1 can be computed after s2 def test_twiss_with_range(self): beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command(beam) params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, columns=['betx', 'bety'], sequence='s1') # Compute TWISS on full sequence, then on a sub-range, then again on # the full sequence. This checks that none of the range selections # have side-effects on each other: betx_full1 = self.mad.twiss(**params)['betx'] betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx'] betx_full2 = self.mad.twiss(**params)['betx'] # Check that the results have the expected lengths: self.assertEqual(len(betx_full1), 9) self.assertEqual(len(betx_range), 4) self.assertEqual(len(betx_full2), 9) # Check numeric results. Since the first 3 elements of range and full # sequence are identical, equal results are expected. And non-equal # results afterwards. self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1 self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1 self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2 self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2 def test_range_row_api(self): beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command(beam) params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1') tab = self.mad.twiss(range=('dr[2]', 'sb'), **params) self.assertEqual(tab.range, ('dr[2]', 'sb')) self.assertIn('betx', tab) # def test_survey(self): # def test_aperture(self): # def test_use(self): # def test_match(self): # def test_verbose(self): def test_active_sequence(self): self.mad.command('beam, ex=1, ey=2, particle=electron, sequence=s1;') self.mad.active_sequence = 's1' self.assertEqual(self.mad.active_sequence.name, 's1') def test_get_sequence(self): with self.assertRaises(KeyError): self.mad.sequences['sN'] s1 = self.mad.sequences['s1'] self.assertEqual(s1.name, 's1') def test_get_sequences(self): seqs = self.mad.sequences self.assertItemsEqual(seqs, ['s1', 's2']) def test_evaluate(self): val = self.mad.evaluate("1/QP_K1") self.assertAlmostEqual(val, 0.5) def test_set_value(self): self.mad.set_value('FOO', 1) self.mad.set_value('BAR', 'FOO') self.mad.set_value('FOO', 2) self.assertEqual(self.mad.evaluate('FOO'), 2) self.assertEqual(self.mad.evaluate('BAR'), 1) def test_set_expression(self): self.mad.set_expression('FOO', 'BAR') self.mad.set_value('BAR', 42) self.mad.evaluate('FOO') self.assertEqual(self.mad.evaluate('FOO'), 42) self.mad.set_value('BAR', 43) self.assertEqual(self.mad.evaluate('FOO'), 43) def test_globals(self): g = self.mad.globals # Membership: self.assertNotIn('FOO', g) # Setting values: g['FOO'] = 2 self.assertIn('FOO', g) self.assertEqual(g['FOO'], 2) self.assertEqual(self.mad.evaluate('FOO'), 2) # Re-setting values: g['FOO'] = 3 self.assertEqual(self.mad.evaluate('FOO'), 3) # Setting expressions: g['BAR'] = '3*foo' self.assertEqual(self.mad.evaluate('BAR'), 9) g['FOO'] = 4 self.assertEqual(self.mad.evaluate('BAR'), 12) def test_elements(self): self.assertIn('sb', self.mad.elements) self.assertIn('sb', list(self.mad.elements)) self.assertNotIn('foobar', self.mad.elements) self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14/4) idx = self.mad.elements.index('qp1') elem = self.mad.elements[idx] self.assertEqual(elem['k1'], 3) # def test_sequence_beam(self): # def test_sequence_twiss(self): # def test_sequence_twissname(self): def _get_elems(self, seq_name): elems = self.mad.sequences[seq_name].elements elem_idx = dict((el['name'], i) for i, el in enumerate(elems)) return elems, elem_idx def test_sequence_get_elements_s1(self): s1, idx = self._get_elems('s1') qp1 = s1['qp[1]'] qp2 = s1['qp[2]'] sb1 = s1['sb[1]'] self.assertLess(idx['qp'], idx['qp[2]']) self.assertLess(idx['qp[2]'], idx['sb']) self.assertAlmostEqual(qp1['at'], 1) self.assertAlmostEqual(qp2['at'], 3) self.assertAlmostEqual(sb1['at'], 5) self.assertAlmostEqual(qp1['l'], 1) self.assertAlmostEqual(qp2['l'], 1) self.assertAlmostEqual(sb1['l'], 2) self.assertAlmostEqual(float(qp1['k1']), 2) self.assertAlmostEqual(float(qp2['k1']), 2) self.assertAlmostEqual(float(sb1['angle']), 3.14/4) self.assertEqual(str(qp1['k1']).lower(), "qp_k1") def test_sequence_get_elements_s2(self): s2, idx = self._get_elems('s2') qp1 = s2['qp1[1]'] qp2 = s2['qp2[1]'] self.assertLess(idx['qp1'], idx['qp2']) self.assertAlmostEqual(qp1['at'], 0) self.assertAlmostEqual(qp2['at'], 1) self.assertAlmostEqual(qp1['l'], 1) self.assertAlmostEqual(qp2['l'], 2) self.assertAlmostEqual(float(qp1['k1']), 3) self.assertAlmostEqual(float(qp2['k1']), 2) # def test_sequence_get_expanded_elements(self): def test_crash(self): """Check that a RuntimeError is raised in case MAD-X crashes.""" # a.t.m. MAD-X crashes on this input, because the L (length) # parametere is missing: self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;') def test_sequence_elements(self): elements = self.mad.sequences['s1'].elements iqp2 = elements.index('qp[2]') qp1 = elements['qp[1]'] qp2 = elements[iqp2] self.assertAlmostEqual(qp1['at'], 1) self.assertAlmostEqual(qp2['at'], 3) self.assertEqual(iqp2, elements.at(3.1)) def test_sequence_expanded_elements(self): beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;' self.mad.command(beam) self.mad.use('s1') elements = self.mad.sequences['s1'].expanded_elements iqp2 = elements.index('qp[2]') qp1 = elements['qp[1]'] qp2 = elements[iqp2] self.assertAlmostEqual(qp1['at'], 1) self.assertAlmostEqual(qp2['at'], 3) self.assertEqual(iqp2, elements.at(3.1))