@@ -155,7 +155,7 @@ def clear(self):
155155 self .surface = None
156156 self .cr = None
157157
158- def draw (self , molecule , file_format , target = None ):
158+ def draw (self , molecule , file_format , target = None , use_rdkit = True ):
159159 """
160160 Draw the given `molecule` using the given image `file_format` - pdf, svg, ps, or
161161 png. If `path` is given, the drawing is saved to that location on disk. The
@@ -165,6 +165,9 @@ def draw(self, molecule, file_format, target=None):
165165 This function returns the Cairo surface and context used to create the
166166 drawing, as well as a bounding box for the molecule being drawn as the
167167 tuple (`left`, `top`, `width`, `height`).
168+
169+ If `use_rdkit` is True, then the RDKit 2D coordinate generation is used to generate the coordinates.
170+ If `use_rdkit` is False, then the molecule is drawn using our (deprecated) original algorithm.
168171 """
169172
170173 # The Cairo 2D graphics library (and its Python wrapper) is required for
@@ -219,13 +222,13 @@ def draw(self, molecule, file_format, target=None):
219222 if molecule .contains_surface_site ():
220223 try :
221224 self ._connect_surface_sites ()
222- self ._generate_coordinates ()
225+ self ._generate_coordinates (use_rdkit = use_rdkit )
223226 self ._disconnect_surface_sites ()
224227 except AdsorbateDrawingError as e :
225228 self ._disconnect_surface_sites ()
226- self ._generate_coordinates (fix_surface_sites = False )
229+ self ._generate_coordinates (fix_surface_sites = False , use_rdkit = use_rdkit )
227230 else :
228- self ._generate_coordinates ()
231+ self ._generate_coordinates (use_rdkit = use_rdkit )
229232 self ._replace_bonds (old_bond_dictionary )
230233
231234 # Generate labels to use
@@ -341,7 +344,7 @@ def _find_ring_groups(self):
341344 if not found :
342345 self .ringSystems .append ([cycle ])
343346
344- def _generate_coordinates (self , fix_surface_sites = True ):
347+ def _generate_coordinates (self , fix_surface_sites = True , use_rdkit = True ):
345348 """
346349 Generate the 2D coordinates to be used when drawing the current
347350 molecule. The function uses rdKits 2D coordinate generation.
@@ -372,15 +375,34 @@ def _generate_coordinates(self, fix_surface_sites=True):
372375 self .coordinates [1 , :] = [0.5 , 0.0 ]
373376 return self .coordinates
374377
375- # Decide whether we can use RDKit or have to generate coordinates ourselves
376- for atom in self .molecule .atoms :
377- if atom .charge != 0 :
378- use_rdkit = False
379- break
380- else : # didn't break
381- use_rdkit = True
378+ if use_rdkit == True :
379+ # Use RDKit 2D coordinate generation:
380+
381+ # Generate the RDkit molecule from the RDkit molecule, use geometry
382+ # in order to match the atoms in the rdmol with the atoms in the
383+ # RMG molecule (which is required to extract coordinates).
384+ self .geometry = Geometry (None , None , self .molecule , None )
385+
386+ rdmol , rd_atom_idx = self .geometry .rd_build ()
387+ AllChem .Compute2DCoords (rdmol )
388+
389+ # Extract the coordinates from each atom.
390+ for atom in atoms :
391+ index = rd_atom_idx [atom ]
392+ point = rdmol .GetConformer (0 ).GetAtomPosition (index )
393+ coordinates [index , :] = [point .x * 0.6 , point .y * 0.6 ]
394+
395+ # RDKit generates some molecules more vertically than horizontally,
396+ # Especially linear ones. This will reflect any molecule taller than
397+ # it is wide across the line y=x
398+ ranges = np .ptp (coordinates , axis = 0 )
399+ if ranges [1 ] > ranges [0 ]:
400+ temp = np .copy (coordinates )
401+ coordinates [:, 0 ] = temp [:, 1 ]
402+ coordinates [:, 1 ] = temp [:, 0 ]
382403
383- if not use_rdkit :
404+ else :
405+ logging .warning ("Using deprecated molecule drawing algorithm; undesired behavior may occur. Consider using use_rdkit=True." )
384406 if len (self .cycles ) > 0 :
385407 # Cyclic molecule
386408 backbone = self ._find_cyclic_backbone ()
@@ -438,32 +460,6 @@ def _generate_coordinates(self, fix_surface_sites=True):
438460 # minimize likelihood of overlap
439461 self ._generate_neighbor_coordinates (backbone )
440462
441- else :
442- # Use RDKit 2D coordinate generation:
443-
444- # Generate the RDkit molecule from the RDkit molecule, use geometry
445- # in order to match the atoms in the rdmol with the atoms in the
446- # RMG molecule (which is required to extract coordinates).
447- self .geometry = Geometry (None , None , self .molecule , None )
448-
449- rdmol , rd_atom_idx = self .geometry .rd_build ()
450- AllChem .Compute2DCoords (rdmol )
451-
452- # Extract the coordinates from each atom.
453- for atom in atoms :
454- index = rd_atom_idx [atom ]
455- point = rdmol .GetConformer (0 ).GetAtomPosition (index )
456- coordinates [index , :] = [point .x * 0.6 , point .y * 0.6 ]
457-
458- # RDKit generates some molecules more vertically than horizontally,
459- # Especially linear ones. This will reflect any molecule taller than
460- # it is wide across the line y=x
461- ranges = np .ptp (coordinates , axis = 0 )
462- if ranges [1 ] > ranges [0 ]:
463- temp = np .copy (coordinates )
464- coordinates [:, 0 ] = temp [:, 1 ]
465- coordinates [:, 1 ] = temp [:, 0 ]
466-
467463 # For surface species
468464 if fix_surface_sites and self .molecule .contains_surface_site ():
469465 if len (self .molecule .atoms ) == 1 :
0 commit comments