feat(preconditioner): major work on preconditioner system
first preconditioner MVP
This commit is contained in:
@@ -0,0 +1,402 @@
|
||||
module;
|
||||
|
||||
#include <array>
|
||||
#include <concepts>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
#include <mfem.hpp>
|
||||
|
||||
export module mean_field:preconditioning.equilibrium_coordinates;
|
||||
|
||||
export import :preconditioning.specification_border;
|
||||
|
||||
export namespace mean_field::preconditioning {
|
||||
namespace detail {
|
||||
template <typename CandidateList, typename Universe> struct EquilibriumCoordinateListIsSubset;
|
||||
|
||||
template <typename... Candidates, typename Universe>
|
||||
struct EquilibriumCoordinateListIsSubset<utils::blocks::type_list<Candidates...>, Universe>
|
||||
: std::bool_constant<(utils::blocks::contains_type_v<Candidates, Universe> && ...)> { };
|
||||
} // namespace detail
|
||||
|
||||
template <typename Component, typename Form>
|
||||
concept EquilibriumCoordinateComponentFor =
|
||||
PreconditionerComponent<Component> && utils::blocks::block_form_is_valid_v<Form> &&
|
||||
std::remove_cvref_t<Component>::CorrectionBlocks::size == Form::value_block_count &&
|
||||
std::remove_cvref_t<Component>::ResidualBlocks::size == Form::residual_block_count &&
|
||||
detail::EquilibriumCoordinateListIsSubset<
|
||||
typename std::remove_cvref_t<Component>::CorrectionBlocks,
|
||||
typename Form::value_blocks>::value &&
|
||||
detail::EquilibriumCoordinateListIsSubset<
|
||||
typename std::remove_cvref_t<Component>::ResidualBlocks,
|
||||
typename Form::residual_blocks>::value;
|
||||
|
||||
struct EquilibriumCoordinateRange final {
|
||||
int equilibriumOffset{0};
|
||||
int preconditionerOffset{0};
|
||||
int size{0};
|
||||
|
||||
constexpr bool operator==(const EquilibriumCoordinateRange &) const = default;
|
||||
};
|
||||
|
||||
struct EquilibriumCoordinateMapStatistics final {
|
||||
std::uint64_t residualPacks{0};
|
||||
std::uint64_t residualUnpacks{0};
|
||||
std::uint64_t correctionPacks{0};
|
||||
std::uint64_t correctionUnpacks{0};
|
||||
};
|
||||
|
||||
template <typename Form, typename Component>
|
||||
requires EquilibriumCoordinateComponentFor<Component, Form>
|
||||
class EquilibriumPreconditionerCoordinateMap final {
|
||||
private:
|
||||
using ComponentType = std::remove_cvref_t<Component>;
|
||||
using Layout = utils::blocks::form_layout<Form>;
|
||||
|
||||
static constexpr std::size_t correctionBlockCount = ComponentType::CorrectionBlocks::size;
|
||||
static constexpr std::size_t residualBlockCount = ComponentType::ResidualBlocks::size;
|
||||
|
||||
public:
|
||||
explicit EquilibriumPreconditionerCoordinateMap(const Layout &layout)
|
||||
: m_correctionRanges(MakeCorrectionRanges(
|
||||
layout,
|
||||
typename ComponentType::CorrectionBlocks{}
|
||||
)),
|
||||
m_residualRanges(MakeResidualRanges(
|
||||
layout,
|
||||
typename ComponentType::ResidualBlocks{}
|
||||
)),
|
||||
m_equilibriumStateSize(layout.value_offsets().Last()),
|
||||
m_equilibriumResidualSize(layout.residual_offsets().Last()),
|
||||
m_preconditionerCorrectionSize(TotalSize(m_correctionRanges)),
|
||||
m_preconditionerResidualSize(TotalSize(m_residualRanges)) {
|
||||
if (m_preconditionerCorrectionSize != m_equilibriumStateSize ||
|
||||
m_preconditionerResidualSize != m_equilibriumResidualSize) {
|
||||
throw std::logic_error(
|
||||
"The typed preconditioner coordinate map does not span the complete equilibrium operator."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
void PackResidual(
|
||||
const mfem::Vector &equilibriumResidual,
|
||||
mfem::Vector &preconditionerResidual
|
||||
) const {
|
||||
VerifySizes(
|
||||
equilibriumResidual, m_equilibriumResidualSize, preconditionerResidual, m_preconditionerResidualSize,
|
||||
"residual pack"
|
||||
);
|
||||
EquilibriumToPreconditioner(equilibriumResidual, preconditionerResidual, m_residualRanges);
|
||||
++m_statistics.residualPacks;
|
||||
}
|
||||
|
||||
void UnpackResidual(
|
||||
const mfem::Vector &preconditionerResidual,
|
||||
mfem::Vector &equilibriumResidual
|
||||
) const {
|
||||
VerifySizes(
|
||||
preconditionerResidual, m_preconditionerResidualSize, equilibriumResidual, m_equilibriumResidualSize,
|
||||
"residual unpack"
|
||||
);
|
||||
PreconditionerToEquilibrium(preconditionerResidual, equilibriumResidual, m_residualRanges);
|
||||
++m_statistics.residualUnpacks;
|
||||
}
|
||||
|
||||
void PackCorrection(
|
||||
const mfem::Vector &equilibriumCorrection,
|
||||
mfem::Vector &preconditionerCorrection
|
||||
) const {
|
||||
VerifySizes(
|
||||
equilibriumCorrection, m_equilibriumStateSize, preconditionerCorrection, m_preconditionerCorrectionSize,
|
||||
"correction pack"
|
||||
);
|
||||
EquilibriumToPreconditioner(equilibriumCorrection, preconditionerCorrection, m_correctionRanges);
|
||||
++m_statistics.correctionPacks;
|
||||
}
|
||||
|
||||
void UnpackCorrection(
|
||||
const mfem::Vector &preconditionerCorrection,
|
||||
mfem::Vector &equilibriumCorrection
|
||||
) const {
|
||||
VerifySizes(
|
||||
preconditionerCorrection, m_preconditionerCorrectionSize, equilibriumCorrection, m_equilibriumStateSize,
|
||||
"correction unpack"
|
||||
);
|
||||
PreconditionerToEquilibrium(preconditionerCorrection, equilibriumCorrection, m_correctionRanges);
|
||||
++m_statistics.correctionUnpacks;
|
||||
}
|
||||
|
||||
[[nodiscard]] int EquilibriumStateSize() const noexcept {
|
||||
return m_equilibriumStateSize;
|
||||
}
|
||||
|
||||
[[nodiscard]] int EquilibriumResidualSize() const noexcept {
|
||||
return m_equilibriumResidualSize;
|
||||
}
|
||||
|
||||
[[nodiscard]] int PreconditionerCorrectionSize() const noexcept {
|
||||
return m_preconditionerCorrectionSize;
|
||||
}
|
||||
|
||||
[[nodiscard]] int PreconditionerResidualSize() const noexcept {
|
||||
return m_preconditionerResidualSize;
|
||||
}
|
||||
|
||||
[[nodiscard]] const std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
correctionBlockCount> &
|
||||
GetCorrectionRanges() const noexcept {
|
||||
return m_correctionRanges;
|
||||
}
|
||||
|
||||
[[nodiscard]] const std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
residualBlockCount> &
|
||||
GetResidualRanges() const noexcept {
|
||||
return m_residualRanges;
|
||||
}
|
||||
|
||||
[[nodiscard]] const EquilibriumCoordinateMapStatistics &GetStatistics() const noexcept {
|
||||
return m_statistics;
|
||||
}
|
||||
|
||||
private:
|
||||
template <typename... Blocks>
|
||||
[[nodiscard]] static std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
sizeof...(Blocks)>
|
||||
MakeCorrectionRanges(
|
||||
const Layout &layout,
|
||||
utils::blocks::type_list<Blocks...>
|
||||
) {
|
||||
std::array<EquilibriumCoordinateRange, sizeof...(Blocks)> ranges{};
|
||||
int preconditionerOffset = 0;
|
||||
std::size_t range = 0;
|
||||
(
|
||||
[&] {
|
||||
constexpr int equilibriumBlock = utils::blocks::type_index_v<Blocks, typename Form::value_blocks>;
|
||||
const int size = layout.size(utils::blocks::value_block<equilibriumBlock>{});
|
||||
ranges[range++] = {
|
||||
.equilibriumOffset = layout.offset(utils::blocks::value_block<equilibriumBlock>{}),
|
||||
.preconditionerOffset = preconditionerOffset,
|
||||
.size = size
|
||||
};
|
||||
preconditionerOffset += size;
|
||||
}(),
|
||||
...);
|
||||
return ranges;
|
||||
}
|
||||
|
||||
template <typename... Blocks>
|
||||
[[nodiscard]] static std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
sizeof...(Blocks)>
|
||||
MakeResidualRanges(
|
||||
const Layout &layout,
|
||||
utils::blocks::type_list<Blocks...>
|
||||
) {
|
||||
std::array<EquilibriumCoordinateRange, sizeof...(Blocks)> ranges{};
|
||||
int preconditionerOffset = 0;
|
||||
std::size_t range = 0;
|
||||
(
|
||||
[&] {
|
||||
constexpr int equilibriumBlock =
|
||||
utils::blocks::type_index_v<Blocks, typename Form::residual_blocks>;
|
||||
const int size = layout.size(utils::blocks::residual_block<equilibriumBlock>{});
|
||||
ranges[range++] = {
|
||||
.equilibriumOffset = layout.offset(utils::blocks::residual_block<equilibriumBlock>{}),
|
||||
.preconditionerOffset = preconditionerOffset,
|
||||
.size = size
|
||||
};
|
||||
preconditionerOffset += size;
|
||||
}(),
|
||||
...);
|
||||
return ranges;
|
||||
}
|
||||
|
||||
template <std::size_t Size>
|
||||
[[nodiscard]] static int TotalSize(
|
||||
const std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
Size> &ranges
|
||||
) noexcept {
|
||||
int size = 0;
|
||||
for (const auto &range : ranges) {
|
||||
size += range.size;
|
||||
}
|
||||
return size;
|
||||
}
|
||||
|
||||
template <std::size_t Size>
|
||||
static void EquilibriumToPreconditioner(
|
||||
const mfem::Vector &equilibrium,
|
||||
mfem::Vector &preconditioner,
|
||||
const std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
Size> &ranges
|
||||
) {
|
||||
for (const auto &range : ranges) {
|
||||
for (int index = 0; index < range.size; ++index) {
|
||||
preconditioner(range.preconditionerOffset + index) = equilibrium(range.equilibriumOffset + index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <std::size_t Size>
|
||||
static void PreconditionerToEquilibrium(
|
||||
const mfem::Vector &preconditioner,
|
||||
mfem::Vector &equilibrium,
|
||||
const std::array<
|
||||
EquilibriumCoordinateRange,
|
||||
Size> &ranges
|
||||
) {
|
||||
for (const auto &range : ranges) {
|
||||
for (int index = 0; index < range.size; ++index) {
|
||||
equilibrium(range.equilibriumOffset + index) = preconditioner(range.preconditionerOffset + index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void VerifySizes(
|
||||
const mfem::Vector &source,
|
||||
const int expectedSourceSize,
|
||||
const mfem::Vector &destination,
|
||||
const int expectedDestinationSize,
|
||||
const char *operation
|
||||
) {
|
||||
if (source.Size() != expectedSourceSize || destination.Size() != expectedDestinationSize) {
|
||||
throw std::invalid_argument(
|
||||
std::string("The equilibrium preconditioner ") + operation + " received an incompatible vector."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
std::array<EquilibriumCoordinateRange, correctionBlockCount> m_correctionRanges;
|
||||
std::array<EquilibriumCoordinateRange, residualBlockCount> m_residualRanges;
|
||||
int m_equilibriumStateSize;
|
||||
int m_equilibriumResidualSize;
|
||||
int m_preconditionerCorrectionSize;
|
||||
int m_preconditionerResidualSize;
|
||||
mutable EquilibriumCoordinateMapStatistics m_statistics;
|
||||
};
|
||||
|
||||
struct PreparedStellarPreconditionerStatistics final {
|
||||
std::uint64_t applications{0};
|
||||
std::uint64_t residualCoordinateMappings{0};
|
||||
std::uint64_t correctionCoordinateMappings{0};
|
||||
};
|
||||
|
||||
template <equilibrium::DiscretizedStellarEquilibriumProblem Problem, SpecificationBorderBlockType Block>
|
||||
requires EquilibriumCoordinateComponentFor<Block, typename std::remove_cvref_t<Problem>::FormType>
|
||||
class PreparedStellarPreconditioner final : public mfem::Solver {
|
||||
private:
|
||||
using ProblemType = std::remove_cvref_t<Problem>;
|
||||
using BlockType = std::remove_cvref_t<Block>;
|
||||
|
||||
public:
|
||||
using Form = typename ProblemType::FormType;
|
||||
using BackendType = typename BlockType::BackendType;
|
||||
using GroupedPreconditioner = PreparedSpecificationBorderBlock<ProblemType, BlockType>;
|
||||
using CoordinateMap = EquilibriumPreconditionerCoordinateMap<Form, BlockType>;
|
||||
|
||||
PreparedStellarPreconditioner(
|
||||
const ProblemType &problem,
|
||||
BlockType block
|
||||
)
|
||||
: mfem::Solver(problem.StateSize()),
|
||||
m_grouped(
|
||||
problem,
|
||||
std::move(block)
|
||||
),
|
||||
m_coordinates(problem.GetManifest().layout()),
|
||||
m_groupedResidual(m_coordinates.PreconditionerResidualSize()),
|
||||
m_groupedCorrection(m_coordinates.PreconditionerCorrectionSize()) {
|
||||
if (problem.StateSize() != problem.EquationSize() || m_grouped.Height() != problem.StateSize() ||
|
||||
m_grouped.Width() != problem.EquationSize()) {
|
||||
throw std::logic_error(
|
||||
"The prepared stellar preconditioner is incompatible with the complete equilibrium operator."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
PreparedStellarPreconditioner(const PreparedStellarPreconditioner &) = delete;
|
||||
PreparedStellarPreconditioner &operator=(const PreparedStellarPreconditioner &) = delete;
|
||||
PreparedStellarPreconditioner(PreparedStellarPreconditioner &&) = delete;
|
||||
PreparedStellarPreconditioner &operator=(PreparedStellarPreconditioner &&) = delete;
|
||||
|
||||
void SetOperator(const mfem::Operator &operation) override {
|
||||
if (operation.Height() != Height() || operation.Width() != Width()) {
|
||||
throw std::invalid_argument(
|
||||
"The prepared stellar preconditioner received an incompatible equilibrium operator."
|
||||
);
|
||||
}
|
||||
m_grouped.SetOperator(operation);
|
||||
}
|
||||
|
||||
void Mult(
|
||||
const mfem::Vector &equilibriumResidual,
|
||||
mfem::Vector &equilibriumCorrection
|
||||
) const override {
|
||||
if (equilibriumResidual.Size() != Width() || equilibriumCorrection.Size() != Height()) {
|
||||
throw std::invalid_argument(
|
||||
"The prepared stellar preconditioner requires compatible, preallocated equilibrium vectors."
|
||||
);
|
||||
}
|
||||
m_coordinates.PackResidual(equilibriumResidual, m_groupedResidual);
|
||||
++m_statistics.residualCoordinateMappings;
|
||||
m_grouped.Mult(m_groupedResidual, m_groupedCorrection);
|
||||
m_coordinates.UnpackCorrection(m_groupedCorrection, equilibriumCorrection);
|
||||
++m_statistics.correctionCoordinateMappings;
|
||||
++m_statistics.applications;
|
||||
}
|
||||
|
||||
[[nodiscard]] SpecificationBorderBlockPreparationReport Refresh() {
|
||||
return m_grouped.Refresh();
|
||||
}
|
||||
|
||||
[[nodiscard]] bool IsCurrent() const {
|
||||
return m_grouped.IsCurrent();
|
||||
}
|
||||
|
||||
[[nodiscard]] const BlockType &GetBlock() const noexcept {
|
||||
return m_grouped.GetBlock();
|
||||
}
|
||||
|
||||
[[nodiscard]] const GroupedPreconditioner &GetGroupedPreconditioner() const noexcept {
|
||||
return m_grouped;
|
||||
}
|
||||
|
||||
[[nodiscard]] const CoordinateMap &GetCoordinateMap() const noexcept {
|
||||
return m_coordinates;
|
||||
}
|
||||
|
||||
[[nodiscard]] const PreparedStellarPreconditionerStatistics &GetStatistics() const noexcept {
|
||||
return m_statistics;
|
||||
}
|
||||
|
||||
private:
|
||||
GroupedPreconditioner m_grouped;
|
||||
CoordinateMap m_coordinates;
|
||||
mutable mfem::Vector m_groupedResidual;
|
||||
mutable mfem::Vector m_groupedCorrection;
|
||||
mutable PreparedStellarPreconditionerStatistics m_statistics;
|
||||
};
|
||||
|
||||
template <
|
||||
equilibrium::DiscretizedStellarEquilibriumProblem Problem,
|
||||
SpecificationBorderBlockType Block>
|
||||
requires EquilibriumCoordinateComponentFor<
|
||||
Block,
|
||||
typename std::remove_cvref_t<Problem>::FormType>
|
||||
[[nodiscard]] auto prepare(
|
||||
const Problem &problem,
|
||||
Block block
|
||||
) {
|
||||
return PreparedStellarPreconditioner<Problem, Block>{problem, std::move(block)};
|
||||
}
|
||||
} // namespace mean_field::preconditioning
|
||||
Reference in New Issue
Block a user