plasma.hpp

Go to the documentation of this file.

#ifndef PLASMA_HPP
#define PLASMA_HPP
#include "my_subview.hpp"
#include "species.hpp"
#include "vgrid_distribution.hpp"

extern "C" void f0_init_decomposed_arrays(double* f0_T_ev_cpp, double* f0_inv_grid_vol_cpp, double* f0_grid_vol_cpp,
                                       double* f0_grid_vol_vonly_cpp, double* f0_n_Ta_cpp, double* f0_den_cpp,
                                       double* f0_flow_cpp, double* f0_B_B0_cpp);
extern "C" void f0_init_global_arrays(double* f0_den_global, double* f0_temp_global);
extern "C" void f0_set_ptrs(int nnode, double* f0_delta_n_cpp, double* f0_delta_u_cpp, double* f0_delta_T_cpp);
extern "C" void set_f0_f0g_ptr(int f0_inode1, int f0_inode2, double* f0_f0g_loc);
extern "C" void f0_initialize_wrap();

class Plasma{

    // Variables for managing the device particle memory allocation
    bool species_share_particles_d_ownership; 
    bool particles_d_has_owner; 
    int particles_d_owner; 

    public:

    std::vector<Species<DeviceType>> all_species; 

    // Poloidally decomposed f0 values
    // These will probably become a member of species later, but the fortran code expects
    // the species to be an array index
    VGridDistribution<HostType> f0_f0g;

    private:
    /* Contains f0 values that are poloidally decomposed but don't need to be transferred between ranks during load rebalancing
     * */
    struct DecomposedRecalculableF0Arrays{
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_T_ev;           
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_inv_grid_vol;   
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_grid_vol;       
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_grid_vol_vonly; 
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_n_Ta;           
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_den;            
        Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_flow;           

        // f0_B_B0 is not plasma-related, should be moved elsewhere (e.g. Grid) - ALS
        Kokkos::View<double*,Kokkos::LayoutRight, HostType> f0_B_B0; 

        DecomposedRecalculableF0Arrays(){}

        DecomposedRecalculableF0Arrays(int nsp, const DomainDecomposition<DeviceType>& pol_decomp)
          : f0_T_ev("f0_T_ev", nsp, pol_decomp.nnodes),
            f0_inv_grid_vol("f0_inv_grid_vol", nsp, pol_decomp.nnodes),
            f0_grid_vol("f0_grid_vol", nsp, pol_decomp.nnodes),
            f0_grid_vol_vonly("f0_grid_vol_vonly", nsp, pol_decomp.nnodes),
            f0_n_Ta("f0_n_Ta", nsp, pol_decomp.nnodes),
            f0_den("f0_den", nsp, pol_decomp.nnodes),
            f0_flow("f0_flow", nsp, pol_decomp.nnodes),
            f0_B_B0("f0_B_B0", pol_decomp.nnodes)
        {
            // Fortran arrays are set to point to these Views
            // The calculations are currently done in Fortran too
            f0_init_decomposed_arrays(f0_T_ev.data(),
                                      f0_inv_grid_vol.data(),
                                      f0_grid_vol.data(),
                                      f0_grid_vol_vonly.data(),
                                      f0_n_Ta.data(),
                                      f0_den.data(),
                                      f0_flow.data(),
                                      f0_B_B0.data());
        }
    };

    public:

    DecomposedRecalculableF0Arrays decomposed_recalculable_f0_arrays; 

    // Not poloidally decomposed
    Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_n; 
    Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_u; 
    Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_T; 

    View<double**,CLayout, HostType> f0_den_global; 
    View<double**,CLayout, HostType> f0_temp_global; 

    int nspecies;                 
    int n_nonadiabatic_species;   

    // Constructors
    Plasma()
      : particles_d_has_owner(false),
        species_share_particles_d_ownership(!default_residence_option()),
        nspecies(0), // Initialize with 0 species
        n_nonadiabatic_species(0) // Initialize with 0 species
    {}

