Skip to content

GitLab

Commit 699f23fd authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Add option to select chains in PDB files

parent 5161c0f6
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
openfold/data/data_pipeline.py
View file @ 699f23fd
...@@ -472,7 +472,8 @@ class DataPipeline: ...@@ -472,7 +472,8 @@ class DataPipeline:
self, self,
pdb_path: str, pdb_path: str,
alignment_dir: str, alignment_dir: str,
is_distillation: bool = True is_distillation: bool = True,
chain_id: Optional[str] = None,
) -> FeatureDict: ) -> FeatureDict:
""" """
Assembles features for a protein in a PDB file. Assembles features for a protein in a PDB file.
...@@ -480,7 +481,7 @@ class DataPipeline: ...@@ -480,7 +481,7 @@ class DataPipeline:
with open(pdb_path, 'r') as f: with open(pdb_path, 'r') as f:
pdb_str = f.read() pdb_str = f.read()
protein_object = protein.from_pdb_string(pdb_str) protein_object = protein.from_pdb_string(pdb_str, chain_id)
input_sequence = _aatype_to_str_sequence(protein_object.aatype) input_sequence = _aatype_to_str_sequence(protein_object.aatype)
description = os.path.splitext(os.path.basename(pdb_path))[0].upper() description = os.path.splitext(os.path.basename(pdb_path))[0].upper()
pdb_feats = make_pdb_features( pdb_feats = make_pdb_features(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment