species.hpp

Go to the documentation of this file.

#ifndef SPECIES_HPP
#define SPECIES_HPP
#include <Cabana_AoSoA.hpp>
#include <Cabana_DeepCopy.hpp>
#include <Kokkos_Core.hpp>
#include "NamelistReader.hpp"
#include "timer_macro.hpp"
#include "magnetic_field.hpp"
#include "grid.hpp"
#include "domain_decomposition.hpp"
#include "particles.hpp"
#include "space_settings.hpp"
#include "distribution.hpp"
#include "profile.hpp"
#include "gyro_avg_mat.hpp"
#include "streamed_parallel_for.hpp"
#include "basic_physics.hpp"
#include "memory_prediction.hpp"
#include "xgc_io.hpp"

extern "C" void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles* ptl);
extern "C" void set_min_max_num(int isp, int n_ptl);
extern "C" void adjust_n_ptl_for_core_ptl(int* n_ptl);

inline bool default_streaming_option(){
    return false;
}

inline bool default_residence_option(){
    return false;
}

// Used for Cabana slices (getting one particle property at a time)
namespace PtlSlice{
#ifdef ESC_PTL
enum{Ph=0,Ct,Gid,Flag};
#else
enum{Ph=0,Ct,Gid};
#endif
}

struct PtlMvmt{
    // Options for particle location management when looping over particles
    enum SendOpt{
        NoSend=0,
        Send,
        SendIfNotResident
    };

    enum ReturnOpt{
        NoReturn=0,
        Return,
        ReturnIfNotResident
    };

    SendOpt send_opt;
    ReturnOpt return_opt;

    PtlMvmt(SendOpt send_opt, ReturnOpt return_opt) : send_opt(send_opt), return_opt(return_opt){}
};

enum RKRestorationMethod{
    RestoreOnOriginalRank = 0,
    BringOldPhaseToNewRank
};

// Species class
template<class Device>
class Species{
    public:

    int idx;             
    bool is_electron;    
    bool is_adiabatic;   
    int nonadiabatic_idx;
    KinType kintype;     
    double mass;         
    double charge;       
    double charge_eu;    
    double c_m;          
    double c2_2m;        

    MarkerType marker_type;          
    FAnalyticShape f_analytic_shape; 
    WeightEvoEq weight_evo_eq;      
    bool dynamic_f0;                 

    bool is_deltaf; 

    int ncycles;               
    int ncycles_between_sorts; 

    int minimum_ptl_reservation; 
    int n_ptl; 
    Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> particles; 

    // Device particles
    Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> particles_d; 

    bool owns_particles_d; 

    bool particles_resident_on_device; 
    bool stream_particles; 

    /*** Could be its own class inside species? ***/
    RKRestorationMethod RK_restoration_method; 
    bool particles_are_backed_up; 

    // phase0 (for ion restoration)
    Cabana::AoSoA<PhaseDataTypes,HostType,VEC_LEN> phase0;
    Cabana::AoSoA<PhaseDataTypes,Device,VEC_LEN> phase0_d;

    // For electron restoration
    Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> backup_particles; 
    int n_backup_particles; // Number of particles stored in backup_particles (can't be deduced from its size due to buffer)
    /****/

    Distribution<Device> f0; 

    int collision_grid_index; 

    Eq::Profile<Device> eq_temp; // Equilibrium temperature
    Eq::Profile<Device> eq_den;  // Equilibrium density
    Eq::Profile<Device> eq_flow; // Equilibrium flow
    int eq_flow_type; // Type of Equilibirum flow

    Eq::Profile<Device> eq_fg_temp; // Equilibrium temperature
    Eq::Profile<Device> eq_fg_den;  // Equilibrium density - not used
    Eq::Profile<Device> eq_fg_flow; // Equilibrium flow - not used for now
    int eq_fg_flow_type; // Type of Equilibirum flow

