Add method to Geometry which finds best fit vetctor through a set of coordinates.

cmeutils/geometry.py

@@ -2,6 +2,33 @@
 from numpy.linalg import svd
 
 
+def get_backbone_vector(coordinates):
+    """Calculates the line of bets fit through a set of 3D coordinates.
+
+    Best fit is calculated using numpy's Singular Value Decomposition.
+
+    Parameters
+    ----------
+    coordinates : numpy.ndarray, shape (N,3)
+        Coordinates (x,y,z) through which to fit a plane. Must be at least 3
+        points.
+
+    Returns
+    -------
+    direction_vector : numpy.ndarray, shape (3,)
+        The vector of the best fit line.
+    """
+    if coordinates.shape[0] < 2:
+        raise ValueError("Coordinates must contain at least 2 points.")
+    # Center the coordinates (subtract the mean)
+    centered_coordinates = coordinates - np.mean(coordinates, axis=0)
+    # Use PCA to find the principal components
+    _, _, V = np.linalg.svd(centered_coordinates)
+    # The first principal component (V[0]) is the vec of the best-fit line
+    direction_vector = V[0]
+    return np.abs(direction_vector)
+
+
 def get_plane_normal(points):
     """Calculate the plane which best fits a cloud of points.
 

cmeutils/tests/test_geometry.py

@@ -3,9 +3,11 @@
 import numpy as np
 import pytest
 from base_test import BaseTest
+from mbuild.lib.recipes import Alkane
 
 from cmeutils.geometry import (
     angle_between_vectors,
+    get_backbone_vector,
     get_plane_normal,
     moit,
     radial_grid_positions,
@@ -14,6 +16,24 @@
 
 
 class TestGeometry(BaseTest):
+    def test_backbone_vector(self):
+        z_coords = np.array([[0, 0, 1], [0, 0, 2], [0, 0, 3]])
+        backbone = get_backbone_vector(z_coords)
+        assert np.allclose(backbone, np.array([0, 0, 1]), atol=1e-5)
+
+        x_coords = np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0]])
+        backbone = get_backbone_vector(x_coords)
+        assert np.allclose(backbone, np.array([1, 0, 0]), atol=1e-5)
+
+        mb_chain = Alkane(n=20)
+        chain_backbone = get_backbone_vector(mb_chain.xyz)
+        assert np.allclose(chain_backbone, np.array([0, 1, 0]), atol=1e-2)
+
+    def test_backbone_vector_bad_input(self):
+        with pytest.raises(ValueError):
+            coordinates = np.array([1, 1, 1])
+            get_backbone_vector(coordinates)
+
     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.0, 2.0]))