MeanField/experiments/polytrope_physical_state.hpp

#pragma once

// Include after `import mean_field;`. This observer reconstructs the accepted
// coefficients without mutating the prepared production operator or mesh.
#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <limits>
#include <map>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

namespace experiment::polytrope_validation {
    struct PhysicalPoint final {
        mean_field::mapping::MappingPointContext mapping;
        int element{-1};
        int attribute{-1};
        bool stellarMaterial{false};
        double rho{std::numeric_limits<double>::quiet_NaN()};
        double h{std::numeric_limits<double>::quiet_NaN()};
        double phi{std::numeric_limits<double>::quiet_NaN()};
        // Positive centrifugal potential Psi = |Omega cross (x-center)|^2 / 2.
        // The pointwise Bernoulli balance is h + phi - Psi - C = 0.
        double rotationPotential{std::numeric_limits<double>::quiet_NaN()};
        mfem::Vector gravityGradientPhysical;
        mfem::Vector enthalpyGradientPhysical;
        mfem::Vector potentialGradientPhysical;

        PhysicalPoint()
            : gravityGradientPhysical(3), enthalpyGradientPhysical(3), potentialGradientPhysical(3) {
            const double nan = std::numeric_limits<double>::quiet_NaN();
            gravityGradientPhysical = nan;
            enthalpyGradientPhysical = nan;
            potentialGradientPhysical = nan;
        }
    };

    struct FieldReconstructionReport final {
        std::string field;
        int reducedSize{0};
        int fullTrueSize{0};
        double maximumRoundTripError{0.0};
    };

    // Nonmovable because the mapping evaluator references the owned displacement
    // grid function. The FEM, spaces, compactification field, and mapper remain
    // borrowed: keep this object inside the diagnostic context callback.
    class PhysicalState final {
    public:
        using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

        template <typename State>
        explicit PhysicalState(State& state, const mean_field::fem::FEM& fem)
            : finiteElements(fem),
              density(RequireSpace(fem.densityFes)),
              enthalpy(RequireSpace(fem.enthalpyFes)),
              potential(RequireSpace(fem.gravityPotentialFes)),
              gravityGradientReference(RequireSpace(fem.gravityFluxFes)),
              displacement(RequireSpace(fem.displacementFes)),
              domainMapper(state.Problem().GetDiscretization().domainMapper()),
              mapping(domainMapper, displacement, RequireCompactification(fem)),
              rotation(state.Problem().GetPreparedOperator().GetRotation()) {
            int ranks = 0;
            MPI_Comm_size(fem.mesh->GetComm(), &ranks);
            if (ranks != 1) {
                throw std::invalid_argument("Physical verification currently requires exactly one MPI rank.");
            }
            const auto& problem = state.Problem();
            const auto& accepted = state.AcceptedPhysicalState();
            const auto& prepared = state.normalizedOperator->GetPhysicalState();
            if (accepted.Size() != prepared.Size()) {
                throw std::logic_error("Accepted and prepared physical state sizes differ.");
            }
            for (int index = 0; index < accepted.Size(); ++index) {
                if (!std::isfinite(accepted(index)) || accepted(index) != prepared(index)) {
                    throw std::logic_error("Physical verification requires the accepted state to be prepared.");
                }
            }
            const auto values = problem.GetManifest().valueBlocks();
            ScatterField<mean_field::field::Density>("density", values, accepted, density);
            ScatterField<mean_field::field::Enthalpy>("specific_enthalpy", values, accepted, enthalpy);
            ScatterField<mean_field::field::Gravity>("gravity_potential", values, accepted, potential);
            ScatterField<mean_field::field::Gravity>("gravity_gradient", values, accepted, gravityGradientReference);
            displacement.SetFromTrueDofs(problem.GetPhysicalOperator().GetGeneratedVolumeDisplacement());
            mapping.InvalidateCache();

            bernoulliConstant = Scalar(values, accepted, "fixed_total_mass.multiplier");
            physicalResidualBordered = state.normalizedOperator->GetPhysicalResidual();
            physicalResidualUnbordered = physicalResidualBordered;
            normalizedBorderedResidual = state.acceptedNormalizedResidual;
            normalizedUnborderedResidual.SetSize(physicalResidualUnbordered.Size());

            if constexpr (requires { problem.GetPreparedOperator().GetCentralDensityConstraint(); }) {
                centralBorder = Scalar(values, accepted, "fixed_central_density.border");
                const auto& central = problem.GetPreparedOperator().GetCentralDensityConstraint();
                centralDensityReport = central.GetConstraintReport();
                const auto hydrostatic = RequireBlock(problem.GetManifest().residualBlocks(), "hydrostatic_balance");
                if (central.GetCenterDof().field_size() != hydrostatic.size) {
                    throw std::logic_error("Central-density border and hydrostatic row layouts disagree.");
                }
                for (const int reducedDof : central.GetCenterDof().reduced_dofs()) {
                    if (reducedDof < 0 || reducedDof >= hydrostatic.size) {
                        throw std::logic_error("Central-density border index is outside the hydrostatic block.");
                    }
                    const int row = hydrostatic.offset + reducedDof;
                    centralHydrostaticRows.push_back(row);
                    physicalResidualUnbordered(row) -= centralBorder;
                }
            }
            state.normalizedOperator->NormalizeResidual(physicalResidualUnbordered, normalizedUnborderedResidual);
            normalizedBorderedResidualNorm = normalizedBorderedResidual.Norml2();
            normalizedUnborderedResidualNorm = normalizedUnborderedResidual.Norml2();
            mfem::Vector normalizedBorder(normalizedBorderedResidual);
            normalizedBorder -= normalizedUnborderedResidual;
            normalizedCentralBorderActionNorm = normalizedBorder.Norml2();
        }

