From f0ea24c7a7b7889c69c544f97a71f682b7e2e4c9 Mon Sep 17 00:00:00 2001
From: ax3l <ax3l@users.noreply.github.com>
Date: Fri, 17 May 2024 06:10:20 +0000
Subject: [PATCH] Update Stub Files

---
 src/python/impactx/__init__.pyi               | 17 +++--
 .../impactx/extensions/ImpactXParIter.pyi     | 69 +++++++++++++++++++
 .../extensions/ImpactXParticleContainer.pyi   | 38 ++++++++++
 src/python/impactx/extensions/__init__.pyi    |  5 ++
 4 files changed, 120 insertions(+), 9 deletions(-)
 create mode 100644 src/python/impactx/extensions/ImpactXParIter.pyi
 create mode 100644 src/python/impactx/extensions/ImpactXParticleContainer.pyi
 create mode 100644 src/python/impactx/extensions/__init__.pyi

diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi
index b3a8b9bce..e1e10cd8a 100644
--- a/src/python/impactx/__init__.pyi
+++ b/src/python/impactx/__init__.pyi
@@ -17,28 +17,26 @@ from __future__ import annotations
 import os as os
 
 from amrex import space3d as amr
+from impactx.extensions.ImpactXParIter import register_ImpactXParIter_extension
+from impactx.extensions.ImpactXParticleContainer import (
+    register_ImpactXParticleContainer_extension,
+)
 from impactx.impactx_pybind import (
     Config,
     CoordSystem,
     ImpactX,
     ImpactXParConstIter,
+    ImpactXParIter,
+    ImpactXParticleContainer,
     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,
-)
+from . import MADXParser, extensions, impactx_pybind, madx_to_impactx
 
 __all__ = [
     "Config",
@@ -54,6 +52,7 @@ __all__ = [
     "cxx",
     "distribution",
     "elements",
+    "extensions",
     "impactx_pybind",
     "madx_to_impactx",
     "os",
diff --git a/src/python/impactx/extensions/ImpactXParIter.pyi b/src/python/impactx/extensions/ImpactXParIter.pyi
new file mode 100644
index 000000000..5757d68dd
--- /dev/null
+++ b/src/python/impactx/extensions/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.
+
+    """
diff --git a/src/python/impactx/extensions/ImpactXParticleContainer.pyi b/src/python/impactx/extensions/ImpactXParticleContainer.pyi
new file mode 100644
index 000000000..dfb5ebefa
--- /dev/null
+++ b/src/python/impactx/extensions/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
+    """
diff --git a/src/python/impactx/extensions/__init__.pyi b/src/python/impactx/extensions/__init__.pyi
new file mode 100644
index 000000000..0aa5be122
--- /dev/null
+++ b/src/python/impactx/extensions/__init__.pyi
@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from . import ImpactXParIter, ImpactXParticleContainer
+
+__all__ = ["ImpactXParIter", "ImpactXParticleContainer"]