Function to return Total Energy in physics/mechanics

sympy/physics/mechanics/particle.py

@@ -223,3 +223,38 @@ def potential_energy(self, scalar):
         """
 
         self._pe = sympify(scalar)
+
+    def total_energy(self, frame):
+        """The total energy of the Particle.
+
+        The total energy, E, of a particle, P, is given by
+
+        'E = K.E + P.E.'
+
+        where P.E. is the potential energy of the particle P, and
+        K.E. is kinetic energy of the particle P, in the supplied ReferenceFrame.
+
+        Parameters
+        ==========
+
+        frame : ReferenceFrame
+            The Particle's velocity is typically defined with respect to
+            an inertial frame but any relevant frame in which the velocity is
+            known can be supplied.
+
+        Examples
+        ========
+
+        >>> from sympy.physics.mechanics import Particle, Point, ReferenceFrame
+        >>> from sympy import symbols
+        >>> m, v, r, g, h = symbols('m v r g h')
+        >>> N = ReferenceFrame('N')
+        >>> O = Point('O')
+        >>> P = Particle('P', O, m)
+        >>> P.point.set_vel(N, v * N.y)
+        >>> P.potential_energy = m * g * h
+        >>> P.total_energy(N)
+        g*h*m + m*v**2/2
+
+        """
+        return self.kinetic_energy(frame) + self.potential_energy

sympy/physics/mechanics/rigidbody.py

@@ -199,6 +199,89 @@ def angular_momentum(self, point, frame):
 
         return I.dot(w) + r.cross(m * v)
 
+    def rotational_kinetic_energy(self, frame):
+        """Rotational kinetic energy of the rigid body
+
+        The rotational kinetic energy, T, of a rigid body, B, is given by
+
+        'T = 1/2 (I omega^2)'
+
+        where I is the central inertia dyadic of rigid body B,
+        omega is the angular velocity of the body's
+        mass center in the supplied ReferenceFrame.
+
+        Parameters
+        ==========
+
+        frame : ReferenceFrame
+            The RigidBody's angular velocity and the velocity of it's mass
+            center are typically defined with respect to an inertial frame but
+            any relevant frame in which the velocities are known can be supplied.
+
+        Examples
+        ========
+
+        >>> from sympy.physics.mechanics import Point, ReferenceFrame, outer
+        >>> from sympy.physics.mechanics import RigidBody
+        >>> from sympy import symbols
+        >>> M, v, r, omega = symbols('M v r omega')
+        >>> N = ReferenceFrame('N')
+        >>> b = ReferenceFrame('b')
+        >>> b.set_ang_vel(N, omega * b.x)
+        >>> P = Point('P')
+        >>> P.set_vel(N, v * N.x)
+        >>> I = outer (b.x, b.x)
+        >>> inertia_tuple = (I, P)
+        >>> B = RigidBody('B', P, b, M, inertia_tuple)
+        >>> B.rotational_kinetic_energy(N)
+        omega**2/2
+
+        """
+
+        return (self.frame.ang_vel_in(frame) & (self.central_inertia &
+                self.frame.ang_vel_in(frame)) / sympify(2))
+
+    def translational_kinetic_energy(self, frame):
+        """Translational kinetic energy of the rigid body
+
+        The translational kinetic energy, T, of a rigid body, B, is given by
+
+        'T = 1/2 (m v^2)'
+
+        where m is the mass of rigid body B, v is the velocity of the body's
+        mass center in the supplied ReferenceFrame.
+
+        Parameters
+        ==========
+
+        frame : ReferenceFrame
+            The RigidBody's angular velocity and the velocity of it's mass
+            center are typically defined with respect to an inertial frame but
+            any relevant frame in which the velocities are known can be supplied.
+
+        Examples
+        ========
+
+        >>> from sympy.physics.mechanics import Point, ReferenceFrame, outer
+        >>> from sympy.physics.mechanics import RigidBody
+        >>> from sympy import symbols
+        >>> M, v, r, omega = symbols('M v r omega')
+        >>> N = ReferenceFrame('N')
+        >>> b = ReferenceFrame('b')
+        >>> b.set_ang_vel(N, omega * b.x)
+        >>> P = Point('P')
+        >>> P.set_vel(N, v * N.x)
+        >>> I = outer (b.x, b.x)
+        >>> inertia_tuple = (I, P)
+        >>> B = RigidBody('B', P, b, M, inertia_tuple)
+        >>> B.translational_kinetic_energy(N)
+        M*v**2/2
+
+        """
+
+        return (self.mass * (self.masscenter.vel(frame) &
+            self.masscenter.vel(frame)) / sympify(2))
+
     def kinetic_energy(self, frame):
         """Kinetic energy of the rigid body
 
