test_twodim_base.py 18 KB

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