Построение амплитуды вектора в полярных

Question

Я просмотрел много разных тем и не могу найти подходящий вариант для своей проблемы. У меня есть векторная функция в 2D, для которой я хочу построить амплитуду. Естественно, я определил функцию, которая выглядит довольно некрасиво, но дело в том, что это функция от r и theta. Я хотел бы раскрасить амплитуду и должен признаться, я чувствую себя немного растерянным. Мне нужно иметь определенное значение функции в каждой точке моего 2D-пространства, и я не уверен, как с этим справиться. Код выглядит примерно так:

import numpy as np
import matplotlib.pyplot as plt

#==============================================================================
# DEFINING PARAMETERS AND FUNCTIONS 
#==============================================================================

kb = 1.38e-23 #Boltzman [m²kg/(K*s²)]
e0 = 8.854e-12 #Vacuum permittivity, epsilon0 [F/m]
P = 1e-14 #Value of the static dipole moment
dpr = 0.1 #Dipole ratio, p/P, p = dpr*P
pi = np.pi
T = 300 #Temperature in Kelvins [K]

gamma = P*P*dpr/(4*pi*e0*T*kb)

#Average force in 2D space:

def F(r,theta):
    return [-1*kb*T/(2*pi + 0.5*(np.sin(theta)*gamma/(r**3))**2)*3*gamma**2*(np.sin(theta))**2/(r**7), kb*T/(2*pi + 0.5*(np.sin(theta)*gamma/(r**3))**2)*0.5*gamma**2 * np.sin(2*theta)/(r**7)]

#Amplitude of the vector:

def amp_2d(v):
    return np.sqrt(v[0]**2+v[1]**2)

#==============================================================================
# Plotting  
#==============================================================================

n = 150 #Number of points

r = np.linspace(0,1,n)
theta = np.linspace(0,2*pi,n)

Моя идея состояла в том, чтобы создать своего рода список кортежей, чтобы у меня было 150 * 150 кортежей с формой (r, theta, amp_2d (F (r, тэта))), и на этом все и заканчивается.

Спасибо за помощь!

Answer 1

Вы должны взглянуть по этой ссылке https://matplotlib.org/3.1.0/gallery/mplot3d/surface3d_radial.html.

Вам необходимо использовать сетку для 3d графика. Вы можете сделать что-то вроде:

r = np.linspace(0.00001, 1, 100)
p = np.linspace(0, 2*np.pi, 100)
R, P = np.meshgrid(r, p)

Z = amp_2d([F(R,P)])

X, Y = R*np.cos(P), R*np.sin(P)

#Set the Z limits here
ax.set_zlim(0, 1)
ax.plot_surface(X, Y, Z, cmap=plt.cm.YlGnBu_r)
plt.show()