diff --git a/CHANGELOG.md b/CHANGELOG.md index d311117..cccdd80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,10 @@ +0.7.0 (2024/12/05) + +Rewrite documentation +Update to MAD-NG 1.1.1 +Handle opening and closing of MAD-NG process more robustly + + 0.6.0 (2024/12/05) Remove `debug` input variable functionality, now it is only a boolean, and dictates whether the debug information is printed to the console. \ Add `stdout` input to the `MAD` object, this allows the user to redirect the output of the MAD-NG process to a file. \ @@ -53,7 +60,7 @@ Set pymadng to now be in beta. Fix MADX issue Move binaries to the bin folder \ Update MAD-NG binaries \ -Rename files to start with madp_... \ +Rename files to start with madp\_... \ Completely refactor underlying process and remove reliance on mad objects, mad strings and `__last__`, now the process is completely self contained and can be separated into MAD-NG itself. List of changes to PyMAD-NG @@ -70,4 +77,4 @@ Update MAD-NG binaries Update MAD-NG binaries \ Fix bug with negative integer values \ -Initialise CHANGELOG +Initialise CHANGELOG diff --git a/README.md b/README.md index 53efe79..97454c5 100644 --- a/README.md +++ b/README.md @@ -1,30 +1,100 @@ -# pymadng -Python interface to MAD-NG running as subprocess +# PyMAD-NG -Install using below, see [The Python Package Index (PyPi)](https://pypi.org/project/pymadng/); +**Python interface to MAD-NG running as a subprocess** -`pip install pymadng` +[![PyPI version](https://img.shields.io/pypi/v/pymadng.svg)](https://pypi.org/project/pymadng/) +[![Documentation Status](https://readthedocs.org/projects/pymadng/badge/?version=latest)](https://pymadng.readthedocs.io/en/latest/) +[![License](https://img.shields.io/github/license/MethodicalAcceleratorDesign/MAD-NG.py)](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/blob/main/LICENSE) -Familiarising yourself with pymadng -=================================== +--- -First, we recommend familiarising yourself with MAD-NG, documentation can be found [here](https://madx.web.cern.ch/releases/madng/html/). +## 🚀 Installation -Then reading through the Low-Level Example Explained on the [pymadng documentation](https://pymadng.readthedocs.io/en/latest/) should be sufficient (alongside knowledge of MAD-NG), assuming you are not planning to use any "syntactic sugar". If you plan to use the available pythonic looking code, there are plenty of examples to look at. +Install via pip from [PyPI](https://pypi.org/project/pymadng/): -In the documentation, [FODO Examples Explained](https://pymadng.readthedocs.io/en/latest/ex-fodo.html), is a chapter that goes into detail on what is happening on each line of the [FODO example](https://github.com/MethodicalAcceleratorDesign/MADpy/blob/main/examples/ex-fodo/ex-fodos.py), while [LHC Example](https://pymadng.readthedocs.io/en/latest/ex-lhc-couplingLocal.html) gives an example of loading the LHC and how to grab intermediate results from a match. +```bash +pip install pymadng +``` -The only other example that may be of use is the [ps-twiss](https://github.com/MethodicalAcceleratorDesign/MADpy/blob/main/examples/ex-ps-twiss/ps-twiss.py) example. This is an extremely simple example, extending the FODO example to perform a twiss on the PS sequence. -If anything does not seem fully explained, initially check the [API Reference](https://pymadng.readthedocs.io/en/latest/pymadng.html#module-pymadng) and/or the [MAD-NG Documentation](https://mad.web.cern.ch/mad/releases/madng/html/), then feel free to open an [issue](https://github.com/MethodicalAcceleratorDesign/MADpy/issues) so improvements can be made. +--- -Documentation -============= +## 🧠 Getting Started -Documentation, including explanation of a couple of examples and the limitations of the API can be found [here](https://pymadng.readthedocs.io/en/latest/). +Before diving into PyMAD-NG, we recommend you: -The API reference is also included in [this documentation](https://pymadng.readthedocs.io/en/latest/). You can also compile to documentation yourself by cloning the repository and running ``make html`` in the docs folder. +1. Familiarise yourself with [MAD-NG](https://madx.web.cern.ch/releases/madng/html/) — understanding MAD-NG is essential. +2. Read the [Quick Start Guide](https://pymadng.readthedocs.io/en/latest/) to see how to control MAD-NG from Python. -Getting the examples working -============================ +### Explore Key Examples -You can run the example with `python3 EXAMPLE_NAME.py` +- **[Low-Level Example Explained](https://pymadng.readthedocs.io/en/latest/)** – Learn the fundamentals line-by-line. +- **[FODO Example Breakdown](https://pymadng.readthedocs.io/en/latest/ex-fodo.html)** – Annotated walkthrough of a FODO cell simulation. +- **[LHC Matching Example](https://pymadng.readthedocs.io/en/latest/ex-lhc-couplingLocal.html)** – Real-world optics matching with intermediate feedback. +- **[PS Twiss Example](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/blob/main/examples/ex-ps-twiss/ps-twiss.py)** – Minimal example applying `twiss()` to the Proton Synchrotron. + +If anything seems unclear: +- Refer to the [API Reference](https://pymadng.readthedocs.io/en/latest/pymadng.html#module-pymadng) +- Check the [MAD-NG Docs](https://madx.web.cern.ch/releases/madng/html/) +- Or open an [issue](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/issues) + +--- + +## 📚 Documentation + +Full documentation and example breakdowns are hosted at: +[https://pymadng.readthedocs.io/en/latest/](https://pymadng.readthedocs.io/en/latest/) + +To build locally: + +```bash +git clone https://github.com/MethodicalAcceleratorDesign/MAD-NG.py.git +cd MAD-NG.py/docs +make html +``` + +--- + +## 🧪 Running Examples + +Examples are stored in the `examples/` folder. +Run any script with: + +```bash +python3 examples/ex-fodos.py +``` + +You can also batch-run everything using: + +```bash +python3 runall.py +``` + +--- + +## 💡 Features + +- High-level Python interface to MAD-NG +- Access to MAD-NG functions, sequences, optics, and tracking +- Dynamic `send()` and `recv()` communication +- Python-native handling of MAD tables and expressions +- Optional integration with `pandas` and `tfs-pandas` + +--- + +## 🤝 Contributing + +We welcome contributions! See [`CONTRIBUTING.md`](docs/source/contributing.md) or the [Contributing Guide](https://pymadng.readthedocs.io/en/latest/contributing.html) in the docs. + +Bug reports, feature requests, and pull requests are encouraged. + +--- + +## 📜 License + +PyMAD-NG is licensed under the [MIT License](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/blob/main/LICENSE). + +--- + +## 🙌 Acknowledgements + +Built on top of MAD-NG, developed at CERN. This interface aims to bring MAD's power to the Python ecosystem with minimal friction. \ No newline at end of file diff --git a/docs/source/advanced_features.md b/docs/source/advanced_features.md new file mode 100644 index 0000000..e6e0304 --- /dev/null +++ b/docs/source/advanced_features.md @@ -0,0 +1,204 @@ +```{eval-rst} +.. currentmodule:: pymadng +``` + +# Advanced Features in PyMAD-NG + +This section covers some of the most powerful capabilities in PyMAD-NG. These features allow you to create scalable and complex accelerator workflows by combining the performance of MAD-NG with Python's expressiveness. + +```{contents} +:depth: 1 +:local: +``` + +--- + +## Understanding `_last[]` Temporary Variables + +In MAD-NG, when a command returns a value, it is not automatically captured unless explicitly assigned. PyMAD-NG handles this by assigning results to a set of reserved variables: `_last[1]`, `_last[2]`, etc. + +These are managed internally by PyMAD-NG using a helper class {class}`madp_last.last_counter`, and accessed in Python via references. This allows expressions like: + +```python +result = mad.math.sqrt(2) + mad.math.log(10) +``` + +Behind the scenes, each intermediate operation is stored in a new `_last[i]` reference, then combined. You can access or evaluate the result using `.eval()`: + +```python +print(result.eval()) +``` + +These temporary variables are recycled unless manually stored using: + +```python +mad["my_var"] = result +``` + +This is particularly useful in expressions, multi-step computations, and avoiding naming clutter. + +--- + +## Function and Object References in MAD-NG + +In PyMAD-NG, accessing or calling any MAD-NG function or object returns a Python reference to that MAD-NG entity, rather than immediately executing or resolving it. This enables symbolic chaining and precise control over execution. + +### Example: +```python +r = mad.math.exp(1) +print(type(r)) # high_level_mad_ref +print(r.eval()) # 2.718... +``` + +You can delay evaluation until needed, allowing reuse: +```python +mad["result"] = mad.math.log(10) + mad.math.sin(1) +``` + +This keeps Python responsive and lets MAD-NG do the heavy lifting. + +--- + +## Real-Time Feedback with Python During Matching + +MAD-NG supports callbacks and iterative evaluations, which can be tied into Python logic. One common use is during `match` procedures, where you want to receive intermediate updates. + +### Example Workflow: +In MAD: +```lua +function twiss_and_send() + local tbl, flow = twiss {sequence=seq, method=4} + py:send({tbl.s, tbl.beta11}) + return tbl, flow +end +``` + +In Python: +```python +mad.match( + command=mad.twiss_and_send, + variables=[...], + equalities=[...], + objective={"fmin": 1e-3}, + maxcall=100 +) + +while True: + data = mad.recv() + if data is None: + break + update_plot(data) +``` + +This is ideal for live visualization, feedback loops, or diagnostics during optimization. + +--- + +## Using PyMAD-NG with Multiprocessing + +Because PyMAD-NG communicates with MAD-NG via pipes (not shared memory), you can launch multiple independent MAD processes using `os.fork()` or `multiprocessing`. + +### When to Use This: +- Run parallel simulations or parameter scans +- Avoid reloading large sequences repeatedly + +### Example: +```python +import os +if os.fork() == 0: + mad = MAD() + mad.send("... long running setup ...") + os._exit(0) +``` + +Each process maintains its own MAD instance and data pipeline. + +--- + +## Loading and Using External MAD Files and Modules + +MAD-X and MAD-NG models often consist of `.seq`, `.mad`, `.madx`, or `.str` files. You can load these via the high-level interface: + +```python +mad.MADX.load("'lhc.seq'", "'lhc.mad'") +mad.load("MADX", "lhcb1") +``` + +Or load additional MAD-NG modules: +```python +mad.load("MAD.gphys", "melmcol") +``` + +This loads extended libraries for magnet properties, tracking models, or optics algorithms. + +--- + +## Exporting Results for External Use + +After running a Twiss or Survey, the results are stored in an `mtable`, which can be exported to a TFS file: + +```python +mad.tbl.write("'results.tfs'", mad.quote_strings(["s", "beta11", "mu1"])) +``` + +You can read this file with `tfs-pandas` or use it as input to another tool. + +--- + +## Combining with NumPy and Pandas + +PyMAD-NG integrates cleanly with Python’s data ecosystem: + +- Pass `numpy` arrays to MAD-NG using {func}`MAD.send` +- Use {func}`.to_df` on MAD tables to get Pandas DataFrames +- Use `tfs-pandas` for rich metadata support + +### Example: +```python +import numpy as np +mad.send("my_array = py:recv()") +mad.send(np.linspace(0, 1, 100)) +``` + +This allows direct use of scientific computation tools in tandem with accelerator modeling. + +--- + +## Managing Larger Workflows + +PyMAD-NG supports: +- Loading full files with `mad.loadfile("mysetup.mad")` +- Organising expressions using Python variables +- Retaining command history using: + +```python +print(mad.history()) +``` + +For clean resource management, always use context blocks: +```python +with MAD() as mad: + mad.MADX.load("'lhc.seq'", "'lhc.mad'") +``` + +This ensures the MAD process is correctly shut down when finished. + +--- + +## Summary of Advanced Features + +| Feature | Purpose | +|---------------------------------|--------------------------------------------------| +| `_last[]` Variables | Track intermediate return values symbolically | +| Reference Objects | Access MAD-NG objects with delayed evaluation | +| Matching Feedback | Monitor intermediate results during match | +| Multiprocessing | Run multiple MAD-NG simulations in parallel | +| File and Module Loading | Import sequences, optics files, and Lua modules | +| Table Export | Write TFS files from MAD tables | +| NumPy / Pandas Interoperability | Pass data between Python and MAD-NG seamlessly | +| Project Structuring | Use {func}`MAD.loadfile`, {func}`MAD.history`, and `with` block | + +These tools are designed to give you complete control over your simulations while staying fast and maintainable. + +Next: head over to **Debugging & Troubleshooting** to diagnose and resolve common issues in real-world workflows. + diff --git a/docs/source/architecture.md b/docs/source/architecture.md new file mode 100644 index 0000000..738b1f7 --- /dev/null +++ b/docs/source/architecture.md @@ -0,0 +1,137 @@ +```{eval-rst} +.. currentmodule:: pymadng +``` + +# PyMAD-NG Architecture Overview + +This section explains the internal architecture of PyMAD-NG, outlining how Python interacts with MAD-NG. + +```{note} +Understanding the architecture will help advanced users contribute to development, extend functionality, and debug issues more effectively. +``` + +```{contents} +:depth: 2 +:local: +``` + +--- + +## High-Level Overview + +PyMAD-NG is a **Python wrapper for MAD-NG**, using standard UNIX pipes to manage two-way communication between Python and the MAD-NG subprocess. It provides both: + +- A **low-level string-based API** (like scripting MAD-NG manually) +- A **high-level object-based API** that emulates Pythonic behavior + +At the heart of PyMAD-NG is the {class}`MAD` class, which: + +- Spawns the MAD-NG binary as a subprocess +- Manages sending commands and receiving data +- Handles variable binding, temporary variables, and object references + +--- + +## Communication Pipeline + +The communication between Python and MAD-NG follows this pipeline: + +``` ++-----------+ send() +---------+ +| Python | -------------------> | MAD-NG | +| Script | <------------------- | Process | ++-----------+ recv() +---------+ +``` + +This is implemented using `os.pipe()` and `select.select()` to manage reads and writes asynchronously. + +### Communication Protocol + +- Commands are sent as strings to MAD-NG (e.g., `mad.send("a = 1 + 2")`) +- Responses are requested explicitly via `{func}`MAD.recv` +- Data can be sent and received in binary or string formats, depending on the type + +--- + +## Core Components + +### {class}`MAD` Class ({class}`madp_object`) + +- Main interface for users +- Automatically loads common MAD-NG modules (e.g. `twiss`, `element`, `sequence`) +- Handles naming of Python process in MAD-NG via `py_name` +- Configures subprocess behavior (debug mode, stdout redirection) + +### {class}`madp_pymad.mad_process` + +- Manages low-level pipe setup and subprocess launch +- Provides `send`, `recv`, and `close` methods +- Serialises and deserialises Python <-> MAD data + +### {class}`madp_classes` + +- Wraps MAD-NG references returned to Python +- Allows chained method calls like `mad.math.sin(1)` +- Supports `.eval()` to convert a MAD value into a native Python object +- Implements `__getattr__`, `__getitem__`, `__call__` to emulate Lua table behavior + +### {class}`madp_last` + +- Implements `_last[]` variable management for symbolic return values +- Ensures unique temporary variable usage +- Used in constructing expression chains + +### {class}`madp_strings` + +- Contains helpers for quoting strings for MAD-NG consumption +- Supports quote-conversion for use in file writing and table exporting + +--- + +## Data Types and Serialization + +PyMAD-NG supports a wide range of Python-native and NumPy types, which are automatically serialised and mapped to MAD-NG equivalents. + +| Python Type | MAD-NG Type | Transport Format | +| --------------------- | --------------- | ----------------- | +| `int`, `float`, `str` | number, string | text or binary | +| `bool` | bool | text | +| `complex` | complex | binary | +| `list` | table | serialised list | +| `numpy.ndarray` | matrix, cmatrix | binary | +| `range`, `np.geomspace` | range types | encoded structure | + +Conversion is handled by {func}`MAD.send` and `{func}`MAD.recv`, both methods on the {class}`MAD` class. +--- + +## Dynamic Attributes & Autocompletion + +Because MAD-NG entities are only known at runtime, PyMAD-NG uses dynamic attribute access via `__getattr__`. This means: + +- Tab completion (`dir(mad)`) only works for preloaded or cached attributes +- Use {func}`MAD.globals()` to list current MAD-NG global variables +- {class}`madp_classes.high_level_mad_ref` objects will not show introspectable properties until evaluated + +--- + +## Protected Execution + +All sends are "protected" by default: + +- If MAD-NG returns an error, PyMAD-NG raises a Python `RuntimeError` +- The return type `err_` is automatically detected and escalated +- This avoids crashes and lets users handle exceptions gracefully + +--- + +## Summary + +| Component | Role | +| ---------------------------------------------- | ---------------------------------------------------------- | +| {class}`MAD` class | Main interface and environment for interacting with MAD-NG | +| {class}`madp_pymad.mad_process` | Launches and communicates with MAD-NG subprocess | +| {class}`madp_classes.high_level_mad_ref` | Wraps MAD objects for Pythonic access | +| {class}`madp_last` | Tracks temporary intermediate results from MAD-NG | +| Type dispatch | Maps Python objects to MAD-NG-compatible messages | + +This modular, pipe-based architecture ensures PyMAD-NG remains flexible, efficient, and closely integrated with MAD-NG’s scripting model. \ No newline at end of file diff --git a/docs/source/communication.md b/docs/source/communication.md new file mode 100644 index 0000000..79164ad --- /dev/null +++ b/docs/source/communication.md @@ -0,0 +1,183 @@ +# Communication with MAD-NG + +```{contents} +:depth: 2 +:local: +``` + +## Protocol Overview + +PyMAD-NG communicates with MAD-NG using a **pipe-based protocol**, ensuring efficient, direct, two-way communication between the Python and MAD-NG processes. + +Key points: + +- Data is sent through FIFO pipes (first-in, first-out). +- Commands are sent as MAD-NG script strings (Lua-like). +- Data is retrieved via `{func}`MAD.recv`()` after explicit instruction to send it. +- MAD-NG stdout is redirected to Python, but not intercepted. + +```{important} +You must always **send instructions before sending data**, and **send a request before receiving data**. +``` + +### Example: Basic Communication +```python +from pymadng import MAD +mad = MAD() + +mad.send("a = py:recv()") # Tell MAD-NG to receive +mad.send(42) # Send the value +mad.send("py:send(a)") # Request it back +mad.recv() # Receive the value → 42 +``` + +Both {func}`MAD.send` and {func}`MAD.recv` are the core communication methods. +See the {class}`pymadng.MAD` reference for more details. + +--- + +## Supported Data Types + +The following types can be sent from Python to MAD-NG: + +```{list-table} Supported Send Types +:header-rows: 1 + +* - Python Type + - MAD-NG Type + - Function to Use +* - `None`, `str`, `int`, `float`, `complex`, `bool`, `list` + - Various + - {func}`MAD.send` +* - `numpy.ndarray (float64)` + - `matrix` + - {func}`MAD.send` +* - `numpy.ndarray (complex128)` + - `cmatrix` + - {func}`MAD.send` +* - `numpy.ndarray (int32)` + - `imatrix` + - {func}`MAD.send` +* - `range` + - `irange` + - {func}`MAD.send` +* - `start, stop, size` as float, int + - `range`, `logrange` + - `mad.send_rng()`, `mad.send_lrng()` +* - Complex structures (e.g., TPSA, CTPSA) + - `TPSA`, `CTPSA` + - `mad.send_tpsa()`, `mad.send_ctpsa()` +``` + +For full compatibility, see the {mod}`pymadng.MAD` documentation. + +--- + +## Converting TFS Tables to DataFrames + +If you use {func}`twiss` or {func}`survey`, MAD-NG returns an `mtable`, which can be converted to a Pandas or TFS-style DataFrame: + +```python +mtbl = mad.twiss(...) +df = mtbl.to_df() # Either DataFrame or TfsDataFrame +``` + +If the object is not an `mtable`, a `TypeError` will be raised. + +```{note} +`tfs-pandas` (if installed) will enhance the output with headers and metadata. +``` + +See: +```{literalinclude} ../../examples/ex-ps-twiss/ps-twiss.py +:lines: 18, 24, 41-49 +:linenos: +``` + +--- + +## Avoiding Deadlocks + +Deadlocks can occur if Python and MAD-NG wait on each other to send/receive large data without syncing. + +### Example of a Deadlock +```python +mad.send('arr = py:recv()') +mad.send(arr0) # Large matrix +mad.send('py:send(arr)') # Sends data to Python +mad.send('arr2 = py:recv()') # Asks for new data +mad.send(arr2) # DEADLOCK if previous data not yet received +``` + +```{warning} +Always ensure each {func}`MAD.send` has a matching {func}`MAD.recv` if data is expected back. +``` + +--- + +## Scope: Local vs Global + +MAD-NG uses Lua-style scoping: + +- Variables declared with `local` are temporary. +- Variables without `local` persist across {func}`MAD.send` calls. + +### Example: +```python +mad.send(""" +a = 10 +local b = 20 +print(a + b) +""") + +mad.send("print(a + (b or 5))") # b is nil → 10 + 5 = 15 +``` + +```{tip} +Use `local` to avoid polluting the global MAD-NG namespace. +``` + +--- + +## Customising the Environment + +You can configure the `MAD()` instance with options to better suit your environment: + +### Change the Python alias used inside MAD-NG: +```python +mad = MAD(py_name="python") +``` + +### Specify a custom MAD-NG binary: +```python +mad = MAD(mad_path="/custom/path/to/mad") +``` + +### Enable debug mode: +```python +mad = MAD(debug=True) +``` + +### Increase the number of temporary variables: +```python +mad = MAD(num_temp_vars=10) +``` + +See {meth}`pymadng.MAD.__init__` for all configuration options. + +--- + +```{eval-rst} +.. currentmodule:: pymadng +``` + +## Summary + +- Always match {func}`MAD.send` with {func}`MAD.recv` when data is expected. +- Use `mad.to_df()` for table conversion. +- Avoid deadlocks by receiving before sending again. +- Manage scope using `local` wisely. +- Use configuration flags to tailor behavior. + +For more, see the {doc}`advanced_features`, {doc}`debugging`, and {doc}`function_reference` sections. + diff --git a/docs/source/conf.py b/docs/source/conf.py index 5f70f51..7624228 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -9,7 +9,7 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information project = 'PyMAD-NG' -copyright = '2023, Joshua Gray, Laurent Deniau' +copyright = '2025, Joshua Gray, Laurent Deniau' author = 'Joshua Gray, Laurent Deniau' import pymadng release = pymadng.__version__ @@ -20,7 +20,8 @@ # Add napoleon to the extensions list extensions = [ 'sphinx.ext.autodoc', - 'sphinx.ext.napoleon' + 'sphinx.ext.napoleon', + 'myst_parser' ] templates_path = ['_templates'] diff --git a/docs/source/contributing.md b/docs/source/contributing.md new file mode 100644 index 0000000..a60d143 --- /dev/null +++ b/docs/source/contributing.md @@ -0,0 +1,89 @@ +# Contributing and Extending PyMAD-NG + +This guide is for developers who want to contribute to PyMAD-NG or extend its capabilities. It outlines the structure of the codebase, how to develop new features, and best practices for writing and testing code. + +--- + +## Overview of Repository Structure + +| File / Module | Purpose | +|--------------------------------------|--------------------------------------------------------| +| `src/pymadng/madp_object.py` | Contains the main {class}`pymadng.MAD` class and high-level interface | +| `src/pymadng/madp_pymad.py` | Manages subprocess communication and type handling | +| `src/pymadng/madp_classes.py` | Defines reference object wrappers (`mad_ref`, etc.) | +| `src/pymadng/madp_last.py` | Manages temporary variables like `_last[]` | +| `src/pymadng/madp_strings.py` | Utility for quoting and formatting MAD-compatible text | +| `examples/` | Contains working scripts and tutorials | +| `tests/` | Contains unit tests for the codebase | +| `docs/` | Documentation files | + +--- + +## Getting Started + +### Prerequisites +- Python 3.7+ +- `numpy`, `pandas`, and optionally `tfs-pandas` +- A valid MAD-NG executable (either system or bundled) + +### Install PyMAD-NG in Editable Mode +To contribute to PyMAD-NG, clone the repository and install it in editable mode. This allows you to make changes and test them without reinstalling. +In specific cases, you may be allowed write access to the MAD-NG.py repository. If so, you only need to clone the repository and install it in editable mode. + +```bash +git clone https://github.com//MAD-NG.py.git +cd pymadng +pip install -e . +``` + +--- + +## Development Workflow + +1. **Create a Branch** +```bash +git checkout -b your-feature-name +``` + +2. **Write Code** + +3. **Test Code** + - All tests are located in the `tests/` directory, currently `unittests` are used. + +4. **Run Tests** +```bash +python -m unittest tests/*.py +``` + +5. **Submit a Pull Request** + - Include a clear description of the change + - Reference any related issues + +--- + +## Best Practices + +### Code Style +- Follow PEP8 (enforced via linters) +- Use descriptive names for MAD objects (e.g. `tw`, `flow`, `seq`) +- Keep high-level user APIs separate from internal helpers +- Use Ruff for code and import formatting. + +--- + +## Documentation + +- Update or extend `.md` or `.rst` files as appropriate +- Document new examples in the `examples/` directory +- Ensure all new features are at least in the API reference + +--- + +## Questions or Issues? + +- Open a GitHub Issue +- Tag maintainers in your Pull Request +- See the [Debugging Guide] or [Architecture Overview] for internals + +Thanks for contributing to PyMAD-NG! + diff --git a/docs/source/dataframes.rst b/docs/source/dataframes.rst deleted file mode 100644 index 5c2c74e..0000000 --- a/docs/source/dataframes.rst +++ /dev/null @@ -1,12 +0,0 @@ -Converting TFS tables to Pandas DataFrames ------------------------------------------- - -The package `pandas` is an optional module, that has an inbuilt function to convert TFS tables (called ``mtable`` in MAD-NG) to a `pandas` ``DataFrame`` or a ``TfsDataFrame`` if you have `tfs-pandas` installed. In the example below, we generate an ``mtable`` by doing a survey and twiss on the Proton Synchrotron lattice, and then convert these to a ``DataFrame`` (or ``TfsDataFrame``). - -.. literalinclude:: ../../examples/ex-ps-twiss/ps-twiss.py - :lines: 18, 24, 41-49 - :linenos: - -In this script, we create the variables ``srv`` and ``mtbl`` which are ``mtable``\ s created by ``survey`` and ``twiss`` respectively. Then first, we convert the ``mtbl`` to a ``DataFrame`` and print it, before checking if you have `tfs-pandas` installed to check if we need to print out the header of the TFS table, which is stored in the attrs attribute of the ``DataFrame``, but is automatically printed when using `tfs-pandas`. Then we convert the ``srv`` to a ``DataFrame`` and print it. - -Note: If your object is not an ``mtable`` then this function will raise a ``TypeError``, but it is available to call on all ``object`` types in MAD-NG. \ No newline at end of file diff --git a/docs/source/debugging.md b/docs/source/debugging.md new file mode 100644 index 0000000..3265d33 --- /dev/null +++ b/docs/source/debugging.md @@ -0,0 +1,179 @@ +```{eval-rst} +.. currentmodule:: pymadng +``` + +# Debugging & Troubleshooting in PyMAD-NG + +This guide explains how to diagnose and fix common issues in PyMAD-NG’s communication with MAD-NG. You’ll learn how to enable debug output, inspect logs, redirect streams, and avoid typical pitfalls (like deadlocks or type mismatches). + +```{contents} +:depth: 2 +:local: +``` + +--- + +## 1. Enable Debug Mode + +When you initialize the {class}`MAD` object with `debug=True`, MAD-NG runs in a verbose mode. This prints extra information about each command you send, as well as diagnostic messages from the Lua side. + +```python +from pymadng import MAD + +mad = MAD(debug=True) +``` + +### 1.1 Redirecting Output +By default, PyMAD-NG writes MAD-NG’s standard output to Python’s `sys.stdout`. To redirect it: + +```python +# Send MAD-NG stdout to a file +mad = MAD(debug=True, stdout="mad_debug.log") +``` + +If you need to redirect standard error as well: + +```python +mad = MAD(debug=True, stdout="mad_debug.log", redirect_stderr=True) +``` + +This helps keep logs organized, especially when running long scripts. + +--- + +## 2. Inspecting Command History + +PyMAD-NG keeps track of all **string-based commands** sent to MAD-NG in a history buffer. To review them: + +```python +print(mad.history()) +``` + +This is invaluable for tracing unexpected behavior. For instance, if you suspect a weird command or an incorrect syntax was sent, you can look up the last lines in the history. + +```{note} +Binary data (like large NumPy arrays) won’t appear in `mad.history()`. Only textual commands are recorded. +``` + +--- + +## 3. Communication Rules + +### 3.1 Send Before Receive + +PyMAD-NG uses pipes for **first-in-first-out** communication. If you call: + +```python +mad.recv() # This will block forever if there's nothing to read! +``` + +without telling MAD-NG to `py:send(...)` first, your script will hang. + +**Correct Sequence**: +1. `mad.send(...)` +2. `mad.recv()` + +Any mismatch in these calls can lead to deadlocks. + +### 3.2 Matching Data Transfers + +If you instruct MAD-NG to receive data (`arr = py:recv()`), you must ensure Python **actually sends** that data: + +```python +mad.send("arr = py:recv()") +mad.send(my_array) # Actually transmit the data +``` + +Failing to do so can cause indefinite blocking or partial reads. + +--- + +## 4. Handling Errors + +### 4.1 Protected Sends + +All sends in PyMAD-NG are automatically “protected” by default. If MAD-NG issues an error (`err_`), PyMAD-NG raises a `RuntimeError` on the Python side. + +```python +try: + mad.send("invalid_lua_code").recv() +except RuntimeError as e: + print("Caught MAD-NG error:", e) +``` + +If you need to **ignore** an error and continue, you can instantiate: + +```python +mad = MAD(raise_on_madng_error=False) +``` + +and manually check for failures. + +--- + +## 5. Debugging Subprocess Behavior + +### 5.1 Starting MAD-NG + +During initialization, PyMAD-NG calls: + +``` +mad_binary -q -e "MAD.pymad 'py' {_dbg = true} :__ini(fd)" +``` + +- If `mad_binary` is missing or not executable, you’ll get `FileNotFoundError`. +- If it fails to run, an `OSError` is raised. + +### 5.2 Checking Streams + +- **stdout**: By default prints to Python’s standard output unless you pass `stdout=...`. +- **stderr**: Remains attached to Python’s own stderr, unless you specify `redirect_sterr=True`. + +--- + +## 6. Common Pitfalls & Solutions + +| Issue | Possible Cause | Recommended Fix | +|------------------------------------------|------------------------------------------------------------------|-----------------------------------------------------------| +| **Hang / Deadlock** | Called `mad.recv()` without `mad.send(...)`, or vice versa | Always pair `send()` → `recv()`. Use `debug=True` to see if MAD is expecting data. | +| **BrokenPipeError** | MAD-NG crashed or closed unexpectedly | Re-initialize `MAD()`. Check logs for the underlying error. | +| **“Unsupported data type”** error | Attempted to `send()` an object that PyMAD-NG can’t serialize | Limit data to `str`, `int`, `float`, `bool`, `list`, or `np.ndarray`. | +| **AttributeError / KeyError** accessing a field | Tried to read a reference property without evaluating it first | Call `.eval()` if you need the actual value. | +| **Exceeding `_last[]`** references | Too many temp variables stored in `_last[]` | Manually name them in MAD, or increase `num_temp_vars`. | + +--- + +## 7. Cleaning Up + +If you’re done using MAD-NG, **close** the session: + +```python +mad.close() +``` + +or use Python’s **context manager**: + +```python +with MAD(debug=True) as mad: + mad.send("a = 1 + 2").recv() + ... +# Subprocess automatically ends here +``` + +--- + +## 8. Summary + +- **Enable** `debug=True` to see more logs. +- **Check** `mad.history()` to identify incorrect or unexpected commands. +- **Balance** each `mad.send()` with a `mad.recv()` to avoid deadlocks. +- **Catch** `RuntimeError` to handle failures gracefully. +- **Evaluate** references (`.eval()`) if you need real values from objects in MAD. + +If you still have trouble: +- Look at the [Architecture Overview](architecture.md) for the internal design. +- See [Contributing](contributing.md) for details on how to extend or fix PyMAD-NG’s internals. +- Open an issue on GitHub if you suspect a bug in the code. + +Happy debugging! + diff --git a/docs/source/developer.rst b/docs/source/developer.rst new file mode 100644 index 0000000..f7dc64c --- /dev/null +++ b/docs/source/developer.rst @@ -0,0 +1,37 @@ +Developer API Reference +======================= + + +PyMAD-NG Low Level API (Able to be run independently of madp_object) +-------------------------------------------------------------------- + +.. automodule:: pymadng.madp_pymad + :members: + :undoc-members: + :show-inheritance: + + +PyMAD-NG Reference Layer for Object Communication +------------------------------------------------- + +.. automodule:: pymadng.madp_classes + :members: + :undoc-members: + :show-inheritance: + + +Helper Functions for Communicating Strings to MAD-NG +---------------------------------------------------- + +.. automodule:: pymadng.madp_strings + :members: + :undoc-members: + :show-inheritance: + +Special Last Object Handler +--------------------------- + +.. automodule:: pymadng.madp_last + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/ex-fodo.rst b/docs/source/ex-fodo.rst deleted file mode 100644 index 9ae72ad..0000000 --- a/docs/source/ex-fodo.rst +++ /dev/null @@ -1,105 +0,0 @@ -FODO Examples -============= - -The file :ref:`ex-fodo/ex-fodos.py ` has several methods of loading the same FODO cell and then plotting :math:`s` and :math:`\beta_{xx}` (known as beta11 in MAD-NG). - -For the simplest low level example, see :ref:`here ` - -Simple (higher level) example ------------------------------ - -The first example (below) shows how to load a MADX sequence, perform a twiss on it in MAD-NG and then plot :math:`s` and :math:`\beta_{xx}`. - -Important points from this examples includes the ``mad[*args]`` notation and the use of double quotations in the function call ``mad.MADX.load(...)``. The notation ``mad[*args]`` is a method of creating variables within the MAD-NG environment from python variables. The double quotes in the load function are required because strings are interpreted by MAD-NG as scripts, allowing the use of variables in the string. - -So going through the example line by line; - - 1. Start MAD-NG and name the communication object ``mad`` - 2. Load the MADX sequence in the file `fodo.seq `_ and store the translation into `fodo.mad `_. - 3. Grab the variable ``seq`` from the MADX environment and give it the name ``seq`` in the MAD-NG environment. - 4. Create the default beam object and attach this to the sequence ``seq``. - 5. Run a twiss (the last three arguments are equivalent to the MADX ``plot "interpolate"``) and name the return values ``mtbl`` and ``mflw`` in the MAD-NG environment. - 6. Create the plot of :math:`s` and :math:`\beta_{xx}`. - 7. Display the plot. - -.. literalinclude:: ../../examples/ex-fodo/ex-fodos.py - :lines: 7-13 - :linenos: - -Different (higher level) example --------------------------------- - -Going through the example line by line (ignoring lines that have been explained above); - - - 3. Essentially the same as line 3 in the example above, yet this time you don't control the naming in the MAD-NG environment, equivalent to ``seq = MADX.seq`` in MAD-NG. - - 6. Create a variable cols, for use in the next line. The use of ``py_strs_to_mad_strs`` adds quotes around every string in the list, so that they are interpreted as strings in MAD-NG. - 7. Write the columns above to the file twiss_py.tfs - 8. Loop through the elements of the sequence - 9. Print the name and kind of element in the sequence - - -.. literalinclude:: ../../examples/ex-fodo/ex-fodos.py - :lines: 15-25 - :linenos: - -Lines 8 - 9 are equivalent to - -.. code-block:: - - for i in range(len(mad.seq)): - x = mad.seq[i] - print(x.name, x.kind) - -Creating a sequence from Python -------------------------------- - -In this example, we demonstrate creating a sequence from scratch in MAD-NG from python. - -Going through the example line by line; - - 1. Load the element ``quadrupole`` from the module ``element`` in MAD-NG, making it available directly, equivalent to the python syntax ``from element import quadrupole``. - 2. Create two variables in the MAD-NG environment, ``circum = 60``, ``lcell = 20``. - - 3. Create two variables in the MAD-NG environment, ``sin = math.sin``, ``pi = math.pi`` (in Python this is equivalent to ``from math import sin, pi``). - 4. Create a deferred expression equivalent to ``v := 1/(lcell/sin(pi/4)/4)`` in MAD-X. - - 5. Create a quadrupole named ``qf`` with ``k1 = v.k`` and length of 1. - 6. Create a quadrupole named ``qd`` with ``k1 = -v.k`` and length of 1. - 7. Create a sequence using ``qf`` and ``qd``, specifying the positions and a length of 60. - - 8. Attach a default beam to the sequence. - 9. Run a twiss and name the return values ``mtbl`` and ``mflw`` in the MAD-NG environment. - - 10. Write the columns ``col`` to the file twiss_py.tfs. - - 11. Plot the result (``mad.mtbl["beta11"]`` is equivalent to ``mad.mtbl.beta11``). - -.. literalinclude:: ../../examples/ex-fodo/ex-fodos.py - :lines: 28-54 - :linenos: - - -.. _low-level: - -Pure low level example ----------------------- - -This is the simplest way to communicate with MAD-NG, but it requires the user to actively, deal with references and sending and receiving of specific data. In this example, three methods of receiving data are shown. - -Going through the example line by line; - - 2. Send a string with the instructions (from the above examples) to load a sequence, attach the default beam and perform a twiss, then send to Python ``mtbl``. - - 9. Recieve the object ``mtbl`` from MAD-NG. As this is an object that cannot be replicated in Python, you receive a **reference**. This **reference** must be named the name of the variable in the MAD-NG environment (done within the call of the ``recv`` function), so that you can use the reference to communicate with MAD-NG. - 10. Using the **reference**, grab the attributes ``s`` and ``beta11`` and plot them - - 12. Ask MAD-NG to send the attributes ``s`` and ``beta11``. (Instead of dealing with the reference, we send the attributes that can be received in Python directly, see this :ref:`table ` for more information on these types) - 13. Receive and plot the attributes ``s`` and ``beta11``. - - 15. Using the MAD-NG object, grab the variable ``mtbl`` and then the attributes ``s`` and ``beta11`` to plot them (as above) - -.. literalinclude:: ../../examples/ex-fodo/ex-fodos.py - :lines: 56-73 - :linenos: \ No newline at end of file diff --git a/docs/source/ex-lhc-couplingLocal.rst b/docs/source/ex-lhc-couplingLocal.rst index 59fb5bf..c291134 100644 --- a/docs/source/ex-lhc-couplingLocal.rst +++ b/docs/source/ex-lhc-couplingLocal.rst @@ -1,6 +1,10 @@ LHC Example =========== +.. contents:: + :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. Loading the LHC @@ -25,7 +29,7 @@ From lines 10 - 21, we run a match, with a **reference** to the match result ret 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 :doc:`ex-managing-refs` 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 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. .. 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. diff --git a/docs/source/ex-lowlevel.rst b/docs/source/ex-lowlevel.rst deleted file mode 100644 index 52e23a8..0000000 --- a/docs/source/ex-lowlevel.rst +++ /dev/null @@ -1,113 +0,0 @@ -Low-Level PyMAD-NG -================== - - -Sending and Receiving Multiple types ------------------------------------- - -To start, the file :ref:`ex-LowLevel/ex-send-multypes.py ` imports all the necessary modules, creates a large numpy array, setups the ``mad`` object to communicate with MAD-NG, and then creates a string to send to MAD. - -The string starting on line 9 below creates a 1000x1000 matrix within MAD-NG and generates the matrix into a sequence, see MAD-NG documentation on `mat:seq `_ for more details. Then on line 11, MAD-NG is asked to send the matrix back. - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-multypes.py - :lines: 1-13 - :linenos: - - -The next section creates a complex matrix in MAD-NG named ``cm1``, then creates and asks MAD-NG to send back to python multiple variations of this complex matrix with calculations performed on them. - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-multypes.py - :lines: 15-26 - :linenos: - - -Then the same is done for a single vector of length 45. - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-multypes.py - :lines: 28-32 - :linenos: - -We then receive all of the variables in the same order they were sent, time the transfer and check that the correct calculations we performed. If everything goes well, a time will be followed by four ``True``\ s - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-multypes.py - :lines: 33-49 - :linenos: - -.. important:: The examples above break what we describe as *sequencing*, therefore, if you are not careful with the procedure, you may end up with a **deadlock**. This is because both MAD-NG and python can wait for the other to receive the data that is being sent. See :ref:`below ` for more details. - -The rest of the ex-send-multypes.py file shows and tests sending and receiving some of the available types. - -.. _deadlock: - -Creating Deadlocks ------------------- - -To cause a deadlock, the easiest way is to send and ask for lots of large data without handling the data in python. For example, if we were to ask for a large matrix, and then send another large matrix, without handling the first matrix, then MAD-NG will be waiting for python to receive the first matrix, while python is waiting for MAD-NG to send the second matrix (see below). - -.. code-block:: python - - # Create an array much bigger than 65 kB buffer - arr0 = np.zeros((10000, 1000)) + 1j - # Ask MAD-NG to receive data from python - mad.send('arr = py:recv()') - # Send data to MAD-NG - mad.send(arr0) - - # Ask MAD-NG to send data back to python. - # Since this fills the buffer, it will hang until python receives the data - mad.send('py:send(arr)') # Danger starts here, as MAD-NG is hanging, waiting for python to receive the data - arr2 = arr0 * 2 # Manipulate data in python - - # Ask MAD-NG to receive more data from python - # (this will not be executed until python receives the data from MAD-NG) - mad.send('arr2 = py:recv()') - - # DEADLOCK! You have now filled the buffer on both sides, - # and both are waiting for the other to receive the data, so nothing happens - mad.send(arr2) - - -Sending and Receiving Large Datasets ------------------------------------- - -The file :ref:`ex-LowLevel/ex-send-recv.py ` shows sending 960 MB arrays and then receiving manipulated versions of these arrays and verifies the manipulations. -This example uses very large datasets, that, if not careful, can cause a deadlock. - -On the first line below, we ask MAD-NG to receive some data from python, and then on the second line, we send the data, if we forget to do this, MAD-NG will error or receive gobbledygook as the contents of the matrix, as it tries to interpret the string on line 4 as a matrix. - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-recv.py - :lines: 12-16 - :linenos: - -.. important:: - - Therefore from this and the previous section, the general rules are: - - * If you ask MAD-NG to receive data from python, you must immediately send the data to MAD-NG - * If you ask MAD-NG to send data to python, you should immediately receive the data from MAD-NG - -Sending and Receiving Scripts -""""""""""""""""""""""""""""" - -Finally we test the receiving scripts and execution, where we find the first script sent to MAD-NG must be executed twice by python, as the first command sends a command back to MAD-NG, which sends a command to be executed by python. - -Then the second script explicitly sent to MAD-NG, extends the first script, so that python is required to execute the script 3 times. - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-recv.py - :lines: 61-80 - -On success of both scripts, the print command sent by MAD-NG will only be executed after the explicit print commands, visible in python - -Local and Global Scopes ------------------------ - -If you are used to coding in MAD-NG, you will be familiar with the use of the statement ``local``, which allows you to define variables in the local scope. This is done because as you increase the number of global variables, the slower the code will run, and because otherwise you will get a warning from MAD-NG. In PyMAD-NG, it is not entirely necessary to use, but it is recommended in order to ensure scope management, however you will not get a warning if you do not use it. This is because every use of the ``mad.send()`` function is an separate scope, so if a variable is made local in one scope, it will not be available in another scope. If the ``local`` is ommitted, then the variable will exist within the scope of the ``MAD`` instance, and will be available to all subsequent calls to ``mad.send()``. To see this in action, see below: - -.. code:: python - - mad.send(""" - a = 10 - local b = 20 - print(a+b) -- 30 - """) - mad.send("print(a+(b or 5))") # b is not defined, so it becomes a + 5 = 15 diff --git a/docs/source/ex-managing-refs.rst b/docs/source/ex-managing-refs.rst deleted file mode 100644 index ebcbb71..0000000 --- a/docs/source/ex-managing-refs.rst +++ /dev/null @@ -1,41 +0,0 @@ -Managing References -=================== - -Creating temporary variables ----------------------------- - -Within :numref:`init-listing`, we perform a twiss to begin with (we assume all the correct variables have been defined and setup, so that the twiss runs successfully (see :ref:`ex-managing-refs/ex-managing-refs.py `) - -However, in this example, instead of using the MAD object to define the variable, with ``mad[*args] = twiss(...)``, we simply store the returned variable into a Python variable. - -.. important:: When a function is called from python and executed in MAD-NG, it will **always** return a **reference**. - -When we evaluate the twiss, we create a reference to a temporary return variable (a temporary reference), as is the case whenever you use a high level mad object as a function. [#f1]_ There is only a limited amount of these temporary references, by default the limit is set to 8, this should be plenty, yet it is possibleto increase this number with setting ``num_temp_vars`` in the ``__init__``. The more temporary variables you use at once, the slower MAD will become to retreive and set variables, therefore 8 is only a limit so you cannot slow the code by creating too many temporary variables. - -The twiss function creates the temporary variable as the return of twiss in the MAD-NG environment, then Python returns a **reference** to this temporary variable from the Python function. When we perform ``reim(1.42 + 0.62j)``, the result is stored in a different temporary variable, but is cleared immediately since it has no python reference. When ``mad["mtbl", "mflw"] = twissrtrn`` is performed, in the MAD-NG environment, ``mtbl`` and ``mflw`` are set to the variables returned from the twiss. However, in Python, ``twissrtrn`` is still a reference to the temporary variable created by the twiss function, so it is recommended to clear the temporary variable by deleting the python object with ``del twissrtrn``. - -.. important:: In general, we recommend not storing temporary variables in python, instead set them in the MAD-NG environment using the syntax ``mad[*args]``. Temporary variables are only useful when you wish to delay communication with MAD-NG, see an example on line 48 `here `_. - -.. literalinclude:: ../../examples/ex-managing-refs/ex-managing-refs.py - :lines: 2, 7-8, 16-27 - :caption: Initializing MAD-NG and performing a twiss - :name: init-listing - -For cases such as the ``reim`` function, you might expect, since it only returns plain data, for the function to return plain data (data that has an identical type in Python) and not a **reference**. However, in pymadng, the functions used are extremely generic and are unaware of the input and output data types. *For Python to know about the data types, Python is required to ask MAD-NG about the return values.* If this was done automatically, then one of the more powerful parts of pymadng, receiving data during function calls, would not be possible as Python instead would be waiting for details on what the variable type is from MAD-NG, and MAD-NG would not be able to tell Python anything until it had completed the function. - -Retreiving plain data variables from reference ----------------------------------------------- - -In this part, we create a matrix, where we are calling a function, ``matrix``, therefore, Python only returns a **reference**. A benefit of this scenario is that, as we get a reference, not a numpy array, we can still perform MAD-NG functions on this reference. - -It is for cases like this that the ``eval()`` function has been created. For any plain data type, this will return the plain data from a function call, however if it is any other object in MAD that cannot be transferred normally, you will still receive a reference that you can continue to use to communicate with MAD-NG. - -If we use the normal syntax of ``mad["myMatrix"] = mad.MAD.matrix(4).seq()``, then ``myMatrix`` is set to the temporary variable returned by the functions on the right (which is cleared immediately due to the lack of python reference), then, when we retrieve ``myMatrix`` from the MAD environment, we get the plain data as expected. - -.. important:: The use of ``eval()`` and ``mad[*args] = ...`` requires communication with MAD-NG, so Python will have to wait until MAD-NG has finished executing the function. - -.. literalinclude:: ../../examples/ex-managing-refs/ex-managing-refs.py - :lines: 29- - :caption: Creating a matrices and retrieving the data - -.. [#f1] This reference has **no** information on the result of the function call, instead, we receive the information when the reference is queried. \ No newline at end of file diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 69afc1d..9320723 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -1,21 +1,9 @@ List of Examples ================ -Sending and Receiving Multiple Types ------------------------------------- - -.. _ex-send-multypes: - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-multypes.py - :linenos: - -Sending and Receiving Large Datasets ------------------------------------- - -.. _ex-send-recv: - -.. literalinclude:: ../../examples/ex-LowLevel/ex-send-recv.py - :linenos: +.. contents:: + :local: + :depth: 2 FODO Cell Example ----------------- diff --git a/docs/source/function_reference.md b/docs/source/function_reference.md new file mode 100644 index 0000000..2a49695 --- /dev/null +++ b/docs/source/function_reference.md @@ -0,0 +1,201 @@ +```{eval-rst} +.. currentmodule:: pymadng +``` + +# Useful Functions & Objects in PyMAD-NG + +This page provides a reference-style overview of functions, libraries, and patterns available through a {class}`MAD()` instance. It complements the API reference by highlighting common tools and objects exposed in the high-level interface for scripting and interacting with MAD-NG. + +```{contents} +:depth: 2 +:local: +``` + +--- + +## Top-Level Objects in `pymadng` + +When you create a {class}`MAD()` instance, PyMAD-NG exposes many MAD-NG global libraries and functions directly as attributes of the object. + +### Example: +```python +from pymadng import MAD +mad = MAD() +print(mad.math.sin(1).eval()) +``` + +These objects mirror MAD-NG’s Lua-style modules and can be accessed directly via `mad.`. + +--- + +## Global Utility Modules + +| Name | Description | +|------------|--------------------------------------------------| +| `math` | Basic math functions: `sin`, `cos`, `sqrt`, etc. | +| `os` | Operating system functions (limited) | +| `io` | Input/output and file manipulation | +| `table` | Table and list utilities | +| `string` | String manipulation | +| `MAD` | Core MAD-NG functions and module loader | +| `MADX` | Legacy-style MAD-X environment adapter | + +All returned values are *references* until explicitly evaluated with `.eval()`. + +--- + +## Global Functions from Lua +These functions are available globally in the MAD-NG environment: + +| Name | Description | +|-------------|--------------------------------------------------| +| `assert` | Assert a condition (throws error if false) | +| `print` | Print to the console | +| `error` | Throw an error | +| `ipairs` | Iterate over indexed part of a table | +| `pairs` | Iterate over all parts of a table | +| `tonumber` | Convert a string to a number | +| `tostring` | Convert a number to a string | +| `type` | Get the type of a variable | +| `load` | Load a Lua chunk (string) | +| `loadfile` | Load a Lua file | +| `require` | Load a Lua module | +--- + +## MAD Physics Libraries + +These libraries expose MAD-NG’s physics simulation functionality, and are made available automatically by PyMAD-NG when you initialise {class}`MAD`. + +| Name | Purpose | +|-------------|---------------------------------------------------------| +| `beam` | Define beam particle and energy | +| `beta0` | Create Beta0 optics state | +| `element` | Create MAD-NG elements (e.g. quadrupoles, dipoles) | +| `sequence` | Build beamline sequences | +| `object` | General MAD-NG object interface | +| `track` | Perform particle tracking | +| `match` | Launch matching optimizations | +| `twiss` | Compute Twiss parameters | +| `survey` | Calculate survey tables and coordinates | +| `mtable` | Handle MAD-NG table output | + +These are accessible as attributes of the {class}`MAD` instance: + +```python +mad["seq"] = mad.sequence(...) # Create a sequence +mad.seq.beam = mad.beam(...) # Attach a beam to the sequence +mad["tbl", "flw"] = mad.twiss(sequence=mad.seq) # Run a Twiss +``` +See [Setting Variables in MAD-NG](#setting-variables-in-mad-ng) for more details on how to set variables in MAD-NG and explanation of the code above. + +--- + +## Importing MAD-NG Modules via PyMAD-NG + +You can load MAD-NG modules dynamically using the {func}`MAD.load` function. +This is useful for loading built-in modules that are not automatically available. + +### Example: + +You can load modules directly from MAD-NG namespaces: + +```python +mad.load("MAD", "gphys") # Load a module such as MAD.gphys +mad.load("element", "quadrupole") # Load a specific element type +``` + +The first argument is the module path. Optional additional arguments allow importing specific submodules or functions. + +--- + + +## Core Communication Methods + +| Method | Description | +|----------------------|--------------------------------------------------| +| {func}`MAD.send` | Send Python data or MAD-NG code to MAD-NG | +| {func}`MAD.recv` | Receive results or values from MAD-NG | +| {func}`MAD.eval` | Evaluate an expression and return the result | +| {func}`MAD.recv_and_exec`| Execute Python code sent from MAD-NG | + +These tools are essential for controlling the MAD subprocess directly. + +--- + +## Reference Objects and Evaluation + +(setting-variables-in-mad-ng)= +### Setting Variables in MAD-NG + +To assign values in MAD-NG, always use square bracket syntax on the {class}`MAD` object. Dot syntax does not trigger assignment inside MAD-NG. + +```python +mad["energy"] = 6500 # Sets a global MAD-NG variable +mad["tw"] = mad.twiss(...) # Stores the result in MAD-NG +``` + +Using `mad.energy = 6500` only affects Python, not MAD-NG, and will not produce the expected behavior. + + +Functions like `mad.math.sin()` return reference objects that defer computation. +You must call `.eval()` to obtain the actual value. + +```python +r = mad.math.sin(1) +print(r.eval()) # returns float result +``` + +References can be composed symbolically: +```python +expr = mad.math.sin(1) + mad.math.cos(1) +print(expr.eval()) +``` + +--- + +## Table Conversion + +MAD tables returned from `twiss()`, `survey()`, etc., can be converted into Pandas-style data frames: + +```python +df = mad.tbl.to_df() +``` + +- If `tfs-pandas` is installed, a `TfsDataFrame` is returned (with headers). +- Otherwise, a regular `pandas.DataFrame` is returned, with the headers in the attrs. + +Raises `TypeError` if called on a non-table object. + +--- + + +## Listing Available Globals + +PyMAD-NG provides functions to explore the MAD-NG runtime environment: + +### `mad.globals()` +Returns all top-level variable names currently defined in the MAD-NG session. + +### `dir(mad)` +Returns cached or known Python-accessible attributes. + +--- + +## Quoting Strings + +### `mad.quote_strings(list_of_str)` +Converts a Python list of strings to a MAD-NG-compatible quoted list. + +#### Example: +```python +columns = ["s", "beta11", "mu1"] +mad.tbl.write("'output.tfs'", mad.quote_strings(columns)) +``` + +--- + +## Summary + +PyMAD-NG makes many MAD-NG modules and tools accessible from Python in an intuitive way. This reference highlights the most useful global objects, simulation modules, utility functions, and internal submodules available for scripting or extension work. + +For additional control or automation, refer to the API reference or the examples directory. \ No newline at end of file diff --git a/docs/source/gettingstarted.rst b/docs/source/gettingstarted.rst deleted file mode 100644 index b624f3f..0000000 --- a/docs/source/gettingstarted.rst +++ /dev/null @@ -1,143 +0,0 @@ -Getting Started -=============== - -Installation ------------- - -.. code-block:: - - pip install pymadng - -Each example can run by executing ``python3 ex-name.py``, using a console with the same working directory as the example. These examples give a range of available uses and show some techniques that are available with MAD-NG and pymadng. You can also run all the examples with ``runall.py`` - -To begin communication with MAD-NG, you simply are required to do: - -.. code-block:: - - from pymadng import MAD - mad = MAD() - -Communication protocol ----------------------- - -The communication protocol has been presented previously in a meeting that can be found here: `MAD-NG Python interface `_. The essential points are: - -- PyMAD-NG communicates through pipes (first in, first out) -- Communication occurs by sending MAD-NG scripts (as strings) to MAD -- Retrieve data from MAD to Python pipe. -- The stdout of MAD is redirected to the stdout of Python (not intercepted by PyMAD-NG) - -The first point is the most consequential for the user, as it means that the order in which you send data to MAD-NG is the order in which it will be received and vice versa for retrieving data. Therefore, you must adhere to the following rules: - -.. important:: - - **Before you receive any data from MAD-NG, you must always ask MAD-NG to send the data.** - - **Before you send data to MAD-NG, you must always send MAD-NG the instructions to read the data.** - -.. code-block:: - :caption: An example of using the MAD object to communicate with MAD-NG - - #Load MAD from pymadng - from pymadng import MAD - mad = MAD() - - #Tell mad that it should expect data and then place it in 'a' - mad.send("a = py:recv()") - - #Send the data - mad.send(42) - - #Ask mad to send the data back - mad.send("py:send(a)") - - #Read the data - mad.recv() #-> 42 - - -:meth:`mad.send() ` and :meth:`mad.recv() ` are the main ways to communicate with MAD-NG and is extremely simple, for specific details on what data can be sent see the :class:`API Reference `. - -For types that are not naturally found in numpy or python, you will be required to use a different function to *send* data (see below). The functions used in these specific cases can be found in the :mod:`MAD ` documentation. To *receive* any data just use :meth:`mad.recv() `. - -.. _typestbl: - -.. table:: Types that can be sent to MAD-NG and the function to use to send them - - +----------------------------------------+------------------------+----------------------------------------------+ - | Type in Python | Type in MAD | Function to send from Python | - +========================================+========================+==============================================+ - | None | nil | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | str | string | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | int | number :math:`<2^{31}` | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | float | number | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | complex | complex | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | list | table | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | bool | bool | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | NumPy ndarray (dtype = np.float64) | matrix | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | NumPy ndarray (dtype = np.complex128) | cmatrix | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | NumPy ndarray (dtype = np.int32) | imatrix | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | range | irange | :meth:`send ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | start(float), stop(float), size(int) | range | :meth:`send_rng ` | - +----------------------------------------+------------------------+----------------------------------------------+ - | start(float), stop(float), size(int) | logrange | :meth:`send_lrng ` | - +----------------------------------------+------------------------+----------------------------------------------+ - || NumPy ndarray (dtype = np.uint8) and || TPSA || :meth:`send_tpsa ` | - || NumPy ndarray (dtype = np.float64) || || | - +----------------------------------------+------------------------+----------------------------------------------+ - || NumPy ndarray (dtype = np.uint8) and || CTPSA || :meth:`send_ctpsa ` | - || NumPy ndarray (dtype = np.complex128) || || | - +----------------------------------------+------------------------+----------------------------------------------+ - -Recommended reading -------------------- - -First, we recommend familiarising yourself with MAD-NG, documentation can be found `here `_. - -Then reading through :doc:`ex-lowlevel` should be sufficient (alongside knowledge of MAD-NG), assuming you are not planning to use any "syntactic sugar". If you plan to use the available pythonic looking code, there are plenty of examples to look at. - -In the documentation, :doc:`ex-fodo` is a chapter that goes into detail on what is happening on each line of the :ref:`FODO example `, while :doc:`ex-lhc-couplingLocal` gives an example of loading the LHC and how to grab intermediate results from a match. - -The only other example that may be of use is the :ref:`ps-twiss ` example. This is an extremely simple example, extending the FODO example to perform a twiss on the PS sequence. -If anything does not seem fully explained, initially check the :mod:`MAD ` module and/or the `MAD-NG Documentation `_, then feel free to open an `issue `_ so improvements can be made. - -Customising your environment ----------------------------- - -Few things can be changed about the setup of your communication with MAD-NG, below lists a couple of use cases that may be of use. See also :meth:`__init__`. - -To change how you refer to your python prcess from within MAD-NG, by default, we use ``py`` (which may conflict with some variables you intend to define): - -.. code-block:: - - from pymadng import MAD - mad = MAD(py_name = "python") - -To change the MAD-NG executable used when pymadng is run: - -.. code-block:: - - from pymadng import MAD - mad = MAD(mad_path = r"/path/to/mad") - -To enable debugging mode: - -.. code-block:: - - from pymadng import MAD - mad = MAD(debug = True) - -To increase the number of temporary variables available to you (see :doc:`ex-managing-refs` for more information): - -.. code-block:: - - from pymadng import MAD - mad = MAD(num_temp_vars = 10) diff --git a/docs/source/highlevel.rst b/docs/source/highlevel.rst deleted file mode 100644 index 178b4ce..0000000 --- a/docs/source/highlevel.rst +++ /dev/null @@ -1,76 +0,0 @@ -High-Level PyMAD-NG -=================== - -This page describes how the pythonic interface to MAD-NG works. This way of using PyMAD-NG is far more limited than the low-level interface, and everything that can be done with the high-level interface can also be done with the low-level interface, usually with more flexibility and better performance. However, the high-level interface is much easier to use for people used to python. - -The "high-level" describes essentially every function, method, class, variable, etc. accessible from the top-level namespace of the `pymadng` module not listed in the :class:`API Reference `. - -When a user initialises a :class:`MAD ` object, as follows, we say that the user has created a "MAD object", which creates a new instance of the MAD-NG program, and from this object we can interact with the MAD-NG program. - -.. code-block:: python - - import pymadng - mad = pymadng.MAD() - -The MAD object is used to access the global environment of the MAD instance, there does not exist a significant amount within the global environment, the full list of variables can be found on the `MAD-NG GitHub `_, but a few of the more important ones are listed below, for a better understanding and examples see the `MAD-NG manual `_. - -Global Libraries ----------------- - -- :class:`io` - A lua library for reading and writing files. -- :class:`os` - A lua library for interacting with the operating system. -- :class:`math` - A lua library for mathematical functions. -- :class:`table` - A lua library for manipulating tables. -- :class:`string` - A lua library for manipulating strings. -- :class:`MAD` - The main library for accessing all the modules and functions of MAD-NG. -- :class:`MADX` - The MAD-X environment in MAD-NG - explicitly not MAD-X and so cannot be used to run MAD-X code, used for importing MAD-X sequences and adjusting variables and knobs. - -Global Functions ----------------- - -- :func:`assert` - A lua function for asserting a condition. -- :func:`print` - A lua function for printing to the console. -- :func:`error` - A lua function for throwing an error. -- :func:`ipairs` - A lua function for iterating over the indexed part of a table (the array part). -- :func:`pairs` - A lua function for iterating over the whole table, including indices and keys (the array and hash parts). -- :func:`tonumber` - A lua function for converting a string to a number. -- :func:`tostring` - A lua function for converting a number to a string. -- :func:`type` - A lua function for getting the type of a variable. - -Other Global Variables - Specific to PyMAD-NG ---------------------------------------------- - -The functions that are required for physics are all contained within the :class:`MAD` library, a few of the more important ones are immediately exposed to the user, only in PyMAD-NG, and are listed below: - -- :class:`beam` - See `Beams `_ -- :class:`beta0` - See `Beta0 `_ -- :class:`element` - See `Elements `_ -- :class:`match` - See `Match `_ -- :class:`mtable` - See `MTables `_ -- :class:`object` - See `Objects `_ -- :class:`sequence` - See `Sequences `_ -- :class:`survey` - See `Survey `_ -- :class:`track` - See `Track `_ -- :class:`twiss` - See `Twiss `_ - -Accessing the Global Environment --------------------------------- - -To access anything in the global environment, we can use the :class:`MAD ` object, for example, to access the :class:`math` library, we can use the following: - -.. code-block:: python - - import pymadng - mad = pymadng.MAD() - print(mad.math.sin(1).eval(), mad.math.cos(1).eval()) # 0.8414709848078965 0.5403023058681398 - -The reason we have to use the :func:`eval` function is because when a function is called from python and executed in MAD-NG, it will always return a reference, allowing us to continue to use the returned value in MAD-NG and it is not automatically converted to a python object. See :doc:`/ex-managing-refs` for more information on these references and how to use/deal with them. This evaulation can also be done using the low level interface, as follows: - -.. code-block:: python - - import pymadng - mad = pymadng.MAD() - mad.send("py:send(math.sin(1)):send(math.cos(1))") - print(mad.recv(), mad.recv()) # 0.8414709848078965 0.5403023058681398 - -To conclude, the global environment is accessed through the :class:`MAD ` object, several of libraries and functions are listed above, most of the other necessary functions required can be accessed through the :class:`MAD` library in MAD-NG, by doing ``mad.MAD``. With knowledge of MAD-NG, the user can access and interact with almost everything in MAD-NG through the high-level interface, however, the low-level interface is much more flexible and powerful, so it is recommended to use for more complex tasks. \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst index b343163..5e90288 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -3,23 +3,49 @@ You can adapt this file completely to your liking, but it should at least contain the root `toctree` directive. -Welcome to the documentation for PyMAD-NG! -========================================== +.. PyMAD-NG documentation master file + +Welcome to the PyMAD-NG documentation! +====================================== + +.. toctree:: + :maxdepth: 2 + :caption: Getting Started + + intro + installation + quickstartguide + +.. toctree:: + :maxdepth: 2 + :caption: Core Concepts + + communication + function_reference + architecture .. toctree:: :maxdepth: 2 - :caption: Contents: + :caption: Advanced Topics + + advanced_features + debugging + +.. toctree:: + :maxdepth: 2 + :caption: Examples - gettingstarted - ex-lowlevel - highlevel - ex-managing-refs - ex-fodo ex-lhc-couplingLocal - dataframes - modules examples +.. toctree:: + :maxdepth: 2 + :caption: Reference + + reference + developer + contributing + Indices and tables ================== diff --git a/docs/source/installation.md b/docs/source/installation.md new file mode 100644 index 0000000..6e3b7ef --- /dev/null +++ b/docs/source/installation.md @@ -0,0 +1,177 @@ +# Installation & Setup + +```{contents} +:depth: 2 +:local: +``` + +## System Requirements +Before installing PyMAD-NG, ensure that your system meets the following requirements: + +- **Python**: Version 3.7 or later +- **Operating System**: Linux, macOS, or Windows +- **Dependencies**: + - `numpy` + - `pandas` (optional, for data handling) + - `matplotlib` (optional, for visualization) + +--- + +## Installing PyMAD-NG + +PyMAD-NG is available on PyPI and can be installed using `pip`: + +```bash +pip install pymadng +``` + +If you need to install optional dependencies for advanced features: + +```bash +pip install pymadng[tfs] +``` + +This installs `tfs-pandas` for handling MAD-NG TFS tables. + +### Verifying Installation + +To check if PyMAD-NG is installed correctly, open Python and run: + +```python +import pymadng +print(pymadng.__version__) +``` + +If this prints a version number, the installation was successful. + +--- + +## Updating PyMAD-NG + +To update PyMAD-NG to the latest version: + +```bash +pip install --upgrade pymadng +``` + +You can check your current version using: + +```bash +pip show pymadng +``` + +--- + +## Uninstalling PyMAD-NG + +If you need to remove PyMAD-NG from your system: + +```bash +pip uninstall pymadng +``` + +--- + +## Using a Custom MAD-NG Executable + +If you have a specific MAD-NG binary, you can specify its path: + +```python +from pymadng import MAD +mad = MAD(mad_path="/path/to/mad") +``` +## Understanding {class}`pymadng.MAD` Object Initialization + +When you create a {class}`pymadng.MAD` object, an instance of MAD-NG is launched. The `__init__` method in `mad_object.py` handles the following: + +- **Process Creation**: Starts MAD-NG and establishes a communication channel. +- **Module Imports**: By default, imports essential MAD-NG modules, see [Useful Modules](function_reference.md) for the complete list. +- **Environment Configuration**: Sets up the Python-MAD environment, including handling `py_name` and error handling. +- **Temporary Variable Management**: Uses `last_counter` to track temporary variables (`_last[]`) for improved performance. + +This ensures that users can quickly interface with MAD-NG without extensive manual setup. + +--- + +## Debugging Tips + +Debugging PyMAD-NG can involve both Python-side and MAD-NG-side issues. Below are some strategies and tools built into PyMAD-NG to help troubleshoot effectively: + +### 1. Enable Debug Mode + +You can enable verbose output during initialization by setting `debug=True`. This enables MAD-NG's debug mode and prints useful messages: + +```python +mad = MAD(debug=True) +``` + +### 2. Redirect Standard Output and Standard Error + +You can log all MAD-NG output to a file: + +```python +mad = MAD(stdout="mad_debug.log") +``` + +This is useful for reviewing MAD-NG output, especially when running large scripts. + +If you want to redirect the standard error to a different file (as by default it is redirected to the standard output file), you can use the `stderr` parameter: + +```python +mad = MAD(stdout="mad_debug.log", stderr="mad_error.log") +``` + +### 3. Catch Errors from MAD-NG + +Internally, MAD-NG returns an `err_` token to signal an error. PyMAD-NG raises a `RuntimeError` with a message when this happens. If you encounter unexpected crashes, wrap your calls in try/except: + +```python +try: + mad.send("a = 1/'0'").recv() +except RuntimeError as e: + print("MAD-NG Error:", e) +``` + +Or it is possible to ignore the error and continue execution of the python script: + +```python +mad = MAD(raise_on_error=False) +mad.send("a = 1/'0'") +assert mad.send("py:send(1)").recv() == 1 +``` + +### 4. View Communication History + +You can inspect the full history of MAD-NG commands sent from Python, *only if **debug mode** is enabled*: + +```python +print(mad.history()) +``` + +### 5. Interactive Execution + +Use `mad.recv_and_exec()` to execute MAD-NG commands returned as Python strings. This is helpful when MAD sends back executable code: + +```python +mad.send("py:send('print(\'debug line\')')") +mad.recv_and_exec() +``` + +### 6. Ensure Communication Order + +MAD-NG uses FIFO communication. Always follow the rule: + +- **Before calling `recv()`**, ensure MAD was told to `send()` something. +- **Before calling `send(data)`**, ensure MAD expects to `recv()` something. + +Failing to follow this can cause deadlocks or hangs. + +--- + +## Next Steps + +Now that PyMAD-NG is installed, you can move on to: + +- **[Quick Start Guide →](quickstartguide.md)** *(Learn the basics and run your first PyMAD-NG script)* +- **[API Reference →](reference.rst)** *(Explore the functions and classes available in PyMAD-NG)* + diff --git a/docs/source/intro.rst b/docs/source/intro.rst new file mode 100644 index 0000000..38f1f5a --- /dev/null +++ b/docs/source/intro.rst @@ -0,0 +1,90 @@ +.. _introduction: + +======================== +Introduction to PyMAD-NG +======================== + +What is PyMAD-NG? +----------------- +PyMAD-NG is a **Python interface** for **MAD-NG** (Methodical Accelerator Design - Next Generation), a powerful software for simulating and analysing particle accelerators. PyMAD-NG enables seamless communication between Python and MAD-NG, allowing users to **script, automate, and interactively control** MAD-NG simulations from Python. + +PyMAD-NG provides a **pythonic API** that simplifies interaction with MAD-NG while maintaining high performance. Whether you're performing **optics calculations, beam dynamics simulations, or machine tuning**, PyMAD-NG offers the flexibility to work efficiently with MAD-NG from within Python. + +Why Use PyMAD-NG? +----------------- +PyMAD-NG is designed for **scientists, engineers, and researchers** working on accelerator physics and beam dynamics. It offers several advantages over traditional MAD-NG workflows: + +- **Pythonic Interface** - Write MAD-NG scripts using intuitive Python commands. +- **Efficient Communication** - Uses **pipes** for fast data exchange between Python and MAD-NG. +- **Seamless Data Handling** - Convert MAD-NG tables into **Pandas DataFrames** for analysis. +- **High Performance** - Designed for handling **large datasets** and computationally intensive simulations. +- **Flexible APIs** - Use either a **high-level API** (more Pythonic) or a **low-level API** (more control). +- **Jupyter Notebook Support** - Work interactively with MAD-NG in a Python notebook. +- **MAD-X Compatibility** - Load MAD-X sequences and interact with them in MAD-NG. + +How PyMAD-NG Works +------------------ +PyMAD-NG operates by **launching a MAD-NG process** in the background and establishing a **two-way communication channel** between Python and MAD-NG. + +- **Sending Commands** - You can send MAD-NG commands as **Python strings**. +- **Receiving Results** - Data from MAD-NG can be retrieved into Python for further analysis. +- **MAD Objects in Python** - PyMAD-NG exposes MAD-NG objects as Python objects for easy manipulation. + +Example Workflow +---------------- + +1. **Initialise PyMAD-NG**:: + + from pymadng import MAD + mad = MAD() + +2. **Load a Sequence & Perform Calculations**:: + + mad.MADX.load("'lhc_as-built.seq'", "'lhc_as-built.mad'") + mad["tbl", "flw"] = mad.twiss(sequence=mad.MADX.lhcb1, method=4) + +3. **Retrieve Data from MAD-NG**:: + + df = mad.tbl.to_df() # Convert twiss table to a Pandas DataFrame + print(df.head()) + +4. **Visualise Results in Python**:: + + import matplotlib.pyplot as plt + plt.plot(df["s"], df["beta11"]) + plt.xlabel("s (m)") + plt.ylabel("$\beta_x$-function") + plt.show() + +Key Features of PyMAD-NG +------------------------- + ++--------------------------------+----------------------------------------------------------+ +| Feature | Description | ++================================+==========================================================+ +| **Pythonic Interface** | Interact with MAD-NG using Python objects. | ++--------------------------------+----------------------------------------------------------+ +| **High-Level & Low-Level API** | Choose between a simple or customisable approach. | ++--------------------------------+----------------------------------------------------------+ +| **Efficient Data Handling** | Convert MAD-NG tables (`mtable`) into Pandas DataFrames. | ++--------------------------------+----------------------------------------------------------+ +| **Two-Way Communication** | Send commands to MAD-NG and retrieve results. | ++--------------------------------+----------------------------------------------------------+ +| **MAD-X Compatibility** | Import and work with MAD-X sequences. | ++--------------------------------+----------------------------------------------------------+ +| **Performance Optimised** | Supports large datasets and numerical computations. | ++--------------------------------+----------------------------------------------------------+ +| **Jupyter Notebook Support** | Use PyMAD-NG interactively within Jupyter. | ++--------------------------------+----------------------------------------------------------+ + +Who Should Use PyMAD-NG? +------------------------- +If you work with MAD-NG and want to **leverage Python's ecosystem (NumPy, Pandas, Matplotlib, etc.),** PyMAD-NG is the perfect tool for you. + +Next Steps +---------- +Now that you have an overview of PyMAD-NG, you can dive into: + +- :doc:`Installation ` - Set up PyMAD-NG on your system. +- :doc:`Quick Start Guide` - Run your first PyMAD-NG script in minutes. +- :doc:`API Reference` - Explore the available functions and classes. diff --git a/docs/source/quickstartguide.md b/docs/source/quickstartguide.md new file mode 100644 index 0000000..34b0e02 --- /dev/null +++ b/docs/source/quickstartguide.md @@ -0,0 +1,122 @@ +# Quick Start Guide + +This guide will walk you through your first experience with **PyMAD-NG**. By the end, you'll know how to run MAD-NG commands from Python, send and receive data, and perform a basic optics calculation. + +```{contents} +:depth: 2 +:local: +``` + +--- + +## Step 1: Create a MAD Instance + +Start by importing and creating an instance of the {class}`pymadng.MAD` object: + +```python +from pymadng import MAD +mad = MAD() +``` + +This automatically launches a MAD-NG process and connects it to Python. + +--- + +## Step 2: Load a Sequence + +PyMAD-NG supports two approaches: a high-level (pythonic) interface and a low-level (script-driven) interface. + +### High-Level API: + +```python +mad.MADX.load("'fodo.seq'", "'fodo.mad'") +mad["seq"] = mad.MADX.seq +``` + +### Low-Level API: + +```python +mad.send("MADX:load('fodo.seq', 'fodo.mad')") +mad.send("seq = MADX.seq") +``` + +--- + +## Step 3: Set Up a Beam + +### High-Level: + +```python +mad.seq.beam = mad.beam() +``` + +### Low-Level: + +```python +mad.send("seq.beam = beam") +``` + +--- + +## Step 4: Run a Twiss Calculation + +### High-Level: + +```python +mad["tbl", "flow"] = mad.twiss(sequence=mad.seq, method=4) +``` + +### Low-Level: + +```python +mad.send("tbl, flow = twiss {sequence=seq, method=4}") +mad.send("py:send(tbl)") +tbl = mad.recv() +``` + +--- + +## Step 5: Analyse the Results + +Convert the resulting MAD table to a Pandas DataFrame: + +```python +df = mad.tbl.to_df() +print(df.head()) +``` + +--- + +## Step 6: Visualise the Optics + +Plot the beta function using Matplotlib: + +```python +import matplotlib.pyplot as plt +plt.plot(df["s"], df["beta11"]) +plt.xlabel("s [m]") +plt.ylabel("Beta Function") +plt.title("Twiss Beta11 vs s") +plt.grid(True) +plt.show() +``` + +--- + +## Additional Tips + +- If you want to send MAD-NG commands directly: + ```python + mad.send("py:send(math.sin(1))") + print(mad.recv()) + ``` + +- Always match `send()` and `recv()` properly to avoid blocking communication. + +## What Next? + +Now that you’ve completed your first PyMAD-NG workflow, explore: +- **[MAD-NG Documentation](https://madx.web.cern.ch/releases/madng/html/)** for details on MAD-NG features +- **[API Reference](reference.rst)** for full documentation of the {class}`pymadng.MAD` class +- **[Examples](examples.rst)** to see real-world scripts. With all the scripts also in the [`examples` folder of the PyMAD-NG repository](https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/tree/main/examples). + diff --git a/docs/source/modules.rst b/docs/source/reference.rst similarity index 92% rename from docs/source/modules.rst rename to docs/source/reference.rst index 888bb80..12af164 100644 --- a/docs/source/modules.rst +++ b/docs/source/reference.rst @@ -13,4 +13,7 @@ Useful functions for MAD References ----------------------------------- .. autofunction:: pymadng.madp_classes.high_level_mad_object.convert_to_dataframe -.. autofunction:: pymadng.madp_classes.high_level_mad_object.to_df \ No newline at end of file + :no-index: + +.. autofunction:: pymadng.madp_classes.high_level_mad_object.to_df + :no-index: \ No newline at end of file diff --git a/examples/ex-fodo/ex-fodos.py b/examples/ex-fodo/ex-fodos.py index 8a018f3..547d7f7 100644 --- a/examples/ex-fodo/ex-fodos.py +++ b/examples/ex-fodo/ex-fodos.py @@ -4,61 +4,13 @@ orginal_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. with MAD() as mad: - mad.MADX.load(f"'fodo.seq'", f"'fodo.mad'") - mad["seq"] = mad.MADX.seq - mad.seq.beam = mad.beam() - mad["mtbl", "mflw"] = mad.twiss(sequence=mad.seq, method=4, implicit=True, nslice=10, save="'atbody'") - plt.plot(mad.mtbl.s, mad.mtbl.beta11) - plt.show() - -with MAD() as mad: - mad.MADX.load(f"'fodo.seq'", f"'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.write("'twiss_py.tfs'", cols) - for x in mad.seq: - print(x.name, x.kind) - plt.plot(mad.mtbl.s, mad.mtbl.beta11) - plt.show() - - -with MAD() as mad: - mad.load("element", "quadrupole") - mad["circum", "lcell"] = 60, 20 - - mad.load("math", "sin", "pi") - mad["v"] = mad.create_deferred_expression(k="1/(lcell/sin(pi/4)/4)") - - mad["qf"] = mad.quadrupole("knl:={0, v.k}", l=1) - mad["qd"] = mad.quadrupole("knl:={0, -v.k}", l=1) - mad["seq"] = mad.sequence(""" - qf { at = 0 }, - qd { at = 0.5 * lcell }, - qf { at = 1.0 * lcell }, - qd { at = 1.5 * lcell }, - qf { at = 2.0 * lcell }, - qd { at = 2.5 * lcell }, - """, refer="'entry'", l=mad.circum,) - 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.write("'twiss_py.tfs'", cols) - - plt.plot(mad.mtbl.s, mad.mtbl["beta11"]) - plt.title("FODO Cell") - plt.xlabel("s") - plt.ylabel("beta11") - plt.show() - -with MAD() as mad: - mad.send(f""" + mad.send(""" 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, method=4, implicit=true, nslice=10, save="atbody"} py:send(mtbl) """) mtbl = mad.recv("mtbl") @@ -72,4 +24,19 @@ 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, method=4, 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(orginal_dir) \ No newline at end of file diff --git a/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py b/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py index 8448729..939d41f 100644 --- a/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py +++ b/examples/ex-lhc-couplingLocal/lhc-couplingLocal.py @@ -1,9 +1,16 @@ -from pymadng import MAD -import time +""" +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 import matplotlib.pyplot as plt -import numpy as np + +from pymadng import MAD + orginal_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) @@ -13,13 +20,18 @@ 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.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.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 @@ -47,25 +59,27 @@ """) 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}, + 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}, + 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, + objective={"fmin": 1e-3}, + maxcall=100, + info=2, ) mad.send("py:send(nil)") - tws_result = mad.recv () + 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-') + (line1,) = ax.plot(x, y, "b-") while tws_result: line1.set_xdata(tws_result[0]) line1.set_ydata(tws_result[1]) @@ -85,4 +99,4 @@ t1 = time.time() print("pre-tracking time: " + str(t1 - t0) + "s") -os.chdir(orginal_dir) \ No newline at end of file +os.chdir(orginal_dir) diff --git a/examples/ex-managing-refs/ex-managing-refs.py b/examples/ex-managing-refs/ex-managing-refs.py index b52a9c0..85bf033 100644 --- a/examples/ex-managing-refs/ex-managing-refs.py +++ b/examples/ex-managing-refs/ex-managing-refs.py @@ -1,40 +1,42 @@ -#Code that does not necessarily work as expected +# Code that does not necessarily work as expected +import os + +import numpy as np + from pymadng import MAD -import os, numpy as np orginal_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) -mad = MAD() #Not being in context manager makes not difference. +mad = MAD() # Not being in context manager makes not difference. -#Just boring setup (in lots of other examples) -mad.MADX.load(f"'fodo.seq'", f"'fodo.mad'") +# 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) +# 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) -#Any high level MAD-NG function will create a reference +# Any high level MAD-NG function will create a reference mad.MAD.gmath.reim(1.42 + 0.62j) -#Try to receive from twiss +# Try to receive from twiss mad["mtbl", "mflw"] = twissrtrn -# mtbl and mflow correctly stored! +# mtbl and mflow correctly stored! print(mad.mtbl) print(mad.mflw) -myMatrix = mad.MAD.matrix(4).seq() #Create 4x4 matrix +myMatrix = mad.MAD.matrix(4).seq() # Create 4x4 matrix -print(type(myMatrix)) #Not a 4x4 matrix! -print(type(myMatrix.eval())) #A 4x4 matrix! +print(type(myMatrix)) # Not a 4x4 matrix! +print(type(myMatrix.eval())) # A 4x4 matrix! -myMatrix = myMatrix.eval() #Store the matrix permanantly +myMatrix = myMatrix.eval() # Store the matrix permanantly mad["myMatrix"] = mad.MAD.matrix(4).seq() print(mad.myMatrix, np.all(myMatrix == mad.myMatrix)) -os.chdir(orginal_dir) \ No newline at end of file +os.chdir(orginal_dir) diff --git a/pyproject.toml b/pyproject.toml index 8f3db1f..d66c1da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,8 +26,8 @@ dependencies = [ ] [project.urls] -"Repository" = "https://github.com/MethodicalAcceleratorDesign/MADpy" -"Bug Tracker" = "https://github.com/MethodicalAcceleratorDesign/MADpy/issues" +"Repository" = "https://github.com/MethodicalAcceleratorDesign/MAD-NG.py" +"Bug Tracker" = "https://github.com/MethodicalAcceleratorDesign/MAD-NG.py/issues" "MAD Source" = "https://github.com/MethodicalAcceleratorDesign/MAD" "Documentation" = "https://pymadng.readthedocs.io/en/latest/" diff --git a/src/pymadng/__init__.py b/src/pymadng/__init__.py index 1fa0d8c..feb70d9 100644 --- a/src/pymadng/__init__.py +++ b/src/pymadng/__init__.py @@ -4,7 +4,7 @@ __version__ = "0.6.0" __summary__ = "Python interface to MAD-NG running as subprocess" -__uri__ = "https://github.com/MethodicalAcceleratorDesign/MADpy" +__uri__ = "https://github.com/MethodicalAcceleratorDesign/MAD-NG.py" __credits__ = """ Creator: Joshua Gray diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 03fec1d..cd131bd 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -19,6 +19,13 @@ # MAD High Level reference class high_level_mad_ref(mad_ref): + """ + A high-level MAD reference. + + Wraps a MAD-NG reference so that attribute and item access are transparent. + Automatically handles indexing differences between Python and MAD-NG. + """ + def __init__(self, name: str, mad_proc: mad_process): super(high_level_mad_ref, self).__init__(name, mad_proc) self._parent = ( @@ -122,6 +129,13 @@ def __deepcopy__(self, memo): class high_level_mad_object(high_level_mad_ref): + """ + A high-level MAD object for complex data and table handling. + + Inherits from high_level_mad_ref and extends functionality for retrieving object keys, + iteration, and conversion to Pandas dataframes. + """ + def __dir__(self) -> Iterable[str]: if not self._mad.ipython_use_jedi: self._mad.protected_send( @@ -261,6 +275,13 @@ def convert_to_dataframe(self, columns: list = None): class high_level_mad_func(high_level_mad_ref): + """ + A high-level MAD function reference. + + Enables calling MAD-NG functions with support for method and function invocation. + Automatically passes instance or object context as needed. + """ + # ----------------------------------Calling/Creating functions--------------------------------------# def __call_func(self, funcName: str, *args): """Call the function funcName and store the result in ``_last``.""" @@ -303,6 +324,12 @@ def __dir__(self): # Separate class for _last objects for simplicity and fewer if statements class mad_last: # The init and del for a _last object + """ + Manage a temporary '_last' variable in MAD-NG. + + This class assigns a unique temporary name for storing return values from MAD-NG functions. + """ + def __init__(self, mad_proc: mad_process): self._mad = mad_proc self._last_counter = mad_proc.last_counter @@ -315,6 +342,12 @@ def __del__(self): class mad_high_level_last_ref(mad_last, high_level_mad_ref): + """ + A combined reference for '_last' objects. + + Inherits from both mad_last and high_level_mad_ref to provide callable behavior on the last result. + """ + def __call__(self, *args: Any, **kwargs: Any) -> Any: obj = self.eval() if isinstance(obj, (high_level_mad_object, high_level_mad_func)): @@ -327,6 +360,12 @@ def __dir__(self): class high_level_last_object(mad_last, high_level_mad_object): + """ + A high-level representation for '_last' objects that are of object type. + + Provides seamless integration of temporary MAD-NG objects with Python functionality. + """ + pass diff --git a/src/pymadng/madp_last.py b/src/pymadng/madp_last.py index b4df51f..0dfd487 100644 --- a/src/pymadng/madp_last.py +++ b/src/pymadng/madp_last.py @@ -1,20 +1,20 @@ class last_counter: - """A class to keep track of the special variable named '__last__' in MAD-NG + """Maintain a counter for anonymous '_last' variables in MAD-NG. - The '__last__' variable contains the last value returned by a function in MAD-NG. - This allows the user to use pythonic syntax and get a return value from a function in MAD-NG, as an object in Python. + The __last__ variable stores the most recent function result in MAD-NG. + This helper class tracks available temporary variable indices. Args: - size (int): The maximum number of '__last__' variables to keep track of in the MAD-NG environment + size (int): The maximum number of temporary '_last' variables available. """ def __init__(self, size: int): self.counter = list(range(size, 0, -1)) def get(self): - assert ( - len(self.counter) > 0 - ), "Assigned too many anonymous variables, increase num_temp_vars or assign the variables into MAD" + assert len(self.counter) > 0, ( + "Assigned too many anonymous variables, increase num_temp_vars or assign the variables into MAD" + ) return self.counter.pop() def set(self, idx): diff --git a/src/pymadng/madp_object.py b/src/pymadng/madp_object.py index 45feb7a..55937d7 100644 --- a/src/pymadng/madp_object.py +++ b/src/pymadng/madp_object.py @@ -72,25 +72,22 @@ def __init__( num_temp_vars: int = 8, ipython_use_jedi: bool = False, ): - """Create a MAD Object to interface with MAD-NG. + """ + Initialise a MAD object for communication with MAD-NG. - The modules MADX, elements, sequence, mtable, twiss, beta0, beam, survey, object, track, match are imported into - the MAD-NG environment by default. + This constructor starts the MAD subprocess, establishes communication pipes, + and imports necessary MAD modules. The mad_path defaults to a bundled executable if not provided. Args: - mad_path (str): The path to the mad executable, for the default value of None, the one that comes with pymadng package will be used - py_name (str): The name used to interact with the python process from MAD, default is 'py'. - raise_on_madng_error (bool): Will _always_ raise an error if MAD-NG errors, default is True. - debug (bool): Sets debug mode on or off. This will output additional information to the stdout. (default is False) - stdout (TextIO | str | Path): Redirect the MAD-NG stdout to a file, useful for debugging. If a TextIO is given, it will write to this file. If a string or Path is given, it is expected to be a file path, and will open this file to write to. (default is None) - redirect_sterr (bool): Redirect the stderr to the stdout, useful for debugging. (default is False) - num_temp_vars (int): The number of unique temporary variables you intend to use, see :doc:`Managing References ` for more information. (default is 8) - ipython_use_jedi (bool): Allow ipython to use jedi in tab completion, will be slower and may result in MAD-NG throwing errors. (default is False) - - Returns: - A MAD object, allowing for communication with MAD-NG + mad_path (str | Path, optional): Path to the MAD executable. + py_name (str, optional): Name used for MAD-to-Python communication. + 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. + 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. """ - # ------------------------- Create the process --------------------------- # mad_path = mad_path or bin_path / ("mad_" + platform.system()) self.__process = mad_process( @@ -152,124 +149,131 @@ def __init__( def recv( self, varname: str = None ) -> str | int | float | np.ndarray | bool | list | high_level_mad_ref: - """Return received data from MAD-NG. + """ + Retrieve data from the MAD-NG process. + + Reads a data type identifier and then receives the corresponding value from MAD-NG. Args: - varname(str): The name of the variable you are receiving (Only useful when receiving references) + varname (str, optional): Variable name used for clarity when receiving references. Returns: - Data from MAD-NG with type str/int/float/ndarray/bool/list/ref, depending what was asked from MAD-NG. - - Raises: - TypeError: If you forget to give a name when receiving a reference, + The value received from MAD-NG. """ return self.__process.recv(varname) def receive( self, varname: str = None ) -> str | int | float | np.ndarray | bool | list | high_level_mad_ref: - """See :meth:`recv`""" + """ + Alias for the recv method. + + Args: + varname (str, optional): Name of the variable to receive. + + Returns: + The received data. + """ return self.__process.recv(varname) def recv_and_exec(self, context: dict = {}) -> dict: - """Receive a string from MAD-NG and execute it. + """ + Receive a string from MAD-NG and execute it. - Note: The class numpy and the instance of this object are available during the execution as ``np`` and ``mad`` respectively + The provided execution context is updated with numpy (np) and the current MAD object. Args: - context (dict): The environment context you would like the string to be executed in. + context (dict, optional): The environment for executing the received code. Returns: - The updated environment after executing the string. + dict: The updated execution environment. """ context["mad"] = self return self.__process.recv_and_exec(context) # --------------------------------Sending data to subprocess------------------------------------# def send(self, data: str | int | float | np.ndarray | bool | list) -> MAD: - """Send data to MAD-NG. + """ + Send data to MAD-NG. + + Accepts various types of data and serialises them for transfer to MAD-NG. Args: - data (str/int/float/ndarray/bool/list): The data to send to MAD-NG. + data (str/int/float/np.ndarray/bool/list): The information to send. Returns: - self (the instance of the mad object) - - Raises: - TypeError: An unsupported type was attempted to be sent to MAD-NG. - AssertionError: If data is a np.ndarray, the matrix must be of two dimensions. + MAD: Returns self to facilitate method chaining. """ self.__process.send(data) return self def protected_send(self, string: str) -> MAD: - """Send a string to MAD-NG, but if any of the command errors, python will be notified. - Then, once python receives the fact that the command errored, it will raise an error. + """ + Send a command to MAD-NG with error protection. - For example, if you send ``mad.send("a = 1/'a'")``, MAD-NG will error, and python will just continue. - But if you send ``mad.protected_send("a = 1/'a'").recv()``, python will raise an error. - Note, if the ``recv`` is not called, the error will not be raised. + Temporarily enables error handling so that any command errors in MAD-NG are caught and raised in Python. Args: - string (str): The string to send to MAD-NG. This must be a string to be evaluated in MAD-NG. Not a string to be sent to MAD-NG. + string (str): The MAD-NG command code to send. Returns: - self (the instance of the mad object) + MAD: Self for method chaining. """ assert isinstance(string, str), "The input to protected_send must be a string" self.__process.protected_send(string) return self def psend(self, string: str) -> MAD: - """See :meth:`protected_send`""" + """Alias for protected_send""" return self.protected_send(string) def send_range(self, start: float, stop: float, size: int): - """Send a range to MAD-NG, equivalent to np.linspace, but in MAD-NG. + """ + Send a linear range to MAD-NG. + + Builds a range analogous to numpy.linspace and sends it based on the specified boundaries and size. Args: - start (float): The start of range - stop (float): The end of range (inclusive) - size (float): The length of range + start (float): The starting value of the range. + stop (float): The ending value of the range (inclusive). + size (int): The total number of elements in the range. """ self.__process.send_range(start, stop, size) def send_logrange(self, start: float, stop: float, size: int): - """Send a numpy array as a logrange to MAD-NG, equivalent to np.geomspace, but in MAD-NG. + """ + Send a logarithmic range to MAD-NG. + + Generates and sends a logarithmic range equivalent to numpy.geomspace. Args: - start (float): The start of range - stop (float): The end of range (inclusive) - size (float): The length of range + start (float): The start value. + stop (float): The stop value. + size (int): The number of points in the range. """ self.__process.send_logrange(start, stop, size) def send_tpsa(self, monos: np.ndarray, coefficients: np.ndarray): - """Send the monomials and coefficients of a TPSA to MAD + """ + Transmit TPSA data for processing in MAD-NG. - The combination of monomials and coefficients creates a table representing the TPSA object in MAD-NG. + Sends a two-dimensional array of monomials and the associated coefficient array. Args: - monos (ndarray): A list of monomials in the TPSA. - coefficients (ndarray): A list of coefficients in the TPSA. - - Raises: - AssertionError: The list of monomials must be a 2-D array (each monomial is 1-D). - AssertionError: The number of monomials and coefficients must be identical. - AssertionError: The monomials must be of type 8-bit unsigned integer + monos (np.ndarray): 2D array containing the monomials. + coefficients (np.ndarray): Array of TPSA coefficients. """ self.__process.send_tpsa(monos, coefficients) def send_cpx_tpsa(self, monos: np.ndarray, coefficients: np.ndarray): - """Send the monomials and coefficients of a complex TPSA to MAD-NG + """ + Transmit complex TPSA data to MAD-NG. - The combination of monomials and coefficients creates a table representing the complex TPSA object in MAD-NG. + Sends complex-valued monomials and coefficients for TPSA table construction. Args: - See: :meth:`send_tpsa`. - - Raises: - See: :meth:`send_tpsa`. + monos (np.ndarray): 2D array of monomials (unsigned bytes). + coefficients (np.ndarray): Array of complex TPSA coefficients. """ self.__process.send_cpx_tpsa(monos, coefficients) @@ -277,45 +281,39 @@ def send_cpx_tpsa(self, monos: np.ndarray, coefficients: np.ndarray): # -------------------------------- Dealing with communication of variables --------------------------------# def send_vars(self, **vars: str | int | float | np.ndarray | bool | list): - """Send variables to the MAD-NG process. + """ + Send multiple named variables to the MAD-NG process. - Send the variables in vars with the names in names to MAD-NG. + Each keyword pair becomes a variable in the MAD-NG environment. Args: - **vars (str/int/float/ndarray/bool/list): The variables to send with the assigned name in MAD-NG by the keyword argument. - - Raises: - See: :meth:`send`. + **vars: Key-value pairs representing variable names and their values. """ self.__process.send_vars(**vars) def recv_vars(self, *names: str) -> Any: - """Receive variables from the MAD-NG process - - Given a list of variable names, receive the variables from the MAD-NG process. + """ + Retrieve one or more variables from MAD-NG. Args: - *names (str): The name(s) of variables that you would like to receive from MAD-NG. + *names (str): The names of the variables to be fetched from MAD-NG. Returns: - See :meth:`recv`. + The retrieved variable value or a tuple of values if multiple names are provided. """ return self.__process.recv_vars(*names) # -------------------------------------------------------------------------------------------------------------# def load(self, module: str, *vars: str): - """Import modules into the MAD-NG environment - - Retrieve the classes in MAD-NG from the module ``module``, while only importing the classes in the variable length list ``vars``. - If no string is provided, it is assumed that you would like to import every class from the module. + """ + Import classes or functions from a specific MAD-NG module. - For example, ``mad.load("MAD.gmath")`` imports all variables from the module ``MAD.gmath``. - But ``mad.load("MAD", "matrix", "cmatrix")```` only imports the modules ``matrix`` and ``cmatrix`` from the module ``MAD.gmath``. + If no specific names are provided, imports all available members from the module. Args: - module (str): The name of the module to import from. - *vars (str): Variable length argument list of the variable(s) to import from module. + module (str): The module name in MAD-NG. + *vars (str): Optional list of members to import. """ script = "" if vars == (): @@ -325,14 +323,14 @@ def load(self, module: str, *vars: str): self.__process.send(script) def loadfile(self, path: str | Path, *vars: str): - """Load a .mad file into the MAD-NG environment. + """ + Load and execute a .mad file in the MAD-NG environment. - If ``vars`` is not provided, this is equivalent to ``assert(loadfile(path))`` in MAD-NG. - If ``vars`` is provided, for each ``var`` in ``vars``, this is equivalent to ``var = require(path).var`` in MAD-NG. + If additional variable names are provided, assign each to the corresponding member of the loaded file. Args: - path (str | Path): The path to the file to import. - *vars (str): Variable length argument list of the variable(s) to import from the file. + path (str | Path): File path for the .mad file. + *vars (str): Optional names to bind to specific elements from the file. """ path: Path = Path(path).resolve() if vars == (): @@ -350,11 +348,30 @@ def loadfile(self, path: str | Path, *vars: str): # ----------------------- Make the class work with dict and dot access ------------------------# def __getattr__(self, item): + """ + Retrieve a MAD-NG variable using attribute access. + + Args: + item (str): The name of the variable to retrieve. + + Returns: + The value of the MAD-NG variable. + """ if is_private(item): raise AttributeError(item) return self.__process.recv_vars(item) def __setitem__(self, var_name: str, var: Any) -> None: + """ + Set one or more variables in the MAD-NG process via item assignment. + + Args: + var_name (str): The name(s) of the variable(s). + var (Any): The value(s) to assign. + + Raises: + ValueError: If the number of names does not match the number of values. + """ if not isinstance(var_name, tuple): var_name = (var_name,) var = (var,) @@ -369,6 +386,15 @@ def __setitem__(self, var_name: str, var: Any) -> None: self.__process.send_vars(**dict(zip(var_name, var))) def __getitem__(self, var_name: str) -> Any: + """ + Retrieve one or more variables from MAD-NG using item access. + + Args: + var_name (str): The name or tuple of names of variables. + + Returns: + The corresponding variable value(s) from MAD-NG. + """ if isinstance(var_name, tuple): return self.__process.recv_vars(*var_name) else: @@ -377,34 +403,41 @@ def __getitem__(self, var_name: str) -> Any: # ----------------------------------------------------------------------------------------------# def eval(self, input: str) -> Any: - """Evaluate an expression in MAD and return the result. + """ + Evaluate an expression in the MAD-NG environment. + + Assigns the result to a temporary variable and returns its value. Args: - input (str): An expression that would like to be evaluated in MAD-NG. + input (str): The expression to evaluate. Returns: - The evaluated result. + The evaluated result. """ rtrn = self.__get_mad_reflast() self.send(f"{rtrn._name} =" + input) return rtrn.eval() def evaluate_in_madx_environment(self, input: str) -> None: - """Open the MAD-X environment in MAD-NG and directly send code. + """ + Execute code in the native MAD-X environment. + + Opens a MAD-X environment, sends the code, and then closes the environment. Args: - input (str): The code that would like to be evaluated in the MAD-X environment in MAD-NG. + input (str): The MAD-X code to execute. """ self.__process.send("MADX:open_env()\n" + input + "\nMADX:close_env()") def quote_strings(self, input: str | list[str]) -> str | list[str]: - """Add ' to either side of a string or each string in a list of strings. + """ + Surround the provided string or list of strings with single quotes. Args: - input(str/list[str]): The string(s) that you would like to add ' either side to each string. + input (str or list[str]): The input string(s) to quote. Returns: - A string or list of strings with ' placed at the beginning and the end of each string. + The quoted string or list of quoted strings. """ if isinstance(input, list): return ["'" + x + "'" for x in input] @@ -416,16 +449,17 @@ def quote_strings(self, input: str | list[str]) -> str | list[str]: # ---------------------------------------------------------------------------------------------------# def create_deferred_expression(self, **kwargs) -> mad_high_level_last_ref: - """Create a deferred expression object + """ + Create a deferred expression object in MAD-NG. - For the deferred expression object, the kwargs are used as the deffered expressions, with ``=`` replaced - with ``:=``. To assign the returned object a name, then use ``mad['name'] = mad.deffered(kwargs)``. + The keyword arguments represent deferred expressions using ':=' syntax. + Intended for assigning a temporary variable that holds a lazy-evaluated result. Args: - **kwargs: A variable list of keyword arguments, keyword as the name of the deffered expression within the object and the value as a string that is sent directly to the MAD-NG environment. + **kwargs: Expression name-value pairs. Returns: - A reference to the deffered expression object. + mad_high_level_last_ref: A reference to the deferred expression. """ rtrn = self.__get_mad_reflast() kwargs_string, vars_to_send = format_kwargs_to_string(self.py_name, **kwargs) @@ -443,18 +477,22 @@ def __dir__(self) -> Iterable[str]: return pyObjs def globals(self) -> list[str]: - """Retreive the list of names of variables in the environment of MAD + """ + Retrieve a list of all global variable names in the MAD-NG environment. Returns: - A list of strings indicating the globals variables and modules within the MAD-NG environment + list[str]: A list containing the names of global variables. """ return dir(self.__process.recv_vars(f"{self.py_name}._env")) def history(self) -> str: - """Retrieve the history of strings that have been sent to MAD-NG + """ + Retrieve the command history sent to MAD-NG. + + Filters out internal error-handling commands. Returns: - A string containing the history of commands that have been sent to MAD-NG + str: A newline-separated string of historical commands. """ # delete all lines that start py:__err and end with __err(false)\n history = self.__process.history diff --git a/src/pymadng/madp_pymad.py b/src/pymadng/madp_pymad.py index d3fe4bf..7b86ba6 100644 --- a/src/pymadng/madp_pymad.py +++ b/src/pymadng/madp_pymad.py @@ -15,11 +15,18 @@ if TYPE_CHECKING: from collections.abc import Callable -__all__ = ["mad_process"] - # TODO: look at cpymad for the suppression of the error messages at exit - copy? (jgray 2024) + def is_private(varname): + """Check if the variable name is considered private. + + Args: + varname (str): The variable name to check. + + Returns: + bool: True if the variable name is private, False otherwise. + """ assert isinstance(varname, str), "Variable name to receive must be a string" if varname[0] == "_" and varname[:6] != "_last[": return True @@ -37,7 +44,9 @@ def __init__( redirect_sterr: bool = False, ) -> None: self.py_name = py_name - self.raise_on_madng_error = False # Set the error handler to be off during initialization + + # Set the error handler to be off during initialization + self.raise_on_madng_error = False mad_path = Path(mad_path) if not mad_path.exists(): @@ -56,13 +65,15 @@ def __init__( elif isinstance(stdout, str) or isinstance(stdout, Path): self.stdout_file = open(Path(stdout), "w") stdout = self.stdout_file - + # Convert stdout to a file descriptor try: stdout = stdout.fileno() except AttributeError as e: - raise TypeError("Stdout must be a file name, file descriptor, or an object with the fileno method") from e - + raise TypeError( + "Stdout must be a file name, file descriptor, or an object with the fileno method" + ) from e + # Redirect stderr to stdout, if specified if redirect_sterr: stderr = stdout @@ -71,7 +82,9 @@ def __init__( # Create a chunk of code to start the process lua_debug_flag = "true" if debug else "false" - startupChunk = f"MAD.pymad '{py_name}' {{_dbg = {lua_debug_flag}}} :__ini({mad_write})" + startupChunk = ( + f"MAD.pymad '{py_name}' {{_dbg = {lua_debug_flag}}} :__ini({mad_write})" + ) original_sigint_handler = signal.getsignal(signal.SIGINT) def delete_process(sig, frame): @@ -107,6 +120,7 @@ def delete_process(sig, frame): # Open the pipe from MAD (this is where MAD will no longer hang) self.mad_read_stream = os.fdopen(self.mad_output_pipe, "rb") self.history = "" # Begin the recording of the history + self.debug = debug # Record debug mode status # stdout should be line buffered by default, but for jupyter notebook, # stdout is redirected and not line buffered by default @@ -120,40 +134,60 @@ def delete_process(sig, frame): # Check if MAD started successfully using select mad_rtrn = self.recv() - if not startup_status_checker[0] or mad_rtrn != 'started': + if not startup_status_checker[0] or mad_rtrn != "started": self.close() - if mad_rtrn == 'started': - raise OSError(f"Could not establish communication with {mad_path} process") + if mad_rtrn == "started": + raise OSError( + f"Could not establish communication with {mad_path} process" + ) else: - raise OSError(f"Could not start {mad_path} process, received: {mad_rtrn}") - + raise OSError( + f"Could not start {mad_path} process, received: {mad_rtrn}" + ) + # Set the error handler to be on by default if raise_on_madng_error: self.set_error_handler(True) self.raise_on_madng_error = True def send_range(self, start: float, stop: float, size: int) -> None: - """Send a numpy array as a range to MAD""" + """Send a linear range (numpy array) to MAD-NG. + + Constructs a numpy.linspace array based on start, stop, and size and sends it over the MAD pipe. + """ self.mad_input_stream.write(b"rng_") send_generic_range(self, start, stop, size) def send_logrange(self, start: float, stop: float, size: int) -> None: - """Send a numpy array as a logrange to MAD""" + """Send a logarithmic range (numpy array) to MAD-NG. + + Builds an array equivalent to numpy.geomspace from start to stop with given size. + """ self.mad_input_stream.write(b"lrng") send_generic_range(self, start, stop, size) def send_tpsa(self, monos: np.ndarray, coefficients: np.ndarray) -> None: - """Send the monomials and coeeficients of a TPSA to MAD, creating a table representing the TPSA object""" + """Transmit TPSA data to MAD-NG. + + Sends the monomials and their corresponding coefficients to construct a TPSA table. + """ self.mad_input_stream.write(b"tpsa") send_generic_tpsa(self, monos, coefficients, send_num) def send_cpx_tpsa(self, monos: np.ndarray, coefficients: np.ndarray) -> None: - """Send the monomials and coeeficients of a complex TPSA to MAD, creating a table representing the complex TPSA object""" + """Transmit a complex TPSA to MAD-NG. + + Sends complex monomials and coefficients to MAD-NG for table creation. + """ 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: - """Send data to MAD, returns self for chaining""" + """Send data to the MAD-NG process. + + Accepts several types (str, int, float, ndarray, bool, list) and sends them using the appropriate serialization. + Returns self to allow method chaining. + """ try: typ = type_str[get_typestr(data)] self.mad_input_stream.write(typ.encode("utf-8")) @@ -165,24 +199,27 @@ def send(self, data: str | int | float | np.ndarray | bool | list) -> mad_proces ) from None def protected_send(self, string: str) -> mad_process: - """Perform a protected send to MAD, by first enabling error handling, so that if an error occurs, an error is returned""" - if self.raise_on_madng_error: - # If the user has specified that they want to raise an error always, skip the error handling on and off + """Send a command string to MAD-NG with temporary error handling. + + If error handling is enabled, any errors in MAD-NG will be reported back. + """ + if self.raise_on_madng_error: + # If the user has specified that they want to raise an error always, skip the error handling on and off return self.send(string) return self.send( f"{self.py_name}:__err(true); {string}; {self.py_name}:__err(false);" ) def protected_variable_retrieval(self, name: str) -> Any: - """Perform a protected variable retrieval from MAD, by first enabling error handling, so that if an error occurs, an error is returned + """Safely retrieve a variable from MAD-NG. + Enables temporary error handling while retrieving a variable. Args: - name (str): The name of the variable to retrieve from MAD - + name (str): The MAD-NG variable name to retrieve. Returns: - Any: The value of the variable retrieved from MAD + The value of the variable. """ - if self.raise_on_madng_error: + if self.raise_on_madng_error: return self.send(f"py:send({name})").recv(name) self.send( f"{self.py_name}:__err(true):send({name}):__err(false)" @@ -190,19 +227,30 @@ def protected_variable_retrieval(self, name: str) -> Any: return self.recv(name) def set_error_handler(self, on_off: bool) -> mad_process: - """Enable or disable error handling""" - 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 + """Toggle error handling in the MAD-NG process. + + This determines whether errors are raised immediately. + """ + 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 self.send(f"{self.py_name}:__err({str(on_off).lower()})") def recv(self, varname: str = None) -> Any: - """Receive data from MAD, if a function is returned, it will be executed with the argument mad_communication""" + """Receive data from MAD-NG. + + Reads 4 bytes to detect the data type and then extracts the corresponding value. + Optional varname is used for reference purposes. + """ typ = self.mad_read_stream.read(4).decode("utf-8") self.varname = varname # For mad reference return type_fun[typ]["recv"](self) def recv_and_exec(self, env: dict = {}) -> dict: - """Read data from MAD and execute it""" + """Receive a command string from MAD-NG and execute it. + + The execution context includes numpy as np and the mad process instance. + Returns the updated execution environment. + """ # Check if user has already defined mad (madp_object will have mad defined), otherwise define it try: env["mad"] @@ -214,6 +262,10 @@ def recv_and_exec(self, env: dict = {}) -> dict: # ----------------- Dealing with communication of variables ---------------- # def send_vars(self, **vars) -> mad_process: + """Send multiple variables to MAD-NG. + + Each keyword argument becomes a variable in the MAD-NG environment. + """ for name, var in vars.items(): if isinstance(var, mad_ref): self.send(f"{name} = {var._name}") @@ -221,6 +273,11 @@ def send_vars(self, **vars) -> mad_process: self.send(f"{name} = {self.py_name}:recv()").send(var) def recv_vars(self, *names) -> 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. + """ if len(names) == 1: if not is_private(names[0]): return self.protected_variable_retrieval(names[0]) @@ -234,7 +291,10 @@ def recv_vars(self, *names) -> Any: # -------------------------------------------------------------------------- # def close(self) -> None: - """Close the pipes and wait for the process to finish""" + """Terminate the MAD-NG process. + + Closes all communication pipes and waits for the subprocess to finish. + """ if self.process.poll() is None: # If process is still running self.send(f"{self.py_name}:__fin()") # Tell the mad side to finish open_pipe = select.select([self.mad_read_stream], [], []) @@ -246,7 +306,7 @@ def close(self) -> None: f"Unexpected message received: {close_msg}, MAD-NG may not have completed properly" ) self.process.terminate() # Terminate the process on the python side - + # Close the debug file if it exists with suppress(AttributeError): self.stdout_file.close() @@ -261,18 +321,35 @@ def close(self) -> None: self.process.wait() def __del__(self): + """Destructor: Close the MAD process gracefully.""" self.close() class mad_ref(object): + """A reference to a variable in MAD-NG. + This class allows for the retrieval of variables from MAD-NG without + having to send them explicitly. + """ def __init__(self, name: str, mad_proc: mad_process): - assert ( - name is not None - ), "Reference must have a variable to reference to. Did you forget to put a name in the receive functions?" + assert name is not None, ( + "Reference must have a variable to reference to." + "Did you forget to put a name in the receive functions?" + ) self._name = name self._mad = mad_proc def __getattr__(self, item): + """Retrieve attribute corresponding to a variable in MAD-NG. + + Args: + item (str): The attribute name. + + Returns: + Any: The value of the variable in MAD-NG. + + Raises: + AttributeError: If the attribute is not found. + """ if not is_private(item): try: return self[item] @@ -281,8 +358,21 @@ def __getattr__(self, item): raise AttributeError(item) # For python def __getitem__(self, item: str | int): + """Retrieve item from MAD-NG using indexing. + + Args: + item (str | int): The key or index. + + Returns: + Any: The value corresponding to the item. + + Raises: + IndexError: If the index is out of range. + KeyError: If the key is not present. + TypeError: If the item type is not valid. + """ if isinstance(item, int): - result = self._mad.protected_variable_retrieval(f"{self._name}[{item+1}]") + result = self._mad.protected_variable_retrieval(f"{self._name}[{item + 1}]") if result is None: raise IndexError(item) # For python elif isinstance(item, str): @@ -295,6 +385,7 @@ def __getitem__(self, item: str | int): return result def eval(self) -> Any: + """Evaluate the reference and return the value.""" return self._mad.recv_vars(self._name) @@ -302,102 +393,284 @@ def eval(self) -> Any: # Data ----------------------------------------------------------------------- # -def write_serial_data(self: mad_process, dat_fmt: str, *dat: Any): - self.mad_input_stream.write(struct.pack(dat_fmt, *dat)) +def write_serial_data(self: mad_process, dat_fmt: str, *dat: Any) -> int: + """Write data to the MAD-NG pipe in a specific format." + + Args: + dat_fmt (str): The format string for struct.pack. + *dat: The data to be packed and written. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ + return self.mad_input_stream.write(struct.pack(dat_fmt, *dat)) -def read_data_stream(self: mad_process, dat_sz: int, dat_typ: np.dtype): +def read_data_stream(self: mad_process, dat_sz: int, dat_typ: np.dtype) -> np.ndarray: + """Read data from the MAD-NG pipe in a specific format. + + Args: + dat_sz (int): The size of the data to read. + dat_typ (np.dtype): The data type for numpy conversion. + + Returns: + np.ndarray: The data read from the MAD-NG input stream, converted to the specified numpy type. + """ return np.frombuffer(self.mad_read_stream.read(dat_sz), dtype=dat_typ) # None ----------------------------------------------------------------------- # def send_nil(self: mad_process, input): + """Send a nil value to MAD-NG. + + This is a placeholder function as nil is not sent down the pipe, but rather + the sending of the type is used to tell MAD-NG that the value is nil. + The function is included for consistency with the other send functions. + + Args: + self (mad_process): The MAD-NG process instance. + input: The input value (not used). + + Returns: + None: No value is sent, but the function is included for consistency. + """ return None def recv_nil(self: mad_process): + """Receive a nil value from MAD-NG. + + This is a placeholder function as nil is not sent down the pipe, but rather + the receiving of the type is used to tell MAD-NG that the value is nil. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + None: No value is received. + """ return None # Boolean -------------------------------------------------------------------- # -def send_bool(self: mad_process, input: bool): +def send_bool(self: mad_process, input: bool) -> int: + """Write a boolean value to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + input (bool): The boolean value to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ return self.mad_input_stream.write(struct.pack("?", input)) def recv_bool(self: mad_process) -> bool: + """Read a boolean value from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + bool: The boolean value received from MAD-NG. + """ return read_data_stream(self, 1, np.bool_)[0] # int32 ---------------------------------------------------------------------- # -def send_int(self: mad_process, input: int): +def send_int(self: mad_process, input: int) -> int: + """Send a 32-bit integer to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + input (int): The integer value to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ return write_serial_data(self, "i", input) def recv_int(self: mad_process) -> int: + """Receive a 32-bit integer from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + int: The integer value received from MAD-NG. + """ return read_data_stream(self, 4, np.int32)[0] # String --------------------------------------------------------------------- # -def send_str(self: mad_process, input: str): - self.history += input + "\n" +def send_str(self: mad_process, input: str) -> int: + """Send a string to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + input (str): The string value to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ + # Only store history if debug mode is enabled + if getattr(self, "debug", False): + self.history += input + "\n" send_int(self, len(input)) - self.mad_input_stream.write(input.encode("utf-8")) + return self.mad_input_stream.write(input.encode("utf-8")) def recv_str(self: mad_process) -> str: + """Receive a string from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + str: The string value received from MAD-NG. + """ res = self.mad_read_stream.read(recv_int(self)).decode("utf-8") return res # number (in lua, float64 in python) ----------------------------------------- # -def send_num(self: mad_process, input: float): +def send_num(self: mad_process, input: float) -> int: + """Send a 64-bit float to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + input (float): The float value to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ return write_serial_data(self, "d", input) -def recv_num(self: mad_process) -> float: +def recv_num(self: mad_process) -> np.float64: + """Receive a 64-bit float from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + np.float64: The float value received from MAD-NG. + """ return read_data_stream(self, 8, np.float64)[0] # Complex (complex128) ------------------------------------------------------- # -def send_cpx(self: mad_process, input: complex): +def send_cpx(self: mad_process, input: complex) -> int: + """Send a complex number to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + input (complex): The complex number to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ return write_serial_data(self, "dd", input.real, input.imag) def recv_cpx(self: mad_process) -> complex: + """Receive a complex number from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + + Returns: + complex: The received complex number. + """ return read_data_stream(self, 16, np.complex128)[0] # Range ---------------------------------------------------------------------- # def send_generic_range(self, start: float, stop: float, size: int): + """Send a generic range to MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + start (float): The starting value of the range. + stop (float): The ending value of the range. + size (int): The number of points in the range. + + Returns: + None + """ write_serial_data(self, "ddi", start, stop, size) def recv_range(self: mad_process) -> np.ndarray: + """Receive a linear range from the MAD-NG pipe. + + Returns: + np.ndarray: The range as a numpy array. + """ return np.linspace(*struct.unpack("ddi", self.mad_read_stream.read(20))) def recv_logrange(self: mad_process) -> np.ndarray: + """Receive a logarithmic range from the MAD-NG pipe. + + Returns: + np.ndarray: The logarithmic range as a numpy array. + """ return np.geomspace(*struct.unpack("ddi", self.mad_read_stream.read(20))) # irange --------------------------------------------------------------------- # def send_int_range(self: mad_process, rng: range): + """Send an integer range to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + rng (range): The range object to send. + + Returns: + int: The number of bytes written. + """ return write_serial_data(self, "iii", rng.start, rng.stop, rng.step) def recv_int_range(self: mad_process) -> range: + """Receive an inclusive integer range from the MAD-NG pipe. + + Returns: + range: The received range, inclusive of both ends. + """ start, stop, step = read_data_stream(self, 12, np.int32) return range(start, stop + 1, step) # MAD is inclusive at both ends # matrix --------------------------------------------------------------------- # def send_generic_matrix(self: mad_process, mat: np.ndarray): + """Send a 2D matrix to MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + mat (np.ndarray): A 2-dimensional numpy array to send. + + Returns: + None + """ assert len(mat.shape) == 2, "Matrix must be of two dimensions" write_serial_data(self, "ii", *mat.shape) self.mad_input_stream.write(mat.tobytes()) def recv_generic_matrix(self: mad_process, dtype: np.dtype) -> str: + """Receive a generic matrix from the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + dtype (np.dtype): The numpy data type of the matrix. + + Returns: + str: A string representation of the matrix (reshaped numpy array). + """ shape = read_data_stream(self, 8, np.int32) return read_data_stream(self, shape[0] * shape[1] * dtype.itemsize, dtype).reshape( shape @@ -405,24 +678,53 @@ def recv_generic_matrix(self: mad_process, dtype: np.dtype) -> str: def recv_matrix(self: mad_process) -> np.ndarray: + """Receive a matrix of 64-bit floats from the MAD-NG pipe. + + Returns: + np.ndarray: The received matrix. + """ return recv_generic_matrix(self, np.dtype("float64")) def recv_cpx_matrix(self: mad_process) -> np.ndarray: + """Receive a matrix of complex numbers from the MAD-NG pipe. + + Returns: + np.ndarray: The received complex matrix. + """ return recv_generic_matrix(self, np.dtype("complex128")) def recv_int_matrix(self: mad_process) -> np.ndarray: + """Receive a matrix of 32-bit integers from the MAD-NG pipe. + + Returns: + np.ndarray: The received integer matrix. + """ return recv_generic_matrix(self, np.dtype("int32")) # monomial ------------------------------------------------------------------- # def send_monomial(self: mad_process, mono: np.ndarray): + """Send a monomial to MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + mono (np.ndarray): A numpy array representing the monomial. + + Returns: + None + """ send_int(self, mono.size) self.mad_input_stream.write(mono.tobytes()) def recv_monomial(self: mad_process) -> np.ndarray: + """Receive a monomial from MAD-NG. + + Returns: + np.ndarray: The received monomial as an array of 8-bit unsigned integers. + """ return read_data_stream(self, recv_int(self), np.ubyte) @@ -433,13 +735,24 @@ def send_generic_tpsa( coefficients: np.ndarray, send_num: Callable[[mad_process, float | complex], None], ): + """Send a generic TPSA table to MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + monos (np.ndarray): 2D array of monomials (must be uint8). + coefficients (np.ndarray): Array of coefficients corresponding to the monomials. + send_num (Callable): The function to send a numeric (float or complex) value. + + Returns: + None + """ assert len(monos.shape) == 2, "The list of monomials must have two dimensions" - assert len(monos) == len( - coefficients - ), "The number of monomials must be equal to the number of coefficients" - assert ( - monos.dtype == np.uint8 - ), "The monomials must be of type 8-bit unsigned integer " + assert len(monos) == len(coefficients), ( + "The number of monomials must be equal to the number of coefficients" + ) + assert monos.dtype == np.uint8, ( + "The monomials must be of type 8-bit unsigned integer " + ) write_serial_data(self, "ii", len(monos), len(monos[0])) for mono in monos: self.mad_input_stream.write(mono.tobytes()) @@ -448,6 +761,15 @@ def send_generic_tpsa( def recv_generic_tpsa(self: mad_process, dtype: np.dtype) -> np.ndarray: + """Receive a generic TPSA table from MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + dtype (np.dtype): The numeric data type for coefficients. + + Returns: + np.ndarray: A tuple (monomial list, coefficients array). + """ num_mono, mono_len = read_data_stream(self, 8, np.int32) mono_list = np.reshape( read_data_stream(self, mono_len * num_mono, np.ubyte), @@ -458,26 +780,50 @@ def recv_generic_tpsa(self: mad_process, dtype: np.dtype) -> np.ndarray: def recv_cpx_tpsa(self: mad_process): + """Receive a complex TPSA table from the MAD-NG pipe. + + Returns: + tuple: A tuple containing the monomial list and coefficients. + """ return recv_generic_tpsa(self, np.dtype("complex128")) def recv_dbl_tpsa(self: mad_process): + """Receive a double TPSA table from the MAD-NG pipe. + + Returns: + tuple: A tuple containing the monomial list and coefficients. + """ return recv_generic_tpsa(self, np.dtype("float64")) # lists ---------------------------------------------------------------------- # # Lists of strings are really slow to send, is there a way to improve this? (jgray 2024) def send_list(self: mad_process, lst: list): + """Send a list to the MAD-NG pipe. + + Args: + self (mad_process): The MAD-NG process instance. + lst (list): The list to send. + + Returns: + None + """ send_int(self, len(lst)) for item in lst: self.send(item) def recv_list(self: mad_process) -> list: + """Receive a list from the MAD-NG pipe. + + Returns: + 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)] + 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) @@ -487,17 +833,37 @@ def recv_list(self: mad_process) -> list: # object (table with metatable are treated as pure reference) ---------------- # def recv_reference(self: mad_process): + """Receive a reference to an object from MAD-NG. + + Returns: + mad_ref: A reference object corresponding to the received variable. + """ return mad_ref(self.varname, self) def send_reference(self, obj: mad_ref): + """Send a reference to an object in MAD-NG. + + Args: + self (mad_process): The MAD-NG process instance. + obj (mad_ref): The reference object to send. + + Returns: + int: The number of bytes written to the MAD-NG input stream. + """ return send_str(self, f"return {obj._name}") # error ---------------------------------------------------------------------- # - - def recv_err(self: mad_process): + """Receive an error message from MAD-NG and raise an exception. + + Args: + self (mad_process): The MAD-NG process instance. + + Raises: + RuntimeError: Always raised with the error message from MAD-NG. + """ self.set_error_handler(False) raise RuntimeError("MAD Errored (see the MAD error output)") @@ -524,13 +890,21 @@ def recv_err(self: mad_process): "tpsa": {"recv": recv_dbl_tpsa}, "ctpa": {"recv": recv_cpx_tpsa}, "err_": {"recv": recv_err}, - "" : {"recv": BrokenPipeError}, + "": {"recv": BrokenPipeError}, } def get_typestr( a: str | int | float | np.ndarray | bool | list | tuple | range | mad_ref, ) -> type: + """Determine the type string for the given input. + + Args: + a: The data for which to determine the type. + + Returns: + type: The corresponding type used for serialization. + """ if isinstance(a, np.ndarray): return a.dtype elif type(a) is int: # Check for signed 32 bit int diff --git a/src/pymadng/madp_strings.py b/src/pymadng/madp_strings.py index 98ef634..da6ac22 100644 --- a/src/pymadng/madp_strings.py +++ b/src/pymadng/madp_strings.py @@ -7,13 +7,15 @@ def format_kwargs_to_string(py_name, **kwargs): """Convert a keyword argument input to a string used by MAD-NG + The function produces a Lua table-like string representation of the arguments and + gathers any non-string items in a list for separate sending. + Args: - py_name (str): The name of the Python variable in the MAD-NG environment - **kwargs: The keyword arguments to be converted to a string + py_name (str): The name of the Python reference variable in MAD-NG. + **kwargs: Arbitrary keyword arguments to be converted. Returns: - str: The string representation of the keyword arguments (may include variables that need to be sent through the pipe) - list: The variables to send to the MAD-NG environment + tuple: A tuple with the formatted string and a list of variables to send. """ formatted_kwargs = "{" vars_to_send = [] @@ -30,13 +32,17 @@ def format_kwargs_to_string(py_name, **kwargs): def format_args_to_string(py_name, *args): """Convert an argument input to a string used by MAD-NG + Convert positional arguments into a MAD-NG formatted string. + + Each argument is processed to produce a string suitable for MAD-NG while collecting any + additional variables that require separate sending. + Args: - py_name (str): The name of the Python variable in the MAD-NG environment - *args: The arguments to be converted to a string + py_name (str): The Python reference name used in MAD-NG. + *args: Positional arguments to be formatted. Returns: - str: The string representation of the arguments (may include variables that need to be sent through the pipe) - list: The variables to send to the MAD-NG environment + tuple: A tuple containing the composed argument string and a list of variables to send. """ mad_string = "" vars_to_send = [] @@ -53,9 +59,17 @@ def create_mad_string(py_name, var: Any): """Convert a list of objects into the required string for MAD-NG. Converting string instead of sending more data is up to 2x faster (therefore last resort). Slowdown is mainly due to sending lists of strings. + Convert a Python variable to its MAD-NG string representation. + + Handles lists, dictionaries, strings, and MAD references. For non-string or non-primitive + types, it falls back on using a receive call. + Args: - py_name (str): The name of the Python variable in the MAD-NG environment - var (Any): The variable to be converted to a string + py_name (str): The Python reference name in MAD-NG. + var (Any): The variable to be converted. + + Returns: + tuple: A tuple containing the formatted string and a list of associated variables. """ if isinstance(var, list): string, vars_to_send = format_args_to_string(py_name, *var)