moments.hpp

Go to the documentation of this file.

#ifndef MOMENTS_HPP
#define MOMENTS_HPP

#include "grid.hpp"
#include "domain_decomposition.hpp"
#include "vgrid_distribution.hpp"
#include "plasma.hpp"
#include "basic_physics.hpp"

extern "C" void set_mom_module_ptrs(double* den_local_cpp, double* upara_local_cpp, double* uperp_local_cpp, double* temp_local_cpp, double* Kperp_local_cpp,
        double* Kpara_local_cpp, double* temp_para_local_cpp, double* den_all_cpp, double* temp_all_cpp, double* temp_perp_all_cpp, double* temp_para_all_cpp, double* upara_all_cpp);

struct Moments{
    View<double**, CLayout, HostType> den_local; 
    View<double**, CLayout, HostType> upara_local; 
    View<double**, CLayout, HostType> uperp_local; 
    View<double**, CLayout, HostType> temp_local;  
    View<double**, CLayout, HostType> Kperp_local; 
    View<double**, CLayout, HostType> Kpara_local; 
    View<double**, CLayout, HostType> temp_para_local;  

    View<double**, CLayout, HostType> den_all;   
    View<double**, CLayout, HostType> upara_all; 
    View<double**, CLayout, HostType> temp_all;  
    View<double**, CLayout, HostType> temp_perp_all; 
    View<double**, CLayout, HostType> temp_para_all; 

    Moments(){}

    Moments(const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid, Plasma& plasma, const VGridDistribution<HostType>& f, bool do_gather)
      : den_local("den_local", f.n_nodes(), f.n_species()),
        upara_local("upara_local", den_local.layout()),
        uperp_local("uperp_local", den_local.layout()),
        temp_local(NoInit("temp_local"), den_local.layout()), // NoInit because it is derived from the other quantities
        Kperp_local("Kperp_local", den_local.layout()),
        Kpara_local("Kpara_local", den_local.layout()),
        temp_para_local(NoInit("temp_para_local"), den_local.layout()), // NoInit because it is derived from the other quantities

        // These don't need to be allocated if do_gather==false, but allocate anyway for now for Fortran pointers
        den_all(NoInit("den_all"), grid.nnode, f.n_species()),
        upara_all(NoInit("upara_all"), den_all.layout()),
        temp_all(NoInit("temp_all"), den_all.layout()),
        temp_perp_all(NoInit("temp_perp_all"), den_all.layout()),
        temp_para_all(NoInit("temp_para_all"), den_all.layout())
    {
        // Calculate moments of f in local domain
        calculate_moments_of_distribution(vgrid, plasma, f);

        // Gather decomposed moments if requested
        if(do_gather) gather(pol_decomp);

        // Set fortran pointers
        set_mom_module_ptrs(den_local.data(), upara_local.data(), uperp_local.data(), temp_local.data(), Kperp_local.data(),
                            Kpara_local.data(), temp_para_local.data(), den_all.data(), temp_all.data(), temp_perp_all.data(),
                            temp_para_all.data(), upara_all.data());
    }

