|
| 1 | +# vim: set expandtab shiftwidth=4 softtabstop=4: |
| 2 | + |
| 3 | +# === UCSF ChimeraX Copyright === |
| 4 | +# Copyright 2022 Regents of the University of California. All rights reserved. |
| 5 | +# The ChimeraX application is provided pursuant to the ChimeraX license |
| 6 | +# agreement, which covers academic and commercial uses. For more details, see |
| 7 | +# <https://www.rbvi.ucsf.edu/chimerax/docs/licensing.html> |
| 8 | +# |
| 9 | +# This particular file is part of the ChimeraX library. You can also |
| 10 | +# redistribute and/or modify it under the terms of the GNU Lesser General |
| 11 | +# Public License version 2.1 as published by the Free Software Foundation. |
| 12 | +# For more details, see |
| 13 | +# <https://www.gnu.org/licenses/old-licenses/lgpl-2.1.html> |
| 14 | +# |
| 15 | +# THIS SOFTWARE IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER |
| 16 | +# EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES |
| 17 | +# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. ADDITIONAL LIABILITY |
| 18 | +# LIMITATIONS ARE DESCRIBED IN THE GNU LESSER GENERAL PUBLIC LICENSE |
| 19 | +# VERSION 2.1 |
| 20 | +# |
| 21 | +# This notice must be embedded in or attached to all copies, including partial |
| 22 | +# copies, of the software or any revisions or derivations thereof. |
| 23 | +# === UCSF ChimeraX Copyright === |
| 24 | + |
| 25 | +# ----------------------------------------------------------------------------- |
| 26 | +# |
| 27 | +from chimerax.core.tools import ToolInstance |
| 28 | +class MutationScoresHeatmap(ToolInstance): |
| 29 | + |
| 30 | + help = 'https://www.rbvi.ucsf.edu/chimerax/data/mutation-scores-oct2024/mutation_scores.html' |
| 31 | + |
| 32 | + def __init__(self, session, tool_name = 'Mutation Scores Heatmap'): |
| 33 | + |
| 34 | + ToolInstance.__init__(self, session, tool_name) |
| 35 | + |
| 36 | + from chimerax.ui import MainToolWindow |
| 37 | + tw = MainToolWindow(self) |
| 38 | + tw.fill_context_menu = self._fill_context_menu |
| 39 | + self.tool_window = tw |
| 40 | + parent = tw.ui_area |
| 41 | + |
| 42 | + from chimerax.ui.widgets import vertical_layout |
| 43 | + layout = vertical_layout(parent, margins = (5,0,0,0)) |
| 44 | + |
| 45 | + self._score_view = gv = ScoreView(parent, self._report_cell_info) |
| 46 | + from Qt.QtWidgets import QSizePolicy |
| 47 | + from Qt.QtCore import Qt |
| 48 | +# gv.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) |
| 49 | +# gv.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) |
| 50 | +# gv.setSizePolicy(QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Fixed) |
| 51 | + layout.addWidget(gv, stretch=1) |
| 52 | + |
| 53 | + from Qt.QtWidgets import QGraphicsScene |
| 54 | + self._scene = gs = QGraphicsScene(gv) |
| 55 | + gs.setSceneRect(0, 0, 500, 500) |
| 56 | + gv.setScene(gs) |
| 57 | + |
| 58 | + from Qt.QtWidgets import QLabel |
| 59 | + self._info_label = info = QLabel(parent) |
| 60 | + layout.addWidget(info) |
| 61 | + |
| 62 | + self._set_heatmap_image() |
| 63 | + |
| 64 | + tw.manage(placement=None) # Start floating |
| 65 | + |
| 66 | + # --------------------------------------------------------------------------- |
| 67 | + # |
| 68 | + def closed(self): |
| 69 | + return self.tool_window.tool_instance is None |
| 70 | + |
| 71 | + # --------------------------------------------------------------------------- |
| 72 | + # |
| 73 | + def _fill_context_menu(self, menu, x, y): |
| 74 | + menu.addAction('Select residue', self._select_residue) |
| 75 | + menu.addAction('Save image', self._save_image) |
| 76 | + |
| 77 | + # --------------------------------------------------------------------------- |
| 78 | + # |
| 79 | + def _set_heatmap_image(self): |
| 80 | + score_matrix = self._score_matrix() |
| 81 | + blue,white,red = (0,0,1,1), (1,1,1,1), (1,0,0,1) |
| 82 | + from chimerax.core.colors import Colormap |
| 83 | + colormap = Colormap((-2.0,-1.0,1.0,2.0), (blue, white, white, red)) |
| 84 | + self._score_view._make_image(score_matrix, colormap) |
| 85 | + |
| 86 | + # --------------------------------------------------------------------------- |
| 87 | + # |
| 88 | + _amino_acids = 'HRKDEFWYNQILCSTVMAGP' |
| 89 | + def _score_matrix(self): |
| 90 | + from .ms_data import mutation_all_scores |
| 91 | + msets = mutation_all_scores(self.session) |
| 92 | + self._mutation_set = mset = msets[0] # TODO: Allow choosing mutation set |
| 93 | + scores = None |
| 94 | + score_names = mset.score_names() |
| 95 | + # TODO: Allow choosing score names |
| 96 | + score_names = [score_name for score_name in score_names if score_name.endswith('_effect')] |
| 97 | + self._score_names = score_names |
| 98 | + self._num_scores = score_count = len(score_names) |
| 99 | + aa_to_index = {aa:i for i, aa in enumerate(self._amino_acids)} |
| 100 | + self._res_aa = res_aa = {} |
| 101 | + for snum, score_name in enumerate(score_names): |
| 102 | + score_values = mset.score_values(score_name) |
| 103 | + if scores is None: |
| 104 | + # TODO: This may not give maximum res number |
| 105 | + self._num_residues = rmax = max(score_values.residue_numbers()) |
| 106 | + from numpy import zeros, float32 |
| 107 | + self._scores = scores = zeros((rmax, 20, score_count), float32) |
| 108 | + sscores = scores[:,:,snum] |
| 109 | + for res_num, from_aa, to_aa, value in score_values.all_values(): |
| 110 | + res_aa[res_num] = from_aa |
| 111 | + aa_index = aa_to_index[to_aa] |
| 112 | + sscores[res_num-1, aa_index] = value |
| 113 | + mean, sdev = score_values.synonymous_mean_and_sdev() |
| 114 | + sscores -= mean |
| 115 | + sscores /= sdev |
| 116 | + |
| 117 | + scores_2d = scores.reshape((rmax, 20*score_count)).transpose() |
| 118 | + return scores_2d |
| 119 | + |
| 120 | + # --------------------------------------------------------------------------- |
| 121 | + # |
| 122 | + def _report_cell_info(self, column_index, row_index): |
| 123 | + num_cols = self._num_residues |
| 124 | + num_rows = 20 * self._num_scores |
| 125 | + if column_index < 0 or row_index < 0 or column_index >= num_cols or row_index >= num_rows or column_index+1 not in self._res_aa: |
| 126 | + msg = '' |
| 127 | + else: |
| 128 | + res_num = column_index + 1 |
| 129 | + from_aa = self._res_aa[res_num] |
| 130 | + score_num = row_index % self._num_scores |
| 131 | + score_name = self._score_names[score_num] |
| 132 | + aa_index = row_index // self._num_scores |
| 133 | + to_aa = self._amino_acids[aa_index] |
| 134 | + score_value = self._scores[res_num-1, aa_index, score_num] |
| 135 | + msg = f'{from_aa}{res_num}{to_aa} {score_name} {"%.2f"%score_value}' |
| 136 | + self._info_label.setText(msg) |
| 137 | + |
| 138 | + # --------------------------------------------------------------------------- |
| 139 | + # |
| 140 | + def _save_image(self, default_suffix = '_heatmap.png'): |
| 141 | + from os.path import dirname, join |
| 142 | + filename = self._mutation_set.name + default_suffix |
| 143 | + dir = dirname(self._mutation_set.path) |
| 144 | + suggested_path = join(dir, filename) |
| 145 | + from Qt.QtWidgets import QFileDialog |
| 146 | + parent = self.tool_window.ui_area |
| 147 | + path, ftype = QFileDialog.getSaveFileName(parent, |
| 148 | + 'Mutation Heatmap Image', |
| 149 | + suggested_path) |
| 150 | + if path: |
| 151 | + self._score_view.save_image(path) |
| 152 | + |
| 153 | + # --------------------------------------------------------------------------- |
| 154 | + # |
| 155 | + def _show_help(self): |
| 156 | + from chimerax.core.commands import run |
| 157 | + run(self.session, 'help %s' % self.help) |
| 158 | + |
| 159 | +# --------------------------------------------------------------------------- |
| 160 | +# |
| 161 | +from Qt.QtWidgets import QGraphicsView |
| 162 | +class ScoreView(QGraphicsView): |
| 163 | + def __init__(self, parent, report_cell_info_cb=None): |
| 164 | + QGraphicsView.__init__(self, parent) |
| 165 | + self._report_cell_info_callback = report_cell_info_cb |
| 166 | + self._pixmap_item = None |
| 167 | + |
| 168 | + # Report cell info as mouse hovers over plot. |
| 169 | + self.setMouseTracking(True) |
| 170 | + |
| 171 | + # Zoom in |
| 172 | + self.scale(2,2) |
| 173 | + |
| 174 | + def sizeHint(self): |
| 175 | + from Qt.QtCore import QSize |
| 176 | + return QSize(500,500) |
| 177 | + |
| 178 | + def mouseMoveEvent(self, event): |
| 179 | + if self._report_cell_info_callback: |
| 180 | + x,y = self._scene_position(event) |
| 181 | + self._report_cell_info_callback(int(x),int(y)) |
| 182 | + |
| 183 | + def _scene_position(self, event): |
| 184 | + p = self.mapToScene(event.pos()) |
| 185 | + return p.x(), p.y() |
| 186 | + |
| 187 | + def _make_image(self, matrix, colormap): |
| 188 | + scene = self.scene() |
| 189 | + pi = self._pixmap_item |
| 190 | + if pi is not None: |
| 191 | + scene.removeItem(pi) |
| 192 | + |
| 193 | + rgb = matrix_to_rgb(matrix, colormap) |
| 194 | + pixmap = rgb_to_pixmap(rgb) |
| 195 | + self._pixmap_item = scene.addPixmap(pixmap) |
| 196 | + scene.setSceneRect(0, 0, pixmap.width(), pixmap.height()) |
| 197 | + |
| 198 | + def save_image(self, path): |
| 199 | + pixmap = self.grab() |
| 200 | + pixmap.save(path) |
| 201 | + |
| 202 | +# ----------------------------------------------------------------------------- |
| 203 | +# |
| 204 | +def matrix_to_rgb(matrix, colormap): |
| 205 | + rgb_flat = colormap.interpolated_rgba8(matrix.ravel())[:,:3] |
| 206 | + n,m = matrix.shape |
| 207 | + rgb = rgb_flat.reshape((n,m,3)).copy() |
| 208 | + return rgb |
| 209 | + |
| 210 | +# ----------------------------------------------------------------------------- |
| 211 | +# |
| 212 | +def rgb_to_pixmap(rgb): |
| 213 | + # Save image to a PNG file |
| 214 | + from Qt.QtGui import QImage, QPixmap |
| 215 | + h, w = rgb.shape[:2] |
| 216 | + im = QImage(rgb.data, w, h, 3*w, QImage.Format_RGB888) |
| 217 | + pixmap = QPixmap.fromImage(im) |
| 218 | + return pixmap |
| 219 | + |
| 220 | +def mutation_heatmap(session): |
| 221 | + hm = MutationScoresHeatmap(session) |
| 222 | + return hm |
| 223 | + |
| 224 | +def register_command(logger): |
| 225 | + from chimerax.core.commands import CmdDesc, register, StringArg, BoolArg |
| 226 | + desc = CmdDesc( |
| 227 | + synopsis = 'Show a heatmap of mutation scores.' |
| 228 | + ) |
| 229 | + register('mutationscores heatmap', desc, mutation_heatmap, logger=logger) |
0 commit comments