Python Twiss Helper: mean & dispersion

Simplify usage & pick-up in dashboard.
Straight forwarding of mean and dispersion parameters.

Author: ax3l

examples/dogleg/run_dogleg_reverse.py

@@ -40,8 +40,8 @@
         alpha_x=1.3385830279518021,
         alpha_y=-1.3479109197361046,
         alpha_t=92.624347459169869,
+        dispersion_x=-0.26669723385388505,
     ),
-    dispX=-0.26669723385388505,
 )
 sim.add_particles(bunch_charge_C, distr, npart)
 

examples/htu_beamline/run_impactx_offenergy1.py

@@ -66,8 +66,8 @@
         alpha_x=0.0,
         alpha_y=0.0,
         alpha_t=0.0,
+        mean_pt=-mu_p,
     ),
-    meanPt=-mu_p,
 )
 sim.add_particles(bunch_charge_C, distr, npart)
 

examples/htu_beamline/run_impactx_offenergy2.py

@@ -66,8 +66,8 @@
         alpha_x=0.0,
         alpha_y=0.0,
         alpha_t=0.0,
+        mean_pt=-mu_p,
     ),
-    meanPt=-mu_p,
 )
 sim.add_particles(bunch_charge_C, distr, npart)
 

src/python/impactx/distribution_input_helpers.py

@@ -19,6 +19,16 @@ def twiss(
     alpha_x: numpy.float64 = 0.0,
     alpha_y: numpy.float64 = 0.0,
     alpha_t: numpy.float64 = 0.0,
+    mean_x: numpy.float64 = 0.0,
+    mean_y: numpy.float64 = 0.0,
+    mean_t: numpy.float64 = 0.0,
+    mean_px: numpy.float64 = 0.0,
+    mean_py: numpy.float64 = 0.0,
+    mean_pt: numpy.float64 = 0.0,
+    dispersion_x: numpy.float64 = 0.0,
+    dispersion_y: numpy.float64 = 0.0,
+    dispersion_px: numpy.float64 = 0.0,
+    dispersion_py: numpy.float64 = 0.0,
 ):
     """
     Helper function to convert Courant-Snyder / Twiss input into phase space ellipse input.
@@ -32,6 +42,16 @@ def twiss(
     :param alpha_x: Alpha function value () in the x dimension, default is 0.0.
     :param alpha_y: Alpha function value in the y dimension, default is 0.0.
     :param alpha_t: Alpha function value in the t dimension, default is 0.0.
+    :param mean_x: offset of the mean (centroid) position in x from that of the reference particle
+    :param mean_y: offset of the mean (centroid) position in y from that of the reference particle
+    :param mean_t: offset of the mean (centroid) position in t from that of the reference particle
+    :param mean_px: offset of the mean (centroid) momentum in x from that of the reference particle
+    :param mean_py: offset of the mean (centroid) momentum in y from that of the reference particle
+    :param mean_pt: offset of the mean (centroid) momentum in t from that of the reference particle
+    :param dispersion_x: dispersion and its derivative in horizontal and vertical directions
+    :param dispersion_y: dispersion and its derivative in horizontal and vertical directions
+    :param dispersion_px: dispersion and its derivative in horizontal and vertical directions
+    :param dispersion_py: dispersion and its derivative in horizontal and vertical directions
     :return: A dictionary containing calculated phase space input: 'lambdaX', 'lambdaY', 'lambdaT', 'lambdaPx', 'lambdaPy', 'lambdaPt', 'muxpx', 'muypy', 'mutpt'.
     """
     import numpy as np
@@ -65,4 +85,14 @@ def twiss(
         "muxpx": alpha_x / np.sqrt(beta_x * gamma_x),
         "muypy": alpha_y / np.sqrt(beta_y * gamma_y),
         "mutpt": alpha_t / np.sqrt(beta_t * gamma_t),
+        "meanX": mean_x,
+        "meanY": mean_y,
+        "meanT": mean_t,
+        "meanPx": mean_px,
+        "meanPy": mean_py,
+        "meanPt": mean_pt,
+        "dispX": dispersion_x,
+        "dispY": dispersion_y,
+        "dispPx": dispersion_px,
+        "dispPy": dispersion_py,
     }