bachelor_thesis/prog/python/qqgg/vegas_diagnostics.py

31 lines
932 B
Python

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from monte_carlo import *
def plot_cubes(f, increments):
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
cubes = generate_cubes(increments)
for cube in cubes:
x_range, y_range = cube
xs, ys = np.mgrid[
x_range[0] : x_range[1] : 0.01, y_range[0] : y_range[1] : 0.01,
]
vol = get_integration_volume(cube)
fun = f(xs, ys)
ax.plot_surface(xs, ys, fun * vol * len(cubes), cmap=cm.coolwarm, linewidth=0)
ax.plot_surface(xs, ys, fun, cmap=cm.coolwarm, linewidth=0)
# ax.plot_surface(xs, ys, np.ones_like(fun) / vol, cmap=cm.coolwarm, linewidth=0)
return fig, ax
def test_f(f, *args, **kwargs):
res = integrate_vegas_nd(f, *args, **kwargs)
fig, ax = plot_cubes(f, res.increment_borders)
plt.show()
return res