1 #!/bin/env python 2 # 3 # File: RDKitCalculateEnergy.py 4 # Author: Manish Sud <msud@san.rr.com> 5 # 6 # Copyright (C) 2026 Manish Sud. All rights reserved. 7 # 8 # The functionality available in this script is implemented using RDKit, an 9 # open source toolkit for cheminformatics developed by Greg Landrum. 10 # 11 # This file is part of MayaChemTools. 12 # 13 # MayaChemTools is free software; you can redistribute it and/or modify it under 14 # the terms of the GNU Lesser General Public License as published by the Free 15 # Software Foundation; either version 3 of the License, or (at your option) any 16 # later version. 17 # 18 # MayaChemTools is distributed in the hope that it will be useful, but without 19 # any warranty; without even the implied warranty of merchantability of fitness 20 # for a particular purpose. See the GNU Lesser General Public License for more 21 # details. 22 # 23 # You should have received a copy of the GNU Lesser General Public License 24 # along with MayaChemTools; if not, see <http://www.gnu.org/licenses/> or 25 # write to the Free Software Foundation Inc., 59 Temple Place, Suite 330, 26 # Boston, MA, 02111-1307, USA. 27 # 28 29 from __future__ import print_function 30 31 import os 32 import sys 33 import time 34 import re 35 import multiprocessing as mp 36 37 # RDKit imports... 38 try: 39 from rdkit import rdBase 40 from rdkit import Chem 41 from rdkit.Chem import AllChem 42 except ImportError as ErrMsg: 43 sys.stderr.write("\nFailed to import RDKit module/package: %s\n" % ErrMsg) 44 sys.stderr.write("Check/update your RDKit environment and try again.\n\n") 45 sys.exit(1) 46 47 # MayaChemTools imports... 48 sys.path.insert(0, os.path.join(os.path.dirname(sys.argv[0]), "..", "lib", "Python")) 49 try: 50 from docopt import docopt 51 import MiscUtil 52 import RDKitUtil 53 except ImportError as ErrMsg: 54 sys.stderr.write("\nFailed to import MayaChemTools module/package: %s\n" % ErrMsg) 55 sys.stderr.write("Check/update your MayaChemTools environment and try again.\n\n") 56 sys.exit(1) 57 58 ScriptName = os.path.basename(sys.argv[0]) 59 Options = {} 60 OptionsInfo = {} 61 62 63 def main(): 64 """Start execution of the script.""" 65 66 MiscUtil.PrintInfo( 67 "\n%s (RDKit v%s; MayaChemTools v%s; %s): Starting...\n" 68 % (ScriptName, rdBase.rdkitVersion, MiscUtil.GetMayaChemToolsVersion(), time.asctime()) 69 ) 70 71 (WallClockTime, ProcessorTime) = MiscUtil.GetWallClockAndProcessorTime() 72 73 # Retrieve command line arguments and options... 74 RetrieveOptions() 75 76 # Process and validate command line arguments and options... 77 ProcessOptions() 78 79 # Perform actions required by the script... 80 CalculateEnergy() 81 82 MiscUtil.PrintInfo("\n%s: Done...\n" % ScriptName) 83 MiscUtil.PrintInfo("Total time: %s" % MiscUtil.GetFormattedElapsedTime(WallClockTime, ProcessorTime)) 84 85 86 def CalculateEnergy(): 87 """Calculate single point energy.""" 88 89 # Setup a molecule reader... 90 MiscUtil.PrintInfo("\nProcessing file %s..." % OptionsInfo["Infile"]) 91 Mols = RDKitUtil.ReadMolecules(OptionsInfo["Infile"], **OptionsInfo["InfileParams"]) 92 93 # Set up a molecule writer... 94 Writer = RDKitUtil.MoleculesWriter(OptionsInfo["Outfile"], **OptionsInfo["OutfileParams"]) 95 if Writer is None: 96 MiscUtil.PrintError("Failed to setup a writer for output fie %s " % OptionsInfo["Outfile"]) 97 MiscUtil.PrintInfo("Generating file %s..." % OptionsInfo["Outfile"]) 98 99 MolCount, ValidMolCount, EnergyFailedCount = ProcessMolecules(Mols, Writer) 100 101 if Writer is not None: 102 Writer.close() 103 104 MiscUtil.PrintInfo("\nTotal number of molecules: %d" % MolCount) 105 MiscUtil.PrintInfo("Number of valid molecules: %d" % ValidMolCount) 106 MiscUtil.PrintInfo("Number of molecules failed during energy calculation: %d" % EnergyFailedCount) 107 MiscUtil.PrintInfo("Number of ignored molecules: %d" % (MolCount - ValidMolCount + EnergyFailedCount)) 108 109 110 def ProcessMolecules(Mols, Writer): 111 """Process and calculate energy of molecules.""" 112 113 if OptionsInfo["MPMode"]: 114 return ProcessMoleculesUsingMultipleProcesses(Mols, Writer) 115 else: 116 return ProcessMoleculesUsingSingleProcess(Mols, Writer) 117 118 119 def ProcessMoleculesUsingSingleProcess(Mols, Writer): 120 """Process and calculate energy of molecules using a single process.""" 121 122 MiscUtil.PrintInfo("\nCalculating energy...") 123 124 (MolCount, ValidMolCount, EnergyFailedCount) = [0] * 3 125 for Mol in Mols: 126 MolCount += 1 127 128 if Mol is None: 129 continue 130 131 if RDKitUtil.IsMolEmpty(Mol): 132 if not OptionsInfo["QuietMode"]: 133 MolName = RDKitUtil.GetMolName(Mol, MolCount) 134 MiscUtil.PrintWarning("Ignoring empty molecule: %s" % MolName) 135 continue 136 137 ValidMolCount += 1 138 139 CalcStatus, Energy = CalculateMoleculeEnergy(Mol, MolCount) 140 if CalcStatus: 141 Energy = "%.2f" % Energy 142 else: 143 if not OptionsInfo["QuietMode"]: 144 MolName = RDKitUtil.GetMolName(Mol, MolCount) 145 MiscUtil.PrintWarning("Failed to calculate energy for molecule %s" % MolName) 146 147 EnergyFailedCount += 1 148 continue 149 150 WriteMolecule(Writer, Mol, Energy) 151 152 return (MolCount, ValidMolCount, EnergyFailedCount) 153 154 155 def ProcessMoleculesUsingMultipleProcesses(Mols, Writer): 156 """Process and calculate energy of molecules using process.""" 157 158 MiscUtil.PrintInfo("\nCalculating energy using multiprocessing...") 159 160 MPParams = OptionsInfo["MPParams"] 161 162 # Setup data for initializing a worker process... 163 InitializeWorkerProcessArgs = ( 164 MiscUtil.ObjectToBase64EncodedString(Options), 165 MiscUtil.ObjectToBase64EncodedString(OptionsInfo), 166 ) 167 168 # Setup a encoded mols data iterable for a worker process... 169 WorkerProcessDataIterable = RDKitUtil.GenerateBase64EncodedMolStrings(Mols) 170 171 # Setup process pool along with data initialization for each process... 172 MiscUtil.PrintInfo( 173 "\nConfiguring multiprocessing using %s method..." 174 % ("mp.Pool.imap()" if re.match("^Lazy$", MPParams["InputDataMode"], re.I) else "mp.Pool.map()") 175 ) 176 MiscUtil.PrintInfo( 177 "NumProcesses: %s; InputDataMode: %s; ChunkSize: %s\n" 178 % ( 179 MPParams["NumProcesses"], 180 MPParams["InputDataMode"], 181 ("automatic" if MPParams["ChunkSize"] is None else MPParams["ChunkSize"]), 182 ) 183 ) 184 185 ProcessPool = mp.Pool(MPParams["NumProcesses"], InitializeWorkerProcess, InitializeWorkerProcessArgs) 186 187 # Start processing... 188 if re.match("^Lazy$", MPParams["InputDataMode"], re.I): 189 Results = ProcessPool.imap(WorkerProcess, WorkerProcessDataIterable, MPParams["ChunkSize"]) 190 elif re.match("^InMemory$", MPParams["InputDataMode"], re.I): 191 Results = ProcessPool.map(WorkerProcess, WorkerProcessDataIterable, MPParams["ChunkSize"]) 192 else: 193 MiscUtil.PrintError( 194 'The value, %s, specified for "--inputDataMode" is not supported.' % (MPParams["InputDataMode"]) 195 ) 196 197 (MolCount, ValidMolCount, EnergyFailedCount) = [0] * 3 198 for Result in Results: 199 MolCount += 1 200 MolIndex, EncodedMol, CalcStatus, Energy = Result 201 202 if EncodedMol is None: 203 continue 204 ValidMolCount += 1 205 206 Mol = RDKitUtil.MolFromBase64EncodedMolString(EncodedMol) 207 if CalcStatus: 208 Energy = "%.2f" % Energy 209 else: 210 if not OptionsInfo["QuietMode"]: 211 MolName = RDKitUtil.GetMolName(Mol, MolCount) 212 MiscUtil.PrintWarning("Failed to calculate energy for molecule %s" % MolName) 213 214 EnergyFailedCount += 1 215 continue 216 217 WriteMolecule(Writer, Mol, Energy) 218 219 return (MolCount, ValidMolCount, EnergyFailedCount) 220 221 222 def InitializeWorkerProcess(*EncodedArgs): 223 """Initialize data for a worker process.""" 224 225 global Options, OptionsInfo 226 227 MiscUtil.PrintInfo("Starting process (PID: %s)..." % os.getpid()) 228 229 # Decode Options and OptionInfo... 230 Options = MiscUtil.ObjectFromBase64EncodedString(EncodedArgs[0]) 231 OptionsInfo = MiscUtil.ObjectFromBase64EncodedString(EncodedArgs[1]) 232 233 234 def WorkerProcess(EncodedMolInfo): 235 """Process data for a worker process.""" 236 237 MolIndex, EncodedMol = EncodedMolInfo 238 239 CalcStatus = False 240 Energy = None 241 242 if EncodedMol is None: 243 return [MolIndex, None, CalcStatus, Energy] 244 245 Mol = RDKitUtil.MolFromBase64EncodedMolString(EncodedMol) 246 if RDKitUtil.IsMolEmpty(Mol): 247 if not OptionsInfo["QuietMode"]: 248 MolName = RDKitUtil.GetMolName(Mol, (MolIndex + 1)) 249 MiscUtil.PrintWarning("Ignoring empty molecule: %s" % MolName) 250 return [MolIndex, None, CalcStatus, Energy] 251 252 CalcStatus, Energy = CalculateMoleculeEnergy(Mol, (MolIndex + 1)) 253 254 return [ 255 MolIndex, 256 RDKitUtil.MolToBase64EncodedMolString( 257 Mol, PropertyPickleFlags=Chem.PropertyPickleOptions.MolProps | Chem.PropertyPickleOptions.PrivateProps 258 ), 259 CalcStatus, 260 Energy, 261 ] 262 263 264 def CalculateMoleculeEnergy(Mol, MolNum=None): 265 "Calculate energy." 266 267 Status = True 268 Energy = None 269 ConfID = -1 270 271 if OptionsInfo["UseUFF"]: 272 UFFMoleculeForcefield = AllChem.UFFGetMoleculeForceField(Mol, confId=ConfID) 273 if UFFMoleculeForcefield is None: 274 Status = False 275 else: 276 Energy = UFFMoleculeForcefield.CalcEnergy() 277 elif OptionsInfo["UseMMFF"]: 278 MMFFMoleculeProperties = AllChem.MMFFGetMoleculeProperties(Mol, mmffVariant=OptionsInfo["MMFFVariant"]) 279 MMFFMoleculeForcefield = AllChem.MMFFGetMoleculeForceField(Mol, MMFFMoleculeProperties, confId=ConfID) 280 if MMFFMoleculeForcefield is None: 281 Status = False 282 else: 283 Energy = MMFFMoleculeForcefield.CalcEnergy() 284 else: 285 MiscUtil.PrintError( 286 "Couldn't calculate nergy: Specified forcefield, %s, is not supported" % OptionsInfo["ForceField"] 287 ) 288 289 return (Status, Energy) 290 291 292 def WriteMolecule(Writer, Mol, Energy): 293 """Write molecule.""" 294 295 Mol.SetProp(OptionsInfo["EnergyLabel"], Energy) 296 Writer.write(Mol) 297 298 299 def ProcessOptions(): 300 """Process and validate command line arguments and options.""" 301 302 MiscUtil.PrintInfo("Processing options...") 303 304 # Validate options... 305 ValidateOptions() 306 307 OptionsInfo["Infile"] = Options["--infile"] 308 OptionsInfo["InfileParams"] = MiscUtil.ProcessOptionInfileParameters( 309 "--infileParams", Options["--infileParams"], Options["--infile"] 310 ) 311 312 OptionsInfo["Outfile"] = Options["--outfile"] 313 OptionsInfo["OutfileParams"] = MiscUtil.ProcessOptionOutfileParameters( 314 "--outfileParams", Options["--outfileParams"] 315 ) 316 317 OptionsInfo["Overwrite"] = Options["--overwrite"] 318 319 if re.match("^UFF$", Options["--forceField"], re.I): 320 ForceField = "UFF" 321 UseUFF = True 322 UseMMFF = False 323 elif re.match("^MMFF$", Options["--forceField"], re.I): 324 ForceField = "MMFF" 325 UseUFF = False 326 UseMMFF = True 327 else: 328 MiscUtil.PrintError( 329 'The value, %s, specified for "--forceField" is not supported.' % (Options["--forceField"],) 330 ) 331 332 MMFFVariant = "MMFF94" if re.match("^MMFF94$", Options["--forceFieldMMFFVariant"], re.I) else "MMFF94s" 333 334 OptionsInfo["ForceField"] = ForceField 335 OptionsInfo["MMFFVariant"] = MMFFVariant 336 OptionsInfo["UseMMFF"] = UseMMFF 337 OptionsInfo["UseUFF"] = UseUFF 338 339 if UseMMFF: 340 OptionsInfo["EnergyLabel"] = "%s_Energy" % MMFFVariant 341 else: 342 OptionsInfo["EnergyLabel"] = "%s_Energy" % ForceField 343 344 OptionsInfo["MPMode"] = True if re.match("^yes$", Options["--mp"], re.I) else False 345 OptionsInfo["MPParams"] = MiscUtil.ProcessOptionMultiprocessingParameters("--mpParams", Options["--mpParams"]) 346 347 OptionsInfo["QuietMode"] = True if re.match("^yes$", Options["--quiet"], re.I) else False 348 349 350 def RetrieveOptions(): 351 """Retrieve command line arguments and options.""" 352 353 # Get options... 354 global Options 355 Options = docopt(_docoptUsage_) 356 357 # Set current working directory to the specified directory... 358 WorkingDir = Options["--workingdir"] 359 if WorkingDir: 360 os.chdir(WorkingDir) 361 362 # Handle examples option... 363 if "--examples" in Options and Options["--examples"]: 364 MiscUtil.PrintInfo(MiscUtil.GetExamplesTextFromDocOptText(_docoptUsage_)) 365 sys.exit(0) 366 367 368 def ValidateOptions(): 369 """Validate option values.""" 370 371 MiscUtil.ValidateOptionTextValue("-f, --forceField", Options["--forceField"], "UFF MMFF") 372 MiscUtil.ValidateOptionTextValue(" --forceFieldMMFFVariant", Options["--forceFieldMMFFVariant"], "MMFF94 MMFF94s") 373 374 MiscUtil.ValidateOptionFilePath("-i, --infile", Options["--infile"]) 375 MiscUtil.ValidateOptionFileExt("-i, --infile", Options["--infile"], "sdf sd mol") 376 377 MiscUtil.ValidateOptionFileExt("-o, --outfile", Options["--outfile"], "sdf sd") 378 MiscUtil.ValidateOptionsOutputFileOverwrite( 379 "-o, --outfile", Options["--outfile"], "--overwrite", Options["--overwrite"] 380 ) 381 MiscUtil.ValidateOptionsDistinctFileNames( 382 "-i, --infile", Options["--infile"], "-o, --outfile", Options["--outfile"] 383 ) 384 385 MiscUtil.ValidateOptionTextValue("--mp", Options["--mp"], "yes no") 386 MiscUtil.ValidateOptionTextValue("-q, --quiet", Options["--quiet"], "yes no") 387 388 389 # Setup a usage string for docopt... 390 _docoptUsage_ = """ 391 RDKitCalculateEnergy.py - Calculate energy 392 393 Usage: 394 RDKitCalculateEnergy.py [--forceField <UFF, or MMFF>] [--forceFieldMMFFVariant <MMFF94 or MMFF94s>] 395 [--infileParams <Name,Value,...>] [--mp <yes or no>] [--mpParams <Name,Value,...>] 396 [ --outfileParams <Name,Value,...> ] [--overwrite] [--quiet <yes or no>] [-w <dir>] -i <infile> -o <outfile> 397 RDKitCalculateEnergy.py -h | --help | -e | --examples 398 399 Description: 400 Calculate single point energy for molecules using a specified forcefield. The 401 molecules must have 3D coordinates in input file. 402 403 The supported input file formats are: Mol (.mol), SD (.sdf, .sd) 404 405 The supported output file formats are: SD (.sdf, .sd) 406 407 Options: 408 -f, --forceField <UFF, MMFF> [default: MMFF] 409 Forcefield method to use for energy calculation. Possible values: Universal Force 410 Field (UFF) [ Ref 81 ] or Merck Molecular Mechanics Force Field [ Ref 83-87 ] . 411 --forceFieldMMFFVariant <MMFF94 or MMFF94s> [default: MMFF94] 412 Variant of MMFF forcefield to use for energy calculation. 413 -e, --examples 414 Print examples. 415 -h, --help 416 Print this help message. 417 -i, --infile <infile> 418 Input file name. 419 --infileParams <Name,Value,...> [default: auto] 420 A comma delimited list of parameter name and value pairs for reading 421 molecules from files. The supported parameter names for different file 422 formats, along with their default values, are shown below: 423 424 SD, MOL: removeHydrogens,yes,sanitize,yes,strictParsing,yes 425 426 --mp <yes or no> [default: no] 427 Use multiprocessing. 428 429 By default, input data is retrieved in a lazy manner via mp.Pool.imap() 430 function employing lazy RDKit data iterable. This allows processing of 431 arbitrary large data sets without any additional requirements memory. 432 433 All input data may be optionally loaded into memory by mp.Pool.map() 434 before starting worker processes in a process pool by setting the value 435 of 'inputDataMode' to 'InMemory' in '--mpParams' option. 436 437 A word to the wise: The default 'chunkSize' value of 1 during 'Lazy' input 438 data mode may adversely impact the performance. The '--mpParams' section 439 provides additional information to tune the value of 'chunkSize'. 440 --mpParams <Name,Value,...> [default: auto] 441 A comma delimited list of parameter name and value pairs to configure 442 multiprocessing. 443 444 The supported parameter names along with their default and possible 445 values are shown below: 446 447 chunkSize, auto 448 inputDataMode, Lazy [ Possible values: InMemory or Lazy ] 449 numProcesses, auto [ Default: mp.cpu_count() ] 450 451 These parameters are used by the following functions to configure and 452 control the behavior of multiprocessing: mp.Pool(), mp.Pool.map(), and 453 mp.Pool.imap(). 454 455 The chunkSize determines chunks of input data passed to each worker 456 process in a process pool by mp.Pool.map() and mp.Pool.imap() functions. 457 The default value of chunkSize is dependent on the value of 'inputDataMode'. 458 459 The mp.Pool.map() function, invoked during 'InMemory' input data mode, 460 automatically converts RDKit data iterable into a list, loads all data into 461 memory, and calculates the default chunkSize using the following method 462 as shown in its code: 463 464 chunkSize, extra = divmod(len(dataIterable), len(numProcesses) * 4) 465 if extra: chunkSize += 1 466 467 For example, the default chunkSize will be 7 for a pool of 4 worker processes 468 and 100 data items. 469 470 The mp.Pool.imap() function, invoked during 'Lazy' input data mode, employs 471 'lazy' RDKit data iterable to retrieve data as needed, without loading all the 472 data into memory. Consequently, the size of input data is not known a priori. 473 It's not possible to estimate an optimal value for the chunkSize. The default 474 chunkSize is set to 1. 475 476 The default value for the chunkSize during 'Lazy' data mode may adversely 477 impact the performance due to the overhead associated with exchanging 478 small chunks of data. It is generally a good idea to explicitly set chunkSize to 479 a larger value during 'Lazy' input data mode, based on the size of your input 480 data and number of processes in the process pool. 481 482 The mp.Pool.map() function waits for all worker processes to process all 483 the data and return the results. The mp.Pool.imap() function, however, 484 returns the the results obtained from worker processes as soon as the 485 results become available for specified chunks of data. 486 487 The order of data in the results returned by both mp.Pool.map() and 488 mp.Pool.imap() functions always corresponds to the input data. 489 -o, --outfile <outfile> 490 Output file name. 491 --outfileParams <Name,Value,...> [default: auto] 492 A comma delimited list of parameter name and value pairs for writing 493 molecules to files. The supported parameter names for different file 494 formats, along with their default values, are shown below: 495 496 SD: kekulize,yes,forceV3000,no 497 498 --overwrite 499 Overwrite existing files. 500 -q, --quiet <yes or no> [default: no] 501 Use quiet mode. The warning and information messages will not be printed. 502 -w, --workingdir <dir> 503 Location of working directory which defaults to the current directory. 504 505 Examples: 506 To calculate single point energy using MMFF forcefield for molecules in a SD file 507 containing 3D structures and write a new SD file, type: 508 509 % RDKitCalculateEnergy.py -i Sample3D.sdf -o Sample3DOut.sdf 510 511 To run the first example in multiprocessing mode on all available CPUs 512 without loading all data into memory and write out a SD file, type: 513 514 % RDKitCalculateEnergy.py --mp yes -i Sample3D.sdf -o Sample3DOut.sdf 515 516 To run the first example in multiprocessing mode on all available CPUs 517 by loading all data into memory and write out a SD file, type: 518 519 % RDKitCalculateEnergy.py --mp yes --mpParams "inputDataMode, 520 InMemory" -i Sample3D.sdf -o Sample3DOut.sdf 521 522 To run the first example in multiprocessing mode on specific number of 523 CPUs and chunk size without loading all data into memory and write out a SD file, 524 type: 525 526 % RDKitCalculateEnergy.py --mp yes --mpParams "inputDataMode,Lazy, 527 numProcesses,4,chunkSize,8" -i Sample3D.sdf -o Sample3DOut.sdf 528 529 To calculate single point energy using UFF forcefield for molecules in a SD file 530 containing 3D structures and write a new SD file, type: 531 532 % RDKitCalculateEnergy.py -f UFF -i Sample3D.sdf -o Sample3DOut.sdf 533 534 To calculate single point energy using MMFF94s variant of MMFF forcefield for 535 molecules in a SD file containing 3D structures and write a new SD file, type: 536 537 % RDKitCalculateEnergy.py --forceField MMFF --forceFieldMMFFVariant 538 MMFF94s -i Sample3D.sdf -o Sample3DOut.sdf 539 540 Author: 541 Manish Sud(msud@san.rr.com) 542 543 See also: 544 RDKitCalculateRMSD.py, RDKitCalculateMolecularDescriptors.py, RDKitCompareMoleculeShapes.py, 545 RDKitConvertFileFormat.py, RDKitGenerateConformers.py, RDKitPerformMinimization.py 546 547 Copyright: 548 Copyright (C) 2026 Manish Sud. All rights reserved. 549 550 The functionality available in this script is implemented using RDKit, an 551 open source toolkit for cheminformatics developed by Greg Landrum. 552 553 This file is part of MayaChemTools. 554 555 MayaChemTools is free software; you can redistribute it and/or modify it under 556 the terms of the GNU Lesser General Public License as published by the Free 557 Software Foundation; either version 3 of the License, or (at your option) any 558 later version. 559 560 """ 561 562 if __name__ == "__main__": 563 main()
