Merge pull request #1 from mihofer/main

add XY grid and set coordinates functions, to_pandas, unit tests, readme

.github/workflows/python-package.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9"]
 
     steps:
     - uses: actions/checkout@v3
@@ -27,8 +27,8 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install flake8 pytest
-        python -m pip install .
+        python -m pip install flake8 pytest poetry
+        poetry install --with gh_actions
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names
@@ -37,4 +37,4 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
-        pytest
+        poetry run pytest

README.md

@@ -1,2 +1,48 @@
 # xdyna
+
 Tools to study beam dynamics in xtrack simulations, like dynamic aperture calculations, PYTHIA integration, dynamic indicators, ...
+
+## Dynamic aperture studies
+
+The `xdyna` package provides the `DA` class which serves as a simple front-end for setting up and running dynamic aperture studies.
+
+To start, a `xtrack.line` object is required.
+The following code then sets up the study and launches the tracking
+
+```python
+
+import xdyna as xd
+
+da = xd.DA(
+    name='name_of_your_study', # used to generate a directory where files are stored
+    normalised_emittance=[1,1], # provide normalized emittance for particle initialization in [m]
+    max_turns=1e5, # number of turns to track
+    use_files=False 
+    # in case DA studies should run on HTC condor, files are used to collect the information
+    # if the tracking is performed locally, no files are needed
+)
+
+# initialize a grid of particles using 5 angles in x-y space, in a range from 0 to 20 sigmas in steps of 5 sigma.
+da.generate_initial_radial(angles=5, r_min=0, r_max=20, r_step=5, delta=0.) 
+
+da.line = line # associate prev. created line, holding the lattice and context, with DA object
+
+da.track_job() # start the tracking
+
+da.survival_data # returns a dataframe with the number of survived turns for the initial position of each particle
+
+```
+
+To use on a platform like HTCondor, perform the same setup as before but using `use_files=True`.
+Each HTCondor job then only requires the following lines
+
+```python
+import xdyna as xd
+# This will load the existing DA based on the meta file with the same name found in the working directory.
+# If the script is ran somewhere else, the path to the metafile can be passed with 'path=...'.
+DA = xd.DA(name=study, use_files=True)
+
+# Do the tracking, here for 100 particles.
+# The code will automatically look for particles that are not-submitted yet and use these.
+DA.track_job(npart=100)
+```

pyproject.toml

@@ -14,10 +14,15 @@ python = ">=3.8"
 scipy = ">=1.9.3"
 pandas = ">=1.5.1"
 
-
 [tool.poetry.dev-dependencies]
 pytest = "^5.2"
 
+[tool.poetry.group.gh_actions]
+optional = true
+
+[tool.poetry.group.gh_actions.dependencies]
+xsuite = ">=0.4.0"
+
 [build-system]
 requires = ["poetry-core>=1.0.8"]  # Needed for pip install -e
 build-backend = "poetry.core.masonry.api"

tests/test_ee.py

@@ -13,7 +13,7 @@
 
 @pytest.mark.parametrize("mode", ['t']) # use t mode for now, since z takes quite long
 def test_simple_radial(mode):
-    with open(TEST_DIR/'input'/'tapered_t_b1_thin.json', 'r', encoding='utf-8') as fid:
+    with open(TEST_DIR/'input'/f'tapered_{mode}_b1_thin.json', 'r', encoding='utf-8') as fid:
         loaded_dct = json.load(fid)
     line = xt.Line.from_dict(loaded_dct)
 
@@ -22,13 +22,36 @@ def test_simple_radial(mode):
     ref_particle = xp.Particles(mass0=xp.ELECTRON_MASS_EV, q0=1, p0c=ENERGY[mode]*10**9, x=0, y=0)
     line.particle_ref = ref_particle
     tracker = xt.Tracker(_context=context, line=line)
-    tracker.configure_radiation(mode='mean')
-    tracker.matrix_stability_tol = 9e-1
+    tracker.configure_radiation(model='mean')
 
     DA = xd.DA(name=f'fcc_ee_{mode}',
-               normalised_emittance=[EMITTANCE[mode]['X'], EMITTANCE[mode]['Y']],
+               normalised_emittance=[EMITTANCE[mode]['X']*ref_particle.beta0[0]*ref_particle.gamma0[0],
+                                     EMITTANCE[mode]['Y']*ref_particle.beta0[0]*ref_particle.gamma0[0]],
                max_turns=TURNS[mode],
                use_files=False)
     DA.generate_initial_radial(angles=1, r_min=2, r_max=20, r_step=4., delta=0.000)
     DA.line = line
     DA.track_job()
