feat(libmeanfield): variadic refactor

also added normaliztion operator
This commit is contained in:
2026-09-06 10:15:00 -04:00
parent 71423d543f
commit 76818f2f82
63 changed files with 28794 additions and 1119 deletions

View File

@@ -0,0 +1,621 @@
module;
#include <array>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <span>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
#include <mfem.hpp>
export module mean_field:normalization.operators;
export import :normalization.physical_riesz;
export namespace mean_field::normalization {
class DiagonalNormalization final {
public:
DiagonalNormalization(
mfem::Vector stateToNormalized,
mfem::Vector residualToNormalized
)
: m_stateToNormalized(std::move(stateToNormalized)),
m_residualToNormalized(std::move(residualToNormalized)) {
ValidateFactors(m_stateToNormalized, "state");
ValidateFactors(m_residualToNormalized, "residual");
}
[[nodiscard]] static DiagonalNormalization Identity(
const int stateSize,
const int residualSize
) {
if (stateSize < 0 || residualSize < 0) {
throw std::invalid_argument("Normalization dimensions cannot be negative.");
}
mfem::Vector state(stateSize);
mfem::Vector residual(residualSize);
state = 1.0;
residual = 1.0;
return {std::move(state), std::move(residual)};
}
[[nodiscard]] int StateSize() const noexcept {
return m_stateToNormalized.Size();
}
[[nodiscard]] int ResidualSize() const noexcept {
return m_residualToNormalized.Size();
}
[[nodiscard]] const mfem::Vector &StateFactors() const noexcept {
return m_stateToNormalized;
}
[[nodiscard]] const mfem::Vector &ResidualFactors() const noexcept {
return m_residualToNormalized;
}
void NormalizeState(
const mfem::Vector &physical,
mfem::Vector &normalized
) const {
Apply(m_stateToNormalized, physical, normalized, false, "state");
}
void DenormalizeState(
const mfem::Vector &normalized,
mfem::Vector &physical
) const {
Apply(m_stateToNormalized, normalized, physical, true, "state");
}
void NormalizeResidual(
const mfem::Vector &physical,
mfem::Vector &normalized
) const {
Apply(m_residualToNormalized, physical, normalized, false, "residual");
}
void DenormalizeResidual(
const mfem::Vector &normalized,
mfem::Vector &physical
) const {
Apply(m_residualToNormalized, normalized, physical, true, "residual");
}
[[nodiscard]] double LocalStateNormSquared(const mfem::Vector &physical) const {
return LocalNormSquared(m_stateToNormalized, physical, "state");
}
[[nodiscard]] double LocalResidualNormSquared(const mfem::Vector &physical) const {
return LocalNormSquared(m_residualToNormalized, physical, "residual");
}
private:
static void ValidateFactors(
const mfem::Vector &factors,
const char *role
) {
for (int index = 0; index < factors.Size(); ++index) {
if (!std::isfinite(factors(index)) || factors(index) <= 0.0) {
throw std::invalid_argument(
std::string("The ") + role + " normalization factors must be finite and positive."
);
}
}
}
static void Apply(
const mfem::Vector &factors,
const mfem::Vector &input,
mfem::Vector &output,
const bool inverse,
const char *role
) {
if (input.Size() != factors.Size()) {
throw std::invalid_argument(std::string("The ") + role + " vector has the wrong size.");
}
const bool exactAlias = input.GetData() == output.GetData() && input.Size() == output.Size();
if (!exactAlias) {
output.SetSize(input.Size());
}
for (int index = 0; index < input.Size(); ++index) {
const double value = input(index);
output(index) = inverse ? value / factors(index) : factors(index) * value;
}
}
[[nodiscard]] static double LocalNormSquared(
const mfem::Vector &factors,
const mfem::Vector &physical,
const char *role
) {
if (physical.Size() != factors.Size()) {
throw std::invalid_argument(std::string("The ") + role + " vector has the wrong size.");
}
double normSquared = 0.0;
for (int index = 0; index < physical.Size(); ++index) {
const double normalized = factors(index) * physical(index);
normSquared += normalized * normalized;
}
return normSquared;
}
mfem::Vector m_stateToNormalized;
mfem::Vector m_residualToNormalized;
};
/* Detection-safe public operation for an ordinary third-party runtime
* policy. The exact policy is recovered from the problem type and must own
* every method in its compiled plan. Its implementation remains beside
* the policy and is found by ADL, so adding a normalization family does
* not edit a library registry or switch. */
template <typename Problem>
concept RuntimePreparedNormalizationOperation =
requires(const std::remove_cvref_t<Problem> &problem) {
typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType;
typename std::remove_cvref_t<Problem>::FormType;
requires RuntimePreparedNormalizationFor<
typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
typename std::remove_cvref_t<Problem>::FormType>;
{
problem.GetNormalizationPrescription()
} -> std::same_as<const typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType &>;
{
prepareStellarNormalization(
problem.GetNormalizationPrescription(),
problem)
} -> std::same_as<DiagonalNormalization>;
};
template <typename Form>
requires utils::blocks::block_form_is_valid_v<Form>
class DiagonalNormalizationBuilder final {
public:
explicit DiagonalNormalizationBuilder(const utils::blocks::form_layout<Form> &layout)
: m_layout(&layout),
m_stateFactors(layout.value_offsets().Last()),
m_residualFactors(layout.residual_offsets().Last()) {
}
explicit DiagonalNormalizationBuilder(
utils::blocks::form_layout<Form> &&
) = delete;
explicit DiagonalNormalizationBuilder(
const utils::blocks::form_layout<Form> &&
) = delete;
template <typename Block>
requires utils::blocks::contains_type_v<Block, typename Form::value_blocks>
void SetValueBlock(
const double physicalScale,
const mfem::Vector &primalGramDiagonal
) {
constexpr int block = utils::blocks::type_index_v<Block, typename Form::value_blocks>;
RequireUnassigned(m_valueAssigned[block], "value");
AssignBlock(
m_stateFactors,
m_layout->value_offsets()[block],
m_layout->value_offsets()[block + 1] - m_layout->value_offsets()[block],
physicalScale,
primalGramDiagonal,
false
);
m_valueAssigned[block] = true;
}
template <typename Block>
requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
void SetResidualBlock(
const double physicalScale,
const mfem::Vector &primalGramDiagonal
) {
constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
RequireUnassigned(m_residualAssigned[block], "residual");
AssignBlock(
m_residualFactors,
m_layout->residual_offsets()[block],
m_layout->residual_offsets()[block + 1] - m_layout->residual_offsets()[block],
physicalScale,
primalGramDiagonal,
true
);
m_residualAssigned[block] = true;
}
template <typename Block>
requires utils::blocks::contains_type_v<Block, typename Form::value_blocks>
void SetValueGlobal(const double physicalScale) {
constexpr int block = utils::blocks::type_index_v<Block, typename Form::value_blocks>;
SetConstantMetricValueBlock<Block>(physicalScale, BlockSize(m_layout->value_offsets(), block));
}
template <typename Block>
requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
void SetResidualGlobal(const double physicalScale) {
constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
mfem::Vector metric(BlockSize(m_layout->residual_offsets(), block));
metric = 1.0;
SetResidualBlock<Block>(physicalScale, metric);
}
template <typename Block>
requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
void SetHybridResidualBlock(
const double physicalScale,
const mfem::Vector &bulkPrimalGramDiagonal,
const std::span<const int> pointRows,
const double pointMetric = 1.0
) {
constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
const int size = BlockSize(m_layout->residual_offsets(), block);
if (bulkPrimalGramDiagonal.Size() != size) {
throw std::invalid_argument("The hybrid residual Gram diagonal has the wrong size.");
}
ValidateMetric(pointMetric);
std::vector<bool> isPointRow(static_cast<std::size_t>(size), false);
for (const int row : pointRows) {
if (row < 0 || row >= size) {
throw std::out_of_range("A hybrid point row lies outside its residual block.");
}
if (isPointRow[static_cast<std::size_t>(row)]) {
throw std::invalid_argument("A hybrid point row was supplied more than once.");
}
isPointRow[static_cast<std::size_t>(row)] = true;
}
mfem::Vector metric(size);
for (int row = 0; row < size; ++row) {
metric(row) = isPointRow[static_cast<std::size_t>(row)]
? pointMetric
: bulkPrimalGramDiagonal(row);
}
SetResidualBlock<Block>(physicalScale, metric);
}
[[nodiscard]] DiagonalNormalization Build() && {
for (const bool assigned : m_valueAssigned) {
if (!assigned) {
throw std::logic_error("The normalization is missing a value block.");
}
}
for (const bool assigned : m_residualAssigned) {
if (!assigned) {
throw std::logic_error("The normalization is missing a residual block.");
}
}
return {std::move(m_stateFactors), std::move(m_residualFactors)};
}
private:
template <typename Block>
void SetConstantMetricValueBlock(
const double physicalScale,
const int size
) {
mfem::Vector metric(size);
metric = 1.0;
SetValueBlock<Block>(physicalScale, metric);
}
[[nodiscard]] static int BlockSize(
const mfem::Array<int> &offsets,
const int block
) noexcept {
return offsets[block + 1] - offsets[block];
}
static void RequireUnassigned(
const bool assigned,
const char *role
) {
if (assigned) {
throw std::logic_error(std::string("The ") + role + " block normalization was assigned twice.");
}
}
static void ValidateMetric(const double metric) {
if (!std::isfinite(metric) || metric <= 0.0) {
throw std::invalid_argument("Every Riesz Gram diagonal entry must be finite and positive.");
}
}
static void AssignBlock(
mfem::Vector &factors,
const int offset,
const int size,
const double physicalScale,
const mfem::Vector &primalGramDiagonal,
const bool dual
) {
if (!std::isfinite(physicalScale) || physicalScale <= 0.0) {
throw std::invalid_argument("A physical normalization scale must be finite and positive.");
}
if (primalGramDiagonal.Size() != size) {
throw std::invalid_argument("A Riesz Gram diagonal has the wrong block size.");
}
for (int index = 0; index < size; ++index) {
const double metric = primalGramDiagonal(index);
ValidateMetric(metric);
const double rieszFactor = std::sqrt(metric);
const double factor = dual
? 1.0 / (physicalScale * rieszFactor)
: rieszFactor / physicalScale;
if (!std::isfinite(factor) || factor <= 0.0) {
throw std::overflow_error("A normalization factor is not finite and positive.");
}
factors(offset + index) = factor;
}
}
const utils::blocks::form_layout<Form> *m_layout;
mfem::Vector m_stateFactors;
mfem::Vector m_residualFactors;
std::array<bool, Form::value_block_count> m_valueAssigned{};
std::array<bool, Form::residual_block_count> m_residualAssigned{};
};
class ScaledJacobianOperator final : public mfem::Operator {
public:
ScaledJacobianOperator(
const mfem::Operator &physicalJacobian,
const DiagonalNormalization &normalization
)
: mfem::Operator(normalization.ResidualSize(), normalization.StateSize()),
m_physicalJacobian(&physicalJacobian),
m_normalization(&normalization),
m_physicalDirection(normalization.StateSize()),
m_physicalAction(normalization.ResidualSize()) {
if (physicalJacobian.Width() != normalization.StateSize() ||
physicalJacobian.Height() != normalization.ResidualSize()) {
throw std::invalid_argument("The physical Jacobian and normalization dimensions do not agree.");
}
}
ScaledJacobianOperator(
mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledJacobianOperator(
const mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledJacobianOperator(
const mfem::Operator &,
DiagonalNormalization &&
) = delete;
ScaledJacobianOperator(
const mfem::Operator &,
const DiagonalNormalization &&
) = delete;
void Mult(
const mfem::Vector &normalizedDirection,
mfem::Vector &normalizedAction
) const override {
m_normalization->DenormalizeState(normalizedDirection, m_physicalDirection);
m_physicalJacobian->Mult(m_physicalDirection, m_physicalAction);
m_normalization->NormalizeResidual(m_physicalAction, normalizedAction);
}
private:
const mfem::Operator *m_physicalJacobian;
const DiagonalNormalization *m_normalization;
mutable mfem::Vector m_physicalDirection;
mutable mfem::Vector m_physicalAction;
};
class ScaledInverseOperator final : public mfem::Operator {
public:
ScaledInverseOperator(
const mfem::Operator &physicalInverse,
const DiagonalNormalization &normalization
)
: mfem::Operator(normalization.StateSize(), normalization.ResidualSize()),
m_physicalInverse(&physicalInverse),
m_normalization(&normalization),
m_physicalResidual(normalization.ResidualSize()),
m_physicalCorrection(normalization.StateSize()) {
if (physicalInverse.Width() != normalization.ResidualSize() ||
physicalInverse.Height() != normalization.StateSize()) {
throw std::invalid_argument("The physical inverse and normalization dimensions do not agree.");
}
}
ScaledInverseOperator(
mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledInverseOperator(
const mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledInverseOperator(
const mfem::Operator &,
DiagonalNormalization &&
) = delete;
ScaledInverseOperator(
const mfem::Operator &,
const DiagonalNormalization &&
) = delete;
void Mult(
const mfem::Vector &normalizedResidual,
mfem::Vector &normalizedCorrection
) const override {
m_normalization->DenormalizeResidual(normalizedResidual, m_physicalResidual);
m_physicalInverse->Mult(m_physicalResidual, m_physicalCorrection);
m_normalization->NormalizeState(m_physicalCorrection, normalizedCorrection);
}
private:
const mfem::Operator *m_physicalInverse;
const DiagonalNormalization *m_normalization;
mutable mfem::Vector m_physicalResidual;
mutable mfem::Vector m_physicalCorrection;
};
struct ScaledPreconditionerStatistics final {
std::uint64_t operatorBindings{0};
std::uint64_t applications{0};
};
/*
* Solver-compatible realization of R^{-1} M^{-1} L^{-1}. The wrapped
* inverse always sees the dimensional Jacobian, even when an MFEM Krylov
* solver binds this object to the normalized Jacobian L J R.
*/
class ScaledPreconditioner final : public mfem::Solver {
public:
ScaledPreconditioner(
mfem::Solver &physicalInverse,
const mfem::Operator &physicalJacobian,
const mfem::Operator &normalizedJacobian,
const DiagonalNormalization &normalization
)
: mfem::Solver(
normalization.StateSize(),
normalization.ResidualSize(),
physicalInverse.iterative_mode
),
m_physicalInverse(&physicalInverse),
m_physicalJacobian(&physicalJacobian),
m_expectedNormalizedJacobian(&normalizedJacobian),
m_normalization(&normalization),
m_physicalResidual(normalization.ResidualSize()),
m_physicalCorrection(normalization.StateSize()) {
if (physicalInverse.Width() != normalization.ResidualSize() ||
physicalInverse.Height() != normalization.StateSize() ||
physicalJacobian.Width() != normalization.StateSize() ||
physicalJacobian.Height() != normalization.ResidualSize()) {
throw std::invalid_argument(
"The physical preconditioner, Jacobian, and normalization dimensions do not agree."
);
}
SetOperator(normalizedJacobian);
}
ScaledPreconditioner(
mfem::Solver &,
mfem::Operator &&,
const mfem::Operator &,
const DiagonalNormalization &
) = delete;
ScaledPreconditioner(
mfem::Solver &,
const mfem::Operator &&,
const mfem::Operator &,
const DiagonalNormalization &
) = delete;
ScaledPreconditioner(
mfem::Solver &,
const mfem::Operator &,
mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledPreconditioner(
mfem::Solver &,
const mfem::Operator &,
const mfem::Operator &&,
const DiagonalNormalization &
) = delete;
ScaledPreconditioner(
mfem::Solver &,
const mfem::Operator &,
const mfem::Operator &,
DiagonalNormalization &&
) = delete;
ScaledPreconditioner(
mfem::Solver &,
const mfem::Operator &,
const mfem::Operator &,
const DiagonalNormalization &&
) = delete;
ScaledPreconditioner(const ScaledPreconditioner &) = delete;
ScaledPreconditioner &operator=(const ScaledPreconditioner &) = delete;
ScaledPreconditioner(ScaledPreconditioner &&) = delete;
ScaledPreconditioner &operator=(ScaledPreconditioner &&) = delete;
void SetOperator(const mfem::Operator &normalizedJacobian) override {
if (normalizedJacobian.Width() != Width() || normalizedJacobian.Height() != Height()) {
throw std::invalid_argument(
"The scaled preconditioner received an incompatible normalized Jacobian."
);
}
if (&normalizedJacobian != m_expectedNormalizedJacobian) {
throw std::invalid_argument(
"The scaled preconditioner cannot be rebound to a different normalized Jacobian."
);
}
m_physicalInverse->SetOperator(*m_physicalJacobian);
m_normalizedJacobian = &normalizedJacobian;
++m_statistics.operatorBindings;
}
void Mult(
const mfem::Vector &normalizedResidual,
mfem::Vector &normalizedCorrection
) const override {
if (m_normalizedJacobian == nullptr) {
throw std::logic_error("The scaled preconditioner has not been bound to a normalized Jacobian.");
}
if (normalizedResidual.Size() != Width() || normalizedCorrection.Size() != Height()) {
throw std::invalid_argument(
"The scaled preconditioner requires compatible, preallocated normalized vectors."
);
}
m_normalization->DenormalizeResidual(normalizedResidual, m_physicalResidual);
m_physicalInverse->Mult(m_physicalResidual, m_physicalCorrection);
m_normalization->NormalizeState(m_physicalCorrection, normalizedCorrection);
++m_statistics.applications;
}
[[nodiscard]] const mfem::Solver &GetPhysicalInverse() const noexcept {
return *m_physicalInverse;
}
[[nodiscard]] const mfem::Operator &GetPhysicalJacobian() const noexcept {
return *m_physicalJacobian;
}
[[nodiscard]] const mfem::Operator &GetNormalizedJacobian() const {
if (m_normalizedJacobian == nullptr) {
throw std::logic_error("The scaled preconditioner has not been bound to a normalized Jacobian.");
}
return *m_normalizedJacobian;
}
[[nodiscard]] const ScaledPreconditionerStatistics &GetStatistics() const noexcept {
return m_statistics;
}
private:
mfem::Solver *m_physicalInverse;
const mfem::Operator *m_physicalJacobian;
const mfem::Operator *m_expectedNormalizedJacobian;
const mfem::Operator *m_normalizedJacobian{nullptr};
const DiagonalNormalization *m_normalization;
mutable mfem::Vector m_physicalResidual;
mutable mfem::Vector m_physicalCorrection;
mutable ScaledPreconditionerStatistics m_statistics;
};
} // namespace mean_field::normalization