Fix and enhance geometric node features

graphein/protein/features/nodes/geometry.py

@@ -42,6 +42,11 @@ def add_sidechain_vector(
         if d["residue_name"] == "GLY":
             # If GLY, set vector to 0
             vec = np.array([0, 0, 0])
+        elif n not in sc_centroid.index:
+            vec = np.array([0, 0, 0])
+            log.warning(
+                f"Non-glycine residue {n} does not have side-chain atoms."
+            )
         else:
             if reverse:
                 vec = d["coords"] - np.array(
@@ -94,6 +99,11 @@ def add_beta_carbon_vector(
     for n, d in g.nodes(data=True):
         if d["residue_name"] == "GLY":
             vec = np.array([0, 0, 0])
+        elif n not in c_beta_coords.index:
+            vec = np.array([0, 0, 0])
+            log.warning(
+                f"Non-glycine residue {n} does not have a beta-carbon."
+            )
         else:
             if reverse:
                 vec = d["coords"] - np.array(

graphein/protein/graphs.py

@@ -537,7 +537,7 @@ def initialise_graph_with_metadata(
         chain_ids=list(protein_df["chain_id"].unique()),
         pdb_df=protein_df,
         raw_pdb_df=raw_pdb_df,
-        rgroup_df=compute_rgroup_dataframe(remove_insertions(raw_pdb_df)),
+        rgroup_df=compute_rgroup_dataframe(raw_pdb_df),
         coords=np.asarray(protein_df[["x_coord", "y_coord", "z_coord"]]),
     )
 

tests/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+from loguru import logger
+
+
+@pytest.fixture
+def caplog(caplog):
+    handler_id = logger.add(caplog.handler, format="{message}")
+    yield caplog
+    logger.remove(handler_id)

tests/protein/nodes/features/test_geometry.py

@@ -8,6 +8,7 @@
 from functools import partial
 
 import numpy as np
+from loguru import logger
 
 from graphein.protein.config import ProteinGraphConfig
 from graphein.protein.features.nodes.geometry import (
@@ -17,7 +18,7 @@
 from graphein.protein.graphs import construct_graph
 
 
-def test_add_beta_carbon_vector():
+def test_add_beta_carbon_vector(caplog):
     config = ProteinGraphConfig(
         edge_construction_functions=[
             partial(add_beta_carbon_vector, scale=True)
@@ -76,6 +77,16 @@ def test_add_beta_carbon_vector():
     for n, d in g.nodes(data=True):
         assert d["c_beta_vector"].shape == (3,)
 
+    # Test handling of missing beta-carbons
+    g = construct_graph(config=config, pdb_code="3se8")
+    assert "H:CYS:104" in caplog.text  # Test warning is raised
+    for n, d in g.nodes(data=True):
+        assert d["c_beta_vector"].shape == (3,)
+        if n == "H:CYS:104":
+            np.testing.assert_equal(
+                d["c_beta_vector"], np.array([0.0, 0.0, 0.0])
+            )
+
 
 def test_add_sidechain_vector():
     config = ProteinGraphConfig(
@@ -127,8 +138,3 @@ def test_add_sidechain_vector():
             np.testing.assert_almost_equal(
                 sc_true, d["coords"] + d["sidechain_vector"]
             )
-
-
-if __name__ == "__main__":
-    test_add_beta_carbon_vector()
-    test_add_sidechain_vector()