    Plasma(NLReader::NamelistReader& nlr, bool use_f0_grid, bool allocate_for_adiabatic_electrons, const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid)
      : particles_d_has_owner(false)
    {
        if(nlr.namelist_present("performance_param")){
            nlr.use_namelist("performance_param");
            species_share_particles_d_ownership = !(nlr.get<bool>("particles_resident_on_device", default_residence_option()));
        }else{
            species_share_particles_d_ownership = !(default_residence_option());
        }
#ifndef USE_GPU
        species_share_particles_d_ownership = false;
#endif

        nlr.use_namelist("ptl_param");
        nspecies = nlr.get<int>("ptl_nsp",1);
        nspecies += 1; // Add electrons

        // Add each species
        int i_nonadiabatic_species = 0;
        for(int isp = 0; isp<nspecies; isp++){
            all_species.push_back(Species<DeviceType>(nlr, isp, i_nonadiabatic_species));
            if(!all_species[isp].is_adiabatic){
                i_nonadiabatic_species++;
            }
        }

        // Save count of non-adiabatic species
        n_nonadiabatic_species = i_nonadiabatic_species;

        /*** f0 arrays ***/
        int grid_nnode = pol_decomp.gvid0_pid_h(pol_decomp.gvid0_pid_h.size()-1) - 1;
        init_global_f0_arrays(grid_nnode, allocate_for_adiabatic_electrons);

        if(use_f0_grid){
            resize_f0_f0g(pol_decomp, vgrid);
            f0_initialize_wrap();

            // Reallocate and recalculate arrays that are decomposed but recalculated rather than
            // transferred by the load balance (f0_T_ev etc)
            update_decomposed_f0_calculations(pol_decomp);
        }
    }

    enum DevicePtlOpt{
        UseDevicePtl=0,
        NoDevicePtl
    };

    enum ParticleType{
        Electrons=0,
        Ions
    };

