Встраивание 2D-графиков в 3D-среды в matplotlib и построение отдельных точек - PullRequest
Новая
       36

Встраивание 2D-графиков в 3D-среды в matplotlib и построение отдельных точек

0 голосов
/ 23 апреля 2020

У меня есть трехмерная поверхность plot_wireframe и контурный график на плоскости xy (добавлено с использованием смещения). Я хочу нанести отдельные точки на контурный график, но когда я использую plt.scatter и устанавливаю координату z на смещение, точки (a.) 3D и (b.) Исчезают, когда я делаю их положение совпадающим с плоскостью xy. содержащий контурный график.

Есть ли способ решить эту проблему? Я приложил автономный (но немного извините) код с графиками внизу.

Спасибо! 

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np

#define function to calcualte potential.
def U(X, Y):
    U = np.zeros((len(X), len(Y)))
    for i in range(0, len(X)):
        for j in range (0, len(Y)):
            p = np.array([X[i], Y[j], 0])
            #gravitational potentials
            R_s = np.linalg.norm(p - r_s)
            R_p = np.linalg.norm(p - r_p)
            U_grav_s = -mu_s/R_s
            U_grav_p = -mu_p/R_p
            #centrifugal potential
            oa = np.cross(Omega, p)
            U_c = - 0.5 * np.dot(oa, oa)
            U[i,j] =  U_c + U_grav_p + U_grav_s
    return(U)

def U_single(X, Y):
    p = np.array([X, Y, 0])
    #gravitational potentials
    R_s = np.linalg.norm(p - r_s)
    R_p = np.linalg.norm(p - r_p)
    U_grav_s = -mu_s/R_s
    U_grav_p = -mu_p/R_p
    #centrifugal potential
    oa = np.cross(Omega, p)
    U_c = - 0.5 * np.dot(oa, oa)
    U_single =  U_c + U_grav_p + U_grav_s
    return(U_single)

#Define unscaled variables for this plot.
G = 4.0*np.pi**2
m_s = 1 #mass of Sun
m_p = 0.1 #mass of plant
mu_s = G*m_s
mu_p = G*m_p
alpha = m_p/(m_s+m_p)
beta  = m_s/(m_s+m_p)
r_s = np.array([-alpha, 0, 0]) #position of Sun
r_p = np.array([beta, 0, 0]) #position of planet
R = alpha + beta
omega = np.sqrt((mu_s+mu_p)/R**3) #define angular velocity of rotating frame
Omega = np.array([0, 0, omega]) #define vector quantity

#Generate data
X = np.arange( -1.5*R, 1.5*R, 0.01)
Y = np.arange(-1.5*R, 1.5*R, 0.01)
U_1 = U(X, Y)
U_2 = U(X, Y)
U_1[U_1<-150] = np.nan
U_2[U_2 < -250] = np.nan
X, Y = np.meshgrid(X, Y)

#plot potential surface
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, U_1, color = 'k', linewidth = 0.05,
        rcount = 500, ccount = 500)
ax.view_init(elev= 20., azim=-10)
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
ax.grid(False)
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
for line in ax.xaxis.get_ticklines():
    line.set_visible(False)
for line in ax.yaxis.get_ticklines():
    line.set_visible(False)
for line in ax.zaxis.get_ticklines():
    line.set_visible(False)
ax.text2D(0.5, 0, r"Combined Potential, $U_{tot.}$",
    horizontalalignment='center', verticalalignment='bottom',
    transform=ax.transAxes)
ax.set_zlim(-250,0)

#plot contour map
levels = np.arange(-255,0,2.5)
CS = ax.contour(X, Y, U_2, levels = levels, linewidths = 0.2, colors = 'k', offset = -250)
cntr1 = ax.contourf(X, Y, U_2, levels=levels, cmap = plt.get_cmap('jet'), offset = -250)

#Example point to plot:
L1 = np.array([0.0, R*(1-((m_p)/(3*(m_s + m_p))**(1/3))-0.315), -250])
ax.scatter(L1[0], L1[1], L1[2])

plt.show()
...