Numpy任意长度的分段

时间:2014-02-12 10:39:06

标签: python numpy piecewise

我需要构建一个具有任意数量的间隔和函数的分段函数,能够在numpy输入数组上运行。

我可以使用for循环和指标数组来实现,如下面的代码片段所示,但是有更多的pythonic方法吗?

我尝试使用numpy.piecewise,但据我所知,需要在源代码中静态定义段和函数的数量。

import numpy as np
import matplotlib.pyplot as plt 

# inputs:
#    -xs: points in which to compute the piecewise function
#    -segments: the extremes of the piecewise intervals (as a list)
#    -funcs: the functions (as a list; len(funcs)==len(segments)-1 )
def calc_piecewise(xs, segments, funcs):
    # prepare indicators and results arrays
    indaseg = np.zeros(len(xs), np.bool)
    ys = np.zeros_like(xs)

    # loop through intervals and compute the ys
    for ii in range(len(funcs)):
        indaseg = np.logical_and(xs>=segments[ii], xs<=segments[ii+1])
        ys[indaseg] = funcs[ii](xs[indaseg])

    return ys

def test_calc_piecewise():
    segments = [0.0, 1.0, 2.5, 4.0, 5.0]
    def f0(xs):
        return xs
    def f1(xs):
        return xs*xs
    def f2(xs):
        return 12.5-xs*xs
    def f3(xs):
        return 4.0*xs-19.5
    funcs = [f0, f1, f2, f3]

    xs = np.linspace(0.0, 5.0, 500)
    ys = calc_piecewise(xs, segments, funcs)

    plt.figure()
    title = "calc_piecewise"
    plt.title(title)
    plt.plot(xs, ys, 'r-')
    plt.show()

    return 


test_calc_piecewise()

1 个答案:

答案 0 :(得分:1)

您可以使用np.piecewise执行以下操作(格式化道歉!):

ys = np.piecewise(
        xs,
        [(xs >= segments[i]) & (xs <= segments[i+1]) for i in range(len(segments)-1)],
        funcs)

结果是一样的。

基本上你的循环和等同于你的行indaseg = np.logical_and(xs>=segments[ii], xs<=segments[ii+1])的测试被移动到调用代码中。