[feat] expose bfactors in protein_to_pyg function #388
Conversation
kierandidi
force-pushed
the
feat/expose_b_factor
branch
from
1d4b17a
to
8fd1eb2
Compare
graphein/protein/tensor/io.py
Outdated
| @@ -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) | |||
There was a problem hiding this comment.
I think torch.from_numpy might be a little better
kierandidi
force-pushed
the
feat/expose_b_factor
branch
from
62c8719
to
b4bf164
Compare
graphein/protein/tensor/io.py
Outdated
| @@ -254,6 +261,10 @@ def protein_to_pyg( | |||
| ) | |||
| if store_het: | |||
| out.hetatms = [het_coords] | |||
|
|
|||
| if store_bfactor: | |||
| out.bfactor = torch.from_numpy(df["b_factor"].values) | |||
There was a problem hiding this comment.
This will be of shape num_atoms x 1 right? I expect this would break batching as all other tensors are of shape n_res x X
There was a problem hiding this comment.
Good catch, solved this via a group by now that averages b factors on a per residue basis, consistent with the plddt information in the predicted datasets.
There was a problem hiding this comment.
Ran more tests now and everything seems to work in a backward compatible way, lmk if there is anything else that needs to happen before merging @a-r-j
kierandidi
force-pushed
the
feat/expose_b_factor
branch
from
6bfb02e
to
8f0a730
Compare
for more information, see https://pre-commit.ci
|
Reference Issues/PRs
What does this implement/fix? Explain your changes
Bfactors are present in df from biopandas, but cannot be saved to pyg object. Now this option is enabled.
What testing did you do to verify the changes in this PR?
Pull Request Checklist
./CHANGELOG.mdfile (if applicable)./graphein/tests/*directories (if applicable)./notebooks/(if applicable)python -m py.test tests/and make sure that all unit tests pass (for small modifications, it might be sufficient to only run the specific test file, e.g.,python -m py.test tests/protein/test_graphs.py)black .andisort .