bond angles

src/beignet/bond_angles.py

@@ -21,15 +21,13 @@ def bond_angles(input: Tensor, angle_indices: Tensor) -> Tensor:
     # Data verification
     num_frames, n_atoms, _ = input.shape
     if torch.any(angle_indices >= n_atoms) or torch.any(angle_indices < 0):
-        raise ValueError(
-            'angle_indices must be between 0 and %d' % (n_atoms - 1))
+        raise ValueError("angle_indices must be between 0 and %d" % (n_atoms - 1))
 
     if angle_indices.shape[0] == 0:
         return torch.zeros((num_frames, 0), dtype=torch.float32)
 
     # Initializing the output tensor
-    angles = torch.zeros((num_frames, angle_indices.shape[0]),
-                         dtype=torch.float32)
+    angles = torch.zeros((num_frames, angle_indices.shape[0]), dtype=torch.float32)
 
     # Gathering vectors related to the angle calculation
     vec1 = input[:, angle_indices[:, 1]] - input[:, angle_indices[:, 0]]

tests/beignet/test_bond_angles.py

@@ -11,44 +11,44 @@ def radians(degrees):
 
 def test_straight_line_angle():
     # Tests three collinear points which must produce an angle of pi radians (180 degrees)
-    traj = torch.tensor([[[0, 0, 0],
-                          [1, 0, 0],
-                          [2, 0, 0]]], dtype=torch.float32)
+    traj = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]]], dtype=torch.float32)
     angle_indices = torch.tensor([[0, 1, 2]])
 
     expected_angles = torch.tensor([[radians(180)]])
     computed_angles = bond_angles(traj, angle_indices)
 
-    assert torch.allclose(computed_angles,
-                          expected_angles), "Should calculate 180 degrees for collinear points."
+    assert torch.allclose(
+        computed_angles, expected_angles
+    ), "Should calculate 180 degrees for collinear points."
 
 
 def test_right_angle():
     # Tests an L shape (right angle, 90 degrees)
-    traj = torch.tensor([[[0, 0, 0],
-                          [1, 0, 0],
-                          [1, 1, 0]]], dtype=torch.float32)
+    traj = torch.tensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0]]], dtype=torch.float32)
     angle_indices = torch.tensor([[0, 1, 2]])
 
     expected_angles = torch.tensor([[radians(90)]])
     computed_angles = bond_angles(traj, angle_indices)
 
-    assert torch.allclose(computed_angles, expected_angles,
-                          atol=1e-5), "Should calculate 90 degrees for orthogonal vectors."
+    assert torch.allclose(
+        computed_angles, expected_angles, atol=1e-5
+    ), "Should calculate 90 degrees for orthogonal vectors."
 
 
 def test_acute_angle():
     # Acute angle test 45 degrees
-    traj = torch.tensor([[[0, 0, 0],
-                          [1, 0, 0],
-                          [1, torch.sqrt(torch.tensor(2.0)), 0]]], dtype=torch.float32)
+    traj = torch.tensor(
+        [[[0, 0, 0], [1, 0, 0], [1, torch.sqrt(torch.tensor(2.0)), 0]]],
+        dtype=torch.float32,
+    )
     angle_indices = torch.tensor([[0, 1, 2]])
 
     expected_angles = torch.tensor([[radians(45)]])
     computed_angles = bond_angles(traj, angle_indices)
 
-    assert torch.allclose(computed_angles, expected_angles,
-                          atol=1e-5), "Should calculate 45 degrees for acute angle."
+    assert torch.allclose(
+        computed_angles, expected_angles, atol=1e-5
+    ), "Should calculate 45 degrees for acute angle."
 
 
 def test_no_indices_provided():
@@ -57,4 +57,6 @@ def test_no_indices_provided():
     angle_indices = torch.empty((0, 3), dtype=torch.int32)
 
     computed_angles = bond_angles(traj, angle_indices)
-    assert computed_angles.nelement() == 0, "Providing no indices should result in an empty tensor."
+    assert (
+        computed_angles.nelement() == 0
+    ), "Providing no indices should result in an empty tensor."