-
Notifications
You must be signed in to change notification settings - Fork 66
Expand file tree
/
Copy pathtest_pyvectorize.cpp
More file actions
57 lines (50 loc) · 1.87 KB
/
test_pyvectorize.cpp
File metadata and controls
57 lines (50 loc) · 1.87 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
/***************************************************************************
* Copyright (c) Wolf Vollprecht, Johan Mabille and Sylvain Corlay *
* Copyright (c) QuantStack *
* *
* Distributed under the terms of the BSD 3-Clause License. *
* *
* The full license is in the file LICENSE, distributed with this software. *
****************************************************************************/
#include "gtest/gtest.h"
#include "test_common.hpp"
#include "xtensor-python/pytensor.hpp"
#include "xtensor-python/pyvectorize.hpp"
#include "pybind11/pybind11.h"
#include "pybind11/numpy.h"
namespace xt
{
double f1(double a, double b)
{
return a + b;
}
using shape_type = std::vector<std::size_t>;
TEST(pyvectorize, function)
{
auto vecf1 = pyvectorize(f1);
shape_type shape = { 3, 2 };
pyarray<double> a(shape, 1.5);
pyarray<double> b(shape, 2.3);
pyarray<double> c = vecf1(a, b);
EXPECT_EQ(a(0, 0) + b(0, 0), c(0, 0));
}
TEST(pyvectorize, lambda)
{
auto vecf1 = pyvectorize([](double a, double b) { return a + b; });
shape_type shape = { 3, 2 };
pyarray<double> a(shape, 1.5);
pyarray<double> b(shape, 2.3);
pyarray<double> c = vecf1(a, b);
EXPECT_EQ(a(0, 0) + b(0, 0), c(0, 0));
}
TEST(pyvectorize, complex)
{
using complex_t = std::complex<double>;
shape_type shape = { 3, 2 };
pyarray<complex_t> a(shape, complex_t(1.2, 2.5));
auto f = pyvectorize([](complex_t x) { return std::abs(x); });
auto res = f(a);
double exp = std::abs(a(1, 1));
EXPECT_EQ(exp, res(1, 1));
}
}