Mode: Reference Orbit Tracking (#831)

Add a simulation mode that only tracks the reference orbit.

docs/source/usage/parameters.rst

@@ -21,8 +21,8 @@ Tracking Modes
   Mode that specifies how the beam is tracked:
 
   * ``particles`` (default): symplectic particle tracking
-  * ``reference_orbit``: *not yet implemented*, only tracking of the reference particle orbit
   * ``envelope``: beam envelop (covariance matrix) tracking, through linearized transport maps
+  * ``reference_orbit``: only tracking of the reference particle orbit
 
 
 .. _running-cpp-parameters-particle:

docs/source/usage/python.rst

@@ -233,6 +233,10 @@ Collective Effects & Overall Simulation Parameters
 
       Run the envelope tracking simulation loop.
 
+   .. py:method:: track_reference()
+
+      Run the reference orbit tracking simulation loop.
+
    .. py:method:: resize_mesh()
 
       Resize the mesh :py:attr:`~domain` based on the :py:attr:`~dynamic_size` and related parameters.

src/ImpactX.H

@@ -114,6 +114,10 @@ namespace impactx
          */
         void track_envelope ();
 
+        /** Run the reference orbit tracking simulation loop
+         */
+        void track_reference (RefPart & ref);
+
         /** Query input for warning logger variables and set up warning logger accordingly
          *
          * Input variables are: ``always_warn_immediately`` and ``abort_on_warning_threshold``.

src/ImpactX.cpp

@@ -136,7 +136,16 @@ namespace impactx {
         else if (track == "envelope") {
             track_envelope();
         }
-        // TODO: reference_orbit only tracking
+        else if (track == "reference_orbit") {
+            if (!amr_data->track_reference.m_ref.has_value())
+            {
+                throw std::runtime_error("evolve: Reference particle not set.");
+            }
+            track_reference(amr_data->track_reference.m_ref.value());
+        }
+        else {
+            throw std::runtime_error("Unknown tracking algorithm: algo.track=" + track);
+        }
     }
 
 } // namespace impactx

src/initialization/AmrCoreData.H

@@ -93,6 +93,14 @@ namespace impactx::initialization
 
         } track_envelope;
 
+        /** Data needed for the reference orbit tracking mode */
+        struct TrackModeReference
+        {
+            /** The reference particle */
+            std::optional<RefPart> m_ref;
+
+        } track_reference;
+
         void ErrorEst (
             [[maybe_unused]] int lev,
             [[maybe_unused]] amrex::TagBoxArray& tags,

src/initialization/InitDistribution.cpp

@@ -467,5 +467,9 @@ namespace impactx
             auto dist = initialization::read_distribution(pp_dist);
             amr_data->track_envelope.m_cm = impactx::initialization::create_covariance_matrix(dist);
         }
+        else if (track == "reference_orbit")
+        {
+            amr_data->track_reference.m_ref = initialization::read_reference_particle(pp_dist);
+        }
     }
 } // namespace impactx

src/python/ImpactX.cpp

@@ -470,6 +470,9 @@ void init_ImpactX (py::module& m)
         .def("track_envelope", &ImpactX::track_envelope,
              "Run the envelope tracking simulation loop."
         )
