multichain protein

Author: lucidrains

alphafold3_pytorch/cli.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import click
 from pathlib import Path
 
@@ -15,19 +17,25 @@
 
 @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('-prot', '--protein', type = str, multiple = True, help = 'protein sequences')
+@click.option('-rna', '--rna', type = str, multiple = True, help = 'single stranded rna sequences')
+@click.option('-dna', '--dna', type = str, multiple = True, help = 'single stranded dna sequences')
 @click.option('-o', '--output', type = str, help = 'output path', default = 'output.cif')
 def cli(
     checkpoint: str,
-    protein: str,
+    protein: list[str],
+    rna: list[str],
+    dna: list[str],
     output: str
 ):
 
     checkpoint_path = Path(checkpoint)
     assert checkpoint_path.exists(), f'AlphaFold 3 checkpoint must exist at {str(checkpoint_path)}'
 
     alphafold3_input = Alphafold3Input(
-        proteins = [protein],
+        proteins = protein,
+        ss_rna = rna,
+        ss_dna = dna,
     )
 
     alphafold3 = Alphafold3.init_and_load(checkpoint_path)

pyproject.toml

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

tests/test_cli.py

@@ -22,8 +22,14 @@ def test_cli():
     checkpoint_path = './test-folder/test-cli-alphafold3.pt'
     alphafold3.save(checkpoint_path, overwrite = True)
 
-    cli(['--checkpoint', checkpoint_path, '--protein', 'AG', '--output', './test-folder/output.pdb'], standalone_mode = False)
-
-    assert Path('./test-folder/output.pdb').exists()
+    cli([
+        '--checkpoint', checkpoint_path,
+        '-prot', 'AG',
+        '-prot', 'TC',
+        '--output',
+        './test-folder/output.mmcif'
+    ], standalone_mode = False)
+
+    assert Path('./test-folder/output.mmcif').exists()
 
     rmtree('./test-folder')