Багги векторов в колчане (градиент напряжения) с matplotlib

Question

Я отредактировал несколько примеров, чтобы сделать моделирование для наложения напряжения двух точечных зарядов, и сделал трехмерный график поверхности, код выглядит следующим образом:

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib import cm

    q1 = 2e-9
    q2 = -2e-9
    K = 9e9

    #Charge1 position
    x1 = 2.0
    y1 = 4.0
    #Charge2 position
    x2 = 6.0 
    y2 = 4.0

    x = np.linspace(0,8,50)
    y = np.linspace(0,8,50)
    x, y = np.meshgrid(x,y)

    r1 = np.sqrt((x - x1)**2 + (y - y1)**2)
    r2 = np.sqrt((x - x2)**2 + (y - y2)**2)

    V = K*(q1/r1 + q2/r2)

    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(x, y, V, rstride=1, cstride=1, cmap=cm.rainbow,
        linewidth=0, antialiased=False)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.show()

3D поверхность

Теперь то, что я хочу сделать, это контурный график с векторным (колчаном) графиком поверх него. Я попробовал следующий код, но я получаю кучу глючных векторов из обоих зарядов, даже отрицательных:

fig2, ax2 = plt.subplots(1,1)
cp = ax2.contourf(x, y, V, cmap=cm.coolwarm)
fig2.colorbar(cp)

v,u = np.gradient(-V, 0.2, 0.2) #E = -∇V

ax2.quiver(x, y, u, v)  
ax2.set_title("Point Charges")
plt.show()

Глючные векторы

Я подозреваю что длинные векторы связаны с делением на ноль. Векторы должны выйти из положительного заряда и попасть в отрицательный. Но как мне go исправить их? Заранее спасибо.

Answer 1

Добро пожаловать на ТАК, очень приятно MWE. Одним из вариантов будет исключение всех векторов за пределами определенной длины, установив для них значение NaN. Здесь я использую 95-й процентиль.

r = np.sqrt(u**2 + v**2)
is_valid = r < np.percentile(r, 95)
u[~is_valid] = np.nan
v[~is_valid] = np.nan
x[~is_valid] = np.nan
y[~is_valid] = np.nan

fig2, ax2 = plt.subplots(1,1)
cp = ax2.contourf(x, y, V, cmap=cm.coolwarm)
fig2.colorbar(cp)

ax2.quiver(x, y, u, v)
ax2.set_title("Point Charges")
ax2.set_xlim(0, 8)
ax2.set_ylim(0, 8)
plt.show()