Work for newly released AF3

[1326093445a]

Hello All

Thanks for your wonderful work. I wonder with the newly released AlphaFold3(AF3), will the unmask still work for that?

[clami66]

Hi,

I have recently tested the "unmasked" approach on AF3 and it does seem to work. Just need to rewrite the template preparation script for it and make sure I can release it without breaking the terms of their license.

[clami66]

Hi,

I have now pushed a draft version of AF_unmasked that is compatible with AlphaFold3. You can find it in the af3 branch.

You first need to run prepare_templates_af3.py similarly to how you do for AF2, this will output a fold_input.json file to use as input in AF3.

Then, you need to run a patched version of AF3 to include the --cross_chain_templates flag necessary to use complete multimeric templates. The patch file is here.

So far the template preparation does not include the possibility to add e.g. MSA (paired or unpaired) from another json file, but it is easy to do so manually by e.g. setting `unpairedMsaPath" to a desired a3m file.

Let me know if this helps and if you have questions or comments

[ntnn19]

Hi,

Thank you for developing AF_unmasked and providing compatibility with alphafold3.

I tried to follow the instructions here but was not able to run af_unmasked with alphafold3. I built the alphafold3 singularity container with patched run_alphafold.py, and try to run it with the json file that is generated by running prepare_templates_af3.py with the example as input, but got the following error:

0527 14:13:10.540634 22542279331840 xla_bridge.py:895] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0527 14:13:10.548691 22542279331840 xla_bridge.py:895] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
Traceback (most recent call last):
  File "/app/alphafold/run_alphafold.py", line 754, in <module>
    app.run(main)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 735, in main
    for fold_input in fold_inputs:
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/common/folding_input.py", line 1195, in load_fold_inputs_from_path
    yield Input.from_json(json_str, json_path)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/common/folding_input.py", line 876, in from_json
    chains.append(ProteinChain.from_dict(sequence, json_path, seq_id))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/common/folding_input.py", line 261, in from_dict
    zip(raw_template['queryIndices'], raw_template['templateIndices'])
        ~~~~~~~~~~~~^^^^^^^^^^^^^^^^
KeyError: 'queryIndices'

The content of the json file AF_models/H1142/fold_input.json is:
{ "dialect": "alphafold3", "modelSeeds": [ 0 ], "name": "H1142", "sequences": [ { "protein": { "id": "A", "pairedMsa": "", "sequence": "GLEKDFLPLYFGWFLTKKSSETLRKAGQVFLEELGNHKAFKKELRHFISGDEPKEKLELVSYFGKRPPGVLHCTTKFCDYKAAGAEEYAQQEVVKRSYGKAFKLSISALFVTPKTAGAQVVLTDQELQLWPSDLDKPSASEGLPPGSRAHVTLGCAADVQPVQTGLDLLDILQQVKGGSQGEAVGELPRGKLYSLGKGRWMLSLTKKMEVKAIFTGYYG", "templates": [ { "mmcifPath": "<abs_path_to_this_repo>/H1142/template_data/mmcif_files/0000_B.cif" } ], "unpairedMsa": "" } }, { "protein": { "id": "B", "pairedMsa": "", "sequence": "EVQLEESGGGLVQAGGSLTLSCAASGFTFDDYAMGWYRQAPGKERVGVSCISRTDGYTYYLDSVKGRFTISTDHAKHTVYLQMNNLKPDDTGLYYCAADADPEYGSRCPDPYYGMDYWGKGILVTVSS", "templates": [ { "mmcifPath": "<abs_path_to_this_repo>/H1142/template_data/mmcif_files/0000_C.cif" } ], "unpairedMsa": "" } } ], "version": 2 }

the command I used to run the container is:
singularity run \ --nv -B $(realpath AF_models/H1142):/root/af_input \ -B $(realpath AF_models/H1142/af_output_unmasked):/root/af_output \ -B $(realpath ../../singularity_containers/alphafold3/weights):/root/models \ -B $(realpath ../../af3db:/root/public_databases \ ../../natan/singularity_containers/af2_unmasked/patched_af3_commit_9befb639292fb53fa36b5a3e6ee0273e5a1d1094_with_af_unmasked.sif \ python /app/alphafold/run_alphafold.py --json_path=/root/af_input/fold_input.json --model_dir=/root/models \ --output_dir=/root/af_output --db_dir=/root/public_databases --run_data_pipeline=true --run_inference=true --cross_chain_templates=true

I also tried to apply the patch on the run_alphafold.py found on the branch af3 of this repo but got the following error:

git apply AF3_cross_chain_templates.patch run_alphafold.py

error: patch failed: run_alphafold.py:199
error: run_alphafold.py: patch does not apply`
error: src/alphafold3/data/featurisation.py: No such file or directory
error: src/alphafold3/model/diffusion/model.py: No such file or directory                                                                                                                                         
error: src/alphafold3/model/features.py: No such file or directory          
error: src/alphafold3/model/pipeline/pipeline.py: No such file or directory 

