import os
import io
import math
import uuid
import tempfile
from collections import defaultdict
from typing import List, Tuple, Dict, Literal, Final, Optional, Set
from .saturate_tracer import SaturateTracer, UNTRACKED_ORIGIN
type DrawStyle = Literal['name', 'pid', 'oid', 'image']
DefaultDrawStyle: Final[Tuple[DrawStyle, ...]] = ('name', 'pid', 'oid', 'image')
type DrawLayout = Literal['dot', 'neato', 'circo']
[docs]
def get_origin_id_color_map(ids: List[int] | Tuple[int, ...]) -> Dict[int, str]:
if len(ids) == 1:
colors = ['#F88379']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 2:
colors = ['#F88379', '#FFF275']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 3:
colors = ['#F88379', '#FFF275', '#87CEFA']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 4:
colors = ['#F88379', '#FFF275', '#87CEFA', '#CBAACB']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 5:
colors = ['#F88379', '#FFB570', '#FFF275', '#87CEFA', '#CBAACB']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 6:
colors = ['#F88379', '#FFB570', '#FFF275', '#A8E6CF', '#87CEFA', '#CBAACB']
return {_id: colors[k] for k, _id in enumerate(ids)}
if len(ids) == 7:
colors = ['#F88379', '#FFB570', '#FFF275', '#A8E6CF', '#87CEFA', '#CBAACB', '#FF6F61']
return {_id: colors[k] for k, _id in enumerate(ids)}
return {_id: '#FFFF00' for _id in ids}
def _hex_to_rgb_float(hex_color) -> Tuple[float, float, float]:
hex_color = hex_color.lstrip('#')
return int(hex_color[0:2], 16) / 255.0, int(hex_color[2:4], 16) / 255.0, int(hex_color[4:6], 16) / 255.0
def _create_molecule_image(smiles: str, file_path: str, atom_highlights: Dict[int, Tuple[float, float, float]],
atom_labelings: Dict[int, str]):
"""
Create a PNG image file of a molecule with additional labels and highlights
:param smiles: The molecule SMILES to be drawn
:param file_path: The file path to save the image to
:param atom_highlights: Colors for atom IDs defined in the SMILES to be highlighted
:param atom_labelings: Custom labels for atom IDs defined in the SMILES
"""
from rdkit import Chem
from rdkit.Chem import Draw
import cairosvg
from PIL import Image
# noinspection PyUnresolvedReferences
mol = Chem.MolFromSmiles(smiles, sanitize=False)
mol.UpdatePropertyCache(strict=False)
atom_label_index_map = {
a.GetAtomMapNum(): a.GetIdx()
for i, a in enumerate(mol.GetAtoms())
if a.GetAtomMapNum() is not None and a.GetAtomMapNum() > 0 and atom_labelings[a.GetAtomMapNum()] is not None
}
custom_labels = {
a.GetIdx(): atom_labelings[a.GetAtomMapNum()] if atom_labelings[a.GetAtomMapNum()] is not None else ''
for i, a in enumerate(mol.GetAtoms())
if a.GetAtomMapNum() is not None and a.GetAtomMapNum() in atom_labelings
}
# noinspection PyUnresolvedReferences
drawer = Draw.MolDraw2DSVG(512, 512)
opts = drawer.drawOptions()
opts.clearBackground = False
opts.highlightRadius = 0.4
opts.padding = 0.04
opts.additionalAtomLabelPadding = 0.1
for idx, label in custom_labels.items():
opts.atomLabels[idx] = label
drawer.DrawMolecule(mol, highlightAtoms=[atom_label_index_map[k] for k in atom_highlights.keys() if
k in atom_label_index_map],
highlightAtomColors={atom_label_index_map[k]: v for k, v in atom_highlights.items() if
k in atom_label_index_map},
highlightBonds=[])
drawer.FinishDrawing()
svg = drawer.GetDrawingText()
cairosvg.svg2png(bytestring=svg.encode("utf-8"), write_to=file_path)
with Image.open(file_path) as im:
bb = im.getbbox()
im2 = im.crop((bb[0] - 8, bb[1] - 8, bb[2] + 8, bb[3] + 8))
im2.save(file_path)
def internal_draw(tracer: SaturateTracer, nodes, edges, file_path: Optional[str] = None,
origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = DefaultDrawStyle, horizontal: Optional[bool] = False,
layout: Optional[DrawLayout] = 'dot', graphviz_graph_attr: Optional[Dict[str, str]] = None):
import matplotlib.pyplot as plt
import networkx as nx
from networkx.drawing.nx_agraph import to_agraph
import cairosvg
if origin_id_colors_map is None:
origin_id_colors_map = {}
if highlight_position_ids is None:
highlight_position_ids = set()
temp_dir = tempfile.gettempdir()
temp_sentinel_image_path = os.path.join(temp_dir, f'{uuid.uuid4().hex}_sentinel.png')
cairosvg.svg2png(bytestring='<svg xmlns="http://www.w3.org/2000/svg" width="32" height="32"><circle cx="16" cy="16" r="14" stroke="#000000" fill="none" stroke-width="2"/></svg>'.encode('utf-8'), write_to=temp_sentinel_image_path)
if tracer.verbose:
print('[INFO] Generating %s molecule images...' % len(nodes))
species_image_paths = {}
for node in nodes:
key = tuple(x[1] for x in node)
if key not in tracer.templates_lookup:
continue
template = tracer.templates_lookup[key]
if len(template.ids) == 0 or template.smiles is None:
continue
temp_path = os.path.join(temp_dir, f'{uuid.uuid4().hex}_molecule.png')
atom_highlights = {
p: _hex_to_rgb_float(
origin_id_colors_map[o] if o != UNTRACKED_ORIGIN and o in origin_id_colors_map else '#FFFFFF'
) for (o, p) in node
}
atom_labelings = {
# An untracked origin is left off the structure rather than marked: RDKit drops
# non-ASCII atom labels, and any ASCII stand-in would be read as an element symbol
# ('O' on a carbon reads as oxygen). The origin row of the node table carries the
# sentinel instead, and an unlabelled vertex already depicts an atom with no origin.
p: '' if o == UNTRACKED_ORIGIN
else (tracer.label_display_map[o] if o in tracer.label_display_map else str(o))
for (o, p) in node
}
_create_molecule_image(tracer.templates_lookup[key].smiles, temp_path, atom_highlights, atom_labelings)
species_image_paths[node] = temp_path
def create_position_cell(x):
cell_color = 'yellow' if x[1] in highlight_position_ids else 'lightgray'
return f"""<td bgcolor="{cell_color}">{x[1] if x[1] is not None else ''}</td>"""
def create_origin_cell(x):
# Checked before the colour and display-name lookups, so an entry keyed on the untracked
# origin cannot repaint or rename the sentinel.
if x[0] == UNTRACKED_ORIGIN:
return f'<td><IMG SRC="{temp_sentinel_image_path}" SCALE="FALSE"/></td>'
cell_color = origin_id_colors_map[x[0]] if x[0] in origin_id_colors_map else 'white'
label = tracer.label_display_map[x[0]] if x[0] in tracer.label_display_map else x[0]
return f'<td bgcolor="{cell_color}">{label}</td>'
ID_COLUMN_MAX = 12
def create_ids_row_padding(n, i):
if ID_COLUMN_MAX < len(n) and i * ID_COLUMN_MAX > len(n) - ID_COLUMN_MAX:
return '<td colspan="%s"> </td>' % (ID_COLUMN_MAX - (len(n) % ID_COLUMN_MAX))
return ''
def create_position_ids_row(n):
return ''.join([
f""" <tr><td bgcolor="lightgray">p</td>{''.join(create_position_cell(x) for x in n[i * ID_COLUMN_MAX:(i + 1) * ID_COLUMN_MAX])}{create_ids_row_padding(n, i)}</tr>"""
for i in range(math.ceil(len(n) / ID_COLUMN_MAX))
])
def create_origin_ids_row(n):
return ''.join([
f""" <tr><td bgcolor="lightgray">o</td>{''.join(create_origin_cell(x) for x in n[i * ID_COLUMN_MAX:(i + 1) * ID_COLUMN_MAX])}{create_ids_row_padding(n, i)}</tr>"""
for i in range(math.ceil(len(n) / ID_COLUMN_MAX))
])
def create_image_row(n):
if n in species_image_paths:
return f'<tr><td colspan="{min(13, len(n) + 1)}"><img src="{species_image_paths[n]}" scale="true"/></td></tr>'
return ''
def create_title_row(n):
key = tuple(x[1] for x in n)
if key in tracer.templates_lookup and tracer.templates_lookup[key].name is not None:
return f'<tr><td colspan="{min(13, len(n) + 1)}">{tracer.templates_lookup[key].name}</td></tr>'
return ''
style_func = {
'name': create_title_row,
'pid': create_position_ids_row,
'oid': create_origin_ids_row,
'image': create_image_row,
}
if tracer.verbose:
print('[INFO] Preparing graph...')
graph = nx.MultiDiGraph()
for i, n in enumerate(nodes):
label = '<<table border="0" cellborder="1" cellspacing="0">' + ''.join(
style_func[s](n) for s in style if s in style_func) + '</table>>'
graph.add_node(i, label=label)
# Handle hyperedges and edges differently
if len(edges) > 0 and type(next(edges.__iter__())[0]) is tuple:
for i, edge in enumerate(edges):
color = edge[3] if len(edge) > 3 else '#000000'
if len(edge[0]) == 1 and len(edge[1]) == 1:
graph.add_edge(edge[0][0], edge[1][0], label=str(edge[2]), color=color)
else:
edge_node_id = 'he%s' % i
edge_label = f'<<table border="0" cellborder="1" cellspacing="0"><tr><td color="{color}">{edge[2]}</td></tr></table>>'
graph.add_node(edge_node_id, label=edge_label)
for educt in edge[0]:
graph.add_edge(educt, edge_node_id, color=color)
for product in edge[1]:
graph.add_edge(edge_node_id, product, color=color)
else:
graph.add_edges_from([(e[0], e[1], {'label': str(e[2]) if len(e) > 2 else ''}) for e in edges])
A = to_agraph(graph)
A.graph_attr.update({
'rankdir': 'LR' if horizontal else 'TB',
'splines': 'true',
'overlap': 'false',
'concentrate': 'false',
'pad': '0',
'margin': '0',
'dpi': '300',
'label': f'<{title}>' if title is not None else '',
'labelloc': 't',
'fontsize': '14',
})
if graphviz_graph_attr is not None:
A.graph_attr.update(graphviz_graph_attr)
A.node_attr.update({
'shape': 'plaintext',
'fontsize': '14',
'fontname': 'Helvetica',
})
A.edge_attr.update({
'arrowsize': '1.0',
'arrowhead': 'normal',
'fontsize': '14',
'fontname': 'Helvetica',
'labelfontsize': '14',
'labeldistance': '10',
'decorate': 'false',
'penwidth': '1.2',
})
# noinspection PyArgumentList
for u, v, key, data in graph.edges(keys=True, data=True):
e = A.get_edge(u, v, key=key)
for attr in ('color', 'penwidth', 'style', 'arrowhead', 'arrowsize'):
if attr in data:
e.attr[attr] = str(data[attr])
if 'labeldistance' not in e.attr:
e.attr['labeldistance'] = '1.5'
if 'labelangle' not in e.attr:
e.attr['labelangle'] = '0'
if graph.has_edge(v, u):
e.attr.setdefault('minlen', '2')
if tracer.verbose:
print('[INFO] Running graph layout...')
A.layout(prog=layout)
min_x, min_y, max_x, max_y = map(float, A.graph_attr['bb'].split(','))
figsize = ((max_x - min_x + 20) / 72.0, (max_y - min_y + 20) / 72.0)
if tracer.verbose:
print('[INFO] Drawing graph...')
if file_path:
if file_path.endswith('.dot'):
with open(file_path, 'w', encoding='utf-8') as f:
f.write(A.to_string())
else:
A.draw(file_path, format='pdf')
else:
buf = io.BytesIO()
A.draw(buf, format='png')
data = buf.getvalue()
img = plt.imread(io.BytesIO(data))
plt.figure(figsize=figsize)
plt.imshow(img)
plt.axis('off')
plt.tight_layout()
plt.show()
[docs]
def draw_reaction_network(tracer: SaturateTracer, file_path: Optional[str] = None,
origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = ('name', 'pid', 'image'),
horizontal: Optional[bool] = False, layout: Optional[DrawLayout] = 'dot',
graphviz_graph_attr: Optional[Dict[str, str]] = None):
nodes = []
node_index_map = {}
for key in tracer._template_node_indices.keys():
node_index_map[key] = len(nodes)
nodes.append(tuple((x, x) for x in key))
edges = []
for i, r in enumerate(tracer.reaction_rules):
for e in r.educts:
key = tuple(sorted(e.keys()))
if key not in node_index_map:
node_index_map[key] = len(nodes)
nodes.append(tuple((x, x) for x in key))
for p in r.products:
key = tuple(sorted(p))
if key not in node_index_map:
node_index_map[key] = len(nodes)
nodes.append(tuple((x, x) for x in key))
edges.append((
tuple(node_index_map[tuple(sorted(e.keys()))] for e in r.educts),
tuple(node_index_map[tuple(sorted(p))] for p in r.products),
r.name if r.name is not None else f'<{i}>',
r.color if r.color is not None else '#000000',
))
internal_draw(tracer, nodes, edges, file_path, origin_id_colors_map, highlight_position_ids, title, style,
horizontal, layout, graphviz_graph_attr)
[docs]
def draw(tracer: SaturateTracer, file_path: Optional[str] = None, origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = DefaultDrawStyle, horizontal: Optional[bool] = False,
layout: Optional[DrawLayout] = 'dot', graphviz_graph_attr: Optional[Dict[str, str]] = None):
edges = {(e[0], e[1], _get_edge_label(tracer, e[2]), _get_edge_color(tracer, e[2])) for e in tracer.hyperedges}
internal_draw(tracer, tracer.nodes, edges, file_path, origin_id_colors_map, highlight_position_ids, title, style,
horizontal, layout, graphviz_graph_attr)
def _get_edge_label(tracer: SaturateTracer, rule_index: int):
r = tracer.reaction_rules[rule_index]
return r.name if r.name is not None else f'<{rule_index}>'
def _get_edge_color(tracer: SaturateTracer, rule_index: int):
r = tracer.reaction_rules[rule_index]
return r.color if r.color is not None else f'#000000'
[docs]
def draw_filtered_source(tracer: SaturateTracer, origin_ids: Set[int], file_path: Optional[str] = None,
origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = DefaultDrawStyle,
horizontal: Optional[bool] = False, layout: Optional[DrawLayout] = 'dot',
graphviz_graph_attr: Optional[Dict[str, str]] = None):
nodes = []
node_index_map = {}
node_filtered_index_map = {}
for i, n in enumerate(tracer.nodes):
if any(x[0] in origin_ids for x in n):
filtered_node = tuple(x if x[0] in origin_ids else (UNTRACKED_ORIGIN, x[1]) for x in n)
if filtered_node not in node_index_map:
node_index_map[filtered_node] = len(nodes)
nodes.append(filtered_node)
node_filtered_index_map[i] = node_index_map[filtered_node]
node_indices_with_origin_id = set(node_filtered_index_map.keys())
extra_node_ids = set()
for edge in tracer.hyperedges:
if not set(edge[0]).isdisjoint(node_filtered_index_map):
extra_node_ids.update(edge[0])
extra_node_ids.update(edge[1])
extra_node_ids = extra_node_ids.difference(node_filtered_index_map.keys())
for i in extra_node_ids:
n = tracer.nodes[i]
filtered_node = tuple(x if x[0] in origin_ids else (UNTRACKED_ORIGIN, x[1]) for x in n)
if filtered_node not in node_index_map:
node_index_map[filtered_node] = len(nodes)
nodes.append(filtered_node)
node_filtered_index_map[i] = node_index_map[filtered_node]
edges = {
(
tuple(node_filtered_index_map[x] for x in e[0]),
tuple(node_filtered_index_map[x] for x in e[1]),
_get_edge_label(tracer, e[2]),
_get_edge_color(tracer, e[2])
)
for e in tracer.hyperedges
if any(x in node_indices_with_origin_id for x in e[0])
}
internal_draw(tracer, nodes, edges, file_path, origin_id_colors_map, highlight_position_ids, title, style,
horizontal, layout, graphviz_graph_attr)
[docs]
def draw_filtered_target(tracer: SaturateTracer, target_ids: Set[int], file_path: Optional[str] = None,
origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = DefaultDrawStyle,
horizontal: Optional[bool] = False, layout: Optional[DrawLayout] = 'dot',
graphviz_graph_attr: Optional[Dict[str, str]] = None):
nodes = []
node_index_map = {}
node_filtered_index_map = {}
for i, n in enumerate(tracer.nodes):
if any(x[1] in target_ids for x in n):
if n not in node_index_map:
node_index_map[n] = len(nodes)
nodes.append(n)
node_filtered_index_map[i] = node_index_map[n]
connected_targets = nodes_reachability(tracer.hyperedges, set(node_filtered_index_map.keys()))
edges = set()
for i in range(len(tracer.nodes)):
if i in node_filtered_index_map:
for j in node_filtered_index_map.keys():
if j in connected_targets[i]:
edges.add((node_filtered_index_map[i], node_filtered_index_map[j]))
internal_draw(tracer, nodes, edges, file_path, origin_id_colors_map, highlight_position_ids, title, style,
horizontal, layout, graphviz_graph_attr)
[docs]
def draw_filtered_source_target(tracer: SaturateTracer, origin_ids: Set[int], target_ids: Set[int],
file_path: Optional[str] = None, origin_id_colors_map: Optional[Dict[int, str]] = None,
highlight_position_ids: Optional[Set[int]] = None, title: Optional[str] = None,
style: Optional[Tuple[DrawStyle, ...]] = DefaultDrawStyle,
horizontal: Optional[bool] = False, layout: Optional[DrawLayout] = 'dot',
graphviz_graph_attr: Optional[Dict[str, str]] = None):
nodes = []
node_index_map = {}
node_filtered_index_map = {}
for i, n in enumerate(tracer.nodes):
if any(x[1] in target_ids for x in n):
filtered_node = tuple(
x if x[0] in origin_ids and x[1] in target_ids else (UNTRACKED_ORIGIN, x[1]) for x in n)
if filtered_node not in node_index_map:
node_index_map[filtered_node] = len(nodes)
nodes.append(filtered_node)
node_filtered_index_map[i] = node_index_map[filtered_node]
connected_targets = nodes_reachability(tracer.hyperedges, set(node_filtered_index_map.keys()))
edges = set()
for i in range(len(tracer.nodes)):
if i in node_filtered_index_map:
for j in node_filtered_index_map.keys():
if j in connected_targets[i]:
edges.add((node_filtered_index_map[i], node_filtered_index_map[j]))
internal_draw(tracer, nodes, edges, file_path, origin_id_colors_map, highlight_position_ids, title, style,
horizontal, layout, graphviz_graph_attr)
def nodes_reachability(edges, nodes: Set[int]):
"""
Based on the graph described by ``edges``, for all nodes of interest in ``nodes``, find all other nodes in
``nodes`` that are reachable without traversing any other nodes from ``nodes``.
"""
successors = defaultdict(set)
if len(edges) > 0 and type(next(edges.__iter__())[0]) is tuple:
for edge in edges:
for source in edge[0]:
successors[source].update(edge[1])
else:
for edge in edges:
successors[edge[0]].add(edge[1])
results = {}
for n in nodes:
visited = set()
stack = [n]
while stack:
node = stack.pop()
for neighbor in successors[node]:
if neighbor in visited:
continue
visited.add(neighbor)
if neighbor not in nodes:
stack.append(neighbor)
results[n] = visited & nodes
return results