feat(newton): first newton solver implementation

2026-09-08 06:36:39 -04:00
parent 76818f2f82
commit b3c04d507a
98 changed files with 20397 additions and 11040 deletions
--- a/libmeanfield/interface/solver/newton.cppm
+++ b/libmeanfield/interface/solver/newton.cppm
@@ -0,0 +1,579 @@
+module;
+
+#include <cmath>
+#include <concepts>
+#include <cstdint>
+#include <exception>
+#include <functional>
+#include <optional>
+#include <span>
+#include <stdexcept>
+#include <string_view>
+#include <type_traits>
+#include <utility>
+
+#include <mfem.hpp>
+#include <mpi.h>
+
+export module mean_field:solver.newton;
+
+export import :solver.linear_backend;
+
+export namespace mean_field::solver::nonlinear {
+    /*
+     * Backtracking counts the full Newton trial as its first trial.  A
+     * contraction is applied only after that candidate has been rejected.
+     */
+    struct BacktrackingOptions final {
+        double initialStepLength{1.0};
+        double contractionFactor{0.5};
+        double fractionToBoundarySafety{0.9};
+        double sufficientDecrease{1.0e-4};
+        double minimumStepLength{1.0e-8};
+        int maximumTrials{20};
+
+        void Validate() const {
+            if (!std::isfinite(initialStepLength) || initialStepLength <= 0.0) {
+                throw std::invalid_argument("Newton backtracking requires a finite, positive initial step length.");
+            }
+            if (!std::isfinite(contractionFactor) || contractionFactor <= 0.0 || contractionFactor >= 1.0) {
+                throw std::invalid_argument(
+                    "Newton backtracking requires a finite contraction factor strictly between zero and one."
+                );
+            }
+            if (!std::isfinite(fractionToBoundarySafety) || fractionToBoundarySafety <= 0.0 ||
+                fractionToBoundarySafety >= 1.0) {
+                throw std::invalid_argument(
+                    "Newton backtracking requires a finite fraction-to-boundary safety factor strictly between zero "
+                    "and one."
+                );
+            }
+            if (!std::isfinite(sufficientDecrease) || sufficientDecrease <= 0.0 || sufficientDecrease >= 1.0) {
+                throw std::invalid_argument(
+                    "Newton backtracking requires a finite sufficient-decrease factor strictly between zero and one."
+                );
+            }
+            if (!std::isfinite(minimumStepLength) || minimumStepLength <= 0.0 ||
+                minimumStepLength > initialStepLength) {
+                throw std::invalid_argument(
+                    "Newton backtracking requires a finite, positive minimum step no larger than the initial step."
+                );
+            }
+            if (maximumTrials <= 0) {
+                throw std::invalid_argument("Newton backtracking requires at least one permitted trial.");
+            }
+        }
+    };
+
+    struct NewtonOptions final {
+        double relativeTolerance{1.0e-8};
+        double absoluteTolerance{0.0};
+        int maximumIterations{50};
+        LinearSolveControl linearSolve{};
+        BacktrackingOptions backtracking{};
+
+        void Validate() const {
+            if (!std::isfinite(relativeTolerance) || relativeTolerance < 0.0) {
+                throw std::invalid_argument("A Newton solve requires a finite, non-negative relative tolerance.");
+            }
+            if (!std::isfinite(absoluteTolerance) || absoluteTolerance < 0.0) {
+                throw std::invalid_argument("A Newton solve requires a finite, non-negative absolute tolerance.");
+            }
+            if (maximumIterations <= 0) {
+                throw std::invalid_argument("A Newton solve requires at least one permitted iteration.");
+            }
+            linearSolve.Validate();
+            backtracking.Validate();
+        }
+
+        [[nodiscard]] double ConvergenceThreshold(const double initialResidualNorm) const {
+            Validate();
+            if (!std::isfinite(initialResidualNorm) || initialResidualNorm < 0.0) {
+                throw std::invalid_argument(
+                    "A Newton convergence threshold requires a finite, non-negative initial residual norm."
+                );
+            }
+            const double relativeThreshold = relativeTolerance * initialResidualNorm;
+            if (!std::isfinite(relativeThreshold)) {
+                throw std::invalid_argument("The Newton relative convergence threshold must be finite.");
+            }
+            return absoluteTolerance > relativeThreshold ? absoluteTolerance : relativeThreshold;
+        }
+    };
+
+    /*
+     * The MVP globalization merit is phi(x) = 0.5 ||F_normalized(x)||^2.
+     * A metric customization receives the solve communicator and must return
+     * communicator-consistent values or throw collectively. The Newton engine
+     * reduces every predicate that drives control flow, but it cannot make a
+     * rank-local exception inside an arbitrary callback collective-safe.
+     */
+    struct NormalizedResidualMetric final { };
+
+    struct MetricEvaluation final {
+        double residualNorm{0.0};
+        double merit{0.0};
+    };
+
+    [[nodiscard]] inline MetricEvaluation getMetric(
+        const NormalizedResidualMetric &,
+        const mfem::Vector &normalizedResidual,
+        const MPI_Comm communicator
+    ) {
+        if (communicator == MPI_COMM_NULL) {
+            throw std::invalid_argument("A nonlinear metric requires a valid communicator.");
+        }
+
+        double localSquaredNorm = 0.0;
+        for (int index = 0; index < normalizedResidual.Size(); ++index) {
+            const double value = static_cast<double>(normalizedResidual(index));
+            localSquaredNorm += value * value;
+        }
+
+        double globalSquaredNorm = 0.0;
+        if (MPI_Allreduce(&localSquaredNorm, &globalSquaredNorm, 1, MPI_DOUBLE, MPI_SUM, communicator) != MPI_SUCCESS) {
+            throw std::runtime_error("The nonlinear metric could not reduce the normalized residual norm.");
+        }
+
+        const double residualNorm = std::sqrt(globalSquaredNorm);
+        return {.residualNorm = residualNorm, .merit = 0.5 * residualNorm * residualNorm};
+    }
+
+    template <typename Metric = NormalizedResidualMetric>
+        requires std::move_constructible<std::remove_cvref_t<Metric>>
+    class Newton final {
+    public:
+        using MetricType = std::remove_cvref_t<Metric>;
+
+        Newton()
+            requires std::default_initializable<MetricType>
+            : Newton(
+                  NewtonOptions{},
+                  MetricType{}
+              ) {
+        }
+
+        explicit Newton(NewtonOptions options)
+            requires std::default_initializable<MetricType>
+            : Newton(
+                  std::move(options),
+                  MetricType{}
+              ) {
+        }
+
+        Newton(
+            NewtonOptions options,
+            MetricType metric
+        )
+            : m_options(std::move(options)),
+              m_metric(std::move(metric)) {
+            m_options.Validate();
+        }
+
+        [[nodiscard]] const NewtonOptions &options() const noexcept {
+            return m_options;
+        }
+
+        [[nodiscard]] const MetricType &metric() const noexcept {
+            return m_metric;
+        }
+
+    private:
+        NewtonOptions m_options;
+        [[no_unique_address]] MetricType m_metric;
+    };
+
+    Newton() -> Newton<NormalizedResidualMetric>;
+    Newton(NewtonOptions) -> Newton<NormalizedResidualMetric>;
+
+    template <typename Metric>
+    Newton(
+        NewtonOptions,
+        Metric
+    ) -> Newton<std::remove_cvref_t<Metric>>;
+
+    template <typename Candidate> struct IsNewtonConfiguration : std::false_type { };
+
+    template <typename Metric> struct IsNewtonConfiguration<Newton<Metric>> : std::true_type { };
+
+    template <typename Candidate>
+    concept NewtonConfiguration = IsNewtonConfiguration<std::remove_cvref_t<Candidate>>::value;
+
+    enum class IterationDisposition : std::uint8_t {
+        unspecified,
+        accepted,
+        converged,
+        inadmissible_state,
+        non_finite_state,
+        non_finite_residual,
+        linear_solve_failure,
+        globalization_failure,
+        stagnation,
+        iteration_limit
+    };
+
+    /*
+     * Event spans borrow solver workspaces and are valid only for the duration
+     * of the callback. Copy values that must outlive the callback.
+     */
+    struct BeforeIteration final {
+        int iteration{0};
+        double initialResidualNorm{0.0};
+        double residualNorm{0.0};
+        double relativeResidualNorm{0.0};
+        double merit{0.0};
+        MPI_Comm communicator{MPI_COMM_NULL};
+        std::span<const mfem::real_t> physicalState{};
+        std::span<const mfem::real_t> normalizedState{};
+        std::span<const mfem::real_t> normalizedResidual{};
+    };
+
+    enum class LineSearchTrialDisposition : std::uint8_t {
+        accepted,
+        inadmissible_state,
+        non_finite_state,
+        non_finite_residual,
+        insufficient_decrease
+    };
+
+    struct AfterLineSearchTrial final {
+        int iteration{0};
+        int trial{0};
+        double stepLength{0.0};
+        LineSearchTrialDisposition disposition{LineSearchTrialDisposition::insufficient_decrease};
+        std::string_view rejectionSource{};
+        std::optional<MetricEvaluation> metric{};
+        std::optional<double> minimumJacobianDeterminant{};
+        double preparationSeconds{0.0};
+        double metricSeconds{0.0};
+        MPI_Comm communicator{MPI_COMM_NULL};
+        std::span<const mfem::real_t> candidatePhysicalState{};
+        std::span<const mfem::real_t> candidateNormalizedState{};
+        std::span<const mfem::real_t> candidateNormalizedResidual{};
+    };
+
+    struct AfterIteration final {
+        int iteration{0};
+        IterationDisposition disposition{IterationDisposition::unspecified};
+        bool stepAccepted{false};
+        double acceptedStepLength{0.0};
+        int lineSearchTrials{0};
+        double initialResidualNorm{0.0};
+        double previousResidualNorm{0.0};
+        double residualNorm{0.0};
+        double relativeResidualNorm{0.0};
+        double merit{0.0};
+        double iterationSeconds{0.0};
+        double lineSearchSeconds{0.0};
+        double trialPreparationSeconds{0.0};
+        double metricEvaluationSeconds{0.0};
+        double preconditionerRefreshSeconds{0.0};
+        double rollbackSeconds{0.0};
+        std::optional<LinearSolveReport> linearSolve{};
+        MPI_Comm communicator{MPI_COMM_NULL};
+        std::span<const mfem::real_t> physicalState{};
+        std::span<const mfem::real_t> normalizedState{};
+        std::span<const mfem::real_t> normalizedResidual{};
+    };
+
+    struct NoObserver final { };
+
+    template <typename Candidate>
+    concept BeforeIterationCallback = std::invocable<Candidate &, const BeforeIteration &> &&
+                                      std::same_as<std::invoke_result_t<Candidate &, const BeforeIteration &>, void>;
+
+    template <typename Candidate>
+    concept AfterIterationCallback = std::invocable<Candidate &, const AfterIteration &> &&
+                                     std::same_as<std::invoke_result_t<Candidate &, const AfterIteration &>, void>;
+
+    template <typename Candidate>
+    concept LineSearchTrialCallback =
+        std::invocable<Candidate &, const AfterLineSearchTrial &> &&
+        std::same_as<std::invoke_result_t<Candidate &, const AfterLineSearchTrial &>, void>;
+
+    template <BeforeIterationCallback BeforeCallback, AfterIterationCallback AfterCallback>
+    class CallbackObserver final {
+    public:
+        CallbackObserver(
+            BeforeCallback before,
+            AfterCallback after
+        )
+            : m_before(std::move(before)),
+              m_after(std::move(after)) {
+        }
+
+        void beforeIteration(const BeforeIteration &event) noexcept(std::is_nothrow_invocable_v<
+                                                                    BeforeCallback &,
+                                                                    const BeforeIteration &>) {
+            std::invoke(m_before, event);
+        }
+
+        void afterIteration(const AfterIteration &event) noexcept(std::is_nothrow_invocable_v<
+                                                                  AfterCallback &,
+                                                                  const AfterIteration &>) {
+            std::invoke(m_after, event);
+        }
+
+    private:
+        [[no_unique_address]] BeforeCallback m_before;
+        [[no_unique_address]] AfterCallback m_after;
+    };
+
+    template <
+        typename BeforeCallback,
+        typename AfterCallback>
+        requires BeforeIterationCallback<std::decay_t<BeforeCallback>> &&
+                 AfterIterationCallback<std::decay_t<AfterCallback>> &&
+                 std::constructible_from<
+                     std::decay_t<BeforeCallback>,
+                     BeforeCallback> &&
+                 std::constructible_from<
+                     std::decay_t<AfterCallback>,
+                     AfterCallback>
+    [[nodiscard]] auto makeObserver(
+        BeforeCallback &&before,
+        AfterCallback &&after
+    ) {
+        return CallbackObserver<std::decay_t<BeforeCallback>, std::decay_t<AfterCallback>>{
+            std::forward<BeforeCallback>(before), std::forward<AfterCallback>(after)
+        };
+    }
+
+    template <
+        BeforeIterationCallback BeforeCallback,
+        LineSearchTrialCallback TrialCallback,
+        AfterIterationCallback AfterCallback>
+    class DetailedCallbackObserver final {
+    public:
+        DetailedCallbackObserver(
+            BeforeCallback before,
+            TrialCallback trial,
+            AfterCallback after
+        )
+            : m_before(std::move(before)),
+              m_trial(std::move(trial)),
+              m_after(std::move(after)) {
+        }
+
+        void beforeIteration(const BeforeIteration &event) noexcept(std::is_nothrow_invocable_v<
+                                                                    BeforeCallback &,
+                                                                    const BeforeIteration &>) {
+            std::invoke(m_before, event);
+        }
+
+        void afterLineSearchTrial(const AfterLineSearchTrial &event) noexcept(std::is_nothrow_invocable_v<
+                                                                              TrialCallback &,
+                                                                              const AfterLineSearchTrial &>) {
+            std::invoke(m_trial, event);
+        }
+
+        void afterIteration(const AfterIteration &event) noexcept(std::is_nothrow_invocable_v<
+                                                                  AfterCallback &,
+                                                                  const AfterIteration &>) {
+            std::invoke(m_after, event);
+        }
+
+    private:
+        [[no_unique_address]] BeforeCallback m_before;
+        [[no_unique_address]] TrialCallback m_trial;
+        [[no_unique_address]] AfterCallback m_after;
+    };
+
+    template <
+        typename BeforeCallback,
+        typename TrialCallback,
+        typename AfterCallback>
+        requires BeforeIterationCallback<std::decay_t<BeforeCallback>> &&
+                 LineSearchTrialCallback<std::decay_t<TrialCallback>> &&
+                 AfterIterationCallback<std::decay_t<AfterCallback>> &&
+                 std::constructible_from<
+                     std::decay_t<BeforeCallback>,
+                     BeforeCallback> &&
+                 std::constructible_from<
+                     std::decay_t<TrialCallback>,
+                     TrialCallback> &&
+                 std::constructible_from<
+                     std::decay_t<AfterCallback>,
+                     AfterCallback>
+    [[nodiscard]] auto makeObserver(
+        BeforeCallback &&before,
+        TrialCallback &&trial,
+        AfterCallback &&after
+    ) {
+        return DetailedCallbackObserver<
+            std::decay_t<BeforeCallback>, std::decay_t<TrialCallback>, std::decay_t<AfterCallback>>{
+            std::forward<BeforeCallback>(before), std::forward<TrialCallback>(trial), std::forward<AfterCallback>(after)
+        };
+    }
+
+    namespace detail {
+        [[nodiscard]] inline double NextBacktrackingStepLength(
+            const double rejectedStepLength,
+            const double acceptedMinimumJacobianDeterminant,
+            const std::optional<double> rejectedMinimumJacobianDeterminant,
+            const bool rejectedByInvertedGeometry,
+            const BacktrackingOptions &options
+        ) noexcept {
+            const double contractedStepLength = rejectedStepLength * options.contractionFactor;
+            if (!rejectedByInvertedGeometry || !rejectedMinimumJacobianDeterminant.has_value() ||
+                !std::isfinite(acceptedMinimumJacobianDeterminant) || acceptedMinimumJacobianDeterminant <= 0.0 ||
+                !std::isfinite(*rejectedMinimumJacobianDeterminant) || *rejectedMinimumJacobianDeterminant > 0.0) {
+                return contractedStepLength;
+            }
+
+            const double determinantChange = acceptedMinimumJacobianDeterminant - *rejectedMinimumJacobianDeterminant;
+            if (!std::isfinite(determinantChange) || determinantChange <= 0.0) {
+                return contractedStepLength;
+            }
+
+            const double estimatedBoundaryStep =
+                rejectedStepLength * acceptedMinimumJacobianDeterminant / determinantChange;
+            const double safeguardedStep = options.fractionToBoundarySafety * estimatedBoundaryStep;
+            if (!std::isfinite(safeguardedStep) || safeguardedStep <= 0.0 || safeguardedStep >= rejectedStepLength) {
+                return contractedStepLength;
+            }
+
+            /*
+             * Keep the configured backtracking ladder intact.  The geometry
+             * certificate is used only to skip rungs that its local boundary
+             * estimate says are unsafe; it does not introduce a new trial
+             * length between two rungs.  This preserves the candidates that
+             * ordinary backtracking would eventually test while avoiding the
+             * expensive preparation of the skipped, inverted geometries.
+             */
+            if (contractedStepLength <= safeguardedStep) {
+                return contractedStepLength;
+            }
+
+            const double rung =
+                std::ceil(std::log(safeguardedStep / rejectedStepLength) / std::log(options.contractionFactor));
+            double skippedStep = rejectedStepLength * std::pow(options.contractionFactor, rung);
+            if (!std::isfinite(skippedStep) || skippedStep <= 0.0 || skippedStep >= rejectedStepLength) {
+                return contractedStepLength;
+            }
+            if (skippedStep > safeguardedStep) {
+                skippedStep *= options.contractionFactor;
+            }
+            return skippedStep;
+        }
+
+        template <typename Observer>
+        inline constexpr bool isNoObserver = std::same_as<std::remove_cvref_t<Observer>, NoObserver>;
+
+        template <typename Observer>
+        concept ObservesBeforeIteration =
+            !isNoObserver<Observer> &&
+            requires(std::remove_reference_t<Observer> &observer, const BeforeIteration &event) {
+                { observer.beforeIteration(event) } -> std::same_as<void>;
+            };
+
+        template <typename Observer>
+        concept ObservesAfterIteration =
+            !isNoObserver<Observer> &&
+            requires(std::remove_reference_t<Observer> &observer, const AfterIteration &event) {
+                { observer.afterIteration(event) } -> std::same_as<void>;
+            };
+
+        template <typename Observer>
+        concept ObservesLineSearchTrial =
+            !isNoObserver<Observer> &&
+            requires(std::remove_reference_t<Observer> &observer, const AfterLineSearchTrial &event) {
+                { observer.afterLineSearchTrial(event) } -> std::same_as<void>;
+            };
+
+        template <typename Callback>
+        void InvokeObserverHookCollectively(
+            const MPI_Comm communicator,
+            const char *remoteFailureMessage,
+            Callback &&callback
+        ) {
+            std::exception_ptr localFailure;
+            try {
+                std::invoke(std::forward<Callback>(callback));
+            } catch (...) {
+                localFailure = std::current_exception();
+            }
+
+            const int localFailureFlag = localFailure != nullptr ? 1 : 0;
+            int globalFailureFlag      = 0;
+            if (MPI_Allreduce(&localFailureFlag, &globalFailureFlag, 1, MPI_INT, MPI_MAX, communicator) !=
+                MPI_SUCCESS) {
+                if (localFailure != nullptr) {
+                    std::rethrow_exception(localFailure);
+                }
+                throw std::runtime_error("The nonlinear solver could not synchronize an observer callback.");
+            }
+
+            if (globalFailureFlag != 0) {
+                if (localFailure != nullptr) {
+                    std::rethrow_exception(localFailure);
+                }
+                throw std::runtime_error(remoteFailureMessage);
+            }
+        }
+
+        template <typename Observer>
+        void InvokeBeforeIteration(
+            Observer &observer,
+            const BeforeIteration &event
+        ) {
+            if constexpr (ObservesBeforeIteration<Observer>) {
+                if constexpr (noexcept(observer.beforeIteration(event))) {
+                    observer.beforeIteration(event);
+                } else {
+                    InvokeObserverHookCollectively(
+                        event.communicator, "An observer before-iteration callback failed on another rank.",
+                        [&observer, &event] { observer.beforeIteration(event); }
+                    );
+                }
+            }
+        }
+
+        template <typename Observer>
+        void InvokeAfterIteration(
+            Observer &observer,
+            const AfterIteration &event
+        ) {
+            if constexpr (ObservesAfterIteration<Observer>) {
+                if constexpr (noexcept(observer.afterIteration(event))) {
+                    observer.afterIteration(event);
+                } else {
+                    InvokeObserverHookCollectively(
+                        event.communicator, "An observer after-iteration callback failed on another rank.",
+                        [&observer, &event] { observer.afterIteration(event); }
+                    );
+                }
+            }
+        }
+
+        template <typename Observer>
+        void InvokeAfterLineSearchTrial(
+            Observer &observer,
+            const AfterLineSearchTrial &event
+        ) {
+            if constexpr (ObservesLineSearchTrial<Observer>) {
+                if constexpr (noexcept(observer.afterLineSearchTrial(event))) {
+                    observer.afterLineSearchTrial(event);
+                } else {
+                    InvokeObserverHookCollectively(
+                        event.communicator, "An observer line-search callback failed on another rank.",
+                        [&observer, &event] { observer.afterLineSearchTrial(event); }
+                    );
+                }
+            }
+        }
+    } // namespace detail
+
+    /*
+     * Observers run synchronously on every solve rank. Ordinary callback
+     * exceptions are synchronized before the solver proceeds, so all ranks can
+     * unwind together; explicitly noexcept callbacks bypass that synchronization.
+     * A callback must still not enter an MPI collective on only a subset of
+     * ranks. A before/after pair is guaranteed for iterations that finish by
+     * returning an evaluation report. Infrastructure exceptions unwind
+     * immediately and do not promise an after callback.
+     */
+    template <typename Candidate>
+    concept NewtonObserver = detail::isNoObserver<Candidate> || detail::ObservesBeforeIteration<Candidate> ||
+                             detail::ObservesLineSearchTrial<Candidate> || detail::ObservesAfterIteration<Candidate>;
+} // namespace mean_field::solver::nonlinear