Adding multiprocessing features

pydft/dft.py

@@ -20,7 +20,8 @@ def __init__(self,
                  nangpts:int = 110,
                  lmax:int = 8,
                  fdpts:int=7,
-                 normalize:bool = True):
+                 normalize:bool = True,
+                 parallel:bool = True):
         """
         Constructs the DFT class
 
@@ -45,6 +46,8 @@ def __init__(self,
             density
         verbose : bool, optional
             whether to provide verbose output, by default False
+        parallel : bool, optional
+            whether to use multiprocessing features, by default False
         """
         self.__mol = mol
         self.__integrator = PyQInt()
@@ -57,6 +60,7 @@ def __init__(self,
         self.__lmax = lmax
         self.__functional = functional
         self.__normalize = normalize
+        self.__parallel = parallel
 
         # keep track of time
         self.calctimes = {
@@ -352,7 +356,8 @@ def __setup(self):
                                        nangpts=self.__nangpts,
                                        lmax=self.__lmax,
                                        fdpts=self.__fdpts,
-                                       functional=self.__functional)
+                                       functional=self.__functional,
+                                       parallel=self.__parallel)
         self.__molgrid.initialize() # molecular grid uses late initialization
 
         # build one-electron matrices; because these matrices are Hermetian,

pydft/moleculargrid.py

@@ -7,7 +7,9 @@
 from pyqint import PyQInt
 import pyqint as pq
 import time
+import multiprocessing
 from .xcfunctionals import Functionals
+from functools import partial
 
 class MolecularGrid:
     def __init__(self, 
@@ -17,7 +19,8 @@ def __init__(self,
                  nangpts:int=110, 
                  lmax:int=8,
                  fdpts:int=7,
-                 functional:str='svwn5'):
+                 functional:str='svwn5',
+                 parallel:bool=False):
         """
         Construct MolecularGrid
 
@@ -37,6 +40,8 @@ def __init__(self,
             number of grid point in finite difference scheme, by default 7
         functional : str, optional
             exchange-correlation functional, by default 'svwn5'
+        parallel : bool, optional
+             whether to use multiprocessing features, by default False
         """
         # keep track of build times
         self.construct_times = {}
