MeanField/libmeanfield/impl/operators/prepared_stellar_equilibrium.cpp

module;

#include <array>
#include <cmath>
#include <utility>

#include <mfem.hpp>

module mean_field;

import :operators.prepared_stellar_equilibrium;
import :physics.gravity;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    [[nodiscard]] mean_field::fem::FEM &ensure_gravity_static_operators(mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.mesh != nullptr && f.densityFes != nullptr && f.displacementFes != nullptr &&
                f.gravityFluxFes != nullptr && f.gravityPotentialFes != nullptr && f.enthalpyFes != nullptr,
            "PreparedStellarEquilibriumOperator requires the complete coupled finite-element discretization."
        );

        if (f.gravityContext.b_form == nullptr || f.gravityContext.BT == nullptr) {
            mean_field::physics::update_stiffness_matrix(f);
        }

        MFEM_VERIFY(
            f.gravityContext.b_form != nullptr && f.gravityContext.BT != nullptr,
            "PreparedStellarEquilibriumOperator could not initialize the static gravity divergence operators."
        );

        return f;
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumLayout make_layout(
        const mean_field::field::FieldDofMap &densityMap,
        const mean_field::field::FieldDofMap &displacementMap,
        const mean_field::field::FieldDofMap &gravityFluxMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap,
        const mean_field::field::FieldDofMap &enthalpyMap
    ) {
        using Form = mean_field::utils::blocks::barotropic_equilibrium_form;

        const std::array<int, Form::value_block_count> valueSizes{
            densityMap.reduced_size(),          displacementMap.reduced_size(), gravityFluxMap.reduced_size(),
            gravityPotentialMap.reduced_size(), enthalpyMap.reduced_size(),     1
        };

        const std::array<int, Form::residual_block_count> residualSizes{
            gravityFluxMap.reduced_size(),  gravityPotentialMap.reduced_size(), densityMap.reduced_size(),
            displacementMap.reduced_size(), enthalpyMap.reduced_size(),         1
        };

        return {valueSizes, residualSizes};
    }

    [[nodiscard]] mfem::Array<int> make_gravity_state_offsets(const mean_field::fem::FEM &f) {
        mfem::Array<int> offsets(5);
        offsets[0] = 0;
        offsets[1] = offsets[0] + f.densityFes->GetTrueVSize();
        offsets[2] = offsets[1] + f.displacementFes->GetTrueVSize();
        offsets[3] = offsets[2] + f.gravityFluxFes->GetTrueVSize();
        offsets[4] = offsets[3] + f.gravityPotentialFes->GetTrueVSize();
        return offsets;
    }

    [[nodiscard]] mfem::Array<int> make_gravity_residual_offsets(const mean_field::fem::FEM &f) {
        mfem::Array<int> offsets(3);
        offsets[0] = 0;
        offsets[1] = f.gravityFluxFes->GetTrueVSize();
        offsets[2] = offsets[1] + f.gravityPotentialFes->GetTrueVSize();
        return offsets;
    }

    template <int index>
    [[nodiscard]] mfem::Vector make_value_view(
        const mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::value_block<index> block
    ) {
        MFEM_VERIFY(
            vector.Size() == layout.value_offsets().Last(),
            "The coupled vector does not match the stellar-equilibrium value layout."
        );

        return mfem::Vector(const_cast<mfem::real_t *>(vector.GetData()) + layout.offset(block), layout.size(block));
    }

    template <int index>
    [[nodiscard]] mfem::Vector make_residual_view(
        mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block
    ) {
        MFEM_VERIFY(
            vector.Size() == layout.residual_offsets().Last(),
            "The coupled vector does not match the stellar-equilibrium residual layout."
        );

        return mfem::Vector(vector.GetData() + layout.offset(block), layout.size(block));
    }

    template <int index>
    void assign_residual_block(
        mfem::Vector &coupledResidual,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block,
        const mfem::Vector &blockResidual,
        const char *message
    ) {
        MFEM_VERIFY(layout.size(block) == blockResidual.Size(), message);
        mfem::Vector destination = make_residual_view(coupledResidual, layout, block);
        destination              = blockResidual;
    }

    void assign_gravity_block(
        mfem::Vector &gravityState,
        const mfem::Array<int> &offsets,
        const int blockIndex,
        const mfem::Vector &source,
        const char *message
    ) {
        MFEM_VERIFY(offsets.Size() == 5, "Gravity state offsets are invalid.");
        MFEM_VERIFY(blockIndex >= 0 && blockIndex + 1 < offsets.Size(), "Requested gravity-state block is invalid.");

        const int blockSize = offsets[blockIndex + 1] - offsets[blockIndex];
        MFEM_VERIFY(blockSize == source.Size(), message);
        MFEM_VERIFY(gravityState.Size() == offsets.Last(), "Packed gravity state has the wrong size.");

        mfem::Vector destination(gravityState.GetData() + offsets[blockIndex], blockSize);
        destination = source;
    }

    void pack_gravity_vector(
        mfem::Vector &gravityState,
        const mfem::Array<int> &offsets,
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mfem::Vector &gravityGradient,
        const mfem::Vector &gravityPotential
    ) {
        MFEM_VERIFY(offsets.Size() == 5, "Packed gravity state requires four blocks.");
        if (gravityState.Size() != offsets.Last()) {
            gravityState.SetSize(offsets.Last());
        }

        assign_gravity_block(
            gravityState, offsets, 0, density, "The full density vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 1, displacement, "The displacement vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 2, gravityGradient, "The gravity-gradient vector has the wrong gravity-state size."
        );
        assign_gravity_block(
            gravityState, offsets, 3, gravityPotential, "The gravity-potential vector has the wrong gravity-state size."
        );
    }

    void validate_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        for (int index = 0; index < vector.Size(); ++index) {
            MFEM_VERIFY(std::isfinite(vector(index)), message);
        }
    }

    void validate_dependency_transition(
        const mean_field::operators::StellarEquilibriumDependencyStamp &prepared,
        const mean_field::operators::StellarEquilibriumDependencyStamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(prepared.identity != requested.identity || requested.revision >= prepared.revision, message);
        MFEM_VERIFY(
            prepared.identity == requested.identity || prepared.revision != requested.revision,
            "A new stellar-equilibrium dependency identity must also carry a visibly different revision."
        );
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions
    make_gravity_revisions(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {
            .discretization    = {.value = dependencies.discretization.revision},
            .displacement      = {.value = dependencies.displacement.revision},
            .density           = {.value = dependencies.density.revision},
            .gravity_gradient  = {.value = dependencies.gravityGradient.revision},
            .gravity_potential = {.value = dependencies.gravityPotential.revision}
        };
    }

    [[nodiscard]] mean_field::operators::context::barotropic::BarotropicClosureDependencies
    make_barotropic_closure_dependencies(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density      = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .enthalpy     = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .displacement = {
                .identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision
            }
        };
    }

    [[nodiscard]] mean_field::operators::DisplacementResidualDependencies
    make_displacement_dependencies(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .displacement =
                {.identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision},
            .gravityGradient =
                {.identity = dependencies.gravityGradient.identity, .revision = dependencies.gravityGradient.revision},
            .enthalpy = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .rotation = {.identity = dependencies.rotation.identity, .revision = dependencies.rotation.revision}
        };
    }

    [[nodiscard]] mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies
    make_hydrostatic_dependencies(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .enthalpy = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .gravityPotential =
                {.identity = dependencies.gravityPotential.identity,
                 .revision = dependencies.gravityPotential.revision},
            .displacement =
                {.identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision},
            .rotation = {.identity = dependencies.rotation.identity, .revision = dependencies.rotation.revision},
            .bernoulliConstant = {
                .identity = dependencies.bernoulliConstant.identity, .revision = dependencies.bernoulliConstant.revision
            }
        };
    }

    [[nodiscard]] mean_field::operators::MassNormalizationDependencies
    make_mass_dependencies(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .displacement =
                {.identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision},
            .targetMass = {.identity = dependencies.targetMass.identity, .revision = dependencies.targetMass.revision}
        };
    }
} // namespace

