Licensed under the Apache License, Version 2.0 (the "License");


In [0]:
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Root Search

View source on GitHub

In [0]:
#@title Upgrade to TensorFlow 2.1+
!pip install --upgrade tensorflow

In [0]:
#@title Install TF Quant Finance
!pip install tf-quant-finance

This notebook demonstrates the use of low level Tensorflow Quant Finance tools for root finding using Brent's method with emphasis on the following aspects:

  • Write Once: Tensorflow supports GPUs without any significant code changes. The same model can be run on CPU and GPU
  • XLA Acceleration: The XLA compiler can reduce overhead associated with the Tensorflow graph by fusing operators.

In [0]:
#@title Imports { display-mode: "form" }

import tensorflow as tf

 # tff for Tensorflow Finance
import tf_quant_finance as tff 

root_search = tff.math.root_search

import warnings
warnings.filterwarnings("ignore",
                        category=FutureWarning)  # suppress printing warnings

Brent's Method

Find the risk free rate for an interest rate swap: $$f(x) = \log(\sum_{i=0}^{N} e^{-r_i t_i}) - V_{swap}$$


In [0]:
#@title Search Range

number_of_tenors = 100  #@param
swap_value = 0.5  #@param

dtype = tf.float32

tenors = range(1, number_of_tenors + 1)

tf.compat.v1.reset_default_graph()

t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)

def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v

## TFF on CPU
with tf.device("/cpu:0"):
  brent_result = root_search.brentq(valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))

session = tf.Session()

estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)

print("------------------------")
print("Tensorflow CPU (with auto-threading)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit session.run(brent_result)
print("------------------------")


tf.compat.v1.reset_default_graph()

t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)

def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v

## TFF on GPU
with tf.device("/gpu:0"):
  brent_result = root_search.brentq(valuation_fn, tf.constant(0, dtype=dtype), tf.constant(1, dtype=dtype))

session = tf.Session()

estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)

print("------------------------")
print("Tensorflow GPU")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit session.run(brent_result)
print("------------------------")


------------------------
Tensorflow CPU (with auto-threading)
('Converged:', True)
('Estimated root:', 0.47407696)
('Objective at root:', 1.1920929e-07)
('Number of search steps:', 5)
Timing:
1000 loops, best of 3: 985 µs per loop
------------------------
------------------------
Tensorflow GPU
('Converged:', True)
('Estimated root:', 0.47407705)
('Objective at root:', -1.7881393e-07)
('Number of search steps:', 5)
Timing:
100 loops, best of 3: 6.82 ms per loop
------------------------

Speedup from XLA


In [0]:
#@title Search Range

number_of_tenors = 100  #@param
swap_value = 0.5  #@param

dtype = np.float32

tenors = range(1, number_of_tenors + 1)


tf.compat.v1.reset_default_graph()

t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)

def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v

def brent_search(objective):
  def fn(lower_bound, upper_bound):
    return root_search.brentq(objective, lower_bound, upper_bound)
  return fn

search_fn = brent_search(valuation_fn)

## TFF on CPU compiled with XLA
with tf.device("/cpu:0"):
  brent_result = tf.xla.experimental.compile(
    search_fn,
    [
      tf.constant(0, dtype=dtype),
      tf.constant(1, dtype=dtype),
    ])

session = tf.Session()

estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)

print("------------------------")
print("Tensorflow CPU (compiled with XLA)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit session.run(brent_result)
print("------------------------")


tf.compat.v1.reset_default_graph()

t = tf.constant(tenors, dtype=dtype)
v = tf.constant(swap_value, dtype=dtype)

def valuation_fn(x):
  return tf.reduce_logsumexp(-x * t) - v

def brent_search(objective):
  def fn(lower_bound, upper_bound):
    return root_search.brentq(objective, lower_bound, upper_bound)
  return fn

search_fn = brent_search(valuation_fn)

## TFF on GPU compiled with XLA
with tf.device("/gpu:0"):
  brent_result = tf.xla.experimental.compile(
    search_fn,
    [
      tf.constant(0, dtype=dtype),
      tf.constant(1, dtype=dtype),
    ])


session = tf.Session()

estimated_root, objective_at_root, num_iterations, converged = session.run(brent_result)

print("------------------------")
print("Tensorflow GPU (compiled with XLA)")
print("Converged:", converged)
print("Estimated root:", estimated_root)
print("Objective at root:", objective_at_root)
print("Number of search steps:", num_iterations)
print("Timing:")
%timeit session.run(brent_result)
print("------------------------")


------------------------
Tensorflow CPU (compiled with XLA)
('Converged:', True)
('Estimated root:', 0.47407696)
('Objective at root:', 1.1920929e-07)
('Number of search steps:', 5)
Timing:
1000 loops, best of 3: 264 µs per loop
------------------------
------------------------
Tensorflow GPU (compiled with XLA)
('Converged:', True)
('Estimated root:', 0.47407696)
('Objective at root:', 1.1920929e-07)
('Number of search steps:', 5)
Timing:
1000 loops, best of 3: 1.73 ms per loop
------------------------