Skip to content

Add ITP one dimensional solver #3362

Description

@spinkney

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions