List of Examples
Sending and Receiving Multiple Types
1from pymadng import MAD
2import numpy as np
3import time
4
5arr0 = np.zeros((10000, 1000)) + 1j # 2*10000*1000*8 -> 160 MB
6
7mad = MAD()
8
9## Create a matrix in MAD and send it to Python
10mad.send("""
11 local m1 = MAD.matrix(1000, 1000):seq()
12 py:send(m1)
13 """)
14
15## Create a complex matrix in MAD
16mad.send("cm1 = (MAD.cmatrix(10000, 1000) + 1i)")
17
18# Create a string the manipulates the complex matrix in MAD and sends the result it to Python
19cmatrixString = """
20 {0} = cm1 {1} {2}
21 py:send({0})"""
22
23mad.send(cmatrixString.format("cm4", "*", 1)) ## Set cm4 to cm1 * 1 and send it to Python
24mad.send(cmatrixString.format("cm1", "*", 2)) ## Set cm1 to cm1 * 2 and send it to Python
25mad.send(cmatrixString.format("cm2", "*", 2)) ## Set cm2 to cm1 * 2 and send it to Python
26mad.send(cmatrixString.format("cm3", "/", 3)) ## Set cm3 to cm1 / 3 and send it to Python
27
28## Create a vector in MAD and send it to Python
29mad.send("""
30 local v1 = (MAD.vector(45):seq()*2 + 1)/3
31 py:send(v1)
32 """)
33start_time = time.time() # Start timer
34
35# Receive the matrices and vectors
36m1 = mad.recv()
37cm4 = mad.recv()
38cm1 = mad.recv()
39cm2 = mad.recv()
40cm3 = mad.recv()
41v1 = mad.recv()
42
43print(time.time() - start_time) # Print time
44
45# Check if the matrices have been correctly sent
46print(np.all(cm1 == arr0*2))
47print(np.all(cm2 == arr0*2*2))
48print(np.all(cm3 == arr0*2/3))
49print(np.all(cm4 == arr0))
50
51# Send a list to MAD and receive a changed version back
52myList = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2
53mad.send("""
54local list = py:recv()
55list[1][1] = 10
56list[2][1] = 10
57py:send(list)
58""")
59mad.send(myList)
60myList[0][0] = 10
61myList[1][0] = 10
62print("receiving lists", mad.recv() == myList)
63
64# Send an integer to MAD and receive a changed version back
65myInt = 4
66mad.send("""
67local myInt = py:recv()
68py:send(myInt+2)
69""")
70mad.send(myInt)
71print("Integers", mad.recv() == 6)
72
73# Send a float to MAD and receive a changed version back
74myFloat = 6.612
75mad.send("""
76local myFloat = py:recv()
77py:send(myFloat + 0.56)
78""")
79mad.send(myFloat)
80print("Floats", mad.recv() == 6.612 + 0.56)
81
82# Send a complex number to MAD and receive a changed version back
83myCpx = 6.612 + 4j
84mad.send("""
85local myCpx = py:recv()
86py:send(myCpx + 0.5i)
87""")
88mad.send(myCpx)
89print("Complex", mad.recv() == 6.612 + 4.5j)
90
91# Send None to MAD and it back
92mad.send("""
93local myNil = py:recv()
94py:send(myNil)
95""")
96mad.send(None)
97print("Nil/None", mad.recv() == None)
98
99# Receive ranges from MAD
100mad.send("""
101py:send(3..11..2)
102py:send(MAD.nrange(3.5, 21.4, 12))
103py:send(MAD.nlogrange(1, 20, 20))
104""")
105print("irng", mad.recv() == range(3 , 12 , 2)) #Py not inclusive, mad is
106print("rng", mad.recv() == np.linspace(3.5, 21.4, 12))
107print("lrng", np.allclose(mad.recv(), np.geomspace(1, 20, 20)))
Sending and Receiving Large Datasets
1from pymadng import MAD
2import numpy as np
3import time, sys
4
5arr0 = np.zeros((10000, 1000)) + 1j# 2*10000*1000*8 -> 160 MB
6
7start_time = time.time()
8mad = MAD(debug = False)
9# Number one rule! If you ask mad to send you something, read it!
10# Don't ask mad to send you data, then try to send more data, this will lead to a deadlock!
11
12mad.send("cm1 = py:recv()") # Tell MAD to receive something
13mad.send(arr0) # Send the data to MAD
14
15mad.send("cm2 = py:recv()")
16mad.send(arr0)
17
18mad.send("cm3 = py:recv()")
19mad.send(arr0)
20
21mad.send("cm4 = py:recv()")
22mad.send(arr0)
23
24mad.send("cm4 = py:recv()")
25mad.send(arr0)
26
27mad.send("cm4 = py:recv()")
28mad.send(arr0)
29
30print("time to write to MAD 960 MB", time.time() - start_time)
31start_time = time.time()
32
33# Create a string the manipulates the complex matrix in MAD and sends the result it to Python
34cmatrixString = """
35 {0} = {0} {1} {2}
36 py:send({0})
37 """
38
39start_time = time.time()
40mad.send(cmatrixString.format("cm1", "*", 2)) # Set cm1 to cm1 * 2 and send it to Python
41cm1 = mad.recv()
42mad.send(cmatrixString.format("cm2", "*", 2)) # Set cm2 to cm1 * 2 and send it to Python
43cm2 = mad.recv()
44mad.send(cmatrixString.format("cm3", "/", 3)) # Set cm3 to cm1 / 3 and send it to Python
45cm3 = mad.recv()
46mad.send(cmatrixString.format("cm4", "*", 1)) # Set cm4 to cm1 * 1 and send it to Python
47cm4 = mad.recv()
48mad.send(cmatrixString.format("cm4", "*", 2)) # Set cm4 to cm1 * 1 and send it to Python
49cm4 = mad.recv()
50mad.send(cmatrixString.format("cm4", "*", 2)) # Set cm4 to cm1 * 1 and send it to Python
51cm4 = mad.recv()
52print("time to read 960 MB from MAD", time.time() - start_time)
53
54# Check that the data is correct
55print(np.all(cm1 == arr0*2))
56print(np.all(cm2 == arr0*2))
57print(np.all(cm3 == arr0/3))
58print(np.all(cm4 == arr0*4))
59
60
61# Send a string to MAD that will send a string back to Python
62mad.send("""
63 py:send([=[
64mad.send("py:send([[print('pymad success')]])") # Needs to lack indent for python
65 ]=])""")
66mad.recv_and_exec()
67print("On second read:")
68mad.recv_and_exec()
69
70mad.send("""
71 py:send([==[
72mad.send('''py:send([=[
73mad.send("py:send([[print('Another pymad success') #This MAD string is executed in Python]])")
74 ]=])''')
75 ]==])
76 """)
77mad.recv_and_exec()
78mad.recv_and_exec()
79print("On third read:")
80mad.recv_and_exec()
FODO Cell Example
1from pymadng import MAD
2import matplotlib.pyplot as plt
3import os
4orginal_dir = os.getcwd()
5os.chdir(os.path.dirname(os.path.realpath(__file__)))
6
7with MAD() as mad:
8 mad.MADX.load(f"'fodo.seq'", f"'fodo.mad'")
9 mad["seq"] = mad.MADX.seq
10 mad.seq.beam = mad.beam()
11 mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, implicit=True, nslice=10, save="'atbody'")
12 plt.plot(mad.mtbl.s, mad.mtbl.beta11)
13 plt.show()
14
15with MAD() as mad:
16 mad.MADX.load(f"'fodo.seq'", f"'fodo.mad'")
17 mad.load("MADX", "seq")
18 mad.seq.beam = mad.beam()
19 mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, implicit=True, nslice=10, save="'atbody'")
20 cols = mad.py_strs_to_mad_strs(["name", "s", "beta11", "beta22", "mu1", "mu2", "alfa11", "alfa22"])
21 mad.mtbl.write("'twiss_py.tfs'", cols)
22 for x in mad.seq:
23 print(x.name, x.kind)
24 plt.plot(mad.mtbl.s, mad.mtbl.beta11)
25 plt.show()
26
27
28with MAD() as mad:
29 mad.load("element", "quadrupole")
30 mad["circum", "lcell"] = 60, 20
31
32 mad.load("math", "sin", "pi")
33 mad["v"] = mad.deferred(k="1/(lcell/sin(pi/4)/4)")
34
35 mad["qf"] = mad.quadrupole("knl:={0, v.k}", l=1)
36 mad["qd"] = mad.quadrupole("knl:={0, -v.k}", l=1)
37 mad["seq"] = mad.sequence("""
38 qf { at = 0 },
39 qd { at = 0.5 * lcell },
40 qf { at = 1.0 * lcell },
41 qd { at = 1.5 * lcell },
42 qf { at = 2.0 * lcell },
43 qd { at = 2.5 * lcell },
44 """, refer="'entry'", l=mad.circum,)
45 mad.seq.beam = mad.beam()
46 mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, implicit=True, nslice=10, save="'atbody'")
47 cols = mad.py_strs_to_mad_strs(["name", "s", "beta11", "beta22", "mu1", "mu2", "alfa11", "alfa22"])
48 mad.mtbl.write("'twiss_py.tfs'", cols)
49
50 plt.plot(mad.mtbl.s, mad.mtbl["beta11"])
51 plt.title("FODO Cell")
52 plt.xlabel("s")
53 plt.ylabel("beta11")
54 plt.show()
55
56with MAD() as mad:
57 mad.send(f"""
58 MADX:load("fodo.seq", "fodo.mad")
59 local seq in MADX
60 seq.beam = beam -- use default beam
61 mtbl, mflw = twiss {{sequence=seq, method=4, implicit=true, nslice=10, save="atbody"}}
62 py:send(mtbl)
63 """)
64 mtbl = mad.recv("mtbl")
65 plt.plot(mtbl.s, mtbl.beta11, "r-", label="Method 1")
66
67 mad.send("""py:send(mtbl.s) ; py:send(mtbl.beta11)""")
68 plt.plot(mad.recv(), mad.recv(), "g--", label="Method 2", )
69
70 plt.plot(mad.mtbl.s, mad.mtbl.beta11, "b:", label="Method 3")
71
72 plt.legend()
73 plt.show()
74
75os.chdir(orginal_dir)
PS Twiss Example
1import os, time, pandas
2from pymadng import MAD
3
4orginal_dir = os.getcwd()
5os.chdir(os.path.dirname(os.path.realpath(__file__)))
6
7with MAD(debug=False) as mad:
8 mad["psbeam"] = mad.beam(particle="'proton'", pc=2.794987)
9 mad.MADX.BEAM = mad.psbeam
10 mad.MADX.BRHO = mad.psbeam.brho
11 mad.MADX.load(f"'ps_unset_vars.mad'")
12 mad.MADX.load(f"'ps_mu.seq'")
13 mad.MADX.load(f"'ps_ss.seq'")
14 mad.MADX.load(f"'ps_fb_lhc.str'")
15
16 mad.load("MADX", "ps")
17 mad.ps.beam = mad.psbeam
18 mad["srv", "mflw"] = mad.survey(sequence=mad.ps)
19
20 mad.srv.write("'PS_survey_py.tfs'",
21 mad.py_strs_to_mad_strs(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]),
22 )
23
24 mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3, chrom=True)
25
26 mad.load("MAD.gphys", "melmcol")
27 #Add element properties as columns
28 mad.melmcol(mad.mtbl,
29 mad.py_strs_to_mad_strs(
30 ["angle", "tilt", "k0l", "k1l", "k2l", "k3l", "k4l", "k5l", "k6l", "k0sl",
31 "k1sl", "k2sl", "k3sl", "k4sl", "k5sl", "k6sl", "ksl", "hkick", "vkick" ]),
32 )
33
34 mad.mtbl.write("'PS_twiss_py.tfs'",
35 mad.py_strs_to_mad_strs(
36 ["name", "kind", "s", "x", "px", "beta11", "alfa11", "beta22", "alfa22","dx",
37 "dpx", "mu1", "mu2", "l", "angle", "k0l", "k1l", "k2l", "k3l", "hkick", "vkick"]
38 )
39 )
40
41 df = mad.mtbl.to_df()
42 print(df)
43 try:
44 import tfs
45 except ImportError:
46 print("tfs-pandas not installed, so the header is stored in attrs instead of headers")
47 print(df.attrs)
48
49 print(mad.srv.to_df())
50
51os.chdir(orginal_dir)
LHC Example
1from pymadng import MAD
2import time
3import os
4
5import matplotlib.pyplot as plt
6import numpy as np
7orginal_dir = os.getcwd()
8os.chdir(os.path.dirname(os.path.realpath(__file__)))
9
10
11with MAD() as mad:
12 mad.load("MAD.utility", "assertf")
13
14 mad.MADX.load("'lhc_as-built.seq'", "'lhc_as-built.mad'")
15 mad.MADX.load("'opticsfile.21'", "'opticsfile.21.mad'")
16 mad.MADX.load("'lhc_unset_vars.mad'") # Load a list of unset variables to prevent warnings
17
18 mad.load("MADX", "lhcb1", "nrj")
19
20 mad.assertf("#lhcb1 == 6694",
21 "'invalid number of elements %d in LHCB1 (6694 expected)'", "#lhcb1")
22
23 mad.lhcb1.beam = mad.beam(particle="'proton'", energy=mad.nrj)
24 mad.MADX_env_send("""
25 ktqx1_r2 = -ktqx1_l2 ! remove the link between these 2 vars
26 kqsx3_l2 = -0.0015
27 kqsx3_r2 = +0.0015
28 """)
29 t0 = time.time()
30 mad["tbl", "flw"] = mad.twiss(sequence=mad.lhcb1, method=4)
31 # plt.plot(mad.tbl.s, mad.tbl.beta11)
32 # plt.show()
33 mad.tbl.write("'before_tune_correction_n'")
34
35 print("Values before matching")
36 print("dQx.b1=", mad.MADX.dqx_b1)
37 print("dQy.b1=", mad.MADX.dqy_b1)
38
39 mad.send("""
40 expr1 = \\t, s -> t.q1 - 62.30980
41 expr2 = \\t, s -> t.q2 - 60.32154
42 function twiss_and_send()
43 local mtbl, mflow = twiss {sequence=lhcb1, method=4}
44 py:send({mtbl.s, mtbl.beta11})
45 return mtbl, mflow
46 end
47 """)
48 match_rtrn = mad.match(
49 command=mad.twiss_and_send,
50 variables = [
51 {"var":"'MADX.dqx_b1'", "name":"'dQx.b1'", "'rtol'":1e-6},
52 {"var":"'MADX.dqy_b1'", "name":"'dQy.b1'", "'rtol'":1e-6},
53 ],
54 equalities = [
55 {"expr": mad.expr1, "name": "'q1'", "tol":1e-3},
56 {"expr": mad.expr2, "name": "'q2'", "tol":1e-3},
57 ],
58 objective={"fmin": 1e-3}, maxcall=100, info=2,
59 )
60 mad.send("py:send(nil)")
61 tws_result = mad.recv ()
62 x = tws_result[0]
63 y = tws_result[1]
64
65 plt.ion()
66 fig = plt.figure()
67 ax = fig.add_subplot(111)
68 line1, = ax.plot(x, y, 'b-')
69 while tws_result:
70 line1.set_xdata(tws_result[0])
71 line1.set_ydata(tws_result[1])
72 fig.canvas.draw()
73 fig.canvas.flush_events()
74 tws_result = mad.recv()
75
76 mad["status", "fmin", "ncall"] = match_rtrn
77 del match_rtrn
78
79 print("Values after matching")
80 print("dQx.b1=", mad.MADX.dqx_b1)
81 print("dQy.b1=", mad.MADX.dqy_b1)
82
83 mad.twiss("tbl", sequence=mad.lhcb1, method=4, chrom=True)
84 mad.tbl.write("'after_tune_correction_n'")
85 t1 = time.time()
86 print("pre-tracking time: " + str(t1 - t0) + "s")
87
88os.chdir(orginal_dir)
Managing References Example
1#Code that does not necessarily work as expected
2from pymadng import MAD
3import os, numpy as np
4
5orginal_dir = os.getcwd()
6os.chdir(os.path.dirname(os.path.realpath(__file__)))
7
8mad = MAD() #Not being in context manager makes not difference.
9
10#Just boring setup (in lots of other examples)
11mad.MADX.load(f"'fodo.seq'", f"'fodo.mad'")
12mad["seq"] = mad.MADX.seq
13mad.seq.beam = mad.beam()
14
15
16
17#Only one thing is returned from the twiss, a reference (Nothing in python is ever received from MAD after telling MAD-NG to execute)
18twissrtrn = mad.twiss(sequence=mad.seq, method=4)
19
20#Any high level MAD-NG function will create a reference
21mad.MAD.gmath.reim(1.42 + 0.62j)
22
23#Try to receive from twiss
24mad["mtbl", "mflw"] = twissrtrn
25
26# mtbl and mflow correctly stored!
27print(mad.mtbl)
28print(mad.mflw)
29
30myMatrix = mad.MAD.matrix(4).seq() #Create 4x4 matrix
31
32print(type(myMatrix)) #Not a 4x4 matrix!
33print(type(myMatrix.eval())) #A 4x4 matrix!
34
35myMatrix = myMatrix.eval() #Store the matrix permanantly
36
37mad["myMatrix"] = mad.MAD.matrix(4).seq()
38print(mad.myMatrix, np.all(myMatrix == mad.myMatrix))
39
40os.chdir(orginal_dir)
Receiving Data from the LHC Example
1from pymadng import MAD
2import time, os
3current_dir = os.path.dirname(os.path.realpath(__file__)) + "/"
4
5mad = MAD()
6
7mad.send(f"""
8function LHC_load () --Not local!
9 local beam in MAD
10 local assertf in MAD.utility
11
12 MADX:load('{current_dir}lhc_as-built.seq', '{current_dir}lhc_as-built.mad') -- convert and save on need
13 MADX:load('{current_dir}opt_400_10000_400_3000.madx', '{current_dir}opt_400_10000_400_3000.mad') -- ditto
14 MADX:load('{current_dir}lhc_unset_vars.mad') -- handmade, cleaner than disabling warnings...
15
16 local lhcb1, lhcb2 in MADX
17
18 -- sanity checks
19 local n1, n2 = #lhcb1, #lhcb2
20 assertf(n1 == 6677, "invalid number of elements %d in LHCB1 (6677 expected)", n1)
21 assertf(n2 == 6676, "invalid number of elements %d in LHCB2 (6676 expected)", n2)
22 py:send("done")
23end
24""")
25
26mad.send("""
27local is_number, is_string, is_function, is_table in MAD.typeid
28
29expr = {} -- list of deferred expressions
30
31function reg_expr(k,v) -- collect deferred expressions
32 if is_number(v) or is_string(v) then
33 ;
34 elseif is_function(v) then -- deferred expressions are functions
35 expr[#expr+1] = v
36 elseif is_table(v) then
37 for kk, vv in pairs(v) do reg_expr(kk,vv) end -- recursive call
38 else
39 print(k, v, type(v)) -- unexpected for MAD-X, just for debug
40 end
41end
42""")
43
44t0 = time.time()
45mad.LHC_load ()
46mad.recv() #done
47t1 = time.time()
48print("Load time:", t1-t0, " sec\n")
49
50t0 = time.time()
51mad.reg_expr("MADX", mad.MADX)
52mad.send("py:send('done')").recv() # reg_expr is recursive
53t1 = time.time()
54print("reg_expr time:", t1-t0, " sec\n")
55
56mad.send("for i=1,#expr do expr[i]() end") #So that warnings are performed here and do no affect timing
57
58
59#Methods of evaluation:
60t0 = time.time()
61mad.send("py:send(#expr)")
62for i in range(mad.recv()):
63 mad.send(f"py:send(expr[{i+1}]())").recv()
64t1 = time.time()
65print("eval time method 1:", t1-t0, " sec")
66
67t0 = time.time()
68mad.send("len = #expr py:send(len) for i=1,len do py:send(expr[i]()) end")
69for i in range(mad.recv()):
70 mad.recv()
71t1 = time.time()
72print("eval time method 2:", t1-t0, " sec")
73
74t0 = time.time()
75mad.send("py:send(#expr)")
76exprList1 = [mad.send(f"py:send(expr[{i+1}]())").recv() for i in range(mad.recv())]
77t1 = time.time()
78print("eval time method 3:", t1-t0, " sec")
79
80t0 = time.time()
81mad.send("len = #expr py:send(len) for i=1,len do py:send(expr[i]()) end")
82exprList2 = [mad.recv() for i in range(mad.recv())]
83t1 = time.time()
84print("eval time method 4:", t1-t0, " sec\n")
85
86
87print("sanity check", exprList1 == exprList2, len(exprList1))
88
89t0 = time.time()
90mad["lhcb1"] = mad.MADX.lhcb1
91
92mad.send("""
93py:send(#lhcb1)
94for i, elm, spos, len in lhcb1:iter() do
95 py:send(elm.name)
96end
97""")
98nameList = [mad.recv() for _ in range(mad.recv())]
99t1 = time.time()
100print("time to retrieve every element name in lhcb1 sequence", t1-t0, "sec")
101print(len(nameList))
102
103
104t0 = time.time()
105mad["lhcb2"] = mad.MADX.lhcb2
106mad.send("""
107lhcb2_tbl = {}
108for i, elm, spos, len in lhcb2:iter() do
109 lhcb2_tbl[i] = elm.name
110end
111py:send(lhcb2_tbl)
112""")
113nameList = mad.recv()
114t1 = time.time()
115print("time to retrieve every element name in lhcb2 sequence", t1-t0, "sec")
116print(len(nameList))