col_species.hpp

Go to the documentation of this file.

#ifndef COL_SPECIES_HPP
#define COL_SPECIES_HPP
 
#include <vector>
#include <Kokkos_Core.hpp>
#include <Kokkos_DualView.hpp>
 
#include "vgrid_distribution.hpp"
#include "moments.hpp"
 
struct MomentSet {
    double n; // Density
    double w; // Energy
    double p; // Momentum
 
    KOKKOS_INLINE_FUNCTION constexpr MomentSet() : n(0.0), w(0.0), p(0.0) {}
    //KOKKOS_INLINE_FUNCTION MomentSet(){}
    KOKKOS_INLINE_FUNCTION MomentSet(double n_in, double w_in, double p_in) : n(n_in), w(w_in), p(p_in) {}
 
    // Join operation
    KOKKOS_INLINE_FUNCTION void operator+=(const MomentSet& other) {
        n += other.n;
        w += other.w;
        p += other.p;
    }
};
 
/*struct Triple {
    double a;
    double b;
    double c;
 
    KOKKOS_INLINE_FUNCTION Triple() : a(0.0), b(0.0), c(0.0) {}
 
    // Join operation
    KOKKOS_INLINE_FUNCTION void operator+=(const Triple& other) {
        a += other.a;
        b += other.b;
        c += other.c;
    }
};*/
 
struct CollisionSpeciesScalars{
    int gid;               // Corresponding XGC species
    bool is_electron;      // Whether species is electron
    double T_avg;
    double mass;
    double charge_eu;
    double den_lambda_gamma;
 
    //> Computes the base collision frequency -- Eq. (6) in Hager et al.
    KOKKOS_INLINE_FUNCTION double lambda_gamma_pair(const CollisionSpeciesScalars& sp_b ) const {// C(a->b)
        double lambda;
        // NRL Plasma Formulary 2011 version, p.34
        //(c) Mixed ion-ion collisions (here, same species ion-ion collisions)
        if(is_electron || sp_b.is_electron){
            //if electron is invovled
            if(is_electron && sp_b.is_electron){
                //(a) Thermal electron-electron collisions
                // Note that pointers for colliding and target are same.
                lambda = 2.35e1 - log( sqrt(den_lambda_gamma*1e-6)*pow(T_avg,-1.25))
                                - sqrt(std::abs(1e-5 + pow(log(T_avg)-2,2)/16));
            } else {
                double ti_ev,massi,densi,chargei,te_ev,masse,dense;
                if(is_electron){
                    ti_ev = sp_b.T_avg;
                    massi = sp_b.mass;
                    densi = sp_b.den_lambda_gamma;
                    chargei = sp_b.charge_eu;
                    te_ev = T_avg;
                    masse = mass;
                    dense = den_lambda_gamma;
                } else {
                    te_ev = sp_b.T_avg;
                    masse = sp_b.mass;
                    dense = sp_b.den_lambda_gamma;
                    ti_ev = T_avg;
                    massi = mass;
                    densi = den_lambda_gamma;
                    chargei = charge_eu;
                }
                //(b) Electron-ion (and ion-electron) collisions
                if(ti_ev*masse/massi<te_ev){
                    if(te_ev<1e1*(chargei*chargei)){
                        lambda = 2.3e1 - log(chargei*sqrt(dense*1e-6/(te_ev*te_ev*te_ev)));
                    } else {
                        lambda = 2.4e1 - log(sqrt(dense*1e-6)/te_ev);
                    }
                } else {
                    lambda = 3e1-log((chargei*chargei)*sqrt(densi*1e-6/(ti_ev*ti_ev*ti_ev))*PROTON_MASS/massi);
                }
            }
        } else {
            lambda = 2.3e1 - log(charge_eu*sp_b.charge_eu*(mass+sp_b.mass)
                                      /(mass*sp_b.T_avg+sp_b.mass*T_avg)
                                      *sqrt( den_lambda_gamma*1e-6*(charge_eu*charge_eu)/T_avg
                                            +sp_b.den_lambda_gamma*1e-6*(sp_b.charge_eu*sp_b.charge_eu)/sp_b.T_avg));
        }
        double tmp = lambda*(charge_eu*charge_eu)*(sp_b.charge_eu*sp_b.charge_eu)*pow(UNIT_CHARGE,4)/(pow(8.8542e-12,2)*8*PI);
        return tmp/mass; // colliding : a, target : b
    }
};
 
template<class Device>
class CollisionSpecies{
    public:
    const VGridDistribution<HostType> f;
 
    Kokkos::DualView<CollisionSpeciesScalars**,Device> s;
 
