Minor fix to ensure virtual angles are returned in correct shape (#315)

* add PSW to nonstandard residues

* improve insertion and non-standard residue handling

* refactor chain selection

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove unused verbosity arg

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix chain selection in tests

* fix chain selection in tutorial notebook

* fix notebook chain selection

* fix chain selection typehint

* Update changelog

* Add NLW to non-standard residues

* add .ent support

* add entry for construction from dataframe

* add missing stage arg

* improve obsolete mapping retrieving to include entries with no replacement

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update changelog

* add transforms to foldcomp datasets

* fix jaxtyping syntax

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update changelog

* fix double application of transforms

* improve foldcomp data loading performance

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove unused imports

* linting

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update changelog

* add B factors to FC parsing output

* bugfix to alpha & kappa angle embedding

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update changelog

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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

@@ -226,7 +226,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 +288,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