如何使用pylab/matplotlib绘制分段sin(x)函数并解决运行报错问题

报错原因 #

• 你使用的math.sin仅支持单个数值输入，而代码中传入的x是pylab.linspace生成的numpy数组，无法被math模块的函数直接处理
• 普通if/else判断无法直接对整个数组做逐元素的条件校验，只能处理单个值的判断逻辑

最优修正方案（向量化实现，性能最高） #

直接使用numpy内置的向量化函数实现分段逻辑，无需循环和逐值判断：

import numpy as np
import pylab as plb

def f(x):
    # 逐元素取sin(x)和0的较大值，完全等价于你原本的分段逻辑
    return np.maximum(np.sin(x), 0)

x = plb.linspace(-4, 4, 10000)
plb.plot(x, f(x))
plb.show()

兼容原有写法的方案 #

如果你希望保留原有的if/else编写逻辑，可以用np.vectorize包装函数，让它支持数组的逐元素处理：

from math import sin
import numpy as np
import pylab as plb

@np.vectorize
def f(x):
    if sin(x) > 0:
        return sin(x)
    else:
        return 0

x = plb.linspace(-4, 4, 10000)
plb.plot(x, f(x))
plb.show()

内容的提问来源于stack exchange，提问作者fredericoamigo