make format

netam/dasm.py

@@ -134,8 +134,9 @@ def train_val_datasets_of_pcp_df(pcp_df, branch_length_multiplier=5.0):
 def zero_predictions_along_diagonal(predictions, aa_parents_idxs):
     """Zero out the diagonal of a batch of predictions.
 
-    We do this so that we can sum then have the same type of predictions as for
-    the DNSM."""
+    We do this so that we can sum then have the same type of predictions as for the
+    DNSM.
+    """
     # We would like to do
     # predictions[torch.arange(len(aa_parents_idxs)), aa_parents_idxs] = 0.0
     # but we have a batch dimension. Thus the following.
@@ -199,7 +200,7 @@ def loss_of_batch(self, batch):
         predictions = torch.cat(
             [predictions, torch.zeros_like(predictions[:, :, :1])], dim=-1
         )
-        
+
         predictions = zero_predictions_along_diagonal(predictions, aa_parents_idxs)
 
         predictions_of_mut = torch.sum(predictions, dim=-1)

tests/test_dasm.py

@@ -10,7 +10,11 @@
     add_shm_model_outputs_to_pcp_df,
 )
 from netam.models import TransformerBinarySelectionModelWiggleAct
-from netam.dasm import DASMBurrito, train_val_datasets_of_pcp_df, zero_predictions_along_diagonal
+from netam.dasm import (
+    DASMBurrito,
+    train_val_datasets_of_pcp_df,
+    zero_predictions_along_diagonal,
+)
 
 
 # TODO code dup
@@ -87,12 +91,17 @@ def test_zero_diagonal(dasm_burrito):
     )
     aa_parents_idxs = batch["aa_parents_idxs"].to(dasm_burrito.device)
     zeroed_predictions = predictions.clone()
-    zeroed_predictions = zero_predictions_along_diagonal(zeroed_predictions, aa_parents_idxs)
+    zeroed_predictions = zero_predictions_along_diagonal(
+        zeroed_predictions, aa_parents_idxs
+    )
     L = predictions.shape[1]
     for batch_idx in range(2):
         for i in range(L):
             for j in range(20):
                 if j == aa_parents_idxs[batch_idx, i]:
                     assert zeroed_predictions[batch_idx, i, j] == 0.0
                 else:
-                    assert zeroed_predictions[batch_idx, i, j] == predictions[batch_idx, i, j]
+                    assert (
+                        zeroed_predictions[batch_idx, i, j]
+                        == predictions[batch_idx, i, j]
+                    )