@@ -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,42 @@ 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+ # Generate the RDkit molecule from the RMG molecule, saving mapping
381+ # in order to match the atoms in the rdmol with the atoms in the
382+ # RMG molecule (which is required to extract coordinates).
383+ #
384+ # partial sanitization is used to allow molecules to fail RDKit's
385+ # bond order and implicit hydrogen assignments, and possibly also
386+ # RDKit's aromaticity perception, while still drawing the molecule.
387+ # this can happen because RMG uses partial bond orders and atom types
388+ # that RDKit doesn't understand, though RDKit can still generate
389+ # coordinates for the molecule.
390+ rdmol , rd_atom_idx = self .molecule .to_rdkit_mol (remove_h = False ,
391+ return_mapping = True ,
392+ sanitize = "partial" )
393+
394+ AllChem .Compute2DCoords (rdmol )
395+
396+ # Extract the coordinates from each atom.
397+ rd_conformer = rdmol .GetConformer (0 )
398+ for atom in atoms :
399+ index = rd_atom_idx [atom ]
400+ point = rd_conformer .GetAtomPosition (index )
401+ coordinates [index , :] = [point .x * 0.6 , point .y * 0.6 ]
402+
403+ # RDKit generates some molecules more vertically than horizontally,
404+ # especially linear ones. This will reflect any molecule taller than
405+ # it is wide across the line y=x.
406+ ranges = np .ptp (coordinates , axis = 0 )
407+ if ranges [1 ] > ranges [0 ]:
408+ temp = np .copy (coordinates )
409+ coordinates [:, 0 ] = temp [:, 1 ]
410+ coordinates [:, 1 ] = temp [:, 0 ]
382411
383- if not use_rdkit :
412+ else :
413+ logging .warning ("Using deprecated molecule drawing algorithm; undesired behavior may occur. Consider using use_rdkit=True." )
384414 if len (self .cycles ) > 0 :
385415 # Cyclic molecule
386416 backbone = self ._find_cyclic_backbone ()
@@ -438,32 +468,6 @@ def _generate_coordinates(self, fix_surface_sites=True):
438468 # minimize likelihood of overlap
439469 self ._generate_neighbor_coordinates (backbone )
440470
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-
467471 # For surface species
468472 if fix_surface_sites and self .molecule .contains_surface_site ():
469473 if len (self .molecule .atoms ) == 1 :
0 commit comments