LHC Example

The file ex-lhc-couplingLocal/lhc-couplingLocal.py contains an example of loading the required files to use and run the LHC, while including a method to receive and plot intermediate results of a match.

Loading the LHC

The following lines loads the required variables and files for the example. assertf is not required, but it is used to check if the loading of the LHC is successful. The two string inputs to mad.MADX.load is also not required, but it is used to specify the final destination of the translated files from MAD-X to MAD-NG. Also, line 5 is not required as this is just to prevent MAD-NG from reporting warnings of unset variables from loading the LHC. Once again, in this extract, the important point to note is that to input strings into functions in MAD-NG, they must be enclosed in quotes. This is because any string input from the side of python is evaluated by MAD-NG, therefore the input can be a variable or expressions.

To grab variables from the MAD-X environment, we use mad.load("MADX", ...).

 1    mad.load("MAD.utility", "assertf")
 2
 3    mad.MADX.load("'lhc_as-built.seq'", "'lhc_as-built.mad'")
 4    mad.MADX.load("'opticsfile.21'", "'opticsfile.21.mad'")
 5    mad.MADX.load("'lhc_unset_vars.mad'") # Load a list of unset variables to prevent warnings
 6
 7    mad.load("MADX", "lhcb1", "nrj")
 8
 9    mad.assertf("#lhcb1 == 6694",
10        "'invalid number of elements %d in LHCB1 (6694 expected)'", "#lhcb1")
11    

Receving intermediate results

The most complicated part of the example includes the following set of lines.

From lines 4 - 8 below, we define a function that will be invoked during the optimization process at each iteration. Within this function, we perform a twiss for the match function to use, while also sending some information on the twiss to python, on line 6.

From lines 10 - 21, we run a match, with a reference to the match result returned to the variable match_rtrn. Line 22 is a very important line, as this is something you place in the pipe to MAD-NG for MAD-NG to execute once the match is done. Lines 23-25 receive the first result returned during the match, so that we can start plotting the results.

The plotting occurs between lines 27 - 36, wtih the while loop continuing until twiss result is None, which occurs when the match is done, as requested on line 22.

Finally, on lines 38 and 39, we retrieve the results of the match from the variable match_rtrn. Since match_rtrn is a temporary variable, there is a limit to how many of these that can be stored (see Managing References for more information on these), we delete the reference in python to clear the temporary variable so that is is available for future use.

Important

As MAD-NG is running in the background, the variable match_rtrn contains no information and instead must be queried for the results. During the query, python will then have to wait for MAD-NG to finish the match, and then return the results. On the other hand, if we do not query for the results, the match will continue to run in the background, we can do other things in python, and then query for the results later.

 1    mad.send("""
 2    expr1 = \\t, s -> t.q1 - 62.30980
 3    expr2 = \\t, s -> t.q2 - 60.32154
 4    function twiss_and_send()
 5        local mtbl, mflow = twiss {sequence=lhcb1, method=4}
 6        py:send({mtbl.s, mtbl.beta11})
 7        return mtbl, mflow
 8    end
 9    """)
10    match_rtrn = mad.match(
11        command=mad.twiss_and_send,
12        variables = [
13            {"var":"'MADX.dqx_b1'", "name":"'dQx.b1'", "'rtol'":1e-6},
14            {"var":"'MADX.dqy_b1'", "name":"'dQy.b1'", "'rtol'":1e-6},
15        ],
16        equalities = [
17            {"expr": mad.expr1, "name": "'q1'", "tol":1e-3},
18            {"expr": mad.expr2, "name": "'q2'", "tol":1e-3},
19        ],
20        objective={"fmin": 1e-3}, maxcall=100, info=2,
21    )
22    mad.send("py:send(nil)")
23    tws_result = mad.recv ()
24    x = tws_result[0]
25    y = tws_result[1]
26
27    plt.ion()
28    fig = plt.figure()
29    ax = fig.add_subplot(111)
30    line1, = ax.plot(x, y, 'b-')
31    while tws_result:
32        line1.set_xdata(tws_result[0])
33        line1.set_ydata(tws_result[1])
34        fig.canvas.draw()
35        fig.canvas.flush_events()
36        tws_result = mad.recv()
37
38    mad["status", "fmin", "ncall"] = match_rtrn
39    del match_rtrn

