Multimer data pipeline merging sequences that are from different taxonomy?

Expected Behavior

I am using colabfold.batch.generate_input_feature function for a protein complex using parameters msa_mode = "unpaired+paired" and model_type = "AlphaFold2-multimer-v2". I expected the output input_features to have a feature called input_features["msa_all_seqs"] which pairs up the multiple chain msas into one msa at axis 1, as done by function call to colabfold.batch.process_multimer_features

Current Behavior

The function generate_input_feature calls build_multimer_feature on paired msas. The function then calls alphafold.data.pipeline.make_msa_features. However, make_msa_features removes any duplicate sequences from paired msas of each chain separately resulting in different lengths of paired msas per chain and hence sequences no longer aligning the "paired sequences". I am also getting an error message at merging features stage where the two msas do not have the same length paired msas but only for paired msas that are shorter than 1024 as alphafold.data.feature_processing.crop_chains called by colabfold.batch.process_multimer_features crops out msas at msa_crop_size // 2 which is by default 2048 from the model config. But even when it is cropped out so that the lenghts of the msas are matched, this still could have paired sequences that are no longer from the same branch of taxonomy due to that independent sequence duplicate remove step in make_msa_features

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [31], in <cell line: 1>()
      1 for job_number, (raw_jobname, query_sequence, a3m_lines) in enumerate(queries):
      2     (
      3         unpaired_msa, 
      4         paired_msa, 
   (...)
      7         template_features
      8     ) = unserialize_msa(a3m_lines, query_sequence)
     10     (input_features, domain_names) = \
---> 11         generate_input_feature(
     12             query_seqs_unique = query_seqs_unique,
     13             query_seqs_cardinality = query_seqs_cardinality,
     14             unpaired_msa = unpaired_msa,
     15             paired_msa = paired_msa,
     16             template_features = template_features,
     17             is_complex = is_complex,
     18             model_type = model_type,
     19         )
     20     break

File /lustre07/scratch/sunyun/workspace/colabfold-pipeline/src/colabfold/batch.py:1058, in generate_input_feature(query_seqs_unique, query_seqs_cardinality, unpaired_msa, paired_msa, template_features, is_complex, model_type)
   1050     domain_names = {
   1051         protein.PDB_CHAIN_IDS[0]: [
   1052             name.decode("UTF-8")
   (...)
   1055         ]
   1056     }
   1057 elif model_type.startswith("AlphaFold2-multimer"):
-> 1058     input_feature = process_multimer_features(features_for_chain)
   1059     domain_names = {
   1060         chain: [
   1061             name.decode("UTF-8")
   (...)
   1065         for (chain, feature) in features_for_chain.items()
   1066     }
   1067 elif is_complex and model_type == "AlphaFold2-ptm":

File /lustre07/scratch/sunyun/workspace/colabfold-pipeline/src/colabfold/batch.py:943, in process_multimer_features(features_for_chain)
    938 common_features = set([*np_chains_list[0]]).intersection(*np_chains_list)
    939 np_chains_list = [
    940     {key: value for (key, value) in chain.items() if key in common_features}
    941     for chain in np_chains_list
    942 ]
--> 943 np_example = feature_processing.msa_pairing.merge_chain_features(
    944     np_chains_list=np_chains_list,
    945     pair_msa_sequences=pair_msa_sequences,
    946     max_templates=feature_processing.MAX_TEMPLATES,
    947 )
    948 np_example = feature_processing.process_final(np_example)
    950 # Pad MSA to avoid zero-sized extra_msa.

File /lustre07/scratch/sunyun/workspace/colabfold-pipeline/src/alphafold/data/msa_pairing.py:435, in merge_chain_features(np_chains_list, pair_msa_sequences, max_templates)
    432 np_chains_list = _merge_homomers_dense_msa(np_chains_list)
    433 # Unpaired MSA features will be always block-diagonalised; paired MSA
    434 # features will be concatenated.
--> 435 np_example = _merge_features_from_multiple_chains(
    436     np_chains_list, pair_msa_sequences=False)
    437 if pair_msa_sequences:
    438   np_example = _concatenate_paired_and_unpaired_features(np_example)

File /lustre07/scratch/sunyun/workspace/colabfold-pipeline/src/alphafold/data/msa_pairing.py:360, in _merge_features_from_multiple_chains(chains, pair_msa_sequences)
    358 if feature_name_split in MSA_FEATURES:
    359   if pair_msa_sequences or '_all_seq' in feature_name:
--> 360     merged_example[feature_name] = np.concatenate(feats, axis=1)
    361   else:
    362     merged_example[feature_name] = block_diag(
    363         *feats, pad_value=MSA_PAD_VALUES[feature_name])

File <__array_function__ internals>:5, in concatenate(*args, **kwargs)

ValueError: all the input array dimensions for the concatenation axis must match exactly, but along dimension 0, the array at index 0 has size 268 and the array at index 1 has size 362

Thank you in advance!

Start in colabfold/batch.py at generate_input_feature, build_multimer_feature, and process_multimer_features, then inspect alphafold/data/pipeline.py's make_msa_features and msa_pairing.py's merge_chain_features. Reproduce the mismatch with paired MSAs and trace duplicate removal and merging. Done means paired sequences retain their alignment and multimer feature generation no longer fails on differing chain MSA lengths.