run_pretrained_openfold.py 18 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
3
#
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
4
5
6
7
8
9
10
11
12
13
14
# 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
import pickle
import random
import time
23
import json
24

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

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
29
import torch
30
31
32
torch_versions = torch.__version__.split(".")
torch_major_version = int(torch_versions[0])
torch_minor_version = int(torch_versions[1])
33
if (
34
    torch_major_version > 1 or
35
36
37
38
39
40
41
    (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)

42
from openfold.config import model_config
43
from openfold.data import templates, feature_pipeline, data_pipeline
44
45
46
47
48
from openfold.data.tools import hhsearch, hmmsearch
from openfold.np import protein
from openfold.utils.script_utils import (load_models_from_command_line, parse_fasta, run_model,
                                         prep_output, relax_protein)
from openfold.utils.tensor_utils import tensor_tree_map
49
50
51
52
from openfold.utils.trace_utils import (
    pad_feature_dict_seq,
    trace_model_,
)
53

54
from scripts.precompute_embeddings import EmbeddingGenerator
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
55
from scripts.utils import add_data_args
56

57

58
TRACING_INTERVAL = 50
59
60


61
def precompute_alignments(tags, seqs, alignment_dir, args):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
62
63
64
65
66
    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}")

67
        local_alignment_dir = os.path.join(alignment_dir, tag)
68

69
        if args.use_precomputed_alignments is None:
70
            logger.info(f"Generating alignments for {tag}...")
71

72
73
74
75
76
77
78
79
80
81
82
83
84
            os.makedirs(local_alignment_dir, exist_ok=True)

            if "multimer" in args.config_preset:
                template_searcher = hmmsearch.Hmmsearch(
                    binary_path=args.hmmsearch_binary_path,
                    hmmbuild_binary_path=args.hmmbuild_binary_path,
                    database_path=args.pdb_seqres_database_path,
                )
            else:
                template_searcher = hhsearch.HHSearch(
                    binary_path=args.hhsearch_binary_path,
                    databases=[args.pdb70_database_path],
                )
85

86
            # In seqemb mode, use AlignmentRunner only to generate templates
87
88
89
90
            if args.use_single_seq_mode:
                alignment_runner = data_pipeline.AlignmentRunner(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
91
                    template_searcher=template_searcher,
92
93
                    no_cpus=args.cpus,
                )
94
                embedding_generator = EmbeddingGenerator()
95
                embedding_generator.run(tmp_fasta_path, alignment_dir)
96
97
98
99
100
101
102
            else:
                alignment_runner = data_pipeline.AlignmentRunner(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    hhblits_binary_path=args.hhblits_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
                    mgnify_database_path=args.mgnify_database_path,
                    bfd_database_path=args.bfd_database_path,
103
                    uniref30_database_path=args.uniref30_database_path,
104
                    uniclust30_database_path=args.uniclust30_database_path,
105
                    uniprot_database_path=args.uniprot_database_path,
106
107
                    template_searcher=template_searcher,
                    use_small_bfd=args.bfd_database_path is None,
108
                    no_cpus=args.cpus
109
                )
110

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
111
            alignment_runner.run(
112
                tmp_fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
113
            )
114
115
116
117
        else:
            logger.info(
                f"Using precomputed alignments for {tag} at {alignment_dir}..."
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
118
119
120
121
122

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


123
124
125
126
def round_up_seqlen(seqlen):
    return int(math.ceil(seqlen / TRACING_INTERVAL)) * TRACING_INTERVAL


127
128
129
130
131
132
133
def generate_feature_dict(
    tags,
    seqs,
    alignment_dir,
    data_processor,
    args,
):
134
    tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
135
136
137
138
139
140
141
142
143
144

    if "multimer" in args.config_preset:
        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_fasta(
            fasta_path=tmp_fasta_path, alignment_dir=alignment_dir,
        )
    elif len(seqs) == 1:
145
        tag = tags[0]
146
147
148
149
150
151
        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(
152
153
154
            fasta_path=tmp_fasta_path,
            alignment_dir=local_alignment_dir,
            seqemb_mode=args.use_single_seq_mode,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
155
        )
156
157
158
159
160
161
162
163
164
165
166
167
    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)

168
    return feature_dict
169

170

171
172
def list_files_with_extensions(dir, extensions):
    return [f for f in os.listdir(dir) if f.endswith(extensions)]
173

174

175
def main(args):
176
    # Create the output directory
177
178
    os.makedirs(args.output_dir, exist_ok=True)

179
180
    if args.config_preset.startswith("seq"):
        args.use_single_seq_mode = True
181

182
183
184
185
186
    config = model_config(
        args.config_preset, 
        long_sequence_inference=args.long_sequence_inference,
        use_deepspeed_evoformer_attention=args.use_deepspeed_evoformer_attention,
        )