namespace mean_field::operators {
    struct PreparedStellarEquilibriumOperator::ConstructionData {
        field::FieldDofMap densityMap;
        field::FieldDofMap displacementMap;
        field::FieldDofMap gravityFluxMap;
        field::FieldDofMap gravityPotentialMap;
        field::FieldDofMap enthalpyMap;

        StellarEquilibriumLayout layout;
        mfem::Array<int> gravityStateOffsets;
        mfem::Array<int> gravityResidualOffsets;

        explicit ConstructionData(fem::FEM &f)
            : densityMap(
                  field::make_field_dof_map<
                      field::Density,
                      DomainSchema>(*f.densityFes)
              ),
              displacementMap(
                  field::make_field_dof_map<
                      field::Displacement,
                      DomainSchema>(*f.displacementFes)
              ),
              gravityFluxMap(
                  field::make_field_dof_map<
                      field::Gravity,
                      DomainSchema>(*f.gravityFluxFes)
              ),
              gravityPotentialMap(
                  field::make_field_dof_map<
                      field::Gravity,
                      DomainSchema>(*f.gravityPotentialFes)
              ),
              enthalpyMap(
                  field::make_field_dof_map<
                      field::Enthalpy,
                      DomainSchema>(*f.enthalpyFes)
              ),
              layout(make_layout(
                  densityMap,
                  displacementMap,
                  gravityFluxMap,
                  gravityPotentialMap,
                  enthalpyMap
              )),
              gravityStateOffsets(make_gravity_state_offsets(f)),
              gravityResidualOffsets(make_gravity_residual_offsets(f)) {
        }
    };

