gpu4pyscf 1.0.2__tar.gz → 1.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (146) hide show
  1. {gpu4pyscf-1.0.2/gpu4pyscf.egg-info → gpu4pyscf-1.2.0}/PKG-INFO +1 -1
  2. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/README.md +9 -6
  3. gpu4pyscf-1.2.0/gpu4pyscf/__config__.py +34 -0
  4. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/__init__.py +3 -3
  5. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/cc/ccsd_incore.py +8 -8
  6. gpu4pyscf-1.2.0/gpu4pyscf/df/df.py +353 -0
  7. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/df_jk.py +234 -113
  8. gpu4pyscf-1.2.0/gpu4pyscf/df/grad/jk.py +82 -0
  9. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/rhf.py +58 -73
  10. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/rks.py +6 -7
  11. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/uhf.py +59 -44
  12. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/uks.py +2 -3
  13. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/rhf.py +73 -81
  14. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/rks.py +20 -17
  15. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/uhf.py +60 -73
  16. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/uks.py +26 -24
  17. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/int3c2e.py +628 -543
  18. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/__init__.py +2 -3
  19. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/gks.py +1 -0
  20. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/libxc.py +14 -28
  21. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/numint.py +629 -523
  22. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/rks.py +27 -17
  23. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/roks.py +24 -20
  24. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/uks.py +23 -17
  25. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/xc_deriv.py +1 -1
  26. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/dft_driver.py +3 -3
  27. gpu4pyscf-1.2.0/gpu4pyscf/grad/__init__.py +6 -0
  28. gpu4pyscf-1.2.0/gpu4pyscf/grad/rhf.py +376 -0
  29. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/rks.py +171 -100
  30. gpu4pyscf-1.2.0/gpu4pyscf/grad/uhf.py +135 -0
  31. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/uks.py +158 -112
  32. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/gto/mole.py +13 -10
  33. gpu4pyscf-1.2.0/gpu4pyscf/gto/moleintor.py +482 -0
  34. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/__init__.py +2 -8
  35. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/dispersion.py +1 -2
  36. gpu4pyscf-1.2.0/gpu4pyscf/hessian/rhf.py +804 -0
  37. gpu4pyscf-1.2.0/gpu4pyscf/hessian/rks.py +702 -0
  38. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/uhf.py +72 -239
  39. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/uks.py +106 -132
  40. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cublas.py +2 -0
  41. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cupy_helper.py +203 -7
  42. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cusolver.py +70 -16
  43. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cutensor.py +21 -21
  44. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/logger.py +3 -3
  45. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/dfmp2.py +48 -18
  46. gpu4pyscf-1.2.0/gpu4pyscf/pbc/__init__.py +19 -0
  47. gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/__init__.py +25 -0
  48. gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/fft.py +272 -0
  49. gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/fft_jk.py +346 -0
  50. gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/__init__.py +31 -0
  51. gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/gen_grid.py +79 -0
  52. gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/numint.py +433 -0
  53. gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/rks.py +228 -0
  54. gpu4pyscf-1.2.0/gpu4pyscf/pbc/scf/__init__.py +32 -0
  55. gpu4pyscf-1.2.0/gpu4pyscf/pbc/scf/hf.py +268 -0
  56. gpu4pyscf-1.2.0/gpu4pyscf/pbc/tools/__init__.py +18 -0
  57. gpu4pyscf-1.2.0/gpu4pyscf/pbc/tools/pbc.py +103 -0
  58. gpu4pyscf-1.2.0/gpu4pyscf/pop/esp.py +221 -0
  59. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/shielding.py +25 -31
  60. gpu4pyscf-1.2.0/gpu4pyscf/qmmm/chelpg.py +127 -0
  61. gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/__init__.py +16 -0
  62. gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/itrf.py +1171 -0
  63. gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/mm_mole.py +425 -0
  64. gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/tools.py +307 -0
  65. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/_response_functions.py +65 -38
  66. gpu4pyscf-1.2.0/gpu4pyscf/scf/dispersion.py +22 -0
  67. gpu4pyscf-1.2.0/gpu4pyscf/scf/hf.py +538 -0
  68. gpu4pyscf-1.2.0/gpu4pyscf/scf/hf_symm.py +22 -0
  69. gpu4pyscf-1.2.0/gpu4pyscf/scf/int2c2e.py +141 -0
  70. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/int4c2e.py +310 -4
  71. gpu4pyscf-1.2.0/gpu4pyscf/scf/j_engine.py +230 -0
  72. gpu4pyscf-1.2.0/gpu4pyscf/scf/jk.py +777 -0
  73. gpu4pyscf-1.2.0/gpu4pyscf/scf/rohf.py +229 -0
  74. gpu4pyscf-1.2.0/gpu4pyscf/scf/soscf.py +704 -0
  75. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/uhf.py +46 -16
  76. gpu4pyscf-1.2.0/gpu4pyscf/scf/uhf_symm.py +20 -0
  77. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/__init__.py +2 -2
  78. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/_attach_solvent.py +1 -0
  79. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/pcm.py +48 -48
  80. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/smd.py +5 -10
  81. gpu4pyscf-1.2.0/gpu4pyscf/solvent/hessian/__init__.py +0 -0
  82. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/pcm.py +1 -0
  83. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/pcm.py +4 -1
  84. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/smd.py +6 -2
  85. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/__init__.py +22 -0
  86. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/_lr_eig.py +994 -0
  87. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/_uhf_resp_sf.py +217 -0
  88. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/rhf.py +368 -0
  89. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/rks.py +156 -0
  90. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/uhf.py +777 -0
  91. gpu4pyscf-1.2.0/gpu4pyscf/tdscf/uks.py +197 -0
  92. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0/gpu4pyscf.egg-info}/PKG-INFO +1 -1
  93. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/SOURCES.txt +34 -1
  94. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/requires.txt +1 -1
  95. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/setup.py +1 -1
  96. gpu4pyscf-1.0.2/gpu4pyscf/__config__.py +0 -38
  97. gpu4pyscf-1.0.2/gpu4pyscf/df/df.py +0 -299
  98. gpu4pyscf-1.0.2/gpu4pyscf/grad/__init__.py +0 -6
  99. gpu4pyscf-1.0.2/gpu4pyscf/grad/rhf.py +0 -702
  100. gpu4pyscf-1.0.2/gpu4pyscf/grad/uhf.py +0 -353
  101. gpu4pyscf-1.0.2/gpu4pyscf/hessian/rhf.py +0 -695
  102. gpu4pyscf-1.0.2/gpu4pyscf/hessian/rks.py +0 -701
  103. gpu4pyscf-1.0.2/gpu4pyscf/qmmm/chelpg.py +0 -271
  104. gpu4pyscf-1.0.2/gpu4pyscf/scf/dispersion.py +0 -32
  105. gpu4pyscf-1.0.2/gpu4pyscf/scf/hf.py +0 -1063
  106. gpu4pyscf-1.0.2/gpu4pyscf/scf/rohf.py +0 -76
  107. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/LICENSE +0 -0
  108. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/MANIFEST.in +0 -0
  109. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/cc/__init__.py +0 -0
  110. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/__init__.py +0 -0
  111. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/cderi.py +0 -0
  112. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/__init__.py +0 -0
  113. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/__init__.py +0 -0
  114. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/gen_grid.py +0 -0
  115. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/libxc_structs.py +0 -0
  116. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/radi.py +0 -0
  117. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/xc_alias.py +0 -0
  118. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/__init__.py +0 -0
  119. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/opt_driver.py +0 -0
  120. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/fci/__init__.py +0 -0
  121. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/fci/direct_spin1.py +0 -0
  122. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/dispersion.py +0 -0
  123. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/gto/__init__.py +0 -0
  124. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/__init__.py +0 -0
  125. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/diis.py +0 -0
  126. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/utils.py +0 -0
  127. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/__init__.py +0 -0
  128. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/mp2.py +0 -0
  129. {gpu4pyscf-1.0.2/gpu4pyscf/solvent/hessian → gpu4pyscf-1.2.0/gpu4pyscf/pop}/__init__.py +0 -0
  130. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/__init__.py +0 -0
  131. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/ir.py +0 -0
  132. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/polarizability.py +0 -0
  133. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/qmmm/__init__.py +0 -0
  134. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/__init__.py +0 -0
  135. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/cphf.py +0 -0
  136. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/diis.py +0 -0
  137. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/ghf.py +0 -0
  138. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/ucphf.py +0 -0
  139. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/__init__.py +0 -0
  140. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/smd_experiment.py +0 -0
  141. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/smd.py +0 -0
  142. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/smd_experiment.py +0 -0
  143. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/smd_experiment.py +0 -0
  144. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/dependency_links.txt +0 -0
  145. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/top_level.txt +0 -0
  146. {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: gpu4pyscf
3
- Version: 1.0.2
3
+ Version: 1.2.0
4
4
  Summary: GPU extensions for PySCF
5
5
  Home-page: UNKNOWN
6
6
  Author: Qiming Sun
@@ -43,6 +43,7 @@ Features
43
43
  - Density fitting scheme and direct SCF scheme;
44
44
  - SCF, analytical Gradient, and analytical Hessian calculations for Hartree-Fock and DFT;
45
45
  - LDA, GGA, mGGA, hybrid, and range-separated functionals via [libXC](https://gitlab.com/libxc/libxc/-/tree/master/);
46
+ - Spin-conserved and spin-flip TDA and TDDFT for excitated states
46
47
  - Geometry optimization and transition state search via [geomeTRIC](https://geometric.readthedocs.io/en/latest/);
47
48
  - Dispersion corrections via [DFTD3](https://github.com/dftd3/simple-dftd3) and [DFTD4](https://github.com/dftd4/dftd4);
48
49
  - Nonlocal functional correction (vv10) for SCF and gradient;
@@ -50,7 +51,10 @@ Features
50
51
  - PCM models, SMD model, their analytical gradients, and semi-analytical Hessian matrix;
51
52
  - Unrestricted Hartree-Fock and Unrestricted DFT, gradient, and Hessian;
52
53
  - MP2/DF-MP2 and CCSD (experimental);
53
- - Polarizability, IR, and NMR shielding
54
+ - Polarizability, IR, and NMR shielding (experimental);
55
+ - QM/MM with PBC;
56
+ - CHELPG, ESP, and RESP atomic charge;
57
+ - Multi-GPU for density fitting (experimental)
54
58
 
55
59
  Limitations
56
60
  --------
@@ -58,7 +62,6 @@ Limitations
58
62
  - Atomic basis up to g orbitals;
59
63
  - Auxiliary basis up to i orbitals;
60
64
  - Density fitting scheme up to ~168 atoms with def2-tzvpd basis, bounded by CPU memory;
61
- - Hessian is unavailable for Direct SCF yet;
62
65
  - meta-GGA without density laplacian;
63
66
  - Double hybrid functionals are not supported;
64
67
 
@@ -132,22 +135,22 @@ References
132
135
  ---------
133
136
  ```
134
137
  @misc{li2024introducting,
135
- title={Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework},
138
+ title={Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework},
136
139
  author={Rui Li and Qiming Sun and Xing Zhang and Garnet Kin-Lic Chan},
137
140
  year={2024},
138
141
  eprint={2407.09700},
139
142
  archivePrefix={arXiv},
140
143
  primaryClass={physics.comp-ph},
141
- url={https://arxiv.org/abs/2407.09700},
144
+ url={https://arxiv.org/abs/2407.09700},
142
145
  }
143
146
 
144
147
  @misc{wu2024enhancing,
145
- title={Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework},
148
+ title={Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework},
146
149
  author={Xiaojie Wu and Qiming Sun and Zhichen Pu and Tianze Zheng and Wenzhi Ma and Wen Yan and Xia Yu and Zhengxiao Wu and Mian Huo and Xiang Li and Weiluo Ren and Sheng Gong and Yumin Zhang and Weihao Gao},
147
150
  year={2024},
148
151
  eprint={2404.09452},
149
152
  archivePrefix={arXiv},
150
153
  primaryClass={physics.comp-ph},
151
- url={https://arxiv.org/abs/2404.09452},
154
+ url={https://arxiv.org/abs/2404.09452},
152
155
  }
153
156
  ```
@@ -0,0 +1,34 @@
1
+ import cupy
2
+
3
+ _num_devices = cupy.cuda.runtime.getDeviceCount()
4
+
5
+ # TODO: switch to non_blocking stream (currently blocked by libxc)
6
+ _streams = [None] * _num_devices
7
+ for device_id in range(_num_devices):
8
+ with cupy.cuda.Device(device_id):
9
+ _streams[device_id] = cupy.cuda.stream.Stream(non_blocking=False)
10
+
11
+ props = cupy.cuda.runtime.getDeviceProperties(0)
12
+ GB = 1024*1024*1024
13
+ min_ao_blksize = 128
14
+ min_grid_blksize = 128*128
15
+ ao_aligned = 32
16
+ grid_aligned = 256
17
+
18
+ # Use smaller blksize for old gaming GPUs
19
+ if props['totalGlobalMem'] < 16 * GB:
20
+ min_ao_blksize = 64
21
+ min_grid_blksize = 64*64
22
+
23
+ # Use 90% of the global memory for CuPy memory pool
24
+ mem_fraction = 0.9
25
+ cupy.get_default_memory_pool().set_limit(fraction=mem_fraction)
26
+
27
+ # Check P2P data transfer is available
28
+ _p2p_access = True
29
+ if _num_devices > 1:
30
+ for src in range(_num_devices):
31
+ for dst in range(_num_devices):
32
+ if src != dst:
33
+ can_access_peer = cupy.cuda.runtime.deviceCanAccessPeer(src, dst)
34
+ _p2p_access &= can_access_peer
@@ -1,7 +1,7 @@
1
- from . import lib, grad, hessian, solvent, scf, dft
2
-
3
- __version__ = '1.0.2'
1
+ __version__ = '1.2.0'
4
2
 
5
3
  # monkey patch libxc reference due to a bug in nvcc
6
4
  from pyscf.dft import libxc
7
5
  libxc.__reference__ = 'unable to decode the reference due to https://github.com/NVIDIA/cuda-python/issues/29'
6
+
7
+ from . import lib, grad, hessian, solvent, scf, dft
@@ -31,7 +31,7 @@ from pyscf.ao2mo import _ao2mo
31
31
  from pyscf.cc import ccsd
32
32
  from pyscf.cc import _ccsd
33
33
  from pyscf import __config__
34
- from gpu4pyscf.scf import hf as gpu_hf
34
+ from gpu4pyscf.scf import int4c2e
35
35
  from gpu4pyscf.lib.cupy_helper import load_library
36
36
  from gpu4pyscf.lib import logger
37
37
 
@@ -171,7 +171,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
171
171
  ((mem_avail-Ht2_mem)*.5/8/nao_cart**2)**.5)))
172
172
  logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
173
173
 
174
- vhfopt = gpu_hf._VHFOpt(mycc.mol, 'int2e')
174
+ vhfopt = int4c2e._VHFOpt(mycc.mol, 'int2e')
175
175
  vhfopt.build(group_size=blksize, diag_block_with_triu=True)
176
176
  mol = vhfopt.mol
177
177
 
@@ -219,7 +219,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
219
219
  log_qs = vhfopt.log_qs
220
220
  cp_idx, cp_jdx = np.tril_indices(len(vhfopt.uniq_l_ctr))
221
221
 
222
- if vhfopt.uniq_l_ctr[:,0].max() <= gpu_hf.LMAX_ON_GPU:
222
+ if vhfopt.uniq_l_ctr[:,0].max() <= int4c2e.LMAX_ON_GPU:
223
223
  # Computing ERIs on GPU
224
224
  idx, idy = cupy.tril_indices(nao)
225
225
  #eribuf = cupy.empty(blksize**2*nao**2)
@@ -279,7 +279,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
279
279
  cpj = cp_jdx[cp_ij_id]
280
280
  li = vhfopt.uniq_l_ctr[cpi,0]
281
281
  lj = vhfopt.uniq_l_ctr[cpj,0]
282
- if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
282
+ if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
283
283
  continue
284
284
 
285
285
  ish0 = l_ctr_offsets[cpi]
@@ -375,7 +375,7 @@ def _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id):
375
375
  cpl = cp_jdx[cp_kl_id]
