test_legendre.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556
  1. """Tests for legendre module.
  2. """
  3. from functools import reduce
  4. import numpy as np
  5. import numpy.polynomial.legendre as leg
  6. from numpy.polynomial.polynomial import polyval
  7. from numpy.testing import (
  8. assert_almost_equal, assert_raises, assert_equal, assert_,
  9. )
  10. L0 = np.array([1])
  11. L1 = np.array([0, 1])
  12. L2 = np.array([-1, 0, 3])/2
  13. L3 = np.array([0, -3, 0, 5])/2
  14. L4 = np.array([3, 0, -30, 0, 35])/8
  15. L5 = np.array([0, 15, 0, -70, 0, 63])/8
  16. L6 = np.array([-5, 0, 105, 0, -315, 0, 231])/16
  17. L7 = np.array([0, -35, 0, 315, 0, -693, 0, 429])/16
  18. L8 = np.array([35, 0, -1260, 0, 6930, 0, -12012, 0, 6435])/128
  19. L9 = np.array([0, 315, 0, -4620, 0, 18018, 0, -25740, 0, 12155])/128
  20. Llist = [L0, L1, L2, L3, L4, L5, L6, L7, L8, L9]
  21. def trim(x):
  22. return leg.legtrim(x, tol=1e-6)
  23. class TestConstants:
  24. def test_legdomain(self):
  25. assert_equal(leg.legdomain, [-1, 1])
  26. def test_legzero(self):
  27. assert_equal(leg.legzero, [0])
  28. def test_legone(self):
  29. assert_equal(leg.legone, [1])
  30. def test_legx(self):
  31. assert_equal(leg.legx, [0, 1])
  32. class TestArithmetic:
  33. x = np.linspace(-1, 1, 100)
  34. def test_legadd(self):
  35. for i in range(5):
  36. for j in range(5):
  37. msg = f"At i={i}, j={j}"
  38. tgt = np.zeros(max(i, j) + 1)
  39. tgt[i] += 1
  40. tgt[j] += 1
  41. res = leg.legadd([0]*i + [1], [0]*j + [1])
  42. assert_equal(trim(res), trim(tgt), err_msg=msg)
  43. def test_legsub(self):
  44. for i in range(5):
  45. for j in range(5):
  46. msg = f"At i={i}, j={j}"
  47. tgt = np.zeros(max(i, j) + 1)
  48. tgt[i] += 1
  49. tgt[j] -= 1
  50. res = leg.legsub([0]*i + [1], [0]*j + [1])
  51. assert_equal(trim(res), trim(tgt), err_msg=msg)
  52. def test_legmulx(self):
  53. assert_equal(leg.legmulx([0]), [0])
  54. assert_equal(leg.legmulx([1]), [0, 1])
  55. for i in range(1, 5):
  56. tmp = 2*i + 1
  57. ser = [0]*i + [1]
  58. tgt = [0]*(i - 1) + [i/tmp, 0, (i + 1)/tmp]
  59. assert_equal(leg.legmulx(ser), tgt)
  60. def test_legmul(self):
  61. # check values of result
  62. for i in range(5):
  63. pol1 = [0]*i + [1]
  64. val1 = leg.legval(self.x, pol1)
  65. for j in range(5):
  66. msg = f"At i={i}, j={j}"
  67. pol2 = [0]*j + [1]
  68. val2 = leg.legval(self.x, pol2)
  69. pol3 = leg.legmul(pol1, pol2)
  70. val3 = leg.legval(self.x, pol3)
  71. assert_(len(pol3) == i + j + 1, msg)
  72. assert_almost_equal(val3, val1*val2, err_msg=msg)
  73. def test_legdiv(self):
  74. for i in range(5):
  75. for j in range(5):
  76. msg = f"At i={i}, j={j}"
  77. ci = [0]*i + [1]
  78. cj = [0]*j + [1]
  79. tgt = leg.legadd(ci, cj)
  80. quo, rem = leg.legdiv(tgt, ci)
  81. res = leg.legadd(leg.legmul(quo, ci), rem)
  82. assert_equal(trim(res), trim(tgt), err_msg=msg)
  83. def test_legpow(self):
  84. for i in range(5):
  85. for j in range(5):
  86. msg = f"At i={i}, j={j}"
  87. c = np.arange(i + 1)
  88. tgt = reduce(leg.legmul, [c]*j, np.array([1]))
  89. res = leg.legpow(c, j)
  90. assert_equal(trim(res), trim(tgt), err_msg=msg)
  91. class TestEvaluation:
  92. # coefficients of 1 + 2*x + 3*x**2
  93. c1d = np.array([2., 2., 2.])
  94. c2d = np.einsum('i,j->ij', c1d, c1d)
  95. c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d)
  96. # some random values in [-1, 1)
  97. x = np.random.random((3, 5))*2 - 1
  98. y = polyval(x, [1., 2., 3.])
  99. def test_legval(self):
  100. #check empty input
  101. assert_equal(leg.legval([], [1]).size, 0)
  102. #check normal input)
  103. x = np.linspace(-1, 1)
  104. y = [polyval(x, c) for c in Llist]
  105. for i in range(10):
  106. msg = f"At i={i}"
  107. tgt = y[i]
  108. res = leg.legval(x, [0]*i + [1])
  109. assert_almost_equal(res, tgt, err_msg=msg)
  110. #check that shape is preserved
  111. for i in range(3):
  112. dims = [2]*i
  113. x = np.zeros(dims)
  114. assert_equal(leg.legval(x, [1]).shape, dims)
  115. assert_equal(leg.legval(x, [1, 0]).shape, dims)
  116. assert_equal(leg.legval(x, [1, 0, 0]).shape, dims)
  117. def test_legval2d(self):
  118. x1, x2, x3 = self.x
  119. y1, y2, y3 = self.y
  120. #test exceptions
  121. assert_raises(ValueError, leg.legval2d, x1, x2[:2], self.c2d)
  122. #test values
  123. tgt = y1*y2
  124. res = leg.legval2d(x1, x2, self.c2d)
  125. assert_almost_equal(res, tgt)
  126. #test shape
  127. z = np.ones((2, 3))
  128. res = leg.legval2d(z, z, self.c2d)
  129. assert_(res.shape == (2, 3))
  130. def test_legval3d(self):
  131. x1, x2, x3 = self.x
  132. y1, y2, y3 = self.y
  133. #test exceptions
  134. assert_raises(ValueError, leg.legval3d, x1, x2, x3[:2], self.c3d)
  135. #test values
  136. tgt = y1*y2*y3
  137. res = leg.legval3d(x1, x2, x3, self.c3d)
  138. assert_almost_equal(res, tgt)
  139. #test shape
  140. z = np.ones((2, 3))
  141. res = leg.legval3d(z, z, z, self.c3d)
  142. assert_(res.shape == (2, 3))
  143. def test_leggrid2d(self):
  144. x1, x2, x3 = self.x
  145. y1, y2, y3 = self.y
  146. #test values
  147. tgt = np.einsum('i,j->ij', y1, y2)
  148. res = leg.leggrid2d(x1, x2, self.c2d)
  149. assert_almost_equal(res, tgt)
  150. #test shape
  151. z = np.ones((2, 3))
  152. res = leg.leggrid2d(z, z, self.c2d)
  153. assert_(res.shape == (2, 3)*2)
  154. def test_leggrid3d(self):
  155. x1, x2, x3 = self.x
  156. y1, y2, y3 = self.y
  157. #test values
  158. tgt = np.einsum('i,j,k->ijk', y1, y2, y3)
  159. res = leg.leggrid3d(x1, x2, x3, self.c3d)
  160. assert_almost_equal(res, tgt)
  161. #test shape
  162. z = np.ones((2, 3))
  163. res = leg.leggrid3d(z, z, z, self.c3d)
  164. assert_(res.shape == (2, 3)*3)
  165. class TestIntegral:
  166. def test_legint(self):
  167. # check exceptions
  168. assert_raises(TypeError, leg.legint, [0], .5)
  169. assert_raises(ValueError, leg.legint, [0], -1)
  170. assert_raises(ValueError, leg.legint, [0], 1, [0, 0])
  171. assert_raises(ValueError, leg.legint, [0], lbnd=[0])
  172. assert_raises(ValueError, leg.legint, [0], scl=[0])
  173. assert_raises(TypeError, leg.legint, [0], axis=.5)
  174. # test integration of zero polynomial
  175. for i in range(2, 5):
  176. k = [0]*(i - 2) + [1]
  177. res = leg.legint([0], m=i, k=k)
  178. assert_almost_equal(res, [0, 1])
  179. # check single integration with integration constant
  180. for i in range(5):
  181. scl = i + 1
  182. pol = [0]*i + [1]
  183. tgt = [i] + [0]*i + [1/scl]
  184. legpol = leg.poly2leg(pol)
  185. legint = leg.legint(legpol, m=1, k=[i])
  186. res = leg.leg2poly(legint)
  187. assert_almost_equal(trim(res), trim(tgt))
  188. # check single integration with integration constant and lbnd
  189. for i in range(5):
  190. scl = i + 1
  191. pol = [0]*i + [1]
  192. legpol = leg.poly2leg(pol)
  193. legint = leg.legint(legpol, m=1, k=[i], lbnd=-1)
  194. assert_almost_equal(leg.legval(-1, legint), i)
  195. # check single integration with integration constant and scaling
  196. for i in range(5):
  197. scl = i + 1
  198. pol = [0]*i + [1]
  199. tgt = [i] + [0]*i + [2/scl]
  200. legpol = leg.poly2leg(pol)
  201. legint = leg.legint(legpol, m=1, k=[i], scl=2)
  202. res = leg.leg2poly(legint)
  203. assert_almost_equal(trim(res), trim(tgt))
  204. # check multiple integrations with default k
  205. for i in range(5):
  206. for j in range(2, 5):
  207. pol = [0]*i + [1]
  208. tgt = pol[:]
  209. for k in range(j):
  210. tgt = leg.legint(tgt, m=1)
  211. res = leg.legint(pol, m=j)
  212. assert_almost_equal(trim(res), trim(tgt))
  213. # check multiple integrations with defined k
  214. for i in range(5):
  215. for j in range(2, 5):
  216. pol = [0]*i + [1]
  217. tgt = pol[:]
  218. for k in range(j):
  219. tgt = leg.legint(tgt, m=1, k=[k])
  220. res = leg.legint(pol, m=j, k=list(range(j)))
  221. assert_almost_equal(trim(res), trim(tgt))
  222. # check multiple integrations with lbnd
  223. for i in range(5):
  224. for j in range(2, 5):
  225. pol = [0]*i + [1]
  226. tgt = pol[:]
  227. for k in range(j):
  228. tgt = leg.legint(tgt, m=1, k=[k], lbnd=-1)
  229. res = leg.legint(pol, m=j, k=list(range(j)), lbnd=-1)
  230. assert_almost_equal(trim(res), trim(tgt))
  231. # check multiple integrations with scaling
  232. for i in range(5):
  233. for j in range(2, 5):
  234. pol = [0]*i + [1]
  235. tgt = pol[:]
  236. for k in range(j):
  237. tgt = leg.legint(tgt, m=1, k=[k], scl=2)
  238. res = leg.legint(pol, m=j, k=list(range(j)), scl=2)
  239. assert_almost_equal(trim(res), trim(tgt))
  240. def test_legint_axis(self):
  241. # check that axis keyword works
  242. c2d = np.random.random((3, 4))
  243. tgt = np.vstack([leg.legint(c) for c in c2d.T]).T
  244. res = leg.legint(c2d, axis=0)
  245. assert_almost_equal(res, tgt)
  246. tgt = np.vstack([leg.legint(c) for c in c2d])
  247. res = leg.legint(c2d, axis=1)
  248. assert_almost_equal(res, tgt)
  249. tgt = np.vstack([leg.legint(c, k=3) for c in c2d])
  250. res = leg.legint(c2d, k=3, axis=1)
  251. assert_almost_equal(res, tgt)
  252. class TestDerivative:
  253. def test_legder(self):
  254. # check exceptions
  255. assert_raises(TypeError, leg.legder, [0], .5)
  256. assert_raises(ValueError, leg.legder, [0], -1)
  257. # check that zeroth derivative does nothing
  258. for i in range(5):
  259. tgt = [0]*i + [1]
  260. res = leg.legder(tgt, m=0)
  261. assert_equal(trim(res), trim(tgt))
  262. # check that derivation is the inverse of integration
  263. for i in range(5):
  264. for j in range(2, 5):
  265. tgt = [0]*i + [1]
  266. res = leg.legder(leg.legint(tgt, m=j), m=j)
  267. assert_almost_equal(trim(res), trim(tgt))
  268. # check derivation with scaling
  269. for i in range(5):
  270. for j in range(2, 5):
  271. tgt = [0]*i + [1]
  272. res = leg.legder(leg.legint(tgt, m=j, scl=2), m=j, scl=.5)
  273. assert_almost_equal(trim(res), trim(tgt))
  274. def test_legder_axis(self):
  275. # check that axis keyword works
  276. c2d = np.random.random((3, 4))
  277. tgt = np.vstack([leg.legder(c) for c in c2d.T]).T
  278. res = leg.legder(c2d, axis=0)
  279. assert_almost_equal(res, tgt)
  280. tgt = np.vstack([leg.legder(c) for c in c2d])
  281. res = leg.legder(c2d, axis=1)
  282. assert_almost_equal(res, tgt)
  283. class TestVander:
  284. # some random values in [-1, 1)
  285. x = np.random.random((3, 5))*2 - 1
  286. def test_legvander(self):
  287. # check for 1d x
  288. x = np.arange(3)
  289. v = leg.legvander(x, 3)
  290. assert_(v.shape == (3, 4))
  291. for i in range(4):
  292. coef = [0]*i + [1]
  293. assert_almost_equal(v[..., i], leg.legval(x, coef))
  294. # check for 2d x
  295. x = np.array([[1, 2], [3, 4], [5, 6]])
  296. v = leg.legvander(x, 3)
  297. assert_(v.shape == (3, 2, 4))
  298. for i in range(4):
  299. coef = [0]*i + [1]
  300. assert_almost_equal(v[..., i], leg.legval(x, coef))
  301. def test_legvander2d(self):
  302. # also tests polyval2d for non-square coefficient array
  303. x1, x2, x3 = self.x
  304. c = np.random.random((2, 3))
  305. van = leg.legvander2d(x1, x2, [1, 2])
  306. tgt = leg.legval2d(x1, x2, c)
  307. res = np.dot(van, c.flat)
  308. assert_almost_equal(res, tgt)
  309. # check shape
  310. van = leg.legvander2d([x1], [x2], [1, 2])
  311. assert_(van.shape == (1, 5, 6))
  312. def test_legvander3d(self):
  313. # also tests polyval3d for non-square coefficient array
  314. x1, x2, x3 = self.x
  315. c = np.random.random((2, 3, 4))
  316. van = leg.legvander3d(x1, x2, x3, [1, 2, 3])
  317. tgt = leg.legval3d(x1, x2, x3, c)
  318. res = np.dot(van, c.flat)
  319. assert_almost_equal(res, tgt)
  320. # check shape
  321. van = leg.legvander3d([x1], [x2], [x3], [1, 2, 3])
  322. assert_(van.shape == (1, 5, 24))
  323. class TestFitting:
  324. def test_legfit(self):
  325. def f(x):
  326. return x*(x - 1)*(x - 2)
  327. def f2(x):
  328. return x**4 + x**2 + 1
  329. # Test exceptions
  330. assert_raises(ValueError, leg.legfit, [1], [1], -1)
  331. assert_raises(TypeError, leg.legfit, [[1]], [1], 0)
  332. assert_raises(TypeError, leg.legfit, [], [1], 0)
  333. assert_raises(TypeError, leg.legfit, [1], [[[1]]], 0)
  334. assert_raises(TypeError, leg.legfit, [1, 2], [1], 0)
  335. assert_raises(TypeError, leg.legfit, [1], [1, 2], 0)
  336. assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[[1]])
  337. assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[1, 1])
  338. assert_raises(ValueError, leg.legfit, [1], [1], [-1,])
  339. assert_raises(ValueError, leg.legfit, [1], [1], [2, -1, 6])
  340. assert_raises(TypeError, leg.legfit, [1], [1], [])
  341. # Test fit
  342. x = np.linspace(0, 2)
  343. y = f(x)
  344. #
  345. coef3 = leg.legfit(x, y, 3)
  346. assert_equal(len(coef3), 4)
  347. assert_almost_equal(leg.legval(x, coef3), y)
  348. coef3 = leg.legfit(x, y, [0, 1, 2, 3])
  349. assert_equal(len(coef3), 4)
  350. assert_almost_equal(leg.legval(x, coef3), y)
  351. #
  352. coef4 = leg.legfit(x, y, 4)
  353. assert_equal(len(coef4), 5)
  354. assert_almost_equal(leg.legval(x, coef4), y)
  355. coef4 = leg.legfit(x, y, [0, 1, 2, 3, 4])
  356. assert_equal(len(coef4), 5)
  357. assert_almost_equal(leg.legval(x, coef4), y)
  358. # check things still work if deg is not in strict increasing
  359. coef4 = leg.legfit(x, y, [2, 3, 4, 1, 0])
  360. assert_equal(len(coef4), 5)
  361. assert_almost_equal(leg.legval(x, coef4), y)
  362. #
  363. coef2d = leg.legfit(x, np.array([y, y]).T, 3)
  364. assert_almost_equal(coef2d, np.array([coef3, coef3]).T)
  365. coef2d = leg.legfit(x, np.array([y, y]).T, [0, 1, 2, 3])
  366. assert_almost_equal(coef2d, np.array([coef3, coef3]).T)
  367. # test weighting
  368. w = np.zeros_like(x)
  369. yw = y.copy()
  370. w[1::2] = 1
  371. y[0::2] = 0
  372. wcoef3 = leg.legfit(x, yw, 3, w=w)
  373. assert_almost_equal(wcoef3, coef3)
  374. wcoef3 = leg.legfit(x, yw, [0, 1, 2, 3], w=w)
  375. assert_almost_equal(wcoef3, coef3)
  376. #
  377. wcoef2d = leg.legfit(x, np.array([yw, yw]).T, 3, w=w)
  378. assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T)
  379. wcoef2d = leg.legfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w)
  380. assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T)
  381. # test scaling with complex values x points whose square
  382. # is zero when summed.
  383. x = [1, 1j, -1, -1j]
  384. assert_almost_equal(leg.legfit(x, x, 1), [0, 1])
  385. assert_almost_equal(leg.legfit(x, x, [0, 1]), [0, 1])
  386. # test fitting only even Legendre polynomials
  387. x = np.linspace(-1, 1)
  388. y = f2(x)
  389. coef1 = leg.legfit(x, y, 4)
  390. assert_almost_equal(leg.legval(x, coef1), y)
  391. coef2 = leg.legfit(x, y, [0, 2, 4])
  392. assert_almost_equal(leg.legval(x, coef2), y)
  393. assert_almost_equal(coef1, coef2)
  394. class TestCompanion:
  395. def test_raises(self):
  396. assert_raises(ValueError, leg.legcompanion, [])
  397. assert_raises(ValueError, leg.legcompanion, [1])
  398. def test_dimensions(self):
  399. for i in range(1, 5):
  400. coef = [0]*i + [1]
  401. assert_(leg.legcompanion(coef).shape == (i, i))
  402. def test_linear_root(self):
  403. assert_(leg.legcompanion([1, 2])[0, 0] == -.5)
  404. class TestGauss:
  405. def test_100(self):
  406. x, w = leg.leggauss(100)
  407. # test orthogonality. Note that the results need to be normalized,
  408. # otherwise the huge values that can arise from fast growing
  409. # functions like Laguerre can be very confusing.
  410. v = leg.legvander(x, 99)
  411. vv = np.dot(v.T * w, v)
  412. vd = 1/np.sqrt(vv.diagonal())
  413. vv = vd[:, None] * vv * vd
  414. assert_almost_equal(vv, np.eye(100))
  415. # check that the integral of 1 is correct
  416. tgt = 2.0
  417. assert_almost_equal(w.sum(), tgt)
  418. class TestMisc:
  419. def test_legfromroots(self):
  420. res = leg.legfromroots([])
  421. assert_almost_equal(trim(res), [1])
  422. for i in range(1, 5):
  423. roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2])
  424. pol = leg.legfromroots(roots)
  425. res = leg.legval(roots, pol)
  426. tgt = 0
  427. assert_(len(pol) == i + 1)
  428. assert_almost_equal(leg.leg2poly(pol)[-1], 1)
  429. assert_almost_equal(res, tgt)
  430. def test_legroots(self):
  431. assert_almost_equal(leg.legroots([1]), [])
  432. assert_almost_equal(leg.legroots([1, 2]), [-.5])
  433. for i in range(2, 5):
  434. tgt = np.linspace(-1, 1, i)
  435. res = leg.legroots(leg.legfromroots(tgt))
  436. assert_almost_equal(trim(res), trim(tgt))
  437. def test_legtrim(self):
  438. coef = [2, -1, 1, 0]
  439. # Test exceptions
  440. assert_raises(ValueError, leg.legtrim, coef, -1)
  441. # Test results
  442. assert_equal(leg.legtrim(coef), coef[:-1])
  443. assert_equal(leg.legtrim(coef, 1), coef[:-3])
  444. assert_equal(leg.legtrim(coef, 2), [0])
  445. def test_legline(self):
  446. assert_equal(leg.legline(3, 4), [3, 4])
  447. def test_leg2poly(self):
  448. for i in range(10):
  449. assert_almost_equal(leg.leg2poly([0]*i + [1]), Llist[i])
  450. def test_poly2leg(self):
  451. for i in range(10):
  452. assert_almost_equal(leg.poly2leg(Llist[i]), [0]*i + [1])
  453. def test_weight(self):
  454. x = np.linspace(-1, 1, 11)
  455. tgt = 1.
  456. res = leg.legweight(x)
  457. assert_almost_equal(res, tgt)