def find_center_of_mass(coordinates: np.ndarray=None, masses: np.ndarray=None, cluster=None) -> np.ndarray:
"""Returns the center of mass of the cluster
Args:
coordinates: np.array, coordinate array
masses: np.array, array of masses
cluster: object relating to the cluster. Must have masses and coordinates attributes
Returns:
np.array(3), coordinates corresponding to center of mass
"""
if cluster is not None:
assert hasattr(cluster, "masses") and hasattr(cluster, "coordinates")
coordinates = cluster.coordinates
if cluster.masses is not None:
masses = cluster.masses
else:
masses = get_masses(cluster.particle_names)
com = np.zeros(3)
for c, m in zip(coordinates, masses):
if m == 0:
log.warning("Particle with mass of 0 at {} not included in determining center of mass".format(c))
continue
wc = np.array(c) * m
com += wc
return com / sum(masses)
def get_center_of_mass(self) -> np.ndarray:
"""Returns the center of mass of the cluster
Returns:
np.array(3), coordinates of the center of mass
"""
if self.masses is None:
try:
self.masses = get_masses(self.particle_names)
except KeyError: # TODO: test the error checking in get_center_of_mass
raise
return find_center_of_mass(coordinates=self.coordinates,
masses=self.masses)
def get_masses(self) -> np.ndarray:
"""Wraps get_masses()"""
return get_masses(self.particle_names)
def test_TypeError(self) -> None:
"""Should raise a TypeError when no atom labels are supplied"""
with self.assertRaises(TypeError):
# noinspection PyArgumentList
get_masses()
def test_KeyError(self) -> None:
with self.assertRaises(KeyError):
get_masses(["InvalidKey"])
def test_get_masses_benzene(self) -> None:
self.assertListEqual([12.01, 12.01, 12.01, 12.01, 12.01, 12.01, 1.01, 1.01, 1.01, 1.01, 1.01, 1.01],
list(get_masses(["C", "C", "C", "C", "C", "C", "H", "H", "H", "H", "H", "H"])))
def test_get_masses_H2(self) -> None:
self.assertListEqual([1.01, 1.01], list(get_masses(["H", "H"])))