<jupyter_start><jupyter_text>BSDEHuré, Côme, <NAME>, and <NAME>. ‘Deep Backward Schemes for High-Dimensional Nonlinear PDEs’. ArXiv:1902.01599 [Cs, Math, Stat], 5 June 2020. http://arxiv.org/abs/1902.01599.Note that this is pure BSDE (without forward component), which might not be useful for our control problem.<jupyter_code>import math

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
from tqdm import tqdm

import tensorflow as tf
from keras import Model
from keras.layers import Dense, Input, Lambda
print("Num GPUs Available: ", len(tf.config.list_physical_devices("GPU")))
# n_paths
N = 1000000

# n_timesteps
M = 10

# time horizon
T = 1.0

# space dimension
d = 1

dt = T / M
BATCH_SIZE = 128
EPOCHS = 1
# generate paths
dW = np.sqrt(dt) * np.random.standard_normal((N, M))
# S0 = np.full((N, 1), X0)
S0 = np.random.random((N, 1)) * 10 - 5
S = np.cumsum(np.concatenate([S0, dW], axis=1), axis=1)
def f(t, x, y, z):
    return 0

def F(t, x, y, z, h, Delta):
    return y - f(t, x, y, z) * h + tf.transpose(z) * Delta
def build_model(t, weights=None):

    inputs_x = Input(shape=(d,))
    inputs_dw = Input(shape=(d,))

    x = Dense(d + 10, activation="tanh")(inputs_x)
    x = Dense(d + 10, activation="tanh")(x)
    outputs_u = Dense(1)(x)
    model_u = Model(inputs_x, outputs_u)

    y = Dense(d + 10, activation="tanh")(inputs_x)
    y = Dense(d + 10, activation="tanh")(y)
    outputs_z = Dense(d)(y)
    model_z = Model(inputs_x, outputs_z)

    outputs = Lambda(lambda x: F(t, x[0], x[2], x[3], dt, x[1]))([inputs_x, inputs_dw, outputs_u, outputs_z])

    model_f = Model([inputs_x, inputs_dw], outputs)
    model_f.compile(loss="mse", optimizer="adam")

    # reuse weights from previous layer
    if weights is not None:
        model_f.set_weights(weights)

    return model_u, model_z, model_f
strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))
us = [lambda x: tf.exp(-x ** 2 / 4) / tf.sqrt(4 * tf.constant(math.pi, dtype=tf.float64))]
zs = [lambda x: -x / 2 * tf.exp(-x ** 2 / 4)]
fs = []

# with strategy.scope():
for i in tqdm(range(M - 1, -1, -1)):
    weights = fs[0].get_weights() if len(fs) > 0 else None
    model_u, model_z, model_f = build_model(dt * i, weights)
    history = model_f.fit([S[:, i], dW[:, i]], us[0](S[:, i + 1]), batch_size=BATCH_SIZE, epochs=EPOCHS, verbose=1)
    us = [model_u] + us
    zs = [model_z] + zs
    fs = [model_f] + fs
fig = plt.figure(figsize=(15, 3))
ax0 = fig.add_subplot(131)
ax1 = fig.add_subplot(132)
ax2 = fig.add_subplot(133)

ax0.plot(np.linspace(0, T, M + 1), S[:100].T)

xs = np.linspace(-3, 3, 101)

# numerical
for u in us:
    ax1.plot(xs, u(xs))
    
# analytic
for i in range(M, -1, -1):
    t = T + dt * i
    ys = np.exp(-xs ** 2 / 4 / t) / np.sqrt(4 * math.pi * t)
    ax2.plot(xs, ys)
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection='3d')

x = np.stack([xs for _ in range(1, M + 1)])
y = np.array([np.full((101,), dt * i) for i in range(1, M + 1)])
z = np.array([u(xs).numpy().reshape(-1, 1) for u in us[1:]])
ax.plot_trisurf(x.reshape(-1), y.reshape(-1), z.reshape(-1), cmap='viridis', edgecolor='#00000055')

fig.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1750913520097442, 'fraction_numerical': 0.0438489646772229, 'mean_word_length': 3.7677793904209, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '457417', 'license': 'MIT', 'n_tokens_mistral': 1454, 'n_tokens_neox': 1354, 'n_words': 354, 'path': 'heat_equation.ipynb', 'repo_name': 'starovoitovs/makers', 'size': '3294'}