<jupyter_start><jupyter_text>Brachistochrone Oracle solutionTrue parametric curve$$ \left\{\begin{array}{c}x(\theta) = R(\theta - \sin \theta)\\y(\theta) = R(1 - \cos \theta)\end{array}\right. $$<jupyter_code>from scipy.optimize import newton
import numpy as np

x2, y2 = 4., 2.
def f(theta):
    return y2 / x2 - (1 - np.cos(theta)) / (theta - np.sin(theta))
newton(f, np.pi / 2)
import matplotlib.pyplot as plt

theta2 = 3.508368768524476
N = 50
R = 2 / (1 - np.cos(theta2))
theta = np.linspace(0, theta2, N)
x = R * (theta - np.sin(theta))
y = R * (1 - np.cos(theta))

plt.plot(x, -y), R<jupyter_output><empty_output><jupyter_text>ODE solver$$ y'^2 = \frac{2 R}y - 1 $$<jupyter_code>from scipy.integrate import odeint, solve_ivp
from math import sqrt

def func(x, y):
    return sqrt(2 * r / y - 1)

r = 1.0343998433730988
bunch = solve_ivp(func, (0.105, 2.9), [0.309], method='RK23')
plt.plot(bunch.t, -bunch.y[0])<jupyter_output><empty_output><jupyter_text>Backpropagating through the ODE solver<jupyter_code>def nn_dynamics(state, time, params):
    for r in params:
        output = jnp.sqrt(2 * r / state - 1)
    return output
import jax.numpy as jnp
from jax.experimental.ode import odeint

def odenet(params, inp):
    start_end_times = jnp.array([0.105, 4])
    init_state, final_state = odeint(nn_dynamics, 0.309, start_end_times, params)
    return final_state
from jax import vmap, jit, grad
batched_odenet = vmap(odenet, in_axes=(None, 0))

step_size = 0.01

def odenet_loss(params, inputs, targets):
    preds = batched_odenet(params, inputs)
    print(preds, targets)
    return jnp.mean(jnp.sum((preds - targets) ** 2, axis=1))

@jit
def odenet_update(params, inputs, targets):
    grads = grad(odenet_loss)(params, inputs, targets)
    return [(r - step_size * dr) for (r, dr) in zip(params, grads)]
# import numpy.random as npr
# odenet_params = [npr.RandomState(0).randn(1)]
odenet_params = [2.1]
train_iters = 25

for i in range(train_iters):
    odenet_params = odenet_update(odenet_params,
                                  np.array([[x2]]), np.array([[y2]]))
    print(odenet_params)
odenet_params
this_x = np.linspace(0.105, 4, 100)
# plt.plot(this_x, -odeint(nn_dynamics, 0.309, this_x, [1.15]))
plt.plot(this_x, -odeint(nn_dynamics, 0.309, this_x, odenet_params))
plt.plot(x, -y, label='truth')
plt.plot(bunch.t, -bunch.y[0], label='best solver')
plt.plot(x, -odeint(nn_dynamics, 0.309, x, odenet_params), label='backprop odenet')
plt.legend()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15787370103916867, 'fraction_numerical': 0.04716227018385292, 'mean_word_length': 4.1929460580912865, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '521680', 'license': 'Apache-2.0', 'n_tokens_mistral': 1106, 'n_tokens_neox': 1015, 'n_words': 229, 'path': 'Module17/Brachistochrone.ipynb', 'repo_name': 'Hmkhalla/notebooks', 'size': '2502'}