Merge pull request #40 from chrisjonesBSU/add/moit

Add moment of inertia function

cmeutils/geometry.py

@@ -91,3 +91,33 @@ def angle_between_vectors(u, v, min_angle=True, degrees=True):
     if degrees:
         return np.rad2deg(angle)
     return angle
+
+
+def moit(points, masses, center=np.zeros(3)):
+    """Calculates moment of inertia tensor (moit) for rigid bodies. 
+    
+    Assumes rigid body center is at origin unless center is provided.
+    Only calculates diagonal elements.
+    
+    Parameters
+    ----------
+    points : numpy.ndarray (N,3)
+        x, y, and z coordinates of the rigid body constituent particles
+    masses : numpy.ndarray (N,)
+        masses of the constituent particles
+    center : numpy.ndarray (3,), default np.array([0,0,0])
+        x, y, and z coordinates of the rigid body center
+        
+    Returns
+    -------
+    numpy.ndarray (3,)
+        moment of inertia tensor for the rigid body center
+    """
+    points -= center
+    x = points[:, 0]
+    y = points[:, 1]
+    z = points[:, 2]
+    I_xx = np.sum((y ** 2 + z ** 2) * masses)
+    I_yy = np.sum((x ** 2 + z ** 2) * masses)
+    I_zz = np.sum((x ** 2 + y ** 2) * masses)
+    return np.array((I_xx, I_yy, I_zz))

cmeutils/tests/test_geometry.py

@@ -4,10 +4,14 @@
 import pytest
 
 from base_test import BaseTest
-from cmeutils.geometry import get_plane_normal, angle_between_vectors
+from cmeutils.geometry import get_plane_normal, angle_between_vectors, moit
 
 
 class TestGeometry(BaseTest):
+    def test_moit(self):
+        _moit = moit(points=[(-1, 0, 0), (1, 0, 0)], masses=[1, 1])
+        assert np.array_equal(_moit, np.array([0, 2., 2.]))
+
     def test_get_plane_normal(self):
         points = np.array(
             [[ 1, 0, 0],