tfp.substrates.jax.math.find_root_secant

Finds root(s) of a function of single variable using the secant method.

The secant method is a root-finding algorithm that uses a succession of roots of secant lines to better approximate a root of a function. The secant method can be thought of as a finite-difference approximation of Newton's method. If the objective function's value becomes NaN in a position, NaN is returned in that position as the estimated root.

objective_fn Python callable for which roots are searched. It must be a callable of a single variable. objective_fn must return a Tensor of the same shape and dtype as initial_position.
initial_position Tensor or Python float representing the starting position. The function will search for roots in the neighborhood of each point. The shape of initial_position should match that of the input to objective_fn.
next_position Optional Tensor representing the next position in the search. If specified, this argument must broadcast with the shape of initial_position and have the same dtype. It will be used to compute the first step to take when searching for roots. If not specified, a default value will be used instead. Default value: initial_position * (1 + 1e-4) + sign(initial_position) * 1e-4.
value_at_position Optional Tensor or Python float representing the value of objective_fn at initial_position. If specified, this argument must have the same shape and dtype as initial_position. If not specified, the value will be evaluated during the search. Default value: None.
position_tolerance Optional Tensor representing the tolerance for the estimated roots. If specified, this argument must broadcast with the shape of initial_position and have the same dtype. Default value: 1e-8.
value_tolerance Optional Tensor representing the tolerance used to check for roots. If the absolute value of objective_fn is smaller than value_tolerance at a given position, then that position is considered a root for the function. If specified, this argument must broadcast with the shape of initial_position and have the same dtype. Default value: 1e-8.
max_iterations Optional Tensor or Python integer specifying the maximum number of steps to perform for each initial position. Must broadcast with the shape of initial_position. Default value: 50.
stopping_policy_fn Python callable controlling the algorithm termination. It must be a callable accepting a Tensor of booleans with the shape of initial_position (each denoting whether the search is finished for each starting point), and returning a scalar boolean Tensor (indicating whether the overall search should stop). Typical values are tf.reduce_all (which returns only when the search is finished for all points), and tf.reduce_any (which returns as soon as the search is finished for any point). Default value: tf.reduce_all (returns only when the search is finished for all points).
validate_args Python bool indicating whether to validate arguments such as position_tolerance, value_tolerance, and max_iterations. Default value: False.
name Python str name prefixed to ops created by this function.

root_search_results A Python namedtuple containing the following items: estimated_root: Tensor containing the last position explored. If the search was successful within the specified tolerance, this position is a root of the objective function. objective_at_estimated_root: Tensor containing the value of the objective function at position. If the search was successful within the specified tolerance, then this is close to 0. num_iterations: The number of iterations performed.

ValueError if a non-callable stopping_policy_fn is passed.

Examples

from tensorflow_probability.python.internal.backend import jax as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax
tf.enable_eager_execution()

# Example 1: Roots of a single function from two different starting points.

f = lambda x: (63 * x**5 - 70 * x**3 + 15 * x) / 8.
x = tf.constant([-1, 10], dtype=tf.float64)

tfp.math.secant_root(objective_fn=f, initial_position=x))
# ==> RootSearchResults(
    estimated_root=array([-0.90617985, 0.90617985]),
    objective_at_estimated_root=array([-4.81727769e-10, 7.44957651e-10]),
    num_iterations=array([ 7, 24], dtype=int32))

tfp.math.secant_root(objective_fn=f,
                     initial_position=x,
                     stopping_policy_fn=tf.reduce_any)
# ==> RootSearchResults(
    estimated_root=array([-0.90617985, 3.27379206]),
    objective_at_estimated_root=array([-4.81727769e-10, 2.66058312e+03]),
    num_iterations=array([7, 8], dtype=int32))

# Example 2: Roots of a multiplex function from a single starting point.

def f(x):
  return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
      + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
      + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

x = tf.constant([-1, -1], dtype=tf.float64)

tfp.math.secant_root(objective_fn=f, initial_position=x)
# ==> RootSearchResults(
    estimated_root=array([-0.77459667, -0.90617985]),
    objective_at_estimated_root=array([-7.81339438e-11, -4.81727769e-10]),
    num_iterations=array([7, 7], dtype=int32))

# Example 3: Roots of a multiplex function from two starting points.

def f(x):
  return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
      + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
      + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

x = tf.constant([[-1, -1], [10, 10]], dtype=tf.float64)

tfp.math.secant_root(objective_fn=f, initial_position=x)
# ==> RootSearchResults(
    estimated_root=array([
        [-0.77459667, -0.90617985],
        [ 0.77459667, 0.90617985]]),
    objective_at_estimated_root=array([
        [-7.81339438e-11, -4.81727769e-10],
        [6.66025013e-11, 7.44957651e-10]]),
    num_iterations=array([
        [7, 7],
        [16, 24]], dtype=int32))