imstkNewtonSolver.cpp

/*
** This file is part of the Interactive Medical Simulation Toolkit (iMSTK)
** iMSTK is distributed under the Apache License, Version 2.0.
** See accompanying NOTICE for details.
*/

#include "imstkNewtonSolver.h"
#include "imstkConjugateGradient.h"
#include "imstkDirectLinearSolver.h"
#include "imstkLogger.h"

#include <iostream>

namespace imstk
{
template<>
NewtonSolver<SparseMatrixd>::NewtonSolver() :
    m_linearSolver(std::make_shared<ConjugateGradient>()),
    m_forcingTerm(0.9),
    m_absoluteTolerance(1e-3),
    m_relativeTolerance(1e-6),
    m_gamma(0.9),
    m_etaMax(0.9),
    m_maxIterations(1),
    m_useArmijo(true)
{
}

template<>
NewtonSolver<Matrixd>::NewtonSolver() :
    m_linearSolver(std::make_shared<DirectLinearSolver<Matrixd>>()),
    m_forcingTerm(0.9),
    m_absoluteTolerance(1e-3),
    m_relativeTolerance(1e-6),
    m_gamma(0.9),
    m_etaMax(0.9),
    m_maxIterations(1),
    m_useArmijo(true)
{
}

template<typename SystemMatrix>
void
NewtonSolver<SystemMatrix>::solveGivenState(Vectord& x)
{
    if (!this->m_nonLinearSystem)
    {
        LOG(WARNING) << "NewtonMethod::solve - nonlinear system is not set to the nonlinear solver";
        return;
    }

    // Compute norms, set tolerances and other temporaries
    double fnorm = this->m_nonLinearSystem->evaluateF(x, this->m_isSemiImplicit).norm();
    double stopTolerance = m_absoluteTolerance + m_relativeTolerance * fnorm;

    m_linearSolver->setTolerance(stopTolerance);

    Vectord dx = x;

    const int maxIters = this->m_isSemiImplicit ? 1 : static_cast<int>(m_maxIterations);

    for (int i = 0; i < maxIters; ++i)
    {
        if (fnorm < stopTolerance)
        {
            return;
        }
        this->updateJacobian(x);

        m_linearSolver->solve(dx);
        this->m_updateIterate(-dx, x);

        double newNorm = this->armijo(dx, x, fnorm);

        if (m_forcingTerm > 0.0 && newNorm > stopTolerance)
        {
            double ratio = newNorm / fnorm; // Ratio of successive residual norms
            this->updateForcingTerm(ratio, stopTolerance, fnorm);

            // Reset tolerance in the linear solver according to the new forcing term
            // to avoid over solving of the system.
            m_linearSolver->setTolerance(m_forcingTerm);
        }

        fnorm = newNorm;
    }
}

template<typename SystemMatrix>
void
NewtonSolver<SystemMatrix>::solve()
{
    if (!this->m_nonLinearSystem)
    {
        LOG(WARNING) << "NewtonMethod::solve - nonlinear system is not set to the nonlinear solver";
        return;
    }

    size_t      iterNum;
    const auto& u      = this->m_nonLinearSystem->getUnknownVector();
    Vectord     du     = u; // make this a class member in future
    double      error0 = MAX_D;

    double epsilon = m_relativeTolerance * m_relativeTolerance;
    for (iterNum = 0; iterNum < m_maxIterations; ++iterNum)
    {
        double error = updateJacobian(u);

        if (iterNum == 0)
        {
            error0 = error;
        }

        // std::cout << "Num. of Newton Iterations: " << iterNum << "\tError ratio: " << error/error0 << ", " << error << " " << error0 << std::endl;
        if (error / error0 < epsilon && iterNum > 0)
        {
            // std::cout << "Num. of Newton Iterations: " << iterNum << "\tError ratio: " << error/error0 << ", " << error << " " << error0 << std::endl;
            break;
        }

        m_linearSolver->solve(du);
        this->m_nonLinearSystem->m_FUpdate(du, this->m_isSemiImplicit);
    }

    this->m_nonLinearSystem->m_FUpdatePrevState();

    if (iterNum == m_maxIterations && !this->m_isSemiImplicit)
    {
        LOG(WARNING) << "NewtonMethod::solve - The solver did not converge after max. iterations";
    }
}

template<typename SystemMatrix>
double
NewtonSolver<SystemMatrix>::updateJacobian(const Vectord& x)
{
    // Evaluate the Jacobian and sets the matrix
    if (!this->m_nonLinearSystem)
    {
        LOG(WARNING) << "NewtonMethod::updateJacobian - nonlinear system is not set to the nonlinear solver";
        return -1;
    }

    // auto& A = this->m_nonLinearSystem->m_dF(x);
    // if (A.innerSize() == 0)
    // {
    //     LOG(WARNING) << "NewtonMethod::updateJacobian - Size of matrix is 0!";
    //     return -1;
    // }
    //
    // auto& b = this->m_nonLinearSystem->m_F(x, this->m_isSemiImplicit);

    const auto& vecAndMat = this->m_nonLinearSystem->m_F_dF(x, this->m_isSemiImplicit);
    auto&       b = *vecAndMat.first;
    auto&       A = *vecAndMat.second;
    if (A.innerSize() == 0)
    {
        LOG(WARNING) << "NewtonMethod::updateJacobian - Size of matrix is 0!";
        return -1;
    }

    auto linearSystem = std::make_shared<typename LinearSolverType::LinearSystemType>(A, b);
    //linearSystem->setLinearProjectors(this->m_nonLinearSystem->getLinearProjectors()); /// \todo Left for near future reference. Clear in future.
    m_linearSolver->setSystem(linearSystem);

    return std::sqrt(b.dot(b));
}

template<typename SystemMatrix>
void
NewtonSolver<SystemMatrix>::updateForcingTerm(const double ratio, const double stopTolerance, const double fnorm)
{
    double eta = m_gamma * ratio * ratio;
    double forcingTermSqr = m_forcingTerm * m_forcingTerm;

    // Save guard to prevent the forcing term to become too small for far away iterates
    if (m_gamma * forcingTermSqr > 0.1)
    {
        eta = std::max(eta, m_gamma * forcingTermSqr);
    }

    m_forcingTerm = std::max(std::min(eta, m_etaMax), 0.5 * stopTolerance / fnorm);
}

template<typename SystemMatrix>
// std::shared_ptr<NewtonSolver<SystemMatrix>::LinearSolverType>
auto
NewtonSolver<SystemMatrix>::getLinearSolver() const->std::shared_ptr<LinearSolverType>
{
    return m_linearSolver;
}

template<typename SystemMatrix>
void
NewtonSolver<SystemMatrix>::setAbsoluteTolerance(const double aTolerance)
{
    m_absoluteTolerance = aTolerance;
}

template<typename SystemMatrix>
double
NewtonSolver<SystemMatrix>::getAbsoluteTolerance() const
{
    return m_absoluteTolerance;
}

template class NewtonSolver<SparseMatrixd>;
template class NewtonSolver<Matrixd>;
} // namespace imstk