stan-dev / stan-dev/math

Add ITP one dimensional solver

Open
#3,362 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature good first issue new function
Dominant language
C++
Stars
839
Forks
220
Avg merge
2d 4h
Merged PRs (30d)
14

Description

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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the attached stan-math-itp.patch and the solve_itp_prototype entry point, including its validation and iteration behavior. Done means turning this prototype into the requested Stan Math one-dimensional ITP solver, with the appropriate automatic-differentiation behavior described in the issue.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend-api-design
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.