        // Experiment-only postprocessing replay, not a solver restart. The
        // caller must construct fem from savedDirectory/input.smesh. Historical
        // residuals are not recomputed: only the saved finite-element fields
        // are loaded, using exactly the same Evaluate implementation below.
        explicit PhysicalState(
            const mean_field::fem::FEM& fem, const std::filesystem::path& savedDirectory
        )
            : finiteElements(fem),
              density(RequireSpace(fem.densityFes)),
              enthalpy(RequireSpace(fem.enthalpyFes)),
              potential(RequireSpace(fem.gravityPotentialFes)),
              gravityGradientReference(RequireSpace(fem.gravityFluxFes)),
              displacement(RequireSpace(fem.displacementFes)),
              domainMapper(RequireDomainMapper(fem)),
              mapping(domainMapper, displacement, RequireCompactification(fem)),
              rotation(RequireNonrotatingReplay(fem, savedDirectory)) {
            LoadField(savedDirectory / "density.gf", density);
            LoadField(savedDirectory / "enthalpy.gf", enthalpy);
            LoadField(savedDirectory / "potential.gf", potential);
            LoadField(savedDirectory / "gravity_gradient_reference.gf", gravityGradientReference);
            LoadField(savedDirectory / "displacement.gf", displacement);
            mapping.InvalidateCache();
            centralBorder = std::numeric_limits<double>::quiet_NaN();
            normalizedCentralBorderActionNorm = std::numeric_limits<double>::quiet_NaN();
        }

        PhysicalState(const PhysicalState&) = delete;
        PhysicalState& operator=(const PhysicalState&) = delete;
        PhysicalState(PhysicalState&&) = delete;
        PhysicalState& operator=(PhysicalState&&) = delete;

        [[nodiscard]] bool isStellar(int element) const {
            CheckElement(element);
            return DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Stellar>(
                finiteElements.mesh->GetAttribute(element)
            );
        }

        [[nodiscard]] bool isVacuum(int element) const {
            CheckElement(element);
            return DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(
                finiteElements.mesh->GetAttribute(element)
            );
        }

