run_pretrained_openfold.py 14.7 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 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.
15
import argparse
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
16
import logging
17
import math
18
import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
19
import os
20

21
22
23
from openfold.utils.script_utils import load_models_from_command_line, parse_fasta, run_model, prep_output, \
    update_timings, relax_protein

24
25
26
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
27

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
28
import pickle
29

30
import random
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
31
32
33
import time
import torch

34
35
36
37
38
39
40
41
42
43
44
45
torch_versions = torch.__version__.split(".")
torch_major_version = int(torch_versions[0])
torch_minor_version = int(torch_versions[1])
if(
    torch_major_version > 1 or 
    (torch_major_version == 1 and torch_minor_version >= 12)
):
    # Gives a large speedup on Ampere-class GPUs
    torch.set_float32_matmul_precision("high")

torch.set_grad_enabled(False)

46
from openfold.config import model_config
47
from openfold.data import templates, feature_pipeline, data_pipeline
48
from openfold.np import residue_constants, protein
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
49
import openfold.np.relax.relax as relax
50

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
51
from openfold.utils.tensor_utils import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
52
53
    tensor_tree_map,
)
54
55
56
57
from openfold.utils.trace_utils import (
    pad_feature_dict_seq,
    trace_model_,
)
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
from datetime import date

# from scripts.utils import add_data_args

def add_data_args(parser: argparse.ArgumentParser):
    parser.add_argument(
        '--uniref90_database_path', type=str, default=None,
    )
    parser.add_argument(
        '--mgnify_database_path', type=str, default=None,
    )
    parser.add_argument(
        '--pdb70_database_path', type=str, default=None,
    )
    parser.add_argument(
        '--uniclust30_database_path', type=str, default=None,
    )
    parser.add_argument(
        '--bfd_database_path', type=str, default=None,
    )
    parser.add_argument(
        '--jackhmmer_binary_path', type=str, default='/usr/bin/jackhmmer'
    )
    parser.add_argument(
        '--hhblits_binary_path', type=str, default='/usr/bin/hhblits'
    )
    parser.add_argument(
        '--hhsearch_binary_path', type=str, default='/usr/bin/hhsearch'
    )
    parser.add_argument(
        '--kalign_binary_path', type=str, default='/usr/bin/kalign'
    )
    parser.add_argument(
        '--max_template_date', type=str,
        default=date.today().strftime("%Y-%m-%d"),
    )
    parser.add_argument(
        '--obsolete_pdbs_path', type=str, default=None
    )
    parser.add_argument(
        '--release_dates_path', type=str, default=None
    )
100

101

102
TRACING_INTERVAL = 50
103
104


Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
105
106
107
108
109
110
111
def precompute_alignments(tags, seqs, alignment_dir, args):
    for tag, seq in zip(tags, seqs):
        tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
112
        if(args.use_precomputed_alignments is None and not os.path.isdir(local_alignment_dir)):
113
            logger.info(f"Generating alignments for {tag}...")
114
115
                
            os.makedirs(local_alignment_dir)
116

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
117
118
119
120
121
122
123
124
125
126
127
128
            alignment_runner = data_pipeline.AlignmentRunner(
                jackhmmer_binary_path=args.jackhmmer_binary_path,
                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,
                pdb70_database_path=args.pdb70_database_path,
                no_cpus=args.cpus,
            )
            alignment_runner.run(
129
                tmp_fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
130
            )
