From 5ddb4efff96faa3467d47479b6d60707427873ed Mon Sep 17 00:00:00 2001 From: ifilot Date: Sat, 14 Dec 2024 22:22:31 +0100 Subject: [PATCH] Adding multiprocessing features --- pydft/dft.py | 9 ++- pydft/moleculargrid.py | 85 +++++++++++++++++++++------- tests/test_finite_difference_grid.py | 6 +- tests/test_parallel.py | 29 ++++++++++ 4 files changed, 106 insertions(+), 23 deletions(-) create mode 100644 tests/test_parallel.py diff --git a/pydft/dft.py b/pydft/dft.py index eaf6933..cc1b7f5 100644 --- a/pydft/dft.py +++ b/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, diff --git a/pydft/moleculargrid.py b/pydft/moleculargrid.py index 2eacdbc..fe61fb9 100644 --- a/pydft/moleculargrid.py +++ b/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 diff --git a/tests/test_finite_difference_grid.py b/tests/test_finite_difference_grid.py index 29c07e2..159995d 100644 --- a/tests/test_finite_difference_grid.py +++ b/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') diff --git a/tests/test_parallel.py b/tests/test_parallel.py new file mode 100644 index 0000000..cfe8e2f --- /dev/null +++ b/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()