    // Loop over all species
    template<typename F>
    inline void for_all_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        for(int isp = 0; isp<all_species.size(); isp++){
            manage_particle_ownership(isp, device_ptl_opt);
            func(all_species[isp]);
        }
    }

    // Loop over all non-adiabatic species
    template<typename F>
    inline void for_all_nonadiabatic_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        for(int isp = 0; isp<all_species.size(); isp++){
            if(isp>10) { // ISP ERROR
                printf("ISP ERROR in for_all_nonadiabatic_sepcies: isp=%d",isp);
                fflush(stdout);
            }
            if(!all_species[isp].is_adiabatic){
                manage_particle_ownership(isp, device_ptl_opt);
                func(all_species[isp]);
            }
        }
    }

    // Loop over electrons
    template<typename F>
    inline void for_electrons(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        for(int isp = 0; isp<all_species.size(); isp++){
            if(all_species[isp].is_electron){
                manage_particle_ownership(isp, device_ptl_opt);
                func(all_species[isp]);
            }
        }
    }

    // Loop over ions
    template<typename F>
    inline void for_all_ions(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        for(int isp = 0; isp<all_species.size(); isp++){
            if(!all_species[isp].is_electron){
                manage_particle_ownership(isp, device_ptl_opt);
                func(all_species[isp]);
            }
        }
    }

    // Loop over electrons or ions as specified
    template<typename F>
    inline void for_all(ParticleType particle_type, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        for(int isp = 0; isp<all_species.size(); isp++){
            if((particle_type==Electrons && all_species[isp].is_electron) ||
               (particle_type==Ions && !all_species[isp].is_electron)){
                manage_particle_ownership(isp, device_ptl_opt);
                func(all_species[isp]);
            }
        }
    }

    // Operate on one species
    template<typename F>
    inline void for_one_species(int isp, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
        manage_particle_ownership(isp, device_ptl_opt);
        func(all_species[isp]);
    }

    int largest_n_ptl(bool check_backup){
        int max_n_ptl = 0;

        // Loop over all species
        for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            if(species.is_electron && check_backup){
                // Check on backup particles instead, if that's what's getting used
                max_n_ptl = std::max(max_n_ptl,species.n_backup_particles);
            } else {
                // For now, need to access fortran object for n_ptl
                max_n_ptl = std::max(max_n_ptl,species.n_ptl);
            }
        }, NoDevicePtl);

        return max_n_ptl;
    }

    /* Deallocates the device particles and resets ownership tracking variables
     * */
    void deallocate_device_ptl(){
        for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            if(species.owns_particles_d){
                species.particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
                species.owns_particles_d = false;
            }
        });
        particles_d_has_owner = false;
    }

    /* Reallocate and recalculate arrays that are decomposed but recalculated rather than
     * transferred by the load balance (f0_T_ev etc) */
    void update_decomposed_f0_calculations(const DomainDecomposition<DeviceType>& pol_decomp){
        decomposed_recalculable_f0_arrays = DecomposedRecalculableF0Arrays(n_nonadiabatic_species, pol_decomp);

        // Set all_species unmanaged views
        int f0_species_cnt=0;
        for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_T_ev, f0_species_cnt, species.f0.temp_ev_h);
            f0_species_cnt++;
        }, NoDevicePtl);
    }

    private:

    /* Allocate and calculate global f0 arrays. Should be part of a constructor */
    void init_global_f0_arrays(int grid_nnode, bool is_XGCa){
        // f0_delta_{n,u,T} are sized differently in XGC1 vs XGCa (the latter uses index 0 even if
        // electrons are adiabatic)
        int nsp = (is_XGCa ? nspecies : n_nonadiabatic_species);

        f0_delta_n = View<double**,CLayout, HostType>("f0_delta_n", nsp, grid_nnode);
        f0_delta_u = View<double**,CLayout, HostType>("f0_delta_u", nsp, grid_nnode);
        f0_delta_T = View<double**,CLayout, HostType>("f0_delta_T", nsp, grid_nnode);
        f0_set_ptrs(grid_nnode, f0_delta_n.data(), f0_delta_u.data(), f0_delta_T.data());

        // These are constant for the simulation:
        f0_den_global = View<double**,CLayout, HostType>("f0_den_global", n_nonadiabatic_species, grid_nnode);
        f0_temp_global = View<double**,CLayout, HostType>("f0_temp_global", n_nonadiabatic_species, grid_nnode);

        // Fortran arrays are set to point to these Views
        // The calculations are currently done in Fortran too
        f0_init_global_arrays(f0_den_global.data(), f0_temp_global.data());

        for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            set_unmanaged_f0_species_view(f0_den_global, species.nonadiabatic_idx, species.f0.den_global_h);
            set_unmanaged_f0_species_view(f0_temp_global, species.nonadiabatic_idx, species.f0.temp_global_h);
        }, NoDevicePtl);
    }

    // Creates an unmanaged view pointing to the subview specified by isp
    template<typename T_in, typename T_out>
    void set_unmanaged_f0_species_view(const T_in& view_in, int isp, T_out& view_out){
        auto view_in_subview = my_subview(view_in, isp);
        view_out = T_out(view_in_subview.data(), view_in_subview.layout());
    }

    public:

    void resize_f0_f0g(const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid){
        // Create new object rather than resizing since we don't need to preserve data
        f0_f0g = VGridDistribution<HostType>(n_nonadiabatic_species, vgrid, pol_decomp);

        int f0_inode1 = pol_decomp.first_node; // 1-indexed
        int f0_inode2 = f0_inode1 + pol_decomp.nnodes - 1; // 1-indexed
        set_f0_f0g_ptr(f0_inode1, f0_inode2, f0_f0g.data());

        // Set all_species unmanaged views
        int f0_species_cnt=0;
        for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            set_unmanaged_f0_species_view(f0_f0g.f, f0_species_cnt, species.f0.f0g_h);
            f0_species_cnt++;
        }, NoDevicePtl);
    }

    private:

    void transfer_particles_d_ownership(int isp){
        if(!particles_d_has_owner){
            // Set new owner and initialize with 0 particles
            particles_d_owner = isp;

            all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
            all_species[particles_d_owner].owns_particles_d = true;
            particles_d_has_owner = true;
        }else{
            // No-op if species is handing off to itself
            if(particles_d_owner == isp) return;

            // Shallow copy to new owner
            all_species[isp].particles_d = all_species[particles_d_owner].particles_d;

            // Delete original
            all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
            all_species[particles_d_owner].owns_particles_d = false;

            // Set new owner
            particles_d_owner = isp;
            all_species[particles_d_owner].owns_particles_d = true;
        }
    }

    void manage_particle_ownership(int isp, DevicePtlOpt device_ptl_opt){
        if((!all_species[isp].is_adiabatic) && device_ptl_opt==UseDevicePtl){
            if(species_share_particles_d_ownership){
                transfer_particles_d_ownership(isp);
            }else{
                if(!all_species[isp].owns_particles_d){
                    all_species[isp].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
                    all_species[isp].owns_particles_d = true;
                }
            }

            // Resize device particles if they are used
            all_species[isp].resize_device_particles();
        }
    }

};

#endif