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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions