[fix] calculate bfactor per residue instead of per atom

a-r-j · Apr 21, 2024 · 8f0a730 · 8f0a730
1 parent b4bf164
commit 8f0a730
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/graphein/protein/tensor/io.py b/graphein/protein/tensor/io.py
@@ -263,7 +263,9 @@ def protein_to_pyg(
         out.hetatms = [het_coords]
 
     if store_bfactor:
-        out.bfactor = torch.from_numpy(df["b_factor"].values)
+        # group by residue_id and average b_factor per residue
+        residue_bfactors = df.groupby("residue_id")["b_factor"].mean()
+        out.bfactor = torch.from_numpy(residue_bfactors.values)
 
     return out