    PreparedStellarEquilibriumOperator::ConstructionData
    PreparedStellarEquilibriumOperator::MakeConstructionData(fem::FEM &f) {
        ensure_gravity_static_operators(f);
        return ConstructionData(f);
    }

    PreparedStellarEquilibriumOperator::PreparedStellarEquilibriumOperator(
        fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const eos::Polytrope &equationOfState,
        const models::StellarModel &stellarModel
    )
        : PreparedStellarEquilibriumOperator(
              f,
              domainMapper,
              equationOfState,
              stellarModel.targetMass()
          ) {
    }

    PreparedStellarEquilibriumOperator::PreparedStellarEquilibriumOperator(
        fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const eos::Polytrope &equationOfState,
        const double targetMass
    )
        : PreparedStellarEquilibriumOperator(
              f,
              domainMapper,
              equationOfState,
              targetMass,
              MakeConstructionData(f)
          ) {
    }

    PreparedStellarEquilibriumOperator::PreparedStellarEquilibriumOperator(
        fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const eos::Polytrope &equationOfState,
        const double targetMass,
        ConstructionData constructionData
    )
        : mfem::Operator(
              constructionData.layout.residual_offsets().Last(),
              constructionData.layout.value_offsets().Last()
          ),
          m_layout(constructionData.layout),
          m_gravityStateOffsets(constructionData.gravityStateOffsets),
          m_gravityContext(
              f,
              domainMapper
          ),
          m_gravityJacobianOperator(
              f,
              domainMapper,
              m_gravityContext,
              m_gravityStateOffsets,
              constructionData.gravityResidualOffsets
          ),
          m_gravityOperator(
              f,
              domainMapper,
              m_gravityContext,
              m_gravityStateOffsets,
              m_gravityJacobianOperator
          ),
          m_barotropicClosureOperator(
              f,
              domainMapper,
              equationOfState
          ),
          m_hydrostaticOperator(
              f,
              domainMapper
          ),
          m_displacementOperator(
              f,
              domainMapper,
              equationOfState,
              m_gravityContext
          ),
          m_massNormalizationOperator(
              f,
              domainMapper,
              m_gravityContext
          ),
          m_targetMass(targetMass),
          m_densityMap(std::move(constructionData.densityMap)),
          m_displacementMap(std::move(constructionData.displacementMap)),
          m_gravityFluxMap(std::move(constructionData.gravityFluxMap)),
          m_gravityPotentialMap(std::move(constructionData.gravityPotentialMap)),
          m_enthalpyMap(std::move(constructionData.enthalpyMap)) {
        MFEM_VERIFY(
            std::isfinite(m_targetMass) && m_targetMass > 0.0,
            "PreparedStellarEquilibriumOperator requires a finite, positive target mass."
        );

        MFEM_VERIFY(
            Width() == m_layout.value_offsets().Last() && Height() == m_layout.residual_offsets().Last(),
            "PreparedStellarEquilibriumOperator has inconsistent block dimensions."
        );

        MFEM_VERIFY(
            m_displacementMap.is_identity(), "PreparedStellarEquilibriumOperator currently requires Displacement "
                                             "support to span the full MFEM true-DOF space."
        );
        MFEM_VERIFY(
            m_gravityFluxMap.is_identity(), "PreparedStellarEquilibriumOperator currently requires gravity-flux "
                                            "support to span the full MFEM true-DOF space."
        );
        MFEM_VERIFY(
            m_gravityPotentialMap.is_identity(), "PreparedStellarEquilibriumOperator currently requires "
                                                 "gravity-potential support to span the full MFEM true-DOF space."
        );

        m_fullDensity.SetSize(m_densityMap.full_size());
        m_fullEnthalpy.SetSize(m_enthalpyMap.full_size());
        m_fullGravityState.SetSize(m_gravityStateOffsets.Last());

        m_fullDensityVariation.SetSize(m_densityMap.full_size());
        m_fullEnthalpyVariation.SetSize(m_enthalpyMap.full_size());
        m_fullGravityDirection.SetSize(m_gravityStateOffsets.Last());
        m_fullEnthalpyAction.SetSize(m_enthalpyMap.full_size());

        m_fullDensity           = 0.0;
        m_fullEnthalpy          = 0.0;
        m_fullGravityState      = 0.0;
        m_fullDensityVariation  = 0.0;
        m_fullEnthalpyVariation = 0.0;
        m_fullGravityDirection  = 0.0;
        m_fullEnthalpyAction    = 0.0;
    }

