Как построить трехмерный график в области, определенной функцией в Python?

Question

Итак, я пытаюсь построить 3 размерно с помощью matplotlib.pyplot, но только в определенной области, граница которой задается функцией.Вот код

import numpy
from scipy.integrate import quad
import matplotlib.pyplot
import matplotlib.ticker
from mpl_toolkits.mplot3d import Axes3D

def get_x1min(pT, eta, rS):
    return (pT * numpy.exp(eta)) / (rS - pT * numpy.exp(- eta))

def get_x2(x1, pT, eta, rS):
    return (x1 * pT * numpy.exp(- eta)) / (x1 * rS - pT * numpy.exp(eta))

rS = 1960.0
eta = 0.0
pT = [66.0, 77.5, 89.5, 92.0, 94.0, 96.0, 98.0, 100.0, 103.0, 105.0, 107.0, 109.0, 111.0, 113.0, 118.5, 136.5, 157.5, 182.0, 200.0, 209.5, 241.5, 278.5, 321.0, 370.0, 426.5, 492.0]

fig = matplotlib.pyplot.figure()
ax = Axes3D(fig)
x1_test = numpy.linspace(0.1, 1.0, 20)
x1_test, pT = numpy.meshgrid(x1_test, pT)
x2_test = get_x2(x1_test, pT, eta, rS)
ax.plot_surface(x1_test, pT, x2_test, rstride=1, cstride=1, cmap=matplotlib.pyplot.get_cmap('rainbow'))
matplotlib.pyplot.show()

Это прекрасно работает, НО домен является прямоугольником, и единственное необходимое изменение - это x1_test просто начать с минимума, который можно получить из функции get_x1min.Есть ли способ сделать это?

Заранее спасибо!

Answer 1

Вы можете установить нежелательные значения на np.nan, тогда они не будут нанесены на график.Обратите внимание, что, кажется, есть некоторые проблемы с отображением цвета.Вы можете обойти это, используя ключевые слова vmin и vmax:

mask = x1_test< get_x1min(pT, eta, rS)
x2_test[mask] = numpy.nan

ax.plot_surface(
    x1_test, pT, x2_test, rstride=1, cstride=1,
    cmap=matplotlib.pyplot.get_cmap('rainbow'),
    vmin = numpy.nanmin(x2_test),
    vmax = numpy.nanmax(x2_test),    
    )

matplotlib.pyplot.show()

Остальная часть кода останется прежней.Результат выглядит примерно так: