pseudo_inverse.hpp

Go to the documentation of this file.

#ifndef PSEUDO_INVERSE_HPP
#define PSEUDO_INVERSE_HPP
#include <Kokkos_Core.hpp>
#include "space_settings.hpp"
#include "sml.hpp"
#include "plasma.hpp"
#include "DM_wrapper.hpp"

#ifndef NO_PETSC
#include <petscversion.h>
#include <petscsys.h>
#include <petscdmplex.h>
#if PETSC_VERSION_GE(3,14,0)
#include <petscds.h>
#include "petscdmswarm.h"
#include "petscksp.h"
#endif
#endif

#ifndef NO_PETSC
// For some reason Kokkos cannot handle Views of PETSc objects, but it works when putting them in a wrapper
// Pseudo-inverse PETSc objects
struct PseudoInversePetscObjects{
    DM  swarm; 
    Vec rho;   
    Vec rhs;   
    Mat M_p;   
};
#endif

template<class Device>
class Pseudo_inverse{

    public:

#ifndef NO_PETSC
    // All PETSc objects on host
    Kokkos::View<PseudoInversePetscObjects*, Kokkos::LayoutRight, HostType> obj; 
#endif

    // Arrays on Device and Host:
    Kokkos::View<double***, Kokkos::LayoutRight, Device> ptl_coords; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> ptl_coords_h; 

    Kokkos::View<double***, Kokkos::LayoutRight, Device> ptl_weights; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> ptl_weights_h; 

    Kokkos::View<double***, Kokkos::LayoutRight, Device> ptl_alpha_weights; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> ptl_alpha_weights_h; 

    Kokkos::View<int***, Kokkos::LayoutRight, Device> ptl_indices; 
    Kokkos::View<int***, Kokkos::LayoutRight, Kokkos::HostSpace> ptl_indices_h; 

    Kokkos::View<int**, Kokkos::LayoutRight, Device> n_ptl; 
    Kokkos::View<int**, Kokkos::LayoutRight, Kokkos::HostSpace> n_ptl_h; 

    Kokkos::View<double***, Kokkos::LayoutRight, Device> grid_coords; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> grid_coords_h; 

    Kokkos::View<double***, Kokkos::LayoutRight, Device> grid_weights; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> grid_weights_h; 

    Kokkos::View<double***, Kokkos::LayoutRight, Device> grid_alpha_weights; 
    Kokkos::View<double***, Kokkos::LayoutRight, Kokkos::HostSpace> grid_alpha_weights_h; 

    Kokkos::View<bool**, Kokkos::LayoutRight, Device> vgrid_filled; 
    Kokkos::View<bool**, Kokkos::LayoutRight, Kokkos::HostSpace> vgrid_filled_h; 

    int expected_n_grid_points; 
    bool force; 
    bool do_particles_to_grid_petsc; 
    bool use_delta_grid; 
    int n_nonadiabatic_species; 

    // Initialize these Views but don't allocate any memory yet; Instead, resize when needed
    Pseudo_inverse()
        : ptl_coords("ptl_coords", 0, 0, 0),
          ptl_coords_h("ptl_coords_h", 0, 0, 0),
          ptl_weights("ptl_weights", 0, 0, 0),
          ptl_weights_h("ptl_weights_h", 0, 0, 0),
          ptl_alpha_weights("ptl_alpha_weights", 0, 0, 0),
          ptl_alpha_weights_h("ptl_alpha_weights_h", 0, 0, 0),
          ptl_indices("ptl_indices", 0, 0, 0),
          ptl_indices_h("ptl_indices_h", 0, 0, 0),
          n_ptl("n_ptl", 0, 0),
          n_ptl_h("n_ptl_h", 0, 0),
          grid_coords("grid_coords", 0, 0, 0),
          grid_coords_h("grid_coords_h", 0, 0, 0),
          grid_weights("grid_weights", 0, 0, 0),
          grid_weights_h("grid_weights_h", 0, 0, 0),
          grid_alpha_weights("grid_alpha_weights", 0, 0, 0),
          grid_alpha_weights_h("grid_alpha_weights_h", 0, 0, 0),
          vgrid_filled("vgrid_filled", 0, 0),
          vgrid_filled_h("vgrid_filled_h", 0, 0),
          expected_n_grid_points(0),
          force(false),
          do_particles_to_grid_petsc(false),
          use_delta_grid(false),
          n_nonadiabatic_species(0){
#ifndef NO_PETSC
        Kokkos::resize(obj, 0);
#endif
    }

