Training code and saving multiple poses

[rakeshr10]

Hi, I was wondering if the training code and documentation on how to train this model will be made public. I also noticed that the pinder checkpoint file doesn't seem to work.

Is there a way to save top 10 number of docked poses from the sampling process. In your paper you describe that the sampling is pretty good but the scoring is poor. Are there efforts to use something like an ipTM metric to rank the solutions which in the case of AF2 has shown promising results.

Regards
Rakesh

[aggelos-michael-papadopoulos]

Hi @rakeshr10. I have created a PR for the number of output poses, so you can follow it there. What I did pretty much was to add the following def inside the inference_base.py:

def inference_multiple_poses(in_pdb_1, in_pdb_2, num_samples=40, output_dir="output_poses"):
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Load ESM model
    esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
    batch_converter = alphabet.get_batch_converter()
    esm_model = esm_model.to(device).eval()

    # Load score model
    model = Score_Model.load_from_checkpoint(
        str(Path("./checkpoints/dips/model_0.ckpt")),
        map_location=device,
    )
    model.to(device).eval()

    # Load PDBs
    receptor = get_info_from_pdb(in_pdb_1)
    ligand = get_info_from_pdb(in_pdb_2)

    # Prepare inputs
    inputs = {"receptor": receptor, "ligand": ligand}
    batch = get_batch_from_inputs(inputs, batch_converter, esm_model, device)
    batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

    # Define parameters
    num_steps = 40
    use_clash_force = False

    all_outputs = []

    # Run sampling
    for i in range(num_samples):
        rec_pos, lig_pos, rot_update, tr_update, outputs = Euler_Maruyama_sampler(
            model=model,
            batch=batch.copy(),
            num_steps=num_steps,
            device=device,
            use_clash_force=use_clash_force,
        )

        lig_aa_coords = modify_aa_coords(ligand["aa_coords"], inputs["ligand"]["bb_coords"], rot_update, tr_update)
        rec_structure = receptor["structure"]
        lig_structure = ligand["structure"]
        lig_structure.coord = lig_aa_coords

        complex_structure = combine_atom_arrays(rec_structure, lig_structure)

        file = PDBFile()
        file.set_structure(complex_structure)
        output_pdb_path = os.path.join(output_dir, f"output_pose_{i + 1}.pdb")
        file.write(output_pdb_path)

        all_outputs.append({"pose_id": i + 1, "energy": outputs["energy"].item(), "pdb_path": output_pdb_path})

    return all_outputs

then I have created a src/inference_multiple_poses.py with the below content:

from inference_base import inference_multiple_poses

def parse_args():
    parser = argparse.ArgumentParser(description="Process two required PDB files.")
    parser.add_argument("pdb_1", type=str, help="Path to the first PDB file")
    parser.add_argument("pdb_2", type=str, help="Path to the second PDB file")
    parser.add_argument("--num_samples", type=int, default=40, help="Number of output poses/samples, default=40")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    inference_multiple_poses(args.pdb_1, args.pdb_2, args.num_samples)

and finally you can execute this by:

python src/inference_multiple_poses.py path_to_input_pdb_1 path_to_input_pdb_2 --num_samples 16