File: test_multi_colours.py

package info (click to toggle)
rdkit 202503.6-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 222,000 kB
  • sloc: cpp: 411,111; python: 78,482; ansic: 26,181; java: 8,285; javascript: 4,404; sql: 2,393; yacc: 1,626; lex: 1,267; cs: 1,090; makefile: 581; xml: 229; fortran: 183; sh: 121
file content (74 lines) | stat: -rw-r--r-- 1,843 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
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
#!/usr/bin/env python

from json import dumps

from rdkit import Chem, rdBase
from rdkit.Chem import AllChem, Draw, rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D

COLS = [(1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0), (1.0, 0.55, 1.0)]


def get_hit_atoms_and_bonds(mol, smt):
  alist = []
  blist = []
  q = Chem.MolFromSmarts(smt)
  for match in mol.GetSubstructMatches(q):
    alist.extend(match)

  for ha1 in alist:
    for ha2 in alist:
      if ha1 > ha2:
        b = mol.GetBondBetweenAtoms(ha1, ha2)
        if b:
          blist.append(b.GetIdx())

  return alist, blist


def add_colours_to_map(els, cols, col_num):
  for el in els:
    if el not in cols:
      cols[el] = []
    if COLS[col_num] not in cols[el]:
      cols[el].append(COLS[col_num])


def do_a_picture(smi, smarts, filename, label, fmt='svg'):

  with rdDepictor.UsingCoordGen(True):
    mol = Chem.MolFromSmiles(smi)
    mol = Draw.PrepareMolForDrawing(mol)

    acols = {}
    bcols = {}
    h_rads = {}
    h_lw_mult = {}

    for i, smt in enumerate(smarts):
      alist, blist = get_hit_atoms_and_bonds(mol, smt)
      col = i % 4
      add_colours_to_map(alist, acols, col)
      add_colours_to_map(blist, bcols, col)

    if fmt == 'svg':
      d = rdMolDraw2D.MolDraw2DSVG(300, 300)
      mode = 'w'
    elif fmt == 'png':
      d = rdMolDraw2D.MolDraw2DCairo(300, 300)
      mode = 'wb'
    else:
      print('unknown format {}'.format(fmt))
      return

    d.drawOptions().fillHighlights = False
    d.DrawMoleculeWithHighlights(mol, label, acols, bcols, h_rads, h_lw_mult, -1)
    d.FinishDrawing()

    with open(filename, mode) as f:
      f.write(d.GetDrawingText())


smi = 'CO[C@@H](O)C1=C(O[C@H](F)Cl)C(C#N)=C1ONNC[NH3+]'
smarts = ['CONN', 'N#CC~CO', 'C=CON', 'CONNCN']
do_a_picture(smi, smarts, 'atom_highlights_3.png', '', fmt='png')