gradient_matrices.hpp

Go to the documentation of this file.

#ifndef GRADIENT_MATRICES_HPP
#define GRADIENT_MATRICES_HPP

#include "matrix.hpp"
#include "grid.hpp"

extern "C" void get_grad_mat_data(int ind, int* m, int* n, int* w, int* nnz, int* is_csr_int, int** csr_ridx_or_eindex, int** csr_cidx_or_nelement, double** values);

template<class Device>
class GradientMatrices{
    Matrix<Device> gradx;
    Matrix<Device> grady;

    public:

    /* Set up gradx and grady on GPU by (for now) copying from the fortran matrix */
    GradientMatrices(){
        int m, n, w, nnz, is_csr_int;
        int* csr_ridx_or_eindex;
        int* csr_cidx_or_nelement;
        double* values;
        const int X=0;
        const int Y=1;
        // gradx
        get_grad_mat_data(X, &m, &n, &w, &nnz, &is_csr_int, &csr_ridx_or_eindex, &csr_cidx_or_nelement, &values);
        gradx = Matrix<Device>(m, n, w, nnz, (is_csr_int==1), csr_ridx_or_eindex, csr_cidx_or_nelement, values);

        // grady
        get_grad_mat_data(Y, &m, &n, &w, &nnz, &is_csr_int, &csr_ridx_or_eindex, &csr_cidx_or_nelement, &values);
        grady = Matrix<Device>(m, n, w, nnz, (is_csr_int==1), csr_ridx_or_eindex, csr_cidx_or_nelement, values);
    }

    void grid_deriv(const Grid<Device>& grid, const Kokkos::View<double*,Kokkos::LayoutRight,Device>& qty, const Kokkos::View<double*,Kokkos::LayoutRight,Device>& qty_deriv_x, const Kokkos::View<double*,Kokkos::LayoutRight,Device>& qty_deriv_y, bool discard_when_basis_is_one) const{
        // Take the gradient of qty
        gradx.mult(qty, qty_deriv_x);
        grady.mult(qty, qty_deriv_y);

        // Discard poloidal electric field (only for vertices using the psi-theta basis)
        if (discard_when_basis_is_one){
            Kokkos::parallel_for("grid_deriv", Kokkos::RangePolicy<ExSpace>(0,grid.nnode), KOKKOS_LAMBDA( const int i ){
                if ( grid.basis(i)==0 ){
                    qty_deriv_y(i) = 0.0;
                }
            });
        }

    }

};
#endif