<jupyter_start><jupyter_text>Exercise 9.3.4 / 10.3.2<jupyter_code>a = 1.5
b = 1.0#0.1

problem = pychastic.sde_problem.SDEProblem(
  lambda x: a*x,
  lambda x: b*x,
  x0 = 1.0,
  tmax=1,
  exact_solution = lambda x0, t, w: x0*np.exp((a-b**2/2)*t+b*w)
)

solvers = [
  SDESolver(adaptive=False, scheme='euler'),
  SDESolver(adaptive=False, scheme='milstein'),
  SDESolver(adaptive=False, scheme='wagner_platen'),
]

dts = [2**-2, 2**-3, 2**-4, 2**-5, 2**-6, 2**-7]
n_rows = 20
n_wieners_per_cell = 50
%%time

end_errors = jax.numpy.zeros((len(solvers), len(dts), n_rows, n_wieners_per_cell))

for n_solver, solver in enumerate(solvers):
  for n_dt, dt in enumerate(dts):
    solver.dt = dt
    solutions = solver.solve_many(problem, n_trajectories=n_wieners_per_cell*n_rows)

    end_times = solutions['time_values'][:, [-1]]
    end_values = solutions['solution_values'][:, -1]
    end_wieners = solutions['wiener_values'][:, -1]
    exact_values = problem.exact_solution(problem.x0, end_times, end_wieners)

    end_errors = end_errors.at[n_solver, n_dt].set( (end_values - exact_values).reshape(n_rows, n_wieners_per_cell) )

end_errors = jax.numpy.abs(end_errors)
end_errors.shape
table_euler = end_errors[0].mean(axis=-1).T
table_milstein = end_errors[1].mean(axis=-1).T
table_wp = end_errors[2].mean(axis=-1).T

mu_euler = table_euler.mean(axis=0)
sd_euler = table_euler.std(axis=0)/np.sqrt(n_wieners_per_cell)
mu_milstein = table_milstein.mean(axis=0)
sd_milstein = table_milstein.std(axis=0)/np.sqrt(n_wieners_per_cell)
mu_wp = table_wp.mean(axis=0)
sd_wp = table_wp.std(axis=0)/np.sqrt(n_wieners_per_cell)

plt.errorbar(dts, mu_euler, 2*sd_euler, capsize=5)
plt.errorbar(dts, mu_milstein, 2*sd_milstein, capsize=5)
plt.errorbar(dts, mu_wp, 2*sd_wp, capsize=5)
#plt.xlim(0, 0.02)
#plt.ylim(0, 0.1)
plt.show()

plt.plot(np.log2(dts), np.log2(mu_euler), marker='x')
plt.plot(np.log2(dts), np.log2(mu_milstein), marker='x')
plt.plot(np.log2(dts), np.log2(mu_wp), marker='x')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1752988047808765, 'fraction_numerical': 0.035856573705179286, 'mean_word_length': 5.88013698630137, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '98711', 'license': 'MIT', 'n_tokens_mistral': 966, 'n_tokens_neox': 919, 'n_words': 128, 'path': 'benchmarks/benchmark.ipynb', 'repo_name': 'RadostW/stochastic', 'size': '2008'}