test_sympy_Lambdify.py 文件源码

def test_callback_factory__broadcast():
    args = x, y = symbols('x y')
    expr = x + atan(y)
    cb = _callback_factory(args, [expr], 'numpy', np.float64, 'C')
    inp = np.array([[17, 1], [18, 2]])
    ref = [17 + np.arctan(1), 18 + np.arctan(2)]
    assert np.allclose(cb(inp), ref)

    inp2 = np.array([
        [[17, 1], [18, 2]],
        [[27, 21], [28, 22]]
    ])
    ref2 = [
        [17 + np.arctan(1), 18 + np.arctan(2)],
        [27 + np.arctan(21), 28 + np.arctan(22)]
    ]
    assert np.allclose(cb(inp2), ref2)