@@ -238,13 +321,7 @@ def kinetic_energy(self, frame):
 
         """
 
-        rotational_KE = (self.frame.ang_vel_in(frame) & (self.central_inertia &
-                self.frame.ang_vel_in(frame)) / sympify(2))
-
-        translational_KE = (self.mass * (self.masscenter.vel(frame) &
-            self.masscenter.vel(frame)) / sympify(2))
-
-        return rotational_KE + translational_KE
+        return self.rotational_kinetic_energy(frame) + self.translational_kinetic_energy(frame)
 
     @property
     def potential_energy(self):
@@ -296,3 +373,40 @@ def potential_energy(self, scalar):
         """
 
         self._pe = sympify(scalar)
+
+    def total_energy(self, frame):
+        """The total energy of the RigidBody.
+
+        The total energy E, of rigidbody B, is given by
+        E = K.E. + P.E.
+
+        where P.E. is the potential energy of the rigid body B and
+        K.E. is kinetic energy of the rigid body B in the supplied ReferenceFrame.
+
+        Parameters
+        ==========
+
+        frame : ReferenceFrame
+            The RigidBody's angular velocity and the velocity of it's mass
+            center are typically defined with respect to an inertial frame but
+            any relevant frame in which the velocities are known can be supplied.
+
+        Examples
+        ========
+
+        >>> from sympy.physics.mechanics import RigidBody, Point, outer, ReferenceFrame
+        >>> from sympy import symbols
+        >>> M, v, r, omega, g, h = symbols('M v r omega g h')
+        >>> N = ReferenceFrame('N')
+        >>> b = ReferenceFrame('b')
+        >>> b.set_ang_vel(N, omega * b.x)
+        >>> P = Point('P')
+        >>> P.set_vel(N, v * N.x)
+        >>> I = outer (b.x, b.x)
+        >>> Inertia_tuple = (I, P)
+        >>> B = RigidBody('B', P, b, M, Inertia_tuple)
+        >>> B.potential_energy = M * g * h
+        >>> B.total_energy(N)
+        M*g*h + M*v**2/2 + omega**2/2
+        """
+        return self.kinetic_energy(frame) + self.potential_energy

sympy/physics/mechanics/tests/test_particle.py

@@ -33,7 +33,7 @@ def test_particle():
     assert p.angular_momentum(O, N) == m2 * r * (v3 * N.x - v1 * N.z)
     p.potential_energy = m * g * h
     assert p.potential_energy == m * g * h
-    # TODO make the result not be system-dependent
+    assert p.total_energy(
+        N) == m*g*h + m2 * v1**2 / 2 + m2 * v2**2 / 2 + m2 * v3**2 / 2
     assert p.kinetic_energy(
-        N) in [m2*(v1**2 + v2**2 + v3**2)/2,
-        m2 * v1**2 / 2 + m2 * v2**2 / 2 + m2 * v3**2 / 2]
+        N) == m2 * v1**2 / 2 + m2 * v2**2 / 2 + m2 * v3**2 / 2

sympy/physics/mechanics/tests/test_rigidbody.py

@@ -53,7 +53,10 @@ def test_rigidbody2():
     assert B.angular_momentum(O, N) == omega * b.x - M*v*r*b.z
     B.potential_energy = M * g * h
     assert B.potential_energy == M * g * h
+    assert expand(2 * B.translational_kinetic_energy(N)) == M * v**2
+    assert expand(2 * B.rotational_kinetic_energy(N)) == omega**2
     assert expand(2 * B.kinetic_energy(N)) == omega**2 + M * v**2
+    assert B.total_energy(N) == M * g * h + M * v**2/2 + omega**2/2
 
 def test_rigidbody3():
     q1, q2, q3, q4 = dynamicsymbols('q1:5')