Add ITP one dimensional solver

Commencez par examiner le fichier joint stan-math-itp.patch et le point d’entrée solve_itp_prototype, y compris son comportement de validation et d’itération. Le travail est considéré comme terminé lorsque ce prototype aura été transformé en solveur ITP unidimensionnel demandé pour Stan Math, avec le comportement approprié de différentiation automatique décrit dans l’issue.

The code was generated by chatgpt sol. The math patch code is attached. The algorithm is from https://dl.acm.org/doi/10.1145/3423597. @WardBrian this is the solver we discussed about.

functions {
  // Replace this body to prototype a different scalar equation f(x, ...)=0.
  // This example solves x^3 - theta[1] = 0.
  real itp_residual(real x, vector theta,
                    array[] real x_r, array[] int x_i) {
    return x * x * x - theta[1];
  }

  int itp_n_half(real width, real x_tolerance, int max_num_steps) {
    int n_half = 0;
    real bisection_width = 2.0 * x_tolerance;

    while (bisection_width < width) {
      if (n_half >= max_num_steps) {
        reject("solve_itp_prototype: max_num_steps is too small for the ",
               "requested bracket width and x_tolerance");
      }
      bisection_width *= 2.0;
      n_half += 1;
    }
    return n_half;
  }

  // Overflow-resistant regula-falsi point for opposite-signed fa and fb.
  real itp_false_position(real a, real b, real fa, real fb) {
    real abs_fa = abs(fa);
    real abs_fb = abs(fb);
    real lambda;

    if (abs_fa <= abs_fb) {
      real q = abs_fa / abs_fb;
      lambda = q / (1.0 + q);
    } else {
      real q = abs_fb / abs_fa;
      lambda = 1.0 / (1.0 + q);
    }
    return a + lambda * (b - a);
  }

  // Pure-Stan value prototype of ITP. The function being solved is fixed to
  // itp_residual because ordinary Stan user functions cannot accept callbacks.
  // For parameter-dependent roots, the Stan Math implementation should be used:
  // it supplies the implicit derivative rather than differentiating through the
  // finite, branch-dependent sequence of ITP iterations.
  real solve_itp_prototype(real lower, real upper,
                           real x_tolerance,
                           real kappa1, real kappa2,
                           int n0, int max_num_steps,
                           vector theta,
                           array[] real x_r, array[] int x_i) {
    real phi = 0.5 * (1.0 + sqrt(5.0));
    real a = lower;
    real b = upper;
    real width;
    real fa;
    real fb;
    int n_half;
    int n_max;
    int iteration = 0;

    if (is_nan(a) || is_inf(a) || is_nan(b) || is_inf(b)) {
      reject("solve_itp_prototype: interval endpoints must be finite");
    }
    if (!(a < b)) {
      reject("solve_itp_prototype: lower must be strictly less than upper");
    }
    if (!(x_tolerance > 0.0) || is_nan(x_tolerance)
        || is_inf(x_tolerance)) {
      reject("solve_itp_prototype: x_tolerance must be positive and finite");
    }
    if (!(kappa1 > 0.0) || is_nan(kappa1) || is_inf(kappa1)) {
      reject("solve_itp_prototype: kappa1 must be positive and finite");
    }
    if (!(kappa2 >= 1.0 && kappa2 < 1.0 + phi)
        || is_nan(kappa2) || is_inf(kappa2)) {
      reject("solve_itp_prototype: kappa2 must satisfy 1 <= kappa2 < 1 + phi");
    }
    if (n0 < 0) {
      reject("solve_itp_prototype: n0 must be nonnegative");
    }
    if (max_num_steps <= 0) {
      reject("solve_itp_prototype: max_num_steps must be positive");
    }

    width = b - a;
    if (is_inf(width)) {
      reject("solve_itp_prototype: upper - lower must be finite");
    }

    fa = itp_residual(a, theta, x_r, x_i);
    fb = itp_residual(b, theta, x_r, x_i);
    if (is_nan(fa) || is_inf(fa) || is_nan(fb) || is_inf(fb)) {
      reject("solve_itp_prototype: endpoint function values must be finite");
    }
    if (fa == 0.0) {
      return a;
    }
    if (fb == 0.0) {
      return b;
    }
    if ((fa < 0.0 && fb < 0.0) || (fa > 0.0 && fb > 0.0)) {
      reject("solve_itp_prototype: endpoint values must have opposite signs");
    }

    n_half = itp_n_half(width, x_tolerance, max_num_steps);
    if (n0 > max_num_steps - n_half) {
      reject("solve_itp_prototype: theoretical ITP iteration bound exceeds ",
             "max_num_steps = ", max_num_steps);
    }
    n_max = n_half + n0;

    while (b - a > 2.0 * x_tolerance) {
      real midpoint;
      real x_f;
      real gap;
      real sigma;
      real x_t;
      real log_delta;
      real projection_radius;
      real projected_distance;
      real x_itp;
      real f_itp;

      if (iteration > n_max) {
        reject("solve_itp_prototype: exceeded the exact-arithmetic bound ",
               "plus one floating-point repair step");
      }
      if (iteration >= max_num_steps) {
        reject("solve_itp_prototype: reached max_num_steps before convergence");
      }

      width = b - a;
      midpoint = 0.5 * a + 0.5 * b;
      if (!(midpoint > a && midpoint < b)) {
        reject("solve_itp_prototype: no representable floating-point value ",
               "remains strictly inside the bracket at the requested tolerance");
      }

      // If the exact-arithmetic bound has been reached but the floating-point
      // width rounded just above 2 * x_tolerance, take one midpoint repair.
      x_itp = midpoint;
      if (iteration < n_max) {
        // Interpolate.
        x_f = itp_false_position(a, b, fa, fb);

        // Truncate the interpolation point toward the midpoint.
        gap = abs(midpoint - x_f);
        sigma = midpoint > x_f ? 1.0 : (midpoint < x_f ? -1.0 : 0.0);
        x_t = midpoint;
        if (gap > 0.0) {
          log_delta = log(kappa1) + kappa2 * log(width);
          if (log_delta < log(gap)) {
            x_t = x_f + sigma * exp(log_delta);
          }
        }

        // Project onto the minmax-safe interval around the midpoint.
        projection_radius
          = x_tolerance * pow(2.0, n_max - iteration) - 0.5 * width;
        if (projection_radius < 0.0) {
          projection_radius = 0.0;  // possible only through roundoff
        }
        projected_distance = projection_radius < abs(x_t - midpoint)
                             ? projection_radius
                             : abs(x_t - midpoint);
        x_itp = midpoint - sigma * projected_distance;

        // The midpoint is the final safeguarded fallback.
        if (is_nan(x_itp) || is_inf(x_itp) || !(x_itp > a && x_itp < b)) {
          x_itp = midpoint;
        }
      }

      f_itp = itp_residual(x_itp, theta, x_r, x_i);
      if (is_nan(f_itp) || is_inf(f_itp)) {
        reject("solve_itp_prototype: function returned a non-finite value");
      }
      if (f_itp == 0.0) {
        return x_itp;
      }

      if ((fa < 0.0 && f_itp > 0.0)
          || (fa > 0.0 && f_itp < 0.0)) {
        b = x_itp;
        fb = f_itp;
      } else {
        a = x_itp;
        fa = f_itp;
      }
      iteration += 1;
    }

    return 0.5 * a + 0.5 * b;
  }

  real solve_itp_default_prototype(real lower, real upper,
                                   vector theta,
                                   array[] real x_r, array[] int x_i) {
    real width = upper - lower;
    real kappa1 = width > 1e-300 ? 0.2 / width : 1.0;

    return solve_itp_prototype(
      lower, upper,
      1e-10, kappa1, 2.0,
      1, 200,
      theta, x_r, x_i
    );
  }
}

// Minimal value-only demonstration.
data {
  real<lower=0> theta_value;
}
transformed data {
  vector[1] theta;
  array[0] real x_r;
  array[0] int x_i;
  real root;

  theta[1] = theta_value;
  root = solve_itp_default_prototype(
    0.0, 1.0 + theta_value,
    theta, x_r, x_i
  );

  print("ITP root = ", root,
        "; residual = ", root * root * root - theta_value);
}
model {}

stan-math-itp.patch