|
2 | 2 | scaffoldgraph.vis.utils |
3 | 3 | """ |
4 | 4 |
|
| 5 | +import matplotlib.pyplot as plt |
| 6 | +import matplotlib as mpl |
| 7 | + |
5 | 8 | from rdkit.Chem.Draw import rdMolDraw2D |
6 | 9 | from rdkit import Chem |
7 | 10 |
|
@@ -117,6 +120,140 @@ def remove_node_mol_images(graph): |
117 | 120 | _ = data.pop('img', None) |
118 | 121 |
|
119 | 122 |
|
| 123 | +def rgba_to_hex(scalar_mappable, value): |
| 124 | + """str: rgba to hex.""" |
| 125 | + rgba = scalar_mappable.to_rgba(value) |
| 126 | + c_hex = mpl.colors.to_hex(rgba, keep_alpha=False) |
| 127 | + return c_hex |
| 128 | + |
| 129 | + |
| 130 | +def cmap_to_scalar_mappable(cmap, vmin, vmax): |
| 131 | + """Convert matplotlib Colormap to a ScalarMappable. |
| 132 | +
|
| 133 | + Parameters |
| 134 | + ---------- |
| 135 | + cmap : matplotlib.colors.Colormap |
| 136 | + vmin : float |
| 137 | + Minimum value for normalization. |
| 138 | + vmax : float |
| 139 | + Maximum value for normalization. |
| 140 | +
|
| 141 | + Returns |
| 142 | + ------- |
| 143 | + matplolib.cm.ScalarMappable |
| 144 | +
|
| 145 | + """ |
| 146 | + cnorm = mpl.colors.Normalize(vmin, vmax) |
| 147 | + scalar = mpl.cm.ScalarMappable(norm=cnorm, cmap=cmap) |
| 148 | + return scalar |
| 149 | + |
| 150 | + |
| 151 | +def color_nodes_by_attribute(graph, attribute, cmap, node_type, label='color'): |
| 152 | + """ |
| 153 | + Add an attribute to nodes in a ScaffoldGraph containing a color hex code, |
| 154 | + calculated from a paticular node attribute and a matplotlib cmap. The |
| 155 | + operation is perfomred in-place. |
| 156 | +
|
| 157 | + Can be used for adding colors to ScaffoldGraph visualizations. |
| 158 | +
|
| 159 | + Parameters |
| 160 | + ---------- |
| 161 | + graph : ScaffoldGraph |
| 162 | + Input ScaffoldGraph |
| 163 | + attribute : str |
| 164 | + Key for the attibute from which to calculate a color. |
| 165 | + cmap : str or matplotlib.colors.Colormap |
| 166 | + A matplotlib cmap or name of a cmap e.g. 'BuPu' for |
| 167 | + calculating a nodes colour. |
| 168 | + node_type : str |
| 169 | + The type of node to process e.g. 'scaffold' / 'molecule' |
| 170 | + label : str, optional |
| 171 | + The attribute label to use for storing the color. |
| 172 | + The default is 'color'. |
| 173 | +
|
| 174 | + """ |
| 175 | + # Cmap may be a string or a Colormap |
| 176 | + if isinstance(cmap, str): |
| 177 | + cmap = plt.get_cmap(cmap) |
| 178 | + else: |
| 179 | + if not issubclass(type(cmap), mpl.colors.Colormap): |
| 180 | + raise ValueError('cmap must be a string or a matplotlib Colormap') |
| 181 | + |
| 182 | + # Get attribute range. |
| 183 | + _, attr = zip(*graph._get_nodes_with_type(node_type, attribute, None)) |
| 184 | + attr = list(filter(lambda x: x is not None, attr)) |
| 185 | + attr = list(map(float, attr)) |
| 186 | + vmin, vmax = min(attr), max(attr) |
| 187 | + |
| 188 | + # Assign colors to each node. |
| 189 | + scalar_mappable = cmap_to_scalar_mappable(cmap, vmin, vmax) |
| 190 | + for node, data in graph._get_nodes_with_type(node_type, True, None): |
| 191 | + attr_val = data.get(attribute, None) |
| 192 | + if not attr_val: |
| 193 | + color = '#EEEEEE' # Set a neutral default. |
| 194 | + else: |
| 195 | + color = rgba_to_hex(scalar_mappable, attr_val) |
| 196 | + data[label] = color |
| 197 | + |
| 198 | + |
| 199 | +def color_scaffold_nodes_by_attribute(graph, attribute, cmap, label='color'): |
| 200 | + """ |
| 201 | + Add an attribute to scaffold nodes in a ScaffoldGraph containing a color hex code, |
| 202 | + calculated from a paticular scaffold node attribute and a matplotlib cmap. The |
| 203 | + operation is perfomred in-place. |
| 204 | +
|
| 205 | + Can be used for adding colors to ScaffoldGraph visualizations. |
| 206 | +
|
| 207 | + Parameters |
| 208 | + ---------- |
| 209 | + graph : ScaffoldGraph |
| 210 | + Input ScaffoldGraph |
| 211 | + attribute : str |
| 212 | + Key for the attibute from which to calculate a color. |
| 213 | + cmap : str or matplotlib.colors.Colormap |
| 214 | + A matplotlib cmap or name of a cmap e.g. 'BuPu' for |
| 215 | + calculating a nodes colour. |
| 216 | + label : str, optional |
| 217 | + The attribute label to use for storing the color. |
| 218 | + The default is 'color'. |
| 219 | +
|
| 220 | + See Also |
| 221 | + -------- |
| 222 | + color_molecule_nodes_by_attribute |
| 223 | +
|
| 224 | + """ |
| 225 | + color_nodes_by_attribute(graph, attribute, cmap, 'scaffold', label) |
| 226 | + |
| 227 | + |
| 228 | +def color_molecule_nodes_by_attribute(graph, attribute, cmap, label='color'): |
| 229 | + """ |
| 230 | + Add an attribute to molecule nodes in a ScaffoldGraph containing a color hex code, |
| 231 | + calculated from a paticular molecule node attribute and a matplotlib cmap. The |
| 232 | + operation is perfomred in-place. |
| 233 | +
|
| 234 | + Can be used for adding colors to ScaffoldGraph visualizations. |
| 235 | +
|
| 236 | + Parameters |
| 237 | + ---------- |
| 238 | + graph : ScaffoldGraph |
| 239 | + Input ScaffoldGraph |
| 240 | + attribute : str |
| 241 | + Key for the attibute from which to calculate a color. |
| 242 | + cmap : str or matplotlib.colors.Colormap |
| 243 | + A matplotlib cmap or name of a cmap e.g. 'BuPu' for |
| 244 | + calculating a nodes colour. |
| 245 | + label : str, optional |
| 246 | + The attribute label to use for storing the color. |
| 247 | + The default is 'color'. |
| 248 | +
|
| 249 | + See Also |
| 250 | + -------- |
| 251 | + color_scaffold_nodes_by_attribute |
| 252 | +
|
| 253 | + """ |
| 254 | + color_nodes_by_attribute(graph, attribute, cmap, 'molecule', label) |
| 255 | + |
| 256 | + |
120 | 257 | def add_root_node(graph): |
121 | 258 | """Add a root node to a scaffoldgraph. |
122 | 259 |
|
|
0 commit comments