structure-prediction-boltz-2

Boltz-2 Structure Prediction

Prerequisites

How to run

Local installation

pip install boltz[cuda] -U -i

Predict protein complex structure

import os, yaml, subprocess

def predict_protein_complex_structure(sequence_1, sequence_2, project_dir):
    """
    :param sequence_1: sequence of the first protein
    :param sequence_2: sequence of the second protein
    :param project_dir: path to the project
    :return structure of the protein complex in PDB format
    """
    # init project dir
    os.makedirs(project_dir, exist_ok=True)
    log_file = os.join(project_dir, 'log.txt')
    # init input yaml file
    data = {
        "sequences": [
            {
                "protein": {
                    "id": "A",
                    "sequence": sequence_1,
                    "msa": "empty"
                }
            },
            {
                "protein": {
                    "id": "B",
                    "sequence": sequence_2,
                    "msa": "empty"
                }
            },
        ]
    }
    input_file = os.path.join(project_dir, "input.yaml")
    with open(input_file, "w") as f:
        yaml.dump(data, f)
    # init output file
    output_dir = os.path.join(project_dir, "boltz")
    # prediction
    command = [
        'boltz', 'predict', input_file,
        "--out_dir", output_dir,
        '--use_msa_server',
        '--output_format', "pdb",
        "--seed", "42"
    ]
    with open(log_file, 'a') as f:
        process = subprocess.Popen(command, stdout=f, stderr=f, env=self.env)
        process.communicate()
    process.terminate()
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        process.kill() 
        process.wait()

    # extract structure
    with open(os.path.join(output_dir, "boltz_results_input", "predictions", "input", "input_model_0.pdb"), 'r') as f:
        pred_struc = f.read()

    return pred_struc

# Predict protein complex structure for sequence_1 and sequence_2
# pred_structure is the structure prediction in PDB format
pred_structure = predict_protein_complex_structure(sequence_1, sequence_2, project_dir)