MeanField_SERiF/src/include/serif/discretization/backend/mfem/basis_evaluation.hpp

#pragma once

/**
 * @file basis_evaluation.hpp
 * @author Emily M. Boudreaux
 * @brief Basis evaluation, pointwise mapping, and transpose application using dense or tensor-product representations.
 */
#include <memory>

#include "mfem/fem/fe/fe_base.hpp"
#include "mfem/fem/intrules.hpp"

#include "serif/discretization/backend/mfem/tensor_basis.hpp"

namespace serif::discretization::backend::mfem {
    /**
     * @brief Reusable scratch storage for basis contractions, reducing repeated allocations.
     */
    struct BasisWorkspace {
        TensorWorkspace tensor;
        ::mfem::Vector reference_values;
    };

    /**
     * @class BasisEvaluation
     * @brief Represents an element-local basis evaluation operator, supporting dense or tensor contractions, pointwise mappings, and transpose application.
     *
     * Consider a scalar field represented on one finite element:
     *
     * \f[
     * \rho_h(\xi) = \sum_{j=1}^{n} c_j \phi_j(\xi).
     * \f]
     *
     * Here, \f$c_{j}\f$ is a scalar coefficient, \f$\phi_{j}\f$ is its associated basis function,
     * and \f$\xi\f$ is the canonical reference coordinate. The field is expressed in
     * these coordinates through the element mapping.
     *
     * For a quadrature rule with \f$N_{q}\f$ points, define the evaluation matrix \f$\mathbf{E}\f$ by
     *
     * \f[
     * E_{qj} = \phi_j(\xi_q).
     * \f]
     *
     * \f$\mathbf{E}\f$ therefore has \f$N_{q}\f$ rows and \f$n\f$ columns. The field value at one point is
     *
     * \f[
     * \rho_h(\xi_q) = \sum_{j=1}^{n} E_{qj} c_j.
     * \f]
     *
     * Collecting the samples into a vector gives
     *
     * \f[
     * \boldsymbol{\rho}_h = \mathbf{E}\mathbf{c}.
     * \f]
     *
     * This class represents that linear operation without necessarily storing
     * the full matrix \f$\mathbf{E}\f$: it can instead apply equivalent tensor contractions.
     *
     * @par Example
     * Evaluate a constant scalar field on the reference cube, then compute
     * its integral against each basis function.
     *
     * @code{.cpp}
     * namespace backend = serif::discretization::backend::mfem;
     *
     * // First-order nodal H1 basis on the reference cube.
     * ::mfem::H1_HexahedronElement element(1);
     * const ::mfem::IntegrationRule& rule =
     *     ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 2);
     *
     * backend::TensorBasisCache cache;
     * backend::BasisWorkspace workspace;
     * backend::BasisEvaluation basis;
     *
     * // Prepare E: one scalar basis value per quadrature point and DOF.
     * basis.sample(
     *     element,
     *     rule,
     *     backend::ReferenceOperation::value,
     *     1, // One copy of the scalar basis.
     *     cache,
     *     backend::ContractionMode::automatic
     * );
     *
     * // For this nodal basis, equal coefficients represent a constant field.
     * ::mfem::Vector coefficients(element.GetDof());
     * coefficients = 2.0;
     *
     * // Compute values = E * coefficients.
     * // Every quadrature-point value is approximately 2.0.
     * ::mfem::Vector values;
     * basis.evaluate(coefficients, values, workspace);
     *
     * // Supply quadrature weights explicitly: integrate does not add them.
     * ::mfem::Vector weighted_values(rule.GetNPoints());
     * for (int q = 0; q < rule.GetNPoints(); ++q) {
     *     weighted_values[q] = rule.IntPoint(q).weight * values[q];
     * }
     *
     * // Compute element_integrals = E^T * weighted_values.
     * ::mfem::Vector element_integrals;
     * basis.integrate(weighted_values, element_integrals, workspace);
     * @endcode
     *
     * Each output entry approximates
     * \f[
     *     \int_{\hat K} \phi_j(\xi)\rho_h(\xi)\,d\xi.
     * \f]
     * These are element-DOF integral contributions, not reconstructed field
     * coefficients. Integration over a physical element additionally requires
     * the appropriate volume Jacobian in the quadrature weights.
     */
    class BasisEvaluation final {
    public:
        /**
         * @brief Prepare basis values or derivatives at quadrature points using a dense matrix or tensor-product representation.
         * @param element Finite element to sample.
         * @param rule Integration rule to use for sampling.
         * @param operation Reference operation to perform.
         * @param copies Number of independent copies of the finite-element basis; for example, three scalar basis copies for a vector H1 field.
         * @param cache Tensor basis cache to use.
         * @param mode Contraction mode to use.
         */
        void sample(
            const ::mfem::FiniteElement& element,
            const ::mfem::IntegrationRule& rule,
            ReferenceOperation operation,
            int copies,
            TensorBasisCache& cache,
            ContractionMode mode
        );

