MeanField/tests/operators/prepared_hydrostatic_equilibrium_complete_jacobian.cpp

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_complete_test_utils {
    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 503, .revision = 2},
            .enthalpy          = {.identity = 509, .revision = 3},
            .gravityPotential  = {.identity = 521, .revision = 5},
            .displacement      = {.identity = 523, .revision = 7},
            .rotation          = {.identity = 541, .revision = 11},
            .bernoulliConstant = {.identity = 547, .revision = 13}
        };
    }

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumStateView make_state(
        const mfem::Vector &enthalpy,
        const mfem::Vector &gravityPotential,
        const mfem::Vector &displacement,
        const double bernoulliConstant
    ) {
        return {
            .enthalpy          = enthalpy,
            .gravityPotential  = gravityPotential,
            .displacement      = displacement,
            .bernoulliConstant = bernoulliConstant
        };
    }

    mean_field::physics::RigidRotation make_rotation() {
        mfem::Vector angularVelocity(3);

        angularVelocity(0) = 0.18;
        angularVelocity(1) = -0.13;
        angularVelocity(2) = 0.49;

        mfem::Vector center(3);

        center(0) = 0.031;
        center(1) = -0.024;
        center(2) = 0.017;

        return mean_field::physics::RigidRotation(angularVelocity, center);
    }

    mfem::Vector make_displacement_direction(
        const mean_field::fem::FEM &f,
        const double firstPhase,
        const double secondPhase
    ) {
        mfem::Vector direction         = field_dof_test_utils::make_supported_displacement(f, firstPhase);

        const mfem::Vector secondField = field_dof_test_utils::make_supported_displacement(f, secondPhase);

        direction -= secondField;
        return direction;
    }

    void centered_complete_difference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpy,
        const mfem::Vector &baseGravityPotential,
        const mfem::Vector &baseDisplacement,
        const double baseBernoulliConstant,
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &gravityPotentialVariation,
        const mfem::Vector &displacementVariation,
        const double bernoulliConstantVariation,
        const double step,
        mfem::Vector &difference
    ) {
        mfem::Vector enthalpyPlus(baseEnthalpy);
        mfem::Vector enthalpyMinus(baseEnthalpy);
        mfem::Vector gravityPotentialPlus(baseGravityPotential);
        mfem::Vector gravityPotentialMinus(baseGravityPotential);
        mfem::Vector displacementPlus(baseDisplacement);
        mfem::Vector displacementMinus(baseDisplacement);

        enthalpyPlus.Add(step, enthalpyVariation);
        enthalpyMinus.Add(-step, enthalpyVariation);

        gravityPotentialPlus.Add(step, gravityPotentialVariation);

        gravityPotentialMinus.Add(-step, gravityPotentialVariation);

        displacementPlus.Add(step, displacementVariation);
        displacementMinus.Add(-step, displacementVariation);

        const double bernoulliConstantPlus  = baseBernoulliConstant + step * bernoulliConstantVariation;

        const double bernoulliConstantMinus = baseBernoulliConstant - step * bernoulliConstantVariation;

        mfem::Vector residualPlus;
        mfem::Vector residualMinus;

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpyPlus, gravityPotentialPlus, displacementPlus, bernoulliConstantPlus, residualPlus
        );

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpyMinus, gravityPotentialMinus, displacementMinus, bernoulliConstantMinus, residualMinus
        );

        difference = residualPlus;
        difference -= residualMinus;
        difference /= 2.0 * step;
    }

    void set_block(
        mfem::Vector &packedDirection,
        const mean_field::operators::HydrostaticJacobianBlockLayout &layout,
        const mean_field::operators::HydrostaticJacobianInputBlock block,
        const mfem::Vector &values
    ) {
        REQUIRE(layout.Size(block) == values.Size());

        const int offset = layout.Offset(block);

        for (int entry = 0; entry < values.Size(); ++entry) {
            packedDirection(offset + entry) = values(entry);
        }
    }
} // namespace prepared_hydrostatic_complete_test_utils

