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_ptrs(double* f0_T_ev_cpp, double* f0_fg_T_ev_cpp, double* f0_inv_grid_vol_cpp, double* f0_grid_vol_cpp,
                                       double* f0_grid_vol_vonly_cpp, double* f0_den_cpp,
                                       double* f0_flow_cpp);
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);

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; 

    double main_ion_characteristic_energy;

    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;
    View<double*,CLayout,HostType> f0_node_cost;

    private:

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


        DecomposedRecalculableF0Arrays(){}

        DecomposedRecalculableF0Arrays(int nsp, const DomainDecomposition<DeviceType>& pol_decomp)
          : f0_T_ev("f0_T_ev", nsp, pol_decomp.nnodes),
            f0_fg_T_ev("f0_fg_T_ev", nsp, pol_decomp.nnodes),
            f0_fg_vth_inv("f0_fg_vth_inv", 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_den("f0_den", nsp, pol_decomp.nnodes),
            f0_flow("f0_flow", nsp, pol_decomp.nnodes)
        {
#ifndef NO_FORTRAN_MODULES
            // Fortran arrays are set to point to these Views
            f0_init_decomposed_ptrs(f0_T_ev.data(),
                                      f0_fg_T_ev.data(),
                                      f0_inv_grid_vol.data(),
                                      f0_grid_vol.data(),
                                      f0_grid_vol_vonly.data(),
                                      f0_den.data(),
                                      f0_flow.data());
#endif
        }
    };


    public:

    DecomposedRecalculableF0Arrays decomposed_recalculable_f0_arrays; 

    // Not poloidally decomposed
    View<double**,CLayout, HostType> f0_delta_n; 
    View<double**,CLayout, HostType> f0_delta_u; 
    View<double**,CLayout, HostType> f0_delta_T; 
    View<double*,CLayout, HostType> spitzer_resistivity; // Pre-calculated for equilibrium density and temperature

    int nspecies;                 
    int n_nonadiabatic_species;   

    std::vector<std::string> sp_names;

    // 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
        sp_names{"e", "i", "i2", "i3", "i4", "i5", "i6"}
    {}

    Plasma(NLReader::NamelistReader& nlr, const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid);

    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]);
    }

    // Loop over all species, return true if the input condition is met for any of them
    template<typename F>
    inline bool true_for_some_species(F func) const{
        for(int isp = 0; isp<all_species.size(); isp++){
            if(func(all_species[isp])) return true;
        }
        return false;
    }

    template<typename F>
    inline bool true_for_all_species(F func) const{
        for(int isp = 0; isp<all_species.size(); isp++){
            if(!func(all_species[isp])) return false;
        }
        return true;
    }

    // Replace !sml_electron_on
    inline bool electrons_are_adiabatic() const{
        return all_species[ELECTRON].is_adiabatic;
    }

    inline bool f0_grid() const{
        return true_for_some_species([&](const Species<DeviceType>& s){return s.dynamic_f0;});
    }

    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;
    }

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

    /* 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,
                                           const Grid<DeviceType>& grid,
                                           const MagneticField<DeviceType>& magnetic_field,
                                           const VelocityGrid& vgrid);

    private:

    void initialize_spitzer_res();

    /* Allocate and calculate global f0 arrays. Should be part of a constructor */
    void init_global_f0_arrays(const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field){
        // Allocate for all species (even though most simulations don't need it for adiabatic species)
        f0_delta_n = View<double**,CLayout, HostType>("f0_delta_n", nspecies, grid.nnode);
        f0_delta_u = View<double**,CLayout, HostType>("f0_delta_u", nspecies, grid.nnode);
        f0_delta_T = View<double**,CLayout, HostType>("f0_delta_T", nspecies, grid.nnode);
        // Set all_species unmanaged views
        for_all_species([&](Species<DeviceType>& species){
            set_unmanaged_f0_species_view(f0_delta_n, species.idx, species.f0.delta_n_h);
            set_unmanaged_f0_species_view(f0_delta_u, species.idx, species.f0.delta_u_h);
            set_unmanaged_f0_species_view(f0_delta_T, species.idx, species.f0.delta_T_h);
        }, NoDevicePtl);
#ifndef NO_FORTRAN_MODULES
        // Set fortran pointers
        f0_set_ptrs(grid.nnode, f0_delta_n.data(), f0_delta_u.data(), f0_delta_T.data());
#endif

        // These are constant for the simulation:
        for_all_species([&](Species<DeviceType>& species){
            species.initialize_global_f0_arrays(grid, magnetic_field);
        }, NoDevicePtl);

        initialize_spitzer_res();
    }

    // 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 remap_f0_f0g(const DomainDecomposition<DeviceType>& pol_decomp){
#ifndef NO_FORTRAN_MODULES
        int f0_inode1 = pol_decomp.node_offset + 1; // 1-indexed
        int f0_inode2 = f0_inode1 + pol_decomp.nnodes - 1; // 1-indexed
        set_f0_f0g_ptr(f0_inode1, f0_inode2, f0_f0g.data());
#endif

        // 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();
        }
    }

    public:
    void validate_f0_checkpoint_file_dims(const XGC_IO_Stream& stream, const Grid<DeviceType>& grid, const VelocityGrid& vgrid);
    void write_checkpoint_files(const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream);
    void read_checkpoint_files(const Grid<DeviceType>& grid, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, const XGC_IO_Stream& f0_stream, bool n_ranks_is_same, int version);

    inline double get_main_ion_toroidal_transit_time(const MagneticField<DeviceType>& magnetic_field) const{
        double main_ion_mass = all_species[MAIN_ION].mass;
        double axis_length = TWOPI * magnetic_field.equil.axis.r;
        double characteristic_velocity = sqrt(2*main_ion_characteristic_energy/main_ion_mass);
        return axis_length / characteristic_velocity;
    }

    void read_initial_distribution_all(NLReader::NamelistReader& nlr, const DomainDecomposition<DeviceType>& pol_decomp){

        for_all_species([&](Species<DeviceType>& species){
            if(!species.is_adiabatic && !species.maxwellian_init){
                species.read_initial_distribution(nlr, pol_decomp);
            }
        }, NoDevicePtl);
    }

    bool electrons_are_reduced_deltaf() const;
};

#endif