Merge pull request #7 from josan82/improvements

Improvements

chimera_p4/aux_functions.py

@@ -3,6 +3,7 @@
 from __future__ import print_function
 
 import chimera
+from chimera import Vector
 
 import numpy as np
 
@@ -405,7 +406,7 @@ def _fix_nitro(molecule):
 				atom.SetIsAromatic(aromHolder)
 
 #Function to convert from Chimera to RDKit
-def _chimera_to_rdkit(molecule, sanitize=True):
+def _chimera_to_rdkit(molecule, sanitize=True, useXformCoord=True):
 	from rdkit.Chem import Mol, RWMol, Atom, Conformer
 	mol = Mol()
 	emol = RWMol(mol)
@@ -415,7 +416,10 @@ def _chimera_to_rdkit(molecule, sanitize=True):
 	for atom in molecule.atoms:
 		a = Atom(atom.element.number)
 		atom_map[atom] = i = emol.AddAtom(a)
-		emol.GetConformer().SetAtomPosition(i, atom.coord().data())
+		if useXformCoord:
+			emol.GetConformer().SetAtomPosition(i, atom.xformCoord())
+		else:
+			emol.GetConformer().SetAtomPosition(i, atom.coord())
 	for bond in molecule.bonds:
 		a1, a2 = bond.atoms
 		if hasattr(bond, 'order') and bond.order:
@@ -433,9 +437,10 @@ def _chimera_to_rdkit(molecule, sanitize=True):
 	return mol, atom_map
 
 #### Open3Align code
-def _apply_atom_positions(rdkit_mol, chimera_mol, atom_map, rdkit_confId=0):
+def _return_atom_positions(rdkit_mol, chimera_mol, atom_map, rdkit_confId=0):
 	from chimera import Coord, Point
 
+	atom_positions = {}
 	for rdkit_atom in rdkit_mol.GetAtoms():
 		chimera_atom = list(atom_map.keys())[list(atom_map.values()).index(rdkit_atom.GetIdx())]
 
@@ -444,9 +449,16 @@ def _apply_atom_positions(rdkit_mol, chimera_mol, atom_map, rdkit_confId=0):
 		atom_position[1] = list(rdkit_mol.GetConformer(id=rdkit_confId).GetAtomPosition(rdkit_atom.GetIdx()))[1]
 		atom_position[2] = list(rdkit_mol.GetConformer(id=rdkit_confId).GetAtomPosition(rdkit_atom.GetIdx()))[2]
 
-		chimera_atom.setCoord(Coord(atom_position))
-
+		atom_positions[chimera_atom] = Coord(atom_position)
+
+	return atom_positions
+
+def _apply_atom_positions(chimera_mol, new_pos_dict):
+	for atom in new_pos_dict.keys():
+		atom.setCoord(new_pos_dict[atom])
 	chimera_mol.computeIdatmTypes()
+	#Undo possible previous rotations/translations of the molecule
+	chimera_mol.openState.xform = chimera.Xform.identity()
 
 def _del_chimeraHs(molecule):
 	for atom in molecule.atoms:

chimera_p4/core.py

@@ -23,8 +23,8 @@
 from rdkit.Chem.Features import FeatDirUtilsRD as FeatDirUtils
 from rdkit.Chem.AllChem import BuildFeatureFactory, MMFFGetMoleculeProperties, EmbedMultipleConfs, AddHs, RemoveHs
 
-from aux_functions import _chimera_to_rdkit, _GetAcceptor1FeatVects, _GetDonor2FeatVects, _MergeFeatPoints, _apply_atom_positions, _del_chimeraHs
-#import aux_functions as aux
+from aux_functions import _chimera_to_rdkit, _GetAcceptor1FeatVects, _GetDonor2FeatVects, _MergeFeatPoints, _return_atom_positions, _apply_atom_positions, _del_chimeraHs
+
 
 FEATURES_FILE = os.path.join(os.path.dirname(__file__), 'BaseFeatures.fdef')
 
@@ -103,7 +103,7 @@ def _build_vrml(self, bild, name=None):
 		except chimera.NotABug:
 			print(bild)
 		else:
-			chimera.openModels.add(vrml, baseId=self._id, subid=self._subid)
+			chimera.openModels.add(vrml, baseId=self._id, subid=self._subid, shareXform=0)
 			self._subid += 1
 		return vrml
 