        // Grid-function vector values include the MFEM/reference-mesh Piola
        // transform already. Applying J_map/det(J_map) here completes, rather
        // than repeats, the transformation to the deformed physical geometry.
        [[nodiscard]] mean_field::mapping::MappingStatus Evaluate(
            int element, const mfem::IntegrationPoint& point, PhysicalPoint& output
        ) {
            using mean_field::mapping::MappingStatus;
            CheckElement(element);
            output.element = element;
            output.attribute = finiteElements.mesh->GetAttribute(element);
            output.stellarMaterial = isStellar(element);
            const double nan = std::numeric_limits<double>::quiet_NaN();
            output.rho = output.h = output.phi = output.rotationPotential = nan;
            output.gravityGradientPhysical = nan;
            output.enthalpyGradientPhysical = nan;
            output.potentialGradientPhysical = nan;
            auto* transformation = finiteElements.mesh->GetElementTransformation(element);
            const auto status = mapping.EvaluatePoint(*transformation, point, output.mapping);
            if (status != MappingStatus::valid) return status;

            // GetValue/GetGradient evaluate the FE functions on the reference
            // physical mesh. Scalar values pull back unchanged under DomainMapper.
            output.phi = potential.GetValue(element, point);
            transformation->SetIntPoint(&point);
            potential.GetGradient(*transformation, m_referenceGradient);
            output.mapping.inverse_mapping_jacobian.MultTranspose(m_referenceGradient, output.potentialGradientPhysical);
            gravityGradientReference.GetVectorValue(element, point, m_referenceGravity);
            output.mapping.mapping_jacobian.Mult(m_referenceGravity, output.gravityGradientPhysical);
            output.gravityGradientPhysical /= output.mapping.mapping_determinant;
            output.rotationPotential = rotation.potential(output.mapping.physical_position);
            if (output.stellarMaterial) {
                output.rho = density.GetValue(element, point);
                output.h = enthalpy.GetValue(element, point);
                transformation->SetIntPoint(&point);
                enthalpy.GetGradient(*transformation, m_referenceGradient);
                output.mapping.inverse_mapping_jacobian.MultTranspose(m_referenceGradient, output.enthalpyGradientPhysical);
            }
            // rho/h outside their material support intentionally remain NaN;
            // zeroed unsupported FE coefficients are not physical vacuum data.
            if (!std::isfinite(output.phi) || !std::isfinite(output.rotationPotential) ||
                !AllFinite(output.gravityGradientPhysical) || !AllFinite(output.potentialGradientPhysical) ||
                (output.stellarMaterial && (!std::isfinite(output.rho) || !std::isfinite(output.h) ||
                                           !AllFinite(output.enthalpyGradientPhysical)))) {
                return MappingStatus::non_finite_result;
            }
            return MappingStatus::valid;
        }

        const mean_field::fem::FEM& finiteElements;
        mfem::ParGridFunction density;
        mfem::ParGridFunction enthalpy;
        mfem::ParGridFunction potential;
        mfem::ParGridFunction gravityGradientReference;
        mfem::ParGridFunction displacement;
        const mean_field::mapping::DomainMapper& domainMapper;
        mean_field::mapping::GridFunctionMappingEvaluator mapping;
        mean_field::physics::RigidRotation rotation;
        double bernoulliConstant{std::numeric_limits<double>::quiet_NaN()};
        double centralBorder{0.0};
        std::optional<mean_field::operators::CentralDensityConstraintReport> centralDensityReport;
        std::vector<int> centralHydrostaticRows;
        std::vector<FieldReconstructionReport> reconstructionReports;

        // "Unbordered" removes only the artificial lambda_c*e_c contribution
        // from hydrostatic balance. All scalar constraint rows and the physical
        // Bernoulli constant remain present in the original root layout.
        mfem::Vector physicalResidualBordered;
        mfem::Vector physicalResidualUnbordered;
        mfem::Vector normalizedBorderedResidual;
        mfem::Vector normalizedUnborderedResidual;
        double normalizedBorderedResidualNorm{std::numeric_limits<double>::quiet_NaN()};
        double normalizedUnborderedResidualNorm{std::numeric_limits<double>::quiet_NaN()};
        double normalizedCentralBorderActionNorm{0.0};

