Нормализовать 2d гауссовский, используя numpy

Question

Я создал функцию в 1d, которая создает гауссовский волновой пакет и нормализует его:

def gaussTimeDom1D(x, x0, k0, alpha, dx):
    # setting the initial gaussian wave packet
    gauss = np.exp((- alpha * (x - x0) ** 2) + (np.complex(0, 1) * k0 * (x - x0)))
    # normalize
    normalization_factor = np.sqrt(np.trapz(gauss.conjugate() * gauss) * dx)
    if np.abs(normalization_factor) == 0.0:
        return gauss
    else:
        return gauss / normalization_factor

Работа с ним:

# parameteres
x_size = 250.0
x_num_points = 2.0 ** 10
dx = x_size / x_num_points
x_range = np.arange(x_num_points) * dx
x_start = 90.0
k_init = 7  # initial momentum
alpha = 0.2  # gaussian width

# initial wave packet
psi_init = gaussTimeDom1D(x_range, x_start, k_init, alpha, dx)

Эта функция работает нормально, и я Я могу распространить его с помощью оператора Split Step.

Теперь я хотел применить тот же logi c к 2d гауссиану, я сделал следующее:

def gaussTimeDom2D(corr_val, init_corr, k_vec, alpha_vec, deltas):
    x, y = corr_val[:, 0], corr_val[:, 1]
    x0, y0 = init_corr
    k_x, k_y = k_vec
    a_x, a_y = alpha_vec
    dx, dy = deltas

    xx, yy = np.meshgrid(x, y)
    gauss_x = np.exp((- a_x * (xx - x0) ** 2) + (np.complex(0, 1) * k_x * (xx - x0)))
    gauss_y = np.exp((- a_y * (yy - y0) ** 2) + (np.complex(0, 1) * k_y * (yy - y0)))

    zz = gauss_x * gauss_y
    normalization = np.sqrt(np.trapz(zz.conjugate() * zz) * dx)
    normalization = np.sqrt(np.trapz(normalization.conjugate() * normalization) * dy)
    zz /= normalization
    return xx, yy, zz

Эта функция создает гауссовский, как 2d массив, zz, но результат по какой-то причине неправильно нормализован.

# parameteres
grid_size = 250.0
num_points = 2.0 ** 8
deltas = np.array([grid_size, grid_size]) / num_points
grid_range = np.array([np.arange(num_points), np.arange(num_points)]).T * deltas
start_pos = np.array([90.0, 90.0])
k_init = np.array([15.0, 15.0])  # initial momentum
alphas = np.array([0.01, 0.01])  # gaussian width

x_grid, y_grid, psi_init = gaussTimeDom2D(grid_range, start_pos, k_init, alphas, deltas)

Был бы признателен, если бы кто-нибудь мог указать мне на проблему в нормализации. Я знаю, что могу использовать другие методы интеграции, кроме трапециевидных в numpy, но мне нужно реализовать это без сложных пакетов.