feat(alphafold): custom MSA input for monomer and multimer predictions (#52)

docs/src/en/alphafold.md

@@ -35,6 +35,8 @@ Path to folder to save prediction results in (str). Default: "./[date_time]_gget
 
 `-jhd` `--jackhmmer_savedir`  
 Path to a parent directory in which to store the temporary jackhmmer files (str). By default, `gget alphafold` creates a "tmp" folder in your home directory (`~/tmp/jackhmmer/`), which can take up to ~2 GB of disk space. Use this argument to place these temporary files elsewhere, e.g. on a disk with more free space. Default: None.  
+`-msa` `--msa`  
+Custom multiple sequence alignment (MSA) to use instead of running the internal jackhmmer search. Accepts a3m (`.a3m`) or aligned FASTA (`.fasta`, `.fa`, `.afa`) files. For a monomer, pass a single MSA file; for a multimer, pass one MSA file per chain (in the same order as the sequences). The first sequence in each chain's MSA must be that chain's query (matching the input `sequence`, ignoring gaps). When provided, `gget alphafold` skips the jackhmmer search entirely, so no genetic databases are downloaded — useful for submitting a manually curated MSA to improve folding predictions. For multimer cross-chain pairing to work, the MSA headers should carry species identifiers. Default: None.  
 
 **Flags**  
 `-mfm` `--multimer_for_monomer`  

docs/src/en/updates.md

@@ -5,6 +5,7 @@
 #### *gget* officially became part of [*scverse*](https://scverse.org/) on June 9, 2026. 🥳🥳🥳
 
 **Version ≥ 0.30.9** (XXX XX, 2026):  
+- [`gget alphafold`](alphafold.md): Added a new `msa` argument (`-msa`/`--msa` on the command line) that lets you provide a custom multiple sequence alignment (a3m or aligned FASTA) instead of running the internal jackhmmer search. Pass a single MSA file for a monomer, or one file per chain for a multimer; the first sequence in each chain's MSA must be that chain's query. This enables folding from a manually curated MSA and skips the genetic database download entirely.
 
 **Version ≥ 0.30.8** (Jun 28, 2026):  
 - [`gget g2p`](g2p.md): Either `gene` or `--uniprot_id` is now sufficient — whichever is missing is resolved via UniProt and cached. Gene→UniProt picks the canonical reviewed human Swiss-Prot entry; the resolution and its limitations are logged. The canonical pair is **always** prepended to the result as `gene_name` / `uniprot_id` columns (and stored on `df.attrs`), so the output schema is invariant regardless of input mode. Existing call sites continue to work.

gget/main.py

@@ -1341,6 +1341,20 @@ def main() -> None:
             "which can take up to ~2 GB of disk space. Use this argument to place them elsewhere."
         ),
     )
+    parser_alphafold.add_argument(
+        "-msa",
+        "--msa",
+        type=str,
+        nargs="+",
+        default=None,
+        required=False,
+        help=(
+            "Custom multiple sequence alignment(s) (a3m or aligned FASTA) to use instead of running\n"
+            "the internal jackhmmer search. Pass one MSA file for a monomer, or one file per chain (in\n"
+            "the same order as the sequences) for a multimer. The first sequence in each chain's MSA\n"
+            "must be that chain's query (matching the input sequence, ignoring gaps)."
+        ),
+    )
     parser_alphafold.add_argument(
         "-q",
         "--quiet",
@@ -3716,6 +3730,7 @@ def main() -> None:
             show_sidechains=False,
             verbose=args.quiet,
             jackhmmer_savedir=args.jackhmmer_savedir,
+            msa=args.msa,
         )
 
     ## pdb return

tests/fixtures/test_alphafold_msa.a3m

@@ -0,0 +1,8 @@
+>query
+MKVLAAGSTKDEFGHIKLMN
+>sp|homolog1
+MKVLAAGSTKDEFGHIKLMN
+>sp|homolog2
+MKVLAAGSTKDEF-HIKLMN
+>sp|homolog3_with_insertion
+MKVefLAAGSTKDEFGHIKLMN

tests/fixtures/test_alphafold_msa.fasta

@@ -0,0 +1,6 @@
+>query
+MKVLAAGSTKDEFGHIKLMN
+>homolog1
+MKVLAAGSTKDEF-HIKLMN
+>homolog2
+MKVLAAGSTKDEFGHIKLMN

tests/test_alphafold.py

@@ -5,13 +5,28 @@
 from unittest.mock import patch
 
 import gget.gget_alphafold as gget_alphafold
-from gget.gget_alphafold import clean_up, get_jackhmmer_dir
+from gget.gget_alphafold import (
+    _validate_custom_msa_paths,
+    clean_up,
+    detect_msa_format,
+    get_jackhmmer_dir,
+    parse_custom_msa,
+)
 from gget.gget_setup import UUID
 
 # AlphaFold requires heavy third-party dependencies (alphafold, openmm, jackhmmer, model
 # parameters) that are not available in the CI/test environment, so a full prediction run
 # cannot be exercised here. These tests validate the user-facing jackhmmer save-directory
-# option (https://github.com/scverse/gget/issues/49) at the argument/path-handling level.
+# option (https://github.com/scverse/gget/issues/49) at the argument/path-handling level and
+# the user-provided custom MSA input feature (https://github.com/scverse/gget/issues/52) at
+# the parsing/validation level (the logic that turns a custom a3m/FASTA file into AlphaFold
+# MSA features).
+
+FIXTURE_DIR = os.path.join(os.path.dirname(__file__), "fixtures")
+A3M_FIXTURE = os.path.join(FIXTURE_DIR, "test_alphafold_msa.a3m")
+FASTA_FIXTURE = os.path.join(FIXTURE_DIR, "test_alphafold_msa.fasta")
+
+QUERY = "MKVLAAGSTKDEFGHIKLMN"
 
 
 class TestAlphafoldJackhmmerSavedir(unittest.TestCase):