    private:
        struct Block final { int offset; int size; };
        mfem::Vector m_referenceGradient = mfem::Vector(3);
        mfem::Vector m_referenceGravity = mfem::Vector(3);

        static mfem::ParFiniteElementSpace* RequireSpace(
            const std::unique_ptr<mfem::ParFiniteElementSpace>& space
        ) {
            if (space == nullptr) throw std::invalid_argument("Physical verification received an incomplete FE space.");
            return space.get();
        }

        static const mfem::ParGridFunction& RequireCompactification(const mean_field::fem::FEM& fem) {
            if (fem.mesh == nullptr || fem.compactificationCoordinate == nullptr) {
                throw std::invalid_argument("Physical verification received incomplete mapping data.");
            }
            return *fem.compactificationCoordinate;
        }

        static const mean_field::mapping::DomainMapper& RequireDomainMapper(const mean_field::fem::FEM& fem) {
            if (fem.domainMapperStateless == nullptr) {
                throw std::invalid_argument("Physical replay received an incomplete domain mapper.");
            }
            return *fem.domainMapperStateless;
        }

        static mean_field::physics::RigidRotation RequireNonrotatingReplay(
            const mean_field::fem::FEM& fem, const std::filesystem::path& directory
        ) {
            int ranks = 0;
            MPI_Comm_size(fem.mesh->GetComm(), &ranks);
            if (ranks != 1) throw std::invalid_argument("Physical replay requires exactly one MPI rank.");
            std::ifstream metadata(directory / "metadata.txt");
            if (!metadata) throw std::invalid_argument("Physical replay requires saved metadata.txt.");
            std::map<std::string, std::string> values;
            std::string line;
            while (std::getline(metadata, line)) {
                const auto separator = line.find('=');
                if (separator != std::string::npos) values[line.substr(0, separator)] = line.substr(separator + 1);
            }
            constexpr std::string_view supportedModel =
                "nonrotating_n1_fixed_mass_fixed_central_density_zero_surface_pressure";
            if (metadata.bad() || values["mode"] != "solve" || values["mpi_ranks"] != "1" ||
                values["model"] != supportedModel) {
                throw std::invalid_argument("Physical replay only supports saved single-rank nonrotating n=1 solves.");
            }
            if (!values.contains("elements") || std::stoi(values.at("elements")) != fem.mesh->GetNE()) {
                throw std::invalid_argument("Replay FEM element count disagrees with saved metadata.");
            }
            // The supported model has exactly J=0 and computes Omega=0. If a
            // completed measurement file exists, reject contradictory data;
            // an interrupted postprocess can still replay its already-saved GFs.
            std::ifstream metrics(directory / "physical_metrics.csv");
            if (metrics) {
                while (std::getline(metrics, line)) {
                    constexpr std::string_view prefix = "angular_velocity_norm,";
                    if (!line.starts_with(prefix)) continue;
                    const double angularVelocity = std::stod(line.substr(prefix.size()));
                    if (!std::isfinite(angularVelocity) || angularVelocity != 0.0) {
                        throw std::invalid_argument("Physical replay cannot infer a nonzero saved rotation vector.");
                    }
                }
                if (metrics.bad()) throw std::runtime_error("Could not read saved rotation verification data.");
            }
            mfem::Vector zero(3);
            zero = 0.0;
            return mean_field::physics::RigidRotation{zero, zero};
        }