    Eq::Profile<Device> eq_mk_temp; // Equilibrium temperature
    Eq::Profile<Device> eq_mk_den;  // Equilibrium density
    Eq::Profile<Device> eq_mk_flow; // Equilibrium flow
    int eq_mk_flow_type; // Type of Equilibirum flow


    GyroAverageMatrices<HostType> gyro_avg_matrices; // Gyroaverage matrices

    /*** Constructors ***/

    Species(int idx_in, int nonadiabatic_idx_in, bool is_electron_in, bool is_adiabatic_in, KinType kintype_in, double mass_in, double charge_in, double charge_eu_in, bool is_deltaf_in,
        int ncycles_in);

    Species(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const MagneticField<DeviceType> &magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, int idx_in, int nonadiabatic_idx_in);

    // Electron or ion default constructor
    Species(SpeciesType sp_type, int n_ptl)
      : idx(sp_type==ELECTRON ? 0 : 1),
        is_electron(sp_type==ELECTRON),
        mass(is_electron ? 3.344e-30 : PROTON_MASS),
        charge(is_electron ? -UNIT_CHARGE : UNIT_CHARGE),
        charge_eu(is_electron ? -1.0 : 1.0),
        is_deltaf(true),
        is_adiabatic(false),
        nonadiabatic_idx(idx), // Since is_adiabatic is false above
        kintype(is_electron ? DriftKin : GyroKin),
        ncycles(is_electron ? 70 : 1),
        c_m(charge/mass),
        c2_2m(0.5*charge*charge/mass),
        collision_grid_index(nonadiabatic_idx),
        n_ptl(n_ptl),
        backup_particles("backup_particles", 0),
        particles("particles", add_vec_buffer(n_ptl)),
        minimum_ptl_reservation(0),
        ncycles_between_sorts(5),
        eq_temp(1.0e3,-0.1),
        eq_den(1.0e19,-0.1),
        eq_flow_type(2),
        eq_fg_temp(1.0e3,-0.1),
        eq_fg_den(1.0e19,-0.1),
        eq_fg_flow_type(2),
        eq_mk_temp(1.0e3,-0.1),
        eq_mk_den(1.0e19,-0.1),
        eq_mk_flow_type(2),         
        owns_particles_d(false),
        stream_particles(default_streaming_option()),
        particles_resident_on_device(default_residence_option()),
        RK_restoration_method(is_electron ? RestoreOnOriginalRank : BringOldPhaseToNewRank) {}

    // Special constructor for tests that involve tracking particles in memory
    // The idea is to use these particles to test functions that reorder particles,
    // e.g. sort, shift, and cleaning
    Species(int n_ptl_in)
        : n_ptl(n_ptl_in),
          is_electron(true),
          is_adiabatic(false),
          particles("particles", add_vec_buffer(n_ptl_in)),
          minimum_ptl_reservation(0),
          owns_particles_d(false),
          stream_particles(default_streaming_option()),
          particles_resident_on_device(default_residence_option()),
          RK_restoration_method(is_electron ? RestoreOnOriginalRank : BringOldPhaseToNewRank),
          particles_are_backed_up(false)
    {

        // Slice particle properties
        auto ph = Cabana::slice<PtlSlice::Ph>(particles);
        auto ct = Cabana::slice<PtlSlice::Ct>(particles);
        auto gid = Cabana::slice<PtlSlice::Gid>(particles);
#ifdef ESC_PTL
        auto flag = Cabana::slice<PtlSlice::Flag>(particles);
#endif

        // Offset gid if using MPI
#ifdef USE_MPI
        long long int gid_offset = n_ptl*SML_COMM_RANK;
#else
        long long int gid_offset = 0;
#endif

        // Assign trackable values
        for (int i=0;i<n_ptl;i++){
            // Set GID in order
            gid(i) = gid_offset + i+1; // 1-indexed

            // Value of properties is gid + 0.1*(property index)
            // First particle: (1.0, 1.1, ... 1.8)
            // Second particle: (2.0, 2.1, ... 2.8)
            for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
            for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
        }

        // Buffer particles: same but with gid = -1
        if(n_ptl>0){
            for (int i=n_ptl;i<add_vec_buffer(n_ptl);i++){
                gid(i) = -1;
                for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
                for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
            }
        }
    }

