Thermal.H

Go to the documentation of this file.

/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
 *           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * Authors: Ji Qiang, Axel Huebl, Chad Mitchell
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_DISTRIBUTION_THERMAL
#define IMPACTX_DISTRIBUTION_THERMAL
 
#include "particles/ReferenceParticle.H"
 
#include <ablastr/constant.H>
 
#include <AMReX_GpuContainers.H>
#include <AMReX_Math.H>
#include <AMReX_Print.H>
#include <AMReX_Random.H>
#include <AMReX_REAL.H>
 
#include <cmath>
#include <cstdint>
#include <memory>
#include <random>
#include <utility>
#include <vector>
 
 
namespace impactx
{
namespace distribution
{
    struct ThermalData
    {
        static constexpr amrex::ParticleReal tolerance = 1.0e-3; 
        static constexpr amrex::ParticleReal rin = 1.0e-10;  
        static constexpr amrex::ParticleReal rout = 10.0;  
        static constexpr int nsteps = 2000;  
 
        amrex::ParticleReal m_f1;  
        amrex::ParticleReal m_f2;  
        amrex::ParticleReal m_phi1;  
        amrex::ParticleReal m_phi2; 
        amrex::ParticleReal m_p1;  
        amrex::ParticleReal m_p2; 
        amrex::ParticleReal m_rmin;  
        amrex::ParticleReal m_rmax;  
        int m_nbins;  
        amrex::ParticleReal* m_cdf1 = nullptr;  
        amrex::ParticleReal* m_cdf2 = nullptr;  
        amrex::ParticleReal m_Cintensity;  
        amrex::ParticleReal m_bg; 
        amrex::ParticleReal m_k;  
        amrex::ParticleReal m_T1; 
        amrex::ParticleReal m_T2; 
        amrex::ParticleReal m_w;  
 
        // GPU data
        static inline std::unique_ptr<amrex::Gpu::DeviceVector<amrex::ParticleReal>> m_d_cdf1;
        static inline std::unique_ptr<amrex::Gpu::DeviceVector<amrex::ParticleReal>> m_d_cdf2;
 
        ThermalData (
            amrex::ParticleReal kin,
            amrex::ParticleReal T1in,
            amrex::ParticleReal T2in,
            amrex::ParticleReal p1in,
            amrex::ParticleReal p2in,
            amrex::ParticleReal win
        )
        {
             m_k = kin;
             m_T1 = T1in;
             m_T2 = T2in;
             m_p1 = p1in;
             m_p2 = p2in;
             m_w = win;
        }

