Support dummy sequence MSA

openfold3/core/data/framework/data_module.py

@@ -68,6 +68,7 @@
 from openfold3.core.data.pipelines.preprocessing.template import TemplatePreprocessor
 from openfold3.core.data.tools.colabfold_msa_server import (
     MsaComputationSettings,
+    augment_main_msa_with_query_sequence,
     preprocess_colabfold_msas,
 )
 from openfold3.core.utils.tensor_utils import dict_multimap
@@ -521,6 +522,11 @@ def prepare_data(self) -> None:
                 inference_query_set=self.inference_config.query_set,
                 compute_settings=self.msa_computation_settings,
             )
+        else:
+            self.inference_config.query_set = augment_main_msa_with_query_sequence(
+                inference_query_set=self.inference_config.query_set,
+                compute_settings=self.msa_computation_settings,
+            )
 
         if self.use_templates:
             template_preprocessor = TemplatePreprocessor(

openfold3/core/data/tools/colabfold_msa_server.py

@@ -1044,3 +1044,39 @@ def preprocess_colabfold_msas(
     )
 
     return inference_query_set
+
+
+def augment_main_msa_with_query_sequence(
+    inference_query_set: InferenceQuerySet,
+    compute_settings: MsaComputationSettings,
+) -> InferenceQuerySet:
+    output_directory = compute_settings.msa_output_directory
+    for query_name, query in inference_query_set.queries.items():
+        for chain in query.chains:
+            if (
+                chain.molecule_type == MoleculeType.PROTEIN
+                or chain.molecule_type == MoleculeType.RNA
+            ) and chain.main_msa_file_paths is None:
+                dummy_msa_file_path = (
+                    output_directory
+                    / "dummy"
+                    / f"{get_sequence_hash(chain.sequence)}.npz"
+                )
+                dummy_msa_file_path.parent.mkdir(exist_ok=True, parents=True)
+                dummy_msa = ">query\n" + chain.sequence
+                msas_preparsed = {"dummy": parse_a3m(dummy_msa).to_dict()}
+                np.savez_compressed(dummy_msa_file_path, **msas_preparsed)
+                chain.main_msa_file_paths = [dummy_msa_file_path]
+
+                chain_ids = ",".join(chain.chain_ids)
+                warnings.warn(
+                    (
+                        f"Expected MSA file for chain {chain_ids} of "
+                        f"type {chain.molecule_type.name} in query "
+                        f"{query_name}, but no MSA files found. Query sequence "
+                        "will be used as dummy MSA for this chain."
+                    ),
+                    stacklevel=2,
+                )
+
+    return inference_query_set

openfold3/projects/of3_all_atom/config/dataset_config_components.py

@@ -100,6 +100,7 @@ class MSASettings(BaseModel):
         "nt_hits",
         "concat_cfdb_uniref100_filtered",
         "colabfold_main",
+        "dummy",
     ]
     paired_msa_order: list = ["colabfold_paired"]