    // Things moved out of CollisionSpeciesScalars
    View<int*, CLayout, DeviceType> grid_index_d; // Which grid this species is on
    View<int*, CLayout, HostType> grid_index_h;
    View<MomentSet**, CLayout, HostType> moments; // Set of moments used internally; note this is a different calculation from den_moment and temp_moment passed in!
    View<double**, CLayout, HostType> conv_factor;
 
    View<int*,CLayout,HostType> moments_index; // Index mapping to moments
 
    View<double**, CLayout, HostType> den_moment; // Density moments, shallow copy of moments.den_local
    View<double**, CLayout, HostType> temp_moment; // Temperature moments, shallow copy of moments.temp_local
 
    View<double*,CLayout,DeviceType> inv_mass_d; // Inverse mass of each species
    std::vector<View<double*,CLayout,HostType>> fg_temp_ev_all; // Vector of views
 
    View<double****,CLayout,HostType> pdf_h;    // Input distribution (batched)
    View<double****,CLayout,Device> pdf;    // Input distribution (device, batched).
    View<double****,CLayout,HostType> pdf1_h;  // Updated distribution (batched)
    View<double****,CLayout,Device> pdf1;  // Updated distribution (device, batched)
 
    CollisionSpecies<Device>(){}
 
    CollisionSpecies<Device>(const Moments& moments_in, const VGridDistribution<DeviceType>& f_in, const std::vector<Species<Device>> &all_species, int mb_n_nodes)
        : f(f_in.n_species(), f_in, VGridDistributionOption::NoViewInit),
          s("s", f.n_species(), mb_n_nodes),
          grid_index_d("grid_index_d", f.n_species()),
          grid_index_h("grid_index_h", f.n_species()),
          moments(NoInit("moments"), f.n_species(), mb_n_nodes), // note - unrelated to moments_in !!
          conv_factor(NoInit("conv_factor"), f.n_species(), mb_n_nodes),
          moments_index("moments_index", f.n_species()),
          den_moment(moments_in.den_local),
          temp_moment(moments_in.temp_local),
          inv_mass_d("inv_mass_d", f.n_species()),
          pdf("pdf", mb_n_nodes, f.n_species(), f.n_vz(), f.n_vr()),
          pdf1("pdf1", mb_n_nodes, f.n_species(), f.n_vz(), f.n_vr()) // WARNING: count_negatives assumes this index order!!
    {
        // Copy distribution
        Kokkos::deep_copy(f.f, f_in.f);
 
        std::vector<double> inv_mass;
        int col_species_cnt=0;
        for (int isp = 0; isp<all_species.size(); isp++){
            if (!all_species[isp].is_adiabatic){
                for (int mesh_ind = 0; mesh_ind<mb_n_nodes; mesh_ind++){
                    s.h_view(col_species_cnt, mesh_ind).gid = isp;
                    s.h_view(col_species_cnt, mesh_ind).is_electron = all_species[isp].is_electron;
                    s.h_view(col_species_cnt,mesh_ind).mass = all_species[isp].mass;
                    s.h_view(col_species_cnt,mesh_ind).charge_eu = all_species[isp].charge/UNIT_CHARGE;
                }
                grid_index_h(col_species_cnt) = all_species[isp].collision_grid_index;
                moments_index(col_species_cnt) = all_species[isp].nonadiabatic_idx;
 
                // This is a vector of Views, adding the species one at a time
                inv_mass.push_back(1.0/all_species[isp].mass);
                fg_temp_ev_all.push_back(all_species[isp].f0.fg_temp_ev_h);
 
                col_species_cnt++;
            }
        }
 
        Kokkos::deep_copy(grid_index_d, grid_index_h);
        array_deep_copy(inv_mass_d, inv_mass.data());
 
#ifdef USE_GPU
        pdf_h = View<double****,CLayout,HostType>("pdf_h",pdf.layout());
        pdf1_h = View<double****,CLayout,HostType>("pdf1_h",pdf1.layout());
#else
        pdf_h = pdf;
        pdf1_h = pdf1;
#endif
    }
 