    static std::vector<MemoryPrediction> estimate_memory_usage(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const DomainDecomposition<DeviceType>& pol_decomp, int species_idx);

    static int get_initial_n_ptl(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const DomainDecomposition<DeviceType>& pol_decomp, int sml_special, int species_idx, bool verbose);

    void resize_particles(int new_n_ptl){
        n_ptl = new_n_ptl;

#ifndef USE_GPU
        // If CPU-only, particles_d points to the same location as particles. If particles is resized, then Cabana will not deallocate the first allocation
        // since it is still used by particles_d. So, reset particles_d before resize, and point it back to particles only afterwards 
        particles_d = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>();
#endif

        particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
        particles.resize(add_vec_buffer(n_ptl));

#ifndef USE_GPU
        // Point particles_d back to particles after resize
        particles_d = particles;
#endif

        set_spall_num_and_ptr( idx,n_ptl,divide_and_round_up(n_ptl,VEC_LEN), (VecParticles*)particles.data() );
    }

    void resize_host_particles_to_match_device(){
#ifdef USE_GPU
        particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
        particles.resize(particles_d.size());
#else
        // Point particles back to particles_d after resize
        particles = particles_d;
#endif

        set_spall_num_and_ptr( idx,n_ptl,divide_and_round_up(n_ptl,VEC_LEN), (VecParticles*)particles.data() );
    }

     /* If using CPU-only, then "device" particles are a shallow copy of host particles so that
     * there is no unnecessary duplication. When "device" particles are resized, Cabana will keep the
     * original allocation if there is a second reference (i.e. host particles).
     * To resolve this, we free the host particles here so that there is no second reference
     * */
    void unassign_host_particles(){
        particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>();
    }

    /* Resizes device particles if on GPU, or just creates a shallow copy if CPU only
     * */
    void resize_device_particles(){
        if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");

#ifdef USE_GPU
        particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
        // Resize device particles to match n particles
        particles_d.resize(add_vec_buffer(n_ptl));
#else
        // If kernels are on CPU, do shallow copy
        particles_d = particles;
#endif
    }

    /* Resizes device particles
     * */
    void resize_device_particles(int new_n_ptl){
        if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");

        n_ptl = new_n_ptl;

        particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)

