Add accelerator lattice, starting with quadrupoles (#3063)

* Initial version of accelerator lattice

* Clean up EOL white space

* Small clean up for GPU

* Fixed up consts

* Added hard edge fraction plus other clean ups

* More clean up

* Restructure to work on GPUs

* Now this grabs its own copies of particle info

* Updates, including adding dBdx

* Small cleanup in Quad

* Small fixes for GPU

* More cleanup for GPU

* More GPU cleanup

* Rewrite of the accelerator lattice implementation to better handle GPU

* Fix struct forward definition

* Another forward definition fix

* Bug fix

* Added LatticeElementBase

* Removed zcenters array

* Added CI test case

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Clean up in CI analysis.py

* Cleanup of coding

* Added CI test hard_edged_quadrupoles_moving

* Added Lorentz transform between boosted frame and lab frame

* Fixes for working in the boosted frame

* Added boosted CI test

* Change input name, adding the prefix "lattice."

* Added plasma lens lattice element

This will replace the external field plasma lens

* Fixed CI analysis script to look for "lattice.quad"

* Added checks of lattice element input

* Added documentation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Removed duplicate call to lattice finder UpdateIndices

* Added extensive comments

* Reworked the input to use the MAD like description

This is the same as the method used in ImpactX

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove old lines from inputs_lattice_3d

* Added "lattice" element type

* Fixed some Real and ParticleReals

* [pre-commit.ci] pre-commit autoupdate (#3246)

updates:
- [github.com/hadialqattan/pycln: v2.0.1 → v2.0.3](https://github.com/hadialqattan/pycln/compare/v2.0.1...v2.0.3)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* ABLASTR: Update Poisson Solver API (#3243)

Update the Poisson Solver API to be more usable. Needed for ImpactX.

* Docs: New OLCF Machine (#3228)

* D-T fusion (#3153)

* initial work

* fixed bugs and added species

* update documentation

* delete unused file

* Add properties for neutron, hydrogen isotopes, helium isotopes

* Update code to be more consistent

* Correct typo

* Parse deuterium-tritium fusion

* Start putting in place the files for deuterium-tritium

* Update documentation

* Prepare structures for deuterium tritium

* Fix typo

* Fix compilation

* Add neutron

* Add correct formula for the cross-section

* Correct compilation error

* Fix nuclear fusion

* Reset benchmarks

* Prepare creation functor for 2-product fusion

* First implementation of momentum initialization

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use utility function for fusion

* Minor modification of variable names

* Fix GPU compilation

* Fix single precision compilation

* Update types

* Use util function in P-B fusion

* Correct compilation errors

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Correct errors

* Update values of mass and charge

* Correct compilation error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Correct compilation error

* Correct compilation error

* Correct compilation error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reset benchmark

* Use helium particle in proton-boron, to avoid resetting benchmark

* Fixed proton-boron test

* Revert "Fixed proton-boron test"

This reverts commit 73c8d9d0be8417d5cd08a23daeebbc322c984808.

* Incorporate Neil's recommendations

* Reset benchmarks

* Correct compilation errors

* Add new deuterium tritium automated test

* Correct formula of cross-section

* Correct cross-section

* Improve analysis script

* Add test of energy conservation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add test of conservation of momentum

* Progress in analysis script

* Fix error in the initial energy of the deuterium particles

* Add check of isotropy

* Clean up the test script

* Rewrite p_sq formula in a way to avoids machine-precision negative numbers

* Add checksum

* Clean up code

* Apply suggestions from code review

* Update PR according to comments

* Update benchmark

* Address additional comments

* Numerical Literals

Co-authored-by: Luca Fedeli <luca.fedeli@cea.fr>
Co-authored-by: Neïl Zaim <49716072+NeilZaim@users.noreply.github.com>

* Docs: gaussian beam `q_tot` is not optional (#3249)

* Fix a bug in GPU version of Hankel Transform (#3253)

amrex::Array4 is a 4D array that can be accessed with three spatial indices
plus an optional component index.  We must always provide all three spatial
indices even in 2D.

* Add Python Callback Call when Checkpointing Signal is Received (#3251)

* CI: Add Missing Regression Analysis (NCI corrector) (#3252)

* Fixes to allow mixed precision, ParticleReal float, Real double (#3239)

* Fixes to allow mixed precision, ParticleReal float, Real double

* Fix for the optical depth

* A different way of fixing QuantumSynchrotronEvolveOpticalDepth

* In the QED code, consistently use ParticleReal

* Use ParticleReal type consistently

* Fix typo Docs/source/usage/parameters.rst

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Fix typo Docs/source/usage/parameters.rst

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Fix small error in plasma lens lattice documentation

* Small addition to the documentation

* Fix the residence correction to allow short elements

* Updated CI benchmarks

* Added check of lattice to isNoOp

* Updated the hard_edged_quadrupoles CI benchmarks

It is not clear why there was a change, but the difference is
essentially round off in the E field. The important thing is that the
particles are still correct.

* Update Source/AcceleratorLattice/AcceleratorLattice.H

Add include statements

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Update Source/AcceleratorLattice/LatticeElements/LatticeElementBase.H

Add includes

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Renamed to README.rst and updated headers

* Made d_lattice_element_finder optional type

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Include `<optional>`

* Docs: Developer AccLattice Inclusion

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
Co-authored-by: Luca Fedeli <luca.fedeli@cea.fr>
Co-authored-by: Neïl Zaim <49716072+NeilZaim@users.noreply.github.com>
Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>
Co-authored-by: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>
Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

1 files changed, 93 insertions, 0 deletions

diff --git a/Source/AcceleratorLattice/AcceleratorLattice.cpp b/Source/AcceleratorLattice/AcceleratorLattice.cpp
new file mode 100644
index 000000000..1c22e6244
--- /dev/null
+++ b/Source/AcceleratorLattice/AcceleratorLattice.cpp
@@ -0,0 +1,93 @@
+/* Copyright 2022 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#include "AcceleratorLattice.H"
+#include "LatticeElements/Drift.H"
+#include "LatticeElements/HardEdgedQuadrupole.H"
+#include "LatticeElements/HardEdgedPlasmaLens.H"
+
+#include <AMReX_REAL.H>
+
+AcceleratorLattice::AcceleratorLattice ()
+{
+
+    using namespace amrex::literals;
+
+    /* Get the inputs for and initialize all of the lattice element types */
+    amrex::ParticleReal z_location = 0._prt;
+    ReadLattice("lattice", z_location);
+
+    h_quad.WriteToDevice();
+    h_plasmalens.WriteToDevice();
+}
+
+void
+AcceleratorLattice::ReadLattice (std::string const & root_name, amrex::ParticleReal & z_location)
+{
+    amrex::ParmParse pp_lattice(root_name);
+    std::vector<std::string> lattice_elements;
+    pp_lattice.queryarr("elements", lattice_elements);
+
+    if (!lattice_elements.empty()) {
+        m_lattice_defined = true;
+    }
+
+    // Loop through lattice elements
+    for (std::string const & element_name : lattice_elements) {
+        // Check the element type
+        amrex::ParmParse pp_element(element_name);
+        std::string element_type;
+        pp_element.get("type", element_type);
+
+        // Initialize the corresponding element according to its type
+        if (element_type == "drift") {
+            h_drift.AddElement(pp_element, z_location);
+        }
+        else if (element_type == "quad") {
+            h_quad.AddElement(pp_element, z_location);
+        }
+        else if (element_type == "plasmalens") {
+            h_plasmalens.AddElement(pp_element, z_location);
+        }
+        else if (element_type == "lattice") {
+            ReadLattice(element_name, z_location);
+        }
+        else {
+            WARPX_ALWAYS_ASSERT_WITH_MESSAGE(false,
+                "ERROR: Unknown accelerator lattice element type " + element_type);
+        }
+    }
+}
+
+void
+AcceleratorLattice::InitElementFinder (int const lev, amrex::BoxArray const & ba, amrex::DistributionMapping const & dm)
+{
+    if (m_lattice_defined) {
+        m_element_finder = std::make_unique<amrex::LayoutData<LatticeElementFinder>>(ba, dm);
+        for (amrex::MFIter mfi(*m_element_finder); mfi.isValid(); ++mfi)
+        {
+            (*m_element_finder)[mfi].InitElementFinder(lev, mfi, *this);
+        }
+    }
+}
+
+void
+AcceleratorLattice::UpdateElementFinder (int const lev)
+{
+    if (m_lattice_defined) {
+        for (amrex::MFIter mfi(*m_element_finder); mfi.isValid(); ++mfi)
+        {
+            (*m_element_finder)[mfi].UpdateIndices(lev, mfi, *this);
+        }
+    }
+}
+
+LatticeElementFinderDevice
+AcceleratorLattice::GetFinderDeviceInstance (WarpXParIter const& a_pti, int const a_offset) const
+{
+    LatticeElementFinder & finder = (*m_element_finder)[a_pti];
+    return finder.GetFinderDeviceInstance(a_pti, a_offset, *this);
+}