        /**
         * @brief Set the basis evaluation to a constant basis with the given number of points.
         * @param points Number of quadrature points to use for the constant basis.
         */
        void set_constant(
            int points
        );

        /**
         * @brief Allocate storage for a mapped basis using the reference's coefficient ordering and quadrature points.
         * @param reference Basis evaluation supplying the dimensions and any shared tensor representation.
         * @param components Number of output components per quadrature point.
         *
         * @note Populate the destination at every quadrature point using map_point before evaluation or transpose application.
         */
        void reset_from(
            const BasisEvaluation& reference,
            int components
        );

        /**
         * @brief Apply a pointwise linear map to the reference basis or its geometry variation.
         * @tparam MapT Type of the mapping function to apply.
         * @param point Quadrature point index to apply the mapping at.
         * @param reference Reference basis evaluation to map from.
         * @param map Mapping function to apply.
         *
         * @pre The mapping callback must be linear in its input and fill every output component.
         */
        template <typename MapT>
        void map_point(
            int point,
            const BasisEvaluation& reference,
            MapT&& map
        );

        /**
         * @brief Apply the basis to native coefficients; no integration weight is introduced.
         * @param coefficients Input coefficients to apply the basis to.
         * @param values Output values after applying the basis.
         * @param workspace Workspace to use for intermediate computations.
         */
        void evaluate(
            const ::mfem::Vector& coefficients,
            ::mfem::Vector& values,
            BasisWorkspace& workspace
        ) const;

        /**
         * @brief Apply the coordinate transpose to already weighted quadrature duals.
         * @param[in] values Quadrature data with any required integration weights already applied.
         * @param[out] coefficients Element-DOF contributions computed as E^T times values.
         * @param workspace Workspace to use for intermediate computations.
         */
        void integrate(
            const ::mfem::Vector& values,
            ::mfem::Vector& coefficients,
            BasisWorkspace& workspace
        ) const;

        /// Number of flattened quadrature samples.
        [[nodiscard]] int height() const noexcept;

        /// Number of native element coefficients.
        [[nodiscard]] int width() const noexcept;

        /// Report whether contractions use verified Cartesian factors.
        [[nodiscard]] bool uses_tensor() const noexcept;

    private:
        int m_points{0}; ///< Number of quadrature points in the rule
        int m_components{0}; ///< Number of output components per point
        int m_width{0}; ///< Number of native element coefficients
        std::shared_ptr<const TensorBasis> m_tensor; ///< Canonical tensor representation, when the tensor path is selected
        ::mfem::DenseMatrix m_dense; ///< Full evaluation matrix, including applied mappings, when the dense path is selected
        ::mfem::DenseMatrix m_mapping; ///< Stacked pointwise maps from canonical tensor samples to current output components
    };
} // namespace serif::discretization::backend::mfem

template <typename MapT>
void serif::discretization::backend::mfem::BasisEvaluation::map_point(const int point, const BasisEvaluation& reference, MapT&& map) {
    const ::mfem::DenseMatrix& source = reference.m_tensor ? reference.m_mapping : reference.m_dense;

    ::mfem::DenseMatrix& destination = m_tensor ? m_mapping : m_dense;
    ::mfem::Vector input(reference.m_components);
    ::mfem::Vector output(m_components);

    for (int column = 0; column < source.Width(); ++column) {
        for (int axis = 0; axis < reference.m_components; ++axis) {
            input[axis] = source(point * reference.m_components + axis, column);
        }

        map(input, output);

        for (int axis = 0; axis < m_components; ++axis) {
            destination(point * m_components + axis, column) = output[axis];
        }
    }
}