GeneralizedEvDesign/tradeoff.py at master · SchubertLab/GeneralizedEvDesign · GitHub

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


import csv
import logging
import multiprocessing as mp
import os
import time
from builtins import map
from collections import defaultdict
from random import sample as random_sample

import click
import numpy as np
import pandas as pd
import pyomo.environ as aml
import pyomo.kernel as pmo
from pyomo.core.expr.numeric_expr import SumExpression
from pyomo.opt import SolverFactory, TerminationCondition

import Fred2
import utilities
from Fred2.Core import (Allele, Peptide, Protein,
                        generate_peptides_from_proteins)
from Fred2.Core.Peptide import Peptide
from Fred2.EpitopePrediction import (EpitopePredictionResult,
                                     EpitopePredictorFactory)
from Fred2.EpitopeSelection import OptiTope
from Fred2.IO import FileReader
from team_orienteering_ilp import TeamOrienteeringIlp

LOGGER = None


@click.command()
@click.argument('input-epitopes', type=click.Path())
@click.argument('input-cleavages', type=click.Path())
@click.argument('output-frontier', type=click.Path())
@click.option('--verbose', '-v', is_flag=True, help='print debug messages on the console')
@click.option('--log-file', '-l', type=click.Path(), help='where to store the logs')
@click.option('--pareto-steps', '-s', default=10, help='How many points to discover on the pareto frontier')
@click.option('--cocktail', '-c', default=1, help='How many strains to include in the epitopes cocktail')
@click.option('--max-aminoacids', '-a', default=0, help='Maximum length of the epitopes in aminoacids')
@click.option('--max-epitopes', '-e', default=10, help='Maximum length of the epitopes in epitopes')
@click.option('--min-alleles', default=0.0, help='Vaccine must cover at least this many alleles')
@click.option('--min-proteins', default=0.0, help='Vaccine must cover at least this many proteins')
@click.option('--min-avg-prot-conservation', default=0.0, help='On average, epitopes in the epitopes must cover at least this many proteins')
@click.option('--min-avg-alle-conservation', default=0.0, help='On average, epitopes in the epitopes must cover at least this many alleles')
def main(input_epitopes, input_cleavages, output_frontier, cocktail,
         pareto_steps, verbose, log_file, max_aminoacids, max_epitopes,
         min_alleles, min_proteins, min_avg_prot_conservation,
         min_avg_alle_conservation):
    '''
    Explore the trade-off between the immunogenicity and the cleavage
    likelihood for string-of-beads vaccines. Outputs the Pareto frontier
    of the two objectives.
    '''

    global LOGGER
    LOGGER = utilities.init_logging(verbose=verbose, log_file=log_file)

    # load epitopes
    epitopes = utilities.load_epitopes(input_epitopes)
    LOGGER.info('Loaded %d epitopes', len(epitopes))

    # read cleavage scores
    cleavage_epitopes = set()
    cleavages, spacers = {}, {}
    with open(input_cleavages) as f:
        for row in csv.DictReader(f):
            cleavages[row['from'], row['to']] = float(row['score'])
            spacers[row['from'], row['to']] = row.get('spacer', '')
            cleavage_epitopes.add(row['from'])
            cleavage_epitopes.add(row['to'])
    LOGGER.info('Loaded %d cleavage scores', len(cleavages))

    # compute edge cost
    edge_cost, vertices, vertices_rewards = [], [], []
    vertex_to_epitope = [''] + list(cleavage_epitopes)
    for ep_from in vertex_to_epitope:
        vertices.append(ep_from)
        vertices_rewards.append(0 if ep_from == '' else epitopes[ep_from]['immunogen'])
        edge_cost.append([
            cleavages[ep_from, ep_to] if ep_from != '' and ep_to != '' and ep_from != ep_to else 0.0
            for ep_to in vertex_to_epitope
        ])
    LOGGER.info('Graph has %d vertices', len(vertex_to_epitope))

    # find optimal design
    solver_build_time = time.time()
    solver = TeamOrienteeringIlp(
        num_teams=cocktail, vertex_reward=vertices_rewards, edge_cost=edge_cost,
        type_coverage=[], min_type_coverage=[], min_avg_type_conservation=[],
        max_edge_cost=0, max_vertices=max_epitopes
    )
    solver.build_model()
    reward_cost = solver.explore_edge_cost_vertex_reward_tradeoff(pareto_steps)
    with open(output_frontier, 'w') as f:
        writer = csv.writer(f)
        writer.writerow(('immunogenicity', 'cleavage', 'epitopes', 'spacers'))
        for epitopes_index, immunogen, cleavage in reward_cost:
            vax_epis, vax_spacers = [], []
            for a, b in epitopes_index[0]: # FIXME multiple tours?
                if b != 0:
                    vax_epis.append(vertex_to_epitope[b])
                if a != 0 and b != 0:
                    vax_spacers.append(spacers[vertex_to_epitope[a], vertex_to_epitope[b]])

            writer.writerow((immunogen, cleavage, ';'.join(vax_epis), ';'.join(vax_spacers)))
            f.flush()  # write progress immediately

            LOGGER.info('Immunogenicity: %.3f - Cleavage: %.3f - Epitopes: %s - Spacers: %s',
                        immunogen, cleavage, ', '.join(vax_epis), ', '.join(vax_spacers))


if __name__ == '__main__':
    main()