Python: sim.set_diag_iota_invariants() (#476)

* Python: `sim.set_diag_iota_invariants()` Expose the parameters for the IOTA nonlinear lens invariants diagnostics to Python. Example usage before `sim.evolve()` call: ```py sim.set_diag_iota_invariants( alpha=0.0, beta=1.0, tn=0.4, cn=0.01 ) * Add Test
ECP-WarpX · Dec 1, 2023 · 4e8abf5 · 4e8abf5
1 parent 79fab7c
commit 4e8abf5
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 5 deletions.
diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -116,6 +116,15 @@ General
       The minimum number of digits (default: ``6``) used for the step
       number appended to the diagnostic file names.
 
+   .. py:method:: set_diag_iota_invariants(alpha, beta, tn, cn)
+
+      Set the parameters of the IOTA nonlinear lens invariants diagnostics.
+
+      :param float alpha: Twiss alpha
+      :param float beta: Twiss beta (m)
+      :param float tn: dimensionless strength of the nonlinear insert
+      :param float cn: scale parameter of the nonlinear insert (m^[1/2])
+
    .. py:property:: particle_lost_diagnostics_backend
 
       Diagnostics for particles lost in apertures.

diff --git a/examples/iota_lens/run_iotalens.py b/examples/iota_lens/run_iotalens.py
@@ -20,6 +20,9 @@
 # domain decomposition & space charge mesh
 sim.init_grids()
 
+# diagnostics: IOTA nonlinear lens invariants calculation
+sim.set_diag_iota_invariants(alpha=0.0, beta=1.0, tn=0.4, cn=0.01)
+
 # load a 2.5 MeV proton beam
 kin_energy_MeV = 2.5  # reference energy
 bunch_charge_C = 1.0e-9  # used with space charge

diff --git a/src/particles/diagnostics/NonlinearLensInvariants.H b/src/particles/diagnostics/NonlinearLensInvariants.H
@@ -49,7 +49,7 @@ namespace impactx::diagnostics
          * @param cn - scale parameter of the nonlinear insert (m^[1/2])
          *
          */
-        NonlinearLensInvariants(
+        NonlinearLensInvariants (
             amrex::ParticleReal const alpha,
             amrex::ParticleReal const beta,
             amrex::ParticleReal const tn,
@@ -68,10 +68,11 @@ namespace impactx::diagnostics
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         Data operator() (
-                amrex::ParticleReal const x,
-                amrex::ParticleReal const y,
-                amrex::ParticleReal const px,
-                amrex::ParticleReal const py ) const
+            amrex::ParticleReal const x,
+            amrex::ParticleReal const y,
+            amrex::ParticleReal const px,
+            amrex::ParticleReal const py
+        ) const
         {
 
             using namespace amrex::literals; // for _rt and _prt

diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp
@@ -272,6 +272,20 @@ void init_ImpactX (py::module& m)
              "The minimum number of digits (default: 6) used for the step\n"
              "number appended to the diagnostic file names."
         )
+        .def("set_diag_iota_invariants",
+              [](ImpactX & /* ix */, amrex::ParticleReal alpha, amrex::ParticleReal beta, amrex::ParticleReal tn, amrex::ParticleReal cn) {
+                  amrex::ParmParse pp_diag("diag");
+
+                  pp_diag.add("alpha", alpha);
+                  pp_diag.add("beta", beta);
+                  pp_diag.add("tn", tn);
+                  pp_diag.add("cn", cn);
+              },
+              py::arg("alpha"), py::arg("beta"), py::arg("tn"), py::arg("cn"),
+              "Set the Twiss alpha, beta (m), dimensionless strength of the nonlinear insert and "
+              "scale parameter of the nonlinear insert (m^[1/2]) of the IOTA nonlinear lens "
+              "invariants diagnostics."
+        )
         .def_property("particle_lost_diagnostics_backend",
                       [](ImpactX & /* ix */) {
                           return detail::get_or_throw<std::string>("diag", "backend");