376
376
  lk = vhfopt.uniq_l_ctr[cpk,0]
377
377
  ll = vhfopt.uniq_l_ctr[cpl,0]
378
- if lk > gpu_hf.LMAX_ON_GPU or ll > gpu_hf.LMAX_ON_GPU:
378
+ if lk > int4c2e.LMAX_ON_GPU or ll > int4c2e.LMAX_ON_GPU:
379
379
  raise NotImplementedError
380
380
 
381
381
  stream = cupy.cuda.get_current_stream()
@@ -394,7 +394,7 @@ def _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id):
394
394
  cpj = cp_jdx[cp_ij_id]
395
395
  li = vhfopt.uniq_l_ctr[cpi,0]
396
396
  lj = vhfopt.uniq_l_ctr[cpj,0]
397
- if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
397
+ if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
398
398
  continue
399
399
 
400
400
  t0 = time.perf_counter()
@@ -446,7 +446,7 @@ def _make_eris_incore(mycc, mo_coeff=None):
446
446
  (mem_avail*.5/8/nao_cart**2)**.5)))
447
447
  logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
448
448
 
449
- vhfopt = gpu_hf._VHFOpt(mycc.mol, 'int2e')
449
+ vhfopt = int4c2e._VHFOpt(mycc.mol, 'int2e')
450
450
  vhfopt.build(group_size=blksize, diag_block_with_triu=True)
