Skip to content

Visualization

to_image(safe_str, fragments=None, legend=None, mol_size=(300, 300), use_svg=True, highlight_mode='lasso', highlight_bond_width_multiplier=12, **kwargs)

Display a safe string by highlighting the fragments that make it.

Parameters:

Name Type Description Default
safe_str str

the safe string to display

 required
fragments Optional[Union[str, Mol]]

list of fragment to highlight on the molecules. If None, will use safe decomposition of the molecule.

 None
legend Union[str, None]

A string to use as the legend under the molecule.

 None
mol_size Union[Tuple[int, int], int]

The size of the image to be returned

 (300, 300)
use_svg Optional[bool]

Whether to return an svg or png image

 True
highlight_mode Optional[str]

the highlight mode to use. One of ["lasso", "fill", "color"]. If None, no highlight will be shown

 'lasso'
highlight_bond_width_multiplier int

the multiplier to use for the bond width when using the 'fill' mode

 12
**kwargs Any

Additional arguments to pass to the drawing function. See RDKit documentation related to MolDrawOptions for more details at https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html.

 {}
Source code in safe/viz.py 
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
def to_image(
    safe_str: str,
    fragments: Optional[Union[str, dm.Mol]] = None,
    legend: Union[str, None] = None,
    mol_size: Union[Tuple[int, int], int] = (300, 300),
    use_svg: Optional[bool] = True,
    highlight_mode: Optional[str] = "lasso",
    highlight_bond_width_multiplier: int = 12,
    **kwargs: Any,
):
    """Display a safe string by highlighting the fragments that make it.

    Args:
        safe_str: the safe string to display
        fragments: list of fragment to highlight on the molecules. If None, will use safe decomposition of the molecule.
        legend: A string to use as the legend under the molecule.
        mol_size: The size of the image to be returned
        use_svg: Whether to return an svg or png image
        highlight_mode: the highlight mode to use. One of ["lasso", "fill", "color"]. If None, no highlight will be shown
        highlight_bond_width_multiplier: the multiplier to use for the bond width when using the 'fill' mode
        **kwargs: Additional arguments to pass to the drawing function. See RDKit
            documentation related to `MolDrawOptions` for more details at
            https://www.rdkit.org/docs/source/rdkit.Chem.Draw.rdMolDraw2D.html.

    """

    kwargs["legends"] = legend
    kwargs["mol_size"] = mol_size
    kwargs["use_svg"] = use_svg
    if highlight_bond_width_multiplier is not None:
        kwargs["highlightBondWidthMultiplier"] = highlight_bond_width_multiplier

    if highlight_mode == "color":
        kwargs["continuousHighlight"] = False
        kwargs["circleAtoms"] = kwargs.get("circleAtoms", False) or False

    if isinstance(fragments, (str, dm.Mol)):
        fragments = [fragments]

    if fragments is None and highlight_mode is not None:
        fragments = [
            sf.decode(x, as_mol=False, remove_dummies=False, ignore_errors=False)
            for x in safe_str.split(".")
        ]
    elif fragments and len(fragments) > 0:
        parsed_fragments = []
        for fg in fragments:
            if isinstance(fg, str) and dm.to_mol(fg) is None:
                fg = sf.decode(fg, as_mol=False, remove_dummies=False, ignore_errors=False)
            parsed_fragments.append(fg)
        fragments = parsed_fragments
    else:
        fragments = []
    mol = dm.to_mol(safe_str, remove_hs=False)
    cm = plt.get_cmap("gist_rainbow")
    current_colors = [cm(1.0 * i / len(fragments)) for i in range(len(fragments))]

    if highlight_mode == "lasso":
        return dm.viz.lasso_highlight_image(mol, fragments, **kwargs)

    atom_indices = []
    bond_indices = []
    atom_colors = {}
    bond_colors = {}

    for i, frag in enumerate(fragments):
        frag = dm.from_smarts(frag)
        atom_matches, bond_matches = dm.substructure_matching_bonds(mol, frag)
        atom_matches = list(itertools.chain(*atom_matches))
        bond_matches = list(itertools.chain(*bond_matches))
        atom_indices.extend(atom_matches)
        bond_indices.extend(bond_matches)
        atom_colors.update({x: current_colors[i] for x in atom_matches})
        bond_colors.update({x: current_colors[i] for x in bond_matches})

    return dm.viz.to_image(
        mol,
        highlight_atom=[atom_indices],
        highlight_bond=[bond_indices],
        highlightAtomColors=[atom_colors],
        highlightBondColors=[bond_colors],
        **kwargs,
    )