plasma.hpp

Go to the documentation of this file.

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

// Fortran routine:
extern "C" int get_spall_num_g(int isp);

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; 

    // At some point remove nspecies from class Simulation to only have in class Plasma
    int nspecies;                 
    int n_nonadiabatic_species;   

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

    Plasma(NLReader::NamelistReader& nlr)
      : particles_d_has_owner(false),
        species_share_particles_d_ownership(true)
    {
        nlr.use_namelist("ptl_param");
        nspecies = nlr.get<int>("ptl_nsp",1);
        nspecies += 1; // Add electrons
        n_nonadiabatic_species = 0;

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

    // Loop over all species
    template<typename F>
    inline void for_all_species(F func){
        for(int isp = 0; isp<all_species.size(); isp++){
            if(species_share_particles_d_ownership) transfer_particles_d_ownership(isp);
            func(all_species[isp]);
        }
    }

    // Loop over all non-adiabatic species
    template<typename F>
    inline void for_all_nonadiabatic_species(F func){
        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){
                if(species_share_particles_d_ownership) transfer_particles_d_ownership(isp);
                func(all_species[isp]);
            }
        }
    }

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

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

    // Operate on one species
    template<typename F>
    inline void for_one_species(int isp, F func){
        if(species_share_particles_d_ownership) transfer_particles_d_ownership(isp);
        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,get_spall_num_g(species.idx));
            }
        });

        return max_n_ptl;
    }

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

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

};

#endif