root_finder.hpp

// Copyright (c) 2025--present, The Simons Foundation
// This file is part of TRIQS/modest and is licensed under the terms of GPLv3 or later.
// SPDX-License-Identifier: GPL-3.0-or-later
// See LICENSE in the root of this distribution for details.
 
#pragma once
 
#include "../../utility/report_stream.hpp"
 
#include <fmt/core.h>
 
#include <cmath>
#include <functional>
#include <iostream>
#include <stdexcept>
#include <string>
#include <utility>
 
namespace triqs::experimental::utility {
 
  using std::abs;
  using triqs::utility::report_stream;

 
  //------------------------------------------------------
  inline std::pair<double, double> find_bounds(std::function<double(double)> f, double x_init, double y_value, double delta_x, double precision,
                                               long max_loops = 1000) {
 
    double x = x_init;
    delta_x  = abs(delta_x);
 
    double y1  = f(x);
    double eps = (y1 - y_value >= 0) ? 1.0 : -1.0;
 
    double x1 = x;
    double y2 = y1;
    double x2 = x1;
 
    long nbre_loop = 0;
    for (; nbre_loop <= max_loops && (y2 - y_value) * eps > 0 && abs(y2 - y_value) > precision; ++nbre_loop) {
      x2 -= eps * delta_x;
      y2 = f(x2);
    }
 
    if (x1 > x2) {
      std::swap(x1, x2);
      std::swap(y1, y2);
    }
    return {x1, x2};
  }
  //------------------------------------------------------

  inline std::pair<double, double> dichotomy(std::function<double(double)> f, double x_low, double x_high, double y_target, double precision,
                                             long max_loops, std::string x_name, std::string y_name, bool verbosity) {
 
    auto out = report_stream(std::cout, verbosity);
 
    auto x1 = x_low;
    auto x2 = x_high;
 
    auto y1 = f(x_low);
    auto y2 = f(x_high);
 
    out << fmt::format(" {:>4} | {:^24} | residual\n", "iter", fmt::format("{} interval", x_name));
    out << fmt::format("------+--------------------------+----------\n");
 
    double yfound = (abs(y1 - y_target) < abs(y2 - y_target)) ? y1 : y2;
    double x      = (abs(y1 - y_target) < abs(y2 - y_target)) ? x1 : x2;
 
    long nbre_loop = 0;
    for (; nbre_loop <= max_loops && abs(yfound - y_target) > precision; ++nbre_loop) {
      x      = x1 + (x2 - x1) * (y_target - y1) / (y2 - y1);
      yfound = f(x);
      if ((y1 - y_target) * (yfound - y_target) > 0) {
        x1 = x;
        y1 = yfound;
      } else {
        x2 = x;
        y2 = yfound;
      }
      out << fmt::format(" {:4d} | [{:10g}, {:10g}]   | {:8.2e}\n", nbre_loop + 1, x1, x2, abs(yfound - y_target));
    }
 
    if (abs(yfound - y_target) < precision) {
      out << fmt::format("Converged ({} iters): {} = {:g}, {} = {:g}\n", nbre_loop, x_name, x, y_name, yfound);
      return {x, yfound};
    } else {
      out << fmt::format("Failed: {} did not converge to {:g} after {} iters\n", x_name, y_target, nbre_loop);
      throw std::runtime_error{fmt::format("Dichotomy adjustment for {} failed after {} iterations", x_name, nbre_loop)};
    }
  }
  //------------------------------------------------------

  inline std::pair<double, double> bisection(std::function<double(double)> f, double x_low, double x_high, double y_target, double precision,
                                             long max_loops, std::string x_name, std::string y_name, bool verbosity) {
 
    auto out = report_stream(std::cout, verbosity);
 
    out << fmt::format(" {:>4} | {:^24} | residual\n", "iter", fmt::format("{} interval", x_name));
    out << fmt::format("------+--------------------------+----------\n");
 
    for (auto it = 0; it < max_loops; it++) {
      auto x_mid = (x_high + x_low) / 2.0;
      auto y_mid = f(x_mid);
 
      if (abs(y_mid - y_target) <= precision) {
        out << fmt::format("Converged ({} iters): {} = {:g}, {} = {:g}\n", it + 1, x_name, x_mid, y_name, y_mid);
        return {x_mid, y_mid};
      }
 
      (y_mid - y_target >= 0 ? x_high : x_low) = x_mid;
 
      out << fmt::format(" {:4d} | [{:10g}, {:10g}]   | {:8.2e}\n", it + 1, x_low, x_high, abs(y_mid - y_target));
    }
 
    out << fmt::format("Failed: {} did not converge to {:g} after {} iters\n", x_name, y_target, max_loops);
 
    throw std::runtime_error{fmt::format("Bisection adjustment for {} failed after {} iterations", x_name, max_loops)};
  }
  //------------------------------------------------------

  inline std::pair<double, double> root_finder(std::string method, std::function<double(double)> f, double x_init, double y_value, double precision,
                                               double delta_x, long max_loops = 1000, std::string x_name = "", std::string y_name = "",
                                               bool verbosity = false) {
 
    if (verbosity) fmt::print("Root finder: seeking {} s.t. {} = {:g} \u00b1 {:g}\n", x_name, y_name, y_value, precision);
 
    auto [x1, x2] = find_bounds(f, x_init, y_value, delta_x, precision, max_loops);
 
    if (method == "dichotomy") {
      return dichotomy(f, x1, x2, y_value, precision, max_loops, x_name, y_name, verbosity);
    } else if (method == "bisection") {
      return bisection(f, x1, x2, y_value, precision, max_loops, x_name, y_name, verbosity);
    } else {
      throw std::runtime_error("Not a valid choice of root finder method!");
    }
  }

 
} // namespace triqs::experimental::utility