tpsanewtonmethod.hpp

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
// vi: set et ts=4 sw=4 sts=4:
/*
  Copyright 2025 NORCE AS

  This file is part of the Open Porous Media project (OPM).

  OPM is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 2 of the License, or
  (at your option) any later version.

  OPM is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with OPM.  If not, see <http://www.gnu.org/licenses/>.

  Consult the COPYING file in the top-level source directory of this
  module for the precise wording of the license and the list of
  copyright holders.
*/
#ifndef TPSA_NEWTON_METHOD_HPP
#define TPSA_NEWTON_METHOD_HPP

#include <opm/common/Exceptions.hpp>

#include <opm/models/tpsa/tpsabaseproperties.hpp>
#include <opm/models/tpsa/tpsanewtonmethodparams.hpp>
#include <opm/models/utils/timer.hpp>
#include <opm/models/utils/timerguard.hh>

#include <opm/simulators/linalg/PropertyTree.hpp>

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <exception>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <string>
#include <unistd.h>


namespace Opm {

template <class TypeTag>
class TpsaNewtonMethod
{
    using Scalar = GetPropType<TypeTag, Properties::Scalar>;
    using Simulator = GetPropType<TypeTag, Properties::Simulator>;
    using Problem = GetPropType<TypeTag, Properties::Problem>;
    using Model = GetPropType<TypeTag, Properties::ModelTPSA>;

    using SolutionVector = GetPropType<TypeTag, Properties::SolutionVectorTPSA>;
    using GlobalEqVector = GetPropType<TypeTag, Properties::GlobalEqVectorTPSA>;
    using PrimaryVariables = GetPropType<TypeTag, Properties::PrimaryVariablesTPSA>;
    using EqVector = GetPropType<TypeTag, Properties::EqVectorTPSA>;
    using Linearizer = GetPropType<TypeTag, Properties::LinearizerTPSA>;
    using LinearSolverBackend = GetPropType<TypeTag, Properties::LinearSolverBackendTPSA>;
    using SparseMatrixAdapter = GetPropType<TypeTag, Properties::SparseMatrixAdapterTPSA>;

    using IstlMatrix = typename SparseMatrixAdapter::IstlMatrix;

public:
    explicit TpsaNewtonMethod(Simulator& simulator)
        : simulator_(simulator)
        , linearSolver_(simulator)
        , error_(1e100)
        , lastError_(1e100)
        , numIterations_(0)
        , numLinearizations_(0)
    {
        // Read runtime/default Newton parameters
        params_.read();
    }

    static void registerParameters()
    {
        LinearSolverBackend::registerParameters();
        TpsaNewtonMethodParams<Scalar>::registerParameters();
    }