LHC Speed Tests

This file creates functions within MAD-NG, LHC_load, which loads the LHC and reg_expr, which looks through the MADX environment and places anything that is a deferred expression into the table expr.

The LHC_load function sends a string "done" afterwards, so that Python can stay in sync and time the LHC correctly.

The code below just runs the function and times it.

t0 = time.time()
mad.LHC_load ()
mad.recv() #done
t1 = time.time()

The reg_expr fucntion is recursive so, after calling the function, to ensure Python stays in sync, Python is required to ask MAD-NG for the string "done".

t0 = time.time()
mad.reg_expr("MADX", mad.MADX)
mad.send("py:send('done')").recv() # reg_expr is recursive
t1 = time.time()

Next, we have the first of two methods to evaluate every deferred expression in the LHC and receive them, where Python performs a loop through the number of deferred expressions and has to ask MAD-NG everytime to receive the result.

t0 = time.time()
mad.send("py:send(#expr)")
for i in range(mad.recv()):
    mad.send(f"py:send(expr[{i+1}]())").recv()
t1 = time.time()

The second method involves making MAD-NG do a loop at the same time as python and therefore does not require back on forth communication, which speeds up the transfer of data

t0 = time.time()
mad.send("len = #expr py:send(len) for i=1,len do py:send(expr[i]()) end")
for i in range(mad.recv()):
    mad.recv()
t1 = time.time()

The two methods above do not store the data, so the next bit of code is identical to above, but uses list comprehension to store the data into a list automatically, storing the lists into variables exprList1 and exprList2. The main point of seperating the methods above and below was to identify if storing the variables into a list was a bottleneck.

t0 = time.time()
mad.send("py:send(#expr)")
exprList1 = [mad.send(f"py:send(expr[{i+1}]())").recv() for i in range(mad.recv())]
t1 = time.time()

t0 = time.time()
mad.send("len = #expr py:send(len) for i=1,len do py:send(expr[i]()) end")
exprList2 = [mad.recv() for i in range(mad.recv())]
t1 = time.time()

While we have the LHC loaded, the next example grabs the name of every element in the sequence lhcb1, demonstrating the ability and speed of pymadng and MAD-NG, other attributes could also be grabbed, but for simplicity this code just gets the names. This bit of code also uses list comprehension while making MAD-NG loop at the same time as Python.

t0 = time.time()
mad["lhcb1"] = mad.MADX.lhcb1

mad.send("""
py:send(#lhcb1)
for i, elm, spos, len in lhcb1:iter() do 
    py:send(elm.name)
end
""")
nameList = [mad.recv() for _ in range(mad.recv())]
t1 = time.time()

Another method, not shown above could be to create an entire list on the side of MAD-NG and then send the entire list to python. Which is done below, where instead all the names from the sequence lhcb2 are taken from MAD-NG.

t0 = time.time()
mad["lhcb2"] = mad.MADX.lhcb2
mad.send("""
lhcb2_tbl = {}
for i, elm, spos, len in lhcb2:iter() do 
    lhcb2_tbl[i] = elm.name
end
py:send(lhcb2_tbl)
""")
nameList = mad.recv()
t1 = time.time()

You can run the file yourself to retrieve the timings, but below is one run on an Intel® Core™ i7-8550U CPU @ 1.80GHz × 8 in Ubuntu 22.04.1 LTS

Load time: 0.3955872058868408  sec
reg_expr time: 0.034337759017944336  sec

For evaluating the deferred expressions

eval time method 1: 0.5888900756835938  sec
eval time method 2: 0.2224569320678711  sec
eval time method 3: 0.6652431488037109  sec
eval time method 4: 0.21885156631469727  sec
time to retrieve every element name in lhcb1 sequence 0.024236202239990234 sec
time to retrieve every element name in lhcb2 sequence 0.0245511531829834 sec