Repository: microsoft/bioemu

Issue:
Support for user-provided MSAs
Thank you for the wonderful tool!

Is it technically possible to run bioemu fully locally? I have my own very nice MSAs, and I want to use them for bioemu, and I do not want to use any remote servers. 

Would you consider implementing such functionality that allows users to submit their own MSAs?

Thank you for your time

Discussion:
@sarahnlewis: Coming soon! Thank you for the feedback. #38  would let you use your own local MSA server but that's not quite what you want. 

I am not sure when we will add the feature you ask for. If you urgently need to hack it for yourself, I suggest you clone the repo, make an editable install and change the code so that [here](https://github.com/microsoft/bioemu/blob/3e89886143d164a9676e407db0ff01d88d8668cb/src/bioemu/get_embeds.py#L83) a path to your .a3m gets passed instead of .fasta.@genesup: > Coming soon! Thank you for the feedback. [#38](https://github.com/microsoft/bioemu/pull/38) would let you use your own local MSA server but that's not quite what you want.
> 
> I am not sure when we will add the feature you ask for. If you urgently need to hack it for yourself, I suggest you clone the repo, make an editable install and change the code so that [here](https://github.com/microsoft/bioemu/blob/3e89886143d164a9676e407db0ff01d88d8668cb/src/bioemu/get_embeds.py#L83) a path to your .a3m gets passed instead of .fasta.

Thank you very much for a very fast response. 

For clarity. MSA server calculates only the sequence alignment, and then embeddings are computed locally from that alignment? So, I just need to implement a simple fix in the code, that would cut out the MSA server and instead use my local alignment? Did I understand you correctly?@YuuuXie: > > Coming soon! Thank you for the feedback. [#38](https://github.com/microsoft/bioemu/pull/38) would let you use your own local MSA server but that's not quite what you want.
> > I am not sure when we will add the feature you ask for. If you urgently need to hack it for yourself, I suggest you clone the repo, make an editable install and change the code so that [here](https://github.com/microsoft/bioemu/blob/3e89886143d164a9676e407db0ff01d88d8668cb/src/bioemu/get_embeds.py#L83) a path to your .a3m gets passed instead of .fasta.
> 
> Thank you very much for a very fast response.
> 
> For clarity. MSA server calculates only the sequence alignment, and then embeddings are computed locally from that alignment? So, I just need to implement a simple fix in the code, that would cut out the MSA server and instead use my local alignment? Did I understand you correctly?

Yes, that's correct. So instead of `colabfold_batch <fasta_file> <res_dir> ...`, you can do `colabfold_batch <a3m_file> <res_dir> ...` by modifying the lines Sarah has suggested [here](https://github.com/microsoft/bioemu/blob/3e89886143d164a9676e407db0ff01d88d8668cb/src/bioemu/get_embeds.py#L83)@sarahnlewis: To clarify: you do not need to set up an MSA server or use the changes in #38  if you already have the MSAs you want.@genesup: Thank you for all your comments.

I can confirm that this works. Pre-computed MSA must be provided as `a3m` file into the def `run_colabfold()` method of the `get_embeds.py`, as explained above. Then, colabfold would automatically skip the alignment construction using the remote server, and fast forward directly to computing embedding. Everything else works as in the original version.

Depending on the exact OS/setup that you are using you might run into some folder/file-name problems. E.g., I am running tasks on a cluster and I had to set the `tempdir` to something user-local (instead of a system tmp folder) as well as redefine `pair_rep_tempfile` and `single_rep_tempfile` file names (I had to remove "`__unknown_description_`" in both cases). However, these are minor problems that you could figure out by looking at the output error logs.

Overall, I now have a local version of bioemu that runs exclusively on my own hardware and produces a techinically sound output.

Thank you very much for the great work with bioemu!
@YuuuXie: > Thank you for all your comments.
> 
> I can confirm that this works. Pre-computed MSA must be provided as `a3m` file into the def `run_colabfold()` method of the `get_embeds.py`, as explained above. Then, colabfold would automatically skip the alignment construction using the remote server, and fast forward directly to computing embedding. Everything else works as in the original version.
> 
> Depending on the exact OS/setup that you are using you might run into some folder/file-name problems. E.g., I am running tasks on a cluster and I had to set the `tempdir` to something user-local (instead of a system tmp folder) as well as redefine `pair_rep_tempfile` and `single_rep_tempfile` file names (I had to remove "`__unknown_description_`" in both cases). However, these are minor problems that you could figure out by looking at the output error logs.
> 
> Overall, I now have a local version of bioemu that runs exclusively on my own hardware and produces a techinically sound output.
> 
> Thank you very much for the great work with bioemu!

Thanks for figuring this out! Yes you are right. I Just made a PR #48 that added the support for user's own a3m MSA file.@YuuuXie: Close via #48 

Solution patch:
diff --git a/src/bioemu/get_embeds.py b/src/bioemu/get_embeds.py
index 8c1658a..223f4c4 100644
--- a/src/bioemu/get_embeds.py
+++ b/src/bioemu/get_embeds.py
@@ -77,43 +77,69 @@ def _get_default_embeds_dir() -> StrPath:
     return os.path.join(_get_colabfold_install_dir(), "embeds_cache")
 
 
-def run_colabfold(fasta_file: StrPath, res_dir: StrPath, colabfold_env: dict[str, str]) -> int:
-    return subprocess.call(
-        [
-            "colabfold_batch",
-            fasta_file,
-            res_dir,
-            "--num-models",
-            "1",
-            "--model-order",
-            "3",
-            "--model-type",
-            "alphafold2",
-            "--num-recycle",
-            "0",
-            "--save-single-representations",
-            "--save-pair-representations",
-        ],
-        env=colabfold_env,
-    )
-
+def run_colabfold(
+    input_file: StrPath,
+    res_dir: StrPath,
+    colabfold_env: dict[str, str],
+    msa_host_url: str | None = None,
+) -> int:
+    """
+    Runs colabfold.
+    Args:
+        input_file: Input file path. It can be either a fasta or a3m file. If it is a .fasta file,
+            colabfold will retrieve an MSA from an MSA server. If it is a .a3m file,
+            the MSA in the .a3m file will be used instead.
+        res_dir: Directory to store results.
+        colabfold_env: Environment variables for colabfold.
+        msa_host_url: MSA host URL. If None, defaults to the colabfold default, which is a remote server.
+    """
 
-def get_colabfold_embeds(seq: str, cache_embeds_dir: StrPath | None) -> tuple[StrPath, StrPath]:
+    assert str(input_file).endswith(".fasta") or str(input_file).endswith(".a3m")
+
+    cmd = [
+        "colabfold_batch",
+        input_file,
+        res_dir,
+        "--num-models",
+        "1",
+        "--model-order",
+        "3",
+        "--model-type",
+        "alphafold2",
+        "--num-recycle",
+        "0",
+        "--save-single-representations",
+        "--save-pair-representations",
+    ]
+    if msa_host_url is not None:
+        cmd.extend(["--host-url", msa_host_url])
+    return subprocess.call(cmd, env=colabfold_env)
+
+
+def get_colabfold_embeds(
+    seq: str,
+    cache_embeds_dir: StrPath | None,
+    msa_file: StrPath | None = None,
+    msa_host_url: str | None = None,
+) -> tuple[StrPath, StrPath]:
     """
     Uses colabfold to retrieve embeddings for a given sequence. If the embeddings are already stored under `cache_embeds_dir`,
     the cached embeddings are returned. Otherwise, colabfold is used to compute the embeddings and they are saved under `cache_embeds_dir`.
+    Optionally uses an MSA A3M file if provided.
 
     Args:
         seq: Protein sequence to query
         cache_embeds_dir: Cache directory where embeddings will be stored. If None, defaults to a child of the colabfold install directory.
+        msa_file: MSA A3M file to use as input. If None, the sequence is used as input.
 
     Returns:
-        Path to resulting single and pair embeddings
+        Tuple of paths to single and pair embeddings.
     """
     seqsha = shahexencode(seq)
 
     # Setup embedding cache
     cache_embeds_dir = cache_embeds_dir or _get_default_embeds_dir()
+    cache_embeds_dir = os.path.expanduser(cache_embeds_dir)
     os.makedirs(cache_embeds_dir, exist_ok=True)
 
     # Check whether embeds have already been computed
@@ -121,6 +147,7 @@ def get_colabfold_embeds(seq: str, cache_embeds_dir: StrPath | None) -> tuple[St
     pair_rep_file = os.path.join(cache_embeds_dir, f"{seqsha}_pair.npy")
 
     if os.path.exists(single_rep_file) and os.path.exists(pair_rep_file):
+        logger.info(f"Using cached embeddings in {cache_embeds_dir}.")
         return single_rep_file, pair_rep_file
 
     # If we don't already have embeds, run colabfold
@@ -137,16 +164,28 @@ def get_colabfold_embeds(seq: str, cache_embeds_dir: StrPath | None) -> tuple[St
         res_dir = os.path.join(tempdir, "results")
         os.makedirs(res_dir)
         write_fasta(seqs=[seq], fasta_file=fasta_file, ids=[seqsha])
-        res = run_colabfold(fasta_file, res_dir, colabfold_env)
+        if msa_file is not None:
+            logger.info(
+                "Using user provided MSAs. This might result in suboptimal performance of model generated distributions!\n"
+                "BioEmu has been using MSAs from the ColabFold MSA server, or following:\n"
+                "https://github.com/sokrypton/ColabFold?tab=readme-ov-file#generating-msas-for-large-scale-structurecomplex-predictions.\n"
+                "If your MSA is generated differently, the generated results could be different."
+            )
+            msa_file = Path(msa_file).expanduser()
+            res = run_colabfold(msa_file, res_dir, colabfold_env)
+            embed_prefix = Path(msa_file).stem
+        else:
+            res = run_colabfold(fasta_file, res_dir, colabfold_env, msa_host_url)
+            embed_prefix = f"{seqsha}__unknown_description_"
         assert res == 0, "Failed to run colabfold_batch"
 
         single_rep_tempfile = os.path.join(
             res_dir,
-            f"{seqsha}__unknown_description__single_repr_evo_rank_001_alphafold2_model_3_seed_000.npy",
+            f"{embed_prefix}_single_repr_evo_rank_001_alphafold2_model_3_seed_000.npy",
         )
         pair_rep_tempfile = os.path.join(
             res_dir,
-            f"{seqsha}__unknown_description__pair_repr_evo_rank_001_alphafold2_model_3_seed_000.npy",
+            f"{embed_prefix}_pair_repr_evo_rank_001_alphafold2_model_3_seed_000.npy",
         )
         shutil.copy(single_rep_tempfile, single_rep_file)
         shutil.copy(pair_rep_tempfile, pair_rep_file)
diff --git a/src/bioemu/sample.py b/src/bioemu/sample.py
index b03f8a5..ca9e0ac 100644
--- a/src/bioemu/sample.py
+++ b/src/bioemu/sample.py
@@ -59,7 +59,7 @@ def maybe_download_checkpoint(
 
 @torch.no_grad()
 def main(
-    sequence: str,
+    sequence: str | Path,
     num_samples: int,
     output_dir: str | Path,
     batch_size_100: int = 10,
@@ -68,12 +68,14 @@ def main(
     model_config_path: str | Path | None = None,
     denoiser_config_path: str | Path = DEFAULT_DENOISER_CONFIG_PATH,
     cache_embeds_dir: str | Path | None = None,
+    msa_host_url: str | None = None,
 ) -> None:
     """
     Generate samples for a specified sequence, using a trained model.
 
     Args:
-        sequence: Amino acid sequence for which to generate samples or a path to a .fasta file.
+        sequence: Amino acid sequence for which to generate samples, or a path to a .fasta file, or a path to an .a3m file with MSAs.
+            If it is not an a3m file, then colabfold will be used to generate an MSA and embedding.
         num_samples: Number of samples to generate. If `output_dir` already contains samples, this function will only generate additional samples necessary to reach the specified `num_samples`.
         output_dir: Directory to save the samples. Each batch of samples will initially be dumped as .npz files. Once all batches are sampled, they will be converted to .xtc and .pdb.
         batch_size_100: Batch size you'd use for a sequence of length 100. The batch size will be calculated from this, assuming
@@ -85,6 +87,7 @@ def main(
            Only required if `ckpt_path` is set.
         denoiser_config_path: Path to the denoiser config, defining the denoising process.
         cache_embeds_dir: Directory to store MSA embeddings. If not set, this defaults to `COLABFOLD_DIR/embeds_cache`.
+        msa_host_url: MSA server URL. If not set, this defaults to colabfold's remote server. If sequence is an a3m file, this is ignored.
     """
     output_dir = Path(output_dir).expanduser().resolve()
     output_dir.mkdir(parents=True, exist_ok=True)  # Fail fast if output_dir is non-writeable
@@ -100,7 +103,13 @@ def main(
     with open(model_config_path) as f:
         model_config = yaml.safe_load(f)
 
-    # Parse FASTA file if sequence is a file path
+    # User may have provided an MSA file instead of a sequence. This will be used for embeddings.
+    msa_file = sequence if str(sequence).endswith(".a3m") else None
+
+    if msa_file is not None and msa_host_url is not None:
+        logger.warning(f"msa_host_url is ignored because MSA file {msa_file} is provided.")
+
+    # Parse FASTA or A3M file if sequence is a file path. Extract the actual sequence.
     sequence = parse_sequence(sequence)
 
     fasta_path = output_dir / "sequence.fasta"
@@ -149,6 +158,8 @@ def main(
             seed=seed,
             denoiser=denoiser,
             cache_embeds_dir=cache_embeds_dir,
+            msa_file=msa_file,
+            msa_host_url=msa_host_url,
         )
         batch = {k: v.cpu().numpy() for k, v in batch.items()}
         np.savez(npz_path, **batch, sequence=sequence)
@@ -180,6 +191,8 @@ def generate_batch(
     seed: int,
     denoiser: Callable,
     cache_embeds_dir: str | Path | None,
+    msa_file: str | Path | None = None,
+    msa_host_url: str | None = None,
 ) -> dict[str, torch.Tensor]:
     """Generate one batch of samples, using GPU if available.
 
@@ -190,13 +203,18 @@ def generate_batch(
         embeddings_file: Path to embeddings file.
         batch_size: Batch size.
         seed: Random seed.
+        msa_file: Optional path to an MSA A3M file.
+        msa_host_url: MSA server URL for colabfold.
     """
 
     torch.manual_seed(seed)
     n = len(sequence)
 
     single_embeds_file, pair_embeds_file = get_colabfold_embeds(
-        seq=sequence, cache_embeds_dir=cache_embeds_dir
+        seq=sequence,
+        cache_embeds_dir=cache_embeds_dir,
+        msa_file=msa_file,
+        msa_host_url=msa_host_url,
     )
     single_embeds = np.load(single_embeds_file)
     pair_embeds = np.load(pair_embeds_file)
diff --git a/src/bioemu/seq_io.py b/src/bioemu/seq_io.py
index f276555..6193ace 100644
--- a/src/bioemu/seq_io.py
+++ b/src/bioemu/seq_io.py
@@ -43,6 +43,7 @@ def parse_sequence(sequence: StrPath) -> str:
     """Parse sequence if sequence is a file path. Otherwise just return the input."""
     try:
         if Path(sequence).is_file():
+            # The same parser applies to both fasta and a3m files.
             rec = read_fasta(sequence)[0]
             return str(rec.seq)
     except OSError:


Test patch:
diff --git a/tests/test_embeds.py b/tests/test_embeds.py
index 24c7cae..af96164 100644
--- a/tests/test_embeds.py
+++ b/tests/test_embeds.py
@@ -10,9 +10,11 @@ from bioemu.get_embeds import StrPath, get_colabfold_embeds, shahexencode
 TEST_SEQ = "TESTSEQ"
 
 
-def mock_run_colabfold(fasta_file: StrPath, res_dir: StrPath, colabfold_env: dict[str, str]) -> int:
+def mock_run_colabfold(
+    input_file: StrPath, res_dir: StrPath, colabfold_env: dict[str, str], msa_host_url: str | None
+) -> int:
     seq = TEST_SEQ
-    seqsha = os.path.basename(fasta_file).split(".")[0]
+    seqsha = os.path.basename(input_file).split(".")[0]
     single_rep_tempfile = os.path.join(
         res_dir,
         f"{seqsha}__unknown_description__single_repr_evo_rank_001_alphafold2_model_3_seed_000.npy",