Repository navigation
Expand file tree
/
Copy pathdf.py
More file actions
390 lines (342 loc) · 15.3 KB
/
Copy pathdf.py
File metadata and controls
390 lines (342 loc) · 15.3 KB
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
#!/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>
#
'''
J-metric density fitting
'''
import contextlib
import numpy
import h5py
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.df import incore
from pyscf.df import outcore
from pyscf.df import r_incore
from pyscf.df import addons
from pyscf.df import df_jk
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import _conc_mos, iden_coeffs
from pyscf.ao2mo.outcore import _load_from_h5g
from pyscf import __config__
class DF(lib.StreamObject):
r'''
Object to hold 3-index tensor
Input Attributes:
auxbasis : str or dict
Same format as the input attribute mol.basis. If auxbasis is set to
None, an optimal auxiliary basis set will be selected based on the
AO basis set, according to the records in the basis set exchange
database and the predefined mappings in `pyscf.df.addons.DEFAULT_AUXBASIS`.
For DFT methods, the selection of auxiliary basis sets is also
influenced by the xc functional. The J-FIT basis set will be used for
local and semi-local functionals (LDA, GGA, and meta-GGA). JK-FIT
basis set will be employed for hybrid and RSH functionals.
If optimal auxiliary basis sets are not available, the PySCF built-in
even-tempered Gaussian basis set will be used.
_cderi_to_save : str
If _cderi_to_save is specified, the DF integral tensor will be
saved in this file.
_cderi : str or numpy array
If _cderi is specified, the DF integral tensor will be read from
this HDF5 file (or numpy array). When the DF integral tensor is
provided from the HDF5 file, its dataset name should be consistent
with DF._dataname, which is 'j3c' by default.
The DF integral tensor :math:`V_{x,ij}` should be a 2D array in C
(row-major) convention, where x corresponds to index of auxiliary
basis, and the composite index ij is the orbital pair index. The
hermitian symmetry is assumed for the composite ij index, ie
the elements of :math:`V_{x,i,j}` with :math:`i\geq j` are existed
in the DF integral tensor. Thus the shape of DF integral tensor
is (M,N*(N+1)/2), where M is the number of auxbasis functions and
N is the number of basis functions of the orbital basis.
This attribute is not compatible between the CPU and GPU implementations.
It will not be transferred during the to_cpu() and to_gpu() calls.
blockdim : int
When reading DF integrals from disk the chunk size to load. It is
used to improve IO performance.
Intermediate Attributes (These attributes are generated during calculations
and should not be modified. Additionally, they may not be compatible between
GPU and CPU implementations.)
auxmol : Mole object
Read only Mole object to hold the auxiliary basis. auxmol is
generated automatically in the initialization step based on the
given auxbasis. It is used in the rest part of the code to
determine the problem size, the integral batches etc. This object
should NOT be modified.
_vjopt : _VHFOpt instance
An instance that caches precomputation variables to optimize the
computation of the Coulomb matrix.
_rsh_df : dict
For range-separated DFT functionals, this object stores the density fitting
instances for long-range or short-range Coulomb integrals.
'''
blockdim = getattr(__config__, 'df_df_DF_blockdim', 240)
# Store DF tensor in a format compatible to pyscf-1.1 - pyscf-1.6
_compatible_format = getattr(__config__, 'df_df_DF_compatible_format', False)
_dataname = 'j3c'
_keys = {'mol', 'auxmol'}
def __init__(self, mol, auxbasis=None):
self.mol = mol
self.stdout = mol.stdout
self.verbose = mol.verbose
self.max_memory = mol.max_memory
self._auxbasis = auxbasis
##################################################
# Following are not input options
self.auxmol = None
# If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
self._cderi_to_save = None
# If _cderi is specified, the 3C-integral tensor will be read from this file
self._cderi = None
self._vjopt = None
self._rsh_df = {} # Range separated Coulomb DF objects
__getstate__, __setstate__ = lib.generate_pickle_methods(
excludes=('_cderi_to_save', '_cderi', '_vjopt', '_rsh_df'),
reset_state=True)
@property
def auxbasis(self):
return self._auxbasis
@auxbasis.setter
def auxbasis(self, x):
if self._auxbasis != x:
self.reset()
self.auxmol = None
self._auxbasis = x
def dump_flags(self, verbose=None):
log = logger.new_logger(self, verbose)
log.info('******** %s ********', self.__class__)
if self.auxmol is None:
log.info('auxbasis = %s', self.auxbasis)
else:
log.info('auxbasis = auxmol.basis = %s', self.auxmol.basis)
log.info('max_memory = %s', self.max_memory)
if isinstance(self._cderi, str):
log.info('_cderi = %s where DF integrals are loaded (readonly).',
self._cderi)
return self
def build(self):
t0 = (logger.process_clock(), logger.perf_counter())
log = logger.Logger(self.stdout, self.verbose)
self.check_sanity()
self.dump_flags()
if self._cderi is not None and self.auxmol is None:
log.info('Skip DF.build(). Tensor _cderi will be used.')
return self
mol = self.mol
auxmol = self.auxmol = addons.make_auxmol(self.mol, self.auxbasis)
nao = mol.nao_nr()
naux = auxmol.nao_nr()
nao_pair = nao*(nao+1)//2
is_custom_storage = isinstance(self._cderi_to_save, str)
max_memory = self.max_memory - lib.current_memory()[0]
int3c = mol._add_suffix('int3c2e')
int2c = mol._add_suffix('int2c2e')
if (nao_pair*naux*8/1e6 < .9*max_memory and not is_custom_storage):
self._cderi = incore.cholesky_eri(mol, int3c=int3c, int2c=int2c,
auxmol=auxmol,
max_memory=max_memory, verbose=log)
else:
if self._cderi_to_save is None:
self._cderi_to_save = lib.NamedTemporaryFile(dir=lib.param.TMPDIR)
cderi = self._cderi_to_save
if is_custom_storage:
cderi_name = cderi
else:
cderi_name = cderi.name
if self._cderi is None:
log.info('_cderi_to_save = %s', cderi_name)
else:
# If cderi needs to be saved in
log.warn('Value of _cderi is ignored. DF integrals will be '
'saved in file %s .', cderi_name)
if self._compatible_format:
outcore.cholesky_eri(mol, cderi, dataname=self._dataname,
int3c=int3c, int2c=int2c, auxmol=auxmol,
max_memory=max_memory, verbose=log)
else:
# Store DF tensor in blocks. This is to reduce the
# initialization overhead
outcore.cholesky_eri_b(mol, cderi, dataname=self._dataname,
int3c=int3c, int2c=int2c, auxmol=auxmol,
max_memory=max_memory, verbose=log)
self._cderi = cderi
log.timer_debug1('Generate density fitting integrals', *t0)
return self
def kernel(self, *args, **kwargs):
return self.build(*args, **kwargs)
def reset(self, mol=None):
'''Reset mol and clean up relevant attributes for scanner mode'''
if mol is not None:
self.mol = mol
self.auxmol = None
self._cderi = None
self._vjopt = None
self._rsh_df = {}
return self
def loop(self, blksize=None):
if self._cderi is None:
self.build()
if blksize is None:
blksize = self.blockdim
with addons.load(self._cderi, self._dataname) as feri:
if isinstance(feri, numpy.ndarray):
naoaux = feri.shape[0]
for b0, b1 in self.prange(0, naoaux, blksize):
yield numpy.asarray(feri[b0:b1], order='C')
else:
if isinstance(feri, h5py.Group):
# starting from pyscf-1.7, DF tensor may be stored in
# block format
naoaux = feri['0'].shape[0]
def load(aux_slice):
b0, b1 = aux_slice
return _load_from_h5g(feri, b0, b1)
else:
naoaux = feri.shape[0]
def load(aux_slice):
b0, b1 = aux_slice
return numpy.asarray(feri[b0:b1])
for dat in lib.map_with_prefetch(load, self.prange(0, naoaux, blksize)):
yield dat
dat = None
def prange(self, start, end, step):
for i in range(start, end, step):
yield i, min(i+step, end)
def get_naoaux(self):
# determine naoaux with self._cderi, because DF object may be used as CD
# object when self._cderi is provided.
if self._cderi is None:
self.build()
with addons.load(self._cderi, self._dataname) as feri:
if isinstance(feri, h5py.Group):
return feri['0'].shape[0]
else:
return feri.shape[0]
def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
omega=None):
if omega is None:
return df_jk.get_jk(self, dm, hermi, with_j, with_k, direct_scf_tol)
# A temporary treatment for RSH-DF integrals
with self.range_coulomb(omega) as rsh_df:
return df_jk.get_jk(rsh_df, dm, hermi, with_j, with_k, direct_scf_tol)
def get_eri(self):
nao = self.mol.nao_nr()
nao_pair = nao * (nao+1) // 2
ao_eri = numpy.zeros((nao_pair,nao_pair))
for eri1 in self.loop():
lib.dot(eri1.T, eri1, 1, ao_eri, 1)
return ao2mo.restore(8, ao_eri, nao)
get_ao_eri = get_eri
def ao2mo(self, mo_coeffs,
compact=getattr(__config__, 'df_df_DF_ao2mo_compact', True)):
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
mo_eri = numpy.zeros((nij_pair,nkl_pair))
sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
Lij = Lkl = None
for eri1 in self.loop():
Lij = _ao2mo.nr_e2(eri1, moij, ijslice, aosym='s2', mosym=ijmosym, out=Lij)
if sym:
Lkl = Lij
else:
Lkl = _ao2mo.nr_e2(eri1, mokl, klslice, aosym='s2', mosym=klmosym, out=Lkl)
lib.dot(Lij.T, Lkl, 1, mo_eri, 1)
return mo_eri
get_mo_eri = ao2mo
@contextlib.contextmanager
def range_coulomb(self, omega):
'''Creates a temporary density fitting object for RSH-DF integrals.
In this context, only LR or SR integrals for mol and auxmol are computed.
'''
if omega is None:
omega = 0
key = '%.6f' % omega
if key in self._rsh_df:
rsh_df = self._rsh_df[key]
else:
rsh_df = self._rsh_df[key] = self.copy().reset()
if hasattr(self, '_dataname'):
rsh_df._dataname = f'{self._dataname}-lr/{key}'
logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
mol = self.mol
auxmol = self.auxmol
mol_omega = mol.omega
mol.omega = omega
auxmol_omega = None
if auxmol is not None:
auxmol_omega = auxmol.omega
auxmol.omega = omega
assert rsh_df.mol.omega == omega
if rsh_df.auxmol is not None:
assert rsh_df.auxmol.omega == omega
try:
yield rsh_df
finally:
mol.omega = mol_omega
if auxmol_omega is not None:
auxmol.omega = auxmol_omega
to_gpu = lib.to_gpu
GDF = DF
class DF4C(DF):
'''Relativistic 4-component'''
def build(self):
log = logger.Logger(self.stdout, self.verbose)
mol = self.mol
auxmol = self.auxmol = addons.make_auxmol(self.mol, self.auxbasis)
n2c = mol.nao_2c()
naux = auxmol.nao_nr()
nao_pair = n2c*(n2c+1)//2
max_memory = (self.max_memory - lib.current_memory()[0]) * .8
if nao_pair*naux*3*16/1e6*2 < max_memory:
self._cderi =(r_incore.cholesky_eri(mol, auxmol=auxmol, aosym='s2',
int3c='int3c2e_spinor', verbose=log),
r_incore.cholesky_eri(mol, auxmol=auxmol, aosym='s2',
int3c='int3c2e_spsp1_spinor', verbose=log))
else:
raise NotImplementedError
return self
def loop(self, blksize=None):
if self._cderi is None:
self.build()
if blksize is None:
blksize = self.blockdim
with addons.load(self._cderi[0], self._dataname) as ferill:
naoaux = ferill.shape[0]
with addons.load(self._cderi[1], self._dataname) as feriss: # python2.6 not support multiple with
for b0, b1 in self.prange(0, naoaux, blksize):
erill = numpy.asarray(ferill[b0:b1], order='C')
eriss = numpy.asarray(feriss[b0:b1], order='C')
yield erill, eriss
def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
omega=None):
if omega is None:
return df_jk.r_get_jk(self, dm, hermi, with_j, with_k)
with self.range_coulomb(omega) as rsh_df:
return df_jk.r_get_jk(rsh_df, dm, hermi, with_j, with_k)
def get_eri(self):
raise NotImplementedError
get_ao_eri = get_eri
def ao2mo(self, mo_coeffs):
raise NotImplementedError
get_mo_eri = ao2mo
GDF4C = DF4C