List of Examples

FODO Cell Example

import os
from pathlib import Path

import matplotlib.pyplot as plt

from pymadng import MAD

original_dir = Path.cwd()
os.chdir(Path(__file__).resolve().parent)

# The typical way to communicate with MAD-NG is to use the send and recv methods.
with MAD() as mad:
    mad.send("""
    MADX:load("fodo.seq", "fodo.mad")
    local seq in MADX
    seq.beam = beam -- use default beam
    mtbl, mflw = twiss {sequence=seq, implicit=true, nslice=10, save="atbody"}
    py:send(mtbl)
    """)
    mtbl = mad.recv("mtbl")
    plt.plot(mtbl.s, mtbl.beta11, "r-", label="Method 1")

    mad.send("""py:send(mtbl.s) ; py:send(mtbl.beta11)""")
    plt.plot(
        mad.recv(),
        mad.recv(),
        "g--",
        label="Method 2",
    )

    plt.plot(mad.mtbl.s, mad.mtbl.beta11, "b:", label="Method 3")

    plt.legend()
    plt.show()

# If you prefer to use pythonic syntax, the following code is equivalent to the above (- some plotting)(+ writing to a file)
with MAD() as mad:
    mad.MADX.load("'fodo.seq'", "'fodo.mad'")
    mad.load("MADX", "seq")
    mad.seq.beam = mad.beam()
    mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, implicit=True, nslice=10, save="'atbody'")
    cols = mad.quote_strings(["name", "s", "beta11", "beta22", "mu1", "mu2", "alfa11", "alfa22"])
    mad.mtbl.write("'twiss_py.tfs'", cols)

    for x in mad.seq:  # If an object is iterable, it is possible to loop over it
        print(x.name, x.kind)

    plt.plot(mad.mtbl.s, mad.mtbl.beta11)
    plt.show()

os.chdir(original_dir)

PS Twiss Example

import os
from pathlib import Path

from pymadng import MAD

original_dir = Path.cwd()
os.chdir(Path(__file__).parent)

with MAD(debug=False) as mad:
    mad["psbeam"] = mad.beam(particle="'proton'", pc=2.794987)
    mad.MADX.BEAM = mad.psbeam
    mad.MADX.BRHO = mad.psbeam.brho
    mad.MADX.load("'ps_unset_vars.mad'")
    mad.MADX.load("'ps_mu.seq'")
    mad.MADX.load("'ps_ss.seq'")
    mad.MADX.load("'ps_fb_lhc.str'")

    mad.load("MADX", "ps")
    mad.ps.beam = mad.psbeam
    mad["srv", "mflw"] = mad.survey(sequence=mad.ps)

    mad.srv.write(
        "'PS_survey_py.tfs'",
        mad.quote_strings(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]),
    )

    mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3)

    mad.load("MAD.gphys", "melmcol")
    # Add element properties as columns
    # fmt: off
    mad.melmcol(mad.mtbl,
        mad.quote_strings(
            ["angle", "tilt", "k0l", "k1l", "k2l", "k3l", "k4l", "k5l", "k6l", "k0sl",
            "k1sl", "k2sl", "k3sl", "k4sl", "k5sl", "k6sl", "ksl", "hkick", "vkick" ]),
        )

    mad.mtbl.write("'PS_twiss_py.tfs'",
        mad.quote_strings(
            ["name", "kind", "s", "x", "px", "beta11", "alfa11", "beta22", "alfa22","dx",
            "dpx", "mu1", "mu2", "l", "angle", "k0l", "k1l", "k2l", "k3l", "hkick", "vkick"]
            )
        )
    # fmt: on
    df = mad.mtbl.to_df()
    try:
        import tfs

        assert type(df) is tfs.TfsDataFrame
        print("tfs-pandas installed, so the header is stored in headers")
        print(df.headers)
    except ImportError:
        print("tfs-pandas not installed, so the header is stored in attrs instead of headers")
        print(df.attrs)

    print(df)
    print(mad.srv.to_df())

os.chdir(original_dir)

LHC Example

"""
This script demonstrates the usage of the MAD class to perform a local coupling correction in the LHC.

The aim of this script is demonstrate a combination of pythonic and MAD-NG syntax.
Also, it demonstrates how you can retrieve data from MAD-NG and plot it in real-time, as it is being calculated, preventing the need to store it in memory.
"""