I'm probably missing something basic and would appreciate your help with this.

[clami66]

Hi,

run it with the json file that is generated by running prepare_templates_af3.py with the example as input, but got the following error:

It looks like it's missing the alignment indices. I believe you need to run prepare_templates_af3.py with the --align flag, otherwise the indices will not be in the json file. Please try again and see what happens.

I also tried to apply the patch on the run_alphafold.py found on the branch af3 of this repo but got the following error:

The patch should be applied on your installation of AlphaFold3 rather than on the AF_unmasked installation. If you have a singularity image for AF3 with the patch applied, then you should be already set

[ntnn19]

It works perfectly with the --align flag. Thank you very much.

[clami66]

@ntnn19 Great! If you can, let me know if the templates work as expected

[ntnn19]

I believe it works well. The predicted structure is the green one, the other structures are the templates from the example:

[ntnn19]

Hi @clami66,

I'm encountering an issue I hope you can help with. When I run with the --cross_chain_templates=true flag, the error below occurs. The same setup works fine when the flag is omitted.
For context, I'm using af_unmasked in af3 to predict multimeric structures for all pairwise and three-way combinations of three protein chains, i.e., 1A1B, 1A1C, 1B1C, …, 3C. The issue only appears with a specific stoichiometry, 2A1B.

Activating singularity image ../../../../../singularity_containers/af2_unmasked/patched_af3_commit_9befb639292fb53fa36b5a3e6ee0273e5a1d1094_with_af_unmasked.sif
I0604 09:05:44.355641 22625054720000 xla_bridge.py:895] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0604 09:05:44.363523 22625054720000 xla_bridge.py:895] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
/alphafold3_venv/lib/python3.11/site-packages/jax/_src/compiler.py:691: UserWarning: Error reading persistent compilation cache entry for 'jit_convert_element_type': OSError: [Errno 30] Read-only file system: '/path'
  warnings.warn(
I0604 09:05:48.040314 22625054720000 pipeline.py:468] Skipping MSA and template search for protein chain A because it already has MSAs and templates.
I0604 09:05:48.040501 22625054720000 pipeline.py:474] Using empty unpaired MSA for protein chain A
I0604 09:05:48.040540 22625054720000 pipeline.py:476] Using empty paired MSA for protein chain A
I0604 09:05:48.041118 22625054720000 pipeline.py:468] Skipping MSA and template search for protein chain B because it already has MSAs and templates.
I0604 09:05:48.041161 22625054720000 pipeline.py:474] Using empty unpaired MSA for protein chain B
I0604 09:05:48.041198 22625054720000 pipeline.py:476] Using empty paired MSA for protein chain B
I0604 09:05:48.041803 22625054720000 pipeline.py:468] Skipping MSA and template search for protein chain C because it already has MSAs and templates.
I0604 09:05:48.041846 22625054720000 pipeline.py:474] Using empty unpaired MSA for protein chain C
I0604 09:05:48.041873 22625054720000 pipeline.py:476] Using empty paired MSA for protein chain C
I0604 09:05:53.361323 22625054720000 pipeline.py:164] processing 2a_1b_without_ligand, random_seed=0
I0604 09:05:53.468577 22625054720000 pipeline.py:257] Calculating bucket size for input with 2198 tokens.
I0604 09:05:53.468702 22625054720000 pipeline.py:263] Got bucket size 2560 for input with 2198 tokens, resulting in 362 padded tokens.
Running AlphaFold 3. Please note that standard AlphaFold 3 model parameters are
only available under terms of use provided at
https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md.
If you do not agree to these terms and are using AlphaFold 3 derived model
parameters, cancel execution of AlphaFold 3 inference with CTRL-C, and do not
use the model parameters.
Found local devices: [CudaDevice(id=0)]
Building model from scratch...
Processing fold inputs.
Processing fold input #1
Processing fold input 2a_1b_without_ligand
Checking we can load the model parameters...
Running data pipeline...
Processing chain A
Processing chain A took 0.00 seconds
Processing chain B
Processing chain B took 0.00 seconds
Processing chain C
Processing chain C took 0.00 seconds
Output directory: /root/af_output/rule_AF3_DATA_PLUS_INFERENCE/without_ligand/2a_1b_without_ligand
Writing model input JSON to /root/af_output/rule_AF3_DATA_PLUS_INFERENCE/without_ligand/2a_1b_without_ligand
Predicting 3D structure for 2a_1b_without_ligand for seed(s) (0,)...
Featurising data for seeds (0,)...
Featurising 2a_1b_without_ligand with rng_seed 0.
Traceback (most recent call last):
  File "/app/alphafold/run_alphafold.py", line 754, in <module>
    app.run(main)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 737, in main
    process_fold_input(
  File "/app/alphafold/run_alphafold.py", line 593, in process_fold_input
    all_inference_results = predict_structure(
                            ^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 386, in predict_structure
    featurised_examples = featurisation.featurise_input(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/data/featurisation.py", line 77, in featurise_input
    batch = data_pipeline.process_item(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/model/pipeline/pipeline.py", line 342, in process_item
    batch_templates = features.Templates.compute_features(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alphafold3_venv/lib/python3.11/site-packages/alphafold3/model/features.py", line 816, in compute_features
    np_example[feature_name] = v[:max_templates, standard_token_idxs, ...]
                               ~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 2022 is out of bounds for axis 1 with size 2022
 

[clami66]

Hi again,

Is this the first time you try this with a homomeric template (i.e. A2 in this case)? Here it looks like you "run out" of templates, so I wonder if the second template for chain A is ignored somehow.

[ntnn19]

Thank you for your message.

Just to clarify, my template is actually a heterotrimer, not a homomeric A2 complex. I'm trying to predict all pairwise and three-way combinations of three protein chains, i.e., 1A1B, 1A1C, 1B1C, …, 3C , using a template that consists of chains A, B, and C. Is it possible to use AF unmasked for this case?

[clami66]

Yes, it is possible, but possibly untested with the AF3 variant. I was just wondering if the A2B1 was the only case where this failed. Did it work, for example, with B2C1 or A1C2 and so on?

[ntnn19]

It worked with all the other three-way combinations. Interestingly, no error was raised for A2B1, when I manually changed the value of the "id" key in the fold_input.json file for the second copy of chain A, from "B" to "X".

[clami66]

It probably depends on how you run prepare_templates. In your case, you have templates for A, B and C but then you would like to have two copies for, say, chain A. In this case I would run:

python prepare_templates.py --target target.fasta \
    --template template.pdb \
    --output_dir AF_models/ \
    --target_chains A B C \
    --template_chains A - B

Where in target.fasta you have:

> chain_A
SEQA...
> chain_B
SEQA...
> chain_C
SEQB...

and in template.pdb you have the three chains in the A B C order.

This is because for the second copy of the target chain A you don't have a correspondent template chain, so you map it to the special - symbol in --template_chains A - B

[ntnn19]

Thank you for your suggestion. I will try it and let you know if it works.