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,
+        dispX=-0.26669723385388505,
     ),
-    dispX=-0.26669723385388505,
 )
 sim.add_particles(bunch_charge_C, distr, npart)
 

src/python/impactx/distribution_input_helpers.py

@@ -19,6 +19,16 @@ def twiss(
     alpha_x: np.float64 = 0.0,
     alpha_y: np.float64 = 0.0,
     alpha_t: np.float64 = 0.0,
+    meanX: np.float64 = 0.0,
+    meanY: np.float64 = 0.0,
+    meanT: np.float64 = 0.0,
+    meanPx: np.float64 = 0.0,
+    meanPy: np.float64 = 0.0,
+    meanPt: np.float64 = 0.0,
+    dispX: np.float64 = 0.0,
+    dispY: np.float64 = 0.0,
+    dispPx: np.float64 = 0.0,
+    dispPy: np.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 meanX: offset of the mean (centroid) position in x from that of the reference particle
+    :param meanY: offset of the mean (centroid) position in y from that of the reference particle
+    :param meanT: offset of the mean (centroid) position in t from that of the reference particle
+    :param meanPx: offset of the mean (centroid) momentum in x from that of the reference particle
+    :param meanPy: offset of the mean (centroid) momentum in y from that of the reference particle
+    :param meanPt: offset of the mean (centroid) momentum in t from that of the reference particle
+    :param dispX: dispersion and its derivative in horizontal and vertical directions
+    :param dispY: dispersion and its derivative in horizontal and vertical directions
+    :param dispPx: dispersion and its derivative in horizontal and vertical directions
+    :param dispPy: 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'.
     """
     if beta_x <= 0.0 or beta_y <= 0.0 or beta_t <= 0.0:
@@ -63,4 +83,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": meanX,
+        "meanY": meanY,
+        "meanT": meanT,
+        "meanPx": meanPx,
+        "meanPy": meanPy,
+        "meanPt": meanPt,
+        "dispX": dispX,
+        "dispY": dispY,
+        "dispPx": dispPx,
+        "dispPy": dispPy,
     }