Mercurial > repos > galaxy-australia > alphafold2

comparison gen_extra_outputs.py @ 13:c0e71cb2bd1b draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/usegalaxy-au/tools-au commit 830b5bbf9c5375e714e1b7b9a3e8eec1e584e6b2
author galaxy-australia
date Wed, 12 Oct 2022 22:25:20 +0000
parents 3bd420ec162d
children
comparison
equal deleted inserted replaced
12:7fbec959cf2b 13:c0e71cb2bd1b
33 "--plddts", 33 "--plddts",
34 help="output per-residue confidence scores (pLDDTs)", 34 help="output per-residue confidence scores (pLDDTs)",
35 action="store_true" 35 action="store_true"
36 ) 36 )
37 parser.add_argument( 37 parser.add_argument(
38 "-m",
38 "--multimer", 39 "--multimer",
39 help="parse output from AlphaFold multimer", 40 help="parse output from AlphaFold multimer",
40 action="store_true" 41 action="store_true"
41 ) 42 )
42 args = parser.parse_args() 43 args = parser.parse_args()
48 49
49 class ExecutionContext: 50 class ExecutionContext:
50 """uses program settings to get paths to files etc""" 51 """uses program settings to get paths to files etc"""
51 def __init__(self, settings: Settings): 52 def __init__(self, settings: Settings):
52 self.settings = settings 53 self.settings = settings
54
55 def get_model_key(self, ix):
56 """Return json key for model index."""
57 if self.settings.is_multimer:
58 return f'model_{ix}_multimer'
59 return f'model_{ix}'
53 60
54 @property 61 @property
55 def ranking_debug(self) -> str: 62 def ranking_debug(self) -> str:
56 return f'{self.settings.workdir}/ranking_debug.json' 63 return f'{self.settings.workdir}/ranking_debug.json'
57 64
108 model_pkls = self.context.model_pkls 115 model_pkls = self.context.model_pkls
109 for i in range(len(model_pkls)): 116 for i in range(len(model_pkls)):
110 pklfile = model_pkls[i] 117 pklfile = model_pkls[i]
111 with open(pklfile, 'rb') as fp: 118 with open(pklfile, 'rb') as fp:
112 data = pickle.load(fp) 119 data = pickle.load(fp)
113 plddts[f'model_{i+1}'] = [ 120 plddts[self.context.get_model_key(i+1)] = [
114 float(f'{x:.2f}') 121 float(f'{x:.2f}')
115 for x in data['plddt'] 122 for x in data['plddt']
116 ] 123 ]
117 return plddts 124 return plddts
118 125
119 126
120 class OutputGenerator: 127 class OutputGenerator:
121 """generates the output data we are interested in creating""" 128 """generates the output data we are interested in creating"""
122 def __init__(self, loader: FileLoader): 129 def __init__(self, loader: FileLoader):
123 self.loader = loader 130 self.loader = loader
131 self.context = loader.context
124 132
125 def gen_conf_scores(self): 133 def gen_conf_scores(self):
126 mapping = self.loader.get_model_mapping() 134 mapping = self.loader.get_model_mapping()
127 scores = self.loader.get_conf_scores() 135 scores = self.loader.get_conf_scores()
128 ranked = list(scores.items()) 136 ranked = list(scores.items())
129 ranked.sort(key=lambda x: x[1], reverse=True) 137 ranked.sort(key=lambda x: x[1], reverse=True)
130 return {f'model_{mapping[name]}': score 138 return {
131 for name, score in ranked} 139 self.context.get_model_key(mapping[name]): score
140 for name, score in ranked
141 }
132 142
133 def gen_residue_scores(self) -> Dict[str, List[float]]: 143 def gen_residue_scores(self) -> Dict[str, List[float]]:
134 mapping = self.loader.get_model_mapping() 144 mapping = self.loader.get_model_mapping()
135 model_plddts = self.loader.get_model_plddts() 145 model_plddts = self.loader.get_model_plddts()
136 return {f'model_{mapping[name]}': plddts 146 return {
137 for name, plddts in model_plddts.items()} 147 self.context.get_model_key(mapping[name]): plddts
148 for name, plddts in model_plddts.items()
149 }
138 150
139 151
140 class OutputWriter: 152 class OutputWriter:
141 """writes generated data to files""" 153 """writes generated data to files"""
142 def __init__(self, context: ExecutionContext): 154 def __init__(self, context: ExecutionContext):