        // Resize device particles to match host particles
        particles_d.resize(add_vec_buffer(n_ptl));
    }

    /* Copies particles to device - deep copy if using GPU, otherwise shallow copy
     * Also takes the opportunity to set the buffer particles to realistic values
     * */
    void copy_particles_to_device(){
        if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles to device, but doesn't own the device array.");

#ifdef USE_GPU
        // Copy to device
        Cabana::deep_copy(particles_d, particles);
#else
        // No operation required if CPU-only
#endif

        // Copy last particle to fill remainder of trailing vector in AoSoA
        set_buffer_particles_d();
    }

    /* Copies particles from device - deep copy if using GPU, otherwise no copy is necessary
     * */
    void copy_particles_from_device(){
        if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles from device, but doesn't own the device array.");

#ifdef USE_GPU
        // Copy particles to host
        Cabana::deep_copy(particles, particles_d);
#else
        // No operation required if CPU-only
#endif
    }

    /* Copies particles to device if they are resident on the device
     * */
    void copy_particles_to_device_if_resident(){
        if(particles_resident_on_device) copy_particles_to_device();
    }

    /* Copies particles from device if they are resident on the device
     * */
    void copy_particles_from_device_if_resident(){
        if(particles_resident_on_device) copy_particles_from_device();
    }

    /* Copies particles to device if they are NOT resident on the device
     * */
    void copy_particles_to_device_if_not_resident(){
        if(!particles_resident_on_device) copy_particles_to_device();
    }

    /* Copies particles from device if they are NOT resident on the device
     * */
    void copy_particles_from_device_if_not_resident(){
        if(!particles_resident_on_device) copy_particles_from_device();
    }

    void set_buffer_particles_d(){
        if (n_ptl>0){
            int last_ptl_index = n_ptl - 1;
            auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
            auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
            auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
#ifdef ESC_PTL
            auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
#endif

            Kokkos::parallel_for("set_buffer_particles_d", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
                // Buffer particles: same as last particle, gid = -1
                for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
                for (int j=0;j<3;j++) ct(i, j) = ct(last_ptl_index, j);
                gid(i) = -1;
#ifdef ESC_PTL
                flag(i) = flag(last_ptl_index);
#endif
            });
        }
    }

    void set_buffer_phase0_d(){
        if (n_ptl>0){
            int last_ptl_index = n_ptl - 1;
            auto ph = Cabana::slice<PtlSlice::Ph>(phase0_d);

            Kokkos::parallel_for("set_buffer_phase0", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
                // copy final real particle
                for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
            });
        }
    }

    // Options for custom launch bounds since kokkos defaults are suboptimal for electron push kernel
    enum class LaunchBounds{
        Default,
        Custom
    };

    template<typename F>
    inline void for_all_particles(const std::string label, F lambda_func) const {
        Kokkos::RangePolicy<ExSpace> particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
        Kokkos::parallel_for(label,  Opt::require(particle_range_policy, Async), lambda_func);
    }

    inline void back_up_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
        auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
        auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
        auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
#ifdef ESC_PTL
        auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
#endif

        auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
        auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
        auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
#ifdef ESC_PTL
        auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
#endif

        Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
            int i_offset = i + offset;
            // Make backup copy
            for (int j=0;j<6;j++) ph_b(i, j) = ph(i_offset, j);
            for (int j=0;j<3;j++) ct_b(i, j) = ct(i_offset, j);
            gid_b(i) = gid(i_offset);
#ifdef ESC_PTL
            flag_b(i) = flag(i_offset);
#endif
            // Deactivate
            gid(i_offset) = -1;
        });
    }

    inline void restore_backup_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
        auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
        auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
        auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
#ifdef ESC_PTL
        auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
#endif

        auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
        auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
        auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
#ifdef ESC_PTL
        auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
#endif

        Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
            int i_offset = i + offset;
            // Restore from backup copy
            for (int j=0;j<6;j++) ph(i_offset, j) = ph_b(i, j); 
            for (int j=0;j<3;j++) ct(i_offset, j) = ct_b(i, j); 
            gid(i_offset) = gid_b(i);
#ifdef ESC_PTL
            flag(i_offset) = flag_b(i);
