From 054f45bf1f097113771d50dbb01eeaaf7f9143d8 Mon Sep 17 00:00:00 2001 From: Ksenia Date: Tue, 18 Jun 2024 19:38:05 +0200 Subject: [PATCH] Improve docs and fix minor bugs (2) Co-authored-by: Liam Marsh --- qstack/fields/dori.py | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/qstack/fields/dori.py b/qstack/fields/dori.py index fb682eb..b56da24 100644 --- a/qstack/fields/dori.py +++ b/qstack/fields/dori.py @@ -11,7 +11,7 @@ def eval_rho_dm(mol, ao, dm, deriv=2): Taken from pyscf/dft/numint.py and modified to return second derivative matrices. Args: - mol : an instance of :class:`Mole` + mol : an instance of :class:`pyscf.gto.Mole` ao : 3D array of shape (*,ngrids,nao): ao[0] : atomic oribitals values on the grid ao[1:4] : atomic oribitals derivatives values (if deriv>=1) @@ -57,7 +57,7 @@ def eval_rho_df(mol, ao, c, deriv=2): for a fitted density. Args: - mol : an instance of :class:`Mole` + mol : an instance of :class:`pyscf.gto.Mole` ao : 3D array of shape (*,ngrids,nao): ao[0] : atomic oribitals values on the grid ao[1:4] : atomic oribitals derivatives values (if deriv>=1) @@ -74,7 +74,8 @@ def eval_rho_df(mol, ao, c, deriv=2): 3D array of (3,3,ngrids) to store 2nd derivatives (if deriv==2) ''' - rho_all = np.tensordot(np.atleast_3d(ao), c, 1) # corresponds to np.einsum('xip,p->xi', np.atleast_3d(ao), c) + maxdim = 1 if deriv==0 else (4 if deriv==1 else 10) + rho_all = np.tensordot(ao[:maxdim], c, 1) # corresponds to np.einsum('xip,p->xi', ao[:maxdim], c) if deriv==0: return rho_all[0] if deriv==1: @@ -90,7 +91,7 @@ def compute_rho(mol, coords, dm=None, c=None, deriv=2, eps=1e-4): r'''Wrapper to calculate the electron density and the density derivatives. Args: - mol : an instance of :class:`Mole` + mol : an instance of :class:`pyscf.gto.Mole` coords : 2D array of (ngrids,3) Grid coordinates (in Bohr) Kwargs: @@ -121,7 +122,8 @@ def compute_rho(mol, coords, dm=None, c=None, deriv=2, eps=1e-4): return rho good_idx = np.where(rho>=eps)[0] drho_dr = np.zeros((3,len(coords))) - d2rho_dr2 = np.zeros((3,3,len(coords))) + if deriv==2: + d2rho_dr2 = np.zeros((3,3,len(coords))) if len(good_idx)>0: ao = eval_ao(mol, coords[good_idx], deriv=deriv) ret = eval_rho(ao, deriv=deriv) @@ -210,7 +212,7 @@ def compute_dori_num(mol, coords, dm=None, c=None, eps=1e-4, dx=1e-4): See documentation to compute_dori(). Args: - mol : an instance of :class:`Mole` + mol : an instance of :class:`pyscf.gto.Mole` coords : 2D array of (ngrids,3) Grid coordinates (in Bohr) Kwargs: @@ -261,7 +263,7 @@ def dori_on_grid(mol, coords, dm=None, c=None, eps=1e-4, alg='analytical', mem=1 """Wrapper to compute DORI on a given grid Args: - mol (pyscf Mole): pyscf Mole object. + mol : an instance of :class:`pyscf.gto.Mole` coords : 2D array of (ngrids,3) Grid coordinates (in Bohr) Kwargs: @@ -302,12 +304,7 @@ def dori_on_grid(mol, coords, dm=None, c=None, eps=1e-4, alg='analytical', mem=1 d2rho_dr2 = np.zeros((3, 3, len(coords))) for i in grid_chunks: s = np.s_[i:i+dgrid] - rho_i, drho_dr_i, d2rho_dr2_i = compute_rho(mol, coords[s], dm=dm, c=c, eps=eps) - # Yes the data is copied around too much (three times). - # We tried to make compute_rho() write directly to these but didn't gain a lot. - rho[s] = rho_i - drho_dr[:,s] = drho_dr_i - d2rho_dr2[:,:,s] = d2rho_dr2_i + rho[s], drho_dr[:,s], d2rho_dr2[:,:,s] = compute_rho(mol, coords[s], dm=dm, c=c, eps=eps) dori = compute_dori(rho, drho_dr, d2rho_dr2, eps=eps) s2rho = compute_s2rho(rho, d2rho_dr2, eps=eps) return dori, rho, s2rho @@ -332,7 +329,7 @@ def dori(mol, dm=None, c=None, """Compute DORI Args: - mol : an instance of :class:`Mole`. + mol : an instance of :class:`pyscf.gto.Mole` Kwargs: dm : 2D array of (nao,nao) Density matrix (confilicts with c)