sokrypton / ColabDesign

Making Protein Design accessible to all via Google Colab!
529 stars 118 forks source link

How to continue a design run? #157

Open amin-sagar opened 1 year ago

amin-sagar commented 1 year ago

Hello. I am working on generating binders for a target which seems like a difficult one as most runs give binders with very low i_ptm. I thought I could increase the throughput by checking the i_ptm after a certain number of iterations and only continue those which have reached a certain value of i_ptm. My code is as follows.

for i in tqdm(range (0,100)):
    print (i)
    af_model.restart()
    af_model.design_pssm_semigreedy(80,0,num_recycles=num_recycle, dropout=True)
    #Get the mean iPTM after soft iterations
    iptm_value = (af_model._tmp['log'][-1]['i_ptm'])
    if iptm_value < 0.35:
        print(f"The i_PTM has a value of {iptm_value}\n")
        print ("Stopping this iteration\n")
        continue
    else:
        print(f"The i_PTM has a value of {iptm_value}\n")
        print ("Continuing this iteration\n")
        af_model.design_pssm_semigreedy(80,32,num_recycles=num_recycle, dropout=True)

However, this seems to restart the optimization. For example, please have a look at the following output

73 models [2] recycles 4 hard 0 soft 0.91 temp 1 loss 3.66 i_con 3.60 plddt 0.33 ptm 0.88 i_ptm 0.39
74 models [4] recycles 4 hard 0 soft 0.93 temp 1 loss 3.03 i_con 2.97 plddt 0.40 ptm 0.89 i_ptm 0.48
75 models [0] recycles 4 hard 0 soft 0.94 temp 1 loss 2.76 i_con 2.70 plddt 0.40 ptm 0.89 i_ptm 0.52
76 models [3] recycles 4 hard 0 soft 0.95 temp 1 loss 3.89 i_con 3.83 plddt 0.32 ptm 0.88 i_ptm 0.33
77 models [2] recycles 4 hard 0 soft 0.96 temp 1 loss 3.67 i_con 3.60 plddt 0.34 ptm 0.88 i_ptm 0.38
78 models [4] recycles 4 hard 0 soft 0.97 temp 1 loss 3.06 i_con 3.00 plddt 0.40 ptm 0.89 i_ptm 0.49
79 models [4] recycles 4 hard 0 soft 0.99 temp 1 loss 2.90 i_con 2.84 plddt 0.43 ptm 0.89 i_ptm 0.50
80 models [0] recycles 4 hard 0 soft 1 temp 1 loss 3.09 i_con 3.03 plddt 0.38 ptm 0.88 i_ptm 0.44
The i_PTM has a value of 0.44231119751930237

Continuing this iteration

Stage 1: running (logits → soft)
81 models [0] recycles 0 hard 0 soft 0.01 temp 1 loss 4.83 i_con 4.75 plddt 0.25 ptm 0.85 i_ptm 0.11
82 models [4] recycles 0 hard 0 soft 0.03 temp 1 loss 4.19 i_con 4.11 plddt 0.18 ptm 0.27 i_ptm 0.34
83 models [0] recycles 0 hard 0 soft 0.04 temp 1 loss 4.90 i_con 4.82 plddt 0.24 ptm 0.85 i_ptm 0.12
84 models [4] recycles 0 hard 0 soft 0.05 temp 1 loss 5.68 i_con 5.60 plddt 0.21 ptm 0.84 i_ptm 0.10
85 models [2] recycles 0 hard 0 soft 0.06 temp 1 loss 5.44 i_con 5.36 plddt 0.23 ptm 0.82 i_ptm 0.08
86 models [0] recycles 0 hard 0 soft 0.07 temp 1 loss 4.82 i_con 4.74 plddt 0.23 ptm 0.85 i_ptm 0.10
87 models [0] recycles 0 hard 0 soft 0.09 temp 1 loss 4.67 i_con 4.59 plddt 0.23 ptm 0.85 i_ptm 0.12
88 models [1] recycles 0 hard 0 soft 0.10 temp 1 loss 4.78 i_con 4.71 plddt 0.22 ptm 0.86 i_ptm 0.10

The number of iteration continues but the loss increases abruptly indicating the optimization is not continuing from the last step. How can I continue the optimization from the last step? I would be really grateful for any suggestions. Best, Amin.