#endif
        }); 
    }

    template<typename F>
    inline void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const {
        if(end_idx <= begin_idx) return; // Return if range is 0 or less

        // Still need the subset to line up with the AoSoA vector length
        int first_soa = begin_idx/VEC_LEN;
        int n_other_ptl_in_first_soa = begin_idx - first_soa*VEC_LEN;
        bool first_soa_is_partial = (n_other_ptl_in_first_soa>0);
        int last_soa = (end_idx-1)/VEC_LEN;
        int n_other_ptl_in_last_soa = (last_soa+1)*VEC_LEN - end_idx;
        bool last_soa_is_partial = (n_other_ptl_in_last_soa>0);

#ifdef USE_GPU
        int first_item_in_shifted_range = first_soa*VEC_LEN;
#else
        int first_item_in_shifted_range = first_soa;
#endif

        Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_first_soa;
        Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_last_soa;
        if(first_soa_is_partial){
            // Make a backup of the first SoA and set the particles_d GIDs to -1
            ptl_first_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_first_soa", n_other_ptl_in_first_soa);
            back_up_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
        }
        if(last_soa_is_partial){
            // Make a backup of the last SoA and set the particles_d GIDs to -1
            ptl_last_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_last_soa", n_other_ptl_in_last_soa);
            back_up_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
        }

        // Finally, do the parallel_for
        Kokkos::RangePolicy<ExSpace> particle_range_policy( first_item_in_shifted_range, p_range<DeviceType>(end_idx) );
        Kokkos::parallel_for(label,  Opt::require(particle_range_policy, Async), lambda_func);

        // Restore from backup
        if(first_soa_is_partial){
            restore_backup_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
        }
        if(last_soa_is_partial){
            restore_backup_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
        }
    }

    template<typename F>
    inline void for_all_particles(const std::string label, F lambda_func,
        const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default) {
        if(!owns_particles_d) exit_XGC("\nSpecies tried to loop over particles on device, but doesn't own the device array.");

        bool use_streaming = stream_particles;
#ifndef USE_STREAMS
        use_streaming = false; // Just to be safe, turn streams off here
#endif

        bool send_ptl = ( mvmt.send_opt==PtlMvmt::Send ||
                         (mvmt.send_opt==PtlMvmt::SendIfNotResident && !particles_resident_on_device));
        bool return_ptl = ( mvmt.return_opt==PtlMvmt::Return ||
                           (mvmt.return_opt==PtlMvmt::ReturnIfNotResident && !particles_resident_on_device));

        // Don't need to stream if particles are already present and don't need to be returned
        if((!send_ptl) && (!return_ptl)) use_streaming = false;

        if(use_streaming){
#ifdef USE_STREAMS
            Streamed::Option stream_option = Streamed::Normal; // Send to device and back
            if(!send_ptl) stream_option = Streamed::NoSend;
            if(!return_ptl) stream_option = Streamed::NoReturn;

            // Execute streaming parallel_for
            Streamed::parallel_for(label, n_ptl, lambda_func, stream_option, particles, particles_d);
#endif
        }else{
            if(send_ptl) TIMER("copy_ptl_to_device",copy_particles_to_device() );

            if (launch_bounds==LaunchBounds::Custom) {
#ifdef USE_EPUSH_LAUNCH_BOUNDS
# if !defined(PUSH_MAX_THREADS_PER_BLOCK) || !defined(PUSH_MIN_WARPS_PER_EU)
#  error "USE_EPUSH_LAUNCH_BOUNDS requires PUSH_MAX_THREADS_PER_BLOCK and PUSH_MIN_WARPS_PER_EU to be defined"
# endif
                Kokkos::RangePolicy<ExSpace, Kokkos::LaunchBounds<PUSH_MAX_THREADS_PER_BLOCK, PUSH_MIN_WARPS_PER_EU>>
                    particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
                Kokkos::parallel_for(label,  Opt::require(particle_range_policy, Async), lambda_func);
#else
                exit_XGC("\nERROR: LaunchBounds::Custom specified, but USE_EPUSH_LAUNCH_BOUNDS is not defined\n");
#endif
            } else {
                Kokkos::RangePolicy<ExSpace>
                    particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
                Kokkos::parallel_for(label,  Opt::require(particle_range_policy, Async), lambda_func);
            }

            if(return_ptl) TIMER("copy_ptl_from_device", copy_particles_from_device() );
        }
    }

    KOKKOS_INLINE_FUNCTION VecParticles* ptl() const{
        return (VecParticles*)(&particles_d.access(0));
    }

    KOKKOS_INLINE_FUNCTION VecPhase* ph0() const{
        return (VecPhase*)(&phase0_d.access(0));
    }

    void copy_to_phase0(Species<Device>& species){
        for_all_particles("copy_to_phase0", KOKKOS_LAMBDA( const int idx ){
            AoSoAIndices<DeviceType> inds(idx);
            VecParticles* ptl_loc = species.ptl();
            VecPhase* ph0_loc = species.ph0();
            for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
                int p_vec = inds.a + i_simd;
                ph0_loc[inds.s].r[p_vec]   = ptl_loc[inds.s].ph.r[p_vec];
                ph0_loc[inds.s].z[p_vec]   = ptl_loc[inds.s].ph.z[p_vec];
                ph0_loc[inds.s].phi[p_vec] = ptl_loc[inds.s].ph.phi[p_vec];
                ph0_loc[inds.s].rho[p_vec] = ptl_loc[inds.s].ph.rho[p_vec];
                ph0_loc[inds.s].w1[p_vec]  = ptl_loc[inds.s].ph.w1[p_vec];
                ph0_loc[inds.s].w2[p_vec]  = ptl_loc[inds.s].ph.w2[p_vec];
            }
        });
    }

    void save_backup_particles(){
        if(RK_restoration_method==RestoreOnOriginalRank){
            // If particles are resident on the device, then use the host particle allocation as the backup
            // Otherwise, resize the backup_particle array
            if(particles_resident_on_device){
                backup_particles = particles;
            }else{
                backup_particles.resize(particles_d.size());
            }

            // Copy particle data to backup
            Cabana::deep_copy(backup_particles, particles_d);
            n_backup_particles = n_ptl;
        } else {
            // For ions, save phase to phase0
            phase0_d.resize(particles_d.size());
            copy_to_phase0(*this); // Separate function to avoid implicit copy into lambda
        }
        particles_are_backed_up = true;
    }

    void restore_particles_from_backup(){
        if(RK_restoration_method==RestoreOnOriginalRank){
            // If particles are resident on device, don't need to resize host particles since they are already used as the backup.
            // If not, resize host particle array to fit backup particles
            if(particles_resident_on_device){
                n_ptl = n_backup_particles;
                set_spall_num_and_ptr( idx,n_ptl,divide_and_round_up(n_ptl,VEC_LEN), (VecParticles*)particles.data() );
            }else{
                // If particles are not resident on device, resize host particle array to fit backup particles
                resize_particles(n_backup_particles);
            }

            // Resize device particle array to fit backed up particles
            resize_device_particles();

            // Restore particle data from backup
            Cabana::deep_copy(particles_d, backup_particles);

            if(particles_resident_on_device){
                // If the backup particles are simply pointing to the host particles, then
                // point them to their own 0-sized allocation when finished using them
                // so that they don't make a copy when host particles get resized
                backup_particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>("backup_particles", 0);
            }
        } else {
            // When ipc==2, copy phase0 to device
            // First, resize device phase0 and reset pointer
            phase0_d.resize(phase0.size());

            // Next copy data to phase0 on device
            Cabana::deep_copy(phase0_d, phase0);

            // Fill buffer with realistic ptl data
            set_buffer_phase0_d();
        }
        particles_are_backed_up = false;
    }

    KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices<Device>& inds, SimdParticles& part_one) const {
        VecPhase* ph0_loc = ph0();
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            int p_vec = inds.a + i_simd;
            part_one.ph.r[i_simd]   = ph0_loc[inds.s].r[p_vec];
            part_one.ph.z[i_simd]   = ph0_loc[inds.s].z[p_vec];
            part_one.ph.phi[i_simd] = ph0_loc[inds.s].phi[p_vec];
            part_one.ph.rho[i_simd] = ph0_loc[inds.s].rho[p_vec];
            part_one.ph.w1[i_simd]  = ph0_loc[inds.s].w1[p_vec];
            part_one.ph.w2[i_simd]  = ph0_loc[inds.s].w2[p_vec];
        }
    }

    long long int get_total_n_ptl(){
#ifdef USE_MPI
        long long int tmp_n_ptl = n_ptl;
        long long int out_n_ptl = 0;
        MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_LONG_LONG_INT, MPI_SUM, SML_COMM_WORLD);
        return out_n_ptl;