187

188
    if args.experiment_config_json:
189
190
191
192
        with open(args.experiment_config_json, 'r') as f:
            custom_config_dict = json.load(f)
        config.update_from_flattened_dict(custom_config_dict)

193
194
    if args.trace_model:
        if not config.data.predict.fixed_size:
195
196
197
            raise ValueError(
                "Tracing requires that fixed_size mode be enabled in the config"
            )
Christina Floristean's avatar
Christina Floristean committed
198
199

    is_multimer = "multimer" in args.config_preset
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
200

201
    if is_multimer:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
        template_featurizer = templates.HmmsearchHitFeaturizer(
            mmcif_dir=args.template_mmcif_dir,
            max_template_date=args.max_template_date,
            max_hits=config.data.predict.max_templates,
            kalign_binary_path=args.kalign_binary_path,
            release_dates_path=args.release_dates_path,
            obsolete_pdbs_path=args.obsolete_pdbs_path
        )
    else:
        template_featurizer = templates.HhsearchHitFeaturizer(
            mmcif_dir=args.template_mmcif_dir,
            max_template_date=args.max_template_date,
            max_hits=config.data.predict.max_templates,
            kalign_binary_path=args.kalign_binary_path,
            release_dates_path=args.release_dates_path,
            obsolete_pdbs_path=args.obsolete_pdbs_path
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
220
    data_processor = data_pipeline.DataPipeline(
221
222
223
        template_featurizer=template_featurizer,
    )

224
    if is_multimer:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
225
226
227
228
        data_processor = data_pipeline.DataPipelineMultimer(
            monomer_data_pipeline=data_processor,
        )

229
    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
230
    random_seed = args.data_random_seed
231
    if random_seed is None:
232
        random_seed = random.randrange(2 ** 32)
233

234
235
    np.random.seed(random_seed)
    torch.manual_seed(random_seed + 1)
236

237
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
238
239
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
240
    if args.use_precomputed_alignments is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
241
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
242
243
    else:
        alignment_dir = args.use_precomputed_alignments
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
244

245
246
    tag_list = []
    seq_list = []
247
    for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
248
        # Gather input sequences
Christina Floristean's avatar
Christina Floristean committed
249
        fasta_path = os.path.join(args.fasta_dir, fasta_file)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
250
251
        with open(fasta_path, "r") as fp:
            data = fp.read()
252

253
        tags, seqs = parse_fasta(data)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
254

255
        if not is_multimer and len(tags) != 1:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
256
257
258
            print(
                f"{fasta_path} contains more than one sequence but "
                f"multimer mode is not enabled. Skipping..."
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
259
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
260
            continue
Christina Floristean's avatar
Christina Floristean committed
261

262
        # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
263
        tag = '-'.join(tags)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
264

265
        tag_list.append((tag, tags))
266
267
268
269
270
        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 = {}
271
272
273
274
275

    if is_multimer and args.openfold_checkpoint_path:
        raise ValueError(
            '`openfold_checkpoint_path` was specified, but no OpenFold checkpoints are available for multimer mode')

276
277
278
279
280
281
    model_generator = load_models_from_command_line(
        config,
        args.model_device,
        args.openfold_checkpoint_path,
        args.jax_param_path,
        args.output_dir)
282

283
    for model, output_directory in model_generator:
284
        cur_tracing_interval = 0
285
        for (tag, tags), seqs in sorted_targets:
286
287
288
            output_name = f'{tag}_{args.config_preset}'
            if args.output_postfix is not None:
                output_name = f'{output_name}_{args.output_postfix}'
289

290
            # Does nothing if the alignments have already been computed
291
            precompute_alignments(tags, seqs, alignment_dir, args)
292

293
            feature_dict = feature_dicts.get(tag, None)
294
            if feature_dict is None:
295
296
297
298
299
300
                feature_dict = generate_feature_dict(
                    tags,
                    seqs,
                    alignment_dir,
                    data_processor,
                    args,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
301
302
                )

303
                if args.trace_model:
304
305
306
307
308
                    n = feature_dict["aatype"].shape[-2]
                    rounded_seqlen = round_up_seqlen(n)
                    feature_dict = pad_feature_dict_seq(
                        feature_dict, rounded_seqlen,
                    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
309

310
311
312
                feature_dicts[tag] = feature_dict

            processed_feature_dict = feature_processor.process_features(
Christina Floristean's avatar
Christina Floristean committed
313
                feature_dict, mode='predict', is_multimer=is_multimer
314
315
316
            )

            processed_feature_dict = {
317
318
                k: torch.as_tensor(v, device=args.model_device)
                for k, v in processed_feature_dict.items()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
319
            }
320

321
322
            if args.trace_model:
                if rounded_seqlen > cur_tracing_interval:
323
324
325
326
327
                    logger.info(
                        f"Tracing model at {rounded_seqlen} residues..."
                    )
                    t = time.perf_counter()
                    trace_model_(model, processed_feature_dict)
328
                    tracing_time = time.perf_counter() - t
329
                    logger.info(
330
                        f"Tracing time: {tracing_time}"
331
332
                    )
                    cur_tracing_interval = rounded_seqlen
Sam DeLuca's avatar
Sam DeLuca committed
333

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

336
            # Toss out the recycling dimensions --- we don't need them anymore
337
            processed_feature_dict = tensor_tree_map(
338
                lambda x: np.array(x[..., -1].cpu()),
339
340
                processed_feature_dict
            )
341
342
343
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            unrelaxed_protein = prep_output(
344
345
346
347
                out,
                processed_feature_dict,
                feature_dict,
                feature_processor,
348
                args.config_preset,
349
350
                args.multimer_ri_gap,
                args.subtract_plddt
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
351
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
352

353
354
355
            unrelaxed_file_suffix = "_unrelaxed.pdb"
            if args.cif_output:
                unrelaxed_file_suffix = "_unrelaxed.cif"
356
            unrelaxed_output_path = os.path.join(
357
                output_directory, f'{output_name}{unrelaxed_file_suffix}'
358
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
359

360
            with open(unrelaxed_output_path, 'w') as fp:
361
362
363
364
                if args.cif_output:
                    fp.write(protein.to_modelcif(unrelaxed_protein))
                else:
                    fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
365

366
            logger.info(f"Output written to {unrelaxed_output_path}...")
367

368
369
            if not args.skip_relaxation:
                # Relax the prediction.
370
                logger.info(f"Running relaxation on {unrelaxed_output_path}...")
371
372
                relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name,
                              args.cif_output)
373

374
375
            if args.save_outputs:
                output_dict_path = os.path.join(
376
                    output_directory, f'{output_name}_output_dict.pkl'
377
378
379
                )
                with open(output_dict_path, "wb") as fp:
                    pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
380

Sam DeLuca's avatar
Sam DeLuca committed
381
                logger.info(f"Model output written to {output_dict_path}...")
382
383
384
385


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
386
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
387
        "fasta_dir", type=str,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
388
        help="Path to directory containing FASTA files, one sequence per file"
389
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
390
391
392
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
393
394
395
396
    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
397
    )
398
399
400
401
    parser.add_argument(
        "--use_single_seq_mode", action="store_true", default=False,
        help="""Use single sequence embeddings instead of MSAs."""
    )
402
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
403
404
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
405
406
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
407
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
408
409
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
410
411
    )
    parser.add_argument(
412
        "--config_preset", type=str, default="model_1",
413
        help="""Name of a model config preset defined in openfold/config.py"""
414
415
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
416
417
418
419
        "--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"""
420
421
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
422
423
424
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
425
    )
426
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
427
        "--save_outputs", action="store_true", default=False,
428
429
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
430
431
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
432
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
433
    )
434
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
435
        "--preset", type=str, default='full_dbs',
436
437
438
        choices=('reduced_dbs', 'full_dbs')
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
439
440
441
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
442
    parser.add_argument(
443
        "--data_random_seed", type=int, default=None
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
444
445
446
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
447
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
448
449
450
451
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
452
453
454
455
456
457
    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."""
    )
458
459
460
461
462
    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"""
    )
463
464
465
    parser.add_argument(
        "--long_sequence_inference", action="store_true", default=False,
        help="""enable options to reduce memory usage at the cost of speed, helps longer sequences fit into GPU memory, see the README for details"""
466
    )
467
468
469
470
    parser.add_argument(
        "--cif_output", action="store_true", default=False,
        help="Output predicted models in ModelCIF format instead of PDB format (default)"
    )
471
472
473
    parser.add_argument(
        "--experiment_config_json", default="", help="Path to a json file with custom config values to overwrite config setting",
    )
474
475
476
477
    parser.add_argument(
        "--use_deepspeed_evoformer_attention", action="store_true", default=False, 
        help="Whether to use the DeepSpeed evoformer attention layer. Must have deepspeed installed in the environment.",
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
478
    add_data_args(parser)
479
480
    args = parser.parse_args()

481
    if args.jax_param_path is None and args.openfold_checkpoint_path is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
482
        args.jax_param_path = os.path.join(
483
            "openfold", "resources", "params",
484
            "params_" + args.config_preset + ".npz"
485
486
        )

487
    if args.model_device == "cpu" and torch.cuda.is_available():
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
488
489
490
491
492
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )

493
    main(args)