fix for 1 heavy atom case

Author: ioanaapapa

CodeEntropy/levels/axes.py

@@ -293,7 +293,9 @@ def get_UA_axes(self, data_container, index: int, res_position):
         else:
             # only one heavy atom or hydrogen molecule
             make_whole(data_container.atoms)
+            residue = data_container
             trans_axes = data_container.atoms.principal_axes()
+
         residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
         # look for heavy atoms in residue of interest
         heavy_atom_indices = []

tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py

@@ -401,7 +401,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes():
     node = FrameCovarianceNode()
 
     bead = _BeadGroup(1)
-    residue_atoms = MagicMock()
+    residue_group = MagicMock()
 
     axes_manager = MagicMock()
     axes_manager.get_UA_axes.return_value = (
@@ -416,7 +416,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes():
 
     force_vecs, torque_vecs = node._build_ua_vectors(
         bead_groups=[bead],
-        residue_atoms=residue_atoms,
+        residue_group=residue_group,
         axes_manager=axes_manager,
         box=np.array([10.0, 10.0, 10.0]),
         force_partitioning=1.0,
@@ -451,7 +451,7 @@ def test_build_ua_vectors_vanilla_path_uses_principal_axes_and_vanilla_axes(
 
     force_vecs, torque_vecs = node._build_ua_vectors(
         bead_groups=[bead],
-        residue_atoms=residue_atoms,
+        residue_group=residue_atoms,
         axes_manager=axes_manager,
         box=np.array([10.0, 10.0, 10.0]),
         force_partitioning=1.0,