        static void LoadField(const std::filesystem::path& path, mfem::ParGridFunction& target) {
            std::ifstream stream(path);
            if (!stream) throw std::invalid_argument("Missing saved physical field: " + path.string());
            // ParGridFunction's stream constructor reverses the local-DOF
            // orientation handling in ParGridFunction::Save (important for RT).
            mfem::ParGridFunction loaded(target.ParFESpace()->GetParMesh(), stream);
            if (stream.fail() || !AllFinite(loaded)) {
                throw std::invalid_argument("Incomplete or non-finite saved physical field: " + path.string());
            }
            const auto& savedSpace = *loaded.ParFESpace();
            const auto& targetSpace = *target.ParFESpace();
            if (std::string_view(savedSpace.FEColl()->Name()) != targetSpace.FEColl()->Name() ||
                savedSpace.GetVDim() != targetSpace.GetVDim() || savedSpace.GetOrdering() != targetSpace.GetOrdering() ||
                savedSpace.GetVSize() != targetSpace.GetVSize() || savedSpace.GetTrueVSize() != targetSpace.GetTrueVSize() ||
                loaded.Size() != target.Size()) {
                throw std::invalid_argument("Saved field FE collection/layout is incompatible with this build: " + path.string());
            }
            mfem::Array<int> savedDofs, targetDofs;
            for (int element = 0; element < targetSpace.GetNE(); ++element) {
                if (savedSpace.GetElementOrder(element) != targetSpace.GetElementOrder(element)) {
                    throw std::invalid_argument("Saved field element order is incompatible with this build: " + path.string());
                }
                savedSpace.GetElementVDofs(element, savedDofs);
                targetSpace.GetElementVDofs(element, targetDofs);
                if (savedDofs.Size() != targetDofs.Size()) {
                    throw std::invalid_argument("Saved field element DOF count is incompatible with this build: " + path.string());
                }
                for (int index = 0; index < targetDofs.Size(); ++index) {
                    if (savedDofs[index] != targetDofs[index]) {
                        throw std::invalid_argument("Saved field element DOF ordering is incompatible with this build: " + path.string());
                    }
                }
            }
            target = loaded;
        }

        void CheckElement(int element) const {
            if (element < 0 || element >= finiteElements.mesh->GetNE()) {
                throw std::out_of_range("Physical verification element index is outside the local mesh.");
            }
        }

        template <typename Blocks>
        static Block RequireBlock(const Blocks& blocks, std::string_view name) {
            for (const auto& block : blocks) {
                if (block.stableId == name) return {block.offset, block.size};
            }
            throw std::invalid_argument("Physical verification requires root block " + std::string(name));
        }

        template <typename Blocks>
        static double Scalar(const Blocks& blocks, const mfem::Vector& values, std::string_view name) {
            const auto block = RequireBlock(blocks, name);
            if (block.size != 1 || block.offset < 0 || block.offset >= values.Size()) {
                throw std::logic_error("Physical verification scalar block has an incompatible layout.");
            }
            return values(block.offset);
        }

        template <typename Field, typename Blocks>
        void ScatterField(std::string_view name, const Blocks& blocks, const mfem::Vector& values, mfem::ParGridFunction& field) {
            const auto block = RequireBlock(blocks, name);
            const auto adapter = mean_field::field::make_field_dof_grid_function_adapter<Field, DomainSchema>(*field.ParFESpace());
            if (block.size != adapter.dof_map().reduced_size() || block.offset < 0 ||
                block.offset + block.size > values.Size()) {
                throw std::logic_error("Physical verification field map disagrees with root block " + std::string(name));
            }
            mfem::Vector reduced(block.size);
            for (int index = 0; index < block.size; ++index) reduced(index) = values(block.offset + index);
            adapter.scatter(reduced, field);
            const auto roundTrip = adapter.gather(field);
            double error = 0.0;
            for (int index = 0; index < block.size; ++index) error = std::max(error, std::abs(roundTrip(index) - reduced(index)));
            reconstructionReports.push_back({std::string(name), block.size, adapter.dof_map().full_size(), error});
        }

        static bool AllFinite(const mfem::Vector& vector) {
            for (int index = 0; index < vector.Size(); ++index) if (!std::isfinite(vector(index))) return false;
            return true;
        }
    };
} // namespace experiment::polytrope_validation