test_numpy_vectorize.py
Go to the documentation of this file.
1 from __future__ import annotations
2 
3 import pytest
4 
5 from pybind11_tests import numpy_vectorize as m
6 
7 np = pytest.importorskip("numpy")
8 
9 
10 def test_vectorize(capture):
11  assert np.isclose(m.vectorized_func3(np.array(3 + 7j)), [6 + 14j])
12 
13  for f in [m.vectorized_func, m.vectorized_func2]:
14  with capture:
15  assert np.isclose(f(1, 2, 3), 6)
16  assert capture == "my_func(x:int=1, y:float=2, z:float=3)"
17  with capture:
18  assert np.isclose(f(np.array(1), np.array(2), 3), 6)
19  assert capture == "my_func(x:int=1, y:float=2, z:float=3)"
20  with capture:
21  assert np.allclose(f(np.array([1, 3]), np.array([2, 4]), 3), [6, 36])
22  assert (
23  capture
24  == """
25  my_func(x:int=1, y:float=2, z:float=3)
26  my_func(x:int=3, y:float=4, z:float=3)
27  """
28  )
29  with capture:
30  a = np.array([[1, 2], [3, 4]], order="F")
31  b = np.array([[10, 20], [30, 40]], order="F")
32  c = 3
33  result = f(a, b, c)
34  assert np.allclose(result, a * b * c)
35  assert result.flags.f_contiguous
36  # All inputs are F order and full or singletons, so we the result is in col-major order:
37  assert (
38  capture
39  == """
40  my_func(x:int=1, y:float=10, z:float=3)
41  my_func(x:int=3, y:float=30, z:float=3)
42  my_func(x:int=2, y:float=20, z:float=3)
43  my_func(x:int=4, y:float=40, z:float=3)
44  """
45  )
46  with capture:
47  a, b, c = (
48  np.array([[1, 3, 5], [7, 9, 11]]),
49  np.array([[2, 4, 6], [8, 10, 12]]),
50  3,
51  )
52  assert np.allclose(f(a, b, c), a * b * c)
53  assert (
54  capture
55  == """
56  my_func(x:int=1, y:float=2, z:float=3)
57  my_func(x:int=3, y:float=4, z:float=3)
58  my_func(x:int=5, y:float=6, z:float=3)
59  my_func(x:int=7, y:float=8, z:float=3)
60  my_func(x:int=9, y:float=10, z:float=3)
61  my_func(x:int=11, y:float=12, z:float=3)
62  """
63  )
64  with capture:
65  a, b, c = np.array([[1, 2, 3], [4, 5, 6]]), np.array([2, 3, 4]), 2
66  assert np.allclose(f(a, b, c), a * b * c)
67  assert (
68  capture
69  == """
70  my_func(x:int=1, y:float=2, z:float=2)
71  my_func(x:int=2, y:float=3, z:float=2)
72  my_func(x:int=3, y:float=4, z:float=2)
73  my_func(x:int=4, y:float=2, z:float=2)
74  my_func(x:int=5, y:float=3, z:float=2)
75  my_func(x:int=6, y:float=4, z:float=2)
76  """
77  )
78  with capture:
79  a, b, c = np.array([[1, 2, 3], [4, 5, 6]]), np.array([[2], [3]]), 2
80  assert np.allclose(f(a, b, c), a * b * c)
81  assert (
82  capture
83  == """
84  my_func(x:int=1, y:float=2, z:float=2)
85  my_func(x:int=2, y:float=2, z:float=2)
86  my_func(x:int=3, y:float=2, z:float=2)
87  my_func(x:int=4, y:float=3, z:float=2)
88  my_func(x:int=5, y:float=3, z:float=2)
89  my_func(x:int=6, y:float=3, z:float=2)
90  """
91  )
92  with capture:
93  a, b, c = (
94  np.array([[1, 2, 3], [4, 5, 6]], order="F"),
95  np.array([[2], [3]]),
96  2,
97  )
98  assert np.allclose(f(a, b, c), a * b * c)
99  assert (
100  capture
101  == """
102  my_func(x:int=1, y:float=2, z:float=2)
103  my_func(x:int=2, y:float=2, z:float=2)
104  my_func(x:int=3, y:float=2, z:float=2)
105  my_func(x:int=4, y:float=3, z:float=2)
106  my_func(x:int=5, y:float=3, z:float=2)
107  my_func(x:int=6, y:float=3, z:float=2)
108  """
109  )
110  with capture:
111  a, b, c = np.array([[1, 2, 3], [4, 5, 6]])[::, ::2], np.array([[2], [3]]), 2
112  assert np.allclose(f(a, b, c), a * b * c)
113  assert (
114  capture
115  == """
116  my_func(x:int=1, y:float=2, z:float=2)
117  my_func(x:int=3, y:float=2, z:float=2)
118  my_func(x:int=4, y:float=3, z:float=2)
119  my_func(x:int=6, y:float=3, z:float=2)
120  """
121  )
122  with capture:
123  a, b, c = (
124  np.array([[1, 2, 3], [4, 5, 6]], order="F")[::, ::2],
125  np.array([[2], [3]]),
126  2,
127  )
128  assert np.allclose(f(a, b, c), a * b * c)
129  assert (
130  capture
131  == """
132  my_func(x:int=1, y:float=2, z:float=2)
133  my_func(x:int=3, y:float=2, z:float=2)
134  my_func(x:int=4, y:float=3, z:float=2)
135  my_func(x:int=6, y:float=3, z:float=2)
136  """
137  )
138 
139 
140 def test_type_selection():
141  assert m.selective_func(np.array([1], dtype=np.int32)) == "Int branch taken."
142  assert m.selective_func(np.array([1.0], dtype=np.float32)) == "Float branch taken."
143  assert (
144  m.selective_func(np.array([1.0j], dtype=np.complex64))
145  == "Complex float branch taken."
146  )
147 
148 
149 def test_docs(doc):
150  assert (
151  doc(m.vectorized_func)
152  == """
153  vectorized_func(arg0: numpy.ndarray[numpy.int32], arg1: numpy.ndarray[numpy.float32], arg2: numpy.ndarray[numpy.float64]) -> object
154  """
155  )
156 
157 
158 def test_trivial_broadcasting():
159  trivial, vectorized_is_trivial = m.trivial, m.vectorized_is_trivial
160 
161  assert vectorized_is_trivial(1, 2, 3) == trivial.c_trivial
162  assert vectorized_is_trivial(np.array(1), np.array(2), 3) == trivial.c_trivial
163  assert (
164  vectorized_is_trivial(np.array([1, 3]), np.array([2, 4]), 3)
165  == trivial.c_trivial
166  )
167  assert trivial.c_trivial == vectorized_is_trivial(
168  np.array([[1, 3, 5], [7, 9, 11]]), np.array([[2, 4, 6], [8, 10, 12]]), 3
169  )
170  assert (
171  vectorized_is_trivial(np.array([[1, 2, 3], [4, 5, 6]]), np.array([2, 3, 4]), 2)
172  == trivial.non_trivial
173  )
174  assert (
175  vectorized_is_trivial(np.array([[1, 2, 3], [4, 5, 6]]), np.array([[2], [3]]), 2)
176  == trivial.non_trivial
177  )
178  z1 = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype="int32")
179  z2 = np.array(z1, dtype="float32")
180  z3 = np.array(z1, dtype="float64")
181  assert vectorized_is_trivial(z1, z2, z3) == trivial.c_trivial
182  assert vectorized_is_trivial(1, z2, z3) == trivial.c_trivial
183  assert vectorized_is_trivial(z1, 1, z3) == trivial.c_trivial
184  assert vectorized_is_trivial(z1, z2, 1) == trivial.c_trivial
185  assert vectorized_is_trivial(z1[::2, ::2], 1, 1) == trivial.non_trivial
186  assert vectorized_is_trivial(1, 1, z1[::2, ::2]) == trivial.c_trivial
187  assert vectorized_is_trivial(1, 1, z3[::2, ::2]) == trivial.non_trivial
188  assert vectorized_is_trivial(z1, 1, z3[1::4, 1::4]) == trivial.c_trivial
189 
190  y1 = np.array(z1, order="F")
191  y2 = np.array(y1)
192  y3 = np.array(y1)
193  assert vectorized_is_trivial(y1, y2, y3) == trivial.f_trivial
194  assert vectorized_is_trivial(y1, 1, 1) == trivial.f_trivial
195  assert vectorized_is_trivial(1, y2, 1) == trivial.f_trivial
196  assert vectorized_is_trivial(1, 1, y3) == trivial.f_trivial
197  assert vectorized_is_trivial(y1, z2, 1) == trivial.non_trivial
198  assert vectorized_is_trivial(z1[1::4, 1::4], y2, 1) == trivial.f_trivial
199  assert vectorized_is_trivial(y1[1::4, 1::4], z2, 1) == trivial.c_trivial
200 
201  assert m.vectorized_func(z1, z2, z3).flags.c_contiguous
202  assert m.vectorized_func(y1, y2, y3).flags.f_contiguous
203  assert m.vectorized_func(z1, 1, 1).flags.c_contiguous
204  assert m.vectorized_func(1, y2, 1).flags.f_contiguous
205  assert m.vectorized_func(z1[1::4, 1::4], y2, 1).flags.f_contiguous
206  assert m.vectorized_func(y1[1::4, 1::4], z2, 1).flags.c_contiguous
207 
208 
209 def test_passthrough_arguments(doc):
210  assert doc(m.vec_passthrough) == (
211  "vec_passthrough("
212  + ", ".join(
213  [
214  "arg0: float",
215  "arg1: numpy.ndarray[numpy.float64]",
216  "arg2: numpy.ndarray[numpy.float64]",
217  "arg3: numpy.ndarray[numpy.int32]",
218  "arg4: int",
219  "arg5: m.numpy_vectorize.NonPODClass",
220  "arg6: numpy.ndarray[numpy.float64]",
221  ]
222  )
223  + ") -> object"
224  )
225 
226  b = np.array([[10, 20, 30]], dtype="float64")
227  c = np.array([100, 200]) # NOT a vectorized argument
228  d = np.array([[1000], [2000], [3000]], dtype="int")
229  g = np.array([[1000000, 2000000, 3000000]], dtype="int") # requires casting
230  assert np.all(
231  m.vec_passthrough(1, b, c, d, 10000, m.NonPODClass(100000), g)
232  == np.array(
233  [
234  [1111111, 2111121, 3111131],
235  [1112111, 2112121, 3112131],
236  [1113111, 2113121, 3113131],
237  ]
238  )
239  )
240 
241 
242 def test_method_vectorization():
243  o = m.VectorizeTestClass(3)
244  x = np.array([1, 2], dtype="int")
245  y = np.array([[10], [20]], dtype="float32")
246  assert np.all(o.method(x, y) == [[14, 15], [24, 25]])
247 
248 
249 def test_array_collapse():
250  assert not isinstance(m.vectorized_func(1, 2, 3), np.ndarray)
251  assert not isinstance(m.vectorized_func(np.array(1), 2, 3), np.ndarray)
252  z = m.vectorized_func([1], 2, 3)
253  assert isinstance(z, np.ndarray)
254  assert z.shape == (1,)
255  z = m.vectorized_func(1, [[[2]]], 3)
256  assert isinstance(z, np.ndarray)
257  assert z.shape == (1, 1, 1)
258 
259 
260 def test_vectorized_noreturn():
261  x = m.NonPODClass(0)
262  assert x.value == 0
263  m.add_to(x, [1, 2, 3, 4])
264  assert x.value == 10
265  m.add_to(x, 1)
266  assert x.value == 11
267  m.add_to(x, [[1, 1], [2, 3]])
268  assert x.value == 18
test_numpy_vectorize.test_trivial_broadcasting
def test_trivial_broadcasting()
Definition: test_numpy_vectorize.py:158
test_numpy_vectorize.test_passthrough_arguments
def test_passthrough_arguments(doc)
Definition: test_numpy_vectorize.py:209
test_numpy_vectorize.test_vectorized_noreturn
def test_vectorized_noreturn()
Definition: test_numpy_vectorize.py:260
test_numpy_vectorize.test_docs
def test_docs(doc)
Definition: test_numpy_vectorize.py:149
doc
Annotation for documentation.
Definition: attr.h:45
isinstance
bool isinstance(handle obj)
Definition: pytypes.h:842
tree::f
Point2(* f)(const Point3 &, OptionalJacobian< 2, 3 >)
Definition: testExpression.cpp:218
test_numpy_vectorize.test_array_collapse
def test_array_collapse()
Definition: test_numpy_vectorize.py:249
test_numpy_vectorize.test_type_selection
def test_type_selection()
Definition: test_numpy_vectorize.py:140
test_numpy_vectorize.test_vectorize
def test_vectorize(capture)
Definition: test_numpy_vectorize.py:10
test_numpy_vectorize.test_method_vectorization
def test_method_vectorization()
Definition: test_numpy_vectorize.py:242

gtsam
Author(s):
autogenerated on Wed May 28 2025 03:06:00