#!/usr/bin/env python3 # (C) 2023 Cadence Design Systems, Inc. (Cadence) # All rights reserved. # TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is # provided to current licensees or subscribers of Cadence products or # SaaS offerings (each a "Customer"). # Customer is hereby permitted to use, copy, and modify the Sample Code, # subject to these terms. Cadence claims no rights to Customer's # modifications. Modification of Sample Code is at Customer's sole and # exclusive risk. Sample Code may require Customer to have a then # current license or subscription to the applicable Cadence offering. # THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, # EXPRESS OR IMPLIED. OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT # NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A # PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be # liable for any damages or liability in connection with the Sample Code # or its use. """Calculates and visualizes the dipole moment of a molecule using property map.""" import argparse import io import os import sys from pathlib import Path import numpy as np from openeye import oechem, oedepict, oegrapheme from PIL import Image from rich_argparse import HelpPreviewAction, RichHelpFormatter __SCRIPT_NAME__ = Path(__file__).absolute().stem __SCRIPT_DESC__ = "Depict dipole moment of a molecule." __SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme"] __SCRIPT_CATEGORIES__ = ["visualization"] def parse_options() -> argparse.Namespace: """Set up command line options.""" parser = argparse.ArgumentParser( add_help=True, formatter_class=RichHelpFormatter, description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]", ) parser.add_argument("--help-image", action=HelpPreviewAction) # input options input_group = parser.add_argument_group("Input options") input_group.add_argument( "--mol", type=str, required=True, metavar="MOL-FILE", help="input molecule file (oeb, sdf)", ) image_group = parser.add_argument_group("Image options") image_group.add_argument( "--image", type=str, required=False, metavar="IMAGE-FILE", help="output image file (PNG, SVG) (required: %(required)s) -- if no output is provided the image will be displayed on the screen", ) image_group.add_argument( "--width", type=int, default=800, help="width of output image (default: %(default)s)", ) image_group.add_argument( "--height", type=int, default=400, help="height of output image (default: %(default)s)", ) return parser.parse_args() def main() -> int: """Visualizes dipole moment.""" args = parse_options() _check_image_file(args) mol = _get_molecule(args) image = oedepict.OEImage(args.width, args.height) opts = oedepict.OE2DMolDisplayOptions( args.width, args.height, oedepict.OEScale_AutoScale ) opts.SetTitleLocation(oedepict.OETitleLocation_Hidden) tag = oechem.OEGetTag("dipole moment") if calculate_dipole_moment(mol, tag): depict_molecule_with_dipole(image, mol, opts, "dipole moment") else: depict_molecule(image, mol, opts) if args.image: oedepict.OEWriteImage(args.image, image) else: _img = Image.open(io.BytesIO(oedepict.OEWriteImageToBytes("png", image))) _img.show() return os.EX_OK def depict_molecule_with_dipole( image: oedepict.OEImage, mol: oechem.OEMolBase, opts: oedepict.OE2DMolDisplayOptions, tag: str, ) -> None: """Depict the molecule property map.""" oegrapheme.OEPrepareDepictionFrom3D(mol, True) opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome) opts.SetScale(oegrapheme.OEGetMoleculeSurfaceScale(mol, opts)) disp = oedepict.OE2DMolDisplay(mol, opts) prop_map = oegrapheme.OE2DPropMap(opts.GetBackgroundColor()) prop_map.SetLegendLocation(oegrapheme.OELegendLocation_Left) prop_map.Render(disp, tag) oedepict.OERenderMolecule(image, disp) def depict_molecule( image: oedepict.OEImage, mol: oechem.OEMolBase, opts: oedepict.OE2DMolDisplayOptions ) -> None: """Depict the molecule.""" oegrapheme.OEPrepareDepictionFrom3D(mol, True) disp = oedepict.OE2DMolDisplay(mol, opts) oedepict.OERenderMolecule(image, disp) def calculate_dipole_moment( # noqa: C901, PLR0912 mol: oechem.OEMolBase, tag: int ) -> bool: """Calculate dipole moment for each atom.""" oechem.OEMMFFAtomTypes(mol) oechem.OEMMFF94PartialCharges(mol) negative_charge = 0.0 positive_charge = 0.0 negative_center = [0.0, 0.0, 0.0] positive_center = [0.0, 0.0, 0.0] for atom in mol.GetAtoms(): charge = atom.GetPartialCharge() x, y, z = mol.GetCoords(atom) if charge < 0.0: negative_center[0] -= charge * x negative_center[1] -= charge * y negative_center[2] -= charge * z negative_charge -= charge elif charge > 0.0: positive_center[0] += charge * x positive_center[1] += charge * y positive_center[2] += charge * z positive_charge += charge for idx in range(3): negative_center[idx] = negative_center[idx] / negative_charge positive_center[idx] = positive_center[idx] / positive_charge positive_charge = min(positive_charge, negative_charge) dipole_mag = 0.0 for i in range(3): dipole_mag += (negative_center[i] - positive_center[i]) * ( negative_center[i] - positive_center[i] ) if abs(dipole_mag) < 0.001: # noqa: PLR2004 # no dipole moment return False dipole_mag *= 4.80324 * positive_charge dipole_direction = [0.0, 0.0, 0.0] for idx in range(3): dipole_direction[idx] = ( 4.80324 * positive_charge * ((positive_center[idx] - negative_center[idx]) / dipole_mag) ) dipole_center = [0.0, 0.0, 0.0] for idx in range(3): dipole_center[idx] = 0.5 * (positive_center[idx] + negative_center[idx]) dipole_values = np.zeros(mol.GetMaxAtomIdx()) for atom in mol.GetAtoms(): x, y, z = mol.GetCoords(atom) dipole_values[atom.GetIdx()] = sum( [ (x - dipole_center[0]) * dipole_direction[0], (y - dipole_center[1]) * dipole_direction[1], (z - dipole_center[2]) * dipole_direction[2], ] ) max_dipole_value = max(dipole_values) min_dipole_value = min(dipole_values) norm_dipole_values = np.zeros(mol.GetMaxAtomIdx()) for i, dipole in enumerate(dipole_values): if dipole < 0.0: norm_dipole_values[i] = -i / min_dipole_value if dipole > 0.0: norm_dipole_values[i] - +i / max_dipole_value for atom, dipole in zip(mol.GetAtoms(), dipole_values, strict=False): atom.SetData(tag, dipole) for bond in mol.GetBonds(): avg = (bond.GetBgn().GetData(tag) + bond.GetEnd().GetData(tag)) / 2.0 bond.SetData(tag, avg) return True def _check_image_file(args: argparse.Namespace) -> None: # script will terminate if there is some issues if not args.image: # image will be displayed on the screen return ext = Path(args.image).suffix[1:].upper() if not oedepict.OEIsRegisteredImageFile(ext): oechem.OEThrow.Fatal("Unknown image output type!") ofs = oechem.oeofstream() if not ofs.open(args.image): oechem.OEThrow.Fatal("Cannot open output image file!") def _get_molecule(args: argparse.Namespace) -> oechem.OEMolBase: ifs = oechem.oemolistream() if not ifs.open(args.mol): oechem.OEThrow.Fatal(f"Cannot open {args.mol} input file!") mol = oechem.OEGraphMol() if not oechem.OEReadMolecule(ifs, mol): oechem.OEThrow.Fatal(f"Cannot read molecule from {args.mol} input file!") if mol.GetDimension() != 3: # noqa: PLR2004 oechem.OEThrow.Fatal("3D coordinates are requires!") return mol setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__) setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__) setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__) setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__) if __name__ == "__main__": sys.exit(main())