131
132
133
134
        else:
            logger.info(
                f"Using precomputed alignments for {tag} at {alignment_dir}..."
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
135
136
137
138
139

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


140
141
142
143
def round_up_seqlen(seqlen):
    return int(math.ceil(seqlen / TRACING_INTERVAL)) * TRACING_INTERVAL


144
145
146
147
148
149
150
def generate_feature_dict(
    tags,
    seqs,
    alignment_dir,
    data_processor,
    args,
):
151
152
    tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
    if len(seqs) == 1:
153
        tag = tags[0]
154
155
156
157
158
159
160
        seq = seqs[0]
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
        feature_dict = data_processor.process_fasta(
            fasta_path=tmp_fasta_path, alignment_dir=local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
161
        )
162
163
164
165
166
167
168
169
170
171
172
173
    else:
        with open(tmp_fasta_path, "w") as fp:
            fp.write(
                '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
            )
        feature_dict = data_processor.process_multiseq_fasta(
            fasta_path=tmp_fasta_path, super_alignment_dir=alignment_dir,
        )

    # Remove temporary FASTA file
    os.remove(tmp_fasta_path)

174
    return feature_dict
175

176
177
def list_files_with_extensions(dir, extensions):
    return [f for f in os.listdir(dir) if f.endswith(extensions)]
178

179

180
181
182
183
def main(args):
    # Create the output directory
    os.makedirs(args.output_dir, exist_ok=True)

184
    config = model_config(args.config_preset)
185
186
187
188
189
190
191
    
    if(args.trace_model):
        if(not config.data.predict.fixed_size):
            raise ValueError(
                "Tracing requires that fixed_size mode be enabled in the config"
            )
    
192
193
194
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
195
        max_hits=config.data.predict.max_templates,
196
        kalign_binary_path=args.kalign_binary_path,
197
        release_dates_path=args.release_dates_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
198
199
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
200

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
201
    data_processor = data_pipeline.DataPipeline(
202
203
204
205
        template_featurizer=template_featurizer,
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
206
    random_seed = args.data_random_seed
207
    if random_seed is None:
208
209
210
211
212
        random_seed = random.randrange(2**32)
    
    np.random.seed(random_seed)
    torch.manual_seed(random_seed + 1)
    
213
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
214
215
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
216
    if args.use_precomputed_alignments is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
217
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
218
219
    else:
        alignment_dir = args.use_precomputed_alignments
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
220

221
222
    tag_list = []
    seq_list = []
223
    for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
224
        # Gather input sequences
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
225
        with open(os.path.join(args.fasta_dir, fasta_file), "r") as fp:
226
227
228
229
            data = fp.read()
    
        tags, seqs = parse_fasta(data)
        # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
230
        tag = '-'.join(tags)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
231

232
        tag_list.append((tag, tags))
233
234
235
236
237
        seq_list.append(seqs)

    seq_sort_fn = lambda target: sum([len(s) for s in target[1]])
    sorted_targets = sorted(zip(tag_list, seq_list), key=seq_sort_fn)
    feature_dicts = {}
238
239
240
241
242
243
244
    model_generator = load_models_from_command_line(
        config,
        args.model_device,
        args.openfold_checkpoint_path,
        args.jax_param_path,
        args.output_dir)
    for model, output_directory in model_generator:
245
        cur_tracing_interval = 0
246
        for (tag, tags), seqs in sorted_targets:
247
248
249
            output_name = f'{tag}_{args.config_preset}'
            if args.output_postfix is not None:
                output_name = f'{output_name}_{args.output_postfix}'
250
    
251
252
253
254
255
256
257
258
259
260
261
262
            # Does nothing if the alignments have already been computed
            precompute_alignments(tags, seqs, alignment_dir, args)
        
            feature_dict = feature_dicts.get(tag, None)
            if(feature_dict is None):
                feature_dict = generate_feature_dict(
                    tags,
                    seqs,
                    alignment_dir,
                    data_processor,
                    args,
                )
Sam DeLuca's avatar
Sam DeLuca committed
263

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
                if(args.trace_model):
                    n = feature_dict["aatype"].shape[-2]
                    rounded_seqlen = round_up_seqlen(n)
                    feature_dict = pad_feature_dict_seq(
                        feature_dict, rounded_seqlen,
                    )

                feature_dicts[tag] = feature_dict

            processed_feature_dict = feature_processor.process_features(
                feature_dict, mode='predict',
            )

            processed_feature_dict = {
                k:torch.as_tensor(v, device=args.model_device) 
                for k,v in processed_feature_dict.items()
            }

            if(args.trace_model):
                if(rounded_seqlen > cur_tracing_interval):
                    logger.info(
                        f"Tracing model at {rounded_seqlen} residues..."
                    )
                    t = time.perf_counter()
                    trace_model_(model, processed_feature_dict)
289
                    tracing_time = time.perf_counter() - t
290
                    logger.info(
291
                        f"Tracing time: {tracing_time}"
292
293
                    )
                    cur_tracing_interval = rounded_seqlen
Sam DeLuca's avatar
Sam DeLuca committed
294

295
            out = run_model(model, processed_feature_dict, tag, args.output_dir)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
296

297
            # Toss out the recycling dimensions --- we don't need them anymore
298
299
300
301
            processed_feature_dict = tensor_tree_map(
                lambda x: np.array(x[..., -1].cpu()), 
                processed_feature_dict
            )
302
303
304
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            unrelaxed_protein = prep_output(
305
306
307
308
                out, 
                processed_feature_dict, 
                feature_dict, 
                feature_processor, 
309
                args.config_preset,
310
311
                args.multimer_ri_gap,
                args.subtract_plddt
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
312
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
313

314
315
316
317
            unrelaxed_output_path = os.path.join(
                output_directory, f'{output_name}_unrelaxed.pdb'
            )

318
319
            with open(unrelaxed_output_path, 'w') as fp:
                fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
320

321
            logger.info(f"Output written to {unrelaxed_output_path}...")
322
            
323
324
            if not args.skip_relaxation:
                # Relax the prediction.
325
                logger.info(f"Running relaxation on {unrelaxed_output_path}...")
326
                relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name)
327

328
329
            if args.save_outputs:
                output_dict_path = os.path.join(
330
                    output_directory, f'{output_name}_output_dict.pkl'
331
332
333
                )
                with open(output_dict_path, "wb") as fp:
                    pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
334

Sam DeLuca's avatar
Sam DeLuca committed
335
                logger.info(f"Model output written to {output_dict_path}...")
336

337

338
339
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
340
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
341
342
        "fasta_dir", type=str,
        help="Path to directory containing FASTA files, one sequence per file"
343
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
344
345
346
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
347
348
349
350
    parser.add_argument(
        "--use_precomputed_alignments", type=str, default=None,
        help="""Path to alignment directory. If provided, alignment computation 
                is skipped and database path arguments are ignored."""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
351
    )
352
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
353
354
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
355
356
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
357
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
358
359
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
360
361
    )
    parser.add_argument(
362
        "--config_preset", type=str, default="model_1",
363
        help="""Name of a model config preset defined in openfold/config.py"""
364
365
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
366
367
368
369
370
371
372
373
374
        "--jax_param_path", type=str, default=None,
        help="""Path to JAX model parameters. If None, and openfold_checkpoint_path
             is also None, parameters are selected automatically according to 
             the model name from openfold/resources/params"""
    )
    parser.add_argument(
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
375
    )
376
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
377
        "--save_outputs", action="store_true", default=False,
378
379
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
380
381
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
382
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
383
    )
384
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
385
        "--preset", type=str, default='full_dbs',
386
387
        choices=('reduced_dbs', 'full_dbs')
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
388
389
390
391
    parser.add_argument(
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
392
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
393
394
395
396
        "--data_random_seed", type=str, default=None
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
397
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
398
399
400
401
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
402
403
404
405
406
407
    parser.add_argument(
        "--trace_model", action="store_true", default=False,
        help="""Whether to convert parts of each model to TorchScript.
                Significantly improves runtime at the cost of lengthy
                'compilation.' Useful for large batch jobs."""
    )
408
409
410
411
412
    parser.add_argument(
        "--subtract_plddt", action="store_true", default=False,
        help=""""Whether to output (100 - pLDDT) in the B-factor column instead
                 of the pLDDT itself"""
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
413
    add_data_args(parser)
414
415
    args = parser.parse_args()

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
416
417
    if(args.jax_param_path is None and args.openfold_checkpoint_path is None):
        args.jax_param_path = os.path.join(
418
            "openfold", "resources", "params", 
419
            "params_" + args.config_preset + ".npz"
420
421
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
422
423
424
425
426
427
    if(args.model_device == "cpu" and torch.cuda.is_available()):
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )

428
    main(args)