        void
        generate_radial_dist (amrex::ParticleReal bunch_charge, RefPart const & refpart)
        {
            using namespace amrex::literals;  // for _rt and _prt
            using namespace ablastr::constant::math;  // for pi
            using amrex::Math::powi;
 
            // Get relativistic beta*gamma
            amrex::ParticleReal bg = refpart.beta_gamma();
            m_bg = bg;
 
            // Get reference particle rest energy (eV) and charge (q_e)
            amrex::ParticleReal Erest = refpart.mass_MeV()*1.0e6;
            amrex::ParticleReal q_e = refpart.charge_qe();
 
            // Set space charge intensity
            m_Cintensity = q_e*bunch_charge/(powi<2>(bg)*Erest*ablastr::constant::SI::ep0);
 
            // Set minimum and maximum radius
            amrex::ParticleReal r_scale = matched_scale_radius();
            amrex::ParticleReal rmin = rin*r_scale;
            amrex::ParticleReal rmax = rout*r_scale;
            // amrex::PrintToFile("equilibrium_params.out") << r_scale << " " << data.Cintensity << "\n";
 
            // Scale the parameters p1 and p2
            amrex::ParticleReal rt2pi = std::sqrt(2.0_prt*pi);
            amrex::ParticleReal p_scale = powi<-3>(r_scale*rt2pi);
            m_p1 = m_p1*p_scale;
            m_p2 = m_p2*p_scale;
 
            // store integration parameters
            m_nbins = nsteps;
            m_rmin = rmin;
            m_rmax = rmax;
 
            // allocate CDFs
            std::vector<amrex::ParticleReal> cdf1(m_nbins+1);
            std::vector<amrex::ParticleReal> cdf2(m_nbins+1);
            m_cdf1 = cdf1.data();
            m_cdf2 = cdf2.data();
 
            // set initial conditions
            m_f1 = 0.0_prt;
            m_f2 = 0.0_prt;
            m_phi1 = 0.0_prt;
            m_phi2 = 0.0_prt;
            m_cdf1[0] = 0.0_prt;
            m_cdf2[0] = 0.0_prt;
 
            // integrate ODEs for the radial density
            integrate(rmin,rmax,nsteps);
 
            // a search over normalization parameters p1, p2 can be added here
 
            // rescale cdf to ensure the exact range [0,1]
            for (int n = 0; n < m_nbins; ++n) {
                m_cdf1[n] = m_cdf1[n] / m_cdf1[m_nbins];
                m_cdf2[n] = m_cdf2[n] / m_cdf2[m_nbins];
            }
 
            // copy CFD data to device
            m_d_cdf1 = std::make_unique<amrex::Gpu::DeviceVector<amrex::ParticleReal>>(m_nbins+1);
            m_d_cdf2 = std::make_unique<amrex::Gpu::DeviceVector<amrex::ParticleReal>>(m_nbins+1);
            amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
                                  cdf1.begin(), cdf1.end(),
                                  m_d_cdf1->begin());
            amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
                                  cdf2.begin(), cdf2.end(),
                                  m_d_cdf2->begin());
            amrex::Gpu::streamSynchronize();
        }
 
        amrex::ParticleReal
        matched_scale_radius ()
        {
            using namespace amrex::literals;  // for _rt and _prt
            using namespace ablastr::constant::math;  // for pi
 
            amrex::ParticleReal k = m_k;
            amrex::ParticleReal kT = (1_prt - m_w) * m_T1 + m_w * m_T2;
            amrex::ParticleReal a = m_Cintensity/(4_prt*pi*5_prt*std::sqrt(5_prt));
            amrex::ParticleReal rscale = std::sqrt(kT + std::pow(a*k,2_prt/3_prt))/k;
 
            return rscale;
        }
 
        void
        integrate (
            amrex::ParticleReal in,
            amrex::ParticleReal out,
            int steps
        )
        {
            using namespace amrex::literals; // for _rt and _prt
 
            // initialize numerical integration parameters
            amrex::ParticleReal const tau = (out-in)/steps;
            amrex::ParticleReal const half = tau*0.5_prt;
 
            // initialize the value of the independent variable
            amrex::ParticleReal reval = in;
 
            // loop over integration steps
            for (int j=0; j < steps; ++j)
            {
                // for a second-order Strang splitting
                map1(half, reval);
                map2(tau, reval);
                map1(half, reval);
                // store tabulated CDF
                m_cdf1[j+1] = m_f1;
                m_cdf2[j+1] = m_f2;
                // write cumulative density to file for debugging
                /* comment in for debugging
                amrex::PrintToFile("density1.out") << reval << " " << data.f1 << "\n";
                amrex::PrintToFile("density2.out") << reval << " " << data.f2 << "\n";
                amrex::PrintToFile("phi1.out") << reval << " " << data.phi1 << "\n";
                amrex::PrintToFile("phi2.out") << reval << " " << data.phi2 << "\n";
                */
            }
 
        }
 
