adr_operator.h

Go to the documentation of this file.

#ifndef HYBRIDADRSOLVER_ADR_OPERATOR_H
#define HYBRIDADRSOLVER_ADR_OPERATOR_H
 
#include "core/problem_definition.h"
 
#include <deal.II/base/vectorization.h>
#include <deal.II/lac/la_parallel_vector.h>
#include <deal.II/matrix_free/fe_evaluation.h>
#include <deal.II/matrix_free/matrix_free.h>
#include <deal.II/matrix_free/operators.h>
#include <deal.II/multigrid/mg_constrained_dofs.h>
 
namespace HybridADRSolver {
using namespace dealii;

template <int dim, int fe_degree, typename Number = double>
class ADROperator : public MatrixFreeOperators::Base<
                        dim, LinearAlgebra::distributed::Vector<Number>> {
public:
    using VectorType = LinearAlgebra::distributed::Vector<Number>;
    using value_type = Number;
    using size_type = types::global_dof_index;
    using Base =
        MatrixFreeOperators::Base<dim,
                                  LinearAlgebra::distributed::Vector<Number>>;

    ADROperator() : problem_ptr(nullptr), is_mg_level(false), mg_level(-1) {}

    void initialize(std::shared_ptr<const MatrixFree<dim, Number>> data_in,
                    const ProblemInterface<dim>& problem) {
        Base::initialize(data_in);
        this->problem_ptr = &problem;
        this->is_mg_level = false;
        this->mg_level = -1;
        precompute_coefficient_data();
    }

    void initialize(std::shared_ptr<const MatrixFree<dim, Number>> data_in,
                    const MGConstrainedDoFs& mg_constrained_dofs,
                    unsigned int level, const ProblemInterface<dim>& problem) {
        Base::initialize(data_in, mg_constrained_dofs, level);
        this->problem_ptr = &problem;
        this->is_mg_level = true;
        this->mg_level = level;
        precompute_coefficient_data();
    }

    void clear() override {
        Base::clear();
        problem_ptr = nullptr;
        diffusion_coefficients.clear();
        advection_coefficients.clear();
        reaction_coefficients.clear();
        is_mg_level = false;
        mg_level = -1;
    }

    void compute_diagonal() override {
        this->inverse_diagonal_entries.reset(new DiagonalMatrix<VectorType>());
        VectorType& inverse_diagonal =
            this->inverse_diagonal_entries->get_vector();
        this->data->initialize_dof_vector(inverse_diagonal);
 
        MatrixFreeTools::compute_diagonal(*this->data, inverse_diagonal,
                                          &ADROperator::local_compute_diagonal,
                                          this);
 
        this->set_constrained_entries_to_one(inverse_diagonal);
 
        for (unsigned int i = 0; i < inverse_diagonal.locally_owned_size();
             ++i) {
            if (std::abs(inverse_diagonal.local_element(i)) > 1e-14)
                inverse_diagonal.local_element(i) =
                    1.0 / inverse_diagonal.local_element(i);
            else
                inverse_diagonal.local_element(i) = 1.0;
        }
    }

    [[nodiscard]] bool is_multigrid_level() const { return is_mg_level; }

    [[nodiscard]] int get_mg_level() const { return mg_level; }
 
private:
    void apply_add(VectorType& dst, const VectorType& src) const override {
        this->data->cell_loop(&ADROperator::local_apply, this, dst, src);
    }

    void precompute_coefficient_data() {
        Assert(problem_ptr != nullptr, ExcNotInitialized());
 
        const unsigned int n_cells = this->data->n_cell_batches();
        const unsigned int n_q_points = Utilities::pow(fe_degree + 1, dim);
 
        diffusion_coefficients.resize(n_cells);
        advection_coefficients.resize(n_cells);
        reaction_coefficients.resize(n_cells);
 
        // Use FEEvaluation to get quadrature points
        FEEvaluation<dim, fe_degree, fe_degree + 1, 1, Number> phi(*this->data);
 
        for (unsigned int cell = 0; cell < n_cells; ++cell) {
            phi.reinit(cell);
 
            diffusion_coefficients[cell].resize(n_q_points);
            advection_coefficients[cell].resize(n_q_points);
            reaction_coefficients[cell].resize(n_q_points);
 
            for (unsigned int q = 0; q < n_q_points; ++q) {
                // Get the vectorized quadrature point
                Point<dim, VectorizedArray<Number>> q_point =
                    phi.quadrature_point(q);
 
                // Evaluate coefficients for each lane in the vector
                VectorizedArray<Number> mu;
                Tensor<1, dim, VectorizedArray<Number>> beta;
                VectorizedArray<Number> gamma;
 
                for (unsigned int v = 0; v < VectorizedArray<Number>::size();
                     ++v) {
                    Point<dim> p;
                    for (unsigned int d = 0; d < dim; ++d)
                        p[d] = q_point[d][v];
 
                    mu[v] = problem_ptr->diffusion_coefficient(p);
                    gamma[v] = problem_ptr->reaction_coefficient(p);
 
                    Tensor<1, dim> beta_scalar =
                        problem_ptr->advection_field(p);
                    for (unsigned int d = 0; d < dim; ++d)
                        beta[d][v] = beta_scalar[d];
                }
 
                diffusion_coefficients[cell][q] = mu;
                advection_coefficients[cell][q] = beta;
                reaction_coefficients[cell][q] = gamma;
            }
        }
    }

    void
    local_apply(const MatrixFree<dim, Number>& /*mf_data*/, VectorType& dst,
                const VectorType& src,
                const std::pair<unsigned int, unsigned int>& cell_range) const {
        FEEvaluation<dim, fe_degree, fe_degree + 1, 1, Number> phi(*this->data);
 
        for (unsigned int cell = cell_range.first; cell < cell_range.second;
             ++cell) {
            phi.reinit(cell);
            phi.gather_evaluate(src, EvaluationFlags::values |
                                         EvaluationFlags::gradients);
 
            for (unsigned int q = 0; q < phi.n_q_points; ++q) {
                // use precomputed coefficients
                const VectorizedArray<Number>& mu_val =
                    diffusion_coefficients[cell][q];
                const Tensor<1, dim, VectorizedArray<Number>>& beta_val =
                    advection_coefficients[cell][q];
                const VectorizedArray<Number>& gamma_val =
                    reaction_coefficients[cell][q];
 
                const auto u_val = phi.get_value(q);
                const auto grad_u = phi.get_gradient(q);
 
                // Diffusion: mu * grad(u)
                const auto flux = mu_val * grad_u;
 
                // Advection + Reaction
                // weak form: (beta . grad_u, v) + (gamma * u, v)
                const auto val = gamma_val * u_val + beta_val * grad_u;
 
                phi.submit_gradient(flux, q);
                phi.submit_value(val, q);
            }
 
            phi.integrate_scatter(
                EvaluationFlags::values | EvaluationFlags::gradients, dst);
        }
    }

    void local_compute_diagonal(
        FEEvaluation<dim, fe_degree, fe_degree + 1, 1, Number>& phi) const {
        const unsigned int cell = phi.get_current_cell_index();
 
        phi.evaluate(EvaluationFlags::values | EvaluationFlags::gradients);
 
        for (unsigned int q = 0; q < phi.n_q_points; ++q) {
            const VectorizedArray<Number> u_val = phi.get_value(q);
            const Tensor<1, dim, VectorizedArray<Number>> grad_u =
                phi.get_gradient(q);
 
            // Use precomputed coefficients
            const Tensor<1, dim, VectorizedArray<Number>> flux =
                diffusion_coefficients[cell][q] * grad_u;
 
            const VectorizedArray<Number> advection_val =
                advection_coefficients[cell][q] * grad_u;
 
            const VectorizedArray<Number> val =
                reaction_coefficients[cell][q] * u_val + advection_val;
 
            phi.submit_gradient(flux, q);
            phi.submit_value(val, q);
        }
 
        phi.integrate(EvaluationFlags::values | EvaluationFlags::gradients);
    }
 
    // Data members
    const ProblemInterface<dim>* problem_ptr;
    bool is_mg_level;
    unsigned int mg_level;
 
    // Precomputed coefficients at quadrature points
    // These are computed ONCE during initialization and reused for all
    std::vector<std::vector<VectorizedArray<Number>>> diffusion_coefficients;
    std::vector<std::vector<Tensor<1, dim, VectorizedArray<Number>>>>
        advection_coefficients;
    std::vector<std::vector<VectorizedArray<Number>>> reaction_coefficients;
};

template <int dim, int fe_degree, typename Number = double>
class JacobiPreconditioner {
public:
    using VectorType = LinearAlgebra::distributed::Vector<Number>;

    void initialize(const ADROperator<dim, fe_degree, Number>& op) {
        op.compute_diagonal();
        inverse_diagonal = op.get_matrix_diagonal_inverse()->get_vector();
    }

    void vmult(VectorType& dst, const VectorType& src) const {
        dst = src;
        dst.scale(inverse_diagonal);
    }
 
private:
    VectorType inverse_diagonal;
};
 
} // namespace HybridADRSolver
 
#endif // HYBRIDADRSOLVER_ADR_OPERATOR_H