import os
import time
from pathlib import Path

import matplotlib.pyplot as plt

from pymadng import MAD

original_dir = Path.cwd()
os.chdir(Path(__file__).parent)


with MAD() as mad:
    mad.load("MAD.utility", "assertf")

    mad.MADX.load("'lhc_as-built.seq'", "'lhc_as-built.mad'")
    mad.MADX.load("'opticsfile.21'", "'opticsfile.21.mad'")
    mad.MADX.load("'lhc_unset_vars.mad'")  # Load a list of unset variables to prevent warnings

    mad.load("MADX", "lhcb1", "nrj")

    mad.assertf(
        "#lhcb1 == 6694",
        "'invalid number of elements %d in LHCB1 (6694 expected)'",
        "#lhcb1",
    )

    mad.lhcb1.beam = mad.beam(particle="'proton'", energy=mad.nrj)
    mad.evaluate_in_madx_environment("""
    ktqx1_r2 = -ktqx1_l2 ! remove the link between these 2 vars
    kqsx3_l2 = -0.0015
    kqsx3_r2 = +0.0015
    """)
    t0 = time.time()
    mad["tbl", "flw"] = mad.twiss(sequence=mad.lhcb1)
    mad["srv"] = mad.survey(sequence=mad.lhcb1)
    print(mad.srv)
    print(mad.srv[0].to_df())
    mad.tbl.write("'before_tune_correction_n'")

    print("Values before matching")
    print("dQx.b1=", mad.MADX.dqx_b1)
    print("dQy.b1=", mad.MADX.dqy_b1)

    mad.send("""
    expr1 = \\t, s -> t.q1 - 62.30980
    expr2 = \\t, s -> t.q2 - 60.32154
    function twiss_and_send()
        local mtbl, mflow = twiss {sequence=lhcb1}
        py:send({mtbl.s, mtbl.beta11}, true) -- True means that the data table is sent as a shallow copy
        return mtbl, mflow
    end
    """)
    match_rtrn = mad.match(
        command=mad.twiss_and_send,
        variables=[
            {"var": "'MADX.dqx_b1'", "name": "'dQx.b1'", "'rtol'": 1e-6},
            {"var": "'MADX.dqy_b1'", "name": "'dQy.b1'", "'rtol'": 1e-6},
        ],
        equalities=[
            {"expr": mad.expr1, "name": "'q1'", "tol": 1e-3},
            {"expr": mad.expr2, "name": "'q2'", "tol": 1e-3},
        ],
        objective={"fmin": 1e-3},
        maxcall=100,
        info=2,
    )
    mad.send("py:send(nil)")
    tws_result = mad.recv()
    x = tws_result[0]
    y = tws_result[1]

    plt.ion()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    (line1,) = ax.plot(x, y, "b-")
    while tws_result:
        line1.set_xdata(tws_result[0])
        line1.set_ydata(tws_result[1])
        fig.canvas.draw()
        fig.canvas.flush_events()
        tws_result = mad.recv()

    mad["status", "fmin", "ncall"] = match_rtrn
    del match_rtrn

    print("Values after matching")
    print("dQx.b1=", mad.MADX.dqx_b1)
    print("dQy.b1=", mad.MADX.dqy_b1)

    mad.twiss("tbl", sequence=mad.lhcb1)
    mad.tbl.write("'after_tune_correction_n'")
    t1 = time.time()
    print("Matching time: " + str(t1 - t0) + "s")

os.chdir(original_dir)

Managing References Example

# Code that does not necessarily work as expected
import os
from pathlib import Path

import numpy as np

from pymadng import MAD

original_dir = Path.cwd()
os.chdir(Path(__file__).parent)

mad = MAD()  # Not being in context manager makes not difference.

# Just boring setup (in lots of other examples)
mad.MADX.load("'fodo.seq'", "'fodo.mad'")
mad["seq"] = mad.MADX.seq
mad.seq.beam = mad.beam()


# Only one thing is returned from the twiss, a reference (Nothing in python is ever received from MAD after telling MAD-NG to execute)
twissrtrn = mad.twiss(sequence=mad.seq)

