使用 np.linspace() 绘制离散分布

我正在尝试使用 matplotlib 绘制一个简单的离散分布：

• 如果-1
• 如果0
• 如果 1

我该如何开始x = np.linspace(-1, 2)?

到目前为止我尝试过的是：

def mapDiscProb(x):
    if np.any(x < 0):
        return 0.3 + x * 0
    elif np.any(x >= 1):
        return 0.2 + x * 0
    else:
        return 0.5 + x * 0

x = np.linspace(-1, 2)
y = mapDiscProb(x4)

ax.plot(x, y, clip_on = False)

结果只是从 -1 到 2 的 0.3 整条线，就好像elif和else没有被执行。

我的预期输出是三个断开的水平线，这是离散 pmf 的标准。

您可以使用

numpy.piecewise

numpy.piecewise允许根据某些条件定义函数。这里你有三个条件[x<0, x>=1, (x>=0) & (x<1)]，并且您可以定义一个函数来用于每个函数。

import matplotlib.pyplot as plt
import numpy as np

l1 = lambda x: 0.3 + x * 0
l2 = lambda x: 0.2 + x * 0
l3 = lambda x: 0.5 + x * 0

mapDiscProb=lambda x: np.piecewise(x, [x<0, x>=1, (x>=0) & (x<1)],[l1,l2,l3])

x = np.linspace(-1, 2)
y = mapDiscProb(x)

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

numpy.vectorize

numpy.vectorize矢量化一个要使用标量调用的函数，以便对数组中的每个元素进行评估。这允许if/else按预期使用的语句。

import matplotlib.pyplot as plt
import numpy as np

def mapDiscProb(x):
    if x < 0:
        return 0.3
    elif x >= 1:
        return 0.2
    else:
        return 0.5

x = np.linspace(-1, 2)
y = np.vectorize(mapDiscProb)(x)

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

numpy.select

（这个想法归功于 PaulH）numpy.select可以根据条件从不同的数组中选择值。对于分段常数函数，这是一个简单的工具，因为它不需要构建任何附加函数（单行）。

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(-1, 2)
y = np.select([x<0, x<1, x>1], [0.3, 0.5, 0.2])

fig, ax = plt.subplots()
ax.plot(x, y, clip_on = False)

plt.show()

所有情况下的输出：

无垂直线

如果您不希望出现任何垂直线，则根据条件绘制尽可能多的图是有意义的。

import matplotlib.pyplot as plt
import numpy as np

l1 = lambda x: 0.3 + x * 0
l2 = lambda x: 0.2 + x * 0
l3 = lambda x: 0.5 + x * 0

x = np.linspace(-1, 2)
func = [l1,l2,l3]
cond = [x<0, x>=1, (x>=0) & (x<1)]

fig, ax = plt.subplots()
for f,c in zip(func,cond):
    xi = x[c]
    ax.plot(xi, f(xi), color="C0")

plt.show()

或者，使用numpy.select，您可以修改x数组肯定包含值[0,1]，位于条件之间的边缘。选择明确排除这些值的条件，[x<0, (x>0) & (x<1), x>1]（注意缺少任何等号）将允许将这些值设置为 nan。未显示 Nan 值，因此出现间隙。

import matplotlib.pyplot as plt
import numpy as np

x = np.sort(np.append(np.linspace(-1, 2),[0,1]))
y = np.select([x<0, (x>0) & (x<1), x>1], [0.3, 0.5, 0.2], np.nan)

fig, ax = plt.subplots()
ax.plot(x, y)

plt.show()