Python 在极坐标系中绘制矢量振幅

我查看了许多不同的线程，似乎找不到适合我的问题的选项。我有一个2D的向量函数，我想画出它的振幅。很自然地，我定义了一个函数，看起来很难看，但重点是，它是r和θ的函数。我想用彩色绘出振幅图，我必须承认，我感到有点失落。我需要在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，θ，amp_2d（F（r，θ）），这几乎就是它的终点

---

谢谢你的帮助

你应该看看这个链接

您需要使用网格网格进行三维打印。您可以执行以下操作：

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()