[feat] expose bfactors in protein_to_pyg function

graphein/protein/tensor/io.py

@@ -108,6 +108,7 @@ def protein_to_pyg(
     atom_types: List[str] = PROTEIN_ATOMS,
     remove_nonstandard: bool = True,
     store_het: bool = False,
+    store_bfactor: bool = False,
     fill_value_coords: float = 1e-5,
 ) -> Data:
     """
@@ -159,6 +160,12 @@ def protein_to_pyg(
     :param store_het: Whether or not to store heteroatoms in the ``Data``
         object. Default is ``False``.
     :type store_het: bool
+    :param store_bfactor: Whether or not to store bfactors in the ``Data``
+        object. Default is ``False.
+    :type store_bfactor: bool
+    :param fill_value_coords: Fill value to use for positions in atom37 
+        representation that are not filled. Defaults to 1e-5
+    :type fill_value_coords: float
     :returns: ``Data`` object with attributes: ``x`` (AtomTensor), ``residues``
         (list of 3-letter residue codes), id (ID of protein), residue_id (E.g.
         ``"A:SER:1"``), residue_type (torch.Tensor), ``chains`` (torch.Tensor).
@@ -254,6 +261,10 @@ def protein_to_pyg(
     )
     if store_het:
         out.hetatms = [het_coords]
+
+    if store_bfactor:
+        out.bfactor = torch.tensor(df["b_factor"].values)
+
     return out