Hi, I try to solve a triple integral and to obtain a 3d plot in spherical coordinates. Python is running more then an hour and I even do not know, will it finish the calculation or not. Is there something wrong with code and is it possible to make it faster?
import numpy as np
import matplotlib
matplotlib.use('Qt5Agg') # or 'Qt5Agg' for Qt backend
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy import integrate
import multiprocessing as mp
s = 3
n = 1
p = 1
r = 1
L = np.sqrt((4 * p + r ** 2) / 3)
b = np.sqrt(2 * n * p)
# Define the integrand function
def integrand(k, u, c, a, t):
return (
(k ** 2) * np.exp(-1.5 * k ** 2)
* (u ** 2) * (1 - u ** 2) * np.exp(-(b * u) ** 2) * np.cos(s * k * u * np.cos(a) / L)
* (np.cos(c) ** 2) * np.cos(k * np.sqrt(1 - u ** 2) * (s * np.sin(a) * np.cos(t) * np.cos(c)
+ s * np.sin(a) * np.sin(t) * np.sin(c)) / L)
)
def calculate_result(a, t):
result, _ = integrate.nquad(integrand, [[0, np.inf], [-1, 1], [0, 2 * np.pi]], args=(a, t))
return result
if __name__ == '__main__':
# Define the ranges for a and t
a_range = np.linspace(0, 2 * np.pi, 20)
t_range = np.linspace(0, 2 * np.pi, 20)
# Create a meshgrid for a and t
a_mesh, t_mesh = np.meshgrid(a_range, t_range)
# Number of processors to use for parallel processing
num_processors = mp.cpu_count()
# Evaluate the integrand for each combination of a and t
pool = mp.Pool(processes=num_processors)
results = pool.starmap(calculate_result, [(a, t) for a in a_range for t in t_range])
# Reshape the results to match the shape of the meshgrid
R = np.array(results).reshape(a_mesh.shape)
# Convert polar coordinates to Cartesian coordinates
X = R * np.cos(a_mesh) * np.cos(t_mesh)
Y = R * np.sin(a_mesh) * np.cos(t_mesh)
Z = R * np.sin(t_mesh)
# Create the 3D plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z, cmap='viridis')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()