Source code for pyscf.pbc.gto.eval_gto

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import ctypes
import numpy as np
from pyscf import lib
from pyscf.gto import moleintor
from pyscf.gto.eval_gto import _get_intor_and_comp, BLKSIZE
from pyscf.pbc.gto import _pbcintor
from pyscf import __config__

EXTRA_PREC = getattr(__config__, 'pbc_gto_eval_gto_extra_precision', 1e-2)

libpbc = _pbcintor.libpbc

[docs] def eval_gto(cell, eval_name, coords, comp=None, kpts=None, kpt=None, shls_slice=None, non0tab=None, ao_loc=None, cutoff=None, out=None, Ls=None, rcut=None): r'''Evaluate PBC-AO function value on the given grids, Args: eval_name : str:: ========================== ======================= Function Expression ========================== ======================= "GTOval_sph" \sum_T exp(ik*T) |AO> "GTOval_ip_sph" nabla \sum_T exp(ik*T) |AO> "GTOval_cart" \sum_T exp(ik*T) |AO> "GTOval_ip_cart" nabla \sum_T exp(ik*T) |AO> ========================== ======================= atm : int32 ndarray libcint integral function argument bas : int32 ndarray libcint integral function argument env : float64 ndarray libcint integral function argument coords : 2D array, shape (N,3) The coordinates of the grids. Kwargs: shls_slice : 2-element list (shl_start, shl_end). If given, only part of AOs (shl_start <= shell_id < shl_end) are evaluated. By default, all shells defined in cell will be evaluated. non0tab : 2D bool array mask array to indicate whether the AO values are zero. The mask array can be obtained by calling :func:`dft.gen_grid.make_mask` cutoff : float AO values smaller than cutoff will be set to zero. The default cutoff threshold is ~1e-22 (defined in gto/grid_ao_drv.h) out : ndarray If provided, results are written into this array. Returns: A list of 2D (or 3D) arrays to hold the AO values on grids. Each element of the list corresponds to a k-point and it has the shape (N,nao) Or shape (\*,N,nao). Examples: >>> cell = pbc.gto.M(a=numpy.eye(3)*4, atom='He 1 1 1', basis='6-31g') >>> coords = cell.get_uniform_grids([10,10,10]) >>> kpts = cell.make_kpts([3,3,3]) >>> ao_value = cell.pbc_eval_gto("GTOval_sph", coords, kpts) >>> len(ao_value) 27 >>> ao_value[0].shape (1000, 2) >>> ao_value = cell.pbc_eval_gto("GTOval_ig_sph", coords, kpts, comp=3) >>> print(ao_value.shape) >>> len(ao_value) 27 >>> ao_value[0].shape (3, 1000, 2) ''' if eval_name[:3] == 'PBC': # PBCGTOval_xxx eval_name, comp = _get_intor_and_comp(cell, eval_name[3:], comp) else: eval_name, comp = _get_intor_and_comp(cell, eval_name, comp) eval_name = 'PBC' + eval_name atm = np.asarray(cell._atm, dtype=np.int32, order='C') bas = np.asarray(cell._bas, dtype=np.int32, order='C') env = np.asarray(cell._env, dtype=np.double, order='C') natm = atm.shape[0] nbas = bas.shape[0] if kpts is None: if kpt is not None: kpts_lst = np.reshape(kpt, (1,3)) else: kpts_lst = np.zeros((1,3)) else: kpts_lst = np.reshape(kpts, (-1,3)) nkpts = len(kpts_lst) ngrids = len(coords) if non0tab is None: non0tab = np.empty(((ngrids+BLKSIZE-1)//BLKSIZE, nbas), dtype=np.uint8) # non0tab stores the number of images to be summed in real space. # Initializing it to 255 means all images should be included non0tab[:] = 0xff if ao_loc is None: ao_loc = moleintor.make_loc(bas, eval_name) if shls_slice is None: shls_slice = (0, nbas) sh0, sh1 = shls_slice nao = ao_loc[sh1] - ao_loc[sh0] out = np.empty((nkpts,comp,nao,ngrids), dtype=np.complex128) coords = np.asarray(coords, order='F') if rcut is None: rcut = _estimate_rcut(cell) if Ls is None: Ls = get_lattice_Ls(cell, rcut=rcut.max()) Ls = Ls[np.argsort(lib.norm(Ls, axis=1), kind='stable')] expLk = np.exp(1j * np.asarray(np.dot(Ls, kpts_lst.T), order='C')) with cell.with_integral_screen(cutoff): drv = getattr(libpbc, eval_name) drv(ctypes.c_int(ngrids), (ctypes.c_int*2)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p), Ls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(Ls)), expLk.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts), out.ctypes.data_as(ctypes.c_void_p), coords.ctypes.data_as(ctypes.c_void_p), rcut.ctypes.data_as(ctypes.c_void_p), non0tab.ctypes.data_as(ctypes.c_void_p), atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm), bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas), env.ctypes.data_as(ctypes.c_void_p)) ao_kpts = [] for k, kpt in enumerate(kpts_lst): v = out[k] if abs(kpt).sum() < 1e-9: v = np.asarray(v.real, order='C') v = v.transpose(0,2,1) if comp == 1: v = v[0] ao_kpts.append(v) if kpts is None or np.shape(kpts) == (3,): # A single k-point ao_kpts = ao_kpts[0] return ao_kpts
pbc_eval_gto = eval_gto def _estimate_rcut(cell): '''Cutoff raidus, above which each shell decays to a value less than the required precsion''' vol = cell.vol weight_penalty = vol # ~ V[r] * (vol/ngrids) * ngrids precision = cell.precision / max(weight_penalty, 1) rcut = [] for ib in range(cell.nbas): l = cell.bas_angular(ib) es = cell.bas_exp(ib) cs = abs(cell._libcint_ctr_coeff(ib)).max(axis=1) norm_ang = ((2*l+1)/(4*np.pi))**.5 fac = 2*np.pi/vol * cs*norm_ang/es / precision r = cell.rcut r = (np.log(fac * r**(l+1) + 1.) / es)**.5 r = (np.log(fac * r**(l+1) + 1.) / es)**.5 rcut.append(r.max()) return np.array(rcut)
[docs] def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None, discard=True): '''Get lattice-sum vectors for eval_gto ''' if dimension is None: # For atoms near the boundary of the cell, it is necessary (even in low- # dimensional systems) to include lattice translations in all 3 dimensions. if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum': dimension = cell.dimension else: dimension = 3 if rcut is None: rcut = cell.rcut if dimension == 0 or rcut <= 0: return np.zeros((1, 3)) a = cell.lattice_vectors() atom_coords = cell.atom_coords() scaled_atom_coords = np.linalg.solve(a.T, atom_coords.T).T atom_boundary_max = scaled_atom_coords[:,:dimension].max(axis=0) atom_boundary_min = scaled_atom_coords[:,:dimension].min(axis=0) atom_boundary_max[atom_boundary_max > 1] = 1 atom_boundary_min[atom_boundary_min <-1] = -1 atom_bound1 = np.diag(atom_boundary_max).dot(a[:dimension]) atom_bound2 = np.diag(atom_boundary_min).dot(a[:dimension]) def find_boundary(a): aR = np.vstack([a, atom_bound1, atom_bound2]) r = np.linalg.qr(aR.T)[1] ub = (rcut + abs(r[2,3:]).max()) / abs(r[2,2]) return ub xb = find_boundary(a[[1,2,0]]) if dimension > 1: yb = find_boundary(a[[2,0,1]]) else: yb = 0 if dimension > 2: zb = find_boundary(a) else: zb = 0 bounds = np.ceil([xb, yb, zb]).astype(int) Ts = lib.cartesian_prod((np.arange(-bounds[0], bounds[0]+1), np.arange(-bounds[1], bounds[1]+1), np.arange(-bounds[2], bounds[2]+1))) Ls = np.dot(Ts[:,:dimension], a[:dimension]) # grids with wrap_around: grids_edge ~ [-.5, .5] # regular grids: grids_edge ~ [0, 1] grids_edge = lib.cartesian_prod([[-.5, 1.]] * dimension).dot(a[:dimension]) edge_lb = grids_edge.min(axis=0) edge_ub = grids_edge.max(axis=0) grids2atm = Ls + atom_coords[:,None,:] edge_filter1 = grids2atm > edge_lb edge_filter2 = grids2atm < edge_ub grids2atm[~edge_filter1[:,:,0],0] -= edge_lb[0] grids2atm[~edge_filter1[:,:,1],1] -= edge_lb[1] grids2atm[~edge_filter1[:,:,2],2] -= edge_lb[2] grids2atm[~edge_filter2[:,:,0],0] -= edge_ub[0] grids2atm[~edge_filter2[:,:,1],1] -= edge_ub[1] grids2atm[~edge_filter2[:,:,2],2] -= edge_ub[2] grids2atm[edge_filter1 & edge_filter2] = 0. Ls_mask = (np.linalg.norm(grids2atm, axis=2) < rcut).any(axis=0) Ls = Ls[Ls_mask] return np.asarray(Ls, order='C')