    bool apply()
    {
        // Make sure all timers are prestine
        prePostProcessTimer_.halt();
        linearizeTimer_.halt();
        solveTimer_.halt();
        updateTimer_.halt();

        // Get vectors and linearizer
        SolutionVector& nextSolution = model().solution(/*historyIdx=*/0);
        SolutionVector currentSolution(nextSolution);
        GlobalEqVector solutionUpdate(nextSolution.size());

        Linearizer& linearizer = model().linearizer();

        TimerGuard prePostProcessTimerGuard(prePostProcessTimer_);

        // Tell the implementation that we begin solving
        prePostProcessTimer_.start();
        begin_();
        prePostProcessTimer_.stop();

        try {
            TimerGuard innerPrePostProcessTimerGuard(prePostProcessTimer_);
            TimerGuard linearizeTimerGuard(linearizeTimer_);
            TimerGuard updateTimerGuard(updateTimer_);
            TimerGuard solveTimerGuard(solveTimer_);

            // Execute the method as long as the implementation thinks that we should do another iteration
            while (proceed_()) {
                // Notify the implementation that we're about to start a new iteration
                prePostProcessTimer_.start();
                beginIteration_();
                prePostProcessTimer_.stop();

                // Make the current solution to the old one
                currentSolution = nextSolution;

                // Do the actual linearization
                linearizeTimer_.start();
                linearizeDomain_();
                linearizeTimer_.stop();

                // Get residual and Jacobian for convergence check and preparation of linear solver
                solveTimer_.start();
                auto& residual = linearizer.residual();
                const auto& jacobian = linearizer.jacobian();
                linearSolver_.prepare(jacobian, residual);
                linearSolver_.getResidual(residual);
                solveTimer_.stop();

                // The preSolve_() method usually computes the errors, but it can do something else in addition.
                // TODO: should its costs be counted to the linearization or to the update?
                updateTimer_.start();
                preSolve_(currentSolution, residual);
                updateTimer_.stop();

                // Check convergence criteria
                if (converged()) {
                    // Tell the implementation that we're done with this iteration
                    prePostProcessTimer_.start();
                    endIteration_();
                    prePostProcessTimer_.stop();

                    break;
                }

                // Solve A x = b, where b is the residual, A is its Jacobian and x is the update of the solution
                solveTimer_.start();
                solutionUpdate = 0.0;
                const bool conv = linearSolver_.solve(solutionUpdate);
                solveTimer_.stop();

                if (!conv) {
                    solveTimer_.stop();
                    if (verbosity_() > 0) {
                        std::cout << "TPSA: Linear solver did not converge!" << std::endl;
                    }

                    prePostProcessTimer_.start();
                    failed_();
                    prePostProcessTimer_.stop();

                    return false;
                }

                // Update the current solution with the delta
                updateTimer_.start();
                update_(nextSolution, currentSolution, solutionUpdate, residual);
                updateTimer_.stop();

                // End of iteration calculations
                prePostProcessTimer_.start();
                endIteration_();
                prePostProcessTimer_.stop();
            }
        }
        catch (const Dune::Exception& e)
        {
            if (verbosity_() > 0) {
                std::cout << "TPSA: Newton method caught exception: \""
                          << e.what() << "\"\n" << std::flush;
            }

            prePostProcessTimer_.start();
            failed_();
            prePostProcessTimer_.stop();

            return false;
        }
        catch (const NumericalProblem& e)
        {
            if (verbosity_() > 0) {
                std::cout << "TPSA: Newton method caught exception: \""
                          << e.what() << "\"\n" << std::flush;
            }

            prePostProcessTimer_.start();
            failed_();
            prePostProcessTimer_.stop();

            return false;
        }

        // print the timing summary of the time step
        if (verbosity_() > 0) {
            Scalar elapsedTot =
                linearizeTimer_.realTimeElapsed() +
                solveTimer_.realTimeElapsed() +
                updateTimer_.realTimeElapsed();
            const auto default_precision{std::cout.precision()};
            std::cout << std::setprecision(2)
                      << "TPSA: "
                      << "Newton iter = " << numIterations() << " (error="
                      << error_ << ") | "
                      << "linearization = "
                      << linearizeTimer_.realTimeElapsed() << "s ("
                      << 100 * linearizeTimer_.realTimeElapsed() / elapsedTot << "%) | "
                      << "solve = "
                      << solveTimer_.realTimeElapsed() << "s ("
                      << 100 * solveTimer_.realTimeElapsed() / elapsedTot << "%) | "
                      << "update = "
                      << updateTimer_.realTimeElapsed() << "s ("
                      << 100 * updateTimer_.realTimeElapsed() / elapsedTot << "%)"
                      << "\n" << std::flush;
            std::cout << std::setprecision(default_precision); // restore default output width
        }

        // if we're not converged, tell the implementation that we've failed
        if (!converged()) {
            prePostProcessTimer_.start();
            failed_();
            if (verbosity_() > 0) {
                std::cout << "TPSA: Newton iterations did not converge!" << std::endl;
            }
            prePostProcessTimer_.stop();
            return false;
        }
        return true;
    }

    bool converged() const
    { return error_ <= tolerance(); }

    Problem& problem()
    { return simulator_.problem(); }

    const Problem& problem() const
    { return simulator_.problem(); }

    Model& model()
    { return simulator_.problem().geoMechModel(); }

    const Model& model() const
    { return simulator_.problem().geoMechModel(); }

    LinearSolverBackend& linearSolver()
    { return linearSolver_; }

    const LinearSolverBackend& linearSolver() const
    { return linearSolver_; }

    int numIterations() const
    { return numIterations_; }

    int numLinearizations() const
    { return numLinearizations_; }

    Scalar tolerance() const
    { return params_.tolerance_; }

    Scalar minIterations() const
    { return params_.minIterations_; }

    const Timer& prePostProcessTimer() const
    { return prePostProcessTimer_; }

    const Timer& linearizeTimer() const
    { return linearizeTimer_; }