    PreparedStellarEquilibriumReport PreparedStellarEquilibriumOperator::Prepare(
        const mfem::Vector &state,
        const StellarEquilibriumDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        MFEM_VERIFY(
            state.Size() == Width(), "PreparedStellarEquilibriumOperator received a state with the wrong size."
        );
        validate_finite_vector(state, "PreparedStellarEquilibriumOperator received a non-finite state.");

        const bool wasPrepared = m_isPrepared;
        if (wasPrepared) {
            validate_dependency_transition(
                m_preparedDependencies.discretization, dependencies.discretization,
                "The discretization revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.density, dependencies.density, "The density revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.displacement, dependencies.displacement,
                "The displacement revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.gravityGradient, dependencies.gravityGradient,
                "The gravity-gradient revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.gravityPotential, dependencies.gravityPotential,
                "The gravity-potential revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.enthalpy, dependencies.enthalpy, "The enthalpy revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.bernoulliConstant, dependencies.bernoulliConstant,
                "The Bernoulli-constant revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.rotation, dependencies.rotation, "The rotation revision cannot move backwards."
            );
            validate_dependency_transition(
                m_preparedDependencies.targetMass, dependencies.targetMass,
                "The target-mass revision cannot move backwards."
            );
        }

        m_isPrepared                = false;

        using Form                  = utils::blocks::barotropic_equilibrium_form;
        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);
        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravityPotentialValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto enthalpyValue =
            utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto bernoulliValue =
            utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        const mfem::Vector reducedDensity   = make_value_view(state, m_layout, densityValue);
        const mfem::Vector displacement     = make_value_view(state, m_layout, displacementValue);
        const mfem::Vector gravityGradient  = make_value_view(state, m_layout, gravityGradientValue);
        const mfem::Vector gravityPotential = make_value_view(state, m_layout, gravityPotentialValue);
        const mfem::Vector reducedEnthalpy  = make_value_view(state, m_layout, enthalpyValue);
        const mfem::Vector bernoulli        = make_value_view(state, m_layout, bernoulliValue);

        m_densityMap.scatter(reducedDensity, m_fullDensity);
        m_enthalpyMap.scatter(reducedEnthalpy, m_fullEnthalpy);

        pack_gravity_vector(
            m_fullGravityState, m_gravityStateOffsets, m_fullDensity, displacement, gravityGradient, gravityPotential
        );

        PreparedStellarEquilibriumReport report;

        report.gravity           = m_gravityOperator.Prepare(m_fullGravityState, make_gravity_revisions(dependencies));

        report.barotropicClosure = m_barotropicClosureOperator.Prepare(
            {.density = reducedDensity, .enthalpy = reducedEnthalpy, .displacement = displacement},
            make_barotropic_closure_dependencies(dependencies)
        );

        report.hydrostatic = m_hydrostaticOperator.Prepare(
            {.enthalpy          = m_fullEnthalpy,
             .gravityPotential  = gravityPotential,
             .displacement      = displacement,
             .bernoulliConstant = bernoulli(0)},
            make_hydrostatic_dependencies(dependencies), rotation
        );

        report.displacement = m_displacementOperator.Prepare(
            {.enthalpy = reducedEnthalpy}, make_displacement_dependencies(dependencies), rotation
        );

        report.massNormalization =
            m_massNormalizationOperator.Prepare({.targetMass = m_targetMass}, make_mass_dependencies(dependencies));