    // Allocate memory for the arrays by resizing them
    void resize_on_device_and_host(int nthreads, int n_non_ad, int nnode, int max_n_ptl_on_node, int n_grid_points){
#ifndef NO_PETSC
        Kokkos::resize(obj, nthreads);
#endif
        Kokkos::resize(ptl_coords, n_non_ad, nnode, 2*max_n_ptl_on_node);
        Kokkos::resize(ptl_coords_h, n_non_ad, nnode, 2*max_n_ptl_on_node);
        Kokkos::resize(ptl_weights, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(ptl_weights_h, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(ptl_alpha_weights, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(ptl_alpha_weights_h, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(ptl_indices, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(ptl_indices_h, n_non_ad, nnode, max_n_ptl_on_node);
        Kokkos::resize(n_ptl, n_non_ad, nnode);
        Kokkos::resize(n_ptl_h, n_non_ad, nnode);
        Kokkos::resize(grid_coords, n_non_ad, nnode, 2*n_grid_points);
        Kokkos::resize(grid_coords_h, n_non_ad, nnode, 2*n_grid_points);
        Kokkos::resize(grid_weights, n_non_ad, nnode, n_grid_points);
        Kokkos::resize(grid_weights_h, n_non_ad, nnode, n_grid_points);
        Kokkos::resize(grid_alpha_weights, n_non_ad, nnode, n_grid_points);
        Kokkos::resize(grid_alpha_weights_h, n_non_ad, nnode, n_grid_points);
        Kokkos::resize(vgrid_filled, n_non_ad, nnode);
        Kokkos::resize(vgrid_filled_h, n_non_ad, nnode);
        n_nonadiabatic_species = n_non_ad;
    }

    // Deallocate memory of the device arrays by resizing back to 0
    void resize_to_zero(){
#ifndef NO_PETSC
        Kokkos::resize(obj, 0);
#endif
        Kokkos::resize(ptl_coords, 0, 0, 0);
        Kokkos::resize(ptl_coords_h, 0, 0, 0);
        Kokkos::resize(ptl_weights, 0, 0, 0);
        Kokkos::resize(ptl_weights_h, 0, 0, 0);
        Kokkos::resize(ptl_alpha_weights, 0, 0, 0);
        Kokkos::resize(ptl_alpha_weights_h, 0, 0, 0);
        Kokkos::resize(ptl_indices, 0, 0, 0);
        Kokkos::resize(ptl_indices_h, 0, 0, 0);
        Kokkos::resize(n_ptl, 0, 0);
        Kokkos::resize(n_ptl_h, 0, 0);
        Kokkos::resize(grid_coords, 0, 0, 0);
        Kokkos::resize(grid_coords_h, 0, 0, 0);
        Kokkos::resize(grid_weights, 0, 0, 0);
        Kokkos::resize(grid_weights_h, 0, 0, 0);
        Kokkos::resize(grid_alpha_weights, 0, 0, 0);
        Kokkos::resize(grid_alpha_weights_h, 0, 0, 0);
        Kokkos::resize(vgrid_filled, 0, 0);
        Kokkos::resize(vgrid_filled_h, 0, 0);
        n_nonadiabatic_species = 0;
    }

    // Deallocate memory of the alpha weights device arrays by resizing back to 0
    void resize_alpha_weights_to_zero(){
        Kokkos::resize(ptl_alpha_weights, 0, 0, 0);
        Kokkos::resize(ptl_alpha_weights_h, 0, 0, 0);
        Kokkos::resize(grid_alpha_weights, 0, 0, 0);
        Kokkos::resize(grid_alpha_weights_h, 0, 0, 0);
    }

    // Send arrays between device and host
    void send_device_to_host(){
        Kokkos::deep_copy(ptl_coords_h, ptl_coords);
        Kokkos::deep_copy(ptl_weights_h, ptl_weights);
        Kokkos::deep_copy(ptl_alpha_weights_h, ptl_alpha_weights);
        Kokkos::deep_copy(ptl_indices_h, ptl_indices);
        Kokkos::deep_copy(n_ptl_h, n_ptl);
        Kokkos::deep_copy(grid_coords_h, grid_coords);
        Kokkos::deep_copy(grid_weights_h, grid_weights);
        Kokkos::deep_copy(grid_alpha_weights_h, grid_alpha_weights);
        Kokkos::deep_copy(vgrid_filled_h, vgrid_filled);
    }

    void send_host_to_device(){
        Kokkos::deep_copy(ptl_coords, ptl_coords_h);
        Kokkos::deep_copy(ptl_weights, ptl_weights_h);
        Kokkos::deep_copy(ptl_alpha_weights, ptl_alpha_weights_h);
        Kokkos::deep_copy(ptl_indices, ptl_indices_h);
        Kokkos::deep_copy(n_ptl, n_ptl_h);
        Kokkos::deep_copy(grid_coords, grid_coords_h);
        Kokkos::deep_copy(grid_weights, grid_weights_h);
        Kokkos::deep_copy(grid_alpha_weights, grid_alpha_weights_h);
        Kokkos::deep_copy(vgrid_filled, vgrid_filled_h);
    }

    void send_ptl_device_to_host(){
        Kokkos::deep_copy(ptl_coords_h, ptl_coords);
        Kokkos::deep_copy(ptl_weights_h, ptl_weights);
        Kokkos::deep_copy(ptl_alpha_weights_h, ptl_alpha_weights);
        Kokkos::deep_copy(ptl_indices_h, ptl_indices);
        Kokkos::deep_copy(n_ptl_h, n_ptl);
        //Kokkos::deep_copy(vgrid_filled_h, vgrid_filled);
    }

    void send_grid_host_to_device(){
        Kokkos::deep_copy(grid_coords, grid_coords_h);
        Kokkos::deep_copy(grid_weights, grid_weights_h);
        Kokkos::deep_copy(grid_alpha_weights, grid_alpha_weights_h);
    }

    void set_vgrid_filled(Kokkos::View<bool**,Kokkos::LayoutRight,DeviceType> vgrid_filled_in){
        Kokkos::deep_copy(vgrid_filled, vgrid_filled_in);
        Kokkos::deep_copy(vgrid_filled_h, vgrid_filled);
    }

    void zero_out_grid_weights(){
        Kokkos::deep_copy(grid_weights, 0.0);
        Kokkos::deep_copy(grid_weights_h, 0.0);
    }
};

#ifndef NO_PETSC
PetscInt pseudo_inverse_interpolate_particles_to_velocity_grid(const DM &dm,
                                   Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> particle_coords,
                                   Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> particle_weights,
                                   const PetscInt Nparticles, const PetscInt Nspecies,
                                   Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> grid_coords,
                                   Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> grid_weights);

PetscErrorCode pseudo_inv_create_Swarm(const DM dm, DM *sw);

PetscErrorCode pseudo_inv_create_projection(const DM dm, DM sw,
                               const PetscInt thread_id, const PetscInt target_id,
                               Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> particle_coords,
                               Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> particle_weights,
                               const PetscInt Nparticles,
                               Vec rho, Mat *Mp_out);

PetscErrorCode pseudo_inv_set_rhs(const DM dm, Mat MM,
                                  Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> grid_weights,
                                  const PetscInt Ngrid, Vec rho, Vec rhs);

PetscErrorCode pseudo_inv_grid_to_particles(const DM dm, DM sw,
                                            Vec rhs, Mat M_p, Mat PM_p,
                                            const PetscInt Ngrid,
                                            const PetscInt Nparticles,
                                            const PetscInt order,
                                            Kokkos::View<PetscReal*, Kokkos::LayoutRight, Kokkos::HostSpace> particle_weights,
                                            int &ksp_num_its,
                                            bool &ksp_converged);

#endif
//NO_PETSC

void pseudo_inv_allocate(const Simulation<DeviceType>& sml,
                       const Grid<DeviceType>& grid,
                       const MagneticField<DeviceType>& magnetic_field,
                       Plasma& plasma,
                       const VelocityGrid& vgrid,
                       const DomainDecomposition<DeviceType>& pol_decomp,
                       Pseudo_inverse<DeviceType>& pseudo_inv);

void all_species_update_pseudo_inv_weights(const Simulation<DeviceType>& sml,
                       const Grid<DeviceType>& grid,
                       const MagneticField<DeviceType>& magnetic_field,
                       Plasma& plasma,
                       const VelocityGrid& vgrid,
                       const DomainDecomposition<DeviceType>& pol_decomp,
                       DMWrapper& pseudo_inv_dm,
                       Pseudo_inverse<DeviceType>& pseudo_inv);

int max_n_ptl_on_node_and_vgrid_filled(const Simulation<DeviceType>& sml,
                       const Grid<DeviceType>& grid,
                       const MagneticField<DeviceType>& magnetic_field,
                       Plasma& plasma,
                       const VelocityGrid& vgrid,
                       const DomainDecomposition<DeviceType>& pol_decomp,
                       const std::string &msg,
                       Kokkos::View<bool**,Kokkos::LayoutRight,DeviceType> vgrid_filled);

template<class Device> KOKKOS_INLINE_FUNCTION bool pseudo_inv_do_node(const Pseudo_inverse<Device>& pseudo_inv, int isp_non_ad, int node){
    return (pseudo_inv.n_ptl(isp_non_ad, node) > 0 && (pseudo_inv.vgrid_filled(isp_non_ad, node) || pseudo_inv.force));
}

inline bool pseudo_inv_do_node_h(Pseudo_inverse<DeviceType>& pseudo_inv, int isp_non_ad, int node){
    return (pseudo_inv.n_ptl_h(isp_non_ad, node) > 0 && (pseudo_inv.vgrid_filled_h(isp_non_ad, node) || pseudo_inv.force));
}
#endif