@@ -75,5 +90,107 @@ def test_cli_exposes_jackhmmer_savedir_flag(self):
         self.assertIn("--jackhmmer_savedir", result.stdout)
 
 
+class TestAlphafoldCustomMSA(unittest.TestCase):
+    def test_detect_msa_format(self):
+        self.assertEqual(detect_msa_format("alignment.a3m"), "a3m")
+        self.assertEqual(detect_msa_format("/path/to/ALIGNMENT.A3M"), "a3m")
+        self.assertEqual(detect_msa_format("alignment.fasta"), "fasta")
+        self.assertEqual(detect_msa_format("alignment.fa"), "fasta")
+        self.assertEqual(detect_msa_format("alignment.afa"), "fasta")
+
+    def test_detect_msa_format_unsupported(self):
+        with self.assertRaises(ValueError):
+            detect_msa_format("alignment.txt")
+        with self.assertRaises(ValueError):
+            detect_msa_format("alignment.sto")
+
+    def test_parse_a3m_fixture(self):
+        with open(A3M_FIXTURE) as f:
+            aligned_sequences, deletion_matrix, descriptions = parse_custom_msa(f.read())
+
+        # Four sequences, the first of which is the query.
+        self.assertEqual(len(aligned_sequences), 4)
+        self.assertEqual(len(descriptions), 4)
+        self.assertEqual(aligned_sequences[0], QUERY)
+
+        # After removing a3m insertions, every aligned row has the query's length.
+        for seq in aligned_sequences:
+            self.assertEqual(len(seq), len(QUERY))
+
+        # Deletion matrix rows align 1:1 with the aligned columns.
+        for seq, deletions in zip(aligned_sequences, deletion_matrix, strict=True):
+            self.assertEqual(len(deletions), len(seq))
+
+        # The query and the gapped homolog carry no insertions.
+        self.assertEqual(sum(deletion_matrix[0]), 0)
+        self.assertEqual(sum(deletion_matrix[2]), 0)
+
+        # homolog3 has two lowercase insertion characters ("ef") after the third residue.
+        self.assertEqual(sum(deletion_matrix[3]), 2)
+        self.assertEqual(deletion_matrix[3][3], 2)
+
+    def test_parse_fasta_fixture_has_no_insertions(self):
+        with open(FASTA_FIXTURE) as f:
+            aligned_sequences, deletion_matrix, _ = parse_custom_msa(f.read())
+
+        self.assertEqual(aligned_sequences[0], QUERY)
+        # Aligned FASTA contains no lowercase insertions -> deletion matrix is all zeros.
+        for deletions in deletion_matrix:
+            self.assertEqual(sum(deletions), 0)
+
+    def test_query_matches_first_msa_sequence(self):
+        # Mirrors the validation gget performs: first MSA sequence (gaps removed) == query.
+        with open(A3M_FIXTURE) as f:
+            aligned_sequences, _, _ = parse_custom_msa(f.read())
+        query_in_msa = aligned_sequences[0].replace("-", "").upper()
+        self.assertEqual(query_in_msa, QUERY)
+
+    def test_parse_empty_msa_raises(self):
+        with self.assertRaises(ValueError):
+            parse_custom_msa("")
+
+    def test_parse_non_fasta_raises(self):
+        with self.assertRaises(ValueError):
+            parse_custom_msa("MKVLAAG\nNOHEADERLINE\n")
+
+    def test_cli_exposes_msa_flag(self):
+        """The command-line interface exposes the --msa option for gget alphafold."""
+        result = subprocess.run(
+            [sys.executable, "-m", "gget", "alphafold", "--help"],
+            capture_output=True,
+            text=True,
+        )
+        self.assertIn("--msa", result.stdout)
+
+    # --- custom MSA validation, incl. multimer (one MSA file per chain) ---
+
+    def test_validate_monomer_ok(self):
+        # A single MSA file whose first sequence is the query passes.
+        _validate_custom_msa_paths([A3M_FIXTURE], [QUERY])
+
+    def test_validate_multimer_ok(self):
+        # One MSA file per chain (both fixtures have the query as their first sequence).
+        _validate_custom_msa_paths([A3M_FIXTURE, FASTA_FIXTURE], [QUERY, QUERY])
+
+    def test_validate_count_mismatch_raises(self):
+        # Fewer MSA files than chains must raise.
+        with self.assertRaises(ValueError):
+            _validate_custom_msa_paths([A3M_FIXTURE], [QUERY, QUERY])
+        with self.assertRaises(ValueError):
+            _validate_custom_msa_paths([A3M_FIXTURE, FASTA_FIXTURE], [QUERY])
+
+    def test_validate_query_mismatch_raises(self):
+        # A chain whose first MSA sequence does not match that chain's query must raise.
+        with self.assertRaises(ValueError):
+            _validate_custom_msa_paths([A3M_FIXTURE], ["MKVLAAGSTKDEFGHIKLPP"])
+        # Second chain mismatches while the first is fine.
+        with self.assertRaises(ValueError):
+            _validate_custom_msa_paths([A3M_FIXTURE, FASTA_FIXTURE], [QUERY, "TOTALLYDIFFERENT"])
+
+    def test_validate_missing_file_raises(self):
+        with self.assertRaises(FileNotFoundError):
+            _validate_custom_msa_paths([os.path.join(FIXTURE_DIR, "does_not_exist.a3m")], [QUERY])
+
+
 if __name__ == "__main__":
     unittest.main()