test_ode.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. #from mpmath.calculus import ODE_step_euler, ODE_step_rk4, odeint, arange
  2. from mpmath import odefun, cos, sin, mpf, sinc, mp
  3. '''
  4. solvers = [ODE_step_euler, ODE_step_rk4]
  5. def test_ode1():
  6. """
  7. Let's solve:
  8. x'' + w**2 * x = 0
  9. i.e. x1 = x, x2 = x1':
  10. x1' = x2
  11. x2' = -x1
  12. """
  13. def derivs((x1, x2), t):
  14. return x2, -x1
  15. for solver in solvers:
  16. t = arange(0, 3.1415926, 0.005)
  17. sol = odeint(derivs, (0., 1.), t, solver)
  18. x1 = [a[0] for a in sol]
  19. x2 = [a[1] for a in sol]
  20. # the result is x1 = sin(t), x2 = cos(t)
  21. # let's just check the end points for t = pi
  22. assert abs(x1[-1]) < 1e-2
  23. assert abs(x2[-1] - (-1)) < 1e-2
  24. def test_ode2():
  25. """
  26. Let's solve:
  27. x' - x = 0
  28. i.e. x = exp(x)
  29. """
  30. def derivs((x), t):
  31. return x
  32. for solver in solvers:
  33. t = arange(0, 1, 1e-3)
  34. sol = odeint(derivs, (1.,), t, solver)
  35. x = [a[0] for a in sol]
  36. # the result is x = exp(t)
  37. # let's just check the end point for t = 1, i.e. x = e
  38. assert abs(x[-1] - 2.718281828) < 1e-2
  39. '''
  40. def test_odefun_rational():
  41. mp.dps = 15
  42. # A rational function
  43. f = lambda t: 1/(1+mpf(t)**2)
  44. g = odefun(lambda x, y: [-2*x*y[0]**2], 0, [f(0)])
  45. assert f(2).ae(g(2)[0])
  46. def test_odefun_sinc_large():
  47. mp.dps = 15
  48. # Sinc function; test for large x
  49. f = sinc
  50. g = odefun(lambda x, y: [(cos(x)-y[0])/x], 1, [f(1)], tol=0.01, degree=5)
  51. assert abs(f(100) - g(100)[0])/f(100) < 0.01
  52. def test_odefun_harmonic():
  53. mp.dps = 15
  54. # Harmonic oscillator
  55. f = odefun(lambda x, y: [-y[1], y[0]], 0, [1, 0])
  56. for x in [0, 1, 2.5, 8, 3.7]: # we go back to 3.7 to check caching
  57. c, s = f(x)
  58. assert c.ae(cos(x))
  59. assert s.ae(sin(x))