I am plotting a heatmap using pcolormesh, where I create a meshgrid of X and Y values and each point having Z-valued intensities. Here is my code:

x, y, z = data_loaded['eigenenergies'], data_loaded['b_field'], data_loaded['dc_intensities']

direction = data_loaded['b_direction']
    
X = np.array(x)
Y = np.array(y)
Z = np.array(z)

# Create a DataFrame with your data
data = pd.DataFrame(
    {'x': X.flatten(), 'y': np.repeat(Y, dim), 'z': Z.flatten()})

# Pivot the data to create a 2D grid
Z = data.pivot_table(index='x', columns='y', values='z').T.values
X_unique = np.sort(data.x.unique())
Y_unique = np.sort(data.y.unique())
X, Y = np.meshgrid(X_unique, Y_unique)
print(len(X[0]))
    
# Assuming you've already prepared X, Y, and Z as described earlier
# Replace nan values with zeros
Z_cleaned = np.nan_to_num(Z)
    
norm_factor = np.array(np.max(Z_cleaned))
vmax = np.max(Z_cleaned)
#Z_cleaned = Z_cleaned / norm_factor
print(len(Z_cleaned[0]))
    
# Filter non-zero values
non_zero_indices = np.nonzero(Z_cleaned)
x_non_zero = X[non_zero_indices]
y_non_zero = Y[non_zero_indices]
z_non_zero = Z_cleaned[non_zero_indices]
    
# Manually set parameters (adjust as needed)
sigma = 0.2
    
# Compute Gaussian curves for non-zero data
gaussian_curves = gaussian(X, Y, z_non_zero, x_non_zero, y_non_zero, sigma)
    
fig, ax = plt.subplots(figsize=(5, 8))
    
# show image
shw = ax.pcolormesh(X, Y, gaussian_curves, shading='gouraud', norm=mcolors.PowerNorm(gamma=0.2), cmap=plt.get_cmap(cmap))

# make bar
bar = plt.colorbar(shw, ticks=np.linspace(vmin, vmax, 5))
    

This method give my pixels of colors, but I would like these intensities to be fitted with gaussian curves of amplitudes corresponding to their respective Z-values. This code produces the following error message:

Traceback (most recent call last):

  File ~.condaenvsQuSpinlibsite-packagesspyder_kernelspy3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:userslittl.spyder-py3untitled0.py:176
    gaussian_curves = gaussian(X, Y, z_non_zero, x_non_zero, y_non_zero, sigma)

  File c:userslittl.spyder-py3untitled0.py:35 in gaussian
    return Z * np.exp(-((x - X)**2 + (y - Y)**2) / (2 * sigma**2))

ValueError: operands could not be broadcast together with shapes (300,11658) (12000,)