diff --git a/cmeutils/geometry.py b/cmeutils/geometry.py index 089c561..a1a2111 100644 --- a/cmeutils/geometry.py +++ b/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. diff --git a/cmeutils/tests/test_geometry.py b/cmeutils/tests/test_geometry.py index 8426764..c38a577 100644 --- a/cmeutils/tests/test_geometry.py +++ b/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]))