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, 123 insertions, 0 deletions

diff --git a/Source/AcceleratorLattice/LatticeElementFinder.cpp b/Source/AcceleratorLattice/LatticeElementFinder.cpp
new file mode 100644
index 000000000..14ccb8f71
--- /dev/null
+++ b/Source/AcceleratorLattice/LatticeElementFinder.cpp
@@ -0,0 +1,123 @@
+/* Copyright 2022 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#include "WarpX.H"
+#include "LatticeElementFinder.H"
+#include "LatticeElements/HardEdgedQuadrupole.H"
+#include "LatticeElements/HardEdgedPlasmaLens.H"
+
+#include <AMReX_ParmParse.H>
+#include <AMReX_REAL.H>
+
+using namespace amrex::literals;
+
+void
+LatticeElementFinder::InitElementFinder (int const lev, amrex::MFIter const& a_mfi,
+                                         AcceleratorLattice const& accelerator_lattice)
+{
+
+    // The lattice is assumed to extend in the z-direction
+    // Get the number of nodes where indices will be setup
+    amrex::Box box = a_mfi.tilebox();
+    m_nz = box.size()[WARPX_ZINDEX];
+
+    m_dz = WarpX::CellSize(lev)[2];
+
+    m_gamma_boost = WarpX::gamma_boost;
+    m_uz_boost = std::sqrt(WarpX::gamma_boost*WarpX::gamma_boost - 1._prt)*PhysConst::c;
+
+    AllocateIndices(accelerator_lattice);
+
+    UpdateIndices(lev, a_mfi, accelerator_lattice);
+
+}
+
+void
+LatticeElementFinder::AllocateIndices (AcceleratorLattice const& accelerator_lattice)
+{
+    // Allocate the space for the indices for each element type.
+    // Note that this uses m_nz since the information is saved per node.
+
+    if (accelerator_lattice.h_quad.nelements > 0) {
+        d_quad_indices.resize(m_nz);
+    }
+
+    if (accelerator_lattice.h_plasmalens.nelements > 0) {
+        d_plasmalens_indices.resize(m_nz);
+    }
+}
+
+void
+LatticeElementFinder::UpdateIndices (int const lev, amrex::MFIter const& a_mfi,
+                                     AcceleratorLattice const& accelerator_lattice)
+{
+    auto& warpx = WarpX::GetInstance();
+
+    // Update the location of the index grid.
+    // Note that the current box is used since the box may have been updated since
+    // the initialization in InitElementFinder.
+    amrex::Box box = a_mfi.tilebox();
+    m_zmin = WarpX::LowerCorner(box, lev, 0._rt)[2];
+    m_time = warpx.gett_new(lev);
+
+    if (accelerator_lattice.h_quad.nelements > 0) {
+        setup_lattice_indices(accelerator_lattice.h_quad.d_zs,
+                              accelerator_lattice.h_quad.d_ze,
+                              d_quad_indices);
+    }
+
+    if (accelerator_lattice.h_plasmalens.nelements > 0) {
+        setup_lattice_indices(accelerator_lattice.h_plasmalens.d_zs,
+                              accelerator_lattice.h_plasmalens.d_ze,
+                              d_plasmalens_indices);
+    }
+}
+
+LatticeElementFinderDevice
+LatticeElementFinder::GetFinderDeviceInstance (WarpXParIter const& a_pti, int const a_offset,
+                                              AcceleratorLattice const& accelerator_lattice)
+{
+    LatticeElementFinderDevice result;
+    result.InitLatticeElementFinderDevice(a_pti, a_offset, accelerator_lattice, *this);
+    return result;
+}
+
+
+void
+LatticeElementFinderDevice::InitLatticeElementFinderDevice (WarpXParIter const& a_pti, int const a_offset,
+                                                            AcceleratorLattice const& accelerator_lattice,
+                                                            LatticeElementFinder const & h_finder)
+{
+
+    auto& warpx = WarpX::GetInstance();
+
+    int const lev = a_pti.GetLevel();
+
+    m_get_position = GetParticlePosition(a_pti, a_offset);
+    auto& attribs = a_pti.GetAttribs();
+    m_ux = attribs[PIdx::ux].dataPtr() + a_offset;
+    m_uy = attribs[PIdx::uy].dataPtr() + a_offset;
+    m_uz = attribs[PIdx::uz].dataPtr() + a_offset;
+    m_dt = warpx.getdt(lev);
+
+    m_gamma_boost = WarpX::gamma_boost;
+    m_uz_boost = std::sqrt(WarpX::gamma_boost*WarpX::gamma_boost - 1._prt)*PhysConst::c;
+
+    m_zmin = h_finder.m_zmin;
+    m_dz = h_finder.m_dz;
+    m_time = h_finder.m_time;
+
+    if (accelerator_lattice.h_quad.nelements > 0) {
+        d_quad = accelerator_lattice.h_quad.GetDeviceInstance();
+        d_quad_indices_arr = h_finder.d_quad_indices.data();
+    }
+
+    if (accelerator_lattice.h_plasmalens.nelements > 0) {
+        d_plasmalens = accelerator_lattice.h_plasmalens.GetDeviceInstance();
+        d_plasmalens_indices_arr = h_finder.d_plasmalens_indices.data();
+    }
+
+}