+
+
+@pytest.mark.parametrize("mode", ['t']) # use t mode for now, since z takes quite long
+def test_simple_grid(mode):
+    with open(TEST_DIR/'input'/f'tapered_{mode}_b1_thin.json', 'r', encoding='utf-8') as fid:
+        loaded_dct = json.load(fid)
+    line = xt.Line.from_dict(loaded_dct)
+
+    context = xo.ContextCpu()
+
+    ref_particle = xp.Particles(mass0=xp.ELECTRON_MASS_EV, q0=1, p0c=ENERGY[mode]*10**9, x=0, y=0)
+    line.particle_ref = ref_particle
+    tracker = xt.Tracker(_context=context, line=line)
+    tracker.configure_radiation(model='mean')
+
+    DA = xd.DA(name=f'fcc_ee_{mode}',
+               normalised_emittance=[EMITTANCE[mode]['X']*ref_particle.beta0[0]*ref_particle.gamma0[0],
+                                     EMITTANCE[mode]['Y']*ref_particle.beta0[0]*ref_particle.gamma0[0]],
+               max_turns=TURNS[mode],
+               use_files=False)
+    DA.generate_initial_grid(x_min=0, x_max=20, x_step=10, y_min=0, y_max=20, y_step=10, delta=0.000)
+    DA.line = line
+    DA.track_job()

tests/test_initialization.py

@@ -0,0 +1,83 @@
+import pytest
+import pandas as pd
+from pandas.testing import assert_frame_equal
+import xdyna as xd
+
+
+SURV_COLUMNS = ['ang_xy', 'r_xy', 'nturns', 'x_norm_in', 'y_norm_in', 'px_norm_in',
+       'py_norm_in', 'zeta_in', 'delta_in', 'x_out', 'y_out', 'px_out',
+       'py_out', 'zeta_out', 'delta_out', 's_out', 'state', 'submitted',
+       'finished']
+
+
+
+def test_mismatch_user_coordinates():
+
+    DA = xd.DA(name='user_coordinates',
+               normalised_emittance=[1,1],
+               max_turns=2,
+               use_files=False)
+    with pytest.raises(AssertionError):
+        DA.set_coordinates(x=[1], y = [1,2])
+
+
+def test_user_coordinates():
+
+    DA = xd.DA(name='user_coordinates',
+               normalised_emittance=[1,1],
+               max_turns=2,
+               use_files=False)
+    DA.set_coordinates(x=[1,2], y = [3,4], px=[5,6])
+
+    assert_frame_equal(DA.survival_data[['x', 'y', 'px', 'py', 'delta']],
+                       pd.DataFrame(data={
+                                        'x':[1,2],
+                                        'y':[3,4],
+                                        'px':[5,6],
+                                        'py':[0,0],
+                                        'delta':[0,0]
+                                            }) )
+
+
+def test_xy_grid():
+
+    DA = xd.DA(name='user_coordinates',
+               normalised_emittance=[1,1],
+               max_turns=2,
+               use_files=False)
+    DA.generate_initial_grid(
+        x_min=0, x_max=2, x_step=2,
+        y_min=0, y_max=2, y_step=2,
+    )
+    assert_frame_equal(DA.survival_data[['x', 'y']],
+                       pd.DataFrame(data={'x':[0.,2.,0.,2.], 'y':[0.,0.,2.,2.]}) )
+
+
+def test_radial_grid():
+
+    DA = xd.DA(name='user_coordinates',
+               normalised_emittance=[1,1],
+               max_turns=2,
+               use_files=False)
+    DA.generate_initial_radial(
+        r_min=0, r_max=2, r_step=2,
+        ang_min=0, ang_max=90, angles=2,
+    )
+    assert_frame_equal(DA.survival_data[['amplitude', 'angle']],
+                       pd.DataFrame(data={'amplitude':[0.,2.,0.,2.], 'angle':[0.,0.,90.,90.]}) )
+
+
+def test_pandas():
+
+    DA = xd.DA(name='user_coordinates',
+               normalised_emittance=[1,1],
+               max_turns=2,
+               use_files=False)
+    DA.generate_initial_radial(
+        r_min=0, r_max=2, r_step=2,
+        ang_min=0, ang_max=90, angles=2,
+    )
+
+    assert_frame_equal(DA.survival_data,
+                       DA.to_pandas())
+    assert all(elem in DA.to_pandas(full=True).columns for elem in SURV_COLUMNS)