Update Stub Files

examples/aperture/run_aperture.py

@@ -7,6 +7,7 @@
 # -*- coding: utf-8 -*-
 
 import amrex.space3d as amr
+
 from impactx import ImpactX, distribution, elements
 
 # work-around for https://github.com/ECP-WarpX/impactx/issues/499

examples/fodo_tune/run_fodo_tune.py

@@ -7,6 +7,7 @@
 # -*- coding: utf-8 -*-
 
 import amrex.space3d as amr
+
 from impactx import ImpactX, distribution, elements
 
 sim = ImpactX()

examples/initialize_from_array/run_from_array.py

@@ -9,6 +9,7 @@
 import amrex.space3d as amr
 import numpy as np
 import transformation_utilities as pycoord
+
 from impactx import Config, ImpactX, elements
 
 ################

examples/optimize_triplet/run_triplet.py

@@ -7,11 +7,12 @@
 # -*- coding: utf-8 -*-
 
 import amrex.space3d as amr
-import impactx
 import numpy as np
-from impactx import ImpactX, distribution, elements
 from scipy.optimize import minimize
 
+import impactx
+from impactx import ImpactX, distribution, elements
+
 # Call MPI_Init and MPI_Finalize only once:
 if impactx.Config.have_mpi:
     from mpi4py import MPI  # noqa

examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py

@@ -20,6 +20,8 @@
 
 import numpy as np
 import scipy.optimize as opt
+from surrogate_model_definitions import surrogate_model
+
 from impactx import (
     Config,
     CoordSystem,
@@ -29,7 +31,6 @@
     distribution,
     elements,
 )
-from surrogate_model_definitions import surrogate_model
 
 try:
     import torch

src/python/impactx/ImpactXParIter.pyi

@@ -0,0 +1,69 @@
+"""
+
+This file is part of ImpactX
+
+Copyright 2023 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "register_ImpactXParIter_extension",
+    "soa",
+    "soa_int_comps",
+    "soa_real_comps",
+]
+
+def register_ImpactXParIter_extension(impactx_pybind):
+    """
+    ImpactXParIter helper methods
+    """
+
+def soa(self):
+    """
+    Get the StructOfArrays on the current tile
+
+        Parameters
+        ----------
+        self : ImpactXParIter or ImpactXParConstIter
+          used to query particle container component names
+
+    """
+
+def soa_int_comps(pti, num_comps):
+    """
+
+    Name the ImpactX int components in SoA.
+
+    Parameters
+    ----------
+    pti : ImpactXParIter or ImpactXParConstIter
+      used to query particle container component names
+    num_comps : int
+      number of components to generate names for.
+
+    Returns
+    -------
+    A list of length num_comps with values.
+
+    """
+
+def soa_real_comps(pti, num_comps):
+    """
+
+    Name the ImpactX ParticleReal components in SoA.
+
+    Parameters
+    ----------
+    pti : ImpactXParIter or ImpactXParConstIter
+      used to query particle container component names
+    num_comps : int
+      number of components to generate names for.
+
+    Returns
+    -------
+    A list of length num_comps with values.
+
+    """

src/python/impactx/ImpactXParticleContainer.pyi

@@ -0,0 +1,38 @@
+"""
+
+This file is part of ImpactX
+
+Copyright 2023 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+from __future__ import annotations
+
+__all__ = ["ix_pc_plot_mpl_phasespace", "register_ImpactXParticleContainer_extension"]
+
+def ix_pc_plot_mpl_phasespace(self, num_bins=50, root_rank=0):
+    """
+
+    Plot the longitudinal and transverse phase space projections with matplotlib.
+
+    Parameters
+    ----------
+    self : ImpactXParticleContainer_*
+        The particle container class in ImpactX
+    num_bins : int, default=50
+        The number of bins for spatial and momentum directions per plot axis.
+    root_rank : int, default=0
+        MPI root rank to reduce to in parallel runs.
+
+    Returns
+    -------
+    A matplotlib figure with containing the plot.
+    For MPI-parallel ranks, the figure is only created on the root_rank.
+
+    """
+
+def register_ImpactXParticleContainer_extension(ixpc):
+    """
+    ImpactXParticleContainer helper methods
+    """

src/python/impactx/MADXParser.pyi

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import os as os
+import re as re
+import warnings as warnings
+
+__all__ = [
+    "MADXInputError",
+    "MADXInputWarning",
+    "MADXParser",
+    "MADXParserError",
+    "os",
+    "re",
+    "warnings",
+]
+
+class MADXInputError(MADXParserError):
+    def __init__(self, args, with_traceback): ...
+
+class MADXInputWarning(UserWarning):
+    pass
+
+class MADXParser:
+    """
+
+    Simple MADX parser.
+    It expects a single line per element.
+
+    """
+
+    def __init__(self): ...
+    def __str__(self): ...
+    def _combine(self, lattice):
+        """
+
+        Combine to one list of all basic
+        elements.
+
+        return a list of of element dictionaries
+
+        """
+
+    def _flatten(self, line):
+        """
+
+        Find sublines.
+
+
+        """
+
+    def _noWhitespace(self, string):
+        """
+
+        Remove white space from a string.
+
+        14. Oct. 2017,
+        https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string
+
+
+        """
+
+    def getBeamline(self): ...
+    def getEtot(self): ...
+    def getParticle(self): ...
+    def nonblank_lines_to_lowercase(self, f): ...
+    def parse(self, fn):
+        """
+
+        fn (str)    filename
+
+
+        """
+
+class MADXParserError(Exception):
+    pass

src/python/impactx/__init__.pyi

@@ -0,0 +1,76 @@
+"""
+
+            impactx_pybind
+            --------------
+            .. currentmodule:: impactx_pybind
+
+            .. autosummary::
+               :toctree: _generate
+               ImpactX
+               distribution
+               elements
+
+"""
+
+from __future__ import annotations
+
+import os as os
+
+from amrex import space3d as amr
+
+from impactx.impactx_pybind import (
+    Config,
+    CoordSystem,
+    ImpactX,
+    ImpactXParConstIter,
+    RefPart,
+    coordinate_transformation,
+    distribution,
+    elements,
+    push,
+)
+from impactx.ImpactXParIter import register_ImpactXParIter_extension
+from impactx.ImpactXParticleContainer import register_ImpactXParticleContainer_extension
+from impactx.madx_to_impactx import read_beam, read_lattice
+
+from . import (
+    ImpactXParIter,
+    ImpactXParticleContainer,
+    MADXParser,
+    impactx_pybind,
+    madx_to_impactx,
+)
+
+__all__ = [
+    "Config",
+    "CoordSystem",
+    "ImpactX",
+    "ImpactXParConstIter",
+    "ImpactXParIter",
+    "ImpactXParticleContainer",
+    "MADXParser",
+    "RefPart",
+    "amr",
+    "coordinate_transformation",
+    "cxx",
+    "distribution",
+    "elements",
+    "impactx_pybind",
+    "madx_to_impactx",
+    "os",
+    "push",
+    "read_beam",
+    "read_lattice",
+    "register_ImpactXParIter_extension",
+    "register_ImpactXParticleContainer_extension",
+    "s",
+    "t",
+]
+__author__: str = (
+    "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al."
+)
+__license__: str = "BSD-3-Clause-LBNL"
+__version__: str = "24.04"
+s: impactx_pybind.CoordSystem  # value = <CoordSystem.s: 0>
+t: impactx_pybind.CoordSystem  # value = <CoordSystem.t: 1>
+cxx = impactx_pybind