@@ -51,6 +56,7 @@ def __init__(self,
         self.__basis = cgfs
         self.__functionals = Functionals(functional)
         self.__is_initialized = False
+        self.__enable_parallel = True
 
     def initialize(self):
         """
@@ -62,7 +68,7 @@ def initialize(self):
         if self.__is_initialized:
             return
 
-        self.__build_molecular_grid()
+        self.__build_molecular_grid(self.__enable_parallel)
         self.__build_amplitudes()
         self.__is_initialized = True
 
@@ -344,17 +350,17 @@ def calculate_coulombic_matrix(self) -> np.array:
         # for each grid point determine the Hartree potential from the 
         # spherical harmonic coefficients with respect to each atomic center
         self.__ugpts = np.einsum('ijk,ik,ijk->k', ulmgpts, 
-                                                 self.__rigridpoints, 
-                                                 self.__ylmgpts,
-                                                 optimize=True)
+                                                  self.__rigridpoints, 
+                                                  self.__ylmgpts,
+                                                  optimize=True)
 
         # construct coulombic repulsion matrix by integrating the interaction
         # of the hartree potential with the basis function amplitudes
         J = np.einsum('k,ik,jk,k->ij', self.__ugpts, 
-                                     self.__fullgrid_amplitudes,
-                                     self.__fullgrid_amplitudes,
-                                     self.__mgw,
-                                     optimize=True)
+                                       self.__fullgrid_amplitudes,
+                                       self.__fullgrid_amplitudes,
+                                       self.__mgw,
+                                       optimize=True)
 
         return J
 
@@ -541,7 +547,7 @@ def get_gradient_at_points(self,
         #t2 = np.einsum('ij,jkl,ik->kl', P, grads, amps)
         #return t1 + t2
 
-        return 2.0 * np.einsum('ij,ikl,jk->kl', P, grads, amps)
+        return 2.0 * np.einsum('ij,ikl,jk->kl', P, grads, amps, optimize=True)
 
     def calculate_coulomb_potential_at_points(self, 
                                               pts:np.ndarray) -> np.ndarray:
@@ -729,7 +735,7 @@ def calculate_weights_at_points(self,
 
         return mweights
 
-    def __build_molecular_grid(self):
+    def __build_molecular_grid(self, build_parallel=False):
         """
         Build the molecular grid from the atomic grids
         """
@@ -841,20 +847,61 @@ def __build_molecular_grid(self):
         self.__ylmgpts = np.ndarray((len(self.__atoms), 
                                          (self.__lmax+1)**2, 
                                          np.prod(self.__mweights.shape)))
-        for i,at in enumerate(self.__atoms):    # loop over atoms
-            lmctr = 0
-            for l in range(0, self.__lmax+1):   # loop over spherical harmonics
-                for m in range(-l, l+1):
-                    self.__ylmgpts[i,lmctr,:] = spherical_harmonic(l, m, \
-                                                self.__theta_gridpoints[i,:], 
-                                                self.__phi_gridpoints[i,:])
-                    lmctr += 1
+
+        # perform parallellized calculation of spherical harmonics
+        if build_parallel:
+            atoms = list(range(len(self.__atoms)))
+            inputs = zip([self.__theta_gridpoints[i,:] for i in atoms], 
+                         [self.__phi_gridpoints[i,:] for i in atoms])
+            calculate_partial = partial(self.build_ylmgpts_atom, 
+                                        lmax=self.__lmax, 
+                                        npts=np.prod(self.__mweights.shape))
+            with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+                result = pool.map(calculate_partial, inputs)
+            for i,res in enumerate(result):
+                self.__ylmgpts[i,:,:] = res
+        else: # non-parallelized evaluation of spherical harmonics
+            for i,at in enumerate(self.__atoms):    # loop over atoms
+                self.__ylmgpts[i,:,:] = self.__build_ylmgpts_atom(i)
         self.construct_times['spherical_harmonics'] = time.time() - st
 
         # collect complete weights for the full molecular grid
         weights = np.array([an.get_weights() for an in self.__atomgrids]).flatten()
         self.__mgw = np.multiply(weights, self.__mweights.flatten())
 
+    @staticmethod
+    def build_ylmgpts_atom(Omega, lmax, npts):
+        """
+        Get the spherical harmonic parameters for a single atom
+        
+        This function is 'pickable' and can be used in a multiprocessing pool
+        """
+        lmctr = 0
+        theta = Omega[0]
+        phi = Omega[1]
+        res = np.ndarray(((lmax+1)**2, npts))
+        for l in range(0, lmax+1):   # loop over spherical harmonics
+            for m in range(-l, l+1):
+                res[lmctr,:] = spherical_harmonic(l, m, \
+                                                  theta, 
+                                                  phi)
+                lmctr += 1
+        return res
+
+    def __build_ylmgpts_atom(self, atidx):
+        """
+        Get the spherical harmonic parameters for a single atom
+        """
+        lmctr = 0
+        res = np.ndarray(((self.__lmax+1)**2, np.prod(self.__mweights.shape)))
+        for l in range(0, self.__lmax+1):   # loop over spherical harmonics
+            for m in range(-l, l+1):
+                res[lmctr,:] = spherical_harmonic(l, m, \
+                                                  self.__theta_gridpoints[atidx,:], 
+                                                  self.__phi_gridpoints[atidx,:])
+                lmctr += 1
+        return res
+
     def __build_amplitudes(self):
         """
         Precalculate the basis function amplitudes at the selected grid

tests/test_finite_difference_grid.py

@@ -8,11 +8,13 @@
 
 from pydft import MoleculeBuilder, DFT
 
-class TestDFT(unittest.TestCase):
+class TestFiniteDifferenceSchemes(unittest.TestCase):
 
     def test_co(self):
         """
-        Test DFT calculation of CO molecule
+        Test DFT calculation of CO molecule using various finite difference
+        schemes (always diagonal matrices, but with different number of
+                 non-zero diagonals)
         """
         mol_builder = MoleculeBuilder()
         mol = mol_builder.from_name('CO')

tests/test_parallel.py

@@ -0,0 +1,29 @@
+import unittest
+import sys
+import os
+import numpy as np
+
+# add a reference to load the pyDFT module
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+
+from pydft import MoleculeBuilder, DFT
+
+class TestParallel(unittest.TestCase):
+
+
+    def test_benzene(self):
+        """
+        Test DFT calculation of benzene molecule
+        """
+        mol_builder = MoleculeBuilder()
+        mol = mol_builder.from_name('benzene')
+
+        # construct dft object
+        dft = DFT(mol, basis='sto3g', parallel=True)
+        energy = dft.scf()
+
+        answer = -228.06536332917378
+        np.testing.assert_almost_equal(energy, answer, 4)
+
+if __name__ == '__main__':
+    unittest.main()