+        .def("track_reference", &ImpactX::track_reference,
+             "Run the reference orbit tracking simulation loop."
+        )
 
         .def("resize_mesh", &ImpactX::ResizeMesh,
              "Resize the mesh :py:attr:`~domain` based on the :py:attr:`~dynamic_size` and related parameters."

src/tracking/CMakeLists.txt

@@ -2,4 +2,5 @@ target_sources(lib
   PRIVATE
     envelope.cpp
     particles.cpp
+    reference.cpp
 )

src/tracking/envelope.cpp

@@ -12,7 +12,6 @@
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/Push.H"
 #include "diagnostics/DiagnosticOutput.H"
-#include "particles/wakefields/HandleWakefield.H"
 
 #include <AMReX.H>
 #include <AMReX_AmrParGDB.H>
@@ -43,10 +42,12 @@ namespace impactx
         bool early_params_checked = false;
 
         // access beam data
-        if (!amr_data->track_envelope.m_ref.has_value()) {
+        if (!amr_data->track_envelope.m_ref.has_value())
+        {
             throw std::runtime_error("track_envelope: Reference particle not set.");
         }
-        if (!amr_data->track_envelope.m_cm.has_value()) {
+        if (!amr_data->track_envelope.m_cm.has_value())
+        {
             throw std::runtime_error("track_envelope: Envelope (covariance matrix) not set.");
         }
         auto & ref = amr_data->track_envelope.m_ref.value();
@@ -77,45 +78,53 @@ namespace impactx
         bool space_charge = false;
         pp_algo.query("space_charge", space_charge);
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!space_charge, "Space charge not yet implemented for envelope tracking.");
-        if (verbose > 0) {
+        if (verbose > 0)
+        {
             amrex::Print() << " Space Charge effects: " << space_charge << "\n";
         }
 
         bool csr = false;
         pp_algo.query("csr", csr);
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!csr, "CSR not yet implemented for envelope tracking.");
-        if (verbose > 0) {
+        if (verbose > 0)
+        {
             amrex::Print() << " CSR effects: " << csr << "\n";
         }
 
         // periods through the lattice
         int num_periods = 1;
         amrex::ParmParse("lattice").queryAddWithParser("periods", num_periods);
 
-        for (int period=0; period < num_periods; ++period) {
+        for (int period=0; period < num_periods; ++period)
+        {
             // loop over all beamline elements
-            for (auto &element_variant: m_lattice) {
+            for (auto &element_variant: m_lattice)
+            {
                 // update element edge of the reference particle
                 ref.sedge = ref.s;
 
                 // number of slices used for the application of space charge
                 int nslice = 1;
                 amrex::ParticleReal slice_ds; // in meters
-                std::visit([&nslice, &slice_ds](auto &&element) {
+                std::visit([&nslice, &slice_ds](auto &&element)
+                {
                     nslice = element.nslice();
                     slice_ds = element.ds() / nslice;
                 }, element_variant);
 
                 // sub-steps for space charge within the element
-                for (int slice_step = 0; slice_step < nslice; ++slice_step) {
-                    BL_PROFILE("ImpactX::evolve::slice_step");
+                for (int slice_step = 0; slice_step < nslice; ++slice_step)
+                {
+                    BL_PROFILE("ImpactX::track_envelope::slice_step");
                     step++;
-                    if (verbose > 0) {
+                    if (verbose > 0)
+                    {
                         amrex::Print() << " ++++ Starting step=" << step
                                        << " slice_step=" << slice_step << "\n";
                     }
 
-                    std::visit([&ref, &cm](auto&& element){
+                    std::visit([&ref, &cm](auto&& element)
+                    {
                         // push reference particle in global coordinates
                         {
                             BL_PROFILE("impactx::Push::RefPart");
@@ -128,7 +137,8 @@ namespace impactx
                     }, element_variant);
 
                     // just prints an empty newline at the end of the slice_step
-                    if (verbose > 0) {
+                    if (verbose > 0)
+                    {
                         amrex::Print() << "\n";
                     }
 
@@ -137,7 +147,8 @@ namespace impactx
                     pp_diag.queryAdd("slice_step_diagnostics", slice_step_diagnostics);
 
 
-                    if (diag_enable && slice_step_diagnostics) {
+                    if (diag_enable && slice_step_diagnostics)
+                    {
                         // print slice step reference particle to file
                         diagnostics::DiagnosticOutput(ref, "diags/ref_particle", step, true);
 
@@ -162,7 +173,6 @@ namespace impactx
 
             // print the final values of the reduced beam characteristics
             diagnostics::DiagnosticOutput(cm, ref, "diags/reduced_beam_characteristics_final", step);
-
         }
 
         // loop over all beamline elements & finalize them

src/tracking/reference.cpp

@@ -0,0 +1,135 @@
+/* Copyright 2022-2025 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl, Chad Mitchell
+ * License: BSD-3-Clause-LBNL
+ */
+#include "ImpactX.H"
+#include "initialization/InitAmrCore.H"
+#include "particles/ImpactXParticleContainer.H"
+#include "particles/Push.H"
+#include "diagnostics/DiagnosticOutput.H"
+
+#include <AMReX.H>
+#include <AMReX_AmrParGDB.H>
+#include <AMReX_BLProfiler.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_Print.H>
+
+#include <memory>
+
+
+namespace impactx
+{
+    void
+    ImpactX::track_reference (RefPart & ref)
+    {
+        BL_PROFILE("ImpactX::track_reference");
+
+        // verbosity
+        amrex::ParmParse pp_impactx("impactx");
+        int verbose = 1;
+        pp_impactx.queryAddWithParser("verbose", verbose);
+
+        // a global step for diagnostics including space charge slice steps in elements
+        //   before we start the evolve loop, we are in "step 0" (initial state)
+        int step = 0;
+
+        // check typos in inputs after step 1
+        bool early_params_checked = false;
+
+        // output of init state
+        amrex::ParmParse pp_diag("diag");
+        bool diag_enable = true;
+        pp_diag.queryAdd("enable", diag_enable);
+        if (verbose > 0)
+        {
+            amrex::Print() << " Diagnostics: " << diag_enable << "\n";
+        }
+
+        if (diag_enable)
+        {
+            int file_min_digits = 6;
+            pp_diag.queryAddWithParser("file_min_digits", file_min_digits);
+
+            // print initial reference particle to file
+            diagnostics::DiagnosticOutput(ref, "diags/ref_particle");
+
+        }
+
+        // periods through the lattice
+        int num_periods = 1;
+        amrex::ParmParse("lattice").queryAddWithParser("periods", num_periods);
+
+        for (int period=0; period < num_periods; ++period)
+        {
+            // loop over all beamline elements
+            for (auto &element_variant: m_lattice) {
+                // update element edge of the reference particle
+                ref.sedge = ref.s;
+
+                // number of slices through the element
+                int nslice = 1;
+                amrex::ParticleReal slice_ds; // in meters
+                std::visit([&nslice, &slice_ds](auto &&element)
+                {
+                    nslice = element.nslice();
+                    slice_ds = element.ds() / nslice;
+                }, element_variant);
+
+                // sub-steps within the element
+                for (int slice_step = 0; slice_step < nslice; ++slice_step)
+                {
+                    BL_PROFILE("ImpactX::track_reference::slice_step");
+                    step++;
+                    if (verbose > 0) {
+                        amrex::Print() << " ++++ Starting step=" << step
+                                       << " slice_step=" << slice_step << "\n";
+                    }
+
+                    std::visit([&ref](auto&& element)
+                    {
+                        // push reference particle in global coordinates
+                        BL_PROFILE("impactx::Push::RefPart");
+                        element(ref);
+                    }, element_variant);
+
+                    // just prints an empty newline at the end of the slice_step
+                    if (verbose > 0)
+                    {
+                        amrex::Print() << "\n";
+                    }
+
+                    // slice-step diagnostics
+                    bool slice_step_diagnostics = false;
+                    pp_diag.queryAdd("slice_step_diagnostics", slice_step_diagnostics);
+
+
+                    if (diag_enable && slice_step_diagnostics)
+                    {
+                        // print slice step reference particle to file
+                        diagnostics::DiagnosticOutput(ref, "diags/ref_particle", step, true);
+                    }
+
+                    // inputs: unused parameters (e.g. typos) check after step 1 has finished
+                    if (!early_params_checked) { early_params_checked = early_param_check(); }
+
+                } // end in-element slice-step loop
+
+            } // end beamline element loop
+
+        } // end periods though the lattice loop
+
+        if (diag_enable)
+        {
+            // print final reference particle to file
+            diagnostics::DiagnosticOutput(ref, "diags/ref_particle_final", step);
+        }
+
+        // loop over all beamline elements & finalize them
+        finalize_elements();
+    }
+} // namespace impactx