poly_basis.hpp

Go to the documentation of this file.

#ifndef POLYNOMIAL_BASIS_HPP
#define POLYNOMIAL_BASIS_HPP

#include "uniform_range.hpp"
#include "vgrid_distribution.hpp"

/* Struct to store a polynomial representation of a distribution function */
template<class Device>
struct PolynomialBasisDistribution{
    View<double***, CLayout, Device> f;

    PolynomialBasisDistribution(int n_poly, int n)
      : f("f", n_poly, n_poly, n){} // Initialize to 0.0
};

/* Class for storing a set of basis functions for a 2D grid, as well as functions for converting a distribution
 * function to and from a polynomial representation */
template<class Device>
class PolynomialBasis{
    using exec_space = typename Device::execution_space;

    int n_poly;
    UniformRange xrange;
    UniformRange yrange;
    View<double***, CLayout, Device> bases;
    View<double***, CLayout, HostType> bases_h; // Bases construct on host for now
    View<double****, CLayout, Device> matrix;

    // Making these functions public for the moment, due to:
    // The enclosing parent function ("compute_inner_product") for an extended __host__ __device__ lambda cannot have private or protected access within its class
    public:

    View<double***, CLayout, Device> compute_inner_product(const View<double***, CLayout, Device>& dist_in) const{
        int nnode = dist_in.extent(0);
        View<double***, CLayout, Device> inner_prod_matrix("inner_prod_matrix", n_poly, n_poly, nnode);

        const PolynomialBasis pb = *this;

        Kokkos::parallel_for("inner_product", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
            for (int iy = 0; iy < pb.yrange.n; iy++) {
                for (int ix = 0; ix < pb.xrange.n; ix++) {
                    double w = dist_in(inode, ix, iy);

                    for(int ip = 0; ip<pb.n_poly; ip++){
                        for(int jp=ip; jp<pb.n_poly; jp++){
                            inner_prod_matrix(ip, jp, inode) += w*pb.matrix(ip, jp, ix, iy);
                        }
                    }
                }
            }

            // Reflect to create symmetric matrix
            for(int ip = 0; ip<pb.n_poly; ip++){
                for(int jp=ip+1; jp<pb.n_poly; jp++){
                    inner_prod_matrix(jp, ip, inode) = inner_prod_matrix(ip, jp, inode);
                }
            }

            // Normalize
            double den = inner_prod_matrix(0,0,inode);
            for(int i=0; i<pb.n_poly; i++){
                for(int j=0; j<pb.n_poly; j++){
                    inner_prod_matrix(i,j,inode) /= den;
                }
            }
        });
        Kokkos::fence();
        return inner_prod_matrix;
    }

    PolynomialBasisDistribution<Device> get_coefficients(const View<double***, CLayout, Device>& inner_prod_matrix) const{
        int nnode = inner_prod_matrix.extent(2);
        PolynomialBasisDistribution<Device> polynom(n_poly, nnode);

        View<double**, CLayout, Device> norm("norm", n_poly, nnode); // Initialize to 0.0;

        const PolynomialBasis pb = *this;
        Kokkos::parallel_for("get_coeffs", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
            // The first one is trivial
            polynom.f(0, 0, inode) = 1.0;
            for (int j=0; j<pb.n_poly; j++) {
                double sum_value = 0.0;
                for (int k = 0; k < pb.n_poly; k++) {
                    sum_value += polynom.f(k, 0, inode) * inner_prod_matrix(k, j, inode);
                }
                norm(0, inode) += polynom.f(j, 0, inode) * sum_value;
            }

            // Now the rest
            for (int i = 1; i<pb.n_poly; i++){
                polynom.f(i, i, inode) = 1.0;
                for (int j = i-1; j >= 0; j--) {
                    double sum_value = 0.0;
                    for (int k = 0; k < pb.n_poly; ++k) {
                        sum_value += inner_prod_matrix(k, i, inode) * polynom.f(k, j, inode);
                    }
                    for (int l = 0; l < pb.n_poly; ++l) {
                        polynom.f(l, i, inode) -= sum_value / norm(j, inode) * polynom.f(l, j, inode);
                    }
                }
                for (int j = 0; j < pb.n_poly; j++) {
                    double sum_value = 0.0;
                    for (int k = 0; k < pb.n_poly; ++k) {
                        sum_value += polynom.f(k, i, inode) * inner_prod_matrix(k, j, inode);
                    }
                    norm(i, inode) += polynom.f(j, i, inode) * sum_value;
                }
            }

            for (int i = 0; i < pb.n_poly; i++) {
                double norm_sqrt = sqrt(norm(i, inode));
                for (int j = 0; j < pb.n_poly; j++) {
                    polynom.f(j, i, inode) /= norm_sqrt;
                }
            }
        });
        Kokkos::fence();
        return polynom;
    }

    template<typename F>
    void add_basis(int& basis_ind, F func){
        // This should be small - keep on host for now
        for(int i=0; i<xrange.n; i++){
            double x = xrange.get_val(i);
            for(int j=0; j<yrange.n; j++){
                double y = yrange.get_val(j);
                bases_h(basis_ind, i, j) = func(x, y);
            }
        }

        basis_ind += 1; // Move to next basis function
    }

    public:

    template<typename... Fs>
    PolynomialBasis(const UniformRange& xrange_in, const UniformRange& yrange_in, Fs... funcs)
      : xrange(xrange_in),
        yrange(yrange_in),
        n_poly(sizeof...(Fs)), // sizeof counts number of entries when applied to a parameter pack
        bases(NoInit("bases"), n_poly, xrange.n, yrange.n),
        bases_h(NoInit("bases_h"), bases.layout()),
        matrix(NoInit("matrix"), n_poly, n_poly, xrange.n, yrange.n)
    {
        // Step through the parameter pack of functions, populating the individual basis functions
        int basis_ind = 0;
        (..., add_basis(basis_ind, funcs));

        // Compute the matrix - small so keep on host for now
        View<double****, CLayout, HostType> matrix_h(NoInit("matrix_h"), matrix.layout());
        for(int ip=0; ip<n_poly; ip++){
            for(int jp=ip; jp<n_poly; jp++){
                for(int i=0; i<xrange.n; i++){
                    for(int j=0; j<yrange.n; j++){
                        matrix_h(ip, jp, i, j) = bases_h(ip, i, j)*bases_h(jp, i, j);
                    }
                }
            }
        }
        Kokkos::deep_copy(bases, bases_h);
        Kokkos::deep_copy(matrix, matrix_h);
    }

    PolynomialBasisDistribution<Device> get_distribution(const View<double***, CLayout, Device>& dist_in) const{
        auto inner_prod_matrix = compute_inner_product(dist_in);

        return get_coefficients(inner_prod_matrix);
    }

    template<typename F>
    inline View<double**, CLayout, Device> get_moments(const int nnode, F lambda_func) const{
        const View<double**, CLayout, Device> g_mom(NoInit("g_mom"), n_poly, nnode);

        Kokkos::parallel_for("get_moments", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
            lambda_func(inode, g_mom);
        });
        Kokkos::fence();
        return g_mom;
    }

    /* Computes the scaling polynomials with the precomputed basis
     * polynomials and the new moments (n,v_||,mu,v_||^2)
     *
     * Ansatz: g = sum_[i=1]^4 d_i lambda_i f
     * ==> d_i = int(lambda_i g d^3v) / <lambda_i,lambda_i>
     *         = sum_[j=1]^4 c_[i,j] int(gamma_j g d^3v) / <lambda_i,lambda_i>
     * After some manipulation --->
     * g = f sum_[k=1]^4 D_k gamma_k with
     * D_k = sum_[i,j=1]^4 int(gamma_j g d^3v * c_[i,j]*c_[i,k] / <lambda_i,lambda_i>
     *  @param[in] polynom  Polynomial representation of the distribution
     *  @param[in] g_mom    Array of new moments
     *  return   The scaling
     */
    View<double***, CLayout, Device> get_scaling(const PolynomialBasisDistribution<Device>& polynom, const View<double**, CLayout, Device>& g_mom) const{
        int nnode = polynom.f.extent(2);
        View<double***, CLayout, Device> scaling("scaling", nnode, xrange.n, yrange.n);
        View<double**, CLayout, Device> D_k("D_k", n_poly, nnode);

        const PolynomialBasis pb = *this;
        Kokkos::parallel_for("get_scaling", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
            // 1) Compute the polynomial coefficients D_k
            for (int i = 0; i < pb.n_poly; i++) {
                for (int j = 0; j < pb.n_poly; j++) {
                    for (int k = 0; k < pb.n_poly; k++) {
                        double c2 = polynom.f(j,i,inode)*polynom.f(k,i,inode);
                        D_k(k,inode) += g_mom(j,inode)*c2;
                    }
                }
            }

            // 2) Compute the new distribution function using the final
            // scaling polynomial P = sum_[k=1]^4 D_k gamma_k
            for (int iy = 0; iy < pb.yrange.n; iy++) {
                for (int ix = 0; ix < pb.xrange.n; ix++) {
                    double fac = -1.0;
                    for(int ip=0; ip<pb.n_poly; ip++){
                        fac += D_k(ip,inode)*pb.bases(ip,ix,iy);
                    }
                    scaling(inode, ix, iy) = fac;
                }
            }
        });
        Kokkos::fence();
        return scaling;
    }
};

#endif