-
Notifications
You must be signed in to change notification settings - Fork 55
/
calibrate_percentile_ranks_command.py
411 lines (355 loc) · 13.8 KB
/
calibrate_percentile_ranks_command.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
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
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
"""
Calibrate percentile ranks for models. Runs in-place.
"""
import argparse
import os
import signal
import sys
import time
import traceback
import collections
from functools import partial
import pandas
import numpy
import tqdm # progress bar
tqdm.monitor_interval = 0 # see https://github.com/tqdm/tqdm/issues/481
from .class1_affinity_predictor import Class1AffinityPredictor
from .class1_presentation_predictor import Class1PresentationPredictor
from .common import normalize_allele_name
from .encodable_sequences import EncodableSequences
from .common import configure_logging, random_peptides, amino_acid_distribution
from .local_parallelism import (
add_local_parallelism_args,
worker_pool_with_gpu_assignments_from_args,
call_wrapped_kwargs)
from .cluster_parallelism import (
add_cluster_parallelism_args,
cluster_results_from_args)
# To avoid pickling large matrices to send to child processes when running in
# parallel, we use this global variable as a place to store data. Data that is
# stored here before creating the thread pool will be inherited to the child
# processes upon fork() call, allowing us to share large data with the workers
# via shared memory.
GLOBAL_DATA = {}
parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument(
"--predictor-kind",
choices=("class1_affinity", "class1_presentation"),
default="class1_affinity",
help="Type of predictor to calibrate")
parser.add_argument(
"--models-dir",
metavar="DIR",
required=True,
help="Directory to read and write models")
parser.add_argument(
"--allele",
default=None,
nargs="+",
help="Alleles to calibrate percentile ranks for. If not specified all "
"alleles are used")
parser.add_argument(
"--match-amino-acid-distribution-data",
help="Sample random peptides from the amino acid distribution of the "
"peptides listed in the supplied CSV file, which must have a 'peptide' "
"column. If not specified a uniform distribution is used.")
parser.add_argument(
"--alleles-file",
default=None,
help="Use alleles in supplied CSV file, which must have an 'allele' column.")
parser.add_argument(
"--num-peptides-per-length",
type=int,
metavar="N",
default=int(1e5),
help="Number of peptides per length to use to calibrate percent ranks. "
"Default: %(default)s.")
parser.add_argument(
"--num-genotypes",
type=int,
metavar="N",
default=25,
help="Used when calibrrating a presentation predictor. Number of genotypes"
"to sample")
parser.add_argument(
"--alleles-per-genotype",
type=int,
metavar="N",
default=6,
help="Used when calibrating a presentation predictor. Number of alleles "
"per genotype. Use 1 to calibrate for single alleles. Default: %(default)s")
parser.add_argument(
"--motif-summary",
default=False,
action="store_true",
help="Calculate motifs and length preferences for each allele")
parser.add_argument(
"--summary-top-peptide-fraction",
default=[0.0001, 0.001, 0.01, 0.1, 1.0],
nargs="+",
type=float,
metavar="X",
help="The top X fraction of predictions (i.e. tightest binders) to use to "
"generate motifs and length preferences. Default: %(default)s")
parser.add_argument(
"--length-range",
default=(8, 15),
type=int,
nargs=2,
help="Min and max peptide length to calibrate, inclusive. "
"Default: %(default)s")
parser.add_argument(
"--prediction-batch-size",
type=int,
default=4096,
help="Keras batch size for predictions")
parser.add_argument(
"--alleles-per-work-chunk",
type=int,
metavar="N",
default=1,
help="Number of alleles per work chunk. Default: %(default)s.")
parser.add_argument(
"--verbosity",
type=int,
help="Keras verbosity. Default: %(default)s",
default=0)
add_local_parallelism_args(parser)
add_cluster_parallelism_args(parser)
def run(argv=sys.argv[1:]):
global GLOBAL_DATA
# On sigusr1 print stack trace
print("To show stack trace, run:\nkill -s USR1 %d" % os.getpid())
signal.signal(signal.SIGUSR1, lambda sig, frame: traceback.print_stack())
args = parser.parse_args(argv)
args.models_dir = os.path.abspath(args.models_dir)
configure_logging(verbose=args.verbosity > 1)
aa_distribution = None
if args.match_amino_acid_distribution_data:
distribution_peptides = pandas.read_csv(
args.match_amino_acid_distribution_data).peptide.unique()
distribution_peptides = [
x for x in distribution_peptides if 'X' not in x and 'B' not in x and 'U' not in x
]
aa_distribution = amino_acid_distribution(distribution_peptides)
print("Using amino acid distribution:")
print(aa_distribution)
start = time.time()
peptides = []
lengths = range(args.length_range[0], args.length_range[1] + 1)
for length in lengths:
peptides.extend(
random_peptides(
args.num_peptides_per_length,
length,
distribution=aa_distribution))
print("Done generating peptides in %0.2f sec." % (time.time() - start))
if args.predictor_kind == "class1_affinity":
return run_class1_affinity_predictor(args, peptides)
elif args.predictor_kind == "class1_presentation":
return run_class1_presentation_predictor(args, peptides)
else:
raise ValueError("Unsupported kind %s" % args.predictor_kind)
def run_class1_presentation_predictor(args, peptides):
# This will trigger a Keras import - will break local parallelism.
predictor = Class1PresentationPredictor.load(args.models_dir)
if args.allele:
alleles = [normalize_allele_name(a) for a in args.allele]
elif args.alleles_file:
alleles = pandas.read_csv(args.alleles_file).allele.unique()
else:
alleles = predictor.supported_alleles
print("Num alleles", len(alleles))
genotypes = {}
if args.alleles_per_genotype == 6:
gene_to_alleles = collections.defaultdict(list)
for a in alleles:
for gene in ["A", "B", "C"]:
if a.startswith("HLA-%s" % gene):
gene_to_alleles[gene].append(a)
for _ in range(args.num_genotypes):
genotype = []
for gene in ["A", "A", "B", "B", "C", "C"]:
genotype.append(numpy.random.choice(gene_to_alleles[gene]))
genotypes[",".join(genotype)] = genotype
elif args.alleles_per_genotype == 1:
for _ in range(args.num_genotypes):
genotype = [numpy.random.choice(alleles)]
genotypes[",".join(genotype)] = genotype
else:
raise ValueError("Alleles per genotype must be 6 or 1")
print("Sampled genotypes: ", list(genotypes))
print("Num peptides: ", len(peptides))
start = time.time()
print("Generating predictions")
predictions_df = predictor.predict(
peptides=peptides,
alleles=genotypes)
print("Finished in %0.2f sec." % (time.time() - start))
print(predictions_df)
print("Calibrating ranks")
scores = predictions_df.presentation_score.values
predictor.calibrate_percentile_ranks(scores)
print("Done. Saving.")
predictor.save(
args.models_dir,
write_affinity_predictor=False,
write_processing_predictor=False,
write_weights=False,
write_percent_ranks=True,
write_info=False,
write_metdata=False)
print("Wrote predictor to: %s" % args.models_dir)
def run_class1_affinity_predictor(args, peptides):
global GLOBAL_DATA
# It's important that we don't trigger a Keras import here since that breaks
# local parallelism (tensorflow backend). So we set optimization_level=0.
predictor = Class1AffinityPredictor.load(
args.models_dir,
optimization_level=0,
)
if args.allele:
alleles = [normalize_allele_name(a) for a in args.allele]
elif args.alleles_file:
alleles = pandas.read_csv(args.alleles_file).allele.unique()
else:
alleles = predictor.supported_alleles
allele_set = set(alleles)
if predictor.allele_to_sequence:
# Remove alleles that have the same sequence.
new_allele_set = set()
sequence_to_allele = collections.defaultdict(set)
for allele in list(allele_set):
sequence_to_allele[predictor.allele_to_sequence[allele]].add(allele)
for equivalent_alleles in sequence_to_allele.values():
equivalent_alleles = sorted(equivalent_alleles)
keep = equivalent_alleles.pop(0)
new_allele_set.add(keep)
print(
"Sequence comparison reduced num alleles from",
len(allele_set),
"to",
len(new_allele_set))
allele_set = new_allele_set
alleles = sorted(allele_set)
print("Percent rank calibration for %d alleles. " % (len(alleles)))
print("Encoding %d peptides." % len(peptides))
start = time.time()
encoded_peptides = EncodableSequences.create(peptides)
del peptides
# Now we encode the peptides for each neural network, so the encoding
# becomes cached.
for network in predictor.neural_networks:
network.peptides_to_network_input(encoded_peptides)
assert encoded_peptides.encoding_cache # must have cached the encoding
print("Finished encoding peptides in %0.2f sec." % (time.time() - start))
# Store peptides in global variable so they are in shared memory
# after fork, instead of needing to be pickled (when doing a parallel run).
GLOBAL_DATA["calibration_peptides"] = encoded_peptides
GLOBAL_DATA["predictor"] = predictor
GLOBAL_DATA["args"] = {
'motif_summary': args.motif_summary,
'summary_top_peptide_fractions': args.summary_top_peptide_fraction,
'verbose': args.verbosity > 0,
'model_kwargs': {
'batch_size': args.prediction_batch_size,
}
}
del encoded_peptides
serial_run = not args.cluster_parallelism and args.num_jobs == 0
worker_pool = None
start = time.time()
work_items = []
for allele in alleles:
if not work_items or len(
work_items[-1]['alleles']) >= args.alleles_per_work_chunk:
work_items.append({"alleles": []})
work_items[-1]['alleles'].append(allele)
if serial_run:
# Serial run
print("Running in serial.")
results = (
do_class1_affinity_calibrate_percentile_ranks(**item) for item in work_items)
elif args.cluster_parallelism:
# Run using separate processes HPC cluster.
print("Running on cluster.")
results = cluster_results_from_args(
args,
work_function=do_class1_affinity_calibrate_percentile_ranks,
work_items=work_items,
constant_data=GLOBAL_DATA,
result_serialization_method="pickle",
clear_constant_data=True)
else:
worker_pool = worker_pool_with_gpu_assignments_from_args(args)
print("Worker pool", worker_pool)
assert worker_pool is not None
for item in work_items:
item['constant_data'] = GLOBAL_DATA
results = worker_pool.imap_unordered(
partial(call_wrapped_kwargs, do_class1_affinity_calibrate_percentile_ranks),
work_items,
chunksize=1)
summary_results_lists = collections.defaultdict(list)
for work_item in tqdm.tqdm(results, total=len(work_items)):
for (transforms, summary_results) in work_item:
predictor.allele_to_percent_rank_transform.update(transforms)
if summary_results is not None:
for (item, value) in summary_results.items():
summary_results_lists[item].append(value)
print("Done calibrating %d alleles." % len(alleles))
if summary_results_lists:
for (name, lst) in summary_results_lists.items():
df = pandas.concat(lst, ignore_index=True)
predictor.metadata_dataframes[name] = df
print("Including summary result: %s" % name)
print(df)
predictor.save(args.models_dir, model_names_to_write=[])
percent_rank_calibration_time = time.time() - start
if worker_pool:
worker_pool.close()
worker_pool.join()
print("Percent rank calibration time: %0.2f min." % (
percent_rank_calibration_time / 60.0))
print("Predictor written to: %s" % args.models_dir)
def do_class1_affinity_calibrate_percentile_ranks(
alleles, constant_data=GLOBAL_DATA):
if 'predictor' not in constant_data:
raise ValueError("No predictor provided: " + str(constant_data))
result_list = []
for (i, allele) in enumerate(alleles):
print("Processing allele", i + 1, "of", len(alleles))
result_item = class1_affinity_calibrate_percentile_ranks(
allele,
constant_data['predictor'],
peptides=constant_data['calibration_peptides'],
**constant_data["args"])
result_list.append(result_item)
return result_list
def class1_affinity_calibrate_percentile_ranks(
allele,
predictor,
peptides=None,
motif_summary=False,
summary_top_peptide_fractions=[0.001],
verbose=False,
model_kwargs={}):
if verbose:
print("Calibrating", allele)
predictor.optimize() # since we loaded with optimization_level=0
start = time.time()
summary_results = predictor.calibrate_percentile_ranks(
peptides=peptides,
alleles=[allele],
motif_summary=motif_summary,
summary_top_peptide_fractions=summary_top_peptide_fractions,
verbose=verbose,
model_kwargs=model_kwargs)
if verbose:
print("Done calibrating", allele, "in", time.time() - start, "sec")
transforms = {
allele: predictor.allele_to_percent_rank_transform[allele],
}
return (transforms, summary_results)
if __name__ == '__main__':
run()