```# Copyright 2020 The TensorTrade Authors.
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and

import random
import numpy as np
import scipy as sp

# =============================================================================
# Merton Jump Diffusion Stochastic Process
# =============================================================================
def jump_diffusion_process(params: ModelParameters):
"""
Produces a sequence of Jump Sizes which represent a jump
diffusion process. These jumps are combined with a geometric brownian
motion (log returns) to produce the Merton model.

Arguments:
params : ModelParameters
The parameters for the stochastic model.

Returns:
jump sizes for each point in time (mostly zeroes if jumps are infrequent)
"""
s_n = time = 0
small_lamda = -(1.0 / params.lamda)
jump_sizes = []
for k in range(params.all_time):
jump_sizes.append(0.0)
while s_n < params.all_time:
s_n += small_lamda * np.log(np.random.uniform(0, 1))
for j in range(params.all_time):
if time * params.all_delta <= s_n * params.all_delta <= (j + 1) * params.all_delta:
jump_sizes[j] += random.normalvariate(params.jumps_mu, params.jumps_sigma)
break
time += 1
return jump_sizes

def geometric_brownian_motion_jump_diffusion_log_returns(params: ModelParameters):
"""
Constructs combines a geometric brownian motion process
(log returns) with a jump diffusion process (log returns) to produce
a sequence of gbm jump returns.

Arguments:
params : ModelParameters
The parameters for the stochastic model.

Returns:
A GBM process with jumps in it
"""
jump_diffusion = jump_diffusion_process(params)
geometric_brownian_motion = geometric_brownian_motion_log_returns(params)

def geometric_brownian_motion_jump_diffusion_levels(params: ModelParameters):
"""
Converts a sequence of gbm jmp returns into a price sequence
which evolves according to a geometric brownian motion but can contain
jumps at any point in time.

Arguments:
params : ModelParameters
The parameters for the stochastic model.

Returns:
The price levels
"""
return convert_to_prices(params, geometric_brownian_motion_jump_diffusion_log_returns(params))

# =============================================================================
# Heston Stochastic Volatility Process
# =============================================================================
def cox_ingersoll_ross_heston(params):
"""
Constructs the rate levels of a mean-reverting cox ingersoll ross process.
Used to model interest rates as well as stochastic volatility in the Heston
model. The returns between the underlying and the stochastic volatility
should be correlated we pass a correlated Brownian motion process into the
method from which the interest rate levels are constructed. The other
correlated process are used in the Heston model.

Arguments:
params : ModelParameters
The parameters for the stochastic model.

Returns:
The interest rate levels for the CIR process
"""
# We don't multiply by sigma here because we do that in heston
sqrt_delta_sigma = np.sqrt(params.all_delta) * params.all_sigma
brownian_motion_volatility = np.random.normal(loc=0, scale=sqrt_delta_sigma, size=params.all_time)
a, mu, zero = params.heston_a, params.heston_mu, params.heston_vol0
volatilities = [zero]
for i in range(1, params.all_time):
drift = a * (mu - volatilities[i - 1]) * params.all_delta
randomness = np.sqrt(volatilities[i - 1]) * brownian_motion_volatility[i - 1]
volatilities.append(volatilities[i - 1] + drift + randomness)
return np.array(brownian_motion_volatility), np.array(volatilities)

def heston_construct_correlated_path(params: ModelParameters,
brownian_motion_one: np.array):
"""
This method is a simplified version of the Cholesky decomposition method
for just two assets. It does not make use of matrix algebra and is therefore
quite easy to implement.

Arguments:
params : ModelParameters
The parameters for the stochastic model.
brownian_motion_one : np.array
(Not filled)

Returns:
A correlated brownian motion path.
"""
# We do not multiply by sigma here, we do that in the Heston model
sqrt_delta = np.sqrt(params.all_delta)
# Construct a path correlated to the first path
brownian_motion_two = []
for i in range(params.all_time - 1):
term_one = params.cir_rho * brownian_motion_one[i]
term_two = np.sqrt(1 - pow(params.cir_rho, 2)) * random.normalvariate(0, sqrt_delta)
brownian_motion_two.append(term_one + term_two)
return np.array(brownian_motion_one), np.array(brownian_motion_two)

def heston_model_levels(params):
"""
NOTE - this method is dodgy! Need to debug!
The Heston model is the geometric brownian motion model with stochastic
volatility. This stochastic volatility is given by the cox ingersoll ross
process. Step one on this method is to construct two correlated GBM
processes. One is used for the underlying asset prices and the other is used
for the stochastic volatility levels.

Arguments:
params : ModelParameters
The parameters for the stochastic model.

Returns:
The prices for an underlying following a Heston process
"""
# Get two correlated brownian motion sequences for the volatility parameter and the underlying asset
# brownian_motion_market, brownian_motion_vol = get_correlated_paths_simple(param)
brownian, cir_process = cox_ingersoll_ross_heston(params)
brownian, brownian_motion_market = heston_construct_correlated_path(params, brownian)

heston_market_price_levels = [params.all_s0]
for i in range(1, params.all_time):
drift = params.gbm_mu * heston_market_price_levels[i - 1] * params.all_delta
vol = cir_process[i - 1] * heston_market_price_levels[i - 1] * brownian_motion_market[i - 1]
heston_market_price_levels.append(heston_market_price_levels[i - 1] + drift + vol)
return np.array(heston_market_price_levels), np.array(cir_process)

def get_correlated_geometric_brownian_motions(params: ModelParameters,
correlation_matrix: np.array,
n: int):
"""
Constructs a basket of correlated asset paths using the Cholesky
decomposition method.

Arguments:
params : ModelParameters
The parameters for the stochastic model.
correlation_matrix : np.array
An n x n correlation matrix.
n : int
Number of assets (number of paths to return)

Returns:
n correlated log return geometric brownian motion processes
"""
decomposition = sp.linalg.cholesky(correlation_matrix, lower=False)
uncorrelated_paths = []
sqrt_delta_sigma = np.sqrt(params.all_delta) * params.all_sigma
# Construct uncorrelated paths to convert into correlated paths
for i in range(params.all_time):
uncorrelated_random_numbers = []
for j in range(n):
uncorrelated_random_numbers.append(random.normalvariate(0, sqrt_delta_sigma))
uncorrelated_paths.append(np.array(uncorrelated_random_numbers))
uncorrelated_matrix = np.asmatrix(uncorrelated_paths)
correlated_matrix = uncorrelated_matrix * decomposition
assert isinstance(correlated_matrix, np.matrix)
# The rest of this method just extracts paths from the matrix
extracted_paths = []
for i in range(1, n + 1):
extracted_paths.append([])
for j in range(0, len(correlated_matrix) * n - n, n):
for i in range(n):
extracted_paths[i].append(correlated_matrix.item(j + i))
return extracted_paths

def heston(base_price: int = 1,
base_volume: int = 1,
start_date: str = '2010-01-01',
start_date_format: str = '%Y-%m-%d',
times_to_generate: int = 1000,
time_frame: str = '1h',
params: ModelParameters = None):

data_frame = generate(
price_fn=lambda p: heston_model_levels(p)[0],
base_price=base_price,
base_volume=base_volume,
start_date=start_date,
start_date_format=start_date_format,
times_to_generate=times_to_generate,
time_frame=time_frame,
params=params
)
return data_frame

return data_frame
```