@@ -234,8 +234,8 @@ def chimera_p4(molecules_sel, mergeTol=1.5, minRepeats=1, showVectors=True, fami
 
 	return True
 
-#Function from plume with small changes in sanitization
-def align_o3a(reference, probe, transform=True, sanitize=True, nConformers=0, **kwargs):
+#Function based on plume subalign
+def align_o3a(reference, probe, sanitize=True, nConformers=0, **kwargs):
 	_del_chimeraHs(reference)
 	_del_chimeraHs(probe)
 	rdk_reference, ref_map = _chimera_to_rdkit(reference)
@@ -253,27 +253,23 @@ def align_o3a(reference, probe, transform=True, sanitize=True, nConformers=0, **
 		o3as = GetO3AForProbeConfs(rdk_probe, rdk_reference, numThreads=0)
 
 		highest_conf_score = 0.0
-		conf_id = 0
-		for o3a in o3as:
+		for conf_id, o3a in enumerate(o3as):
 			score_conf = o3a.Score()
 			if score_conf > highest_conf_score:
 				highest_conf_score = score_conf
 				o3a_result = o3a
 				highest_conf_id = conf_id
-			conf_id += 1
 	else:
 		o3a_result = GetO3A(rdk_probe, rdk_reference, probe_params, reference_params)
 		highest_conf_id = 0
-
-	#rmsd, xform = o3a_result.Trans()
-	if transform:
-		o3a_result.Align()
-		_apply_atom_positions(rdk_probe, probe, probe_map, rdkit_confId=highest_conf_id)
-
-	return o3a_result   #return the alignment
+
+	o3a_result.Align()
+	new_probe_pos = _return_atom_positions(rdk_probe, probe, probe_map, rdkit_confId=highest_conf_id)
+
+	return o3a_result.Score(), new_probe_pos  #return the alignment score and the new atom positions
 
 #New function
-def open3align(molecules_sel, transform=True, nConformers=0, _gui=None):
+def open3align(molecules_sel, transform=True, nConformers=0, reference=None, _gui=None):
 	#Delete possible previous pharmacophores
 	chimera.openModels.remove(chimera.openModels.list(id=100))
 	runCommand("2dlabels delete *")
@@ -282,38 +278,43 @@ def open3align(molecules_sel, transform=True, nConformers=0, _gui=None):
 
 	if _gui:
 		molecules = molecules_sel
+		references = [reference] if reference else molecules
 	else:
 		molecules = molecules_sel.molecules()
+		references = [reference.molecules()] if reference else molecules
 
 	if not len(molecules) > 1:
 		raise chimera.UserError("At least 2 molecules are needed to do an alignment")
 
 	#Calculating the best scored alignment
 	max_score = 0.0
-	for i, reference in enumerate(molecules):
-		msg = "Processing molecule {} of {}".format(i+1, len(molecules))
+	for i, reference in enumerate(references):
+		msg = "Processing molecule {} of {}".format(i+1, len(references))
 		if _gui:
 			_gui.status(msg, blankAfter=0)
 		else:
 			chimera.statusline.show_message(msg)
 		align_score = 0.0
+		new_atom_pos = {}
 		for probe in molecules:
-			if(molecules.index(probe) is not molecules.index(reference)):
-				o3a = align_o3a(reference, probe, transform=False, nConformers=nConformers)
-				align_score += o3a.Score()
+			score, new_atom_pos[probe] = align_o3a(reference, probe, nConformers=nConformers)
+			align_score += score
 
 		if align_score > max_score:
 			max_score = align_score
-			if transform:
-				for probe in molecules:
-					if(molecules.index(probe) is not molecules.index(reference)):
-						o3a = align_o3a(reference, probe, transform=True, nConformers=nConformers)
+			max_new_atom_pos = new_atom_pos
+
+	if transform:
+		for mol in max_new_atom_pos.keys():
+			_apply_atom_positions(mol, max_new_atom_pos[mol])
+
 	if not _gui:
 		msg = "The score of the best alignment is {}".format(max_score)
 		chimera.statusline.show_message(msg)
 
 	return max_score
 
+
 ### GUI Code
 class Controller(object):
 

chimera_p4/gui.py

@@ -8,7 +8,7 @@
 # Chimera stuff
 import chimera
 from chimera.baseDialog import ModelessDialog
-from chimera.widgets import MoleculeScrolledListBox
+from chimera.widgets import MoleculeScrolledListBox, MoleculeOptionMenu
 
 # Additional 3rd parties
 
@@ -55,6 +55,8 @@ def __init__(self, *args, **kwargs):
 		self._minRepeats = tk.IntVar()
 		self._check1 = tk.BooleanVar()
 		self._check2 = tk.BooleanVar()
+		self._check3 = tk.BooleanVar()
+		self._useRef = tk.BooleanVar()
 		self._showFamily = {}
 
 		# Fire up
@@ -71,6 +73,7 @@ def _set_defaults(self):
 		self._mergeTol.set(1.5)
 		self._showLegend = True
 		self._showVectors = True
+		self._useRef.set(False)
 		default_families = ['Donor','Acceptor','NegIonizable','PosIonizable','Aromatic', 'LumpedHydrophobe']
 		for family in _featColors.keys():
 			if family in default_families:
@@ -90,12 +93,16 @@ def fill_in_ui(self, parent):
 
 		#Second frame to perform alignments
 		self.ui_o3align_frame = tk.LabelFrame(self.canvas, text="Perform an alignment with open3align")
-		tk.Label(self.ui_o3align_frame, text='Number of conformers:').grid(row=0, column=0, padx=5, pady=5, sticky='w')
-		self.ui_nconformers = tk.Entry(self.ui_o3align_frame, textvariable=self._nConformers, bg='white', width=6).grid(row=0, column=1, padx=5, pady=5, sticky='w')
+		self.ui_molecule = MoleculeOptionMenu(self.ui_o3align_frame)
+		self.ui_molecule.grid(row=0, columnspan=3, padx=5, pady=5, sticky='news')
+		self.ui_o3align_frame.ui_useRef = tk.Checkbutton(self.ui_o3align_frame, text="Use this molecule as reference for the alignment", variable=self._useRef)
+		self.ui_o3align_frame.ui_useRef.grid(row=1, columnspan=3, padx=5, pady=5, sticky='n')
+		tk.Label(self.ui_o3align_frame, text='Number of conformers:').grid(row=2, column=0, padx=5, pady=5, sticky='w')
+		self.ui_nconformers = tk.Entry(self.ui_o3align_frame, textvariable=self._nConformers, bg='white', width=6).grid(row=2, column=1, padx=5, pady=5, sticky='w')
 		self.ui_o3align_frame.rowconfigure(0, weight=1)
 		self.ui_o3align_frame.columnconfigure(1, weight=1)
 		self.ui_o3align_btn = tk.Button(self.ui_o3align_frame, text='Align', command=self._cmd_o3align_btn)
-		self.ui_o3align_btn.grid(row=0, column=2, padx=5, pady=5)
+		self.ui_o3align_btn.grid(row=2, column=2, padx=5, pady=5)
 		self.ui_o3align_frame.pack(expand=True, fill='both', padx=5, pady=5)
 
 
@@ -114,8 +121,9 @@ def fill_in_ui(self, parent):
 	def _cmd_o3align_btn(self):
 		molecules = self.ui_molecules.getvalue()
 		nConformers = self._nConformers.get()
+		ref = self.ui_molecule.getvalue() if self._useRef.get() else None
 		try:
-			max_score = open3align(molecules, nConformers=nConformers, _gui=self)
+			max_score = open3align(molecules, nConformers=nConformers, reference=ref, _gui=self)
 			msg = "Alignment done! Score: {}".format(max_score)
 			self.status(msg, color='green', blankAfter=0)
 		except Exception as e:
@@ -149,6 +157,11 @@ def _accept_adv_btn(self):
 		self._showVectors = self._check2.get()
 		for family in _featColors.keys():
 			self._showFamily[family] = self._tempFamilies[family].get()
+		self._useRef = self._tempUseRef.get()
+		if self._useRef:
+			self._baseMol = self.ui_molecule.getvalue()
+		else:
+			self._baseMol = None
 		self.ui_input_opt_window.destroy()
 
 	def _cancel_adv_btn(self):