reuse existing code from unresolved rasa to allow Alphafold3 to return List[Structure] and then have the cli be able to save a pdb file

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -104,6 +104,7 @@
 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
 
 from Bio.PDB.StructureBuilder import StructureBuilder
+from Bio.PDB.Structure import Structure
 from Bio.PDB.PDBIO import PDBIO
 from Bio.PDB.DSSP import DSSP
 import tempfile
@@ -5130,13 +5131,14 @@ def get_cid_molecule_type(
 
 
 @typecheck
-def _protein_structure_from_feature(
+def protein_structure_from_feature(
     asym_id: Int[" n"],  
     molecule_ids: Int[" n"],  
     molecule_atom_lens: Int[" n"],  
     atom_pos: Float["m 3"],  
     atom_mask: Bool[" m"],  
-) -> Bio.PDB.Structure.Structure:
+) -> Structure:
+
     """Create structure for unresolved proteins.
 
     :param atom_mask: True for valid atoms, False for missing/padding atoms
@@ -5626,7 +5628,7 @@ def _compute_unresolved_rasa(
         chain_atom_pos = atom_pos[chain_mask_to_atom]
         chain_atom_mask = atom_mask[chain_mask_to_atom]
 
-        structure = _protein_structure_from_feature(
+        structure = protein_structure_from_feature(
             chain_asym_id,
             chain_molecule_ids,
             chain_molecule_atom_lens,
@@ -6381,6 +6383,7 @@ def forward(
         return_all_diffused_atom_pos: bool = False,
         return_confidence_head_logits: bool = False,
         return_distogram_head_logits: bool = False,
+        return_bio_pdb_structures: bool = False,
         num_rollout_steps: int | None = None,
         rollout_show_tqdm_pbar: bool = False,
         detach_when_recycling: bool = None,
@@ -6390,8 +6393,9 @@ def forward(
         filepaths: List[str] | None = None
     ) -> (
         Float['b m 3'] |
+        List[Structure] |
         Float['ts b m 3'] |
-        Tuple[Float['b m 3'] | Float['ts b m 3'], ConfidenceHeadLogits | Alphafold3Logits] |
+        Tuple[Float['b m 3'] | List[Structure] | Float['ts b m 3'], ConfidenceHeadLogits | Alphafold3Logits] |
         Float[''] |
         Tuple[Float[''], LossBreakdown]
     ):
@@ -6666,6 +6670,22 @@ def forward(
             if return_confidence_head_logits:
                 confidence_head_atom_pos_input = sampled_atom_pos.clone()
 
+            # convert sampled atom positions to bio pdb structures
+
+            if return_bio_pdb_structures:
+                assert not return_all_diffused_atom_pos
+
+                sampled_atom_pos = [
+                    protein_structure_from_feature(*args)
+                    for args in zip(
+                        additional_molecule_feats[..., 2],
+                        molecule_ids,
+                        molecule_atom_lens,
+                        sampled_atom_pos,
+                        atom_mask
+                    )
+                ]
+
             if not return_confidence_head_logits:
                 return sampled_atom_pos
 

alphafold3_pytorch/cli.py

@@ -9,12 +9,14 @@
     alphafold3_inputs_to_batched_atom_input
 )
 
+from Bio.PDB.PDBIO import PDBIO
+
 # simple cli using click
 
 @click.command()
 @click.option('-ckpt', '--checkpoint', type = str, help = 'path to alphafold3 checkpoint')
 @click.option('-p', '--protein', type = str, help = 'one protein sequence')
-@click.option('-o', '--output', type = str, help = 'output path', default = 'atompos.pt')
+@click.option('-o', '--output', type = str, help = 'output path', default = 'output.pdb')
 def cli(
     checkpoint: str,
     protein: str,
@@ -33,11 +35,13 @@ def cli(
     batched_atom_input = alphafold3_inputs_to_batched_atom_input(alphafold3_input, atoms_per_window = alphafold3.atoms_per_window)
 
     alphafold3.eval()
-    sampled_atom_pos = alphafold3(**batched_atom_input.model_forward_dict())
+    structure, = alphafold3(**batched_atom_input.model_forward_dict(), return_bio_pdb_structures = True)
 
     output_path = Path(output)
     output_path.parents[0].mkdir(exist_ok = True, parents = True)
 
-    torch.save(sampled_atom_pos, str(output_path))
+    pdb_writer = PDBIO()
+    pdb_writer.set_structure(structure)
+    pdb_writer.save(str(output_path))
 
-    print(f'atomic positions saved to {str(output_path)}')
+    print(f'pdb saved to {str(output_path)}')

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.4.43"
+version = "0.4.44"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },

tests/test_input.py

@@ -67,7 +67,9 @@ def test_atom_dataset():
 # alphafold3 input
 
 @pytest.mark.parametrize('directed_bonds', (False, True))
-def test_alphafold3_input(directed_bonds):
+def test_alphafold3_input(
+    directed_bonds
+):
 
     alphafold3_input = Alphafold3Input(
         proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
@@ -119,6 +121,46 @@ def test_alphafold3_input(directed_bonds):
 
     alphafold3(**batched_atom_input.model_forward_dict(), num_sample_steps = 1)
 
+def test_return_bio_pdb_structures():
+
+    alphafold3_input = Alphafold3Input(
+        proteins = ['MLEICLKLVGCKSKKGLSSSSSCYLEEALQRPVASDF', 'MGKCRGLRTARKLRSHRRDQKWHDKQYKKAHLGTALKANPFGGASHAKGIVLEKVGVEAKQPNSAIRKCVRVQLIKNGKKITAFVPNDGCLNFIEENDEVLVAGFGRKGHAVGDIPGVRFKVVKVANVSLLALYKGKKERPRS'],
+    )
+
+    batched_atom_input = alphafold3_inputs_to_batched_atom_input(alphafold3_input)
+
+    # feed it into alphafold3
+
+    alphafold3 = Alphafold3(
+        dim_atom_inputs = 3,
+        dim_atompair_inputs = 5,
+        num_atom_embeds = 0,
+        num_atompair_embeds = 0,
+        atoms_per_window = 27,
+        dim_template_feats = 108,
+        num_dist_bins = 64,
+        num_molecule_mods = 0,
+        confidence_head_kwargs = dict(
+            pairformer_depth = 1
+        ),
+        template_embedder_kwargs = dict(
+            pairformer_stack_depth = 1
+        ),
+        msa_module_kwargs = dict(
+            depth = 1
+        ),
+        pairformer_stack = dict(
+            depth = 2
+        ),
+        diffusion_module_kwargs = dict(
+            atom_encoder_depth = 1,
+            token_transformer_depth = 1,
+            atom_decoder_depth = 1,
+        )
+    )
+
+    alphafold3(**batched_atom_input.model_forward_dict(), num_sample_steps = 1, return_bio_pdb_structures = True)
+
 def test_atompos_input():
 
     contrived_protein = 'AG'