run_pretrained_alphafold.py 8.22 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
import argparse
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
17
from datetime import date
18
import pickle
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
19
import os
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
20
21

# A hack to get OpenMM and PyTorch to peacefully coexist
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
22
23
os.environ["OPENMM_DEFAULT_PLATFORM"] = "OpenCL"

24
25
26
import random
import sys

27
28
from openfold.features import templates, feature_pipeline
from openfold.features.np import data_pipeline
29

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
30
import time
31
32

import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
33
34
import torch

35
from openfold.config import model_config
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
36
from openfold.model.model import AlphaFold
37
from openfold.np import residue_constants, protein
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
38
39
import openfold.np.relax.relax as relax
from openfold.utils.import_weights import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
40
41
    import_jax_weights_,
)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
42
from openfold.utils.tensor_utils import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
43
44
45
    tensor_tree_map,
)

46
47
MAX_TEMPLATE_HITS = 20

48
49
50
51
52
53
def main(args):
    config = model_config(args.model_name)
    model = AlphaFold(config.model)
    model = model.eval()
    import_jax_weights_(model, args.param_path)
    model = model.to(args.device)
54
    
55
56
57
58
59
60
61
62
63
64
    # FEATURE COLLECTION AND PROCESSING
    use_small_bfd = args.preset == "reduced_dbs"
    num_ensemble = 1

    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
        max_hits=MAX_TEMPLATE_HITS,
        kalign_binary_path=args.kalign_binary_path,
        release_dates_path=None,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
65
66
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
67

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
68
    alignment_runner = data_pipeline.AlignmentRunner(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
69
        jackhmmer_binary_path=args.jackhmmer_binary_path,
70
71
72
73
74
75
76
77
        hhblits_binary_path=args.hhblits_binary_path,
        hhsearch_binary_path=args.hhsearch_binary_path,
        uniref90_database_path=args.uniref90_database_path,
        mgnify_database_path=args.mgnify_database_path,
        bfd_database_path=args.bfd_database_path,
        uniclust30_database_path=args.uniclust30_database_path,
        small_bfd_database_path=args.small_bfd_database_path,
        pdb70_database_path=args.pdb70_database_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
78
79
80
81
        use_small_bfd=use_small_bfd,
    )

    data_processor = data_pipeline.DataPipeline(
82
83
84
85
86
        template_featurizer=template_featurizer,
        use_small_bfd=use_small_bfd
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
87
    random_seed = args.data_random_seed
88
89
90
91
92
93
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
    config.data.eval.num_ensemble = num_ensemble
    feature_processor = feature_pipeline.FeaturePipeline(config)
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
94
95
96
    alignment_dir = os.path.join(output_dir_base, "alignments")
    if not os.path.exists(alignment_dir):
        os.makedirs(alignment_dir)
97
98

    print("Collecting data...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
99
100
101
102
103
104
105
    alignment_runner.run_from_fasta(
        args.fasta_path, alignment_dir
    )     

    feature_dict = data_processor.process_fasta(
        input_fasta_path=args.fasta_path, alignment_dir=alignment_dir
    )
106
107

    print("Generating features...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
108
109
110
    processed_feature_dict = feature_processor.process_features(
        feature_dict, random_seed
    )
111
112
113

    print("Executing model...")
    batch = processed_feature_dict
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
    with torch.no_grad():
        batch = {
            k:torch.as_tensor(v, device=args.device) 
            for k,v in batch.items()
        }
        
        longs = [
            "aatype", 
            "template_aatype", 
            "extra_msa", 
            "residx_atom37_to_atom14",
            "residx_atom14_to_atom37",
            "true_msa",
            "residue_index",
        ]
        for l in longs:
            batch[l] = batch[l].long()
        
        # Move the recycling dimension to the end
        move_dim = lambda t: t.permute(*range(len(t.shape))[1:], 0)
        batch = tensor_tree_map(move_dim, batch)
        make_contig = lambda t: t.contiguous()
        batch = tensor_tree_map(make_contig, batch)
    
        t = time.time()
        out = model(batch)
        print(f"Inference time: {time.time() - t}")
    
    # Toss out the recycling dimensions --- we don't need them anymore
    batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
    out = tensor_tree_map(lambda x: np.array(x.cpu()), out)
    
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)
    
    plddt_b_factors = np.repeat(
        plddt[..., None], residue_constants.atom_type_num, axis=-1
    )
    
    unrelaxed_protein = protein.from_prediction(
        features=batch,
        result=out,
        b_factors=plddt_b_factors
    )
    
    os.environ["CUDA_VISIBLE_DEVICES"] = "7"
    
    amber_relaxer = relax.AmberRelaxation(
        **config.relax
    )
    
    # Relax the prediction.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
166
    t = time.time()
167
168
169
170
171
172
173
174
175
176
177
178
179
    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
    print(f"Relaxation time: {time.time() - t}")
    
    # Save the relaxed PDB.
    relaxed_output_path = os.path.join(
        args.output_dir, f'relaxed_{args.model_name}.pdb'
    )
    with open(relaxed_output_path, 'w') as f:
        f.write(relaxed_pdb_str)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
180
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
181
        "fasta_path", type=str,
182
183
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
184
        'uniref90_database_path', type=str, 
185
    )
186
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
187
        'mgnify_database_path', type=str, 
188
189
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
190
        'pdb70_database_path', type=str,
191
192
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
193
        'template_mmcif_dir', type=str,
194
    )
195
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
196
        '--bfd_database_path', type=str, default=None,
197
198
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
199
        '--small_bfd_database_path', type=str, default=None
200
201
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
202
        '--uniclust30_database_path', type=str, default=None
203
204
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
205
        '--jackhmmer_binary_path', type=str, default='/usr/bin/jackhmmer'
206
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
207
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
208
        '--hhblits_binary_path', type=str, default='/usr/bin/hhblits'
209
210
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
211
        '--hhsearch_binary_path', type=str, default='/usr/bin/hhsearch'
212
213
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
214
        '--kalign_binary_path', type=str, default='/usr/bin/kalign'
215
216
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
217
218
        '--max_template_date', type=str, 
        default=date.today().strftime("%Y-%m-%d"),
219
220
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
221
        '--obsolete_pdbs_path', type=str, default=None
222
223
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
224
225
226
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
        required=True
227
228
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
229
230
231
        "--device", type=str, default="cpu",
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
232
233
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
234
235
236
        "--model_name", type=str, default="model_1",
        help="""Name of a model config. Choose one of model_{1-5} or 
             model_{1-5}_ptm, as defined on the AlphaFold GitHub."""
237
238
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
239
240
241
242
        "--param_path", type=str, default=None,
        help="""Path to model parameters. If None, parameters are selected
             automatically according to the model name from 
             openfold/resources/params"""
243
244
    )
    parser.add_argument(
245
        '--preset', type=str, default='full_dbs',
246
247
248
        choices=('reduced_dbs', 'full_dbs')
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
249
        '--data_random_seed', type=str, default=None
250
    )
251
252
253
254
255
256
257
258
259

    args = parser.parse_args()

    if(args.param_path is None):
        args.param_path = os.path.join(
            "openfold", "resources", "params", 
            "params_" + args.model_name + ".npz"
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
260
261
262
263
264
265
266
    if(args.bfd_database_path is None and 
       args.small_bfd_database_path is None):
        raise ValueError(
            "At least one of --bfd_database_path or --small_bfd_database_path"
            "must be specified"
        )

267
    main(args)