File: rdkpympi.py

package info (click to toggle)
rdkit 201603.5-2
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 72,364 kB
  • ctags: 18,217
  • sloc: cpp: 167,966; python: 58,855; java: 5,318; ansic: 5,239; sql: 1,908; yacc: 1,553; lex: 1,131; makefile: 418; xml: 229; sh: 192; fortran: 183; cs: 93
file content (67 lines) | stat: -rw-r--r-- 1,856 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
# $Id$
#
# Copyright (C) 2009 Greg Landrum
#  All rights reserved
#
# Demo for using boost.mpi with the RDKit
#
# run this with : mpirun -n 4 python rdkpympi.py
#
from __future__ import print_function
from boost import mpi
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.RDLogger import logger
logger = logger()

def dividetask(data,task,silent=True):
    data=mpi.broadcast(mpi.world,data,0)

    nProcs = mpi.world.size
    chunkSize=len(data)//nProcs
    extraBits =len(data)%nProcs

    res=[]
    allRes=[]
    # the root node handles the extra pieces:
    if mpi.world.rank == 0:
        for i in range(extraBits):
          elem=data[i]
          res.append(task(elem))
          if not silent:
              logger.info('task(%d) done %d'%(mpi.world.rank,i+1))
    pos=extraBits+mpi.world.rank*chunkSize;
    for i in range(chunkSize):
        elem=data[pos]
        pos += 1
        res.append(task(elem))
        if not silent:
            logger.info('task(%d) done %d'%(mpi.world.rank,i+1))
    if mpi.world.rank==0:
        tmp=mpi.gather(mpi.world,res,0)
        for res in tmp: allRes.extend(res)
    else:
        mpi.gather(mpi.world,res,0)
    return allRes
    
if __name__=='__main__':
    from rdkit import RDConfig
    import os
    fName = os.path.join(RDConfig.RDBaseDir,'Projects','DbCLI','testData','bzr.sdf')
    if mpi.world.rank==0:
        data = [x for x in Chem.SDMolSupplier(fName)][:50]
    else:
        data=None

    def generateconformations(m):
        m = Chem.AddHs(m)
        ids=AllChem.EmbedMultipleConfs(m,numConfs=10)
        for id in ids:
            AllChem.UFFOptimizeMolecule(m,confId=id)
        return m

    cms=dividetask(data,generateconformations,silent=False)
    # report:
    if mpi.world.rank==0:
        for i,mol in enumerate(cms):
            print(i,mol.GetNumConformers())