#else
        return (long long int)(n_ptl);
#endif
    }

    int get_max_n_ptl(){
#ifdef USE_MPI
        int tmp_n_ptl = n_ptl;
        int out_n_ptl = 0;
        MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_INT, MPI_MAX, SML_COMM_WORLD);
        return out_n_ptl;
#else
        return n_ptl;
#endif
    }

    // Gets the gyro_radius of a species based on equilibrium temperature
    // inode is the LOCAL (poloidally decomposed) grid node index to get temperature
    // smu_n is the normalized sqrt(mu)
    // bfield is the magnetic field at inode
    KOKKOS_INLINE_FUNCTION double get_fg_gyro_radius(int inode, double smu_n, double bfield) const{
        // Should replace UNIT_CHARGE*charge_eu with charge(?)
        return smu_n*sqrt(mass*f0.fg_temp_ev(inode)*EV_2_J) / (UNIT_CHARGE*charge_eu*bfield);
    }

    // Gets the equilibrium thermal velocity of a species based on f0 temperature
    // inode is the GLOBAL node index to get temperature
    KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const{
        return thermal_velocity(mass, f0.temp_global(inode));
    }

    // Gets the equilibrium thermal velocity of a species based on f0 temperature, on device
    // inode is the local node index to get temperature
    KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode(int inode) const{
        return thermal_velocity(mass, f0.temp_ev(inode));
    }

    // Gets the equilibrium thermal velocity of a species based on f0 temperature, on host
    // inode is the local node index to get temperature
    KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const{
        return thermal_velocity(mass, f0.temp_ev_h(inode));
    }

    KOKKOS_INLINE_FUNCTION double get_f0_fg_unit_velocity_lnode_h(int inode) const{
        return thermal_velocity(mass, f0.fg_temp_ev_h(inode));
    }

    // Get species velocity
    KOKKOS_INLINE_FUNCTION void get_particle_velocity_and_nearest_node(const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, SimdParticles& part, Simd<double>& smu, Simd<double>& vp, Simd<int>& nearest_node, Simd<bool>& not_in_triangle, Simd<bool>& not_in_poloidal_domain) const{

        // This modulo surely doesnt need to be here (at least, should be elsewhere).
        // Modulo phi coordinate
        grid.wedge_modulo_phi(part.ph.phi);

        SimdGridWeights<Order::Zero, PIT_GLOBAL> grid_wts0;
        grid.get_grid_weights(magnetic_field, part.ph.v(), grid_wts0);

        // Output argument
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            not_in_triangle[i_simd] = !grid_wts0.is_valid(i_simd);
        }

        Simd<double> bmag;
        magnetic_field.bmag_interpol(part.ph.v(), bmag);

        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            if(!grid_wts0.is_valid(i_simd)) continue;

            nearest_node[i_simd]=grid_wts0.node[i_simd] - pol_decomp.node_offset;
            not_in_poloidal_domain[i_simd] = (nearest_node[i_simd]<0 || nearest_node[i_simd]>=pol_decomp.nnodes);

            double temp_ev_norm = not_in_poloidal_domain[i_simd] ? f0.fg_temp_ev(0) : f0.fg_temp_ev(nearest_node[i_simd]);

            // get vp and smu
            const double& B = bmag[i_simd];
            vp[i_simd] = normalized_v_para(c_m, mass, B, temp_ev_norm, part.ph.rho[i_simd]);
            smu[i_simd] = normalized_sqrt_mu(B, temp_ev_norm, part.ct.mu[i_simd]);
        }
    }

    void update_decomposed_f0_calculations(const DomainDecomposition<DeviceType>& pol_decomp,
                                           const Grid<DeviceType>& grid,
                                           const MagneticField<DeviceType>& magnetic_field,
                                           const VelocityGrid& vgrid);

    void calculate_global_f0_arrays(const Grid<DeviceType>& grid,
                                    const MagneticField<DeviceType>& magnetic_field);

    void write_ptl_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
    void write_f0_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
    void read_f0_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
    void read_ptl_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name, bool n_ranks_is_same, int version);

    long long int get_max_gid() const;
    void get_ptl_write_total_and_offsets(const DomainDecomposition<DeviceType>& pol_decomp, long long int& inum_total, long long int& ioff) const;
};

#endif