451
451
  mol = vhfopt.mol
452
452
  mo = vhfopt.coeff.dot(mo_coeff)
@@ -470,7 +470,7 @@ def _make_eris_incore(mycc, mo_coeff=None):
470
470
  cpj = cp_jdx[cp_ij_id]
471
471
  li = vhfopt.uniq_l_ctr[cpi,0]
472
472
  lj = vhfopt.uniq_l_ctr[cpj,0]
473
- if li > gpu_hf.LMAX_ON_GPU or lj > gpu_hf.LMAX_ON_GPU or log_q_ij.size == 0:
473
+ if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
474
474
  continue
475
475
 
476
476
  ish0 = l_ctr_offsets[cpi]
@@ -0,0 +1,353 @@
1
+ # Copyright 2023 The GPU4PySCF Authors. All Rights Reserved.
2
+ #
3
+ # This program is free software: you can redistribute it and/or modify
4
+ # it under the terms of the GNU General Public License as published by
5
+ # the Free Software Foundation, either version 3 of the License, or
6
+ # (at your option) any later version.
7
+ #
8
+ # This program is distributed in the hope that it will be useful,
9
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
10
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11
+ # GNU General Public License for more details.
12
+ #
13
+ # You should have received a copy of the GNU General Public License
14
+ # along with this program. If not, see <http://www.gnu.org/licenses/>.
15
+
16
+
17
+ import copy
18
+ import cupy
19
+ from concurrent.futures import ThreadPoolExecutor
20
+ import numpy as np
21
+ from cupyx.scipy.linalg import solve_triangular
22
+ from pyscf import lib
23
+ from pyscf.df import df, addons, incore
24
+ from gpu4pyscf.lib.cupy_helper import cholesky, tag_array, get_avail_mem, cart2sph, p2p_transfer
25
+ from gpu4pyscf.df import int3c2e, df_jk
26
+ from gpu4pyscf.lib import logger
27
+ from gpu4pyscf import __config__
28
+ from gpu4pyscf.__config__ import _streams, _num_devices
29
+
30
+ MIN_BLK_SIZE = getattr(__config__, 'min_ao_blksize', 128)
31
+ ALIGNED = getattr(__config__, 'ao_aligned', 32)
32
+ GB = 1024*1024*1024
33
+
34
+ LINEAR_DEP_THR = incore.LINEAR_DEP_THR
35
+ GROUP_SIZE = 256
36
+
37
+ class DF(lib.StreamObject):
38
+ from gpu4pyscf.lib.utils import to_gpu, device
39
+
40
+ _keys = {'intopt', 'mol', 'auxmol', 'use_gpu_memory'}
41
+
42
+ def __init__(self, mol, auxbasis=None):
43
+ self.mol = mol
44
+ self.stdout = mol.stdout
45
+ self.verbose = mol.verbose
46
+ self.max_memory = mol.max_memory
47
+ self.use_gpu_memory = True
48
+ self._auxbasis = auxbasis
49
+
50
+ self.auxmol = None
51
+ self.intopt = None
52
+ self.nao = None
53
+ self.naux = None
54
+ self.cd_low = None
55
+ self._cderi = None
56
+ self._rsh_df = {}
57
+
58
+ @property
59
+ def auxbasis(self):
60
+ return self._auxbasis
61
+ @auxbasis.setter
62
+ def auxbasis(self, x):
63
+ if self._auxbasis != x:
64
+ self.reset()
65
+ self._auxbasis = x
66
+
67
+ def to_cpu(self):
68
+ from gpu4pyscf.lib.utils import to_cpu
69
+ obj = to_cpu(self)
70
+ return obj.reset()
71
+
72
+ def build(self, direct_scf_tol=1e-14, omega=None):
73
+ mol = self.mol
74
+ auxmol = self.auxmol
75
+ self.nao = mol.nao
76
+ log = logger.new_logger(mol, mol.verbose)
77
+ t0 = log.init_timer()
78
+ if auxmol is None:
79
+ self.auxmol = auxmol = addons.make_auxmol(mol, self.auxbasis)
80
+
81
+ if omega and omega > 1e-10:
82
+ with auxmol.with_range_coulomb(omega):
83
+ j2c_cpu = auxmol.intor('int2c2e', hermi=1)
84
+ else:
85
+ j2c_cpu = auxmol.intor('int2c2e', hermi=1)
86
+ j2c = cupy.asarray(j2c_cpu, order='C')
87
+ t0 = log.timer_debug1('2c2e', *t0)
88
+ intopt = int3c2e.VHFOpt(mol, auxmol, 'int2e')
89
+ intopt.build(direct_scf_tol, diag_block_with_triu=False, aosym=True,
90
+ group_size=GROUP_SIZE, group_size_aux=GROUP_SIZE)
91
+ log.timer_debug1('prepare intopt', *t0)
92
+ self.j2c = j2c.copy()
93
+
94
+ j2c = intopt.sort_orbitals(j2c, aux_axis=[0,1])
95
+ try:
96
+ self.cd_low = cholesky(j2c)
97
+ self.cd_low = tag_array(self.cd_low, tag='cd')
98
+ except Exception:
99
+ w, v = cupy.linalg.eigh(j2c)
100
+ idx = w > LINEAR_DEP_THR
101
+ self.cd_low = (v[:,idx] / cupy.sqrt(w[idx]))
102
+ self.cd_low = tag_array(self.cd_low, tag='eig')
103
+
104
+ v = w = None
105
+ naux = self.naux = self.cd_low.shape[1]
106
+ log.debug('size of aux basis %d', naux)
107
+
108
+ self._cderi = cholesky_eri_gpu(intopt, mol, auxmol, self.cd_low,
109
+ omega=omega, use_gpu_memory=self.use_gpu_memory)
110
+ log.timer_debug1('cholesky_eri', *t0)
111
+ self.intopt = intopt
112
+ return self
113
+
114
+ def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
115
+ direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
116
+ omega=None):
117
+ if omega is None:
118
+ return df_jk.get_jk(self, dm, hermi, with_j, with_k, direct_scf_tol)
119
+ assert omega >= 0.0
120
+
121
+ # A temporary treatment for RSH-DF integrals
122
+ key = '%.6f' % omega
123
+ if key in self._rsh_df:
124
+ rsh_df = self._rsh_df[key]
125
+ else:
126
+ rsh_df = self._rsh_df[key] = self.copy().reset()
127
+ logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
128
+
129
+ return df_jk.get_jk(rsh_df, dm, hermi, with_j, with_k, direct_scf_tol, omega=omega)
130
+
131
+ def get_blksize(self, extra=0, nao=None):
132
+ '''
133
+ extra for pre-calculated space for other variables
134
+ '''
135
+ if nao is None: nao = self.nao
136
+ mem_avail = get_avail_mem()
137
+ blksize = int(mem_avail*0.2/8/(nao*nao + extra) / ALIGNED) * ALIGNED
138
+ blksize = min(blksize, MIN_BLK_SIZE)
139
+ log = logger.new_logger(self.mol, self.mol.verbose)
140
+ device_id = cupy.cuda.Device().id
141
+ log.debug(f"{mem_avail/1e9:.3f} GB memory available on Device {device_id}, block size = {blksize}")
142
+ if blksize < ALIGNED:
143
+ raise RuntimeError("Not enough GPU memory")
144
+ return blksize
145
+
146
+ def loop(self, blksize=None, unpack=True):
147
+ ''' loop over cderi for the current device
148
+ and unpack the CDERI in (Lij) format
149
+ '''
150
+ device_id = cupy.cuda.Device().id
151
+ cderi_sparse = self._cderi[device_id]
152
+ if blksize is None:
153
+ blksize = self.get_blksize()
154
+ nao = self.nao
155
+ naux_slice = cderi_sparse.shape[0]
156
+ rows = self.intopt.cderi_row
157
+ cols = self.intopt.cderi_col
158
+ buf_prefetch = None
159
+ buf_cderi = cupy.zeros([blksize,nao,nao])
160
+ for p0, p1 in lib.prange(0, naux_slice, blksize):
161
+ p2 = min(naux_slice, p1+blksize)
162
+ if isinstance(cderi_sparse, cupy.ndarray):
163
+ buf = cderi_sparse[p0:p1,:]
164
+ if isinstance(cderi_sparse, np.ndarray):
165
+ # first block
166
+ if buf_prefetch is None:
167
+ buf = cupy.asarray(cderi_sparse[p0:p1,:])
168
+ buf_prefetch = cupy.empty([p2-p1,cderi_sparse.shape[1]])
169
+ if isinstance(cderi_sparse, np.ndarray) and p1 < p2:
170
+ buf_prefetch.set(cderi_sparse[p1:p2,:])
171
+ if unpack:
172
+ buf_cderi[:p1-p0,rows,cols] = buf
173
+ buf_cderi[:p1-p0,cols,rows] = buf
174
+ buf2 = buf_cderi[:p1-p0]
175
+ else:
176
+ buf2 = None
177
+ yield buf2, buf.T
178
+ if isinstance(cderi_sparse, np.ndarray):
179
+ cupy.cuda.Device().synchronize()
180
+
181
+ if buf_prefetch is not None:
182
+ buf = buf_prefetch
183
+
184
+ def reset(self, mol=None):
185
+ '''Reset mol and clean up relevant attributes for scanner mode'''
186
+ if mol is not None:
187
+ self.mol = mol
188
+ self.auxmol = None
189
+ self._cderi = None
190
+ self._vjopt = None
191
+ self._rsh_df = {}
192
+ self.intopt = None
193
+ self.nao = None
194
+ self.naux = None
195
+ self.cd_low = None
196
+ return self
197
+
198
+ get_ao_eri = get_eri = NotImplemented
199
+ get_mo_eri = ao2mo = NotImplemented
200
+
201
+ def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
202
+ omega=None, sr_only=False, use_gpu_memory=True):
203
+ '''
204
+ Returns:
205
+ 2D array of (naux,nao*(nao+1)/2) in C-contiguous
206
+ '''
207
+ naux = cd_low.shape[1]
208
+ npairs = len(intopt.cderi_row)
209
+ log = logger.new_logger(mol, mol.verbose)
210
+
211
+ # Available memory on Device 0.
212
+ avail_mem = get_avail_mem()
213
+
214
+ if use_gpu_memory:
215
+ # CDERI will be equally distributed to the devices
216
+ # Other devices usually have more memory available than Device 0
217
+ # CDERI will use up to 40% of the available memory
218
+ use_gpu_memory = naux * npairs * 8 < 0.4 * avail_mem * _num_devices
219
+
220
+ if use_gpu_memory:
221
+ log.debug("Saving CDERI on GPU")
222
+ else:
223
+ log.debug("Saving CDERI on CPU")
224
+
225
+ _cderi = {}
226
+ blksize = (naux + _num_devices - 1) // _num_devices
227
+ for device_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
228
+ if use_gpu_memory:
229
+ with cupy.cuda.Device(device_id), _streams[device_id]:
230
+ _cderi[device_id] = cupy.empty([p1-p0, npairs])
231
+ log.debug(f"CDERI size {_cderi[device_id].nbytes/GB:.3f} on Device {device_id}")
232
+ else:
233
+ mem = cupy.cuda.alloc_pinned_memory((p1-p0) * npairs * 8)
234
+ cderi_blk = np.ndarray([p1-p0, npairs], dtype=np.float64, order='C', buffer=mem)
235
+ _cderi[device_id] = cderi_blk
236
+
237
+ npairs_per_ctr = [len(intopt.ao_pairs_row[cp_ij_id]) for cp_ij_id in range(len(intopt.log_qs))]
238
+
239
+ npairs_per_ctr = np.array(npairs_per_ctr)
240
+ total_task_list = np.argsort(npairs_per_ctr)
241
+ task_list_per_device = []
242
+ for device_id in range(_num_devices):
243
+ task_list_per_device.append(total_task_list[device_id::_num_devices])
244
+
245
+ cd_low_f = cupy.array(cd_low, order='F', copy=False)
246
+ cd_low_f = tag_array(cd_low_f, tag=cd_low.tag)
247
+
248
+ cupy.cuda.get_current_stream().synchronize()
249
+ futures = []
250
+ with ThreadPoolExecutor(max_workers=_num_devices) as executor:
251
+ for device_id in range(_num_devices):
252
+ task_list = task_list_per_device[device_id]
253
+ future = executor.submit(_cderi_task, intopt, cd_low_f, task_list, _cderi,
254
+ omega=omega, sr_only=sr_only, device_id=device_id)
255
+ futures.append(future)
256
+
257
+ for future in futures:
258
+ future.result()
259
+
260
+ if not use_gpu_memory:
261
+ cupy.cuda.Device().synchronize()
262
+
263
+ return _cderi
264
+
265
+ def _cderi_task(intopt, cd_low, task_list, _cderi, omega=None, sr_only=False, device_id=0):
266
+ ''' Execute CDERI tasks on one device
267
+ '''
268
+ nq = len(intopt.log_qs)
269
+ mol = intopt.mol
270
+ naux = cd_low.shape[1]
271
+ naoaux = cd_low.shape[0]
272
+ npairs = [len(intopt.ao_pairs_row[cp_ij]) for cp_ij in range(len(intopt.log_qs))]
273
+ pairs_loc = np.append(0, np.cumsum(npairs))
274
+ blksize = (naux + _num_devices - 1) // _num_devices
275
+ with cupy.cuda.Device(device_id), _streams[device_id]:
276
+ log = logger.new_logger(mol, mol.verbose)
277
+ t1 = log.init_timer()
278
+ cd_low_tag = cd_low.tag
279
+ cd_low = cupy.asarray(cd_low)
280
+
281
+ cart_ao_loc = intopt.cart_ao_loc
282
+ aux_ao_loc = intopt.aux_ao_loc
283
+ ao_loc = intopt.ao_loc
284
+ for cp_ij_id in task_list:
285
+ cpi = intopt.cp_idx[cp_ij_id]
286
+ cpj = intopt.cp_jdx[cp_ij_id]
287
+ li = intopt.angular[cpi]
288
+ lj = intopt.angular[cpj]
289
+ i0, i1 = cart_ao_loc[cpi], cart_ao_loc[cpi+1]
290
+ j0, j1 = cart_ao_loc[cpj], cart_ao_loc[cpj+1]
291
+ ni = i1 - i0
292
+ nj = j1 - j0
293
+ if sr_only:
294
+ # TODO: in-place implementation or short-range kernel
295
+ ints_slices = cupy.zeros([naoaux, nj, ni], order='C')
296
+ for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
297
+ k0 = aux_ao_loc[cp_kl_id]
298
+ k1 = aux_ao_loc[cp_kl_id+1]
299
+ int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1])
300
+ if omega is not None:
301
+ ints_slices_lr = cupy.zeros([naoaux, nj, ni], order='C')
302
+ for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
303
+ k0 = aux_ao_loc[cp_kl_id]
304
+ k1 = aux_ao_loc[cp_kl_id+1]
305
+ int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1], omega=omega)
306
+ ints_slices -= ints_slices_lr
307
+ else:
308
+ # Initialization is required due to cutensor operations later
309
+ ints_slices = cupy.zeros([naoaux, nj, ni], order='C')
310
+ for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
311
+ k0 = aux_ao_loc[cp_kl_id]
312
+ k1 = aux_ao_loc[cp_kl_id+1]
313
+
314
+ int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1], omega=omega)
315
+ if lj>1 and not mol.cart: ints_slices = cart2sph(ints_slices, axis=1, ang=lj)
316
+ if li>1 and not mol.cart: ints_slices = cart2sph(ints_slices, axis=2, ang=li)
317
+
318
+ i0, i1 = ao_loc[cpi], ao_loc[cpi+1]
319
+ j0, j1 = ao_loc[cpj], ao_loc[cpj+1]
320
+
321
+ row = intopt.ao_pairs_row[cp_ij_id] - i0
322
+ col = intopt.ao_pairs_col[cp_ij_id] - j0
323
+
324
+ ints_slices_f= cupy.empty([naoaux,len(row)], order='F')
325
+ ints_slices_f[:] = ints_slices[:,col,row]
326
+ ints_slices = None
327
+ if cd_low_tag == 'eig':
328
+ cderi_block = cupy.dot(cd_low.T, ints_slices_f)
329
+ ints_slices = None
330
+ elif cd_low_tag == 'cd':
331
+ cderi_block = solve_triangular(cd_low, ints_slices_f, lower=True, overwrite_b=True)
332
+ else:
333
+ RuntimeError('Tag is not found in lower triangular matrix.')
334
+ t1 = log.timer_debug1(f'solve {cp_ij_id} / {nq} on Device {device_id}', *t1)
335
+
336
+ # TODO:
337
+ # 1) async data transfer
338
+ # 2) auxiliary basis in the last dimension
339
+
340
+ # if CDERI is saved on CPU
341
+ ij0 = pairs_loc[cp_ij_id]
342
+ ij1 = pairs_loc[cp_ij_id+1]
343
+ if isinstance(_cderi[0], np.ndarray):
344
+ for slice_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
345
+ for i in range(p0,p1):
346
+ cderi_block[i].get(out=_cderi[slice_id][i-p0,ij0:ij1])
347
+ else:
348
+ # Copy data to other Devices
349
+ for slice_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
350
+ #_cderi[slice_id][:,ij0:ij1] = cderi_block[p0:p1]
351
+ p2p_transfer(_cderi[slice_id][:,ij0:ij1], cderi_block[p0:p1])
352
+ t1 = log.timer_debug1(f'transfer data for {cp_ij_id} / {nq} on Device {device_id}', *t1)
353
+ return