3D plot not loading Python

Question

Sorry I'm a beginner with Python and am unable to find a bug with this program because I get an infinite load. I am trying to plot 'abs(psi)' in three dimensions with matplotlib

import numpy as np
import matplotlib.pyplot as plt
import scipy.fft
from matplotlib import cm
from matplotlib.ticker import LinearLocator
#Selects the initial conditions
L=1000 #Selects the width of space
n=1000 #Selects the number of pixels
sigma=10 #Selects the size of the Gaussian
m=1 #Selects the mass of the particle
h_bar=1 #Defines planck's constant value
del_t=200000 #Chooses the change in time for the time evolution wavefunctions
k=1 #Initialises the wavenumber for Part 8

#Start of part 1
x_y_arrays=(L/n)*(np.mgrid[(-n//2):(n//2),(-n//2):(n//2), (-n//2):(n//2)]) #Width L

#can flip mgrid by doing -1*x_y_arrays to make it not upside down
plt.imshow(x_y_arrays[0,:,:]) #Plots the column array 
plt.title("Columns") #Selects the plot title
plt.colorbar()
plt.figure() #Opens another figure
plt.imshow(x_y_arrays[1,:,:]) #Plots the row array
plt.title("Rows")
plt.colorbar()

#Start of part 2
rows_x=x_y_arrays[1,:,:] #Defines the row array
columns_y=x_y_arrays[0,:,:] #Defines the column array
height_z=x_y_arrays[2,:,:]
r=np.sqrt((rows_x**2)+(columns_y**2)+(height_z**2)) #Determines the magnitude of both the column and the rows
psi=n*(np.exp(-(r**2)/(sigma**2)))*(1+0j) #Creates the Gaussian form complex array


fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
surf = ax.plot_surface(abs(psi), cmap=cm.coolwarm, linewidth=0, antialiased=False)
plt.show()