Python定义分段多元周期函数报错：数组真值歧义问题解决

问题根源 #

当输入为numpy数组时，直接用if/else判断数组元素会触发The truth value of an array with more than one element is ambiguous错误——这是因为numpy数组的布尔运算返回布尔数组，而if语句需要单个布尔值，必须改用向量化操作或适配数组的分段函数方法。

以下是两种可行的重定义方案：

方案1：向量化实现（推荐，简洁高效） #

利用numpy的np.mod计算时间在周期内的偏移量，再通过np.where实现向量化分段判断，天然支持标量和数组输入。

函数定义：

import numpy as np

def periodic_sawtooth(x, t, t_off, t_on):
    cycle = t_off + t_on
    t_in_cycle = np.mod(t, cycle)
    # 分段逻辑：t_in_cycle < t_off时返回0，否则返回x的线性锯齿（从0到x）
    return np.where(t_in_cycle < t_off, 0, x * (t_in_cycle - t_off) / t_on)

测试示例 #

固定x，绘制t的函数

import matplotlib.pyplot as plt

x_fixed = 2.0
t_off = 1.0
t_on = 2.0
t_range = np.linspace(0, 6, 1000)  # 覆盖3个完整周期
func_vals = periodic_sawtooth(x_fixed, t_range, t_off, t_on)

plt.plot(t_range, func_vals)
plt.xlabel('时间t')
plt.ylabel('函数值f(x,t)')
plt.title(f'固定x={x_fixed}，周期={t_off}+{t_on}')
plt.grid(True)
plt.show()

固定t，绘制x的函数

t_fixed = 1.5  # 处于t_on阶段（1 < 1.5 < 3）
x_range = np.linspace(0, 5, 100)
func_vals = periodic_sawtooth(x_range, t_fixed, t_off, t_on)

plt.plot(x_range, func_vals)
plt.xlabel('位置x')
plt.ylabel('函数值f(x,t)')
plt.title(f'固定t={t_fixed}，周期={t_off}+{t_on}')
plt.grid(True)
plt.show()

方案2：numpy.piecewise的正确用法 #

如果需要更复杂的多段逻辑，可使用np.piecewise，注意需将额外参数（如x）传入分段函数。

函数定义：

def periodic_sawtooth_piecewise(x, t, t_off, t_on):
    cycle = t_off + t_on
    t_in_cycle = np.mod(t, cycle)
    # 定义分段条件
    conditions = [
        t_in_cycle < t_off,
        t_in_cycle >= t_off
    ]
    # 定义对应条件的计算逻辑，通过x参数传递位置变量
    funcs = [
        lambda tm, x: 0,
        lambda tm, x: x * (tm - t_off) / t_on
    ]
    return np.piecewise(t_in_cycle, conditions, funcs, x=x)

此方案可直接复用上述测试代码，结果与方案1一致。

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