Update Stub Files

src/python/impactx/__init__.pyi

@@ -17,6 +17,7 @@ from __future__ import annotations
 import os as os
 
 from amrex import space3d as amr
+from impactx.distribution_input_helpers import twiss
 from impactx.extensions.ImpactXParIter import register_ImpactXParIter_extension
 from impactx.extensions.ImpactXParticleContainer import (
     register_ImpactXParticleContainer_extension,
@@ -37,7 +38,13 @@ from impactx.impactx_pybind import (
 )
 from impactx.madx_to_impactx import read_beam, read_lattice
 
-from . import MADXParser, extensions, impactx_pybind, madx_to_impactx
+from . import (
+    MADXParser,
+    distribution_input_helpers,
+    extensions,
+    impactx_pybind,
+    madx_to_impactx,
+)
 
 __all__ = [
     "Config",
@@ -52,6 +59,7 @@ __all__ = [
     "coordinate_transformation",
     "cxx",
     "distribution",
+    "distribution_input_helpers",
     "elements",
     "extensions",
     "impactx_pybind",
@@ -64,6 +72,7 @@ __all__ = [
     "register_ImpactXParticleContainer_extension",
     "s",
     "t",
+    "twiss",
     "wakeconvolution",
 ]
 __author__: str = (

src/python/impactx/distribution_input_helpers.pyi

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import numpy
+import numpy as np
+
+__all__ = ["np", "twiss"]
+
+def twiss(
+    beta_x: numpy.float64,
+    beta_y: numpy.float64,
+    beta_t: numpy.float64,
+    emitt_x: numpy.float64,
+    emitt_y: numpy.float64,
+    emitt_t: numpy.float64,
+    alpha_x: numpy.float64 = 0.0,
+    alpha_y: numpy.float64 = 0.0,
+    alpha_t: numpy.float64 = 0.0,
+):
+    """
+
+    Helper function to convert Courant-Snyder / Twiss input into phase space ellipse input.
+
+    :param beta_x: Beta function value (unit: meter) in the x dimension, must be a non-zero positive value.
+    :param beta_y: Beta function value (unit: meter) in the y dimension, must be a non-zero positive value.
+    :param beta_t: Beta function value (unit: meter) in the t dimension (arrival time differences multiplied by light speed), must be a non-zero positive value.
+    :param emitt_x: Emittance value (unit: meter times radian) in the x dimension, must be a non-zero positive value.
+    :param emitt_y: Emittance value (unit: meter times radian) in the y dimension, must be a non-zero positive value.
+    :param emitt_t: Emittance value (unit: meter times radian) in the t dimension (arrival time differences multiplied by light speed), must be a non-zero positive value.
+    :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.
+    :return: A dictionary containing calculated phase space input: 'lambdaX', 'lambdaY', 'lambdaT', 'lambdaPx', 'lambdaPy', 'lambdaPt', 'muxpx', 'muypy', 'mutpt'.
+
+    """