    const Timer& solveTimer() const
    { return solveTimer_; }

    const Timer& updateTimer() const
    { return updateTimer_; }

protected:
    int verbosity_() const
    { return simulator_.gridView().comm().rank() == 0 ? params_.verbosity_ : 0; }

    void begin_()
    {
        numIterations_ = 0;
        numLinearizations_ = 0;
        error_ = 1e100;
    }

    void beginIteration_()
    {
        // Reset last Newton error to previous
        lastError_ = error_;
    }

    void linearizeDomain_()
    {
        model().linearizer().linearizeDomain();
        ++numLinearizations_;
    }

    void preSolve_(const SolutionVector& /*currentSolution*/,
                   const GlobalEqVector& currentResidual)
    {
        lastError_ = error_;
        Scalar newtonMaxError = params_.maxError_;

        // Calculate the error as the maximum weighted tolerance of the solution's residual
        error_ = 0;
        const auto& elemMapper = simulator_.model().elementMapper();
        for (const auto& elem : elements(simulator_.gridView(), Dune::Partitions::interior)) {
            unsigned dofIdx = elemMapper.index(elem);
            const auto& r = currentResidual[dofIdx];
            for (unsigned eqIdx = 0; eqIdx < r.size(); ++eqIdx) {
                error_ = max(std::abs(r[eqIdx] * model().eqWeight(dofIdx, eqIdx)), error_);
            }
        }

        // Take the other processes into account
        error_ = simulator_.gridView().comm().max(error_);

        // Make sure that the error never grows beyond the maximum allowed one
        if (error_ > newtonMaxError) {
            throw NumericalProblem("TPSA: Newton error " + std::to_string(double(error_)) +
                                   " is larger than maximum allowed error of " +
                                   std::to_string(double(newtonMaxError)));
        }
    }

    void update_(SolutionVector& nextSolution,
                 const SolutionVector& currentSolution,
                 const GlobalEqVector& solutionUpdate,
                 const GlobalEqVector& currentResidual)
    {
        // make sure not to swallow non-finite values at this point
        if (!std::isfinite(solutionUpdate.one_norm())) {
            throw NumericalProblem("TPSA: Non-finite update in Newton!");
        }

        std::size_t numGridDof = model().numGridDof();
        for (unsigned dofIdx = 0; dofIdx < numGridDof; ++dofIdx) {
            updatePrimaryVariables_(dofIdx,
                                   nextSolution[dofIdx],
                                   currentSolution[dofIdx],
                                   solutionUpdate[dofIdx],
                                   currentResidual[dofIdx]);
        }

        // Update material state
        model().updateMaterialState(/*timeIdx=*/0);
    }

    void updatePrimaryVariables_(unsigned /*dofIdx*/,
                                 PrimaryVariables& nextValue,
                                 const PrimaryVariables& currentValue,
                                 const EqVector& update,
                                 const EqVector& /*currentResidual*/)
    {
        nextValue = currentValue;
        nextValue -= update;
    }

    void endIteration_()
    {
        // Increase Newton iterations
        ++numIterations_;

        // Output error info
        if (verbosity_() > 1) {
            std::cout << "TPSA: End Newton iteration " << numIterations_ << ""
                      << " with error = " << error_
                      << std::endl;
        }
    }

    bool proceed_() const
    {
        if (numIterations() < params_.minIterations_) {
            return true;
        }
        else if (converged()) {
            // we are below the specified tolerance, so we don't have to
            // do more iterations
            return false;
        }
        else if (numIterations() >= params_.maxIterations_) {
            // we have exceeded the allowed number of steps.  If the
            // error was reduced by a factor of at least 4,
            // in the last iterations we proceed even if we are above
            // the maximum number of steps
            return error_ * 4.0 < lastError_;
        }

        return true;
    }

    void failed_()
    { numIterations_ = params_.targetIterations_ * 2; }

    Simulator& simulator_;
    LinearSolverBackend linearSolver_;

    Timer prePostProcessTimer_;
    Timer linearizeTimer_;
    Timer solveTimer_;
    Timer updateTimer_;

    Scalar error_;
    Scalar lastError_;
    TpsaNewtonMethodParams<Scalar> params_;

    int numIterations_;
    int numLinearizations_;
};  // class TpsaNewtonMethod

}  // namespace Opm

#endif