        void
        map1 (
            amrex::ParticleReal const tau,
            amrex::ParticleReal & reval
        )
        {
            using namespace amrex::literals;  // for _rt and _prt
            using namespace ablastr::constant::math;  // for pi
 
            amrex::ParticleReal const f1 = m_f1;
            amrex::ParticleReal const f2 = m_f2;
            amrex::ParticleReal const phi1 = m_phi1;
            amrex::ParticleReal const phi2 = m_phi2;
            amrex::ParticleReal const r = reval;
 
            // Apply map to update phi1, phi2, and r:
            reval = r + tau;
            m_phi1 = phi1 + f1/(4.0_prt*pi*reval) - f1/(4.0_prt*pi*r);
            m_phi2 = phi2 + f2/(4.0_prt*pi*reval) - f2/(4.0_prt*pi*r);;
            m_f1 = f1;
            m_f2 = f2;
        }
 
        void
        map2 (
            amrex::ParticleReal const tau,
            amrex::ParticleReal & reval
        )
        {
            using namespace amrex::literals;  // for _rt and _prt
            using namespace ablastr::constant::math;  // for pi
 
            amrex::ParticleReal const f1 = m_f1;
            amrex::ParticleReal const f2 = m_f2;
            amrex::ParticleReal const phi1 = m_phi1;
            amrex::ParticleReal const phi2 = m_phi2;
            amrex::ParticleReal const r = reval;
            amrex::ParticleReal const k = m_k;
            amrex::ParticleReal const w = m_w;
            amrex::ParticleReal const T1 = m_T1;
            amrex::ParticleReal const T2 = m_T2;
 
            // Define space charge intensity parameters
            amrex::ParticleReal const c1 = (1_prt-w) * m_Cintensity;
            amrex::ParticleReal const c2 = w * m_Cintensity;
 
            // Define intermediate quantities
            amrex::ParticleReal potential = 0_prt;
            potential = std::pow(k*r,2)*0.5_prt + c1*phi1 + c2*phi2;
            amrex::ParticleReal Pdensity1 = m_p1 * std::exp(-potential/T1);
            amrex::ParticleReal Pdensity2 = m_p2 * std::exp(-potential/T2);
            // amrex::ParticleReal Pdensity_tot = (1.0_prt-w)*Pdensity1 + w*Pdensity2;
            // amrex::PrintToFile("Pdensity.out") << reval << " " << Pdensity_tot << "\n";
 
            // Apply map to update f1 and f2:
            m_phi1 = phi1;
            m_phi2 = phi2;
            m_f1 = f1 + tau*4_prt*pi * std::pow(r,2)*Pdensity1;
            m_f2 = f2 + tau*4_prt*pi * std::pow(r,2)*Pdensity2;
            reval = r;
        }
    };
 
    struct Thermal
    {
        Thermal (
            amrex::ParticleReal k,
            amrex::ParticleReal kT,
            amrex::ParticleReal kT_halo,
            amrex::ParticleReal normalize,
            amrex::ParticleReal normalize_halo,
            amrex::ParticleReal halo
        )
        : m_k(k),
          m_T1(kT),
          m_T2(kT_halo),
          m_normalize(normalize),
          m_normalize_halo(normalize_halo),
          m_halo(halo)
        {
        }

        void initialize (amrex::ParticleReal bunch_charge, RefPart const & ref)
        {
            // Generate the struct "data" containing the radial CDF:
            ThermalData data(m_k, m_T1, m_T2, m_normalize, m_normalize_halo, m_halo);
            data.generate_radial_dist(bunch_charge, ref);
 
            // share data
            m_w = data.m_w;
            m_T1 = data.m_T1;
            m_T2 = data.m_T2;
            m_rmin = data.m_rmin;
            m_rmax = data.m_rmax;
            m_nbins = data.m_nbins;
            m_bg = data.m_bg;
            m_cdf1 = data.m_d_cdf1->data();
            m_cdf2 = data.m_d_cdf2->data();
        }

        void
        finalize ()
        {
            // deallocate
            ThermalData::m_d_cdf1.reset();
            ThermalData::m_d_cdf2.reset();
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator() (
            amrex::ParticleReal & AMREX_RESTRICT x,
            amrex::ParticleReal & AMREX_RESTRICT y,
            amrex::ParticleReal & AMREX_RESTRICT t,
            amrex::ParticleReal & AMREX_RESTRICT px,
            amrex::ParticleReal & AMREX_RESTRICT py,
            amrex::ParticleReal & AMREX_RESTRICT pt,
            amrex::RandomEngine const & engine
        ) const
        {
 
            using namespace amrex::literals;  // for _rt and _prt
            using namespace ablastr::constant::math;  // for pi
 
            amrex::ParticleReal ln1,norm,u1,u2,uhalo;
            amrex::ParticleReal g1,g2,g3,g4,g5,g6;
            amrex::ParticleReal z,pz;
 
            // Use a Bernoulli random variable to select between core and halo:
            // If u < w, the particle is in the halo population.
            // If u >= w, the particle is in the core population.
            uhalo = amrex::Random(engine);
 
            // Generate six standard normal random variables using Box-Muller:
            u1 = amrex::Random(engine);
            u2 = amrex::Random(engine);
            ln1 = std::sqrt(-2_prt*std::log(u1));
            g1 = ln1 * std::cos(2_prt*pi*u2);
            g2 = ln1 * std::sin(2_prt*pi*u2);
            u1 = amrex::Random(engine);
            u2 = amrex::Random(engine);
            ln1 = std::sqrt(-2_prt*std::log(u1));
            g3 = ln1 * std::cos(2_prt*pi*u2);
            g4 = ln1 * std::sin(2_prt*pi*u2);
            u1 = amrex::Random(engine);
            u2 = amrex::Random(engine);
            ln1 = std::sqrt(-2_prt*std::log(u1));
            g5 = ln1 * std::cos(2_prt*pi*u2);
            g6 = ln1 * std::sin(2_prt*pi*u2);
 
            // Scale the last three variables to produce the momenta:
            amrex::ParticleReal kT = (uhalo > m_w) ? m_T1 : m_T2;  // select core or halo value
            px = std::sqrt(kT)*g4;
            py = std::sqrt(kT)*g5;
            pz = std::sqrt(kT)*g6;
 
            // Normalize the first three variables to produce uniform samples on the unit 3-sphere:
            norm = std::sqrt(g1*g1+g2*g2+g3*g3);
            g1 /= norm;
            g2 /= norm;
            g3 /= norm;
 
            // Collect the radial CDF returned by generate_radial_dist:
 
            // select core or halo CDF
            amrex::ParticleReal const * cdf = (uhalo > m_w) ? m_cdf1 : m_cdf2;
 
            // Generate a radial coordinate from the CDF
            amrex::ParticleReal u = amrex::Random(engine);
            amrex::ParticleReal const * ptr = amrex::lower_bound(cdf, cdf + m_nbins + 1, u);
            int const off = amrex::max(0, int(ptr - cdf - 1));
            amrex::ParticleReal tv = (u - cdf[off]) / (cdf[off + 1] - cdf[off]);
            amrex::ParticleReal xv = (off + tv) / amrex::ParticleReal(m_nbins);
            amrex::ParticleReal r = m_rmin + (m_rmax - m_rmin) * xv;
 
            // Scale to produce samples with the correct radial density:
            x = g1*r;
            y = g2*r;
            z = g3*r;
 
            // Transform longitudinal variables into the laboratory frame:
            t = -z / m_bg;
            pt = -pz * m_bg;
        }
 
        amrex::ParticleReal m_k; 
        amrex::ParticleReal m_T1, m_T2; 
        amrex::ParticleReal m_normalize, m_normalize_halo; 
        amrex::ParticleReal m_halo; 
        amrex::ParticleReal m_rmin;  
        amrex::ParticleReal m_rmax;  
        int m_nbins;  
        amrex::ParticleReal m_bg; 
        amrex::ParticleReal m_w;  
 
    private:
        // radial profile data
        amrex::ParticleReal const * m_cdf1 = nullptr; 
        amrex::ParticleReal const * m_cdf2 = nullptr; 
    };
 
} // namespace distribution
} // namespace impactx
 
#endif // IMPACTX_DISTRIBUTION_THERMAL