File size: 6,158 Bytes
9dd777e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
import re

import numpy as np
from rdkit import Chem, DataStructs
from rdkit.Chem import (
    AllChem,
    Draw,
    rdFMCS,
    rdMolAlign,
)


def save_as_svg(svg_content, filename, num_mols):
    """Save SVG content to a file."""
    with open(filename, 'w') as file:
        data = str(svg_content.data)
        data = data.replace('1500', str(500*num_mols))
        file.write(data)


def align_molecules_2D(ref_mol, to_align_mol):
    AllChem.Compute2DCoords(ref_mol)
    AllChem.Compute2DCoords(to_align_mol)
    # Find the maximum common substructure and use it to align molecules
    mcs = rdFMCS.FindMCS([ref_mol, to_align_mol])
    mcs_mol = Chem.MolFromSmarts(mcs.smartsString)
    ref_match = ref_mol.GetSubstructMatch(mcs_mol)
    align_match = to_align_mol.GetSubstructMatch(mcs_mol)
    atom_map = list(zip(align_match, ref_match))
    rdMolAlign.AlignMol(to_align_mol, ref_mol, atomMap=atom_map)
    return to_align_mol


def align_molecules_by_coordinates(ref_mol, to_align_mol):
    # Find the maximum common substructure
    AllChem.Compute2DCoords(to_align_mol)
    mcs = rdFMCS.FindMCS([ref_mol, to_align_mol])
    mcs_mol = Chem.MolFromSmarts(mcs.smartsString)
    ref_match = ref_mol.GetSubstructMatch(mcs_mol)
    align_match = to_align_mol.GetSubstructMatch(mcs_mol)

    # Copy the coordinates from the reference molecule to the molecule to be aligned
    ref_conf = ref_mol.GetConformer()
    align_conf = to_align_mol.GetConformer()
    for ref_idx, align_idx in zip(ref_match, align_match):
        ref_pos = ref_conf.GetAtomPosition(ref_idx)
        align_conf.SetAtomPosition(align_idx, ref_pos)

    return to_align_mol


def draw_molecule_to_svg(mol, size=(500, 500), scale=1.0):
    drawer = Draw.rdMolDraw2D.MolDraw2DSVG(size[0], size[1])
    drawer.drawOptions().fixedBondLength = scale
    drawer.DrawMolecule(mol)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    svg = re.sub(r'\<\?xml.*?\?\>', '', svg)  # Remove XML declaration
    svg = svg.replace('<svg', '<g').replace(
        '</svg>', '</g>')  # Replace svg tags with g tags
    return svg


def combine_svgs(svgs, output_filename, dimensions=None, size=(500, 500), xy_shifts=None):
    if dimensions is None:
        dimensions = (len(svgs), 1)
    if xy_shifts is None:
        xy_shifts = [(0, 0) for i in range(dimensions[0]*dimensions[1])]

    width, height = size
    grid_width, grid_height = dimensions
    # Include only one XML declaration and the opening <svg> tag
    combined_svg = f'<?xml version="1.0" standalone="no"?>\n'
    combined_svg += f'<svg width="{grid_width * width}px" height="{grid_height * height}px" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">\n'

    # Arrange SVGs in a grid
    for i, (svg, xy_shift) in enumerate(zip(svgs, xy_shifts)):
        x = (i % grid_width) * width
        y = (i // grid_width) * height
        combined_svg += f'<g transform="translate({x+xy_shift[0]},{y-xy_shift[1]})">{svg}</g>\n'

    combined_svg += '</svg>'
    with open(output_filename, 'w') as file:
        file.write(combined_svg)


def draw_molecule_with_highlighted_bonds(mol, bonds_to_highlight):
    """
    Draws a molecule with specified atoms and bonds highlighted.

    Parameters:
    - smiles (str): SMILES string for the molecule.
    - atoms_to_highlight (set): Set of atom indices to highlight.
    - bonds_to_highlight (list): List of bond indices to highlight.
    - highlight_bond_colors (dict): Dictionary mapping bond indices to colors.
    """
    # Create molecule from SMILES

    # Initialize drawer
    d2d = Draw.rdMolDraw2D.MolDraw2DSVG(350*2, 300*2)

    # Set drawing options
    d2d.drawOptions().useBWAtomPalette()
    d2d.drawOptions().continuousHighlight = False
    d2d.drawOptions().highlightBondWidthMultiplier = 24
    d2d.drawOptions().setHighlightColour((0, 0, 1))
    d2d.drawOptions().fillHighlights = False

    # Draw the molecule with highlights
    d2d.DrawMolecule(mol,
                     highlightAtoms=[],
                     highlightBonds=bonds_to_highlight)
    d2d.FinishDrawing()

    # Convert drawing to image and display
    svg = d2d.GetDrawingText()
    svg = svg.replace('svg:', '')

    return svg


def align_mol_2D_ver2(template, query):
    mcs = rdFMCS.FindMCS([template, query])
    patt = Chem.MolFromSmarts(mcs.smartsString)

    query_match = query.GetSubstructMatch(patt)
    template_match = template.GetSubstructMatch(patt)

    rms = AllChem.AlignMol(query, template, atomMap=list(
        zip(query_match, template_match)))
    return template, query


def transform_molecule(mol, degrees, translate_x=0, translate_y=0, flip_x_axis=False):
    """Apply rotation, translation, and optionally flip the molecule."""
    radians = np.deg2rad(degrees)
    rotation_matrix = np.array([
        [np.cos(radians), -np.sin(radians), 0],
        [np.sin(radians),  np.cos(radians), 0],
        [0,               0,               1]
    ])
    AllChem.Compute2DCoords(mol)

    conf = mol.GetConformer()
    for i in range(conf.GetNumAtoms()):
        pos = np.array(conf.GetAtomPosition(i))
        new_pos = np.dot(rotation_matrix, pos)
        new_pos[0] += translate_x  # Translate along the x-axis
        new_pos[1] += translate_y  # Translate along the y-axis
        if flip_x_axis:
            new_pos[1] = -new_pos[1]  # Flip along the x-axis
        conf.SetAtomPosition(i, new_pos)


def tailored_framework_example(mol_ms):
    # remove lone atoms
    # define all atoms to be atom number 1
    # define all bonds to be single bonds

    mol_ms_w = Chem.RWMol(mol_ms)
    atom_idx_to_remove = []
    for atom in mol_ms_w.GetAtoms():
        # lone atom. Need to remove it to create the generic framework.
        if atom.GetDegree() == 1:
            atom_idx_to_remove.append(atom.GetIdx())
            continue
        atom.SetAtomicNum(0)

    for bond in mol_ms_w.GetBonds():
        bond.SetBondType(Chem.rdchem.BondType.SINGLE)

    atom_idx_to_remove.sort(reverse=True)
    for atom_idx in atom_idx_to_remove:
        mol_ms_w.RemoveAtom(atom_idx)

    mol_ms_new = mol_ms_w.GetMol()
    return mol_ms_new