From 71df0dd56c8153fe4106ba8c07f2289c74fad68b Mon Sep 17 00:00:00 2001
From: ax3l <1353258+ax3l@users.noreply.github.com>
Date: Tue, 14 Jan 2025 18:26:28 +0000
Subject: [PATCH] Update Stub Files

---
 .../impactx/impactx_pybind/__init__.pyi       |  1 +
 .../impactx_pybind/elements/__init__.pyi      | 52 +++++++++++++++++++
 .../impactx/impactx_pybind/elements/mixin.pyi | 13 ++++-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi
index 1367006a2..acadebb03 100644
--- a/src/python/impactx/impactx_pybind/__init__.pyi
+++ b/src/python/impactx/impactx_pybind/__init__.pyi
@@ -708,6 +708,7 @@ def push(
     | elements.ExactDrift
     | elements.ExactSbend
     | elements.Kicker
+    | elements.LinearMap
     | elements.Marker
     | elements.Multipole
     | elements.NonlinearLens
diff --git a/src/python/impactx/impactx_pybind/elements/__init__.pyi b/src/python/impactx/impactx_pybind/elements/__init__.pyi
index 950e24875..d4fdcc158 100644
--- a/src/python/impactx/impactx_pybind/elements/__init__.pyi
+++ b/src/python/impactx/impactx_pybind/elements/__init__.pyi
@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import typing
 
+import amrex.space3d.amrex_3d_pybind
 import impactx.impactx_pybind
 
 from . import mixin
@@ -27,6 +28,7 @@ __all__ = [
     "ExactSbend",
     "Kicker",
     "KnownElementsList",
+    "LinearMap",
     "Marker",
     "Multipole",
     "NonlinearLens",
@@ -698,6 +700,7 @@ class KnownElementsList:
         | ExactDrift
         | ExactSbend
         | Kicker
+        | LinearMap
         | Marker
         | Multipole
         | NonlinearLens
@@ -734,6 +737,7 @@ class KnownElementsList:
         | ExactDrift
         | ExactSbend
         | Kicker
+        | LinearMap
         | Marker
         | Multipole
         | NonlinearLens
@@ -771,6 +775,7 @@ class KnownElementsList:
         | ExactDrift
         | ExactSbend
         | Kicker
+        | LinearMap
         | Marker
         | Multipole
         | NonlinearLens
@@ -810,6 +815,53 @@ class KnownElementsList:
         Return and remove the last element of the list.
         """
 
+class LinearMap(mixin.Named, mixin.Alignment, mixin.LinearTransport):
+    @staticmethod
+    def _pybind11_conduit_v1_(*args, **kwargs): ...
+    def __init__(
+        self,
+        R: amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double,
+        ds: float = 0,
+        dx: float = 0,
+        dy: float = 0,
+        rotation: float = 0,
+        name: str | None = None,
+    ) -> None:
+        """
+        (A user-provided linear map, represented as a 6x6 transport matrix.)
+        """
+    def __repr__(self) -> str: ...
+    def push(
+        self,
+        pc: impactx.impactx_pybind.ImpactXParticleContainer,
+        step: int = 0,
+        period: int = 0,
+    ) -> None:
+        """
+        Push first the reference particle, then all other particles.
+        """
+    @property
+    def R(self) -> amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double:
+        """
+        linear map as a 6x6 transport matrix
+        """
+    @R.setter
+    def R(
+        self, arg1: amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double
+    ) -> None: ...
+    @property
+    def ds(self) -> float:
+        """
+        segment length in m
+        """
+    @ds.setter
+    def ds(self, arg1: float) -> None: ...
+    @property
+    def nslice(self) -> int:
+        """
+        one, because we do not support slicing of this element
+        """
+
 class Marker(mixin.Named, mixin.Thin):
     @staticmethod
     def _pybind11_conduit_v1_(*args, **kwargs): ...
diff --git a/src/python/impactx/impactx_pybind/elements/mixin.pyi b/src/python/impactx/impactx_pybind/elements/mixin.pyi
index ff84c2db0..da971b855 100644
--- a/src/python/impactx/impactx_pybind/elements/mixin.pyi
+++ b/src/python/impactx/impactx_pybind/elements/mixin.pyi
@@ -4,7 +4,9 @@ Mixin classes for accelerator lattice elements in ImpactX
 
 from __future__ import annotations
 
-__all__ = ["Alignment", "Named", "Thick", "Thin"]
+import amrex.space3d.amrex_3d_pybind
+
+__all__ = ["Alignment", "LinearTransport", "Named", "Thick", "Thin"]
 
 class Alignment:
     @staticmethod
@@ -35,6 +37,15 @@ class Alignment:
     @rotation.setter
     def rotation(self, arg1: float) -> None: ...
 
+class LinearTransport:
+    Map6x6 = amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double
+    @staticmethod
+    def _pybind11_conduit_v1_(*args, **kwargs): ...
+    def __init__(self) -> None:
+        """
+        Mixin class for linear transport approximation via matrices.
+        """
+
 class Named:
     @staticmethod
     def _pybind11_conduit_v1_(*args, **kwargs): ...