# Any high level MAD-NG function will create a reference
mad.MAD.gmath.reim(1.42 + 0.62j)

# Try to receive from twiss
mad["mtbl", "mflw"] = twissrtrn

# mtbl and mflow correctly stored!
print(mad.mtbl)
print(mad.mflw[0])
mad.send(
    "print(mtbl, mflw[1])"
)  # This will print the table and the particle stored in mflw[1] (lua table)

a_matrix = mad.MAD.matrix(4).seq()  # Create 4x4 matrix

print(type(a_matrix))  # Not a 4x4 matrix!
print(type(a_matrix.eval()))  # A 4x4 matrix!

a_matrix = a_matrix.eval()  # Store the matrix permanantly

mad["myMatrix"] = mad.MAD.matrix(4).seq()
print(mad.myMatrix, np.all(a_matrix == mad.myMatrix))

os.chdir(original_dir)

Receiving Data from the LHC Example

import time
from pathlib import Path

from pymadng import MAD

current_dir = str(Path(__file__).parent) + "/"

mad = MAD()

mad.send(f"""
function LHC_load () --Not local!
  local beam in MAD
  local assertf in MAD.utility

  MADX:load('{current_dir}lhc_as-built.seq', '{current_dir}lhc_as-built.mad') -- convert and save on need
  MADX:load('{current_dir}opt_400_10000_400_3000.madx', '{current_dir}opt_400_10000_400_3000.mad') -- ditto
  MADX:load('{current_dir}lhc_unset_vars.mad') -- handmade, cleaner than disabling warnings...

  local lhcb1, lhcb2 in MADX

  -- sanity checks
  local n1, n2 = #lhcb1, #lhcb2
  assertf(n1 == 6677, "invalid number of elements %d in LHCB1 (6677 expected)", n1)
  assertf(n2 == 6676, "invalid number of elements %d in LHCB2 (6676 expected)", n2)
  py:send("done")
end
""")

mad.send("""
local is_number, is_string, is_function, is_table in MAD.typeid

expr = {} -- list of deferred expressions

function reg_expr(k,v) -- collect deferred expressions
  if is_number(v) or is_string(v) then
    ;
  elseif is_function(v) then -- deferred expressions are functions
    expr[#expr+1] = v
  elseif is_table(v) then
    for kk, vv in pairs(v) do reg_expr(kk,vv) end -- recursive call
  else
    print(k, v, type(v)) -- unexpected for MAD-X, just for debug
  end
end
""")

t0 = time.time()
mad.LHC_load()
mad.recv()  # done
t1 = time.time()
print("Load time:", t1 - t0, " sec\n")

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

mad.send(
    "for i=1,#expr do expr[i]() end"
)  # So that warnings are performed here and do no affect timing


# Methods of evaluation:
t0 = time.time()
mad.send("py:send(#expr)")
for i in range(mad.recv()):
    mad.send(f"py:send(expr[{i + 1}]())").recv(f"expr[{i + 1}]()")
t1 = time.time()
print("eval time method 1:", t1 - t0, " sec")

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(f"expr[{i + 1}]()")
t1 = time.time()
print("eval time method 2:", t1 - t0, " sec")

t0 = time.time()
mad.send("py:send(#expr)")
expr_lst1 = [
    mad.send(f"py:send(expr[{i + 1}]())").recv(f"expr[{i + 1}]()") for i in range(mad.recv())
]
t1 = time.time()
print("eval time method 3:", t1 - t0, " sec")

t0 = time.time()
mad.send("len = #expr py:send(len) for i=1,len do py:send(expr[i]()) end")
expr_lst2 = [mad.recv(f"expr[{i + 1}]()") for i in range(mad.recv())]
t1 = time.time()
print("eval time method 4:", t1 - t0, " sec\n")


print("sanity check", expr_lst1 == expr_lst2, len(expr_lst1))

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
""")
list_of_names = [mad.recv() for _ in range(mad.recv())]
t1 = time.time()
print("time to retrieve every element name in lhcb1 sequence", t1 - t0, "sec")
print(len(list_of_names))


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)
""")
list_of_names = mad.recv("lhcb2_tbl")
t1 = time.time()
print("time to retrieve every element name in lhcb2 sequence", t1 - t0, "sec")
print(len(list_of_names))