test_twodim_base.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559
  1. """Test functions for matrix module
  2. """
  3. import pytest
  4. import numpy as np
  5. from numpy import (
  6. add,
  7. arange,
  8. array,
  9. diag,
  10. eye,
  11. fliplr,
  12. flipud,
  13. histogram2d,
  14. mask_indices,
  15. ones,
  16. tri,
  17. tril_indices,
  18. tril_indices_from,
  19. triu_indices,
  20. triu_indices_from,
  21. vander,
  22. zeros,
  23. )
  24. from numpy.testing import (
  25. assert_,
  26. assert_array_almost_equal,
  27. assert_array_equal,
  28. assert_array_max_ulp,
  29. assert_equal,
  30. assert_raises,
  31. )
  32. def get_mat(n):
  33. data = arange(n)
  34. data = add.outer(data, data)
  35. return data
  36. class TestEye:
  37. def test_basic(self):
  38. assert_equal(eye(4),
  39. array([[1, 0, 0, 0],
  40. [0, 1, 0, 0],
  41. [0, 0, 1, 0],
  42. [0, 0, 0, 1]]))
  43. assert_equal(eye(4, dtype='f'),
  44. array([[1, 0, 0, 0],
  45. [0, 1, 0, 0],
  46. [0, 0, 1, 0],
  47. [0, 0, 0, 1]], 'f'))
  48. assert_equal(eye(3) == 1,
  49. eye(3, dtype=bool))
  50. def test_uint64(self):
  51. # Regression test for gh-9982
  52. assert_equal(eye(np.uint64(2), dtype=int), array([[1, 0], [0, 1]]))
  53. assert_equal(eye(np.uint64(2), M=np.uint64(4), k=np.uint64(1)),
  54. array([[0, 1, 0, 0], [0, 0, 1, 0]]))
  55. def test_diag(self):
  56. assert_equal(eye(4, k=1),
  57. array([[0, 1, 0, 0],
  58. [0, 0, 1, 0],
  59. [0, 0, 0, 1],
  60. [0, 0, 0, 0]]))
  61. assert_equal(eye(4, k=-1),
  62. array([[0, 0, 0, 0],
  63. [1, 0, 0, 0],
  64. [0, 1, 0, 0],
  65. [0, 0, 1, 0]]))
  66. def test_2d(self):
  67. assert_equal(eye(4, 3),
  68. array([[1, 0, 0],
  69. [0, 1, 0],
  70. [0, 0, 1],
  71. [0, 0, 0]]))
  72. assert_equal(eye(3, 4),
  73. array([[1, 0, 0, 0],
  74. [0, 1, 0, 0],
  75. [0, 0, 1, 0]]))
  76. def test_diag2d(self):
  77. assert_equal(eye(3, 4, k=2),
  78. array([[0, 0, 1, 0],
  79. [0, 0, 0, 1],
  80. [0, 0, 0, 0]]))
  81. assert_equal(eye(4, 3, k=-2),
  82. array([[0, 0, 0],
  83. [0, 0, 0],
  84. [1, 0, 0],
  85. [0, 1, 0]]))
  86. def test_eye_bounds(self):
  87. assert_equal(eye(2, 2, 1), [[0, 1], [0, 0]])
  88. assert_equal(eye(2, 2, -1), [[0, 0], [1, 0]])
  89. assert_equal(eye(2, 2, 2), [[0, 0], [0, 0]])
  90. assert_equal(eye(2, 2, -2), [[0, 0], [0, 0]])
  91. assert_equal(eye(3, 2, 2), [[0, 0], [0, 0], [0, 0]])
  92. assert_equal(eye(3, 2, 1), [[0, 1], [0, 0], [0, 0]])
  93. assert_equal(eye(3, 2, -1), [[0, 0], [1, 0], [0, 1]])
  94. assert_equal(eye(3, 2, -2), [[0, 0], [0, 0], [1, 0]])
  95. assert_equal(eye(3, 2, -3), [[0, 0], [0, 0], [0, 0]])
  96. def test_strings(self):
  97. assert_equal(eye(2, 2, dtype='S3'),
  98. [[b'1', b''], [b'', b'1']])
  99. def test_bool(self):
  100. assert_equal(eye(2, 2, dtype=bool), [[True, False], [False, True]])
  101. def test_order(self):
  102. mat_c = eye(4, 3, k=-1)
  103. mat_f = eye(4, 3, k=-1, order='F')
  104. assert_equal(mat_c, mat_f)
  105. assert mat_c.flags.c_contiguous
  106. assert not mat_c.flags.f_contiguous
  107. assert not mat_f.flags.c_contiguous
  108. assert mat_f.flags.f_contiguous
  109. class TestDiag:
  110. def test_vector(self):
  111. vals = (100 * arange(5)).astype('l')
  112. b = zeros((5, 5))
  113. for k in range(5):
  114. b[k, k] = vals[k]
  115. assert_equal(diag(vals), b)
  116. b = zeros((7, 7))
  117. c = b.copy()
  118. for k in range(5):
  119. b[k, k + 2] = vals[k]
  120. c[k + 2, k] = vals[k]
  121. assert_equal(diag(vals, k=2), b)
  122. assert_equal(diag(vals, k=-2), c)
  123. def test_matrix(self, vals=None):
  124. if vals is None:
  125. vals = (100 * get_mat(5) + 1).astype('l')
  126. b = zeros((5,))
  127. for k in range(5):
  128. b[k] = vals[k, k]
  129. assert_equal(diag(vals), b)
  130. b = b * 0
  131. for k in range(3):
  132. b[k] = vals[k, k + 2]
  133. assert_equal(diag(vals, 2), b[:3])
  134. for k in range(3):
  135. b[k] = vals[k + 2, k]
  136. assert_equal(diag(vals, -2), b[:3])
  137. def test_fortran_order(self):
  138. vals = array((100 * get_mat(5) + 1), order='F', dtype='l')
  139. self.test_matrix(vals)
  140. def test_diag_bounds(self):
  141. A = [[1, 2], [3, 4], [5, 6]]
  142. assert_equal(diag(A, k=2), [])
  143. assert_equal(diag(A, k=1), [2])
  144. assert_equal(diag(A, k=0), [1, 4])
  145. assert_equal(diag(A, k=-1), [3, 6])
  146. assert_equal(diag(A, k=-2), [5])
  147. assert_equal(diag(A, k=-3), [])
  148. def test_failure(self):
  149. assert_raises(ValueError, diag, [[[1]]])
  150. class TestFliplr:
  151. def test_basic(self):
  152. assert_raises(ValueError, fliplr, ones(4))
  153. a = get_mat(4)
  154. b = a[:, ::-1]
  155. assert_equal(fliplr(a), b)
  156. a = [[0, 1, 2],
  157. [3, 4, 5]]
  158. b = [[2, 1, 0],
  159. [5, 4, 3]]
  160. assert_equal(fliplr(a), b)
  161. class TestFlipud:
  162. def test_basic(self):
  163. a = get_mat(4)
  164. b = a[::-1, :]
  165. assert_equal(flipud(a), b)
  166. a = [[0, 1, 2],
  167. [3, 4, 5]]
  168. b = [[3, 4, 5],
  169. [0, 1, 2]]
  170. assert_equal(flipud(a), b)
  171. class TestHistogram2d:
  172. def test_simple(self):
  173. x = array(
  174. [0.41702200, 0.72032449, 1.1437481e-4, 0.302332573, 0.146755891])
  175. y = array(
  176. [0.09233859, 0.18626021, 0.34556073, 0.39676747, 0.53881673])
  177. xedges = np.linspace(0, 1, 10)
  178. yedges = np.linspace(0, 1, 10)
  179. H = histogram2d(x, y, (xedges, yedges))[0]
  180. answer = array(
  181. [[0, 0, 0, 1, 0, 0, 0, 0, 0],
  182. [0, 0, 0, 0, 0, 0, 1, 0, 0],
  183. [0, 0, 0, 0, 0, 0, 0, 0, 0],
  184. [1, 0, 1, 0, 0, 0, 0, 0, 0],
  185. [0, 1, 0, 0, 0, 0, 0, 0, 0],
  186. [0, 0, 0, 0, 0, 0, 0, 0, 0],
  187. [0, 0, 0, 0, 0, 0, 0, 0, 0],
  188. [0, 0, 0, 0, 0, 0, 0, 0, 0],
  189. [0, 0, 0, 0, 0, 0, 0, 0, 0]])
  190. assert_array_equal(H.T, answer)
  191. H = histogram2d(x, y, xedges)[0]
  192. assert_array_equal(H.T, answer)
  193. H, xedges, yedges = histogram2d(list(range(10)), list(range(10)))
  194. assert_array_equal(H, eye(10, 10))
  195. assert_array_equal(xedges, np.linspace(0, 9, 11))
  196. assert_array_equal(yedges, np.linspace(0, 9, 11))
  197. def test_asym(self):
  198. x = array([1, 1, 2, 3, 4, 4, 4, 5])
  199. y = array([1, 3, 2, 0, 1, 2, 3, 4])
  200. H, xed, yed = histogram2d(
  201. x, y, (6, 5), range=[[0, 6], [0, 5]], density=True)
  202. answer = array(
  203. [[0., 0, 0, 0, 0],
  204. [0, 1, 0, 1, 0],
  205. [0, 0, 1, 0, 0],
  206. [1, 0, 0, 0, 0],
  207. [0, 1, 1, 1, 0],
  208. [0, 0, 0, 0, 1]])
  209. assert_array_almost_equal(H, answer / 8., 3)
  210. assert_array_equal(xed, np.linspace(0, 6, 7))
  211. assert_array_equal(yed, np.linspace(0, 5, 6))
  212. def test_density(self):
  213. x = array([1, 2, 3, 1, 2, 3, 1, 2, 3])
  214. y = array([1, 1, 1, 2, 2, 2, 3, 3, 3])
  215. H, xed, yed = histogram2d(
  216. x, y, [[1, 2, 3, 5], [1, 2, 3, 5]], density=True)
  217. answer = array([[1, 1, .5],
  218. [1, 1, .5],
  219. [.5, .5, .25]]) / 9.
  220. assert_array_almost_equal(H, answer, 3)
  221. def test_all_outliers(self):
  222. r = np.random.rand(100) + 1. + 1e6 # histogramdd rounds by decimal=6
  223. H, xed, yed = histogram2d(r, r, (4, 5), range=([0, 1], [0, 1]))
  224. assert_array_equal(H, 0)
  225. def test_empty(self):
  226. a, edge1, edge2 = histogram2d([], [], bins=([0, 1], [0, 1]))
  227. assert_array_max_ulp(a, array([[0.]]))
  228. a, edge1, edge2 = histogram2d([], [], bins=4)
  229. assert_array_max_ulp(a, np.zeros((4, 4)))
  230. def test_binparameter_combination(self):
  231. x = array(
  232. [0, 0.09207008, 0.64575234, 0.12875982, 0.47390599,
  233. 0.59944483, 1])
  234. y = array(
  235. [0, 0.14344267, 0.48988575, 0.30558665, 0.44700682,
  236. 0.15886423, 1])
  237. edges = (0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1)
  238. H, xe, ye = histogram2d(x, y, (edges, 4))
  239. answer = array(
  240. [[2., 0., 0., 0.],
  241. [0., 1., 0., 0.],
  242. [0., 0., 0., 0.],
  243. [0., 0., 0., 0.],
  244. [0., 1., 0., 0.],
  245. [1., 0., 0., 0.],
  246. [0., 1., 0., 0.],
  247. [0., 0., 0., 0.],
  248. [0., 0., 0., 0.],
  249. [0., 0., 0., 1.]])
  250. assert_array_equal(H, answer)
  251. assert_array_equal(ye, array([0., 0.25, 0.5, 0.75, 1]))
  252. H, xe, ye = histogram2d(x, y, (4, edges))
  253. answer = array(
  254. [[1., 1., 0., 1., 0., 0., 0., 0., 0., 0.],
  255. [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
  256. [0., 1., 0., 0., 1., 0., 0., 0., 0., 0.],
  257. [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]])
  258. assert_array_equal(H, answer)
  259. assert_array_equal(xe, array([0., 0.25, 0.5, 0.75, 1]))
  260. def test_dispatch(self):
  261. class ShouldDispatch:
  262. def __array_function__(self, function, types, args, kwargs):
  263. return types, args, kwargs
  264. xy = [1, 2]
  265. s_d = ShouldDispatch()
  266. r = histogram2d(s_d, xy)
  267. # Cannot use assert_equal since that dispatches...
  268. assert_(r == ((ShouldDispatch,), (s_d, xy), {}))
  269. r = histogram2d(xy, s_d)
  270. assert_(r == ((ShouldDispatch,), (xy, s_d), {}))
  271. r = histogram2d(xy, xy, bins=s_d)
  272. assert_(r, ((ShouldDispatch,), (xy, xy), {'bins': s_d}))
  273. r = histogram2d(xy, xy, bins=[s_d, 5])
  274. assert_(r, ((ShouldDispatch,), (xy, xy), {'bins': [s_d, 5]}))
  275. assert_raises(Exception, histogram2d, xy, xy, bins=[s_d])
  276. r = histogram2d(xy, xy, weights=s_d)
  277. assert_(r, ((ShouldDispatch,), (xy, xy), {'weights': s_d}))
  278. @pytest.mark.parametrize(("x_len", "y_len"), [(10, 11), (20, 19)])
  279. def test_bad_length(self, x_len, y_len):
  280. x, y = np.ones(x_len), np.ones(y_len)
  281. with pytest.raises(ValueError,
  282. match='x and y must have the same length.'):
  283. histogram2d(x, y)
  284. class TestTri:
  285. def test_dtype(self):
  286. out = array([[1, 0, 0],
  287. [1, 1, 0],
  288. [1, 1, 1]])
  289. assert_array_equal(tri(3), out)
  290. assert_array_equal(tri(3, dtype=bool), out.astype(bool))
  291. def test_tril_triu_ndim2():
  292. for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']:
  293. a = np.ones((2, 2), dtype=dtype)
  294. b = np.tril(a)
  295. c = np.triu(a)
  296. assert_array_equal(b, [[1, 0], [1, 1]])
  297. assert_array_equal(c, b.T)
  298. # should return the same dtype as the original array
  299. assert_equal(b.dtype, a.dtype)
  300. assert_equal(c.dtype, a.dtype)
  301. def test_tril_triu_ndim3():
  302. for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']:
  303. a = np.array([
  304. [[1, 1], [1, 1]],
  305. [[1, 1], [1, 0]],
  306. [[1, 1], [0, 0]],
  307. ], dtype=dtype)
  308. a_tril_desired = np.array([
  309. [[1, 0], [1, 1]],
  310. [[1, 0], [1, 0]],
  311. [[1, 0], [0, 0]],
  312. ], dtype=dtype)
  313. a_triu_desired = np.array([
  314. [[1, 1], [0, 1]],
  315. [[1, 1], [0, 0]],
  316. [[1, 1], [0, 0]],
  317. ], dtype=dtype)
  318. a_triu_observed = np.triu(a)
  319. a_tril_observed = np.tril(a)
  320. assert_array_equal(a_triu_observed, a_triu_desired)
  321. assert_array_equal(a_tril_observed, a_tril_desired)
  322. assert_equal(a_triu_observed.dtype, a.dtype)
  323. assert_equal(a_tril_observed.dtype, a.dtype)
  324. def test_tril_triu_with_inf():
  325. # Issue 4859
  326. arr = np.array([[1, 1, np.inf],
  327. [1, 1, 1],
  328. [np.inf, 1, 1]])
  329. out_tril = np.array([[1, 0, 0],
  330. [1, 1, 0],
  331. [np.inf, 1, 1]])
  332. out_triu = out_tril.T
  333. assert_array_equal(np.triu(arr), out_triu)
  334. assert_array_equal(np.tril(arr), out_tril)
  335. def test_tril_triu_dtype():
  336. # Issue 4916
  337. # tril and triu should return the same dtype as input
  338. for c in np.typecodes['All']:
  339. if c == 'V':
  340. continue
  341. arr = np.zeros((3, 3), dtype=c)
  342. assert_equal(np.triu(arr).dtype, arr.dtype)
  343. assert_equal(np.tril(arr).dtype, arr.dtype)
  344. # check special cases
  345. arr = np.array([['2001-01-01T12:00', '2002-02-03T13:56'],
  346. ['2004-01-01T12:00', '2003-01-03T13:45']],
  347. dtype='datetime64')
  348. assert_equal(np.triu(arr).dtype, arr.dtype)
  349. assert_equal(np.tril(arr).dtype, arr.dtype)
  350. arr = np.zeros((3, 3), dtype='f4,f4')
  351. assert_equal(np.triu(arr).dtype, arr.dtype)
  352. assert_equal(np.tril(arr).dtype, arr.dtype)
  353. def test_mask_indices():
  354. # simple test without offset
  355. iu = mask_indices(3, np.triu)
  356. a = np.arange(9).reshape(3, 3)
  357. assert_array_equal(a[iu], array([0, 1, 2, 4, 5, 8]))
  358. # Now with an offset
  359. iu1 = mask_indices(3, np.triu, 1)
  360. assert_array_equal(a[iu1], array([1, 2, 5]))
  361. def test_tril_indices():
  362. # indices without and with offset
  363. il1 = tril_indices(4)
  364. il2 = tril_indices(4, k=2)
  365. il3 = tril_indices(4, m=5)
  366. il4 = tril_indices(4, k=2, m=5)
  367. a = np.array([[1, 2, 3, 4],
  368. [5, 6, 7, 8],
  369. [9, 10, 11, 12],
  370. [13, 14, 15, 16]])
  371. b = np.arange(1, 21).reshape(4, 5)
  372. # indexing:
  373. assert_array_equal(a[il1],
  374. array([1, 5, 6, 9, 10, 11, 13, 14, 15, 16]))
  375. assert_array_equal(b[il3],
  376. array([1, 6, 7, 11, 12, 13, 16, 17, 18, 19]))
  377. # And for assigning values:
  378. a[il1] = -1
  379. assert_array_equal(a,
  380. array([[-1, 2, 3, 4],
  381. [-1, -1, 7, 8],
  382. [-1, -1, -1, 12],
  383. [-1, -1, -1, -1]]))
  384. b[il3] = -1
  385. assert_array_equal(b,
  386. array([[-1, 2, 3, 4, 5],
  387. [-1, -1, 8, 9, 10],
  388. [-1, -1, -1, 14, 15],
  389. [-1, -1, -1, -1, 20]]))
  390. # These cover almost the whole array (two diagonals right of the main one):
  391. a[il2] = -10
  392. assert_array_equal(a,
  393. array([[-10, -10, -10, 4],
  394. [-10, -10, -10, -10],
  395. [-10, -10, -10, -10],
  396. [-10, -10, -10, -10]]))
  397. b[il4] = -10
  398. assert_array_equal(b,
  399. array([[-10, -10, -10, 4, 5],
  400. [-10, -10, -10, -10, 10],
  401. [-10, -10, -10, -10, -10],
  402. [-10, -10, -10, -10, -10]]))
  403. class TestTriuIndices:
  404. def test_triu_indices(self):
  405. iu1 = triu_indices(4)
  406. iu2 = triu_indices(4, k=2)
  407. iu3 = triu_indices(4, m=5)
  408. iu4 = triu_indices(4, k=2, m=5)
  409. a = np.array([[1, 2, 3, 4],
  410. [5, 6, 7, 8],
  411. [9, 10, 11, 12],
  412. [13, 14, 15, 16]])
  413. b = np.arange(1, 21).reshape(4, 5)
  414. # Both for indexing:
  415. assert_array_equal(a[iu1],
  416. array([1, 2, 3, 4, 6, 7, 8, 11, 12, 16]))
  417. assert_array_equal(b[iu3],
  418. array([1, 2, 3, 4, 5, 7, 8, 9,
  419. 10, 13, 14, 15, 19, 20]))
  420. # And for assigning values:
  421. a[iu1] = -1
  422. assert_array_equal(a,
  423. array([[-1, -1, -1, -1],
  424. [5, -1, -1, -1],
  425. [9, 10, -1, -1],
  426. [13, 14, 15, -1]]))
  427. b[iu3] = -1
  428. assert_array_equal(b,
  429. array([[-1, -1, -1, -1, -1],
  430. [6, -1, -1, -1, -1],
  431. [11, 12, -1, -1, -1],
  432. [16, 17, 18, -1, -1]]))
  433. # These cover almost the whole array (two diagonals right of the
  434. # main one):
  435. a[iu2] = -10
  436. assert_array_equal(a,
  437. array([[-1, -1, -10, -10],
  438. [5, -1, -1, -10],
  439. [9, 10, -1, -1],
  440. [13, 14, 15, -1]]))
  441. b[iu4] = -10
  442. assert_array_equal(b,
  443. array([[-1, -1, -10, -10, -10],
  444. [6, -1, -1, -10, -10],
  445. [11, 12, -1, -1, -10],
  446. [16, 17, 18, -1, -1]]))
  447. class TestTrilIndicesFrom:
  448. def test_exceptions(self):
  449. assert_raises(ValueError, tril_indices_from, np.ones((2,)))
  450. assert_raises(ValueError, tril_indices_from, np.ones((2, 2, 2)))
  451. # assert_raises(ValueError, tril_indices_from, np.ones((2, 3)))
  452. class TestTriuIndicesFrom:
  453. def test_exceptions(self):
  454. assert_raises(ValueError, triu_indices_from, np.ones((2,)))
  455. assert_raises(ValueError, triu_indices_from, np.ones((2, 2, 2)))
  456. # assert_raises(ValueError, triu_indices_from, np.ones((2, 3)))
  457. class TestVander:
  458. def test_basic(self):
  459. c = np.array([0, 1, -2, 3])
  460. v = vander(c)
  461. powers = np.array([[0, 0, 0, 0, 1],
  462. [1, 1, 1, 1, 1],
  463. [16, -8, 4, -2, 1],
  464. [81, 27, 9, 3, 1]])
  465. # Check default value of N:
  466. assert_array_equal(v, powers[:, 1:])
  467. # Check a range of N values, including 0 and 5 (greater than default)
  468. m = powers.shape[1]
  469. for n in range(6):
  470. v = vander(c, N=n)
  471. assert_array_equal(v, powers[:, m - n:m])
  472. def test_dtypes(self):
  473. c = array([11, -12, 13], dtype=np.int8)
  474. v = vander(c)
  475. expected = np.array([[121, 11, 1],
  476. [144, -12, 1],
  477. [169, 13, 1]])
  478. assert_array_equal(v, expected)
  479. c = array([1.0 + 1j, 1.0 - 1j])
  480. v = vander(c, N=3)
  481. expected = np.array([[2j, 1 + 1j, 1],
  482. [-2j, 1 - 1j, 1]])
  483. # The data is floating point, but the values are small integers,
  484. # so assert_array_equal *should* be safe here (rather than, say,
  485. # assert_array_almost_equal).
  486. assert_array_equal(v, expected)