TEST_CASE(
    "Prepared Hydrostatic Complete Jacobian Matches Sum And Centered "
    "Differences",
    tags::barotrope_hydrostatic_prepared_jacobian &tags::self_consistency
) {
    auto args              = test_utils::setup_args();

    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    mean_field::operators::PreparedHydrostaticEquilibriumOperator preparedOperator(f, *f.domainMapperStateless);

    const mfem::Vector enthalpy =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 0.34);

    const mfem::Vector gravityPotential =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 0.57
        );

    const mfem::Vector displacement                   = field_dof_test_utils::make_supported_displacement(f, 0.68);

    constexpr double bernoulliConstant                = 0.43;

    const mean_field::physics::RigidRotation rotation = prepared_hydrostatic_complete_test_utils::make_rotation();

    preparedOperator.Prepare(
        prepared_hydrostatic_complete_test_utils::make_state(
            enthalpy, gravityPotential, displacement, bernoulliConstant
        ),
        prepared_hydrostatic_complete_test_utils::make_dependencies(), rotation
    );

    const mfem::Vector enthalpyVariation =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 1.07);

    const mfem::Vector gravityPotentialVariation =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 1.31
        );

    const mfem::Vector displacementVariation =
        prepared_hydrostatic_complete_test_utils::make_displacement_direction(f, 1.19, 0.38);

    constexpr double bernoulliConstantVariation = -0.37;

    mfem::Vector enthalpyAction;
    mfem::Vector gravityPotentialAction;
    mfem::Vector bernoulliConstantAction;
    mfem::Vector displacementAction;
    mfem::Vector completeAction;

    preparedOperator.ApplyEnthalpyJacobianAction(enthalpyVariation, enthalpyAction);

    preparedOperator.ApplyGravityPotentialJacobianAction(gravityPotentialVariation, gravityPotentialAction);

    preparedOperator.ApplyBernoulliConstantJacobianAction(bernoulliConstantVariation, bernoulliConstantAction);

    preparedOperator.ApplyDisplacementJacobianAction(displacementVariation, displacementAction);

    preparedOperator.ApplyCompleteJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, displacementVariation, completeAction
    );

    mfem::Vector summedAction(enthalpyAction);
    summedAction += gravityPotentialAction;
    summedAction += bernoulliConstantAction;
    summedAction += displacementAction;

    constexpr double finiteDifferenceStep = 1.0e-5;

    mfem::Vector centeredDifference;

    prepared_hydrostatic_complete_test_utils::centered_complete_difference(
        f, rotation, enthalpy, gravityPotential, displacement, bernoulliConstant, enthalpyVariation,
        gravityPotentialVariation, displacementVariation, bernoulliConstantVariation, finiteDifferenceStep,
        centeredDifference
    );

    const double summationError =
        gravity_prepared_test_utils::relative_error(completeAction, summedAction, f.mesh->GetComm());

    const double centeredDifferenceError =
        gravity_prepared_test_utils::relative_error(completeAction, centeredDifference, f.mesh->GetComm());

    INFO("Complete-action summation error = " << summationError);

    INFO("Complete-action centered-difference error = " << centeredDifferenceError);

    CHECK(preparedOperator.GetCompleteJacobianStatistics().applications == 1);

    CHECK(summationError < 5.0e-12);
    CHECK(centeredDifferenceError < 2.0e-7);
}