    // Copy col_spall_in scalar data, but resized to nnode_in
    CollisionSpecies<Device>(const CollisionSpecies<Device>& col_spall_in, int nnode_in)
        : f(col_spall_in.f.n_species(), nnode_in, col_spall_in.f, VGridDistributionOption::NoViewInit),
          s(col_spall_in.s), // SHALLOW COPY
          grid_index_d(col_spall_in.grid_index_d), // SHALLOW COPY
          grid_index_h(col_spall_in.grid_index_h), // SHALLOW COPY
          moments_index(col_spall_in.moments_index), // SHALLOW COPY
          moments(NoInit("moments"), f.n_species(), col_spall_in.pdf.extent(0)), // note - unrelated to moments_in !!
          conv_factor(NoInit("conv_factor"), f.n_species(), col_spall_in.pdf.extent(0)),
          den_moment(NoInit("den_moment"), nnode_in, f.n_species()),
          temp_moment(NoInit("temp_moment"), nnode_in, f.n_species()),
          inv_mass_d(col_spall_in.inv_mass_d), // SHALLOW COPY
          pdf("pdf", col_spall_in.pdf.extent(0), f.n_species(), f.n_vz(), f.n_vr()),
          pdf1("pdf1", col_spall_in.pdf1.extent(0), f.n_species(), f.n_vz(), f.n_vr())
    {
        for(int i=0; i<f.n_species(); i++){
            // This is a vector of Views, adding the species one at a time
            fg_temp_ev_all.push_back(View<double*,CLayout,HostType>(NoInit("fg_temp_ev"), nnode_in));
        }
 
#ifdef USE_GPU
        pdf_h = View<double****,CLayout,HostType>("pdf_h",pdf.layout());
        pdf1_h = View<double****,CLayout,HostType>("pdf1_h",pdf1.layout());
#else
        pdf_h = pdf;
        pdf1_h = pdf1;
#endif
    }
 
    //> Convert the distribution function to the units used
    // internally by the collision operator
    // The distribution function set up with f0_module is in normalized velocity space
    // whereas the collision operator uses SI units.
    void setup_one(int isp, int mesh_ind, int local_node_ind) const{
 
        double conv_inv_mass=TWOPI*UNIT_CHARGE/s.h_view(isp,mesh_ind).mass;
        int moment_isp = moments_index(isp);
        s.h_view(isp,mesh_ind).den_lambda_gamma = den_moment(local_node_ind, moment_isp);
        s.h_view(isp,mesh_ind).T_avg = temp_moment(local_node_ind, moment_isp);
 
 
        double conv_fac_temp = fg_temp_ev_all[moment_isp](local_node_ind);
        conv_factor(isp,mesh_ind) = sqrt(conv_fac_temp*(conv_inv_mass*conv_inv_mass*conv_inv_mass));
        double inv_conv_factor = 1.0/conv_factor(isp,mesh_ind);
 
        // Set up distribution
        for (int imu = 0; imu < pdf_h.extent(3); imu++){
            double vnorm_to_SI = inv_conv_factor/f.get_smu_n(imu);
            for (int ivp = 0; ivp<pdf_h.extent(2); ivp++){
                double& var = pdf_h(mesh_ind,isp,ivp,imu);
 
                // Set to assigned f - involves transpose, hence why this can't just be a copy
                var = f(moment_isp, imu,local_node_ind,ivp);
 
                // Cut off all negative values
                var = std::max(var,0.0);
 
                // Convert to SI
                var *= vnorm_to_SI;
            }
        }
    }
 
    void setup_all(Kokkos::View<int*,HostType>& mesh_nodes, int mb_n_nodes) const{
        Kokkos::parallel_for("col_sp_setup", Kokkos::RangePolicy<HostExSpace>( 0, mb_n_nodes ), [=] (const int mesh_ind){
            for (int isp = 0; isp < s.extent(0); isp++){
                setup_one(isp, mesh_ind, mesh_nodes(mesh_ind));
            }
        });
 
#ifdef USE_GPU
        // Copy dual views to device
        Kokkos::deep_copy(s.d_view, s.h_view);
        Kokkos::deep_copy(pdf, pdf_h);
#endif
    }
 
    void lambda_gamma_pair(const View<double*,CLayout,HostType>& dt, Kokkos::View<double***,Device>& gammac, int mb_n_nodes) const{
        // Done on host, so create a mirror
        //typename Kokkos::View<double***,Device>::HostMirror
        auto gammac_h = Kokkos::create_mirror_view(gammac);
 
        for (int mesh_ind = 0; mesh_ind<mb_n_nodes; mesh_ind++){
            // gamma
            for (int isp = 0; isp < s.extent(0); isp++){
                for (int jsp = 0; jsp < s.extent(0); jsp++){
                    gammac_h(mesh_ind,jsp,isp) = dt(mesh_ind)*s.h_view(isp,mesh_ind).lambda_gamma_pair(s.h_view(jsp,mesh_ind));
                }
            }
        }
        Kokkos::deep_copy(gammac, gammac_h);
    }
 
    inline int n() const{
        return f.n_species();
    }
};
 
#endif