diff --git a/CMakeLists.txt b/CMakeLists.txt index 54c104885..d29fd1611 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,6 +39,7 @@ option(IPPL_ENABLE_FFT "Enable FFT support" OFF) option(IPPL_ENABLE_SOLVERS "Enable IPPL solvers" OFF) option(IPPL_ENABLE_ALPINE "Enable building the Alpine module" OFF) option(IPPL_ENABLE_COSMOLOGY "Enable building the Cosmology module" OFF) +option(IPPL_ENABLE_FEL "Enable building the FEL module" OFF) option(IPPL_ENABLE_EXAMPLES "Enable building the Example module" OFF) option(IPPL_ENABLE_TESTS "Build integration tests in test/ directory" OFF) option(IPPL_ENABLE_COVERAGE "Enable code coverage" OFF) @@ -131,6 +132,10 @@ if(IPPL_ENABLE_COSMOLOGY) add_subdirectory(cosmology) endif() +if(IPPL_ENABLE_FEL) + add_subdirectory(fel) +endif() + if(IPPL_ENABLE_EXAMPLES) add_subdirectory(examples) endif() diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index a463d286e..b858f739f 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -319,3 +319,18 @@ if(IPPL_ENABLE_TESTS) "${DOWNLOADED_HEADERS_DIR}/stb_image_write.h") message(STATUS "✅ stb_image_write loaded for testing FDTD solver.") endif() + +# ------------------------------------------------------------------------------ +# FEL module header-only dependencies (nlohmann/json for config parsing, +# stb_image_write for the Poynting-flux visualization). +# ------------------------------------------------------------------------------ +if(IPPL_ENABLE_FEL) + set(DOWNLOADED_HEADERS_DIR "${CMAKE_CURRENT_BINARY_DIR}/downloaded_headers") + file(DOWNLOAD + https://github.com/nlohmann/json/releases/download/v3.11.3/json.hpp + "${DOWNLOADED_HEADERS_DIR}/json.hpp") + file(DOWNLOAD + https://raw.githubusercontent.com/manuel5975p/stb/master/stb_image_write.h + "${DOWNLOADED_HEADERS_DIR}/stb_image_write.hpp") + message(STATUS "✅ nlohmann/json and stb_image_write loaded for the FEL module.") +endif() diff --git a/fel/CMakeLists.txt b/fel/CMakeLists.txt new file mode 100644 index 000000000..541ace2c9 --- /dev/null +++ b/fel/CMakeLists.txt @@ -0,0 +1,27 @@ +file (RELATIVE_PATH _relPath "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}") +message (STATUS "Adding FEL module found in ${_relPath}") + +include_directories ( + ${CMAKE_SOURCE_DIR}/src + ${CMAKE_BINARY_DIR}/downloaded_headers +) + +link_directories ( + ${CMAKE_CURRENT_SOURCE_DIR} + ${Kokkos_DIR}/.. +) + +set (IPPL_LIBS ippl ${MPI_CXX_LIBRARIES}) +set (COMPILE_FLAGS ${OPAL_CXX_FLAGS}) + +add_executable (FreeElectronLaser FreeElectronLaser.cpp) +target_link_libraries (FreeElectronLaser ${IPPL_LIBS}) + +# vi: set et ts=4 sw=4 sts=4: + +# Local Variables: +# mode: cmake +# cmake-tab-width: 4 +# indent-tabs-mode: nil +# require-final-newline: nil +# End: diff --git a/fel/Config.h b/fel/Config.h new file mode 100644 index 000000000..0252507ea --- /dev/null +++ b/fel/Config.h @@ -0,0 +1,272 @@ +#ifndef IPPL_FEL_CONFIG_H +#define IPPL_FEL_CONFIG_H + +// Configuration parsing for the FEL simulation. +// +// Reads a MITHRA-style JSON job file and converts all physical quantities into +// the natural unit system defined in units.h. Ported from the original +// FreeElectronLaser.cpp. + +#include +#include +#include +#include +#include +#include + +#include "Types/Vector.h" + +#include "units.h" + +#define JSON_HAS_RANGES 0 +#include + +struct config { + using scalar = double; + + // GRID PARAMETERS + ippl::Vector resolution; // Grid resolution in 3D + ippl::Vector extents; // Physical extents of the grid in each dimension + scalar total_time; // Total simulation time + scalar timestep_ratio; // Ratio of timestep to some reference value + + scalar length_scale_in_jobfile; // Length scale defined in the jobfile + scalar temporal_scale_in_jobfile; // Temporal scale defined in the jobfile + + // PARTICLE PARAMETERS + scalar charge; // Particle charge in unit_charge + scalar mass; // Particle mass in unit_mass + uint64_t num_particles; // Number of particles in the simulation + bool space_charge; // Flag for considering space charge effects + + // BUNCH PARAMETERS + ippl::Vector mean_position; // Mean initial position of the particle bunch + ippl::Vector sigma_position; // Standard deviation of the initial position distribution + ippl::Vector position_truncations; // Truncations of the position distribution + ippl::Vector sigma_momentum; // Standard deviation of the initial momentum distribution + scalar bunch_gamma; // Relativistic gamma factor of the bunch + + // UNDULATOR PARAMETERS + scalar undulator_K; // Undulator parameter K + scalar undulator_period; // Period of the undulator + scalar undulator_length; // Length of the undulator + + uint32_t output_rhythm; // Frequency of output in timesteps + std::string output_path; // Path to output files + std::unordered_map experiment_options; // Additional experimental options +}; + +template +ippl::Vector getVector(const nlohmann::json& j) { + if (j.is_array()) { + assert(j.size() == Dim); + ippl::Vector ret; + for (unsigned i = 0; i < Dim; i++) + ret[i] = (scalar)j[i]; + return ret; + } else { + std::cerr << "Warning: Obtaining Vector from scalar json\n"; + ippl::Vector ret = (scalar)j; + return ret; + } +} + +// Compile-time / run-time string hashing used to switch on unit-scale names. +template +struct DefaultedStringLiteral { + constexpr DefaultedStringLiteral(const char (&str)[N], const T val) + : value(val) { + std::copy_n(str, N, key); + } + + T value; + char key[N]; +}; +template +struct StringLiteral { + constexpr StringLiteral(const char (&str)[N]) { std::copy_n(str, N, value); } + + char value[N]; + constexpr DefaultedStringLiteral operator>>(int t) const noexcept { + return DefaultedStringLiteral(value, t); + } + constexpr size_t size() const noexcept { return N - 1; } +}; +template +constexpr size_t chash() { + size_t hash = 5381; + int c; + + for (size_t i = 0; i < lit.size(); i++) { + c = lit.value[i]; + hash = ((hash << 5) + hash) + c; // hash * 33 + c + } + + return hash; +} +inline size_t chash(const char* val) { + size_t hash = 5381; + int c; + + while ((c = *val++)) { + hash = ((hash << 5) + hash) + c; // hash * 33 + c + } + + return hash; +} +inline size_t chash(const std::string& _val) { + size_t hash = 5381; + const char* val = _val.c_str(); + int c; + + while ((c = *val++)) { + hash = ((hash << 5) + hash) + c; // hash * 33 + c + } + + return hash; +} +inline std::string lowercase_singular(std::string str) { + // Convert string to lowercase + std::transform(str.begin(), str.end(), str.begin(), ::tolower); + + // Check if the string ends with "s" and remove it if it does + if (!str.empty() && str.back() == 's') { + str.pop_back(); + } + + return str; +} +inline double get_time_multiplier(const nlohmann::json& j) { + std::string length_scale_string = lowercase_singular((std::string)j["mesh"]["time-scale"]); + double time_factor = 1.0; + switch (chash(length_scale_string)) { + case chash<"planck-time">(): + case chash<"plancktime">(): + case chash<"pt">(): + case chash<"natural">(): + time_factor = unit_time_in_seconds; + break; + case chash<"picosecond">(): + time_factor = 1e-12; + break; + case chash<"nanosecond">(): + time_factor = 1e-9; + break; + case chash<"microsecond">(): + time_factor = 1e-6; + break; + case chash<"millisecond">(): + time_factor = 1e-3; + break; + case chash<"second">(): + time_factor = 1.0; + break; + default: + std::cerr << "Unrecognized time scale: " << (std::string)j["mesh"]["time-scale"] + << "\n"; + break; + } + return time_factor; +} +inline double get_length_multiplier(const nlohmann::json& options) { + std::string length_scale_string = + lowercase_singular((std::string)options["mesh"]["length-scale"]); + double length_factor = 1.0; + switch (chash(length_scale_string)) { + case chash<"planck-length">(): + case chash<"plancklength">(): + case chash<"pl">(): + case chash<"natural">(): + length_factor = unit_length_in_meters; + break; + case chash<"picometer">(): + length_factor = 1e-12; + break; + case chash<"nanometer">(): + length_factor = 1e-9; + break; + case chash<"micrometer">(): + length_factor = 1e-6; + break; + case chash<"millimeter">(): + length_factor = 1e-3; + break; + case chash<"meter">(): + length_factor = 1.0; + break; + default: + std::cerr << "Unrecognized length scale: " + << (std::string)options["mesh"]["length-scale"] << "\n"; + break; + } + return length_factor; +} +inline config read_config(const char* filepath) { + std::ifstream cfile(filepath); + nlohmann::json j; + cfile >> j; + config::scalar lmult = get_length_multiplier(j); + config::scalar tmult = get_time_multiplier(j); + config ret; + + ret.extents[0] = ((config::scalar)j["mesh"]["extents"][0] * lmult) / unit_length_in_meters; + ret.extents[1] = ((config::scalar)j["mesh"]["extents"][1] * lmult) / unit_length_in_meters; + ret.extents[2] = ((config::scalar)j["mesh"]["extents"][2] * lmult) / unit_length_in_meters; + ret.resolution = getVector(j["mesh"]["resolution"]); + + if (j.contains("timestep-ratio")) { + ret.timestep_ratio = (config::scalar)j["timestep-ratio"]; + } else { + ret.timestep_ratio = 1; + } + ret.total_time = ((config::scalar)j["mesh"]["total-time"] * tmult) / unit_time_in_seconds; + ret.space_charge = (bool)(j["mesh"]["space-charge"]); + ret.bunch_gamma = (config::scalar)(j["bunch"]["gamma"]); + if (ret.bunch_gamma < config::scalar(1)) { + std::cerr << "Gamma must be >= 1\n"; + exit(1); + } + assert(j.contains("undulator")); + assert(j["undulator"].contains("static-undulator")); + + ret.undulator_K = j["undulator"]["static-undulator"]["undulator-parameter"]; + ret.undulator_period = ((config::scalar)j["undulator"]["static-undulator"]["period"] * lmult) + / unit_length_in_meters; + ret.undulator_length = ((config::scalar)j["undulator"]["static-undulator"]["length"] * lmult) + / unit_length_in_meters; + assert(!std::isnan(ret.undulator_length)); + assert(!std::isnan(ret.undulator_period)); + assert(!std::isnan(ret.extents[0])); + assert(!std::isnan(ret.extents[1])); + assert(!std::isnan(ret.extents[2])); + assert(!std::isnan(ret.total_time)); + ret.length_scale_in_jobfile = get_length_multiplier(j); + ret.temporal_scale_in_jobfile = get_time_multiplier(j); + ret.charge = (config::scalar)j["bunch"]["charge"] * electron_charge_in_unit_charges; + ret.mass = (config::scalar)j["bunch"]["mass"] * electron_mass_in_unit_masses; + ret.num_particles = (uint64_t)j["bunch"]["number-of-particles"]; + ret.mean_position = + getVector(j["bunch"]["position"]) * lmult / unit_length_in_meters; + ret.sigma_position = + getVector(j["bunch"]["sigma-position"]) * lmult / unit_length_in_meters; + ret.position_truncations = getVector(j["bunch"]["distribution-truncations"]) + * lmult / unit_length_in_meters; + ret.sigma_momentum = getVector(j["bunch"]["sigma-momentum"]); + ret.output_rhythm = j["output"].contains("rhythm") ? uint32_t(j["output"]["rhythm"]) : 0; + ret.output_path = "../data/"; + if (j["output"].contains("path")) { + ret.output_path = j["output"]["path"]; + if (!ret.output_path.ends_with('/')) { + ret.output_path.push_back('/'); + } + } + if (j.contains("experimentation")) { + nlohmann::json je = j["experimentation"]; + for (auto it = je.begin(); it != je.end(); it++) { + ret.experiment_options[it.key()] = double(it.value()); + } + } + return ret; +} + +#endif diff --git a/fel/FELFieldContainer.hpp b/fel/FELFieldContainer.hpp new file mode 100644 index 000000000..a12da3154 --- /dev/null +++ b/fel/FELFieldContainer.hpp @@ -0,0 +1,79 @@ +#ifndef IPPL_FEL_FIELD_CONTAINER_H +#define IPPL_FEL_FIELD_CONTAINER_H + +#include + +#include "Manager/BaseManager.h" + +#include "datatypes.h" + +// Define the FELFieldContainer class. +// +// Mirrors alpine/FieldContainer.hpp but holds the fields needed for the +// electromagnetic FEL simulation: the electric field E, the magnetic field B +// (both three-component vector fields), and the four-current source J +// ([0] = charge density, [1..Dim] = current density) that drives the FDTD +// solver. Unlike the alpine container the field layout is not periodic; the +// decomposition flags select which axes are split across MPI ranks. +template +class FELFieldContainer { +public: + FELFieldContainer(Vector_t& hr, Vector_t& rmin, Vector_t& rmax, + std::array decomp, ippl::NDIndex domain, + Vector_t origin, bool isAllPeriodic) + : hr_m(hr) + , rmin_m(rmin) + , rmax_m(rmax) + , decomp_m(decomp) + , mesh_m(domain, hr, origin) + , fl_m(MPI_COMM_WORLD, domain, decomp, isAllPeriodic) {} + + ~FELFieldContainer() {} + +private: + Vector_t hr_m; + Vector_t rmin_m; + Vector_t rmax_m; + std::array decomp_m; + VField_t E_m; + VField_t B_m; + SourceField_t J_m; + Mesh_t mesh_m; + FieldLayout_t fl_m; + +public: + VField_t& getE() { return E_m; } + void setE(VField_t& E) { E_m = E; } + + VField_t& getB() { return B_m; } + void setB(VField_t& B) { B_m = B; } + + SourceField_t& getJ() { return J_m; } + void setJ(SourceField_t& J) { J_m = J; } + + Vector_t& getHr() { return hr_m; } + void setHr(const Vector_t& hr) { hr_m = hr; } + + Vector_t& getRMin() { return rmin_m; } + void setRMin(const Vector_t& rmin) { rmin_m = rmin; } + + Vector_t& getRMax() { return rmax_m; } + void setRMax(const Vector_t& rmax) { rmax_m = rmax; } + + std::array getDecomp() { return decomp_m; } + void setDecomp(std::array decomp) { decomp_m = decomp; } + + Mesh_t& getMesh() { return mesh_m; } + void setMesh(Mesh_t& mesh) { mesh_m = mesh; } + + FieldLayout_t& getFL() { return fl_m; } + void setFL(FieldLayout_t& fl) { fl_m = fl; } + + void initializeFields() { + E_m.initialize(mesh_m, fl_m); + B_m.initialize(mesh_m, fl_m); + J_m.initialize(mesh_m, fl_m); + } +}; + +#endif diff --git a/fel/FELParticleContainer.hpp b/fel/FELParticleContainer.hpp new file mode 100644 index 000000000..802bc9c4b --- /dev/null +++ b/fel/FELParticleContainer.hpp @@ -0,0 +1,68 @@ +#ifndef IPPL_FEL_PARTICLE_CONTAINER_H +#define IPPL_FEL_PARTICLE_CONTAINER_H + +#include + +#include "Manager/BaseManager.h" + +#include "datatypes.h" + +// Define the FELParticleContainer class. +// +// Mirrors alpine/ParticleContainer.hpp but carries the attributes required by a +// relativistic electromagnetic push: +// * Q, mass : per-particle charge and mass. +// * gamma_beta : relativistic momentum (gamma * beta). +// * R_nm1, R_np1 : positions at the previous / next step. R_nm1 is the +// start point for the charge-conserving current +// deposition (X0 -> X1 = R_nm1 -> R). +// * E_gather, B_gather : E and B interpolated to the particle positions. +// +// Boundary conditions: particles leaving the domain are removed +// explicitly by the manager rather than wrapped or reflected. +template +class FELParticleContainer : public ippl::ParticleBase> { + using Base = ippl::ParticleBase>; + +public: + ippl::ParticleAttrib Q; // charge + ippl::ParticleAttrib mass; // mass + ippl::ParticleAttrib> gamma_beta; // relativistic momentum (gamma*beta) + typename Base::particle_position_type R_nm1; // position at previous step + typename Base::particle_position_type R_np1; // position at next step + ippl::ParticleAttrib> E_gather; // E field at particle position + ippl::ParticleAttrib> B_gather; // B field at particle position + +private: + PLayout_t pl_m; + +public: + FELParticleContainer(Mesh_t& mesh, FieldLayout_t& FL) + : pl_m(FL, mesh) { + this->initialize(pl_m); + registerAttributes(); + setupBCs(); + } + + ~FELParticleContainer() {} + + std::shared_ptr> getPL() { return pl_m; } + void setPL(std::shared_ptr>& pl) { pl_m = pl; } + + void registerAttributes() { + this->addAttribute(Q); + this->addAttribute(mass); + this->addAttribute(gamma_beta); + this->addAttribute(R_nm1); + this->addAttribute(R_np1); + this->addAttribute(E_gather); + this->addAttribute(B_gather); + } + + void setupBCs() { setBCAllOpen(); } + +private: + void setBCAllOpen() { this->setParticleBC(ippl::NO); } +}; + +#endif diff --git a/fel/FreeElectronLaser.cpp b/fel/FreeElectronLaser.cpp new file mode 100644 index 000000000..22e3d4406 --- /dev/null +++ b/fel/FreeElectronLaser.cpp @@ -0,0 +1,70 @@ +// Free Electron Laser simulation. +// +// Usage: +// srun ./FreeElectronLaser [] --info [0-5] +// or +// mpirun -np [N] ./FreeElectronLaser [] --info [0-5] +// +// Reads a MITHRA-style JSON job file (default: ../fel/config.json) describing the +// grid, the relativistic electron bunch, and the undulator. The simulation +// runs in a Lorentz frame co-moving with the bunch: a charge-conserving +// current is deposited onto the grid, Maxwell's equations are advanced with a +// standard FDTD solver (absorbing boundaries), and the particles are pushed +// with a relativistic Boris pusher that also feels the (frame-transformed) +// undulator field. Radiated power is written to a CSV and a +// Poynting-flux video is produced via ffmpeg. + +constexpr unsigned Dim = 3; +using T = double; + +#include "Ippl.h" + +#include +#include + +#include "Utility/IpplTimings.h" + +// stb_image_write's implementation must be emitted in exactly one translation +// unit; define the macro before the (transitive) include of the header. +#define STB_IMAGE_WRITE_IMPLEMENTATION +#include "FreeElectronLaserManager.h" + +int main(int argc, char* argv[]) { + ippl::initialize(argc, argv); + { + Inform msg("FreeElectronLaser"); + + static IpplTimings::TimerRef mainTimer = IpplTimings::getTimer("total"); + IpplTimings::startTimer(mainTimer); + + // First positional argument (if any, and not an --option) is the config + // file path; otherwise fall back to the shipped example config. The + // path is relative to the working directory, so this default assumes + // the program is launched from the build directory. + const char* config_path = "../fel/config.json"; + if (argc > 1 && argv[1][0] != '-') { + config_path = argv[1]; + } + msg << "Reading configuration from " << config_path << endl; + config cfg = read_config(config_path); + + // Create the manager for the FEL application. + FreeElectronLaserManager manager(cfg); + + // Pre-run: build mesh, fields, particles, FDTD solver; derive dt and nt. + manager.pre_run(); + + msg << "Starting iterations ..." << endl; + + manager.run(manager.getNt()); + + msg << "End." << endl; + + IpplTimings::stopTimer(mainTimer); + IpplTimings::print(); + IpplTimings::print(std::string("timing.dat")); + } + ippl::finalize(); + + return 0; +} diff --git a/fel/FreeElectronLaserManager.h b/fel/FreeElectronLaserManager.h new file mode 100644 index 000000000..9e884c420 --- /dev/null +++ b/fel/FreeElectronLaserManager.h @@ -0,0 +1,732 @@ +#ifndef IPPL_FREE_ELECTRON_LASER_MANAGER_H +#define IPPL_FREE_ELECTRON_LASER_MANAGER_H + +#include +#include +#include +#include + +#include "Manager/BaseManager.h" + +#include "Interpolation/CurrentDeposition.hpp" + +#include "Config.h" +#include "FELFieldContainer.hpp" +#include "FELParticleContainer.hpp" +#include "LorentzTransform.h" +#include "MithraBunch.h" +#include "Undulator.h" +#include "VideoWriter.h" +#include "datatypes.h" +#include "units.h" + +// FEL simulation manager. +// +// An electromagnetic (FDTD) PIC manager for the Free Electron Laser: it owns the +// field/particle containers and the Maxwell solver, provides the generic +// particle<->grid operations (deposit / gather / relativistic Boris push), and +// implements the FEL-specific run loop and diagnostics. +// +// It works in a Lorentz frame co-moving with the electron bunch (gamma_frame = +// gamma_bunch / sqrt(1 + K^2/2)); the static undulator field is transformed into +// this frame each step and added to the self-consistent FDTD field during the +// particle push. Radiation is transformed back to the lab frame for output. +template +class FreeElectronLaserManager : public ippl::BaseManager { +public: + using ParticleContainer_t = FELParticleContainer; + using FieldContainer_t = FELFieldContainer; + using FDTDSolver_t = ::FDTDSolver_t; + using Base = ippl::ParticleBase>; + + FreeElectronLaserManager(config cfg) + : m_config(cfg) + , totalP_m(cfg.num_particles) + , nt_m(0) + , time_m(0.0) + , dt_m(0.0) + , it_m(0) + , frame_gamma_m(std::max( + T(1), cfg.bunch_gamma + / std::sqrt(1 + cfg.undulator_K * cfg.undulator_K * T(0.5)))) + , uparams_m(cfg.undulator_K, cfg.undulator_period, cfg.undulator_length) + , frame_m(ippl::UniaxialLorentzframe::from_gamma(frame_gamma_m)) + , undulator_m(uparams_m, 2.0 * cfg.sigma_position[2] * frame_gamma_m * frame_gamma_m) {} + + ~FreeElectronLaserManager() { video_m.close(); } + +protected: + config m_config; + + size_type totalP_m; + int nt_m; + Vector_t nr_m; + + double time_m; + double dt_m; + int it_m; + + Vector_t rmin_m; + Vector_t rmax_m; + Vector_t hr_m; + Vector_t origin_m; + ippl::NDIndex domain_m; + std::array decomp_m; + bool isAllPeriodic_m; + + std::shared_ptr fcontainer_m; + std::shared_ptr pcontainer_m; + std::shared_ptr solver_m; + + int nsubsteps_m = 3; ///< Boris sub-steps per FDTD step (reference behaviour). + + T frame_gamma_m; ///< Lorentz factor of the co-moving frame. + ippl::undulator_parameters uparams_m; ///< Undulator parameters. + ippl::UniaxialLorentzframe frame_m; ///< Boost into the co-moving frame (z-axis). + ippl::Undulator undulator_m; ///< Static undulator field model. + FELVideoWriter video_m; ///< Optional ffmpeg Poynting-flux video. + + // --- narrow-band (resonant) radiation power diagnostic state --- + // MITHRA reports the FEL output power as a sliding-window single-frequency + // DFT of the exit-plane fields at the resonant wavelength, not the total + // broadband Poynting flux that dumpRadiation() integrates. These hold the + // rolling time-domain sample buffer and the derived window length / angular + // frequency; they are allocated lazily on the first diagnostic call. + bool rp_init_m = false; ///< whether rp_fdt_m has been allocated yet + int rp_Nf_m = 0; ///< DFT window length [time steps] (~3 resonant cycles) + double rp_omega_m = 0.0; ///< resonant angular frequency [1/unit_time], c = 1 + Kokkos::View rp_fdt_m; ///< ring buffer [Nf][nx][ny][4] = (Ex,Ey,Bx,By)_lab + +public: + size_type getTotalP() const { return totalP_m; } + void setTotalP(size_type totalP_) { totalP_m = totalP_; } + + int getNt() const { return nt_m; } + void setNt(int nt_) { nt_m = nt_; } + + const Vector_t& getNr() const { return nr_m; } + void setNr(const Vector_t& nr_) { nr_m = nr_; } + + double getTime() const { return time_m; } + void setTime(double time_) { time_m = time_; } + + std::shared_ptr getParticleContainer() { return pcontainer_m; } + void setParticleContainer(std::shared_ptr pcontainer) { + pcontainer_m = pcontainer; + } + + std::shared_ptr getFieldContainer() { return fcontainer_m; } + void setFieldContainer(std::shared_ptr fcontainer) { + fcontainer_m = fcontainer; + } + + std::shared_ptr getFieldSolver() { return solver_m; } + void setFieldSolver(std::shared_ptr solver) { solver_m = solver; } + + void pre_step() override { + Inform m("Pre-step"); + m << "Done" << endl; + } + + void post_step() override { + this->time_m += this->dt_m; + this->it_m++; + this->dump(); + + Inform m("Post-step:"); + m << "Finished time step: " << this->it_m << " time: " << this->time_m << endl; + } + + // Particle -> grid: charge-conserving current deposition into the four-current + // source field J ([1..Dim]); optionally the charge density into J[0]. + void par2grid() { depositCurrent(); } + + // Grid -> particle: interpolate E and B to the particle positions. + void grid2par() { gatherFields(); } + + void depositCurrent() { + using value_type = typename SourceField_t::value_type; + this->fcontainer_m->getJ() = value_type(0); + + auto policy = Kokkos::RangePolicy<>(0, this->pcontainer_m->getLocalNum()); + ippl::assemble_current_collocated(this->fcontainer_m->getMesh(), this->pcontainer_m->Q, + this->pcontainer_m->R_nm1, this->pcontainer_m->R, + this->fcontainer_m->getJ(), policy, (T)this->dt_m); + + if (this->m_config.space_charge) { + depositChargeDensity(); + } + + this->fcontainer_m->getJ().accumulateHalo(); + } + + void gatherFields() { + this->fcontainer_m->getE().fillHalo(); + this->fcontainer_m->getB().fillHalo(); + this->pcontainer_m->E_gather.gather(this->fcontainer_m->getE(), this->pcontainer_m->R, + false); + this->pcontainer_m->B_gather.gather(this->fcontainer_m->getB(), this->pcontainer_m->R, + false); + } + + // Relativistic Boris push with an externally supplied (E, B) field, performed + // in sub-steps over one FDTD time step. external_field(pos, time) must return + // a Kokkos::pair{E, B}. Faithful to the original NSFDSolverWithParticles + // particle update. + template + void push(External external_field) { + auto pc = this->pcontainer_m; + auto fc = this->fcontainer_m; + + // Remember the pre-push position; it is the start point for next step's + // current deposition (R_nm1 -> R). + Kokkos::deep_copy(pc->R_nm1.getView(), pc->R.getView()); + + fc->getE().fillHalo(); + fc->getB().fillHalo(); + Kokkos::fence(); + + const T dt = (T)this->dt_m; + const int nsub = nsubsteps_m; + const T bunch_dt = dt / nsub; + const T time = (T)this->time_m; + + for (int bts = 0; bts < nsub; ++bts) { + pc->E_gather.gather(fc->getE(), pc->R, false); + pc->B_gather.gather(fc->getB(), pc->R, false); + Kokkos::fence(); + + auto gbview = pc->gamma_beta.getView(); + auto eview = pc->E_gather.getView(); + auto bview = pc->B_gather.getView(); + auto qview = pc->Q.getView(); + auto mview = pc->mass.getView(); + auto rview = pc->R.getView(); + + Kokkos::parallel_for( + pc->getLocalNum(), KOKKOS_LAMBDA(const size_t i) { + const ippl::Vector pgammabeta = gbview(i); + ippl::Vector E_grid = eview(i); + ippl::Vector B_grid = bview(i); + ippl::Vector bunchpos = rview(i); + + Kokkos::pair, ippl::Vector> external_eb = + external_field(bunchpos, time + bunch_dt * bts); + + ippl::Vector, 2> EB{ + ippl::Vector(E_grid + external_eb.first), + ippl::Vector(B_grid + external_eb.second)}; + + const T charge = qview(i); + const T mass = mview(i); + + const ippl::Vector t1 = + pgammabeta + charge * bunch_dt * EB[0] / (T(2) * mass); + const T alpha = + charge * bunch_dt / (T(2) * mass * Kokkos::sqrt(1 + t1.dot(t1))); + const ippl::Vector t2 = t1 + alpha * ippl::cross(t1, EB[1]); + const ippl::Vector t3 = + t1 + + ippl::cross(t2, T(2) * alpha + * (EB[1] / (1.0 + alpha * alpha * (EB[1].dot(EB[1]))))); + const ippl::Vector ngammabeta = + t3 + charge * bunch_dt * EB[0] / (T(2) * mass); + + rview(i) = rview(i) + + bunch_dt * ngammabeta + / (Kokkos::sqrt(T(1.0) + ngammabeta.dot(ngammabeta))); + gbview(i) = ngammabeta; + }); + Kokkos::fence(); + } + + destroyOutOfBounds(); + pc->update(); + } + + void pre_run() override { + Inform m("Pre Run"); + + // The longitudinal box and the simulated time are measured in the + // co-moving frame: stretch z and shorten the time accordingly. + this->m_config.extents[2] *= frame_gamma_m; + this->m_config.total_time /= frame_gamma_m; + + for (unsigned i = 0; i < Dim; i++) { + this->nr_m[i] = this->m_config.resolution[i]; + this->domain_m[i] = ippl::Index(this->nr_m[i]); + } + + // Open boundaries (absorbing FDTD); decompose along z only. + this->decomp_m.fill(false); + this->decomp_m[Dim - 1] = true; + this->isAllPeriodic_m = false; + + for (unsigned d = 0; d < Dim; d++) { + this->hr_m[d] = this->m_config.extents[d] / this->m_config.resolution[d]; + this->origin_m[d] = -this->m_config.extents[d] * 0.5; + this->rmin_m[d] = this->origin_m[d]; + this->rmax_m[d] = this->origin_m[d] + this->m_config.extents[d]; + } + + // Courant / dispersion condition for the standard FDTD stencil. + const double rzx = this->hr_m[Dim - 1] / this->hr_m[0]; + const double rzy = this->hr_m[Dim - 1] / this->hr_m[1]; + if (rzx * rzx + rzy * rzy >= 1.0) { + m << "Dispersion relation not satisfiable" << endl; + ippl::Comm->abort(); + } + + m << "Discretization:" << endl + << "nt " << "(derived)" << " Np= " << this->totalP_m << " grid = " << this->nr_m + << endl; + + this->setFieldContainer(std::make_shared( + this->hr_m, this->rmin_m, this->rmax_m, this->decomp_m, this->domain_m, + this->origin_m, this->isAllPeriodic_m)); + + this->fcontainer_m->initializeFields(); + + this->setParticleContainer(std::make_shared( + this->fcontainer_m->getMesh(), this->fcontainer_m->getFL())); + + // The FDTD solver derives its own time step from the mesh (CFL); we read + // it back to drive the particle push and diagnostics. + this->setFieldSolver(std::make_shared(this->fcontainer_m->getJ(), + this->fcontainer_m->getE(), + this->fcontainer_m->getB())); + this->dt_m = this->solver_m->getDt(); + this->nt_m = (int)std::ceil(this->m_config.total_time / this->dt_m); + this->it_m = 0; + this->time_m = 0.0; + + m << "dt = " << this->dt_m << " nt = " << this->nt_m << endl; + + initializeParticles(); + + // Open the ffmpeg pipe (rank 0, only if periodic output was requested). + video_m.open(this->m_config); + + this->dump(); + + m << "Done" << endl; + } + + void initializeParticles() { + Inform m("Initialize Particles"); + + BunchInitialize mithra = generate_mithra_config(this->m_config, frame_m); + + // The MITHRA generator produces the whole bunch on rank 0; the first + // particle update (below) scatters it across ranks by position. + if (ippl::Comm->rank() == 0) { + size_t actualP = initialize_bunch_mithra(*this->pcontainer_m, mithra, frame_gamma_m); + this->pcontainer_m->Q = this->m_config.charge / actualP; + this->pcontainer_m->mass = this->m_config.mass / actualP; + } else { + this->pcontainer_m->create(0); + } + + // Center the bunch on the origin. + { + auto rview = this->pcontainer_m->R.getView(); + auto rm1view = this->pcontainer_m->R_nm1.getView(); + ippl::Vector meanpos = + this->pcontainer_m->R.sum() * (1.0 / this->pcontainer_m->getTotalNum()); + Kokkos::parallel_for( + this->pcontainer_m->getLocalNum(), KOKKOS_LAMBDA(const size_t i) { + rview(i) -= meanpos; + rm1view(i) -= meanpos; + }); + Kokkos::fence(); + } + + // Distribute particles to their owning ranks before the first step. + this->pcontainer_m->update(); + + m << "particles created and initial conditions assigned" << endl; + } + + void advance() override { + // 1. Deposit the current produced by last step's motion (R_nm1 -> R). + this->par2grid(); + + // 2. Advance the electromagnetic field one FDTD step. + this->solver_m->solve(); + + // 3. Push particles with the self-consistent field plus the undulator + // field transformed into the co-moving frame. + auto und = undulator_m; + auto lb = frame_m; + this->push(KOKKOS_LAMBDA(ippl::Vector pos, T time) { + lb.primedToUnprimed(pos, time); + auto eb = und(pos); + return lb.transform_EB(eb); + }); + } + + void dump() { + dumpRadiation(); + dumpRadiationBanded(); + dumpFELDiagnostics(); + + // Emit a video frame on the configured rhythm. writeFrame is collective, + // so every rank must reach it under the same condition. + if (this->m_config.output_rhythm != 0 + && (this->it_m % (int)this->m_config.output_rhythm) == 0) { + video_m.writeFrame(this->it_m, *this->fcontainer_m, *this->pcontainer_m, frame_m, + this->m_config, this->nr_m); + } + } + + // Radiated power leaving the downstream end of the domain, transformed back + // to the lab frame and integrated over the transverse exit plane. + void dumpRadiation() { + auto fc = this->fcontainer_m; + auto eview = fc->getE().getView(); + auto bview = fc->getB().getView(); + auto ldom = fc->getFL().getLocalNDIndex(); + auto lb = frame_m; + + const uint32_t nz = (uint32_t)this->nr_m[Dim - 1]; + + double radiation = 0.0; + Kokkos::parallel_reduce( + ippl::getRangePolicy(eview, 1), + KOKKOS_LAMBDA(const size_t i, const size_t j, const size_t k, double& ref) { + Kokkos::pair, ippl::Vector> buncheb{eview(i, j, k), + bview(i, j, k)}; + auto eblab = lb.inverse_transform_EB(buncheb); + uint32_t kg = (uint32_t)(k + ldom.first()[2]); + if (kg == nz - 3) { + ref += ippl::cross(eblab.first, eblab.second)[2]; + } + }, + radiation); + + double power_local = + radiation + * double(unit_powerdensity_in_watt_per_square_meter * unit_length_in_meters + * unit_length_in_meters) + * this->hr_m[0] * this->hr_m[1]; + double power_global = 0.0; + MPI_Reduce(&power_local, &power_global, 1, MPI_DOUBLE, MPI_SUM, 0, + ippl::Comm->getCommunicator()); + + if (ippl::Comm->rank() == 0) { + // Lab-frame longitudinal position the exit plane maps to at this time. + ippl::Vector pos{0, 0, (T)this->m_config.extents[2]}; + lb.primedToUnprimed(pos, (T)this->time_m); + + std::stringstream fname; + fname << this->m_config.output_path << "radiation_" << ippl::Comm->size() << ".csv"; + Inform csvout(NULL, fname.str().c_str(), Inform::APPEND); + csvout.precision(10); + csvout.setf(std::ios::scientific, std::ios::floatfield); + if (std::fabs(this->time_m) < 1e-14) { + csvout << "labframe_z, radiated_power_W" << endl; + } + csvout << pos[2] * unit_length_in_meters << " " << power_global << endl; + } + ippl::Comm->barrier(); + } + + // Narrow-band radiated power at the FEL resonance, computed the way MITHRA's + // powerSample() does. + void dumpRadiationBanded() { + auto fc = this->fcontainer_m; + auto eview = fc->getE().getView(); + auto bview = fc->getB().getView(); + auto ldom = fc->getFL().getLocalNDIndex(); + auto lb = frame_m; + + const uint32_t nz = (uint32_t)this->nr_m[Dim - 1]; + const int extx = (int)eview.extent(0); + const int exty = (int)eview.extent(1); + const int extz = (int)eview.extent(2); + + if (!rp_init_m) { + const double lambda_rad = this->m_config.undulator_period / frame_gamma_m; + rp_omega_m = 2.0 * M_PI / lambda_rad; // [1/unit_time], c = 1 + rp_Nf_m = std::max(1, (int)std::lround(3.0 * lambda_rad / this->dt_m)); + rp_fdt_m = Kokkos::View("FEL banded field buffer", rp_Nf_m, extx, exty, 4); + rp_init_m = true; + } + + const int Nf = rp_Nf_m; + const int m = ((this->it_m % Nf) + Nf) % Nf; // ring slot for this step + + const int kview = (int)(nz - 3) - ldom.first()[2]; + const bool owns = (kview >= 1 && kview < extz - 1); + + auto fdt = rp_fdt_m; + + // 1. Store this step's lab-frame transverse fields into ring slot m. + if (owns) { + Kokkos::parallel_for( + "FEL banded store", + Kokkos::MDRangePolicy>({1, 1}, {extx - 1, exty - 1}), + KOKKOS_LAMBDA(const int i, const int j) { + Kokkos::pair, ippl::Vector> buncheb{ + eview(i, j, kview), bview(i, j, kview)}; + auto eblab = lb.inverse_transform_EB(buncheb); + fdt(m, i, j, 0) = eblab.first[0]; // Ex_lab + fdt(m, i, j, 1) = eblab.first[1]; // Ey_lab + fdt(m, i, j, 2) = eblab.second[0]; // Bx_lab + fdt(m, i, j, 3) = eblab.second[1]; // By_lab + }); + } + + // 2. Single-frequency DFT over the window, summed over the exit plane. + double power_local = 0.0; + if (owns) { + const double omega = rp_omega_m; + const double dt = this->dt_m; + Kokkos::parallel_reduce( + "FEL banded DFT", + Kokkos::MDRangePolicy>({1, 1}, {extx - 1, exty - 1}), + KOKKOS_LAMBDA(const int i, const int j, double& ref) { + double ex_r = 0, ex_i = 0, ey_r = 0, ey_i = 0; + double bx_r = 0, bx_i = 0, by_r = 0, by_i = 0; + for (int mm = 0; mm < Nf; ++mm) { + const double ph = omega * mm * dt; + const double cp = Kokkos::cos(ph); + const double sp = Kokkos::sin(ph); + const double ex = fdt(mm, i, j, 0); + const double ey = fdt(mm, i, j, 1); + const double bx = fdt(mm, i, j, 2); + const double by = fdt(mm, i, j, 3); + ex_r += ex * cp; ex_i += ex * sp; + ey_r += ey * cp; ey_i += ey * sp; + bx_r += bx * cp; bx_i -= bx * sp; + by_r += by * cp; by_i -= by * sp; + } + ref += (ex_r * by_r - ex_i * by_i) - (ey_r * bx_r - ey_i * bx_i); + }, + power_local); + } + + // Convert the windowed sum to a cycle-averaged power in Watts. + power_local *= (2.0 / (double(Nf) * double(Nf))) + * double(unit_powerdensity_in_watt_per_square_meter * unit_length_in_meters + * unit_length_in_meters) + * this->hr_m[0] * this->hr_m[1]; + + double power_global = 0.0; + MPI_Reduce(&power_local, &power_global, 1, MPI_DOUBLE, MPI_SUM, 0, + ippl::Comm->getCommunicator()); + + if (ippl::Comm->rank() == 0) { + ippl::Vector pos{0, 0, (T)this->m_config.extents[2]}; + lb.primedToUnprimed(pos, (T)this->time_m); + + std::stringstream fname; + fname << this->m_config.output_path << "radiation_band_" << ippl::Comm->size() + << ".csv"; + Inform csvout(NULL, fname.str().c_str(), Inform::APPEND); + csvout.precision(10); + csvout.setf(std::ios::scientific, std::ios::floatfield); + if (std::fabs(this->time_m) < 1e-14) { + csvout << "labframe_z, banded_power_W" << endl; + } + csvout << pos[2] * unit_length_in_meters << " " << power_global << endl; + } + ippl::Comm->barrier(); + } + + // FEL gain diagnostics, written against the same lab-frame distance axis as + // the radiation curve so they can be overlaid: + // * bunching : micro-bunching factor || at the resonant + // wavelength lambda* = undulator_period / (2 gamma_frame). + // Exponential growth of this is the signature of FEL gain; + // if it stays at the ~1% seed there is no gain. + // * max|E| : peak electric-field magnitude in the domain. + // * fieldEnergy: total EM field energy (0.5 * sum(E.E + B.B) * cellVolume). + // * N : live particle count (catches runaway out-of-bounds loss). + void dumpFELDiagnostics() { + auto pc = this->pcontainer_m; + auto fc = this->fcontainer_m; + + // --- micro-bunching factor at the resonant wavelength --- + const double lambda_star = this->m_config.undulator_period / (2.0 * frame_gamma_m); + const double kstar = 2.0 * M_PI / lambda_star; + auto rview = pc->R.getView(); + double sumcos = 0.0, sumsin = 0.0; + Kokkos::parallel_reduce( + "FEL bunching", pc->getLocalNum(), + KOKKOS_LAMBDA(const size_t i, double& c, double& s) { + const double phase = kstar * rview(i)[2]; + c += Kokkos::cos(phase); + s += Kokkos::sin(phase); + }, + sumcos, sumsin); + + double gsumcos = 0.0, gsumsin = 0.0; + ippl::Comm->reduce(sumcos, gsumcos, 1, std::plus()); + ippl::Comm->reduce(sumsin, gsumsin, 1, std::plus()); + size_type nLocal = pc->getLocalNum(), nGlobal = 0; + ippl::Comm->reduce(nLocal, nGlobal, 1, std::plus()); + double bunching = (nGlobal > 0) + ? std::sqrt(gsumcos * gsumcos + gsumsin * gsumsin) / (double)nGlobal + : 0.0; + + // --- peak |E| and total EM field energy over the domain --- + const int nghost = fc->getE().getNghost(); + auto Eview = fc->getE().getView(); + auto Bview = fc->getB().getView(); + using index_array_type = typename ippl::RangePolicy::index_array_type; + double localE2 = 0.0, localEmax = 0.0; + ippl::parallel_reduce( + "FEL field stats", ippl::getRangePolicy(Eview, nghost), + KOKKOS_LAMBDA(const index_array_type& args, double& E2, double& Emax) { + ippl::Vector E = ippl::apply(Eview, args); + ippl::Vector B = ippl::apply(Bview, args); + E2 += E.dot(E) + B.dot(B); + double en = Kokkos::sqrt(E.dot(E)); + if (en > Emax) { + Emax = en; + } + }, + Kokkos::Sum(localE2), Kokkos::Max(localEmax)); + + double globalE2 = 0.0, globalEmax = 0.0; + ippl::Comm->reduce(localE2, globalE2, 1, std::plus()); + ippl::Comm->reduce(localEmax, globalEmax, 1, std::greater()); + double cellVolume = 1.0; + for (unsigned d = 0; d < Dim; ++d) { + cellVolume *= this->hr_m[d]; + } + double fieldEnergy = 0.5 * globalE2 * cellVolume; + + // --- bunch centroid, RMS size, and mean longitudinal momentum, to see + // HOW particles leave the domain (all in the co-moving frame, units + // of unit_length). Box half-widths: z in [origin_z, origin_z+L_z], + // transverse +-extents/2. --- + auto gbview = pc->gamma_beta.getView(); + double sz = 0.0, sz2 = 0.0, sperp2 = 0.0, sgbz = 0.0; + Kokkos::parallel_reduce( + "FEL bunch shape", pc->getLocalNum(), + KOKKOS_LAMBDA(const size_t i, double& az, double& az2, double& aperp2, double& agbz) { + const double x = rview(i)[0], y = rview(i)[1], z = rview(i)[2]; + az += z; + az2 += z * z; + aperp2 += x * x + y * y; + agbz += gbview(i)[2]; + }, + sz, sz2, sperp2, sgbz); + double gsz = 0.0, gsz2 = 0.0, gsperp2 = 0.0, gsgbz = 0.0; + ippl::Comm->reduce(sz, gsz, 1, std::plus()); + ippl::Comm->reduce(sz2, gsz2, 1, std::plus()); + ippl::Comm->reduce(sperp2, gsperp2, 1, std::plus()); + ippl::Comm->reduce(sgbz, gsgbz, 1, std::plus()); + double cz = 0.0, rmsz = 0.0, rmsperp = 0.0, meangbz = 0.0; + if (nGlobal > 0) { + cz = gsz / nGlobal; + rmsz = std::sqrt(std::max(0.0, gsz2 / nGlobal - cz * cz)); + rmsperp = std::sqrt(gsperp2 / nGlobal); + meangbz = gsgbz / nGlobal; + } + + Inform m("FELDiag"); + m << "t=" << this->time_m << " bunching=" << bunching << " max|E|=" << globalEmax + << " fieldEnergy=" << fieldEnergy << " Np=" << nGlobal << " cz=" << cz + << " rms_z=" << rmsz << " rms_perp=" << rmsperp << " mean_gbz=" << meangbz << endl; + + if (ippl::Comm->rank() == 0) { + ippl::Vector pos{0, 0, (T)this->m_config.extents[2]}; + frame_m.primedToUnprimed(pos, (T)this->time_m); + + std::stringstream fname; + fname << this->m_config.output_path << "feldiag_" << ippl::Comm->size() << ".csv"; + Inform csvout(NULL, fname.str().c_str(), Inform::APPEND); + csvout.precision(10); + csvout.setf(std::ios::scientific, std::ios::floatfield); + if (std::fabs(this->time_m) < 1e-14) { + csvout << "labframe_z, bunching, max_E, field_energy, num_particles, " + "centroid_z, rms_z, rms_perp, mean_gbz" + << endl; + } + csvout << pos[2] * unit_length_in_meters << " " << bunching << " " << globalEmax << " " + << fieldEnergy << " " << nGlobal << " " << cz << " " << rmsz << " " << rmsperp + << " " << meangbz << endl; + } + ippl::Comm->barrier(); + } + +protected: + // Deposit charge density (CIC) into component [0] of the four-current field, + // needed only when space charge is enabled. Component [1..Dim] is filled by + // assemble_current_collocated. + void depositChargeDensity() { + auto pc = this->pcontainer_m; + auto fc = this->fcontainer_m; + auto view = fc->getJ().getView(); + auto qview = pc->Q.getView(); + auto rview = pc->R.getView(); + const auto origin = fc->getMesh().getOrigin(); + const auto h = fc->getMesh().getMeshSpacing(); + const auto ldom = fc->getFL().getLocalNDIndex(); + const int nghost = fc->getJ().getNghost(); + T volume = T(1); + for (unsigned d = 0; d < Dim; ++d) + volume *= h[d]; + + Kokkos::parallel_for( + "FEL deposit charge density", pc->getLocalNum(), KOKKOS_LAMBDA(const size_t p) { + const ippl::Vector pos = rview(p); + const T value = qview(p) / volume; + + Kokkos::Array cellIdx; + Kokkos::Array xi; + for (unsigned d = 0; d < Dim; ++d) { + // Half-cell shift to match the Cell-centered field and the + // gather (see assemble_current_collocated). + const T gridpos = (pos[d] - origin[d]) / h[d] - T(0.5); + cellIdx[d] = static_cast(Kokkos::floor(gridpos)); + xi[d] = gridpos - T(cellIdx[d]); + } + for (unsigned corner = 0; corner < (1u << Dim); ++corner) { + size_t idx[Dim]; + T weight = T(1); + for (unsigned d = 0; d < Dim; ++d) { + const unsigned offset = (corner >> d) & 1u; + weight *= offset ? xi[d] : (T(1) - xi[d]); + idx[d] = cellIdx[d] - ldom.first()[d] + nghost + offset; + } + Kokkos::atomic_add(&(ippl::apply(view, idx)[0]), value * weight); + } + }); + Kokkos::fence(); + } + + // Mark particles that have left the physical domain and remove them (open BC). + void destroyOutOfBounds() { + auto pc = this->pcontainer_m; + auto rview = pc->R.getView(); + const auto origin = this->fcontainer_m->getMesh().getOrigin(); + ippl::Vector extent; + for (unsigned d = 0; d < Dim; ++d) + extent[d] = this->nr_m[d] * this->hr_m[d]; + + Kokkos::View invalid("OOB particles", pc->getLocalNum()); + size_type invalid_count = 0; + Kokkos::parallel_reduce( + pc->getLocalNum(), + KOKKOS_LAMBDA(const size_t i, size_type& ref) { + bool out_of_bounds = false; + const ippl::Vector ppos = rview(i); + for (unsigned d = 0; d < Dim; ++d) { + out_of_bounds |= (ppos[d] <= origin[d]); + out_of_bounds |= (ppos[d] >= origin[d] + extent[d]); + } + invalid(i) = out_of_bounds; + ref += out_of_bounds; + }, + invalid_count); + Kokkos::fence(); + pc->destroy(invalid, invalid_count); + Kokkos::fence(); + } +}; + +#endif diff --git a/fel/LorentzTransform.h b/fel/LorentzTransform.h new file mode 100644 index 000000000..c8fa30cc1 --- /dev/null +++ b/fel/LorentzTransform.h @@ -0,0 +1,157 @@ +#ifndef LORENTZ_TRANSFORM_H +#define LORENTZ_TRANSFORM_H +#include + +#include "Types/Vector.h" +namespace ippl { + /** + * @brief A template struct representing a uniaxial Lorentz frame. + * + * The UniaxialLorentzframe struct represents a Lorentz transformation along a specific axis + * (default is the z-axis, axis 2). It includes methods for transforming vectors and + * electromagnetic fields between frames, as well as for computing gamma factors and beta + * velocities. The struct uses Kokkos for portability and performance. + * + * @tparam T The scalar type used for computations (e.g., float, double). + * @tparam axis The axis along which the Lorentz transformation is applied (default is 2, the + * z-axis). + */ + template + struct UniaxialLorentzframe { + /// Speed of light, set to 1 for natural units. + constexpr static T c = 1.0; + /// Alias for the scalar type used in the struct. + using scalar = T; + /// Alias for a 3-dimensional vector of the scalar type. + using Vector3 = ippl::Vector; + + /// Beta velocity component along the specified axis. + scalar beta_m; + /// Product of gamma factor and beta velocity. + scalar gammaBeta_m; + /// Gamma factor for the Lorentz transformation. + scalar gamma_m; + + /** + * @brief Create a UniaxialLorentzframe from a given gamma factor. + * + * This static member function constructs a UniaxialLorentzframe object using a given + * gamma factor. It computes the corresponding beta velocity and gamma*beta product. + * + * @param gamma The gamma factor for the Lorentz transformation. + * @return A UniaxialLorentzframe object with computed beta and gamma*beta. + */ + KOKKOS_INLINE_FUNCTION static UniaxialLorentzframe from_gamma(const scalar gamma) { + UniaxialLorentzframe ret; + ret.gamma_m = gamma; + scalar beta = Kokkos::sqrt(1 - double(1) / (gamma * gamma)); + scalar gammabeta = gamma * beta; + ret.beta_m = beta; + ret.gammaBeta_m = gammabeta; + return ret; + } + + /** + * @brief Get a UniaxialLorentzframe with negative beta velocity. + * + * This member function returns a new UniaxialLorentzframe object with the same gamma + * factor but with a negative beta velocity, effectively representing the inverse + * Lorentz transformation. + * + * @return A UniaxialLorentzframe object with negative beta velocity. + */ + KOKKOS_INLINE_FUNCTION UniaxialLorentzframe negative() const noexcept { + UniaxialLorentzframe ret; + ret.beta_m = -beta_m; + ret.gammaBeta_m = -gammaBeta_m; + ret.gamma_m = gamma_m; + return ret; + } + + /// Default constructor. + KOKKOS_INLINE_FUNCTION UniaxialLorentzframe() = default; + + /** + * @brief Construct a UniaxialLorentzframe from a gamma*beta value. + * + * This constructor initializes a UniaxialLorentzframe object using a given gamma*beta + * value. It computes the corresponding beta velocity and gamma factor. + * + * @param gammaBeta The product of the gamma factor and beta velocity. + */ + KOKKOS_INLINE_FUNCTION UniaxialLorentzframe(const scalar gammaBeta) { + using Kokkos::sqrt; + gammaBeta_m = gammaBeta; + beta_m = gammaBeta / sqrt(1 + gammaBeta * gammaBeta); + gamma_m = sqrt(1 + gammaBeta * gammaBeta); + } + + /** + * @brief Transform a spatial vector from the primed frame to the unprimed frame. + * + * This member function transforms a given spatial vector from the primed frame to the + * unprimed frame using the Lorentz transformation along the specified axis and the given + * time. + * + * @param arg The spatial vector to be transformed. + * @param time The time component for the transformation. + */ + KOKKOS_INLINE_FUNCTION void primedToUnprimed(Vector3& arg, scalar time) const noexcept { + arg[axis] = gamma_m * (arg[axis] + beta_m * time); + } + + /** + * @brief Transform electric and magnetic fields from the unprimed to the primed frame. + * + * This member function transforms a pair of electric and magnetic field vectors from the + * unprimed frame to the primed frame using the Lorentz transformation along the specified + * axis. + * + * @param unprimedEB A pair of vectors representing the electric and magnetic fields in the + * unprimed frame. + * @return A pair of vectors representing the electric and magnetic fields in the primed + * frame. + */ + KOKKOS_INLINE_FUNCTION Kokkos::pair, ippl::Vector> transform_EB( + const Kokkos::pair, ippl::Vector>& unprimedEB) const noexcept { + Kokkos::pair, ippl::Vector> ret; + ippl::Vector betavec{0, 0, beta_m}; + ippl::Vector vnorm{axis == 0, axis == 1, axis == 2}; + ret.first = + ippl::Vector(unprimedEB.first + cross(betavec, unprimedEB.second)) * gamma_m + - (vnorm * (gamma_m - 1) * (unprimedEB.first.dot(vnorm))); + ret.second = + ippl::Vector(unprimedEB.second - cross(betavec, unprimedEB.first)) * gamma_m + - (vnorm * (gamma_m - 1) * (unprimedEB.second.dot(vnorm))); + // NOTE: the two assignments above are already the complete boost: the + // -vnorm*(gamma-1)*(field.vnorm) term leaves the boost-axis (parallel) + // component invariant (gamma*B_z - (gamma-1)*B_z = B_z). The original + // code additionally subtracted (gamma-1)*field[axis] here, which + // double-corrects the parallel component to (2-gamma)*field[axis] -- + // a large, wrong-sign spurious longitudinal field that blows the beam + // up once the undulator (which has B_z != 0 off-axis) turns on. That + // erroneous correction is removed. + return ret; + } + + /** + * @brief Transform electric and magnetic fields from the primed to the unprimed frame. + * + * This member function transforms a pair of electric and magnetic field vectors from the + * primed frame to the unprimed frame using the inverse Lorentz transformation along the + * specified axis. + * + * @param primedEB A pair of vectors representing the electric and magnetic fields in the + * primed frame. + * @return A pair of vectors representing the electric and magnetic fields in the unprimed + * frame. + */ + KOKKOS_INLINE_FUNCTION Kokkos::pair, ippl::Vector> + inverse_transform_EB( + const Kokkos::pair, ippl::Vector>& primedEB) const noexcept { + return negative().transform_EB(primedEB); + } + }; + +} // namespace ippl +#endif diff --git a/fel/MithraBunch.h b/fel/MithraBunch.h new file mode 100644 index 000000000..bb828326d --- /dev/null +++ b/fel/MithraBunch.h @@ -0,0 +1,494 @@ +#ifndef IPPL_FEL_MITHRA_BUNCH_H +#define IPPL_FEL_MITHRA_BUNCH_H + +// MITHRA-style relativistic bunch initialization. +// +// Ported faithfully from the original FreeElectronLaser.cpp. Generates an +// ellipsoidal electron bunch (Gaussian/uniform distributions, optional shot +// noise / bunching factor / tail tapering), boosts it into the moving frame, +// and copies the result into an ippl particle bunch. +// +// Behaviour is unchanged from the original; only the surrounding includes and +// the assert_isreal helper were adapted (the old MaxwellSolvers/FDTD.h that +// defined it is no longer part of the tree). + +#include +#include +#include +#include +#include +#include + +#include + +#include "Types/Vector.h" + +#include "Config.h" +#include "LorentzTransform.h" +#include "datatypes.h" +#include "units.h" + +#ifndef assert_isreal +#define assert_isreal(X) assert(!std::isnan(X) && !std::isinf(X)) +#endif + +template +using FieldVector = ippl::Vector; +template +struct BunchInitialize { + /* Type of the distributions (transverse or longitudinal) in the bunch. + */ + std::string distribution_; + + /* Type of the generator for creating the bunch distribution. + */ + std::string generator_; + + /* Total number of macroparticles in the bunch. */ + unsigned int numberOfParticles_; + + /* Total charge of the bunch in pC. */ + scalar cloudCharge_; + + /* Initial energy of the bunch in MeV. */ + scalar initialGamma_; + + /* Initial normalized speed of the bunch. */ + scalar initialBeta_; + + /* Initial movement direction of the bunch, which is a unit vector. */ + FieldVector initialDirection_; + + /* Position of the center of the bunch in the unit of length scale. */ + FieldVector position_; + + /* Number of macroparticles in each direction for 3Dcrystal type. */ + FieldVector numbers_; + + /* Lattice constant in x, y, and z directions for 3D crystal type. */ + FieldVector latticeConstants_; + + /* Spread in position for each of the directions in the unit of length scale. For the 3D crystal + * type, it will be the spread in position for each micro-bunch of the crystal. + */ + FieldVector sigmaPosition_; + + /* Spread in energy in each direction. */ + FieldVector sigmaGammaBeta_; + + /* Store the truncation transverse distance for the electron generation. + */ + scalar tranTrun_; + + /* Store the truncation longitudinal distance for the electron generation. + */ + scalar longTrun_; + + /* Name of the file for reading the electrons distribution from. + */ + std::string fileName_; + + /* The radiation wavelength corresponding to the bunch length outside the undulator + */ + scalar lambda_; + + /* Bunching factor for the initialization of the bunch. + */ + scalar bF_; + + /* Phase of the bunching factor for the initialization of the bunch. + */ + scalar bFP_; + + /* Boolean flag determining the activation of shot-noise. + */ + bool shotNoise_; + + /* Initial beta vector of the bunch, which is obtained as the product of beta and direction. + */ + FieldVector betaVector_; + +}; + +// LORENTZ FRAME AND UNDULATOR + +template +BunchInitialize generate_mithra_config( + const config& cfg, const ippl::UniaxialLorentzframe& /*frame_boost unused*/) { + using vec3 = ippl::Vector; + scalar frame_gamma = cfg.bunch_gamma / std::sqrt(1 + 0.5 * cfg.undulator_K * cfg.undulator_K); + BunchInitialize init; + init.generator_ = "random"; + init.distribution_ = "uniform"; + init.initialDirection_ = vec3{0, 0, 1}; + init.initialGamma_ = cfg.bunch_gamma; + init.initialBeta_ = cfg.bunch_gamma == scalar(1) + ? 0 + : (sqrt(cfg.bunch_gamma * cfg.bunch_gamma - 1) / cfg.bunch_gamma); + init.sigmaGammaBeta_ = vector_cast(cfg.sigma_momentum); + init.sigmaPosition_ = vector_cast(cfg.sigma_position); + + // Initial bunching factor. + init.bF_ = 0.01; + init.bFP_ = 0; + init.shotNoise_ = false; + init.cloudCharge_ = cfg.charge; + init.lambda_ = cfg.undulator_period / (2 * frame_gamma * frame_gamma); + init.longTrun_ = cfg.position_truncations[2]; + init.tranTrun_ = cfg.position_truncations[0]; + init.position_ = vector_cast(cfg.mean_position); + init.betaVector_ = ippl::Vector{0, 0, init.initialBeta_}; + init.numberOfParticles_ = cfg.num_particles; + + init.numbers_ = 0; // UNUSED + init.latticeConstants_ = vec3{0, 0, 0}; // UNUSED + + return init; +} +template +struct Charge { + Double q; /* Charge of the point in the unit of electron charge. */ + FieldVector rnp, rnm; /* Position vector of the charge. */ + FieldVector gb; /* Normalized velocity vector of the charge. */ + + /* Double flag determining if the particle is passing the entrance point of the undulator. This + * flag can be used for better boosting the bunch to the moving frame. We need to consider it to + * be double, because this flag needs to be communicated during bunch update. + */ + Double e; +}; +template +using ChargeVector = std::list>; +template +void initializeBunchEllipsoid(BunchInitialize bunchInit, ChargeVector& chargeVector, + int rank, int size, int ia) { + /* Correct the number of particles if it is not a multiple of four. + */ + if (bunchInit.numberOfParticles_ % 4 != 0) { + unsigned int n = bunchInit.numberOfParticles_ % 4; + bunchInit.numberOfParticles_ += 4 - n; + Inform msg("BunchInitialize"); + msg << "Warning: number of particles in the bunch is not a multiple of four; " + << "corrected to " << bunchInit.numberOfParticles_ << "." << endl; + } + + /* Save the initially given number of particles. */ + unsigned int Np = bunchInit.numberOfParticles_, i, Np0 = chargeVector.size(); + + /* Declare the required parameters for the initialization of charge vectors. */ + Charge charge; + charge.q = bunchInit.cloudCharge_ / Np; + FieldVector gb = bunchInit.initialGamma_ * bunchInit.betaVector_; + FieldVector r(0.0); + FieldVector t(0.0); + Double t0; //, g; + Double zmin = 1e100; + Double Ne, bF, bFi; + unsigned int bmi; + std::vector randomNumbers; + + /* The initialization in group of four particles should only be done if there exists an + * undulator in the interaction. + */ + unsigned int ng = (bunchInit.lambda_ == 0.0) ? 1 : 4; + + /* Check the bunching factor. */ + if (bunchInit.bF_ > 2.0 || bunchInit.bF_ < 0.0) { + throw IpplException("initializeBunchEllipsoid", + "The bunching factor must be between 0 and 2."); + } + + /* If the generator is random we should make sure that different processors do not produce the + * same random numbers. + */ + if (bunchInit.generator_ == "random") { + /* Initialize the random number generator with a fixed seed so the + * generated bunch is reproducible across runs. + */ + srand(42); + /* Np / ng * 20 is the maximum number of particles. + */ + randomNumbers.resize(Np / ng * 20, 0.0); + for (unsigned int ri = 0; ri < Np / ng * 20; ri++) + randomNumbers[ri] = + (float)std::min(1 - 1e-7, std::max(1e-7, ((double)rand()) / RAND_MAX)); + } + + /* Declare the generator function depending on the input. + */ + auto generate = [&](unsigned int n, unsigned int m) { + return (randomNumbers.at(n * 2 * Np / ng + m)); + }; + + /* Declare the function for injecting the shot noise. + */ + auto insertCharge = [&](Charge q) { + for (unsigned int ii = 0; ii < ng; ii++) { + /* The random modulation is introduced depending on the shot-noise being activated. + */ + if (bunchInit.shotNoise_) { + /* Obtain the number of beamlet. + */ + bmi = int((charge.rnp[2] - zmin) / bunchInit.lambda_); + + /* Obtain the phase and amplitude of the modulation. + */ + bFi = bF * sqrt(-2.0 * log(generate(8, bmi))); + + q.rnp[2] = charge.rnp[2] - bunchInit.lambda_ / 4 * ii; + + q.rnp[2] -= bunchInit.lambda_ / M_PI * bFi + * sin(2.0 * M_PI / bunchInit.lambda_ * q.rnp[2] + + 2.0 * M_PI * generate(9, bmi)); + } else if (bunchInit.lambda_ != 0.0) { + q.rnp[2] = charge.rnp[2] - bunchInit.lambda_ / 4 * ii; + + q.rnp[2] -= bunchInit.lambda_ / M_PI * bunchInit.bF_ + * sin(2.0 * M_PI / bunchInit.lambda_ * q.rnp[2] + + bunchInit.bFP_ * M_PI / 180.0); + } + + /* Set this charge into the charge vector. */ + chargeVector.push_back(q); + } + }; + + /* If the shot noise is on, we need the minimum value of the bunch z coordinate to be able to + * calculate the FEL bucket number. */ + if (bunchInit.shotNoise_) { + for (i = 0; i < Np / ng; i++) { + if (bunchInit.distribution_ == "uniform") + zmin = std::min( + Double(2.0 * generate(2, i + Np0) - 1.0) * bunchInit.sigmaPosition_[2], zmin); + else if (bunchInit.distribution_ == "gaussian") + zmin = std::min( + (Double)(bunchInit.sigmaPosition_[2] * sqrt(-2.0 * log(generate(2, i + Np0))) + * sin(2.0 * M_PI * generate(3, i + Np0))), + zmin); + else { + std::cout << std::string( + "The longitudinal type is not correctly given to the code !!!\n"); + exit(1); + } + } + + if (bunchInit.distribution_ == "uniform") + for (; i < unsigned(Np / ng + * (1.0 + + 2.0 * bunchInit.lambda_ * sqrt(2.0 * M_PI) + / (2.0 * bunchInit.sigmaPosition_[2]))); + i++) { + t0 = 2.0 * bunchInit.lambda_ * sqrt(-2.0 * log(generate(2, i + Np0))) + * sin(2.0 * M_PI * generate(3, i + Np0)); + t0 += (t0 < 0.0) ? (-bunchInit.sigmaPosition_[2]) : (bunchInit.sigmaPosition_[2]); + + zmin = std::min(t0, zmin); + } + + zmin = zmin + bunchInit.position_[2]; + + /* Obtain the average number of electrons per FEL beamlet. + */ + Ne = bunchInit.cloudCharge_ * bunchInit.lambda_ / (2.0 * bunchInit.sigmaPosition_[2]); + + /* Set the bunching factor level for the shot noise depending on the given values. + */ + bF = (bunchInit.bF_ == 0.0) ? 1.0 / sqrt(Ne) : bunchInit.bF_; + } + + /* Determine the properties of each charge point and add them to the charge vector. */ + for (i = rank; i < Np / ng; i += size) { + /* Determine the transverse coordinate. */ + r[0] = bunchInit.sigmaPosition_[0] * sqrt(-2.0 * log(generate(0, i + Np0))) + * cos(2.0 * M_PI * generate(1, i + Np0)); + r[1] = bunchInit.sigmaPosition_[1] * sqrt(-2.0 * log(generate(0, i + Np0))) + * sin(2.0 * M_PI * generate(1, i + Np0)); + + /* Determine the longitudinal coordinate. */ + if (bunchInit.distribution_ == "uniform") + r[2] = (2.0 * generate(2, i + Np0) - 1.0) * bunchInit.sigmaPosition_[2]; + else if (bunchInit.distribution_ == "gaussian") + r[2] = bunchInit.sigmaPosition_[2] * sqrt(-2.0 * log(generate(2, i + Np0))) + * sin(2.0 * M_PI * generate(3, i + Np0)); + else { + exit(1); + } + + /* Determine the transverse momentum. */ + t[0] = bunchInit.sigmaGammaBeta_[0] * sqrt(-2.0 * log(generate(4, i + Np0))) + * cos(2.0 * M_PI * generate(5, i + Np0)); + t[1] = bunchInit.sigmaGammaBeta_[1] * sqrt(-2.0 * log(generate(4, i + Np0))) + * sin(2.0 * M_PI * generate(5, i + Np0)); + t[2] = bunchInit.sigmaGammaBeta_[2] * sqrt(-2.0 * log(generate(6, i + Np0))) + * cos(2.0 * M_PI * generate(7, i + Np0)); + + if (fabs(r[0]) < bunchInit.tranTrun_ && fabs(r[1]) < bunchInit.tranTrun_ + && fabs(r[2]) < bunchInit.longTrun_) { + /* Shift the generated charge to the center position and momentum space. + */ + charge.rnp = bunchInit.position_; + charge.rnp += r; + + charge.gb = gb; + charge.gb += t; + if (std::isinf(gb[2])) { + std::cerr << "[Warning] Gammabeta obtained an klonked here\n"; + } + + /* Insert this charge and the mirrored ones into the charge vector. + */ + insertCharge(charge); + } + } + + /* If the longitudinal type of the bunch is uniform a tapered part needs to be added to remove + * the coherent spontaneous emission (CSE) from the tail of the bunch. + */ + if (bunchInit.distribution_ == "uniform") { + for (; i < unsigned(uint32_t(Np / ng) + * (1.0 + + 2.0 * bunchInit.lambda_ * sqrt(2.0 * M_PI) + / (2.0 * bunchInit.sigmaPosition_[2]))); + i += size) { + r[0] = bunchInit.sigmaPosition_[0] * sqrt(-2.0 * log(generate(0, i + Np0))) + * cos(2.0 * M_PI * generate(1, i + Np0)); + r[1] = bunchInit.sigmaPosition_[1] * sqrt(-2.0 * log(generate(0, i + Np0))) + * sin(2.0 * M_PI * generate(1, i + Np0)); + + /* Determine the longitudinal coordinate. */ + r[2] = 2.0 * bunchInit.lambda_ * sqrt(-2.0 * log(generate(2, i + Np0))) + * sin(2.0 * M_PI * generate(3, i + Np0)); + r[2] += (r[2] < 0.0) ? (-bunchInit.sigmaPosition_[2]) : (bunchInit.sigmaPosition_[2]); + + /* Determine the transverse momentum. + */ + t[0] = bunchInit.sigmaGammaBeta_[0] * sqrt(-2.0 * log(generate(4, i + Np0))) + * cos(2.0 * M_PI * generate(5, i + Np0)); + t[1] = bunchInit.sigmaGammaBeta_[1] * sqrt(-2.0 * log(generate(4, i + Np0))) + * sin(2.0 * M_PI * generate(5, i + Np0)); + t[2] = bunchInit.sigmaGammaBeta_[2] * sqrt(-2.0 * log(generate(6, i + Np0))) + * cos(2.0 * M_PI * generate(7, i + Np0)); + if (fabs(r[0]) < bunchInit.tranTrun_ && fabs(r[1]) < bunchInit.tranTrun_ + && fabs(r[2]) < bunchInit.longTrun_) { + /* Shift the generated charge to the center position and momentum space. + */ + charge.rnp = bunchInit.position_[ia]; + charge.rnp += r; + + charge.gb = gb; + + charge.gb += t; + /* Insert this charge and the mirrored ones into the charge vector. + */ + insertCharge(charge); + } + } + } + + /* Reset the value for the number of particle variable according to the installed number of + * macro-particles and perform the corresponding changes. */ + bunchInit.numberOfParticles_ = chargeVector.size(); +} + +template +void boost_bunch(ChargeVector& chargeVectorn_, Double frame_gamma) { + Double frame_beta = std::sqrt((double)frame_gamma * frame_gamma - 1.0) / double(frame_gamma); + Double zmaxL = -1.0e100, zmaxG; + for (auto iterQ = chargeVectorn_.begin(); iterQ != chargeVectorn_.end(); iterQ++) { + Double g = std::sqrt(1.0 + iterQ->gb.dot(iterQ->gb)); + if (std::isinf(g)) { + std::cerr << __FILE__ << ": " << __LINE__ << " inf gb: " << iterQ->gb << ", g = " << g + << "\n"; + abort(); + } + Double bz = iterQ->gb[2] / g; + iterQ->rnp[2] *= frame_gamma; + + iterQ->gb[2] = frame_gamma * g * (bz - frame_beta); + + zmaxL = std::max(zmaxL, iterQ->rnp[2]); + } + zmaxG = zmaxL; + struct { + Double zu_; + Double beta_; + } bunch_; + bunch_.zu_ = zmaxG; + bunch_.beta_ = frame_beta; + + /****************************************************************************************************/ + + for (auto iterQ = chargeVectorn_.begin(); iterQ != chargeVectorn_.end(); iterQ++) { + Double g = std::sqrt(1.0 + iterQ->gb.dot(iterQ->gb)); + iterQ->rnp[0] += iterQ->gb[0] / g * (iterQ->rnp[2] - bunch_.zu_) * frame_beta; + iterQ->rnp[1] += iterQ->gb[1] / g * (iterQ->rnp[2] - bunch_.zu_) * frame_beta; + iterQ->rnp[2] += iterQ->gb[2] / g * (iterQ->rnp[2] - bunch_.zu_) * frame_beta; + if (std::isnan(iterQ->rnp[2])) { + std::cerr << iterQ->gb[2] << ", " << g << ", " << iterQ->rnp[2] << ", " << bunch_.zu_ + << ", " << frame_beta << "\n"; + std::cerr << __FILE__ << ": " << __LINE__ << " Particle has NaN velocity or position\n"; + abort(); + } + } +} + +template +size_t initialize_bunch_mithra(bunch_type& bunch, const BunchInitialize& bunchInit, + scalar frame_gamma) { + ChargeVector temporary_charge_list; + initializeBunchEllipsoid(bunchInit, temporary_charge_list, 0, 1, 0); + for (auto& c : temporary_charge_list) { + if (std::isnan(c.rnp[0]) || std::isnan(c.rnp[1]) || std::isnan(c.rnp[2])) + std::cout << "Pos before boost: " << c.rnp << "\n"; + if (std::isinf(c.rnp[0]) || std::isinf(c.rnp[1]) || std::isinf(c.rnp[2])) + std::cout << "Pos before boost: " << c.rnp << "\n"; + } + boost_bunch(temporary_charge_list, frame_gamma); + for (auto& c : temporary_charge_list) { + if (std::isnan(c.rnp[0]) || std::isnan(c.rnp[1]) || std::isnan(c.rnp[2])) { + std::cout << "Pos after boost: " << c.rnp << "\n"; + break; + } + } + Kokkos::View*, Kokkos::HostSpace> positions("", temporary_charge_list.size()); + Kokkos::View*, Kokkos::HostSpace> gammabetas("", temporary_charge_list.size()); + auto iterQ = temporary_charge_list.begin(); + for (size_t i = 0; i < temporary_charge_list.size(); i++) { + assert_isreal(iterQ->gb[0]); + assert_isreal(iterQ->gb[1]); + assert_isreal(iterQ->gb[2]); + assert(iterQ->gb[2] != 0.0f); + scalar g = std::sqrt(1.0 + iterQ->gb.dot(iterQ->gb)); + assert_isreal(g); + scalar bz = iterQ->gb[2] / g; + assert_isreal(bz); + (void)bz; + positions(i) = iterQ->rnp; + gammabetas(i) = iterQ->gb; + ++iterQ; + } + if (temporary_charge_list.size() > bunch.getLocalNum()) { + bunch.create(temporary_charge_list.size() - bunch.getLocalNum()); + } + Kokkos::View*> dpositions("", temporary_charge_list.size()); + Kokkos::View*> dgammabetas("", temporary_charge_list.size()); + + Kokkos::deep_copy(dpositions, positions); + Kokkos::deep_copy(dgammabetas, gammabetas); + Kokkos::deep_copy(bunch.R_nm1.getView(), positions); + Kokkos::deep_copy(bunch.gamma_beta.getView(), gammabetas); + auto rview = bunch.R.getView(), rm1view = bunch.R_nm1.getView(), + gbview = bunch.gamma_beta.getView(); + ; + Kokkos::parallel_for( + temporary_charge_list.size(), KOKKOS_LAMBDA(size_t i) { + rview(i) = dpositions(i); + rm1view(i) = dpositions(i); + gbview(i) = dgammabetas(i); + }); + Kokkos::fence(); + + return temporary_charge_list.size(); +} + +#endif diff --git a/fel/README.md b/fel/README.md new file mode 100644 index 000000000..8db95d543 --- /dev/null +++ b/fel/README.md @@ -0,0 +1,35 @@ +# FEL — Free Electron Laser mini-app + +An electromagnetic PIC simulation of a free-electron laser: a relativistic +electron bunch is tracked through an undulator in a co-moving Lorentz frame, +with the self-consistent field advanced by an FDTD Maxwell solver. The radiated +power is written to a CSV (and, optionally, a Poynting-flux video). + +## Build + +```sh +cmake -S . -B build -DIPPL_ENABLE_FEL=ON -DCMAKE_CXX_STANDARD=20 +cmake --build build --target FreeElectronLaser +``` + +The executable is built at +`build/fel/FreeElectronLaser`. + +## Run + +```sh +cd build +./fel/FreeElectronLaser ../fel/config.json --info 5 +``` + +The argument is a MITHRA-style JSON job file (defaults to `../fel/config.json`); +see [config.json](config.json) for the available keys. Run on multiple ranks +with `mpirun -np ...`. + +Output is written to the directory given by `output.path` in the config: +`radiation_.csv` holds the radiated power. If `output.rhythm > 0`, a +Poynting-flux video is produced and requires **ffmpeg** on the `PATH`. + +## Acknowledgements + +The relativistic bunch initialization and the resonance-power diagnostic are directly ported from [MITHRA](https://github.com/aryafallahi/mithra), a full-wave free-electron-laser solver by A. Fallahi (GPL-licensed). It was also used as reference for much of the other parts. diff --git a/fel/Undulator.h b/fel/Undulator.h new file mode 100644 index 000000000..74a589c1e --- /dev/null +++ b/fel/Undulator.h @@ -0,0 +1,116 @@ +#ifndef UNDULATOR_H +#define UNDULATOR_H +#include + +#include + +#include "Types/Vector.h" + +#include "units.h" +#include "LorentzTransform.h" +namespace ippl { + template + struct undulator_parameters { + scalar lambda_u; // MITHRA: lambda_u + scalar K; // Undulator parameter + scalar length; + scalar B_magnitude; + undulator_parameters(scalar K_undulator_parameter, scalar lambda_u, scalar _length) + : lambda_u(lambda_u) + , K(K_undulator_parameter) + , length(_length) { + B_magnitude = (2 * M_PI * electron_mass_in_unit_masses * K) + / (electron_charge_in_unit_charges * lambda_u); + } + }; + /** + * @brief Struct representing an undulator. + * + * @tparam scalar Type of the scalar values (e.g., float, double). + */ + template + struct Undulator { + undulator_parameters uparams; ///< Parameters of the undulator. + scalar distance_to_entry; ///< Distance to the entry of the undulator. + scalar k_u; ///< Wavenumber of the undulator. + + /** + * @brief Constructor to initialize undulator parameters and calculate k_u. + * + * @param p Parameters of the undulator. + * @param dte Distance to the entry of the undulator. + */ + KOKKOS_FUNCTION Undulator(const undulator_parameters& p, scalar dte) + : uparams(p) + , distance_to_entry(dte) + , k_u(2 * M_PI / p.lambda_u) {} + + /** + * @brief Overloaded operator() to compute magnetic field components. + * + * @param position_in_lab_frame Position vector in the lab frame. + * @return Kokkos::pair, ippl::Vector> + * Pair containing magnetic field and its derivative. + */ + KOKKOS_INLINE_FUNCTION Kokkos::pair, ippl::Vector> + operator()(const ippl::Vector& position_in_lab_frame) const noexcept { + using Kokkos::cos; + using Kokkos::cosh; + using Kokkos::exp; + using Kokkos::sin; + using Kokkos::sinh; + + Kokkos::pair, ippl::Vector> + ret; // Return pair containing magnetic field and its derivative. + ret.first = scalar(0); // Initialize magnetic field vector. + ret.second = scalar(0); // Initialize derivative vector. + + // If the position is before the undulator entry. + if (position_in_lab_frame[2] < distance_to_entry) { + scalar z_in_undulator = position_in_lab_frame[2] - distance_to_entry; + assert(z_in_undulator < 0); // Ensure we are in the correct region. + scalar scal = exp(-((k_u * z_in_undulator) * (k_u * z_in_undulator) + * 0.5)); // Gaussian decay factor. + + ret.second[0] = 0; // No x-component. + ret.second[1] = uparams.B_magnitude * cosh(k_u * position_in_lab_frame[1]) + * z_in_undulator * k_u * scal; // y-component. + ret.second[2] = uparams.B_magnitude * sinh(k_u * position_in_lab_frame[1]) + * scal; // z-component. + } + // If the position is within the undulator. + else if (position_in_lab_frame[2] > distance_to_entry + && position_in_lab_frame[2] < distance_to_entry + uparams.length) { + scalar z_in_undulator = position_in_lab_frame[2] - distance_to_entry; + assert(z_in_undulator >= 0); // Ensure we are in the correct region. + + ret.second[0] = 0; // No x-component. + ret.second[1] = uparams.B_magnitude * cosh(k_u * position_in_lab_frame[1]) + * sin(k_u * z_in_undulator); // y-component. + ret.second[2] = uparams.B_magnitude * sinh(k_u * position_in_lab_frame[1]) + * cos(k_u * z_in_undulator); // z-component. + } + // If the position is past the undulator exit, ramp the field down with + // the same Gaussian-damped linear profile as the entrance fringe (mirror + // of MITHRA's staticUndulator exit branch, beam.cc). Without this the + // field terminates abruptly at the end of the undulator, leaving each + // electron with its residual transverse wiggle momentum (~K). That + // uncompensated transverse kick ejects the beam from the domain. + else if (position_in_lab_frame[2] >= distance_to_entry + uparams.length) { + scalar z_past_exit = + position_in_lab_frame[2] - (distance_to_entry + uparams.length); + assert(z_past_exit >= 0); // Ensure we are in the correct region. + scalar scal = exp(-((k_u * z_past_exit) * (k_u * z_past_exit) + * 0.5)); // Gaussian decay factor. + + ret.second[0] = 0; // No x-component. + ret.second[1] = uparams.B_magnitude * cosh(k_u * position_in_lab_frame[1]) + * z_past_exit * k_u * scal; // y-component. + ret.second[2] = uparams.B_magnitude * sinh(k_u * position_in_lab_frame[1]) + * scal; // z-component. + } + return ret; + } + }; +} // namespace ippl +#endif // UNDULATOR_H \ No newline at end of file diff --git a/fel/VideoWriter.h b/fel/VideoWriter.h new file mode 100644 index 000000000..01cd09832 --- /dev/null +++ b/fel/VideoWriter.h @@ -0,0 +1,254 @@ +#ifndef IPPL_FEL_VIDEO_WRITER_H +#define IPPL_FEL_VIDEO_WRITER_H + +// Poynting-flux visualization for the FEL simulation. +// +// Renders a slice of the (lab-frame) Poynting field plus the projected particle +// positions to a BMP frame and pipes the stream to ffmpeg. Ported from the +// original FreeElectronLaser.cpp main loop and encapsulated here so the manager +// and driver stay readable. stb_image_write's implementation is emitted by the +// translation unit that defines STB_IMAGE_WRITE_IMPLEMENTATION before including +// this header (FreeElectronLaser.cpp). + +#include +#include +#include // popen / pclose / FILE +#include +#include +#include + +#include +#include + +#include "Config.h" +#include "LorentzTransform.h" +#include "datatypes.h" +#include "units.h" + +// Google "Turbo" colormap (Anton Mikhailov), 256 RGB entries in [0,1]. +inline constexpr float turbo_cm[256][3] = { + {0.18995, 0.07176, 0.23217}, {0.19483, 0.08339, 0.26149}, {0.19956, 0.09498, 0.29024}, + {0.20415, 0.10652, 0.31844}, {0.20860, 0.11802, 0.34607}, {0.21291, 0.12947, 0.37314}, + {0.21708, 0.14087, 0.39964}, {0.22111, 0.15223, 0.42558}, {0.22500, 0.16354, 0.45096}, + {0.22875, 0.17481, 0.47578}, {0.23236, 0.18603, 0.50004}, {0.23582, 0.19720, 0.52373}, + {0.23915, 0.20833, 0.54686}, {0.24234, 0.21941, 0.56942}, {0.24539, 0.23044, 0.59142}, + {0.24830, 0.24143, 0.61286}, {0.25107, 0.25237, 0.63374}, {0.25369, 0.26327, 0.65406}, + {0.25618, 0.27412, 0.67381}, {0.25853, 0.28492, 0.69300}, {0.26074, 0.29568, 0.71162}, + {0.26280, 0.30639, 0.72968}, {0.26473, 0.31706, 0.74718}, {0.26652, 0.32768, 0.76412}, + {0.26816, 0.33825, 0.78050}, {0.26967, 0.34878, 0.79631}, {0.27103, 0.35926, 0.81156}, + {0.27226, 0.36970, 0.82624}, {0.27334, 0.38008, 0.84037}, {0.27429, 0.39043, 0.85393}, + {0.27509, 0.40072, 0.86692}, {0.27576, 0.41097, 0.87936}, {0.27628, 0.42118, 0.89123}, + {0.27667, 0.43134, 0.90254}, {0.27691, 0.44145, 0.91328}, {0.27701, 0.45152, 0.92347}, + {0.27698, 0.46153, 0.93309}, {0.27680, 0.47151, 0.94214}, {0.27648, 0.48144, 0.95064}, + {0.27603, 0.49132, 0.95857}, {0.27543, 0.50115, 0.96594}, {0.27469, 0.51094, 0.97275}, + {0.27381, 0.52069, 0.97899}, {0.27273, 0.53040, 0.98461}, {0.27106, 0.54015, 0.98930}, + {0.26878, 0.54995, 0.99303}, {0.26592, 0.55979, 0.99583}, {0.26252, 0.56967, 0.99773}, + {0.25862, 0.57958, 0.99876}, {0.25425, 0.58950, 0.99896}, {0.24946, 0.59943, 0.99835}, + {0.24427, 0.60937, 0.99697}, {0.23874, 0.61931, 0.99485}, {0.23288, 0.62923, 0.99202}, + {0.22676, 0.63913, 0.98851}, {0.22039, 0.64901, 0.98436}, {0.21382, 0.65886, 0.97959}, + {0.20708, 0.66866, 0.97423}, {0.20021, 0.67842, 0.96833}, {0.19326, 0.68812, 0.96190}, + {0.18625, 0.69775, 0.95498}, {0.17923, 0.70732, 0.94761}, {0.17223, 0.71680, 0.93981}, + {0.16529, 0.72620, 0.93161}, {0.15844, 0.73551, 0.92305}, {0.15173, 0.74472, 0.91416}, + {0.14519, 0.75381, 0.90496}, {0.13886, 0.76279, 0.89550}, {0.13278, 0.77165, 0.88580}, + {0.12698, 0.78037, 0.87590}, {0.12151, 0.78896, 0.86581}, {0.11639, 0.79740, 0.85559}, + {0.11167, 0.80569, 0.84525}, {0.10738, 0.81381, 0.83484}, {0.10357, 0.82177, 0.82437}, + {0.10026, 0.82955, 0.81389}, {0.09750, 0.83714, 0.80342}, {0.09532, 0.84455, 0.79299}, + {0.09377, 0.85175, 0.78264}, {0.09287, 0.85875, 0.77240}, {0.09267, 0.86554, 0.76230}, + {0.09320, 0.87211, 0.75237}, {0.09451, 0.87844, 0.74265}, {0.09662, 0.88454, 0.73316}, + {0.09958, 0.89040, 0.72393}, {0.10342, 0.89600, 0.71500}, {0.10815, 0.90142, 0.70599}, + {0.11374, 0.90673, 0.69651}, {0.12014, 0.91193, 0.68660}, {0.12733, 0.91701, 0.67627}, + {0.13526, 0.92197, 0.66556}, {0.14391, 0.92680, 0.65448}, {0.15323, 0.93151, 0.64308}, + {0.16319, 0.93609, 0.63137}, {0.17377, 0.94053, 0.61938}, {0.18491, 0.94484, 0.60713}, + {0.19659, 0.94901, 0.59466}, {0.20877, 0.95304, 0.58199}, {0.22142, 0.95692, 0.56914}, + {0.23449, 0.96065, 0.55614}, {0.24797, 0.96423, 0.54303}, {0.26180, 0.96765, 0.52981}, + {0.27597, 0.97092, 0.51653}, {0.29042, 0.97403, 0.50321}, {0.30513, 0.97697, 0.48987}, + {0.32006, 0.97974, 0.47654}, {0.33517, 0.98234, 0.46325}, {0.35043, 0.98477, 0.45002}, + {0.36581, 0.98702, 0.43688}, {0.38127, 0.98909, 0.42386}, {0.39678, 0.99098, 0.41098}, + {0.41229, 0.99268, 0.39826}, {0.42778, 0.99419, 0.38575}, {0.44321, 0.99551, 0.37345}, + {0.45854, 0.99663, 0.36140}, {0.47375, 0.99755, 0.34963}, {0.48879, 0.99828, 0.33816}, + {0.50362, 0.99879, 0.32701}, {0.51822, 0.99910, 0.31622}, {0.53255, 0.99919, 0.30581}, + {0.54658, 0.99907, 0.29581}, {0.56026, 0.99873, 0.28623}, {0.57357, 0.99817, 0.27712}, + {0.58646, 0.99739, 0.26849}, {0.59891, 0.99638, 0.26038}, {0.61088, 0.99514, 0.25280}, + {0.62233, 0.99366, 0.24579}, {0.63323, 0.99195, 0.23937}, {0.64362, 0.98999, 0.23356}, + {0.65394, 0.98775, 0.22835}, {0.66428, 0.98524, 0.22370}, {0.67462, 0.98246, 0.21960}, + {0.68494, 0.97941, 0.21602}, {0.69525, 0.97610, 0.21294}, {0.70553, 0.97255, 0.21032}, + {0.71577, 0.96875, 0.20815}, {0.72596, 0.96470, 0.20640}, {0.73610, 0.96043, 0.20504}, + {0.74617, 0.95593, 0.20406}, {0.75617, 0.95121, 0.20343}, {0.76608, 0.94627, 0.20311}, + {0.77591, 0.94113, 0.20310}, {0.78563, 0.93579, 0.20336}, {0.79524, 0.93025, 0.20386}, + {0.80473, 0.92452, 0.20459}, {0.81410, 0.91861, 0.20552}, {0.82333, 0.91253, 0.20663}, + {0.83241, 0.90627, 0.20788}, {0.84133, 0.89986, 0.20926}, {0.85010, 0.89328, 0.21074}, + {0.85868, 0.88655, 0.21230}, {0.86709, 0.87968, 0.21391}, {0.87530, 0.87267, 0.21555}, + {0.88331, 0.86553, 0.21719}, {0.89112, 0.85826, 0.21880}, {0.89870, 0.85087, 0.22038}, + {0.90605, 0.84337, 0.22188}, {0.91317, 0.83576, 0.22328}, {0.92004, 0.82806, 0.22456}, + {0.92666, 0.82025, 0.22570}, {0.93301, 0.81236, 0.22667}, {0.93909, 0.80439, 0.22744}, + {0.94489, 0.79634, 0.22800}, {0.95039, 0.78823, 0.22831}, {0.95560, 0.78005, 0.22836}, + {0.96049, 0.77181, 0.22811}, {0.96507, 0.76352, 0.22754}, {0.96931, 0.75519, 0.22663}, + {0.97323, 0.74682, 0.22536}, {0.97679, 0.73842, 0.22369}, {0.98000, 0.73000, 0.22161}, + {0.98289, 0.72140, 0.21918}, {0.98549, 0.71250, 0.21650}, {0.98781, 0.70330, 0.21358}, + {0.98986, 0.69382, 0.21043}, {0.99163, 0.68408, 0.20706}, {0.99314, 0.67408, 0.20348}, + {0.99438, 0.66386, 0.19971}, {0.99535, 0.65341, 0.19577}, {0.99607, 0.64277, 0.19165}, + {0.99654, 0.63193, 0.18738}, {0.99675, 0.62093, 0.18297}, {0.99672, 0.60977, 0.17842}, + {0.99644, 0.59846, 0.17376}, {0.99593, 0.58703, 0.16899}, {0.99517, 0.57549, 0.16412}, + {0.99419, 0.56386, 0.15918}, {0.99297, 0.55214, 0.15417}, {0.99153, 0.54036, 0.14910}, + {0.98987, 0.52854, 0.14398}, {0.98799, 0.51667, 0.13883}, {0.98590, 0.50479, 0.13367}, + {0.98360, 0.49291, 0.12849}, {0.98108, 0.48104, 0.12332}, {0.97837, 0.46920, 0.11817}, + {0.97545, 0.45740, 0.11305}, {0.97234, 0.44565, 0.10797}, {0.96904, 0.43399, 0.10294}, + {0.96555, 0.42241, 0.09798}, {0.96187, 0.41093, 0.09310}, {0.95801, 0.39958, 0.08831}, + {0.95398, 0.38836, 0.08362}, {0.94977, 0.37729, 0.07905}, {0.94538, 0.36638, 0.07461}, + {0.94084, 0.35566, 0.07031}, {0.93612, 0.34513, 0.06616}, {0.93125, 0.33482, 0.06218}, + {0.92623, 0.32473, 0.05837}, {0.92105, 0.31489, 0.05475}, {0.91572, 0.30530, 0.05134}, + {0.91024, 0.29599, 0.04814}, {0.90463, 0.28696, 0.04516}, {0.89888, 0.27824, 0.04243}, + {0.89298, 0.26981, 0.03993}, {0.88691, 0.26152, 0.03753}, {0.88066, 0.25334, 0.03521}, + {0.87422, 0.24526, 0.03297}, {0.86760, 0.23730, 0.03082}, {0.86079, 0.22945, 0.02875}, + {0.85380, 0.22170, 0.02677}, {0.84662, 0.21407, 0.02487}, {0.83926, 0.20654, 0.02305}, + {0.83172, 0.19912, 0.02131}, {0.82399, 0.19182, 0.01966}, {0.81608, 0.18462, 0.01809}, + {0.80799, 0.17753, 0.01660}, {0.79971, 0.17055, 0.01520}, {0.79125, 0.16368, 0.01387}, + {0.78260, 0.15693, 0.01264}, {0.77377, 0.15028, 0.01148}, {0.76476, 0.14374, 0.01041}, + {0.75556, 0.13731, 0.00942}, {0.74617, 0.13098, 0.00851}, {0.73661, 0.12477, 0.00769}, + {0.72686, 0.11867, 0.00695}, {0.71692, 0.11268, 0.00629}, {0.70680, 0.10680, 0.00571}, + {0.69650, 0.10102, 0.00522}, {0.68602, 0.09536, 0.00481}, {0.67535, 0.08980, 0.00449}, + {0.66449, 0.08436, 0.00424}, {0.65345, 0.07902, 0.00408}, {0.64223, 0.07380, 0.00401}, + {0.63082, 0.06868, 0.00401}, {0.61923, 0.06367, 0.00410}, {0.60746, 0.05878, 0.00427}, + {0.59550, 0.05399, 0.00453}, {0.58336, 0.04931, 0.00486}, {0.57103, 0.04474, 0.00529}, + {0.55852, 0.04028, 0.00579}, {0.54583, 0.03593, 0.00638}, {0.53295, 0.03169, 0.00705}, + {0.51989, 0.02756, 0.00780}, {0.50664, 0.02354, 0.00863}, {0.49321, 0.01963, 0.00955}, + {0.47960, 0.01583, 0.01055}}; + +// Write an RGB buffer as a BMP to an open FILE* (the ffmpeg pipe). +inline bool writeBMPToFD(FILE* fd, int width, int height, const unsigned char* data) { + const int channels = 3; // RGB + const int stride = width * channels; + std::vector flippedData(data, data + stride * height); + + if (!stbi_write_bmp_to_func( + [](void* context, void* data, int size) { + FILE* f = reinterpret_cast(context); + fwrite(data, 1, size, f); + }, + fd, width, height, channels, flippedData.data())) { + return false; + } + return true; +} + +// Encapsulates the ffmpeg pipe and per-frame rendering. +template +class FELVideoWriter { +public: + // Open the ffmpeg pipe on rank 0 when periodic output is requested. + void open(const config& cfg) { + if (ippl::Comm->rank() == 0 && cfg.output_rhythm != 0) { + const char* ffmpegCmd = + "ffmpeg -y -f image2pipe -vcodec bmp -r 30 -i - -vf " + "scale=force_original_aspect_ratio=decrease:force_divisible_by=2,format=yuv420p " + "-c:v libx264 -movflags +faststart ffmpeg_popen.mkv"; + ffmpeg_file_ = popen(ffmpegCmd, "w"); + } + } + + void close() { + if (ffmpeg_file_ != nullptr) { + pclose(ffmpeg_file_); + ffmpeg_file_ = nullptr; + } + } + + // Render and emit one frame for time step `it`. Collective over all ranks + // (the per-rank partial images are reduced onto rank 0 before encoding). + template + void writeFrame(int it, FieldContainer_t& fc, ParticleContainer_t& pc, const Frame_t& frame, + const config& cfg, const Vector_t& nr) { + const int rank = ippl::Comm->rank(); + const int size = ippl::Comm->size(); + + const int img_height = 400; + const int img_width = int(400.0 * cfg.extents[2] / cfg.extents[0]); + const int floatcount = img_width * img_height * 3; + + std::vector imagedata(floatcount, 0.0f); + std::vector recvbuffer(floatcount, 0.0f); + + auto& mesh = fc.getMesh(); + const auto origin = mesh.getOrigin(); + const auto ldom = fc.getFL().getLocalNDIndex(); + + // Project local particles onto the (z, x) plane as green dots. + auto phmirror = pc.R.getHostMirror(); + Kokkos::deep_copy(phmirror, pc.R.getView()); + for (size_t hi = 0; hi < pc.getLocalNum(); hi++) { + ippl::Vector ppos = phmirror(hi); + ppos -= origin; + ppos /= vector_cast(cfg.extents); + int x_imgcoord = ppos[2] * img_width; + int y_imgcoord = ppos[0] * img_height; + if (y_imgcoord >= 0 && x_imgcoord >= 0 && x_imgcoord < img_width + && y_imgcoord < img_height) { + const float intensity = + std::min(255.f, (img_width * img_height * 15.f) / cfg.num_particles); + float& g = imagedata[(y_imgcoord * img_width + x_imgcoord) * 3 + 1]; + g = std::min(255.f, g + intensity); + } + } + + // Color the lab-frame Poynting magnitude on the mid-y plane. + auto eh = fc.getE().getHostMirror(); + auto bh = fc.getB().getHostMirror(); + Kokkos::deep_copy(eh, fc.getE().getView()); + Kokkos::deep_copy(bh, fc.getB().getView()); + + for (int i = 1; i < img_width; i++) { + for (int j = 1; j < img_height; j++) { + int i_remap = (double(i) / (img_width - 1)) * (nr[2] - 4) + 2; + int j_remap = (double(j) / (img_height - 1)) * (nr[0] - 4) + 2; + if (i_remap >= ldom.first()[2] && i_remap <= ldom.last()[2] + && j_remap >= ldom.first()[0] && j_remap <= ldom.last()[0]) { + ippl::Vector E = eh(j_remap + 1 - ldom.first()[0], nr[1] / 2, + i_remap + 1 - ldom.first()[2]); + ippl::Vector B = bh(j_remap + 1 - ldom.first()[0], nr[1] / 2, + i_remap + 1 - ldom.first()[2]); + + auto eblab = frame.inverse_transform_EB( + Kokkos::make_pair, ippl::Vector>( + ippl::Vector(E), ippl::Vector(B))); + ippl::Vector poynting = ippl::cross(eblab.first, eblab.second); + + float normalized = std::sqrt(poynting.Pnorm()) * 0.00001f; + int index = (int)std::max(0.0f, std::min(normalized * 255.0f, 255.0f)); + for (int c = 0; c < 3; c++) { + imagedata[(j * img_width + i) * 3 + c] += turbo_cm[index][c] * 255.0f; + } + } + } + } + + // Butterfly reduction of the partial images onto rank 0. + int mask = 1; + while (mask < size) { + int partner = rank ^ mask; + if ((rank & mask) == 0) { + MPI_Recv(recvbuffer.data(), floatcount, MPI_FLOAT, partner, 0, + ippl::Comm->getCommunicator(), MPI_STATUS_IGNORE); + for (int f = 0; f < floatcount; f++) { + imagedata[f] += recvbuffer[f]; + } + } else { + MPI_Send(imagedata.data(), floatcount, MPI_FLOAT, partner, 0, + ippl::Comm->getCommunicator()); + } + mask <<= 1; + } + + if (rank == 0 && ffmpeg_file_ != nullptr) { + std::vector final_img(floatcount); + std::transform(imagedata.begin(), imagedata.end(), final_img.begin(), + [](float x) { return (uint8_t)std::min(255.0f, std::max(0.0f, x)); }); + writeBMPToFD(ffmpeg_file_, img_width, img_height, final_img.data()); + } + (void)it; + } + +private: + FILE* ffmpeg_file_ = nullptr; +}; + +#endif diff --git a/fel/config.json b/fel/config.json new file mode 100644 index 000000000..d633a7514 --- /dev/null +++ b/fel/config.json @@ -0,0 +1,32 @@ +{ + "timestep-ratio": 1.0, + "mesh": { + "length-scale": "micrometers", + "time-scale": "picoseconds", + "extents": [3200.0, 3200.0, 280.0], + "resolution": [96, 96, 3000], + "total-time": 30000.0, + "space-charge": false + }, + "bunch": { + "charge": 1.846e8, + "mass": 1.846e8, + "number-of-particles": 131072, + "gamma": 100.41, + "position": [0.0, 0.0, 0.0], + "sigma-position": [260.0, 260.0, 50.25], + "sigma-momentum": [1.0e-8, 1.0e-8, 100.41e-4], + "distribution-truncations": [1040.0, 1040.0, 90.0] + }, + "undulator": { + "static-undulator": { + "undulator-parameter": 1.417, + "period": 30000.0, + "length": 9e6 + } + }, + "output": { + "rhythm": 30, + "path": "../renderdata" + } +} diff --git a/fel/datatypes.h b/fel/datatypes.h new file mode 100644 index 000000000..b50561314 --- /dev/null +++ b/fel/datatypes.h @@ -0,0 +1,65 @@ +#ifndef IPPL_FEL_DATATYPES_H +#define IPPL_FEL_DATATYPES_H + +#include "MaxwellSolvers/StandardFDTDSolver.h" +#include "MaxwellSolvers/NonStandardFDTDSolver.h" + +// Type aliases for the FEL module, mirroring alpine/datatypes.h. +// +// The FEL simulation solves Maxwell's equations with an FDTD solver coupled to +// a relativistic electron bunch. Two field flavours are needed: +// * VField_t : E and B, three-component vector fields. +// * SourceField_t : the four-current source (rho, Jx, Jy, Jz) + +template +using Mesh_t = ippl::UniformCartesian; + +template +using PLayout_t = typename ippl::ParticleSpatialLayout>; + +template +using Centering_t = typename Mesh_t::DefaultCentering; + +template +using FieldLayout_t = ippl::FieldLayout; + +using size_type = ippl::detail::size_type; + +template +using Vector_t = ippl::Vector; + +template +using Field = ippl::Field, Centering_t, ViewArgs...>; + +// Electric and magnetic fields: three-component vector field on the mesh. +template +using VField_t = Field, Dim, ViewArgs...>; + +// Four-current source field: [0] = charge density +// [1..Dim] = current density. +template +using SourceField_t = Field, Dim, ViewArgs...>; + +// The Maxwell FDTD solver used by the FEL simulation. Boundary conditions are +// absorbing (second-order Mur) so radiation leaves the domain cleanly. +// +// Both standard finite difference and the mithra-style non-standard finite difference +// are available. +template +using FDTDSolver_t = + ippl::NonStandardFDTDSolver, SourceField_t, ippl::absorbing>; + +//using FDTDSolver_t = +// ippl::StandardFDTDSolver, SourceField_t, ippl::absorbing>; + +// Component-wise cast of a Vector to another scalar type. +template +KOKKOS_INLINE_FUNCTION ippl::Vector vector_cast(const ippl::Vector& v) { + ippl::Vector ret; + for (unsigned k = 0; k < D; ++k) { + ret[k] = static_cast(v[k]); + } + return ret; +} + +#endif diff --git a/fel/units.h b/fel/units.h new file mode 100644 index 000000000..318892db8 --- /dev/null +++ b/fel/units.h @@ -0,0 +1,37 @@ +#ifndef UNITS_HPP +#define UNITS_HPP +constexpr double sqrt_4pi = 3.54490770181103205459; +constexpr double alpha_scaling_factor = 1e30; +constexpr double unit_length_in_meters = 1.616255e-35 * alpha_scaling_factor; +constexpr double unit_charge_in_electron_charges = 11.70623710394218618969 / sqrt_4pi; +constexpr double unit_time_in_seconds = 5.391247e-44 * alpha_scaling_factor; +constexpr double electron_mass_in_kg = 9.1093837015e-31; +constexpr double unit_mass_in_kg = 2.176434e-8 / alpha_scaling_factor; +constexpr double unit_energy_in_joule = unit_mass_in_kg * unit_length_in_meters * unit_length_in_meters / (unit_time_in_seconds * unit_time_in_seconds); +constexpr double kg_in_unit_masses = 1.0 / unit_mass_in_kg; +constexpr double meter_in_unit_lengths = 1.0 / unit_length_in_meters; +constexpr double electron_charge_in_coulombs = 1.602176634e-19; +constexpr double coulomb_in_electron_charges = 1.0 / electron_charge_in_coulombs; + +constexpr double electron_charge_in_unit_charges = 1.0 / unit_charge_in_electron_charges; +constexpr double second_in_unit_times = 1.0 / unit_time_in_seconds; +constexpr double electron_mass_in_unit_masses = electron_mass_in_kg * kg_in_unit_masses; +constexpr double unit_force_in_newtons = unit_mass_in_kg * unit_length_in_meters / (unit_time_in_seconds * unit_time_in_seconds); + +constexpr double coulomb_in_unit_charges = coulomb_in_electron_charges * electron_charge_in_unit_charges; +constexpr double unit_voltage_in_volts = unit_energy_in_joule * coulomb_in_unit_charges; +constexpr double unit_charges_in_coulomb = 1.0 / coulomb_in_unit_charges; +constexpr double unit_current_in_amperes = unit_charges_in_coulomb / unit_time_in_seconds; +constexpr double ampere_in_unit_currents = 1.0 / unit_current_in_amperes; +constexpr double unit_current_length_in_ampere_meters = unit_current_in_amperes * unit_length_in_meters; +constexpr double unit_magnetic_fluxdensity_in_tesla = unit_voltage_in_volts * unit_time_in_seconds / (unit_length_in_meters * unit_length_in_meters); +constexpr double unit_electric_fieldstrength_in_voltpermeters = (unit_voltage_in_volts / unit_length_in_meters); +constexpr double voltpermeter_in_unit_fieldstrengths = 1.0 / unit_electric_fieldstrength_in_voltpermeters; +constexpr double unit_powerdensity_in_watt_per_square_meter = 1.389e122 / (alpha_scaling_factor * alpha_scaling_factor * alpha_scaling_factor * alpha_scaling_factor); +constexpr double volts_in_unit_voltages = 1.0 / unit_voltage_in_volts; +constexpr double epsilon0_in_si = unit_current_in_amperes * unit_time_in_seconds / (unit_voltage_in_volts * unit_length_in_meters); +constexpr double mu0_in_si = unit_force_in_newtons / (unit_current_in_amperes * unit_current_in_amperes); +constexpr double G = unit_length_in_meters * unit_length_in_meters * unit_length_in_meters / (unit_mass_in_kg * unit_time_in_seconds * unit_time_in_seconds); +constexpr double verification_gravity = unit_mass_in_kg * unit_mass_in_kg / (unit_length_in_meters * unit_length_in_meters) * G; +constexpr double verification_coulomb = (unit_charges_in_coulomb * unit_charges_in_coulomb / (unit_length_in_meters * unit_length_in_meters) * (1.0 / (epsilon0_in_si))) / unit_force_in_newtons; +#endif \ No newline at end of file diff --git a/src/Interpolation/CurrentDeposition.hpp b/src/Interpolation/CurrentDeposition.hpp index bc808777a..e62812d68 100644 --- a/src/Interpolation/CurrentDeposition.hpp +++ b/src/Interpolation/CurrentDeposition.hpp @@ -1,10 +1,110 @@ -#ifndef IPPL_PROJECT_CURRENT_FDTD_H -#define IPPL_PROJECT_CURRENT_FDTD_H +#ifndef IPPL_CURRENT_DEPOSITION_H +#define IPPL_CURRENT_DEPOSITION_H #include "FEM/GridPathSegmenter.h" namespace ippl { +/** + * @brief Deposit current density onto a co-located (non-staggered) grid. + * + * For each particle p moving from X0(p) to X1(p) during one time step dt, + * the trajectory is split into sub-segments each lying within a single mesh + * cell (via GridPathSegmenter). For each segment, the contribution is + * scattered onto all 2^Dim surrounding grid nodes using CIC weights + * evaluated at the midpoint (all J components share the same grid + * location on a co-located grid). + * + * Writes q*dp[c]/(dt*cell_volume) into component [c+1] of the field vector. + * Component [0] is reserved for the scalar potential / charge density. + * This matches the co-located 4-potential FDTD formulation where the source + * field stores (phi, Jx, Jy, Jz) at each node. + */ +template > +inline void assemble_current_collocated(const Mesh& mesh, + const ChargeAttrib& q_attrib, + const PosAttrib& X0, + const PosAttrib& X1, + JField& J_field, + policy_type iteration_policy, + typename Mesh::value_type dt) +{ + using T = typename Mesh::value_type; + constexpr unsigned Dim = Mesh::Dimension; + static_assert(Dim == 2 || Dim == 3, + "assemble_current_collocated only supports 2D and 3D"); + + const auto origin = mesh.getOrigin(); + const auto h = mesh.getMeshSpacing(); + auto ldom = J_field.getLayout().getLocalNDIndex(); + const int nghost = J_field.getNghost(); + auto view = J_field.getView(); + + T volume = T(1); + for (unsigned d = 0; d < Dim; ++d) + volume *= h[d]; + + Kokkos::parallel_for("assemble_current_collocated", iteration_policy, + KOKKOS_LAMBDA(const std::size_t p) { + + auto segs = GridPathSegmenter + ::split(X0(p), X1(p), origin, h); + + const T q_over_dt_vol = q_attrib(p) / (dt * volume); + + for (unsigned i = 0; i < Dim + 1; ++i) { + const auto& seg = segs[i]; + + Vector dp{}; + T len_sq = T(0); + for (unsigned d = 0; d < Dim; ++d) { + dp[d] = seg.p1[d] - seg.p0[d]; + len_sq += dp[d] * dp[d]; + } + if (len_sq == T(0)) continue; + + Vector mid{}; + for (unsigned d = 0; d < Dim; ++d) + mid[d] = T(0.5) * (seg.p0[d] + seg.p1[d]); + + // Half-cell shift so the deposit matches the field centering + // (UniformCartesian DefaultCentering == Cell) and IPPL's gather, + // which evaluates CIC weights at (x-origin)/h + 0.5. Without it the + // deposited current sits half a cell from where E,B are gathered, + // breaking scatter/gather reciprocity (spurious self-force). + Kokkos::Array cellIdx; + Kokkos::Array xi; + for (unsigned d = 0; d < Dim; ++d) { + const T gp = (mid[d] - origin[d]) / h[d] - T(0.5); + cellIdx[d] = static_cast(Kokkos::floor(gp)); + xi[d] = gp - T(cellIdx[d]); + } + + // Scatter to all 2^Dim corners with CIC weights. + // All J components share the same weight and index set because + // the grid is co-located (no per-component staggering). + for (unsigned corner = 0; corner < (1u << Dim); ++corner) { + size_t idx[Dim]; + T weight = T(1); + for (unsigned d = 0; d < Dim; ++d) { + const unsigned offset = (corner >> d) & 1u; + weight *= offset ? xi[d] : (T(1) - xi[d]); + idx[d] = cellIdx[d] - ldom.first()[d] + nghost + offset; + } + + for (unsigned c = 0; c < Dim; ++c) { + Kokkos::atomic_add(&(apply(view, idx)[c + 1]), + q_over_dt_vol * dp[c] * weight); + } + } + } + }); +} + /** * @brief Deposit current density onto a Yee-staggered grid. * diff --git a/unit_tests/Interpolation/TestCurrentDeposition.cpp b/unit_tests/Interpolation/TestCurrentDeposition.cpp index 53a627001..261ece293 100644 --- a/unit_tests/Interpolation/TestCurrentDeposition.cpp +++ b/unit_tests/Interpolation/TestCurrentDeposition.cpp @@ -15,6 +15,46 @@ struct Bunch : public ippl::ParticleBase { typename PLayout::particle_position_type R_next; }; +template +class AssembleCurrentTest; + +template +class AssembleCurrentTest>> : public ::testing::Test { +public: + using value_type = T; + static constexpr unsigned dim = Dim; + + using Mesh_t = ippl::UniformCartesian; + using Centering_t = typename Mesh_t::DefaultCentering; + // Dim+1 components: [0] = scalar potential / charge density, [1..Dim] = Jx, Jy, Jz. + using JField_t = ippl::Field, Dim, Mesh_t, Centering_t>; + using Layout_t = ippl::FieldLayout; + + using playout_t = ippl::ParticleSpatialLayout; + using bunch_t = Bunch; + + static ippl::NDIndex make_owned_nd(int nx) { + ippl::Index I0(nx); + if constexpr (Dim == 1) return ippl::NDIndex<1>(I0); + else if constexpr (Dim == 2) return ippl::NDIndex<2>(I0, I0); + else return ippl::NDIndex<3>(I0, I0, I0); + } + + static Layout_t make_layout(const ippl::NDIndex& owned) { + std::array par{}; par.fill(true); + return ippl::FieldLayout(MPI_COMM_WORLD, owned, par); + } + + static Mesh_t make_mesh(const ippl::NDIndex& owned, + const ippl::Vector& h, + const ippl::Vector& origin) { + return Mesh_t(owned, h, origin); + } +}; + +// Fixture for the Yee-staggered deposition (assemble_current_yee). Uses a +// Dim-component J field (index [c]), in contrast to the co-located fixture +// above which reserves component [0] for the scalar slot (index [c+1]). template class AssembleCurrentYeeTest; @@ -55,8 +95,340 @@ using Precisions = TestParams::Precisions; using Ranks = TestParams::Ranks<2, 3>; using Tests = TestForTypes::type>::type; +TYPED_TEST_SUITE(AssembleCurrentTest, Tests); TYPED_TEST_SUITE(AssembleCurrentYeeTest, Tests); +// Pure x-motion inside a single cell. The midpoint lies at the mesh cell centre. +// assemble_current_collocated applies the same half-cell shift used by the +// centered field gather, so this point lands exactly on node (0,...,0). +// Verifies: +// - Jx is deposited at the shifted node location. +// - Jy (and Jz in 3D) remain exactly zero. +// - The scalar slot field[0] is untouched (remains zero). +TYPED_TEST(AssembleCurrentTest, SingleAxis_X_SameCell_ExactValues) { + using T = typename TestFixture::value_type; + constexpr unsigned Dim = TestFixture::dim; + + using bunch_t = typename TestFixture::bunch_t; + using playout_t = typename TestFixture::playout_t; + using JField_t = typename TestFixture::JField_t; + + int nx = 32; + ippl::Vector origin(0.0); + ippl::Vector h(1.0); + + auto owned = TestFixture::make_owned_nd(nx); + auto layout = TestFixture::make_layout(owned); + auto mesh = TestFixture::make_mesh(owned, h, origin); + + JField_t J_field(mesh, layout); + J_field = T(0); + + playout_t playout(layout, mesh); + bunch_t bunch(playout); + bunch.create(ippl::Comm->rank() == 0 ? 1 : 0); + if (ippl::Comm->rank() == 0) { + auto R_host = bunch.R.getHostMirror(); + auto Rn_host = bunch.R_next.getHostMirror(); + auto Q_host = bunch.Q.getHostMirror(); + R_host(0)[0] = T(0.25); + Rn_host(0)[0] = T(0.75); + for (unsigned d = 1; d < Dim; ++d) { R_host(0)[d] = T(0.50); Rn_host(0)[d] = T(0.50); } + Q_host(0) = T(1.0); + Kokkos::deep_copy(bunch.R.getView(), R_host); + Kokkos::deep_copy(bunch.R_next.getView(), Rn_host); + Kokkos::deep_copy(bunch.Q.getView(), Q_host); + } + bunch.update(); + + auto policy = Kokkos::RangePolicy<>(0, bunch.getLocalNum()); + T dt = T(1.0); + ippl::assemble_current_collocated(mesh, bunch.Q, bunch.R, bunch.R_next, + J_field, policy, dt); + J_field.accumulateHalo(); + + auto ldom = J_field.getLayout().getLocalNDIndex(); + int nghost = J_field.getNghost(); + auto view = J_field.getView(); + auto view_host = Kokkos::create_mirror_view(view); + Kokkos::deep_copy(view_host, view); + Kokkos::fence(); + + const T tol = std::numeric_limits::epsilon() * T(100); + + auto in_ldom = [&](const Kokkos::Array& cell) { + for (unsigned d = 0; d < Dim; ++d) + if (cell[d] < ldom.first()[d] || cell[d] > ldom.last()[d]) return false; + return true; + }; + + // c is the 0-indexed spatial component (Jx=0, Jy=1, Jz=2). + // Field component is c+1 because index 0 is the scalar slot. + auto check = [&](const Kokkos::Array& cell, unsigned c, double expected) { + double local_val = 0.0; + if (in_ldom(cell)) { + size_t idx[Dim]; + for (unsigned d = 0; d < Dim; ++d) + idx[d] = static_cast(cell[d] - ldom.first()[d] + nghost); + local_val = static_cast(ippl::apply(view_host, idx)[c + 1]); + } + double global_val = 0.0; + MPI_Allreduce(&local_val, &global_val, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + EXPECT_NEAR(global_val, expected, static_cast(tol)); + }; + + auto check_scalar = [&](const Kokkos::Array& cell) { + double local_val = 0.0; + if (in_ldom(cell)) { + size_t idx[Dim]; + for (unsigned d = 0; d < Dim; ++d) + idx[d] = static_cast(cell[d] - ldom.first()[d] + nghost); + local_val = static_cast(ippl::apply(view_host, idx)[0]); + } + double global_val = 0.0; + MPI_Allreduce(&local_val, &global_val, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + EXPECT_NEAR(global_val, 0.0, static_cast(tol)); + }; + + // One segment. Mid=(0.5,0.5,...), shifted gp=(0,0,...), xi=(0,0,...), + // dp=(0.5,0,...). The full Jx contribution lands on node (0,...,0). + if constexpr (Dim == 2) { + check({0, 0}, 0, 0.5); check({1, 0}, 0, 0.0); + check({0, 1}, 0, 0.0); check({1, 1}, 0, 0.0); + check({0, 0}, 1, 0.0); check({1, 0}, 1, 0.0); + check({0, 1}, 1, 0.0); check({1, 1}, 1, 0.0); + check_scalar({0, 0}); check_scalar({1, 0}); + check_scalar({0, 1}); check_scalar({1, 1}); + } else if constexpr (Dim == 3) { + for (int i = 0; i <= 1; ++i) + for (int j = 0; j <= 1; ++j) + for (int k = 0; k <= 1; ++k) { + check({i, j, k}, 0, (i == 0 && j == 0 && k == 0) ? 0.5 : 0.0); + check({i, j, k}, 1, 0.0); + check({i, j, k}, 2, 0.0); + check_scalar({i, j, k}); + } + } +} + +// Diagonal path crossing two cell boundaries (three segments). 2D only. +// +// Segments (GridPathSegmenter: x-cut t=1/3, y-cut t=2/3), evaluated with the +// half-cell shifted coordinate gp=mid/h-0.5: +// seg0: xi=(0.375,0.125) dp=(0.25,0.25) +// seg1: xi=(0.625,0.375) dp=(0.25,0.25) +// seg2: xi=(0.875,0.625) dp=(0.25,0.25) +// +// Because dp_x == dp_y for every segment, Jx == Jy at every node. +// With h=1 (volume=1), q_over_dt_vol = 1. +TYPED_TEST(AssembleCurrentTest, DiagonalPath_ThreeCells_ExactValues) { + using T = typename TestFixture::value_type; + constexpr unsigned Dim = TestFixture::dim; + + if constexpr (Dim != 2) { + GTEST_SKIP() << "Exact value check only implemented for 2D"; + } else { + using bunch_t = typename TestFixture::bunch_t; + using playout_t = typename TestFixture::playout_t; + using JField_t = typename TestFixture::JField_t; + + int nx = 32; + ippl::Vector origin(0.0); + ippl::Vector h(1.0); + + auto owned = TestFixture::make_owned_nd(nx); + auto layout = TestFixture::make_layout(owned); + auto mesh = TestFixture::make_mesh(owned, h, origin); + + JField_t J_field(mesh, layout); + J_field = T(0); + + playout_t playout(layout, mesh); + bunch_t bunch(playout); + // Three identical particles — exercises atomic accumulation; all nodal + // values must scale by 3 relative to the single-particle baseline. + bunch.create(ippl::Comm->rank() == 0 ? 3 : 0); + if (ippl::Comm->rank() == 0) { + auto R_host = bunch.R.getHostMirror(); + auto Rn_host = bunch.R_next.getHostMirror(); + auto Q_host = bunch.Q.getHostMirror(); + for (int p = 0; p < 3; ++p) { + R_host(p)[0] = T(0.75); R_host(p)[1] = T(0.50); + Rn_host(p)[0] = T(1.50); Rn_host(p)[1] = T(1.25); + Q_host(p) = T(1.0); + } + Kokkos::deep_copy(bunch.R.getView(), R_host); + Kokkos::deep_copy(bunch.R_next.getView(), Rn_host); + Kokkos::deep_copy(bunch.Q.getView(), Q_host); + } + bunch.update(); + + auto policy = Kokkos::RangePolicy<>(0, bunch.getLocalNum()); + T dt = T(1.0); + ippl::assemble_current_collocated(mesh, bunch.Q, bunch.R, bunch.R_next, + J_field, policy, dt); + J_field.accumulateHalo(); + + auto ldom = J_field.getLayout().getLocalNDIndex(); + int nghost = J_field.getNghost(); + auto view = J_field.getView(); + auto view_host = Kokkos::create_mirror_view(view); + Kokkos::deep_copy(view_host, view); + Kokkos::fence(); + + const T tol = std::numeric_limits::epsilon() * T(100); + + // Checks field[c+1] at global node (gi, gj). Broadcasts via SUM so + // every MPI rank asserts regardless of domain decomposition. + auto check = [&](int gi, int gj, unsigned c, double expected) { + double local_val = 0.0; + if (gi >= ldom.first()[0] && gi <= ldom.last()[0] && + gj >= ldom.first()[1] && gj <= ldom.last()[1]) { + int li = gi - ldom.first()[0] + nghost; + int lj = gj - ldom.first()[1] + nghost; + local_val = static_cast(view_host(li, lj)[c + 1]); + } + double global_val = 0.0; + MPI_Allreduce(&local_val, &global_val, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + EXPECT_NEAR(global_val, expected, static_cast(tol)); + }; + + // All single-particle values scaled by 3 (three identical particles). + check(0, 0, 0, 0.62109375); check(0, 0, 1, 0.62109375); + check(1, 0, 0, 0.78515625); check(1, 0, 1, 0.78515625); + check(0, 1, 0, 0.22265625); check(0, 1, 1, 0.22265625); + check(1, 1, 0, 0.62109375); check(1, 1, 1, 0.62109375); + check(2, 0, 0, 0.0); check(2, 0, 1, 0.0); + check(2, 1, 0, 0.0); check(2, 1, 1, 0.0); + check(1, 2, 0, 0.0); check(1, 2, 1, 0.0); + check(2, 2, 0, 0.0); check(2, 2, 1, 0.0); + } +} + +// 3D path where all three axes are crossed simultaneously at t=0.5 (vertex hit). +// Active segments after degenerate zero-length segments are skipped, evaluated +// with the half-cell shifted coordinate gp=mid/h-0.5: +// seg0: xi=(0.45,0.45,0.4) dp=(0.1,0.1,0.2) +// seg3: xi=(0.55,0.55,0.6) dp=(0.1,0.1,0.2) +// +// Symmetry used to reduce the number of assertions: +// Jy == Jx everywhere (dp_y == dp_x) +// Jz == 2*Jx everywhere (dp_z == 2*dp_x) +TYPED_TEST(AssembleCurrentTest, DiagonalPath_VertexHit_3D) { + using T = typename TestFixture::value_type; + constexpr unsigned Dim = TestFixture::dim; + + if constexpr (Dim != 3) { + GTEST_SKIP() << "Vertex-hit crossing test only implemented for 3D"; + } else { + using bunch_t = typename TestFixture::bunch_t; + using playout_t = typename TestFixture::playout_t; + using JField_t = typename TestFixture::JField_t; + + int nx = 32; + ippl::Vector origin(0.0); + ippl::Vector h(1.0); + + auto owned = TestFixture::make_owned_nd(nx); + auto layout = TestFixture::make_layout(owned); + auto mesh = TestFixture::make_mesh(owned, h, origin); + + JField_t J_field(mesh, layout); + J_field = T(0); + + playout_t playout(layout, mesh); + bunch_t bunch(playout); + bunch.create(ippl::Comm->rank() == 0 ? 1 : 0); + if (ippl::Comm->rank() == 0) { + auto R_host = bunch.R.getHostMirror(); + auto Rn_host = bunch.R_next.getHostMirror(); + auto Q_host = bunch.Q.getHostMirror(); + R_host(0)[0] = T(0.9); R_host(0)[1] = T(0.9); R_host(0)[2] = T(0.8); + Rn_host(0)[0] = T(1.1); Rn_host(0)[1] = T(1.1); Rn_host(0)[2] = T(1.2); + Q_host(0) = T(1.0); + Kokkos::deep_copy(bunch.R.getView(), R_host); + Kokkos::deep_copy(bunch.R_next.getView(), Rn_host); + Kokkos::deep_copy(bunch.Q.getView(), Q_host); + } + bunch.update(); + + auto policy = Kokkos::RangePolicy<>(0, bunch.getLocalNum()); + T dt = T(1.0); + ippl::assemble_current_collocated(mesh, bunch.Q, bunch.R, bunch.R_next, + J_field, policy, dt); + J_field.accumulateHalo(); + + auto ldom = J_field.getLayout().getLocalNDIndex(); + int nghost = J_field.getNghost(); + auto view = J_field.getView(); + auto view_host = Kokkos::create_mirror_view(view); + Kokkos::deep_copy(view_host, view); + Kokkos::fence(); + + const T tol = std::numeric_limits::epsilon() * T(100); + + auto in_ldom = [&](const Kokkos::Array& cell) { + for (unsigned d = 0; d < Dim; ++d) + if (cell[d] < ldom.first()[d] || cell[d] > ldom.last()[d]) return false; + return true; + }; + + auto check = [&](const Kokkos::Array& cell, unsigned c, double expected) { + double local_val = 0.0; + if (in_ldom(cell)) { + size_t idx[Dim]; + for (unsigned d = 0; d < Dim; ++d) + idx[d] = static_cast(cell[d] - ldom.first()[d] + nghost); + local_val = static_cast(ippl::apply(view_host, idx)[c + 1]); + } + double global_val = 0.0; + MPI_Allreduce(&local_val, &global_val, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + EXPECT_NEAR(global_val, expected, static_cast(tol)); + }; + + // Jx (c=0): dp[0]=0.1 per segment, weights from full trilinear CIC. + check({0, 0, 0}, 0, 0.02625); + check({1, 0, 0}, 0, 0.02475); + check({0, 1, 0}, 0, 0.02475); + check({1, 1, 0}, 0, 0.02425); + check({0, 0, 1}, 0, 0.02425); + check({1, 0, 1}, 0, 0.02475); + check({0, 1, 1}, 0, 0.02475); + check({1, 1, 1}, 0, 0.02625); + check({2, 1, 1}, 0, 0.0); + check({1, 2, 1}, 0, 0.0); + check({2, 2, 1}, 0, 0.0); + check({1, 1, 2}, 0, 0.0); + check({2, 1, 2}, 0, 0.0); + check({1, 2, 2}, 0, 0.0); + check({2, 2, 2}, 0, 0.0); + + // Jy == Jx (dp_y == dp_x): spot-check a subset of nodes. + check({0, 0, 0}, 1, 0.02625); + check({1, 1, 0}, 1, 0.02425); + check({1, 1, 1}, 1, 0.02625); + check({1, 1, 2}, 1, 0.0); + check({2, 2, 2}, 1, 0.0); + + // Jz (c=2): dp[2]=0.2 == 2*dp[0], so Jz == 2*Jx at every node. + check({0, 0, 0}, 2, 0.0525); + check({1, 0, 0}, 2, 0.0495); + check({0, 1, 0}, 2, 0.0495); + check({1, 1, 0}, 2, 0.0485); + check({0, 0, 1}, 2, 0.0485); + check({1, 0, 1}, 2, 0.0495); + check({0, 1, 1}, 2, 0.0495); + check({1, 1, 1}, 2, 0.0525); + check({2, 1, 1}, 2, 0.0); + check({1, 2, 1}, 2, 0.0); + check({2, 2, 1}, 2, 0.0); + check({1, 1, 2}, 2, 0.0); + check({2, 1, 2}, 2, 0.0); + check({1, 2, 2}, 2, 0.0); + check({2, 2, 2}, 2, 0.0); + } +} + TYPED_TEST(AssembleCurrentYeeTest, DiagonalPath_ThreeCells_ExactValues) { using T = typename TestFixture::value_type; constexpr unsigned Dim = TestFixture::dim; @@ -77,7 +449,7 @@ TYPED_TEST(AssembleCurrentYeeTest, DiagonalPath_ThreeCells_ExactValues) { auto mesh = TestFixture::make_mesh(owned, h, origin); JField_t J_field(mesh, layout); - + J_field = T(0); // path: (0.75, 0.50) -> (1.50, 1.25), same as FEM test @@ -133,13 +505,13 @@ TYPED_TEST(AssembleCurrentYeeTest, DiagonalPath_ThreeCells_ExactValues) { check(0, 0, 1, 0.03125); check(0, 1, 0, 0.15625); check(1, 0, 0, 0.03125); - check(1, 0, 1, 0.4375); - check(1, 1, 0, 0.4375); + check(1, 0, 1, 0.4375); + check(1, 1, 0, 0.4375); check(1, 1, 1, 0.15625); check(1, 2, 0, 0.03125); check(2, 0, 1, 0.03125); check(2, 1, 1, 0.09375); - + } }