test_functions.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. from sympy import sin, cos, tan, pi, symbols, Matrix, S, Function
  2. from sympy.physics.mechanics import (Particle, Point, ReferenceFrame,
  3. RigidBody)
  4. from sympy.physics.mechanics import (angular_momentum, dynamicsymbols,
  5. kinetic_energy, linear_momentum,
  6. outer, potential_energy, msubs,
  7. find_dynamicsymbols, Lagrangian)
  8. from sympy.physics.mechanics.functions import (
  9. center_of_mass, _validate_coordinates, _parse_linear_solver)
  10. from sympy.testing.pytest import raises, warns_deprecated_sympy
  11. q1, q2, q3, q4, q5 = symbols('q1 q2 q3 q4 q5')
  12. N = ReferenceFrame('N')
  13. A = N.orientnew('A', 'Axis', [q1, N.z])
  14. B = A.orientnew('B', 'Axis', [q2, A.x])
  15. C = B.orientnew('C', 'Axis', [q3, B.y])
  16. def test_linear_momentum():
  17. N = ReferenceFrame('N')
  18. Ac = Point('Ac')
  19. Ac.set_vel(N, 25 * N.y)
  20. I = outer(N.x, N.x)
  21. A = RigidBody('A', Ac, N, 20, (I, Ac))
  22. P = Point('P')
  23. Pa = Particle('Pa', P, 1)
  24. Pa.point.set_vel(N, 10 * N.x)
  25. raises(TypeError, lambda: linear_momentum(A, A, Pa))
  26. raises(TypeError, lambda: linear_momentum(N, N, Pa))
  27. assert linear_momentum(N, A, Pa) == 10 * N.x + 500 * N.y
  28. def test_angular_momentum_and_linear_momentum():
  29. """A rod with length 2l, centroidal inertia I, and mass M along with a
  30. particle of mass m fixed to the end of the rod rotate with an angular rate
  31. of omega about point O which is fixed to the non-particle end of the rod.
  32. The rod's reference frame is A and the inertial frame is N."""
  33. m, M, l, I = symbols('m, M, l, I')
  34. omega = dynamicsymbols('omega')
  35. N = ReferenceFrame('N')
  36. a = ReferenceFrame('a')
  37. O = Point('O')
  38. Ac = O.locatenew('Ac', l * N.x)
  39. P = Ac.locatenew('P', l * N.x)
  40. O.set_vel(N, 0 * N.x)
  41. a.set_ang_vel(N, omega * N.z)
  42. Ac.v2pt_theory(O, N, a)
  43. P.v2pt_theory(O, N, a)
  44. Pa = Particle('Pa', P, m)
  45. A = RigidBody('A', Ac, a, M, (I * outer(N.z, N.z), Ac))
  46. expected = 2 * m * omega * l * N.y + M * l * omega * N.y
  47. assert linear_momentum(N, A, Pa) == expected
  48. raises(TypeError, lambda: angular_momentum(N, N, A, Pa))
  49. raises(TypeError, lambda: angular_momentum(O, O, A, Pa))
  50. raises(TypeError, lambda: angular_momentum(O, N, O, Pa))
  51. expected = (I + M * l**2 + 4 * m * l**2) * omega * N.z
  52. assert angular_momentum(O, N, A, Pa) == expected
  53. def test_kinetic_energy():
  54. m, M, l1 = symbols('m M l1')
  55. omega = dynamicsymbols('omega')
  56. N = ReferenceFrame('N')
  57. O = Point('O')
  58. O.set_vel(N, 0 * N.x)
  59. Ac = O.locatenew('Ac', l1 * N.x)
  60. P = Ac.locatenew('P', l1 * N.x)
  61. a = ReferenceFrame('a')
  62. a.set_ang_vel(N, omega * N.z)
  63. Ac.v2pt_theory(O, N, a)
  64. P.v2pt_theory(O, N, a)
  65. Pa = Particle('Pa', P, m)
  66. I = outer(N.z, N.z)
  67. A = RigidBody('A', Ac, a, M, (I, Ac))
  68. raises(TypeError, lambda: kinetic_energy(Pa, Pa, A))
  69. raises(TypeError, lambda: kinetic_energy(N, N, A))
  70. assert 0 == (kinetic_energy(N, Pa, A) - (M*l1**2*omega**2/2
  71. + 2*l1**2*m*omega**2 + omega**2/2)).expand()
  72. def test_potential_energy():
  73. m, M, l1, g, h, H = symbols('m M l1 g h H')
  74. omega = dynamicsymbols('omega')
  75. N = ReferenceFrame('N')
  76. O = Point('O')
  77. O.set_vel(N, 0 * N.x)
  78. Ac = O.locatenew('Ac', l1 * N.x)
  79. P = Ac.locatenew('P', l1 * N.x)
  80. a = ReferenceFrame('a')
  81. a.set_ang_vel(N, omega * N.z)
  82. Ac.v2pt_theory(O, N, a)
  83. P.v2pt_theory(O, N, a)
  84. Pa = Particle('Pa', P, m)
  85. I = outer(N.z, N.z)
  86. A = RigidBody('A', Ac, a, M, (I, Ac))
  87. Pa.potential_energy = m * g * h
  88. A.potential_energy = M * g * H
  89. assert potential_energy(A, Pa) == m * g * h + M * g * H
  90. def test_Lagrangian():
  91. M, m, g, h = symbols('M m g h')
  92. N = ReferenceFrame('N')
  93. O = Point('O')
  94. O.set_vel(N, 0 * N.x)
  95. P = O.locatenew('P', 1 * N.x)
  96. P.set_vel(N, 10 * N.x)
  97. Pa = Particle('Pa', P, 1)
  98. Ac = O.locatenew('Ac', 2 * N.y)
  99. Ac.set_vel(N, 5 * N.y)
  100. a = ReferenceFrame('a')
  101. a.set_ang_vel(N, 10 * N.z)
  102. I = outer(N.z, N.z)
  103. A = RigidBody('A', Ac, a, 20, (I, Ac))
  104. Pa.potential_energy = m * g * h
  105. A.potential_energy = M * g * h
  106. raises(TypeError, lambda: Lagrangian(A, A, Pa))
  107. raises(TypeError, lambda: Lagrangian(N, N, Pa))
  108. def test_msubs():
  109. a, b = symbols('a, b')
  110. x, y, z = dynamicsymbols('x, y, z')
  111. # Test simple substitution
  112. expr = Matrix([[a*x + b, x*y.diff() + y],
  113. [x.diff().diff(), z + sin(z.diff())]])
  114. sol = Matrix([[a + b, y],
  115. [x.diff().diff(), 1]])
  116. sd = {x: 1, z: 1, z.diff(): 0, y.diff(): 0}
  117. assert msubs(expr, sd) == sol
  118. # Test smart substitution
  119. expr = cos(x + y)*tan(x + y) + b*x.diff()
  120. sd = {x: 0, y: pi/2, x.diff(): 1}
  121. assert msubs(expr, sd, smart=True) == b + 1
  122. N = ReferenceFrame('N')
  123. v = x*N.x + y*N.y
  124. d = x*(N.x|N.x) + y*(N.y|N.y)
  125. v_sol = 1*N.y
  126. d_sol = 1*(N.y|N.y)
  127. sd = {x: 0, y: 1}
  128. assert msubs(v, sd) == v_sol
  129. assert msubs(d, sd) == d_sol
  130. def test_find_dynamicsymbols():
  131. a, b = symbols('a, b')
  132. x, y, z = dynamicsymbols('x, y, z')
  133. expr = Matrix([[a*x + b, x*y.diff() + y],
  134. [x.diff().diff(), z + sin(z.diff())]])
  135. # Test finding all dynamicsymbols
  136. sol = {x, y.diff(), y, x.diff().diff(), z, z.diff()}
  137. assert find_dynamicsymbols(expr) == sol
  138. # Test finding all but those in sym_list
  139. exclude_list = [x, y, z]
  140. sol = {y.diff(), x.diff().diff(), z.diff()}
  141. assert find_dynamicsymbols(expr, exclude=exclude_list) == sol
  142. # Test finding all dynamicsymbols in a vector with a given reference frame
  143. d, e, f = dynamicsymbols('d, e, f')
  144. A = ReferenceFrame('A')
  145. v = d * A.x + e * A.y + f * A.z
  146. sol = {d, e, f}
  147. assert find_dynamicsymbols(v, reference_frame=A) == sol
  148. # Test if a ValueError is raised on supplying only a vector as input
  149. raises(ValueError, lambda: find_dynamicsymbols(v))
  150. # This function tests the center_of_mass() function
  151. # that was added in PR #14758 to compute the center of
  152. # mass of a system of bodies.
  153. def test_center_of_mass():
  154. a = ReferenceFrame('a')
  155. m = symbols('m', real=True)
  156. p1 = Particle('p1', Point('p1_pt'), S.One)
  157. p2 = Particle('p2', Point('p2_pt'), S(2))
  158. p3 = Particle('p3', Point('p3_pt'), S(3))
  159. p4 = Particle('p4', Point('p4_pt'), m)
  160. b_f = ReferenceFrame('b_f')
  161. b_cm = Point('b_cm')
  162. mb = symbols('mb')
  163. b = RigidBody('b', b_cm, b_f, mb, (outer(b_f.x, b_f.x), b_cm))
  164. p2.point.set_pos(p1.point, a.x)
  165. p3.point.set_pos(p1.point, a.x + a.y)
  166. p4.point.set_pos(p1.point, a.y)
  167. b.masscenter.set_pos(p1.point, a.y + a.z)
  168. point_o=Point('o')
  169. point_o.set_pos(p1.point, center_of_mass(p1.point, p1, p2, p3, p4, b))
  170. expr = 5/(m + mb + 6)*a.x + (m + mb + 3)/(m + mb + 6)*a.y + mb/(m + mb + 6)*a.z
  171. assert point_o.pos_from(p1.point)-expr == 0
  172. def test_validate_coordinates():
  173. q1, q2, q3, u1, u2, u3, ua1, ua2, ua3 = dynamicsymbols('q1:4 u1:4 ua1:4')
  174. s1, s2, s3 = symbols('s1:4')
  175. # Test normal
  176. _validate_coordinates([q1, q2, q3], [u1, u2, u3],
  177. u_auxiliary=[ua1, ua2, ua3])
  178. # Test not equal number of coordinates and speeds
  179. _validate_coordinates([q1, q2])
  180. _validate_coordinates([q1, q2], [u1])
  181. _validate_coordinates(speeds=[u1, u2])
  182. # Test duplicate
  183. _validate_coordinates([q1, q2, q2], [u1, u2, u3], check_duplicates=False)
  184. raises(ValueError, lambda: _validate_coordinates(
  185. [q1, q2, q2], [u1, u2, u3]))
  186. _validate_coordinates([q1, q2, q3], [u1, u2, u2], check_duplicates=False)
  187. raises(ValueError, lambda: _validate_coordinates(
  188. [q1, q2, q3], [u1, u2, u2], check_duplicates=True))
  189. raises(ValueError, lambda: _validate_coordinates(
  190. [q1, q2, q3], [q1, u2, u3], check_duplicates=True))
  191. _validate_coordinates([q1, q2, q3], [u1, u2, u3], check_duplicates=False,
  192. u_auxiliary=[u1, ua2, ua2])
  193. raises(ValueError, lambda: _validate_coordinates(
  194. [q1, q2, q3], [u1, u2, u3], u_auxiliary=[u1, ua2, ua3]))
  195. raises(ValueError, lambda: _validate_coordinates(
  196. [q1, q2, q3], [u1, u2, u3], u_auxiliary=[q1, ua2, ua3]))
  197. raises(ValueError, lambda: _validate_coordinates(
  198. [q1, q2, q3], [u1, u2, u3], u_auxiliary=[ua1, ua2, ua2]))
  199. # Test is_dynamicsymbols
  200. _validate_coordinates([q1 + q2, q3], is_dynamicsymbols=False)
  201. raises(ValueError, lambda: _validate_coordinates([q1 + q2, q3]))
  202. _validate_coordinates([s1, q1, q2], [0, u1, u2], is_dynamicsymbols=False)
  203. raises(ValueError, lambda: _validate_coordinates(
  204. [s1, q1, q2], [0, u1, u2], is_dynamicsymbols=True))
  205. _validate_coordinates([s1 + s2 + s3, q1], [0, u1], is_dynamicsymbols=False)
  206. raises(ValueError, lambda: _validate_coordinates(
  207. [s1 + s2 + s3, q1], [0, u1], is_dynamicsymbols=True))
  208. _validate_coordinates(u_auxiliary=[s1, ua1], is_dynamicsymbols=False)
  209. raises(ValueError, lambda: _validate_coordinates(u_auxiliary=[s1, ua1]))
  210. # Test normal function
  211. t = dynamicsymbols._t
  212. a = symbols('a')
  213. f1, f2 = symbols('f1:3', cls=Function)
  214. _validate_coordinates([f1(a), f2(a)], is_dynamicsymbols=False)
  215. raises(ValueError, lambda: _validate_coordinates([f1(a), f2(a)]))
  216. raises(ValueError, lambda: _validate_coordinates(speeds=[f1(a), f2(a)]))
  217. dynamicsymbols._t = a
  218. _validate_coordinates([f1(a), f2(a)])
  219. raises(ValueError, lambda: _validate_coordinates([f1(t), f2(t)]))
  220. dynamicsymbols._t = t
  221. def test_parse_linear_solver():
  222. A, b = Matrix(3, 3, symbols('a:9')), Matrix(3, 2, symbols('b:6'))
  223. assert _parse_linear_solver(Matrix.LUsolve) == Matrix.LUsolve # Test callable
  224. assert _parse_linear_solver('LU')(A, b) == Matrix.LUsolve(A, b)
  225. def test_deprecated_moved_functions():
  226. from sympy.physics.mechanics.functions import (
  227. inertia, inertia_of_point_mass, gravity)
  228. N = ReferenceFrame('N')
  229. with warns_deprecated_sympy():
  230. assert inertia(N, 0, 1, 0, 1) == (N.x | N.y) + (N.y | N.x) + (N.y | N.y)
  231. with warns_deprecated_sympy():
  232. assert inertia_of_point_mass(1, N.x + N.y, N) == (
  233. (N.x | N.x) + (N.y | N.y) + 2 * (N.z | N.z) -
  234. (N.x | N.y) - (N.y | N.x))
  235. p = Particle('P')
  236. with warns_deprecated_sympy():
  237. assert gravity(-2 * N.z, p) == [(p.masscenter, -2 * p.mass * N.z)]