test_utils.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. from pytest import raises
  2. from sympy import (
  3. symbols, Expr, Tuple, Integer, cos, solveset, FiniteSet, ImageSet)
  4. from sympy.plotting.utils import (
  5. _create_ranges, _plot_sympify, extract_solution)
  6. from sympy.physics.mechanics import ReferenceFrame, Vector as MechVector
  7. from sympy.vector import CoordSys3D, Vector
  8. def test_plot_sympify():
  9. x, y = symbols("x, y")
  10. # argument is already sympified
  11. args = x + y
  12. r = _plot_sympify(args)
  13. assert r == args
  14. # one argument needs to be sympified
  15. args = (x + y, 1)
  16. r = _plot_sympify(args)
  17. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 2
  18. assert isinstance(r[0], Expr)
  19. assert isinstance(r[1], Integer)
  20. # string and dict should not be sympified
  21. args = (x + y, (x, 0, 1), "str", 1, {1: 1, 2: 2.0})
  22. r = _plot_sympify(args)
  23. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 5
  24. assert isinstance(r[0], Expr)
  25. assert isinstance(r[1], Tuple)
  26. assert isinstance(r[2], str)
  27. assert isinstance(r[3], Integer)
  28. assert isinstance(r[4], dict) and isinstance(r[4][1], int) and isinstance(r[4][2], float)
  29. # nested arguments containing strings
  30. args = ((x + y, (y, 0, 1), "a"), (x + 1, (x, 0, 1), "$f_{1}$"))
  31. r = _plot_sympify(args)
  32. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 2
  33. assert isinstance(r[0], Tuple)
  34. assert isinstance(r[0][1], Tuple)
  35. assert isinstance(r[0][1][1], Integer)
  36. assert isinstance(r[0][2], str)
  37. assert isinstance(r[1], Tuple)
  38. assert isinstance(r[1][1], Tuple)
  39. assert isinstance(r[1][1][1], Integer)
  40. assert isinstance(r[1][2], str)
  41. # vectors from sympy.physics.vectors module are not sympified
  42. # vectors from sympy.vectors are sympified
  43. # in both cases, no error should be raised
  44. R = ReferenceFrame("R")
  45. v1 = 2 * R.x + R.y
  46. C = CoordSys3D("C")
  47. v2 = 2 * C.i + C.j
  48. args = (v1, v2)
  49. r = _plot_sympify(args)
  50. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 2
  51. assert isinstance(v1, MechVector)
  52. assert isinstance(v2, Vector)
  53. def test_create_ranges():
  54. x, y = symbols("x, y")
  55. # user don't provide any range -> return a default range
  56. r = _create_ranges({x}, [], 1)
  57. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 1
  58. assert isinstance(r[0], (Tuple, tuple))
  59. assert r[0] == (x, -10, 10)
  60. r = _create_ranges({x, y}, [], 2)
  61. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 2
  62. assert isinstance(r[0], (Tuple, tuple))
  63. assert isinstance(r[1], (Tuple, tuple))
  64. assert r[0] == (x, -10, 10) or (y, -10, 10)
  65. assert r[1] == (y, -10, 10) or (x, -10, 10)
  66. assert r[0] != r[1]
  67. # not enough ranges provided by the user -> create default ranges
  68. r = _create_ranges(
  69. {x, y},
  70. [
  71. (x, 0, 1),
  72. ],
  73. 2,
  74. )
  75. assert isinstance(r, (list, tuple, Tuple)) and len(r) == 2
  76. assert isinstance(r[0], (Tuple, tuple))
  77. assert isinstance(r[1], (Tuple, tuple))
  78. assert r[0] == (x, 0, 1) or (y, -10, 10)
  79. assert r[1] == (y, -10, 10) or (x, 0, 1)
  80. assert r[0] != r[1]
  81. # too many free symbols
  82. raises(ValueError, lambda: _create_ranges({x, y}, [], 1))
  83. raises(ValueError, lambda: _create_ranges({x, y}, [(x, 0, 5), (y, 0, 1)], 1))
  84. def test_extract_solution():
  85. x = symbols("x")
  86. sol = solveset(cos(10 * x))
  87. assert sol.has(ImageSet)
  88. res = extract_solution(sol)
  89. assert len(res) == 20
  90. assert isinstance(res, FiniteSet)
  91. res = extract_solution(sol, 20)
  92. assert len(res) == 40
  93. assert isinstance(res, FiniteSet)