diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 33ede8e..79eb210 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -37,8 +37,8 @@ jobs: - name: Get MAD Binaries run: | mkdir src/pymadng/bin - curl https://madx.web.cern.ch/releases/madng/1.1/mad-linux-1.1.2 -o src/pymadng/bin/mad_Linux - curl https://madx.web.cern.ch/releases/madng/1.1/mad-macos-1.1.2 -o src/pymadng/bin/mad_Darwin + curl https://madx.web.cern.ch/releases/madng/1.1/mad-linux-1.1.3 -o src/pymadng/bin/mad_Linux + curl https://madx.web.cern.ch/releases/madng/1.1/mad-macos-1.1.3 -o src/pymadng/bin/mad_Darwin chmod +x src/pymadng/bin/mad_Linux src/pymadng/bin/mad_Darwin - name: Build package run: python -m build diff --git a/.github/workflows/test-pymadng.yml b/.github/workflows/test-pymadng.yml index eaa000a..1176d8c 100644 --- a/.github/workflows/test-pymadng.yml +++ b/.github/workflows/test-pymadng.yml @@ -33,8 +33,8 @@ jobs: - name: Get MAD Binaries run: | mkdir ./src/pymadng/bin - curl https://madx.web.cern.ch/releases/madng/1.1/mad-linux-1.1.2 -o ./src/pymadng/bin/mad_Linux - curl https://madx.web.cern.ch/releases/madng/1.1/mad-macos-1.1.2 -o ./src/pymadng/bin/mad_Darwin + curl https://madx.web.cern.ch/releases/madng/1.1/mad-linux-1.1.3 -o ./src/pymadng/bin/mad_Linux + curl https://madx.web.cern.ch/releases/madng/1.1/mad-macos-1.1.3 -o ./src/pymadng/bin/mad_Darwin chmod +x ./src/pymadng/bin/mad_Linux ./src/pymadng/bin/mad_Darwin - name: Install dependencies run: | diff --git a/CHANGELOG.md b/CHANGELOG.md index bd3423c..7947119 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,14 @@ -0.6.3 (2024/30/05) +0.7.0 (2025/06/05) +Update to MAD-NG 1.1.3 \ +Breaking change: tables in lua are always returned as references, so you must use `eval` to get the value of the table. Or use the optional second argument in `py:send` such as `py:send(data, true)` to return the value of the table. \ +Dictionaries in python can now be sent to MAD-NG, and will be converted to a lua table. \ +Add an optional parameter to to_df and convert_to_dataframe methods to allow the user to specify to always return a pandas dataframe, instead of a tfs dataframe, when tfs is installed. \ +Update the documentation and examples to work again. \ +Remove iter restriction on MAD-NG objects that are not sequences, now all objects can be iterated over. \ +Renamed redirect_sterr to redirect_stderr in the MAD object, fixing a typo. \ + + +0.6.3 (2025/04/30) Update to MAD-NG 1.1.2 diff --git a/docs/source/advanced_features.md b/docs/source/advanced_features.md index e6e0304..f3344e4 100644 --- a/docs/source/advanced_features.md +++ b/docs/source/advanced_features.md @@ -67,7 +67,7 @@ MAD-NG supports callbacks and iterative evaluations, which can be tied into Pyth In MAD: ```lua function twiss_and_send() - local tbl, flow = twiss {sequence=seq, method=4} + local tbl, flow = twiss {sequence=seq} py:send({tbl.s, tbl.beta11}) return tbl, flow end diff --git a/docs/source/communication.md b/docs/source/communication.md index 79164ad..0a0db32 100644 --- a/docs/source/communication.md +++ b/docs/source/communication.md @@ -89,8 +89,8 @@ If the object is not an `mtable`, a `TypeError` will be raised. ``` See: -```{literalinclude} ../../examples/ex-ps-twiss/ps-twiss.py -:lines: 18, 24, 41-49 +```{literalinclude} ../../examples/ex-ps-twiss/ex-ps-twiss.py +:lines: 19, 25, 42-53 :linenos: ``` diff --git a/docs/source/contributing.md b/docs/source/contributing.md index a60d143..3621743 100644 --- a/docs/source/contributing.md +++ b/docs/source/contributing.md @@ -64,8 +64,7 @@ python -m unittest tests/*.py ## Best Practices ### Code Style -- Follow PEP8 (enforced via linters) -- Use descriptive names for MAD objects (e.g. `tw`, `flow`, `seq`) +- Use descriptive names for everything - Keep high-level user APIs separate from internal helpers - Use Ruff for code and import formatting. diff --git a/docs/source/ex-lhc-couplingLocal.rst b/docs/source/ex-lhc-couplingLocal.rst index c291134..e378828 100644 --- a/docs/source/ex-lhc-couplingLocal.rst +++ b/docs/source/ex-lhc-couplingLocal.rst @@ -5,7 +5,7 @@ LHC Example :local: :depth: 2 -The file :ref:`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. +The file :ref:`ex-lhc-couplingLocal/ex-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 --------------- @@ -14,8 +14,8 @@ The following lines loads the required variables and files for the example. ``as To grab variables from the MAD-X environment, we use ``mad.load("MADX", ...)``. -.. literalinclude:: ../../examples/ex-lhc-couplingLocal/lhc-couplingLocal.py - :lines: 12-22 +.. literalinclude:: ../../examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py + :lines: 20-28 :linenos: Receving intermediate results @@ -25,80 +25,80 @@ 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. +From lines 10 - 23, we run a match, with a **reference** to the match result returned to the variable ``match_rtrn``. Line 24 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. +The plotting occurs between lines 29 - 38, wtih the while loop continuing until twiss result is ``None``, which occurs when the match is done, as requested on line 24. -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 :doc:`/advanced_features` for more information on these), we delete the reference in python to clear the temporary variable so that is is available for future use. +Finally, on lines 40 and 41, 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 :doc:`/advanced_features` 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. -.. literalinclude:: ../../examples/ex-lhc-couplingLocal/lhc-couplingLocal.py - :lines: 39-77 +.. literalinclude:: ../../examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py + :lines: 50-90 :linenos: -LHC Speed Tests ---------------- +.. 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``. +.. 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 ``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. +.. The code below just runs the function and times it. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 44-47 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 44-47 -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"``. +.. 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"``. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 50-53 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 50-53 -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. +.. 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. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 60-64 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 60-64 -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 +.. 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 -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 67-71 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 67-71 -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. +.. 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. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 74-77, 79-83 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 74-77, 79-83 -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. +.. 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. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 89-99 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 89-99 -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. +.. 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. -.. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py - :lines: 104-114 +.. .. literalinclude:: ../../examples/ex-recv-lhc/ex-defexpr.py +.. :lines: 104-114 -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 +.. 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 -.. code-block:: console +.. .. code-block:: console - Load time: 0.3955872058868408 sec - reg_expr time: 0.034337759017944336 sec +.. Load time: 0.3955872058868408 sec +.. reg_expr time: 0.034337759017944336 sec -For evaluating the deferred expressions +.. For evaluating the deferred expressions -.. code-block:: console +.. .. code-block:: console - 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 +.. 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 -.. code-block:: console +.. .. code-block:: console - time to retrieve every element name in lhcb1 sequence 0.024236202239990234 sec - time to retrieve every element name in lhcb2 sequence 0.0245511531829834 sec \ No newline at end of file +.. time to retrieve every element name in lhcb1 sequence 0.024236202239990234 sec +.. time to retrieve every element name in lhcb2 sequence 0.0245511531829834 sec \ No newline at end of file diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 9320723..bdf22a7 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -19,7 +19,7 @@ PS Twiss Example .. _ex-ps-twiss: -.. literalinclude:: ../../examples/ex-ps-twiss/ps-twiss.py +.. literalinclude:: ../../examples/ex-ps-twiss/ex-ps-twiss.py :linenos: @@ -28,7 +28,7 @@ LHC Example .. _ex-lhc: -.. literalinclude:: ../../examples/ex-lhc-couplingLocal/lhc-couplingLocal.py +.. literalinclude:: ../../examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py :linenos: Managing References Example diff --git a/docs/source/quickstartguide.md b/docs/source/quickstartguide.md index 34b0e02..09fe983 100644 --- a/docs/source/quickstartguide.md +++ b/docs/source/quickstartguide.md @@ -53,7 +53,7 @@ mad.seq.beam = mad.beam() ### Low-Level: ```python -mad.send("seq.beam = beam") +mad.send("seq.beam = beam {}") ``` --- @@ -63,13 +63,13 @@ mad.send("seq.beam = beam") ### High-Level: ```python -mad["tbl", "flow"] = mad.twiss(sequence=mad.seq, method=4) +mad["tbl", "flw"] = mad.twiss(sequence=mad.seq) ``` ### Low-Level: ```python -mad.send("tbl, flow = twiss {sequence=seq, method=4}") +mad.send("tbl, flw = twiss {sequence=seq}") mad.send("py:send(tbl)") tbl = mad.recv() ``` diff --git a/examples/ex-LowLevel/ex-send-multypes.py b/examples/ex-LowLevel/ex-send-multypes.py index b65706a..1517de2 100644 --- a/examples/ex-LowLevel/ex-send-multypes.py +++ b/examples/ex-LowLevel/ex-send-multypes.py @@ -1,7 +1,9 @@ -from pymadng import MAD -import numpy as np import time +import numpy as np + +from pymadng import MAD + arr0 = np.zeros((10000, 1000)) + 1j # 2*10000*1000*8 -> 160 MB mad = MAD() @@ -20,46 +22,59 @@ {0} = cm1 {1} {2} py:send({0})""" -mad.send(cmatrixString.format("cm4", "*", 1)) ## Set cm4 to cm1 * 1 and send it to Python -mad.send(cmatrixString.format("cm1", "*", 2)) ## Set cm1 to cm1 * 2 and send it to Python -mad.send(cmatrixString.format("cm2", "*", 2)) ## Set cm2 to cm1 * 2 and send it to Python -mad.send(cmatrixString.format("cm3", "/", 3)) ## Set cm3 to cm1 / 3 and send it to Python +mad.send( + cmatrixString.format("cm4", "*", 1) +) ## Set cm4 to cm1 * 1 and send it to Python +mad.send( + cmatrixString.format("cm1", "*", 2) +) ## Set cm1 to cm1 * 2 and send it to Python +mad.send( + cmatrixString.format("cm2", "*", 2) +) ## Set cm2 to cm1 * 2 and send it to Python +mad.send( + cmatrixString.format("cm3", "/", 3) +) ## Set cm3 to cm1 / 3 and send it to Python ## Create a vector in MAD and send it to Python mad.send(""" local v1 = (MAD.vector(45):seq()*2 + 1)/3 py:send(v1) """) -start_time = time.time() # Start timer +start_time = time.time() # Start timer # Receive the matrices and vectors -m1 = mad.recv() +m1 = mad.recv() cm4 = mad.recv() cm1 = mad.recv() cm2 = mad.recv() cm3 = mad.recv() v1 = mad.recv() -print(time.time() - start_time) # Print time +print(time.time() - start_time) # Print time # Check if the matrices have been correctly sent -print(np.all(cm1 == arr0*2)) -print(np.all(cm2 == arr0*2*2)) -print(np.all(cm3 == arr0*2/3)) +print(np.all(cm1 == arr0 * 2)) +print(np.all(cm2 == arr0 * 2 * 2)) +print(np.all(cm3 == arr0 * 2 / 3)) print(np.all(cm4 == arr0)) # Send a list to MAD and receive a changed version back -myList = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2 +my_list = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2 mad.send(""" -local list = py:recv() +list = py:recv() -- Note: no local, so it can be accessed outside this block list[1][1] = 10 list[2][1] = 10 py:send(list) """) -mad.send(myList) -myList[0][0] = 10 -myList[1][0] = 10 -print("receiving lists", mad.recv() == myList) +mad.send(my_list) +my_list[0][0] = 10 +my_list[1][0] = 10 +mad_list = mad.recv("list") +for i, inner_list in enumerate(mad_list): + for j, val in enumerate(inner_list): + print( + f"List value at [{i}][{j}]: {val} == {my_list[i][j]}", val == my_list[i][j] + ) # Send an integer to MAD and receive a changed version back myInt = 4 @@ -94,7 +109,7 @@ py:send(myNil) """) mad.send(None) -print("Nil/None", mad.recv() == None) +print("Nil/None", mad.recv() is None) # Receive ranges from MAD mad.send(""" @@ -102,6 +117,6 @@ py:send(MAD.nrange(3.5, 21.4, 12)) py:send(MAD.nlogrange(1, 20, 20)) """) -print("irng", mad.recv() == range(3 , 12 , 2)) #Py not inclusive, mad is +print("irng", mad.recv() == range(3, 12, 2)) # Py not inclusive, mad is print("rng", mad.recv() == np.linspace(3.5, 21.4, 12)) -print("lrng", np.allclose(mad.recv(), np.geomspace(1, 20, 20))) \ No newline at end of file +print("lrng", np.allclose(mad.recv(), np.geomspace(1, 20, 20))) diff --git a/examples/ex-LowLevel/ex-send-recv.py b/examples/ex-LowLevel/ex-send-recv.py index d053049..faf950a 100644 --- a/examples/ex-LowLevel/ex-send-recv.py +++ b/examples/ex-LowLevel/ex-send-recv.py @@ -1,18 +1,20 @@ -from pymadng import MAD +import time + import numpy as np -import time, sys -arr0 = np.zeros((10000, 1000)) + 1j# 2*10000*1000*8 -> 160 MB +from pymadng import MAD + +arr0 = np.zeros((10000, 1000)) + 1j # 2*10000*1000*8 -> 160 MB start_time = time.time() -mad = MAD(debug = False) +mad = MAD(debug=False) # Number one rule! If you ask mad to send you something, read it! # Don't ask mad to send you data, then try to send more data, this will lead to a deadlock! -mad.send("cm1 = py:recv()") # Tell MAD to receive something -mad.send(arr0) # Send the data to MAD +mad.send("cm1 = py:recv()") # Tell MAD to receive something +mad.send(arr0) # Send the data to MAD -mad.send("cm2 = py:recv()") +mad.send("cm2 = py:recv()") mad.send(arr0) mad.send("cm3 = py:recv()") @@ -37,25 +39,37 @@ """ start_time = time.time() -mad.send(cmatrixString.format("cm1", "*", 2)) # Set cm1 to cm1 * 2 and send it to Python +mad.send( + cmatrixString.format("cm1", "*", 2) +) # Set cm1 to cm1 * 2 and send it to Python cm1 = mad.recv() -mad.send(cmatrixString.format("cm2", "*", 2)) # Set cm2 to cm1 * 2 and send it to Python +mad.send( + cmatrixString.format("cm2", "*", 2) +) # Set cm2 to cm1 * 2 and send it to Python cm2 = mad.recv() -mad.send(cmatrixString.format("cm3", "/", 3)) # Set cm3 to cm1 / 3 and send it to Python +mad.send( + cmatrixString.format("cm3", "/", 3) +) # Set cm3 to cm1 / 3 and send it to Python cm3 = mad.recv() -mad.send(cmatrixString.format("cm4", "*", 1)) # Set cm4 to cm1 * 1 and send it to Python +mad.send( + cmatrixString.format("cm4", "*", 1) +) # Set cm4 to cm1 * 1 and send it to Python cm4 = mad.recv() -mad.send(cmatrixString.format("cm4", "*", 2)) # Set cm4 to cm1 * 1 and send it to Python +mad.send( + cmatrixString.format("cm4", "*", 2) +) # Set cm4 to cm1 * 1 and send it to Python cm4 = mad.recv() -mad.send(cmatrixString.format("cm4", "*", 2)) # Set cm4 to cm1 * 1 and send it to Python +mad.send( + cmatrixString.format("cm4", "*", 2) +) # Set cm4 to cm1 * 1 and send it to Python cm4 = mad.recv() print("time to read 960 MB from MAD", time.time() - start_time) # Check that the data is correct -print(np.all(cm1 == arr0*2)) -print(np.all(cm2 == arr0*2)) -print(np.all(cm3 == arr0/3)) -print(np.all(cm4 == arr0*4)) +print(np.all(cm1 == arr0 * 2)) +print(np.all(cm2 == arr0 * 2)) +print(np.all(cm3 == arr0 / 3)) +print(np.all(cm4 == arr0 * 4)) # Send a string to MAD that will send a string back to Python diff --git a/examples/ex-benchmark-and-fork/ex-benchmark-and-fork.py b/examples/ex-benchmark-and-fork/ex-benchmark-and-fork.py index ee3e798..88fa592 100644 --- a/examples/ex-benchmark-and-fork/ex-benchmark-and-fork.py +++ b/examples/ex-benchmark-and-fork/ex-benchmark-and-fork.py @@ -1,13 +1,17 @@ -from pymadng import MAD -import numpy as np -import os, sys, time +import os +import sys +import time + import matplotlib.pyplot as plt +import numpy as np + +from pymadng import MAD -orginal_dir = os.getcwd() +original_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) pid = os.fork() -#Test 1 +# Test 1 if pid > 0: with MAD() as mad: # open mad process, if you just use mad = MAD(), then be sure to close it afterwords; for multiprocessing, os.fork() can work with the with statement arr0 = np.zeros((10000, 1000)) + 1j # 2*10000*1000*8 -> 160 MB @@ -43,7 +47,7 @@ mad.send(f"myvector = MAD.vector({numVars})") start_time = time.time() for i in range(numVars): - mad.send(f"myvector[{i+1}] = py:recv()") + mad.send(f"myvector[{i + 1}] = py:recv()") mad.send(12345.0) print(f"send {numVars} vals", time.time() - start_time) @@ -53,7 +57,7 @@ start_time = time.time() for i in range(numVars): - mad.send(f"py:send(myvector[{i+1}])") + mad.send(f"py:send(myvector[{i + 1}])") mad.recv() print(f"receive {numVars} vals", time.time() - start_time) @@ -65,20 +69,22 @@ with MAD() as mad: mad.load("element", "quadrupole") # METHOD 1 - mad.MADX.load(f"'fodo.seq'",f"'fodo.mad'") + mad.MADX.load("'fodo.seq'", "'fodo.mad'") mad["seq"] = mad.MADX.seq mad.seq.beam = mad.beam() - mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, chrom=True) + mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq) plt.plot(mad.mtbl.s, mad.mtbl["beta11"]) # plt.show() # METHOD 2 mad["circum", "lcell"] = 60, 20 mad["deferred"] = mad.MAD.typeid.deferred - mad["v"] = mad.create_deferred_expression(f = "lcell/math.sin(math.pi/4)/4", k = "1/v.f") + mad["v"] = mad.create_deferred_expression( + f="lcell/math.sin(math.pi/4)/4", k="1/v.f" + ) - mad["qf"] = mad.quadrupole("knl:={0, v.k}", l = 1) - mad["qd"] = mad.quadrupole("knl:={0, -v.k}", l = 1) + mad["qf"] = mad.quadrupole("knl:={0, v.k}", l=1) + mad["qd"] = mad.quadrupole("knl:={0, -v.k}", l=1) mad.send(""" seq2 = sequence 'seq2' { refer='entry', l=circum, -- assign to seq in scope! @@ -90,10 +96,12 @@ qd { at = 2.5 * lcell }, }""") mad.seq2.beam = mad.beam() - mad["mtbl2", "mflw2"] = mad.twiss(sequence=mad.seq2, method=4, nslice=10, implicit=True, save="'atbody'") + mad["mtbl2", "mflw2"] = mad.twiss( + sequence=mad.seq2, nslice=10, implicit=True, save="'atbody'" + ) plt.plot(mad.mtbl2.s, mad.mtbl2["beta11"]) plt.show() print(mad.mtbl2.header) - sys.exit() # exit the child process + sys.exit() # exit the child process -os.chdir(orginal_dir) \ No newline at end of file +os.chdir(original_dir) diff --git a/examples/ex-fodo/ex-fodos.py b/examples/ex-fodo/ex-fodos.py index 547d7f7..5c87ae7 100644 --- a/examples/ex-fodo/ex-fodos.py +++ b/examples/ex-fodo/ex-fodos.py @@ -1,7 +1,10 @@ -from pymadng import MAD +import os + import matplotlib.pyplot as plt -import os -orginal_dir = os.getcwd() + +from pymadng import MAD + +original_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) # The typical way to communicate with MAD-NG is to use the send and recv methods. @@ -10,17 +13,22 @@ MADX:load("fodo.seq", "fodo.mad") local seq in MADX seq.beam = beam -- use default beam - mtbl, mflw = twiss {sequence=seq, method=4, implicit=true, nslice=10, save="atbody"} + 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.recv(), + mad.recv(), + "g--", + label="Method 2", + ) plt.plot(mad.mtbl.s, mad.mtbl.beta11, "b:", label="Method 3") - + plt.legend() plt.show() @@ -29,14 +37,18 @@ mad.MADX.load("'fodo.seq'", "'fodo.mad'") mad.load("MADX", "seq") mad.seq.beam = mad.beam() - mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, implicit=True, nslice=10, save="'atbody'") - cols = mad.quote_strings(["name", "s", "beta11", "beta22", "mu1", "mu2", "alfa11", "alfa22"]) + 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 + + 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(orginal_dir) \ No newline at end of file +os.chdir(original_dir) diff --git a/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py b/examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py similarity index 85% rename from examples/ex-lhc-couplingLocal/lhc-couplingLocal.py rename to examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py index 939d41f..ca2f151 100644 --- a/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py +++ b/examples/ex-lhc-couplingLocal/ex-lhc-couplingLocal.py @@ -1,9 +1,10 @@ """ 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. +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 @@ -11,7 +12,7 @@ from pymadng import MAD -orginal_dir = os.getcwd() +original_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) @@ -39,9 +40,7 @@ kqsx3_r2 = +0.0015 """) t0 = time.time() - mad["tbl", "flw"] = mad.twiss(sequence=mad.lhcb1, method=4) - # plt.plot(mad.tbl.s, mad.tbl.beta11) - # plt.show() + mad["tbl", "flw"] = mad.twiss(sequence=mad.lhcb1) mad.tbl.write("'before_tune_correction_n'") print("Values before matching") @@ -52,8 +51,8 @@ expr1 = \\t, s -> t.q1 - 62.30980 expr2 = \\t, s -> t.q2 - 60.32154 function twiss_and_send() - local mtbl, mflow = twiss {sequence=lhcb1, method=4} - py:send({mtbl.s, mtbl.beta11}) + 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 """) @@ -94,9 +93,9 @@ print("dQx.b1=", mad.MADX.dqx_b1) print("dQy.b1=", mad.MADX.dqy_b1) - mad.twiss("tbl", sequence=mad.lhcb1, method=4, chrom=True) + mad.twiss("tbl", sequence=mad.lhcb1) mad.tbl.write("'after_tune_correction_n'") t1 = time.time() - print("pre-tracking time: " + str(t1 - t0) + "s") + print("Matching time: " + str(t1 - t0) + "s") -os.chdir(orginal_dir) +os.chdir(original_dir) diff --git a/examples/ex-managing-refs/ex-managing-refs.py b/examples/ex-managing-refs/ex-managing-refs.py index 85bf033..3777bd2 100644 --- a/examples/ex-managing-refs/ex-managing-refs.py +++ b/examples/ex-managing-refs/ex-managing-refs.py @@ -5,7 +5,7 @@ from pymadng import MAD -orginal_dir = os.getcwd() +original_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) mad = MAD() # Not being in context manager makes not difference. @@ -17,7 +17,7 @@ # 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, method=4) +twissrtrn = mad.twiss(sequence=mad.seq) # Any high level MAD-NG function will create a reference mad.MAD.gmath.reim(1.42 + 0.62j) @@ -27,7 +27,10 @@ # mtbl and mflow correctly stored! print(mad.mtbl) -print(mad.mflw) +print(mad.mflw[0]) +mad.send( + "print(mtbl, mflw[1])" +) # This will print the table and the particle stored in mflw[1] (lua table) myMatrix = mad.MAD.matrix(4).seq() # Create 4x4 matrix @@ -39,4 +42,4 @@ mad["myMatrix"] = mad.MAD.matrix(4).seq() print(mad.myMatrix, np.all(myMatrix == mad.myMatrix)) -os.chdir(orginal_dir) +os.chdir(original_dir) diff --git a/examples/ex-ps-twiss/ps-twiss.py b/examples/ex-ps-twiss/ex-ps-twiss.py similarity index 79% rename from examples/ex-ps-twiss/ps-twiss.py rename to examples/ex-ps-twiss/ex-ps-twiss.py index 462eb0f..857bd04 100644 --- a/examples/ex-ps-twiss/ps-twiss.py +++ b/examples/ex-ps-twiss/ex-ps-twiss.py @@ -1,17 +1,18 @@ -import os, time, pandas +import os + from pymadng import MAD -orginal_dir = os.getcwd() +original_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) 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(f"'ps_unset_vars.mad'") - mad.MADX.load(f"'ps_mu.seq'") - mad.MADX.load(f"'ps_ss.seq'") - mad.MADX.load(f"'ps_fb_lhc.str'") + 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 @@ -21,7 +22,7 @@ mad.quote_strings(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]), ) - mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3, chrom=True) + mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3) mad.load("MAD.gphys", "melmcol") #Add element properties as columns @@ -39,13 +40,16 @@ ) df = mad.mtbl.to_df() - print(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(orginal_dir) \ No newline at end of file +os.chdir(original_dir) \ No newline at end of file diff --git a/examples/ex-recv-lhc/ex-defexpr.py b/examples/ex-recv-lhc/ex-defexpr.py index 27c2caa..14e5245 100644 --- a/examples/ex-recv-lhc/ex-defexpr.py +++ b/examples/ex-recv-lhc/ex-defexpr.py @@ -1,5 +1,8 @@ +import os +import time + from pymadng import MAD -import time, os + current_dir = os.path.dirname(os.path.realpath(__file__)) + "/" mad = MAD() @@ -42,46 +45,51 @@ """) t0 = time.time() -mad.LHC_load () -mad.recv() #done +mad.LHC_load() +mad.recv() # done t1 = time.time() -print("Load time:", t1-t0, " sec\n") +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 +mad.send("py:send('done')").recv() # reg_expr is recursive t1 = time.time() -print("reg_expr time:", t1-t0, " sec\n") +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 +mad.send( + "for i=1,#expr do expr[i]() end" +) # So that warnings are performed here and do no affect timing -#Methods of evaluation: +# 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() + mad.send(f"py:send(expr[{i + 1}]())").recv(f"expr[{i + 1}]()") t1 = time.time() -print("eval time method 1:", t1-t0, " sec") +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() + mad.recv(f"expr[{i + 1}]()") t1 = time.time() -print("eval time method 2:", t1-t0, " sec") +print("eval time method 2:", t1 - t0, " sec") t0 = time.time() mad.send("py:send(#expr)") -exprList1 = [mad.send(f"py:send(expr[{i+1}]())").recv() for i in range(mad.recv())] +exprList1 = [ + 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") +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") -exprList2 = [mad.recv() for i in range(mad.recv())] +exprList2 = [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("eval time method 4:", t1 - t0, " sec\n") print("sanity check", exprList1 == exprList2, len(exprList1)) @@ -97,7 +105,7 @@ """) nameList = [mad.recv() for _ in range(mad.recv())] t1 = time.time() -print("time to retrieve every element name in lhcb1 sequence", t1-t0, "sec") +print("time to retrieve every element name in lhcb1 sequence", t1 - t0, "sec") print(len(nameList)) @@ -110,7 +118,7 @@ end py:send(lhcb2_tbl) """) -nameList = mad.recv() +nameList = mad.recv("lhcb2_tbl") t1 = time.time() -print("time to retrieve every element name in lhcb2 sequence", t1-t0, "sec") -print(len(nameList)) \ No newline at end of file +print("time to retrieve every element name in lhcb2 sequence", t1 - t0, "sec") +print(len(nameList)) diff --git a/examples/runall.py b/examples/runall.py index 9a41a29..422ff46 100644 --- a/examples/runall.py +++ b/examples/runall.py @@ -2,9 +2,9 @@ current_dir = os.path.dirname(os.path.realpath(__file__)) + "/" filenames = [ - "ex-benchmark-and-fork/ex-benchmark-and-fork.py", "ex-fodo/ex-fodos.py", "ex-lhc-couplingLocal/lhc-couplingLocal.py", + "ex-benchmark-and-fork/ex-benchmark-and-fork.py", "ex-fodo/ex-fodos.py", "ex-lhc-couplingLocal/ex-lhc-couplingLocal.py", "ex-managing-refs/ex-managing-refs.py", "ex-LowLevel/ex-send-multypes.py", "ex-LowLevel/ex-send-recv.py", - "ex-ps-twiss/ps-twiss.py", "ex-recv-lhc/ex-defexpr.py" + "ex-ps-twiss/ex-ps-twiss.py", "ex-recv-lhc/ex-defexpr.py" ] for name in filenames: print(name) diff --git a/src/pymadng/__init__.py b/src/pymadng/__init__.py index d5941d6..ff4810f 100644 --- a/src/pymadng/__init__.py +++ b/src/pymadng/__init__.py @@ -1,7 +1,7 @@ from .madp_object import MAD __title__ = "pymadng" -__version__ = "0.6.3" +__version__ = "0.7.0" __summary__ = "Python interface to MAD-NG running as subprocess" __uri__ = "https://github.com/MethodicalAcceleratorDesign/MAD-NG.py" diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 746b312..e0c33c3 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -99,7 +99,7 @@ def __str__(self) -> str: Returns: str: The string representation of the MAD-NG reference. """ - val = self._mad.recv_vars(self._name) + val = self._mad.recv_vars(self._name, shallow_copy=True) if isinstance(val, high_level_mad_ref): return repr(val) else: @@ -112,7 +112,7 @@ def eval(self) -> Any: Returns: Any: The evaluated value of the reference in MAD-NG. """ - return self._mad.recv_vars(self._name) + return self._mad.recv_vars(self._name, shallow_copy=True) def __repr__(self): """ @@ -136,7 +136,7 @@ def __dir__(self) -> Iterable[str]: self._mad.protected_send(f""" local modList={{}}; local i = 1; for modname, mod in pairs({name}) do modList[i] = modname; i = i + 1; end - {self._mad.py_name}:send(modList) + {self._mad.py_name}:send(modList, true) """) return [x for x in self._mad.recv() if isinstance(x, str) and x[0] != "_"] @@ -164,10 +164,10 @@ class high_level_mad_object(high_level_mad_ref): def __dir__(self) -> Iterable[str]: if not self._mad.ipython_use_jedi: self._mad.protected_send( - f"{self._mad.py_name}:send({self._name}:get_varkeys(MAD.object))" + f"{self._mad.py_name}:send({self._name}:get_varkeys(MAD.object), true)" ) varnames = self._mad.protected_variable_retrieval( - f"{self._name}:get_varkeys(MAD.object, false)" + f"{self._name}:get_varkeys(MAD.object, false)", shallow_copy=True ) if not self._mad.ipython_use_jedi: @@ -189,9 +189,6 @@ def __call__(self, *args, **kwargs): return last_obj def __iter__(self): - self._mad.send(f"{self._mad.py_name}:send({self._name}:is_instanceOf(sequence))") - is_seq = self._mad.recv() - assert is_seq, "Iteration is only supported for sequences for now" self._iterindex = -1 return self @@ -240,7 +237,7 @@ def convert_to_dataframe(self, columns: list = None, force_pandas: bool = False) f""" local is_vector, is_number, is_string in MAD.typeid local colnames = {py_name}:recv() or {obj_name}:colnames() -- Get the column names -{py_name}:send(colnames) -- Send the column names +{py_name}:send(colnames, true) -- Send the column names -- Loop through all the column names and send them with their data for i, colname in ipairs(colnames) do @@ -265,14 +262,14 @@ def convert_to_dataframe(self, columns: list = None, force_pandas: bool = False) col = tbl end - {py_name}:send(col) -- Send the column data + {py_name}:send(col, true) -- Send the column data end local header = {obj_name}.header -- Get the header names -{py_name}:send(header) -- Send the header names +{py_name}:send(header, true) -- Send the header names for i, attr in ipairs(header) do - {py_name}:send({obj_name}[attr]) -- Send the header data + {py_name}:send({obj_name}[attr], true) -- Send the header data end """ ) diff --git a/src/pymadng/madp_object.py b/src/pymadng/madp_object.py index 7dd9c07..514b664 100644 --- a/src/pymadng/madp_object.py +++ b/src/pymadng/madp_object.py @@ -68,7 +68,7 @@ def __init__( raise_on_madng_error: bool = True, debug: bool = False, stdout: TextIO | str | Path = None, - redirect_sterr: bool = False, + redirect_stderr: bool = False, num_temp_vars: int = 8, ipython_use_jedi: bool = False, ): @@ -84,7 +84,7 @@ def __init__( raise_on_madng_error (bool, optional): If True, raises errors from MAD-NG immediately. debug (bool, optional): If True, enables detailed debugging output. stdout (TextIO | str | Path, optional): Destination for MAD-NG's standard output. - redirect_sterr (bool, optional): If True, redirects stderr to stdout. + redirect_stderr (bool, optional): If True, redirects stderr to stdout. num_temp_vars (int, optional): Maximum number of temporary variables to track. ipython_use_jedi (bool, optional): If True, allows IPython to use jedi for autocompletion. """ @@ -96,7 +96,7 @@ def __init__( raise_on_madng_error=raise_on_madng_error, debug=debug, stdout=stdout, - redirect_sterr=redirect_sterr, + redirect_stderr=redirect_stderr, ) self.__process.ipython_use_jedi = ipython_use_jedi self.__process.last_counter = last_counter(num_temp_vars) @@ -291,17 +291,18 @@ def send_vars(self, **vars: str | int | float | np.ndarray | bool | list): """ self.__process.send_vars(**vars) - def recv_vars(self, *names: str) -> Any: + def recv_vars(self, *names: str, shallow_copy: bool = False) -> Any: """ Retrieve one or more variables from MAD-NG. Args: *names (str): The names of the variables to be fetched from MAD-NG. + shallow_copy (bool, optional): If True, returns a shallow copy of the variables. Returns: The retrieved variable value or a tuple of values if multiple names are provided. """ - return self.__process.recv_vars(*names) + return self.__process.recv_vars(*names, shallow_copy=shallow_copy) # -------------------------------------------------------------------------------------------------------------# @@ -471,7 +472,7 @@ def create_deferred_expression(self, **kwargs) -> mad_high_level_last_ref: def __dir__(self) -> Iterable[str]: pyObjs = [x for x in super(MAD, self).__dir__() if x[0] != "_"] pyObjs.extend(self.globals()) - pyObjs.extend(dir(self.recv_vars("_G"))) + pyObjs.extend(dir(self.recv_vars("_G", shallow_copy=True))) return pyObjs def globals(self) -> list[str]: @@ -481,7 +482,7 @@ def globals(self) -> list[str]: Returns: list[str]: A list containing the names of global variables. """ - return dir(self.__process.recv_vars(f"{self.py_name}._env")) + return dir(self.__process.recv_vars(f"{self.py_name}._env", shallow_copy=True)) def history(self) -> str: """ diff --git a/src/pymadng/madp_pymad.py b/src/pymadng/madp_pymad.py index be493ca..22e17a1 100644 --- a/src/pymadng/madp_pymad.py +++ b/src/pymadng/madp_pymad.py @@ -42,7 +42,7 @@ def __init__( raise_on_madng_error: bool = True, debug: bool = False, stdout: str | Path | TextIO = None, - redirect_sterr: bool = False, + redirect_stderr: bool = False, ) -> None: self.py_name = py_name @@ -76,7 +76,7 @@ def __init__( ) from e # Redirect stderr to stdout, if specified - if redirect_sterr: + if redirect_stderr: stderr = stdout else: stderr = sys.stderr.fileno() @@ -186,10 +186,10 @@ def send_cpx_tpsa(self, monos: np.ndarray, coefficients: np.ndarray) -> None: self.mad_input_stream.write(b"ctpa") send_generic_tpsa(self, monos, coefficients, send_cpx) - def send(self, data: str | int | float | np.ndarray | bool | list) -> mad_process: + def send(self, data: str | int | float | np.ndarray | bool | list | dict) -> mad_process: """Send data to the MAD-NG process. - Accepts several types (str, int, float, ndarray, bool, list) and sends them using the appropriate serialization. + Accepts several types (str, int, float, ndarray, bool, list, dict) and sends them using the appropriate serialization. Returns self to allow method chaining. """ try: @@ -214,19 +214,21 @@ def protected_send(self, string: str) -> mad_process: f"{self.py_name}:__err(true); {string}; {self.py_name}:__err(false);" ) - def protected_variable_retrieval(self, name: str) -> Any: + def protected_variable_retrieval(self, name: str, shallow_copy: bool = False) -> Any: """Safely retrieve a variable from MAD-NG. Enables temporary error handling while retrieving a variable. Args: name (str): The MAD-NG variable name to retrieve. + shallow_copy (bool): If True, retrieves a shallow copy of the variable. This has no effect for most types, but tables in MAD-NG are sent as references by default, so if you want to retrieve a copy of the table, set this to True. Returns: The value of the variable. """ + shallow_copy = str(shallow_copy).lower() if self.raise_on_madng_error: - return self.send(f"py:send({name})").recv(name) + return self.send(f"py:send({name}, {shallow_copy})").recv(name) self.send( - f"{self.py_name}:__err(true):send({name}):__err(false)" + f"{self.py_name}:__err(true):send({name}, {shallow_copy}):__err(false)" ) # Enable error handling, ask for the variable, and disable error handling return self.recv(name) @@ -234,6 +236,10 @@ def set_error_handler(self, on_off: bool) -> mad_process: """Toggle error handling in the MAD-NG process. This determines whether errors are raised immediately. + Args: + on_off (bool): If True, errors will be raised immediately; if False, errors will not be raised. + Returns: + mad_process: Returns self to allow method chaining. """ if self.raise_on_madng_error: return # If the user has specified that they want to raise an error always, skip the error handling on and off @@ -244,6 +250,10 @@ def recv(self, varname: str = None) -> Any: Reads 4 bytes to detect the data type and then extracts the corresponding value. Optional varname is used for reference purposes. + Args: + varname (str): The variable name to use for reference in MAD-NG. + Returns: + Any: The value received from MAD-NG, which can be of various types (str, int, float, ndarray, bool, list, dict). """ typ = self.mad_read_stream.read(4).decode("utf-8") self.varname = varname # For mad reference @@ -254,6 +264,10 @@ def recv_and_exec(self, env: dict = {}) -> dict: The execution context includes numpy as np and the mad process instance. Returns the updated execution environment. + Args: + env (dict): The environment dictionary to execute the received command in. + Returns: + dict: The updated environment dictionary after executing the received command. """ # Check if user has already defined mad (madp_object will have mad defined), otherwise define it try: @@ -269,6 +283,11 @@ def send_vars(self, **vars) -> mad_process: """Send multiple variables to MAD-NG. Each keyword argument becomes a variable in the MAD-NG environment. + If a variable is a mad_ref, it is sent as its name; otherwise, the value is sent directly. + Args: + **vars: Keyword arguments representing variable names and their values. + Returns: + mad_process: Returns self to allow method chaining. """ for name, var in vars.items(): if isinstance(var, mad_ref): @@ -276,18 +295,23 @@ def send_vars(self, **vars) -> mad_process: else: self.send(f"{name} = {self.py_name}:recv()").send(var) - def recv_vars(self, *names) -> Any: + def recv_vars(self, *names, shallow_copy: bool = False) -> Any: """Receive one or multiple variables from MAD-NG. For a single variable (excluding internal names) a direct value is returned. For multiple variables, a tuple of values is returned. + Args: + *names: Variable names to retrieve from MAD-NG. + shallow_copy (bool): If True, retrieves a shallow copy of the variable. This has no effect for most types, but tables in MAD-NG are sent as references by default, so if you want to retrieve a copy of the table, set this to True. + Returns: + Any: The value of the variable if a single name is provided, or a tuple of values if multiple names are provided. """ if len(names) == 1: if not is_private(names[0]): - return self.protected_variable_retrieval(names[0]) + return self.protected_variable_retrieval(names[0], shallow_copy) else: return tuple( - self.protected_variable_retrieval(name) + self.protected_variable_retrieval(name, shallow_copy) for name in names if not is_private(name) ) @@ -390,7 +414,7 @@ def __getitem__(self, item: str | int): def eval(self) -> Any: """Evaluate the reference and return the value.""" - return self._mad.recv_vars(self._name) + return self._mad.recv_vars(self._name, shallow_copy=True) # data transfer -------------------------------------------------------------- # @@ -825,14 +849,50 @@ def recv_list(self: mad_process) -> list: list: The received list. """ varname = self.varname # cache - haskeys = recv_bool(self) lstLen = recv_int(self) vals = [self.recv(varname and varname + f"[{i + 1}]") for i in range(lstLen)] self.varname = varname # reset - if haskeys: - return type_fun["ref_"]["recv"](self) - else: - return vals + return vals + +def send_dict(self: mad_process, dct: dict): + """Send a dictionary to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + dct (dict): The dictionary to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + + Raises: + ValueError: If a key in the dictionary is None, as nil keys are not allowed in Lua. + """ + for key, value in dct.items(): + if key is None: + self.send(None) # Stop the communication of the dictionary to MAD-NG + raise ValueError("nil key in a dictionary is not allowed in lua, remove the key None from the dictionary") + self.send(key) + self.send(value) + self.send(None) + +def recv_dict(self: mad_process) -> dict: + """Receive a dictionary from the MAD-NG pipe. + + Returns: + dict: The received dictionary. + """ + varname = self.varname # cache + dct = {} + while True: + key = self.recv() + if key is None: # End of dictionary + break + if isinstance(key, np.int32): + key = int(key) + value = self.recv(varname and f"{varname}['{key}']") + dct[key] = value + self.varname = varname # reset + return dct # object (table with metatable are treated as pure reference) ---------------- # @@ -877,7 +937,8 @@ def recv_err(self: mad_process): "nil_": {"recv": recv_nil, "send": send_nil}, "bool": {"recv": recv_bool, "send": send_bool}, "str_": {"recv": recv_str, "send": send_str}, - "tbl_": {"recv": recv_list, "send": send_list}, + "lst_": {"recv": recv_list, "send": send_list}, + "dct_": {"recv": recv_dict, "send": send_dict}, "ref_": {"recv": recv_reference, "send": send_reference}, "fun_": {"recv": recv_reference, "send": send_reference}, "obj_": {"recv": recv_reference, "send": send_reference}, @@ -899,7 +960,7 @@ def recv_err(self: mad_process): def get_typestr( - a: str | int | float | np.ndarray | bool | list | tuple | range | mad_ref, + a: str | int | float | np.ndarray | bool | list | dict | tuple | range | mad_ref, ) -> type: """Determine the type string for the given input. @@ -924,8 +985,9 @@ def get_typestr( type(None): "nil_", bool: "bool", str: "str_", - list: "tbl_", - tuple: "tbl_", + list: "lst_", + dict: "dct_", + tuple: "lst_", mad_ref: "ref_", int: "int_", np.int32: "int_", diff --git a/tests/inputs/example.log b/tests/inputs/example.log index ddf98d9..cc5e409 100644 --- a/tests/inputs/example.log +++ b/tests/inputs/example.log @@ -18,16 +18,16 @@ track = MAD.track match = MAD.match ] 194 bytes ***pymad.recv: binary data 4 bytes -***pymad.recv: [py:__err(true):send(MAD):__err(false)] 37 bytes +***pymad.recv: [py:__err(true):send(MAD, false):__err(false)] 44 bytes ***pymad.send: [ref_] 4 bytes ***pymad.recv: binary data 4 bytes -***pymad.recv: [py:__err(true):send(MAD['env']):__err(false)] 44 bytes +***pymad.recv: [py:__err(true):send(MAD['env'], false):__err(false)] 51 bytes ***pymad.send: [ref_] 4 bytes ***pymad.recv: binary data 4 bytes -***pymad.recv: [py:__err(true):send(MAD['env']['version']):__err(false)] 55 bytes +***pymad.recv: [py:__err(true):send(MAD['env']['version'], false):__err(false)] 62 bytes ***pymad.send: [str_] 4 bytes ***pymad.send: binary data 4 bytes -***pymad.send: [1.1.2] 5 bytes +***pymad.send: [1.1.3] 5 bytes ***pymad.recv: binary data 4 bytes ***pymad.recv: [ function __mklast__ (a, b, ...) diff --git a/tests/test_communication.py b/tests/test_communication.py index 5b6847b..2e4e64c 100644 --- a/tests/test_communication.py +++ b/tests/test_communication.py @@ -15,7 +15,7 @@ def test_recv_and_exec(self): self.assertEqual(a, 50) def test_err(self): - with MAD(stdout="/dev/null", redirect_sterr=True) as mad: + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: mad.send("py:__err(true)") mad.send("1+1") #Load error self.assertRaises(RuntimeError, mad.recv) @@ -47,7 +47,7 @@ def test_send(self): self.assertEqual(mad.recv(), initString * 2) def test_protected_send(self): - with MAD(stdout="/dev/null", redirect_sterr=True, raise_on_madng_error=False) as mad: + with MAD(stdout="/dev/null", redirect_stderr=True, raise_on_madng_error=False) as mad: mad.send("py:send('hello world'); a = nil/2") self.assertEqual(mad.recv(), "hello world") # python should not crash mad.send("py:send(1)") diff --git a/tests/test_debug.py b/tests/test_debug.py index b27a711..cb98ef6 100644 --- a/tests/test_debug.py +++ b/tests/test_debug.py @@ -54,7 +54,7 @@ def test_err(self): self.assertFalse("***pymad.run:" in file_text) # Run debug with stderr redirection - with MAD(stdout=self.test_log1, redirect_sterr=True) as mad: + with MAD(stdout=self.test_log1, redirect_stderr=True) as mad: mad.psend("a = nil/2") # receive the error before closing the pipe self.assertRaises(RuntimeError, mad.recv) diff --git a/tests/test_io_and_loading.py b/tests/test_io_and_loading.py index 30739b8..435c71b 100644 --- a/tests/test_io_and_loading.py +++ b/tests/test_io_and_loading.py @@ -36,7 +36,7 @@ def test_run_file(self): a = matrix(4, 5):seq() b = cmatrix(2, 3):seq() """) - with MAD(stdout="/dev/null", redirect_sterr=True) as mad: + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: mad.loadfile("test.mad") self.assertIsNone(mad.matrix) self.assertTrue(np.all(mad.a == self.a)) diff --git a/tests/test_misc_types.py b/tests/test_misc_types.py index 7cbe7b9..7448eeb 100644 --- a/tests/test_misc_types.py +++ b/tests/test_misc_types.py @@ -37,10 +37,11 @@ def test_recv(self): py:send(irng) py:send(rng) py:send(lrng) - py:send(irng:totable()) - py:send(rng:totable()) - py:send(lrng:totable()) + py:send(irng:totable(), true) + py:send(rng:totable(), true) + py:send(lrng:totable(), true) """) + self.assertEqual(mad.recv(), range(3 , 12 , 2)) #MAD is inclusive, python is exclusive (on stop) self.assertTrue (np.allclose(mad.recv(), np.linspace(3.5, 21.4, 12))) self.assertTrue (np.allclose(mad.recv(), np.geomspace(1, 20, 20))) @@ -54,9 +55,9 @@ def test_send(self): irng = py:recv() + 1 rng = py:recv() + 2 lrng = py:recv() - py:send(irng:totable()) - py:send(rng:totable()) - py:send(lrng:totable()) + py:send(irng:totable(), true) + py:send(rng:totable(), true) + py:send(lrng:totable(), true) """) mad.send(range(3, 10, 1)) mad.send_range(3.5, 21.4, 14) @@ -157,8 +158,8 @@ def test_send_tpsa(self): for i = 1, #tab do index_part[i] = tab[i] end - py:send(tab) - py:send(index_part) + py:send(tab, true) + py:send(index_part, true) """) monos = np.asarray([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 0, 0], [1, 1, 0]], dtype=np.uint8) coefficients = [11, 6, 4, 2, 1, 1] @@ -169,11 +170,12 @@ def test_send_tpsa(self): whole_tab = mad.recv("tab") index_part = mad.recv("index_part") self.assertEqual(index_part, expected_index) - for i in range(len(whole_tab)): - self.assertTrue(whole_tab[i] == expected_index[i]) - - for i, key in enumerate(expected_index): - self.assertTrue(whole_tab[key] == coefficients[i]) + for key, value in whole_tab.items(): + if isinstance(key, int): + self.assertTrue(value == expected_index[key-1]) + else: + idx = expected_index.index(key) + self.assertTrue(whole_tab[key] == coefficients[idx]) def test_send_ctpsa(self): with MAD() as mad: @@ -183,8 +185,8 @@ def test_send_ctpsa(self): for i = 1, #tab do index_part[i] = tab[i] end - py:send(tab) - py:send(index_part) + py:send(tab, true) + py:send(index_part, true) """) monos = np.asarray([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 0, 0], [1, 1, 0]], dtype=np.uint8) coefficients = [10+6j, 2+14j, 2+9j, 2+4j, -3+4j, -3+4j] @@ -195,10 +197,12 @@ def test_send_ctpsa(self): whole_tab = mad.recv("tab") index_part = mad.recv("index_part") self.assertEqual(index_part, expected_index) - for i in range(len(whole_tab)): - self.assertTrue(whole_tab[i] == expected_index[i]) - for i, key in enumerate(expected_index): - self.assertTrue(whole_tab[key] == coefficients[i]) + for key, value in whole_tab.items(): + if isinstance(key, int): + self.assertTrue(value == expected_index[key-1]) + else: + idx = expected_index.index(key) + self.assertTrue(whole_tab[key] == coefficients[idx]) def test_send_recv_damap(self): with MAD() as mad: diff --git a/tests/test_numeric_types.py b/tests/test_numeric_types.py index fff265d..033c39e 100644 --- a/tests/test_numeric_types.py +++ b/tests/test_numeric_types.py @@ -7,7 +7,7 @@ def test_send_recv(self): with MAD() as mad: myList = [[1, 2, 3, 4, 5, 6, 7, 8, 9]] * 2 mad.send(""" - local list = py:recv() + list = py:recv() list[1][1] = 10 list[2][1] = 10 py:send(list) @@ -15,25 +15,29 @@ def test_send_recv(self): mad.send(myList) myList[0][0] = 10 myList[1][0] = 10 - self.assertEqual(mad.recv(), myList) + mad_list = mad.recv("list") + for i, inner_list in enumerate(mad_list): + for j, val in enumerate(inner_list): + self.assertEqual(val, myList[i][j], f"Mismatch at index [{i}][{j}]: {val} != {myList[i][j]}") def test_send_recv_wref(self): with MAD() as mad: + python_dict = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, "a": 10, "b": 3, "c": 4} mad.send(""" - list = {MAD.object "a" {a = 2}, MAD.object "b" {b = 2}} - list2 = {1, 2, 3, 4, 5, a = 10, b = 3, c = 4} + list = {MAD.object "a" {a = 2}, MAD.object "b" {b = 6}} + list2 = py:recv() py:send(list) py:send(list2) - """) + """).send(python_dict) list1 = mad.recv("list") list2 = mad.recv("list2") self.assertEqual(len(list1), 2) - self.assertEqual([x for x in list2], [1, 2, 3, 4, 5]) - self.assertEqual(list2["a"], 10) - self.assertEqual(list2["b"], 3) - self.assertEqual(list2["c"], 4) + + self.assertEqual(list2.eval().keys(), python_dict.keys()) + self.assertEqual(sorted(list2.eval().values()), sorted(python_dict.values())) + self.assertEqual(list1[0].a, 2) - self.assertEqual(list1[1].b, 2) + self.assertEqual(list1[1].b, 6) class TestNums(unittest.TestCase): eps = 2**-52 diff --git a/tests/test_object_wrapping.py b/tests/test_object_wrapping.py index 8d8d1cf..d61ad58 100644 --- a/tests/test_object_wrapping.py +++ b/tests/test_object_wrapping.py @@ -15,7 +15,7 @@ class TestGetSet(unittest.TestCase): def test_get(self): - with MAD(stdout="/dev/null", redirect_sterr=True) as mad: + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: mad.load("element", "quadrupole") self.assertEqual(mad.asdfg, None) mad.send("""qd = quadrupole {knl={0, 0.25}, l = 1}""") @@ -24,13 +24,13 @@ def test_get(self): self.assertEqual(qd._name, "qd") self.assertEqual(qd._parent, None) self.assertEqual(qd._mad, mad._MAD__process) - self.assertEqual(qd.knl, [0, 0.25]) + self.assertEqual(qd.knl.eval(), [0, 0.25]) self.assertEqual(qd.l, 1) self.assertRaises(AttributeError, lambda: qd.asdfg) self.assertRaises(KeyError, lambda: qd["asdfg"]) self.assertRaises(IndexError, lambda: qd[1]) self.assertTrue(isinstance(qf.qd, high_level_mad_ref)) - self.assertEqual(qf.qd.knl, [0, 0.25]) + self.assertEqual(qf.qd.knl.eval(), [0, 0.25]) self.assertEqual(qf.qd.l, 1) self.assertEqual(qf.qd, qd) @@ -40,12 +40,12 @@ def test_get(self): if i % 2 != 0: self.assertTrue(isinstance(objList[i].qd, high_level_mad_ref)) self.assertEqual(objList[i].qd._parent, f"objList[{i + 1}]") - self.assertEqual(objList[i].qd.knl, [0, 0.25]) + self.assertEqual(objList[i].qd.knl.eval(), [0, 0.25]) self.assertEqual(objList[i].qd.l, 1) self.assertEqual(objList[i].qd, qd) else: - self.assertEqual(objList[i].knl, [0, 0.25]) + self.assertEqual(objList[i].knl.eval(), [0, 0.25]) self.assertEqual(objList[i].l, 1) self.assertEqual(objList[i], qd) self.assertEqual(objList[i]._parent, "objList") @@ -58,7 +58,7 @@ def test_set(self): # Need more? self.assertEqual(mad.qd2._name, "qd2") self.assertEqual(mad.qd2._parent, None) self.assertEqual(mad.qd2._mad, mad._MAD__process) - self.assertEqual(mad.qd2.knl, [0, 0.25]) + self.assertEqual(mad.qd2.knl.eval(), [0, 0.25]) self.assertEqual(mad.qd2.l, 1) self.assertEqual(mad.qd2, mad.qd) mad["a", "b"] = mad.MAD.gmath.reim(9.75 + 1.5j) @@ -79,7 +79,7 @@ def test_send_vars(self): self.assertEqual(mad.a, 1) self.assertEqual(mad.b, 2.5) self.assertEqual(mad.c, "test") - self.assertEqual(mad.d, [1, 2, 3]) + self.assertEqual(mad.d.eval(), [1, 2, 3]) def test_recv_vars(self): with MAD() as mad: @@ -88,7 +88,7 @@ def test_recv_vars(self): self.assertEqual(a, 1) self.assertEqual(b, 2.5) self.assertEqual(c, "test") - self.assertEqual(d, [1, 2, 3]) + self.assertEqual(d.eval(), [1, 2, 3]) def test_quote_strings(self): with MAD() as mad: @@ -122,7 +122,7 @@ def test_call_obj(self): self.assertEqual(mad.qd._name, "qd") self.assertEqual(mad.qd._parent, None) self.assertEqual(mad.qd._mad, mad._MAD__process) - self.assertEqual(mad.qd.knl, [0, 0.25]) + self.assertEqual(mad.qd.knl.eval(), [0, 0.25]) self.assertEqual(mad.qd.l, 1) sdc = sd @@ -131,14 +131,14 @@ def test_call_obj(self): self.assertEqual(mad.sd._name, "sd") self.assertEqual(mad.sd._parent, None) self.assertEqual(mad.sd._mad, mad._MAD__process) - self.assertEqual(mad.sd.knl, [0, 0.25, 0.5]) + self.assertEqual(mad.sd.knl.eval(), [0, 0.25, 0.5]) self.assertEqual(mad.sd.l, 1) # Reference counting qd = mad.quadrupole(knl=[0, 0.3], l=1) self.assertEqual(sdc._name, "_last[2]") self.assertEqual(qd._name, "_last[3]") - self.assertEqual(qd.knl, [0, 0.3]) + self.assertEqual(qd.knl.eval(), [0, 0.3]) qd = mad.quadrupole(knl=[0, 0.25], l=1) self.assertEqual(qd._name, "_last[1]") @@ -149,7 +149,7 @@ def test_call_last(self): self.assertEqual(mad.func_test(1)(2)(3), 7) def test_call_fail(self): - with MAD(stdout="/dev/null", redirect_sterr=True) as mad: + with MAD(stdout="/dev/null", redirect_stderr=True) as mad: mad.send("func_test = \\a-> \\b-> \\c-> 'a'+b") mad.func_test(1)(2)(3) self.assertRaises(RuntimeError, lambda: mad.recv()) @@ -225,7 +225,7 @@ def test_MADX(self): """) mad.MADX.load("'test.seq'") self.assertEqual(mad.MADX.qd.l, 1) - self.assertEqual(mad.MADX.qd.knl, [0, 0.25]) + self.assertEqual(mad.MADX.qd.knl.eval(), [0, 0.25]) os.remove("test.seq") def test_evaluate_in_madx_environment(self): @@ -235,7 +235,7 @@ def test_evaluate_in_madx_environment(self): """ mad.evaluate_in_madx_environment(madx_code) self.assertEqual(mad.MADX.qd.l, 1) - self.assertEqual(mad.MADX.qd.knl, [0, 0.25]) + self.assertEqual(mad.MADX.qd.knl.eval(), [0, 0.25]) class TestOps(unittest.TestCase): @@ -265,9 +265,9 @@ def test_matrix(self): ) self.assertTrue(np.all(list(mad.MAD.matrix(10).seq()) == np.arange(1, 101))) self.assertTrue(np.all(mad.MAD.matrix(10).seq().eval() == pyMat)) - self.assertEqual(np.sin(1), mad.math.sin(1).eval()) + self.assertEqual(np.sin(1), mad.math['sin'](1).eval()) self.assertAlmostEqual( - np.cos(0.5), mad.math.cos(0.5).eval(), None, None, 4e-16 + np.cos(0.5), mad.math['cos'](0.5).eval(), None, None, 4e-16 ) # temp vars @@ -310,9 +310,9 @@ def test_kwargs(self): opposite=False, mat=mad.m1, ) - self.assertEqual(sd.knl, [0, 0.25j, 1 + 1j]) + self.assertEqual(sd.knl.eval(), [0, 0.25j, 1 + 1j]) self.assertEqual(sd.l, 1) - self.assertEqual(sd.alist, [1, 2, 3, 5]) + self.assertEqual(sd.alist.eval(), [1, 2, 3, 5]) self.assertEqual(sd.abool, True) self.assertEqual(sd.opposite, False) self.assertTrue(np.all(sd.mat == np.arange(9).reshape((3, 3)) + 1)) @@ -405,8 +405,10 @@ def generalDataFrame(self, headers, DataFrame, force_pandas=False): self.assertEqual(header["boolean"], True) self.assertEqual(header["list"], [1, 2, 3, 4, 5]) tbl = getattr(df, headers)["table"] - self.assertEqual([x for x in tbl], [1, 2]) + self.assertEqual(tbl[1], 1) + self.assertEqual(tbl[2], 2) self.assertEqual(tbl["key"], "value") + self.assertTrue(isinstance(tbl, dict)) self.assertEqual(df["string"].tolist(), ["a", "b", "c", "d", "e"]) self.assertEqual(df["number"].tolist(), [1.1, 2.2, 3.3, 4.4, 5.5]) @@ -415,7 +417,8 @@ def generalDataFrame(self, headers, DataFrame, force_pandas=False): df["complex"].tolist(), [1 + 2j, 2 + 3j, 3 + 4j, 4 + 5j, 5 + 6j] ) self.assertEqual(df["boolean"].tolist(), [True, False, True, False, True]) - self.assertEqual(df["list"].tolist(), [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]) + lists = [the_list.eval() for the_list in df["list"]] + self.assertEqual(lists, [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]) tbl = df["table"].tolist() for i in range(len(tbl)): lst = tbl[i] @@ -469,7 +472,7 @@ def test_eval(self): def test_eval_class(self): with MAD() as mad: - result = mad.math.sqrt(2) + mad.math.log(10) + result = mad.math['sqrt'](2) + mad.math['log'](10) self.assertTrue(isinstance(result, mad_high_level_last_ref)) self.assertEqual(result.eval(), math.sqrt(2) + math.log(10)) @@ -484,7 +487,8 @@ def test_iterate_through_object(self): for elem in mad.my_obj: if elem.name == "qd": self.assertEqual(elem.l, 1.6) - self.assertEqual(elem.knl, [0, 0.25]) + for i in range(2): + self.assertEqual(elem.knl[i], 0.25 if i == 1 else 0) else: self.assertEqual(elem.kind, "marker")