Fixes padding in pNerF computations

CHANGELOG.md

@@ -10,6 +10,7 @@
 * Fixes bug with exporting PDBs with only one splitting strategy in PDBManager [#311](https://github.com/a-r-j/graphein/pull/311)
 * Fixes incorrect jaxtyping syntax for variable size dimensions [#312](https://github.com/a-r-j/graphein/pull/312)
 * Fixes shape of angle embeddings for `graphein.protein.tesnor.angles.alpha/kappa`. [#315](https://github.com/a-r-j/graphein/pull/315)
+* Fixes initialisation of `Protein` objects. [#317](https://github.com/a-r-j/graphein/issues/317) [#318](https://github.com/a-r-j/graphein/pull/318)
 
 #### Other Changes
 * Adds entry point for biopandas dataframes in `graphein.protein.tensor.io.protein_to_pyg`. [#310](https://github.com/a-r-j/graphein/pull/310)

graphein/ml/datasets/foldcomp_dataset.py

@@ -302,6 +302,7 @@ def fc_to_pyg(data: Dict[str, Any], name: Optional[str] = None) -> Protein:
             residue_type=residue_type.long(),
             b_factor=torch.from_numpy(b_factor).float(),
             id=name,
+            x=torch.zeros(len(res)),
         )
 
     def len(self) -> int:

graphein/protein/tensor/data.py

@@ -132,7 +132,10 @@ class Protein(Data):
     """
 
     def __init__(
-        self, atom_list: List[str] = PROTEIN_ATOMS, fill_value: float = 1e-5
+        self,
+        atom_list: List[str] = PROTEIN_ATOMS,
+        fill_value: float = 1e-5,
+        **kwargs,
     ) -> None:
         """
         .. seealso::
@@ -145,10 +148,13 @@ def __init__(
         :param fill_value: Value to fill in for missing values in the
             AtomTensor.
         :type fill_value: float
+        :param kwargs: Additional keyword arguments to store as attributes.
         """
         super().__init__()
         self.fill_value = fill_value
         self.atom_list = atom_list
+        for k, v in kwargs.items():
+            setattr(self, k, v)
 
     # I/O
     def from_dataframe(

graphein/protein/tensor/pnerf.py

@@ -277,13 +277,14 @@ def extend(prev_three_coords, point, multi_m):
         )
         coords_trans = torch.cat([coords_pretrans[i], transformed_coords], 0)
 
-    # coords = F.pad(
-    #    coords_trans[: total_num_angles - 1], (0, 0, 0, 0, 1, 0)
-    # )  # original
+    coords = F.pad(
+        coords_trans[: total_num_angles - 1], (0, 0, 0, 0, 1, 0)
+    )  # original
     # coords = F.pad(coords_trans[: total_num_angles - 2], (0, 0, 0, 0, 2, 0))
     # coords = F.pad(coords_trans[: total_num_angles - 3], (0, 0, 0, 0, 3, 0))
     # return coords
-    return coords_trans
+    # return coords_trans
+    return coords
 
 
 def sn_nerf(