digichem-core 6.0.0rc1__py3-none-any.whl

digichem/image/structure.py

@@ -0,0 +1,115 @@
+# General imports.
+from PIL import Image
+from io import BytesIO
+
+import rdkit.RDLogger
+from rdkit.Chem.Draw import rdMolDraw2D
+
+# Digichem imports.
+from digichem.image.base import Image_maker
+
+
+class Skeletal_image_maker(Image_maker):
+    """
+    A class for rendering skeletal-style molecule structure images.
+    """
+    
+    def __init__(self, output, atoms, *args, abs_resolution = None, rel_resolution = 100, numbering = "group", numbering_font_size = 0.7, explicit_h = False, **kwargs):
+        """
+        Constructor for Structure_image_maker objects.
+        
+        :param output: A path to an output file to write to. The extension of this path is used to determine the format of the file (eg: png, jpeg).
+        :param atoms: A list of atom of the molecule/system to render.
+        :param resolution: The width and height of the rendered image.
+        :param render_backend: The library to prefer to use for rendering, possible options are 'rdkit' or 'obabel'. If 'rdkit' is chosen but is not available, obabel will be used as a fallback. 
+        :param numbering: Whether to show atom numberings, either False (off), "atomic" (for atom wise) or "group" (for group wise).
+        :param explicit_h: Whether to show explicit H.
+        """
+        self.atoms = atoms
+        self.abs_resolution = abs_resolution
+        self.rel_resolution = rel_resolution
+        self.numbering = numbering
+        self.numbering_font_size = numbering_font_size
+        self.explicit_h = explicit_h
+        super().__init__(output, *args, **kwargs)
+        
+    @classmethod
+    def from_options(self, output, *, atoms, options, **kwargs):
+        """
+        Constructor that takes a dictionary of config like options.
+        """    
+        return self(
+            output,
+            atoms = atoms,
+            abs_resolution = options['skeletal_image']['resolution']['absolute'],
+            rel_resolution = options['skeletal_image']['resolution']['relative'],
+            #render_backend = options['skeletal_image']['render_backend'],
+            **kwargs
+        )
+        
+    def make_files(self):
+        """
+        Make the image described by this object.
+        """
+        if self.abs_resolution:
+            resolution = self.abs_resolution
+        
+        else:
+            resolution = int(self.rel_resolution * self.atoms.X_length)
+        
+        molecule = self.atoms.to_rdkit_molecule()
+        
+        # Calculate atom groupings.
+        atom_groups = self.atoms.groups
+        
+        for atom in molecule.GetAtoms():
+            
+            group = [group for group in atom_groups.values() if atom.GetIdx() +1 in [atom.index for atom in group.atoms]][0]
+        
+            if self.numbering == "both":
+                # Add atom labelling.
+                #atom.SetProp("molAtomMapNumber", str(atom.GetIdx()+1))
+                atom.SetProp("atomNote",  "{} ({})".format(group.id[0], atom.GetIdx()+1))
+            
+            elif self.numbering == "atomic":
+                atom.SetProp("atomNote",  "{}".format(atom.GetIdx()+1))
+            
+            elif self.numbering == "group":
+                atom.SetProp("atomNote",  "{}".format(group.id[0]))
+        
+        # Remove C-H, if we've been asked to.
+        if not self.explicit_h:
+            edit_mol = rdkit.Chem.EditableMol(molecule)
+            atoms = list(edit_mol.GetMol().GetAtoms())
+            for atom_index, atom in enumerate(reversed(atoms)):
+                # If this atom is a hydrogen, and it has a single bond to a carbon, delete it.
+                if atom.GetSymbol() == "H":
+                    bonds = list(atom.GetBonds())
+                    if len(bonds) == 1 and bonds[0].GetOtherAtom(atom).GetSymbol() == "C":
+                        # This is an implicit H.
+                        edit_mol.RemoveAtom(atom.GetIdx())
+                        
+            molecule = edit_mol.GetMol()
+        
+        # Different libraries for generating 2D depictions.
+        # Coordgen is supposedly superior.
+        #rdkit.Chem.AllChem.Compute2DCoords(molecule)
+        
+        ps = rdkit.Chem.rdCoordGen.CoordGenParams()
+        ps.minimizerPrecision = ps.sketcherBestPrecision
+        rdkit.Chem.rdCoordGen.AddCoords(molecule, ps)
+        rdkit.Chem.rdDepictor.NormalizeDepiction(molecule)
+        
+        # Then write the file, setting any options we need to.
+        d = rdMolDraw2D.MolDraw2DCairo(resolution, resolution)
+        d.drawOptions().annotationFontScale = self.numbering_font_size
+        d.DrawMolecule(molecule)
+        d.FinishDrawing()
+        # To save directly, but we're going to open in PIL to crop.
+        #d.WriteDrawingText(str(self.output))
+        
+        # Crop it to remove whitespace.
+        with Image.open(BytesIO(d.GetDrawingText()), "r") as im:
+            cropped_image = self.auto_crop_image(im)
+        
+        cropped_image.save(self.output)