        const bool dependenciesChanged = !wasPrepared || dependencies != m_preparedDependencies;
        if (dependenciesChanged || report.DidAnyChildWork()) {
            AssembleResidual();
            report.assembledResidual = true;
        }

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;
        return report;
    }

    void PreparedStellarEquilibriumOperator::AssembleResidual() {
        using Form = utils::blocks::barotropic_equilibrium_form;

        constexpr auto gravityGradientResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravityPotentialResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto densityResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);
        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto enthalpyResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        mfem::Vector gravity;
        mfem::Vector closure;
        mfem::Vector displacement;
        mfem::Vector mass;

        m_gravityOperator.Mult(m_fullGravityState, gravity);
        m_barotropicClosureOperator.BuildResidual(closure);
        m_displacementOperator.BuildResidual(displacement);
        m_hydrostaticOperator.BuildResidual(m_fullEnthalpyAction);
        m_massNormalizationOperator.BuildResidual(mass);

        m_cachedResidual.SetSize(Height());
        m_cachedResidual = 0.0;

        MFEM_VERIFY(
            gravity.Size() == m_layout.size(gravityGradientResidual) + m_layout.size(gravityPotentialResidual),
            "The gravity residual has the wrong size."
        );

        mfem::Vector gravityGradient(gravity.GetData(), m_layout.size(gravityGradientResidual));
        mfem::Vector gravityPotential(
            gravity.GetData() + m_layout.size(gravityGradientResidual), m_layout.size(gravityPotentialResidual)
        );

        assign_residual_block(
            m_cachedResidual, m_layout, gravityGradientResidual, gravityGradient,
            "The gravity-gradient residual has the wrong size."
        );
        assign_residual_block(
            m_cachedResidual, m_layout, gravityPotentialResidual, gravityPotential,
            "The gravity-potential residual has the wrong size."
        );
        assign_residual_block(
            m_cachedResidual, m_layout, densityResidual, closure, "The closure residual has the wrong size."
        );
        assign_residual_block(
            m_cachedResidual, m_layout, displacementResidual, displacement,
            "The displacement residual has the wrong size."
        );

        {
            mfem::Vector reducedEnthalpyResidual = make_residual_view(m_cachedResidual, m_layout, enthalpyResidual);
            m_enthalpyMap.gather(m_fullEnthalpyAction, reducedEnthalpyResidual);
        }

        assign_residual_block(
            m_cachedResidual, m_layout, massResidual, mass, "The mass-normalization residual has the wrong size."
        );

        ++m_statistics.residualAssemblies;
    }

    void PreparedStellarEquilibriumOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();
        residual = m_cachedResidual;
        ++m_statistics.residualApplications;
    }

    void PreparedStellarEquilibriumOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            direction.Size() == Width(),
            "PreparedStellarEquilibriumOperator received a Jacobian direction with the wrong size."
        );
        validate_finite_vector(
            direction, "PreparedStellarEquilibriumOperator received a non-finite Jacobian direction."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);
        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravityPotentialValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto enthalpyValue =
            utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto bernoulliValue =
            utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        constexpr auto gravityGradientResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravityPotentialResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto densityResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);
        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto enthalpyResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);
        constexpr auto massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        const mfem::Vector reducedDensityDirection   = make_value_view(direction, m_layout, densityValue);
        const mfem::Vector displacementDirection     = make_value_view(direction, m_layout, displacementValue);
        const mfem::Vector gravityGradientDirection  = make_value_view(direction, m_layout, gravityGradientValue);
        const mfem::Vector gravityPotentialDirection = make_value_view(direction, m_layout, gravityPotentialValue);
        const mfem::Vector reducedEnthalpyDirection  = make_value_view(direction, m_layout, enthalpyValue);
        const mfem::Vector bernoulliDirection        = make_value_view(direction, m_layout, bernoulliValue);

        m_densityMap.scatter(reducedDensityDirection, m_fullDensityVariation);
        m_enthalpyMap.scatter(reducedEnthalpyDirection, m_fullEnthalpyVariation);

        pack_gravity_vector(
            m_fullGravityDirection, m_gravityStateOffsets, m_fullDensityVariation, displacementDirection,
            gravityGradientDirection, gravityPotentialDirection
        );

        mfem::Vector gravityAction;
        mfem::Vector closureAction;
        mfem::Vector displacementAction;
        mfem::Vector massAction;

        m_gravityJacobianOperator.Mult(m_fullGravityDirection, gravityAction);

        m_barotropicClosureOperator.Mult(
            reducedDensityDirection, reducedEnthalpyDirection, displacementDirection, closureAction
        );

        m_displacementOperator.ApplyCompleteJacobianAction(
            m_fullDensityVariation, displacementDirection, gravityGradientDirection, reducedEnthalpyDirection,
            displacementAction
        );

        m_hydrostaticOperator.ApplyCompleteJacobianAction(
            m_fullEnthalpyVariation, gravityPotentialDirection, bernoulliDirection(0), displacementDirection,
            m_fullEnthalpyAction
        );

        m_massNormalizationOperator.ApplyCompleteJacobianAction(
            m_fullDensityVariation, displacementDirection, massAction
        );

        action.SetSize(Height());
        action = 0.0;

        MFEM_VERIFY(
            gravityAction.Size() == m_layout.size(gravityGradientResidual) + m_layout.size(gravityPotentialResidual),
            "The gravity Jacobian action has the wrong size."
        );

        mfem::Vector gravityGradientAction(gravityAction.GetData(), m_layout.size(gravityGradientResidual));
        mfem::Vector gravityPotentialAction(
            gravityAction.GetData() + m_layout.size(gravityGradientResidual), m_layout.size(gravityPotentialResidual)
        );

        assign_residual_block(
            action, m_layout, gravityGradientResidual, gravityGradientAction,
            "The gravity-gradient Jacobian action has the wrong size."
        );
        assign_residual_block(
            action, m_layout, gravityPotentialResidual, gravityPotentialAction,
            "The gravity-potential Jacobian action has the wrong size."
        );
        assign_residual_block(
            action, m_layout, densityResidual, closureAction, "The closure Jacobian action has the wrong size."
        );
        assign_residual_block(
            action, m_layout, displacementResidual, displacementAction,
            "The displacement Jacobian action has the wrong size."
        );

        {
            mfem::Vector reducedEnthalpyAction = make_residual_view(action, m_layout, enthalpyResidual);
            m_enthalpyMap.gather(m_fullEnthalpyAction, reducedEnthalpyAction);
        }

        assign_residual_block(
            action, m_layout, massResidual, massAction, "The mass-normalization Jacobian action has the wrong size."
        );

        ++m_statistics.jacobianApplications;
    }

    bool PreparedStellarEquilibriumOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_gravityContext.IsPrepared() && m_barotropicClosureOperator.IsPrepared() &&
               m_hydrostaticOperator.IsPrepared() && m_displacementOperator.IsPrepared() &&
               m_massNormalizationOperator.IsPrepared();
    }

    double PreparedStellarEquilibriumOperator::GetTargetMass() const noexcept {
        return m_targetMass;
    }

    const StellarEquilibriumLayout &PreparedStellarEquilibriumOperator::GetLayout() const noexcept {
        return m_layout;
    }

    const StellarEquilibriumDependencies &PreparedStellarEquilibriumOperator::GetDependencies() const {
        VerifyPrepared();
        return m_preparedDependencies;
    }

    const PreparedStellarEquilibriumStatistics &PreparedStellarEquilibriumOperator::GetStatistics() const noexcept {
        return m_statistics;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    PreparedStellarEquilibriumOperator::GetGravityContext() const noexcept {
        return m_gravityContext;
    }

    const GravityFieldOperator &PreparedStellarEquilibriumOperator::GetGravityOperator() const noexcept {
        return m_gravityOperator;
    }

    const GravityFieldJacobianOperator &
    PreparedStellarEquilibriumOperator::GetGravityJacobianOperator() const noexcept {
        return m_gravityJacobianOperator;
    }

    const PreparedBarotropicClosureOperator &
    PreparedStellarEquilibriumOperator::GetBarotropicClosureOperator() const noexcept {
        return m_barotropicClosureOperator;
    }

    const context::barotropic::BarotropicClosureLinearizationContext &
    PreparedStellarEquilibriumOperator::GetBarotropicClosureContext() const noexcept {
        return m_barotropicClosureOperator.GetContext();
    }

    const PreparedHydrostaticEquilibriumOperator &
    PreparedStellarEquilibriumOperator::GetHydrostaticOperator() const noexcept {
        return m_hydrostaticOperator;
    }

    const PreparedDisplacementResidualOperator &
    PreparedStellarEquilibriumOperator::GetDisplacementOperator() const noexcept {
        return m_displacementOperator;
    }

    const PreparedMassNormalizationOperator &
    PreparedStellarEquilibriumOperator::GetMassNormalizationOperator() const noexcept {
        return m_massNormalizationOperator;
    }

    void PreparedStellarEquilibriumOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedStellarEquilibriumOperator must be prepared before residual or Jacobian application."
        );
    }
} // namespace mean_field::operators