    void calculate_moments_of_distribution(const VelocityGrid& vgrid, Plasma& plasma, const VGridDistribution<HostType>& f){
        // Create local copies of these member views so that they can be passed into the lambda function
        auto den = den_local;
        auto upara = upara_local;
        auto uperp = uperp_local;
        auto temp = temp_local;
        auto Kperp = Kperp_local;
        auto Kpara = Kpara_local;
        auto temp_para = temp_para_local;

        plasma.for_all_nonadiabatic_species([&](Species<DeviceType>& species){
            int isp = species.nonadiabatic_idx;

            // Calculate sum of f, E*f, vp*f
            Kokkos::parallel_for("calculate_f0_charge_den", Kokkos::RangePolicy<HostExSpace>( 0, f.n_nodes()), KOKKOS_LAMBDA(const int inode){
                for (int ivr=0; ivr<vgrid.nvr; ivr++){
                    // normalized volume element
                    double mu_vol = (ivr==0 || ivr==vgrid.nvr-1) ? 0.5 : 1.0;
                    double smu = ivr*vgrid.dsmu;
                    double mu = smu*smu; // <--- This is (v_perp/v_th)^2 for v_perp grid and mu_N for sqrt(mu) grid!!!!
                    double en_perp= 0.5 * mu;    // energy normalized by T, 0.5*(vperp/vth)^2

                    // Sum over parallel velocity
                    for(int ivz=0; ivz<vgrid.nvz; ivz++){
                        int ivp = ivz - vgrid.nvp;
                        double vp = ivp*vgrid.dvp;
                        double en_para2 = vp*vp; // energy normalized by T, 0.5*(vpara/vth)^2

                        double wt = mu_vol*f(isp, ivr, inode, ivz);
                        den(inode,isp)   += wt;
                        upara(inode,isp) += wt*vp;
                        uperp(inode,isp) += wt*smu;
                        Kperp(inode,isp) += wt*en_perp;
                        Kpara(inode,isp) += wt*en_para2;
                    }
                }
            });
            Kokkos::fence();

            // Normalization of all moments and compute temp_para and temp
            const double UPERP_DEFAULT_COEFF = sqrt(0.25*TWOPI);
            Kokkos::parallel_for("calculate_f0_charge_den", Kokkos::RangePolicy<HostExSpace>( 0, f.n_nodes()), KOKKOS_LAMBDA(const int inode){

                // Impose lower limit on density
                double den0 = 1.0e-4*species.f0.den_h(inode); // For safety, to prevent inf and NaN
                if (!std::isfinite(den(inode,isp))) den(inode,isp) = den0;
                den(inode,isp) = max(den0,den(inode,isp));

                // Normalized v
                double t_norm_ev = species.f0.temp_ev_h(inode);
                double vth = species.get_f0_eq_thermal_velocity_lnode_h(inode);

                // Normalize by density
                double inv_density = 1.0/den(inode,isp);
                Kperp(inode,isp) *= t_norm_ev*inv_density;
                Kpara(inode,isp) *= t_norm_ev*inv_density;
                upara(inode,isp) *= vth*inv_density;
                uperp(inode,isp) *= vth*inv_density;
                den(inode,isp)   *= species.f0.grid_vol_vonly_h(inode);

                // Calculate temp_para and temp
                temp_para(inode,isp) = Kpara(inode,isp) - 2.0*kinetic_energy(species.mass, upara(inode,isp))*J_2_EV;
                temp(inode,isp) = ( temp_para(inode,isp) + 2.0*Kperp(inode,isp) )/3.0;

                // Set possible NaNs and Infs to default values
                if (!std::isfinite(den(inode,isp)))       den(inode,isp)       = species.f0.den_h(inode);
                if (!std::isfinite(Kperp(inode,isp)))     Kperp(inode,isp)     = t_norm_ev;
                if (!std::isfinite(Kpara(inode,isp)))     Kpara(inode,isp)     = t_norm_ev;
                if (!std::isfinite(temp_para(inode,isp))) temp_para(inode,isp) = t_norm_ev;
                if (!std::isfinite(temp(inode,isp)))      temp(inode,isp)      = t_norm_ev;
                if (!std::isfinite(upara(inode,isp)))     upara(inode,isp)     = species.f0.flow_h(inode)*vth;
                if (!std::isfinite(uperp(inode,isp)))     uperp(inode,isp)     = UPERP_DEFAULT_COEFF*vth;

                // Impose lower and upper limits
                den       (inode,isp) = min(max(den       (inode,isp),1.0e-2*species.f0.den_h(inode)),1.0e2*species.f0.den_h(inode));
                Kperp     (inode,isp) = min(max(Kperp     (inode,isp),1.0e-2*t_norm_ev              ),1.0e2*t_norm_ev              );
                Kpara     (inode,isp) = min(max(Kpara     (inode,isp),1.0e-2*t_norm_ev              ),1.0e2*t_norm_ev              );
                temp      (inode,isp) = min(max(temp      (inode,isp),1.0e-2*t_norm_ev              ),1.0e2*t_norm_ev              );
                temp_para (inode,isp) = min(max(temp_para (inode,isp),1.0e-2*t_norm_ev              ),1.0e2*t_norm_ev              );
                upara     (inode,isp) = min(max(upara     (inode,isp),-0.75*vgrid.vp_max*vth        ),0.75*vgrid.vp_max*vth        );
                uperp     (inode,isp) = min(max(uperp     (inode,isp),0.05*vgrid.vp_max*vth         ),0.75*vgrid.vp_max*vth        );
            });
            Kokkos::fence();
        }, Plasma::NoDevicePtl);
    }

    // Gathers all moments on local nodes, and distributes copies to all compute nodes
    void gather(const DomainDecomposition<DeviceType>& pol_decomp){
#ifdef USE_MPI
        View<int*,CLayout,HostType> cnts(NoInit("cnts"), pol_decomp.mpi.n_plane_ranks);
        View<int*,CLayout,HostType> displs(NoInit("cnts"), pol_decomp.mpi.n_plane_ranks);

        // Create cnts array with # vertex nodes on each processor
        int nsp = den_local.extent(1);
        for(int i = 0; i<cnts.size(); i++){
            displs(i) = nsp*(pol_decomp.gvid0_pid_h(i) - 1); // gvid0_pid_h is 1-indexed
            cnts(i)   = nsp*(pol_decomp.gvid0_pid_h(i+1) - pol_decomp.gvid0_pid_h(i));
        }

        int my_send_count = cnts(pol_decomp.mpi.my_plane_rank);
        MPI_Allgatherv(den_local.data(),      my_send_count, MPI_DOUBLE, den_all.data(),       cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
        MPI_Allgatherv(temp_local.data(),     my_send_count, MPI_DOUBLE, temp_all.data(),      cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
        MPI_Allgatherv(temp_para_local.data(),my_send_count, MPI_DOUBLE, temp_para_all.data(), cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
        MPI_Allgatherv(Kperp_local.data(),    my_send_count, MPI_DOUBLE, temp_perp_all.data(), cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
        MPI_Allgatherv(upara_local.data(),    my_send_count, MPI_DOUBLE, upara_all.data(),     cnts.data(), displs.data(), MPI_DOUBLE, pol_decomp.mpi.plane_comm);
#endif
    }
};

#endif