Accounts for selenocysteine in sidechain torsion angle computation

CHANGELOG.md

@@ -9,6 +9,7 @@
 * Ensures exproting groups of PDB chains with PDBManager selects the first model for multu-model structures. [#311](https://github.com/a-r-j/graphein/pull/311)
 * 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)
 
 #### 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)
@@ -18,7 +19,7 @@
 * Improved handling of non-standard residues in the `graphein.protein.tensor` module. [#307](https://github.com/a-r-j/graphein/pull/307)
 * Insertions retained by default in the `graphein.protein.tensor` module. I.e. `insertions=True` is now the default behaviour.[#307](https://github.com/a-r-j/graphein/pull/307)
 * Adds transform composition to FoldComp Dataset [#312](https://github.com/a-r-j/graphein/pull/312)
-* Improve FoldComp dataloading performance [#313](https://github.com/a-r-j/graphein/pull/313)
+* Improve FoldComp dataloading performance and include B factors (pLDDT) in output. [#313](https://github.com/a-r-j/graphein/pull/313) [#315](https://github.com/a-r-j/graphein/pull/315)
 
 ### 1.7.0 - UNRELEASED
 

graphein/ml/datasets/foldcomp_dataset.py

@@ -292,13 +292,15 @@ def fc_to_pyg(data: Dict[str, Any], name: Optional[str] = None) -> Protein:
 
         res_idx = np.repeat(res_num, atom_counts)
         coords[res_idx, atom_idx, :] = np.array(data["coordinates"])
+        b_factor = np.array(data["b_factors"]) / 100
 
         return Protein(
             coords=torch.from_numpy(coords).float(),
             residues=res,
             residue_id=[f"A:{m}:{str(n)}" for m, n in zip(res, res_num)],
             chains=torch.zeros(len(res)),
             residue_type=residue_type.long(),
+            b_factor=torch.from_numpy(b_factor).float(),
             id=name,
         )
 

graphein/protein/tensor/angles.py

@@ -70,10 +70,18 @@ def _extract_torsion_coords(
     res_atoms = []
     idxs = []
 
+    # Whether or not the protein contains selenocysteine
+    selenium = coords.shape[1] == 38
+
     # Iterate over residues and grab indices of the atoms for each Chi angle
     for i, res in enumerate(res_types):
         res_coords = []
-        for angle_coord_set in CHI_ANGLES_ATOMS[res]:
+
+        angle_groups = CHI_ANGLES_ATOMS[res]
+        if not selenium and res == "SEC":
+            angle_groups = []
+
+        for angle_coord_set in angle_groups:
             res_coords.append([ATOM_NUMBERING[i] for i in angle_coord_set])
             idxs.append(i)
         res_atoms.append(torch.tensor(res_coords, device=coords.device))
@@ -115,6 +123,9 @@ def sidechain_torsion(
     :return: _description_
     :rtype: Union[TorsionTensor, Tuple[TorsionTensor, torch.Tensor]]
     """
+    # Whether or not the protein contains selenocysteine
+    selenium = coords.shape[1] == 38
+
     idxs, coords = _extract_torsion_coords(coords, res_types)
     angles = _dihedral_angle(
         coords[:, 0, :].unsqueeze(1),
@@ -139,7 +150,11 @@ def sidechain_torsion(
     res_types = copy.deepcopy(res_types)
     res_types.reverse()
     for res in res_types:
-        if res in ["ALA", "GLY", "UNK"]:
+        PAD_RESIDUES = ["ALA", "GLY", "UNK"]
+        if not selenium:
+            PAD_RESIDUES.append("SEC")
+
+        if res in PAD_RESIDUES:
             post_pad_len += 1
         else:
             break
@@ -226,7 +241,7 @@ def kappa(
         angles = angles[mask]
 
     if embed:
-        angles = angle_to_unit_circle(angles)
+        angles = torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1)
 
     return angles
 
@@ -288,7 +303,7 @@ def alpha(
         angles = angles[mask]
 
     if embed:
-        angles = angle_to_unit_circle(angles)
+        angles = torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1)
 
     return angles