TEST_CASE(
    "Prepared Hydrostatic MFEM Adapter Uses Four Block Layout And Reuses "
    "Preparation",
    tags::barotrope_hydrostatic_prepared_jacobian &tags::mfem_operators &tags::unit
) {
    auto args              = test_utils::setup_args();

    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    mean_field::operators::PreparedHydrostaticEquilibriumOperator preparedOperator(f, *f.domainMapperStateless);

    const mfem::Vector enthalpy =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 0.41);

    const mfem::Vector gravityPotential =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 0.63
        );

    const mfem::Vector displacement    = field_dof_test_utils::make_supported_displacement(f, 0.74);

    constexpr double bernoulliConstant = 0.38;

    preparedOperator.Prepare(
        prepared_hydrostatic_complete_test_utils::make_state(
            enthalpy, gravityPotential, displacement, bernoulliConstant
        ),
        prepared_hydrostatic_complete_test_utils::make_dependencies(),
        prepared_hydrostatic_complete_test_utils::make_rotation()
    );

    mean_field::operators::PreparedHydrostaticEquilibriumJacobianOperator adapter(f, preparedOperator);

    const auto &layout = adapter.GetLayout();

    CHECK(layout.Offset(mean_field::operators::HydrostaticJacobianInputBlock::enthalpy) == 0);

    CHECK(
        layout.Offset(mean_field::operators::HydrostaticJacobianInputBlock::gravityPotential) ==
        preparedOperator.GetEnthalpyMap().reduced_size()
    );

    CHECK(layout.Size(mean_field::operators::HydrostaticJacobianInputBlock::bernoulliConstant) == 1);

    CHECK(adapter.Width() == layout.GetTotalSize());
    CHECK(adapter.Height() == layout.GetResidualSize());
    CHECK(adapter.Height() == preparedOperator.GetEnthalpyMap().reduced_size());

    CHECK(preparedOperator.GetEnthalpyMap().inactive_size() > 0);

    const mfem::Vector enthalpyVariation =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 1.12);

    const mfem::Vector gravityPotentialVariation =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 1.39
        );

    const mfem::Vector displacementVariation =
        prepared_hydrostatic_complete_test_utils::make_displacement_direction(f, 1.28, 0.49);

    constexpr double bernoulliConstantVariation = 0.29;

    mfem::Vector packedDirection(adapter.Width());
    packedDirection = 0.0;

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout, mean_field::operators::HydrostaticJacobianInputBlock::enthalpy, enthalpyVariation
    );

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout, mean_field::operators::HydrostaticJacobianInputBlock::gravityPotential,
        gravityPotentialVariation
    );

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout, mean_field::operators::HydrostaticJacobianInputBlock::displacement,
        displacementVariation
    );

    packedDirection(layout.Offset(mean_field::operators::HydrostaticJacobianInputBlock::bernoulliConstant)) =
        bernoulliConstantVariation;

    mfem::Vector directAction;
    mfem::Vector adapterAction;

    preparedOperator.ApplyCompleteJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, displacementVariation, directAction
    );

    adapter.Mult(packedDirection, adapterAction);

    const double adapterError =
        gravity_prepared_test_utils::relative_error(adapterAction, directAction, f.mesh->GetComm());

    const auto contextStatisticsBefore      = preparedOperator.GetContextPreparationStatistics();

    const auto algebraicStatisticsBefore    = preparedOperator.GetAlgebraicJacobianStatistics();

    const auto displacementStatisticsBefore = preparedOperator.GetDisplacementJacobianStatistics();

    mfem::Vector secondPackedDirection(packedDirection);
    secondPackedDirection *= -0.61;

    mfem::Vector secondAdapterAction;
    adapter.Mult(secondPackedDirection, secondAdapterAction);

    mfem::Vector expectedSecondAction(adapterAction);
    expectedSecondAction *= -0.61;

    const double adapterLinearityError =
        gravity_prepared_test_utils::relative_error(secondAdapterAction, expectedSecondAction, f.mesh->GetComm());

    const auto &contextStatisticsAfter      = preparedOperator.GetContextPreparationStatistics();

    const auto &algebraicStatisticsAfter    = preparedOperator.GetAlgebraicJacobianStatistics();

    const auto &displacementStatisticsAfter = preparedOperator.GetDisplacementJacobianStatistics();

    INFO("MFEM adapter/direct-action error = " << adapterError);

    INFO("MFEM adapter linearity error = " << adapterLinearityError);

    CHECK(adapterError < 5.0e-12);
    CHECK(adapterLinearityError < 5.0e-12);
    CHECK(secondAdapterAction.Size() == adapter.Height());

    CHECK(contextStatisticsAfter == contextStatisticsBefore);

    CHECK(algebraicStatisticsAfter.preparations == algebraicStatisticsBefore.preparations);

    CHECK(displacementStatisticsAfter.preparations == displacementStatisticsBefore.preparations);

    CHECK(preparedOperator.GetCompleteJacobianStatistics().applications == 3);
}