pyfock.XC

18def lda_x(rho):
19    """
20    Compute the LDA exchange energy and potential using the Slater exchange functional (spin-unpolarized).
21
22    This is the standard Local Density Approximation (LDA) exchange functional corresponding to 
23    `LDA_X` with ID 1 in the LibXC functional library.
24
25    Parameters
26    ----------
27    rho : ndarray
28        Electron density array (assumed to be spin-paired / spin-unpolarized).
29
30    Returns
31    -------
32    ex : ndarray
33        Exchange energy density at each point.
34
35    vx : ndarray
36        Exchange potential (functional derivative of exchange energy with respect to density).
37
38    Notes
39    -----
40    This implementation is based on:
41        - Phys. Rev. 81, 385 (1951) — the original Slater exchange.
42    Code adapted from the `eminus` repository:
43        https://github.com/wangenau/eminus/blob/main/eminus/xc/lda_x.py
44        Licensed under the Apache License, Version 2.0.
45
46    """
47    
48    pi34 = (3 / (4 * np.pi))**(1 / 3)
49    f = -3 / 4 * (3 / (2 * np.pi))**(2 / 3)
50    rs = pi34 * np.power(rho, -1 / 3)
51    ex = f / rs
52    vx = 4 / 3 * ex
53    # return {'zk':ex, 'vrho':vx}
54    return ex, vx

56@fuse(kernel_name='lda_x_cupy')
57def lda_x_cupy(rho):
58    """
59    GPU-accelerated version of the LDA exchange functional using CuPy.
60
61    This is the same as `lda_x` but leverages CuPy for GPU computation. It corresponds to 
62    `LDA_X` with ID 1 in LibXC for spin-unpolarized electron density.
63
64    Parameters
65    ----------
66    rho : cupy.ndarray
67        Electron density array (spin-unpolarized), on GPU.
68
69    Returns
70    -------
71    ex : cupy.ndarray
72        Exchange energy density at each point.
73
74    vx : cupy.ndarray
75        Exchange potential (functional derivative of exchange energy with respect to density).
76
77    Notes
78    -----
79    Based on the Slater exchange:
80        - Phys. Rev. 81, 385 (1951)
81    Adapted from: https://github.com/wangenau/eminus/blob/main/eminus/xc/lda_x.py
82    Licensed under the Apache License, Version 2.0.
83
84    This version is fused using CuPy's `@fuse` decorator for better performance.
85
86    """
87    
88    pi34 = (3 / (4 * cp.pi))**(1 / 3)
89    f = -3 / 4 * (3 / (2 * cp.pi))**(2 / 3)
90    rs = pi34 * cp.power(rho, -1 / 3)
91    ex = f / rs
92    vx = 4 / 3 * ex
93    # return {'zk':ex, 'vrho':vx}
94    return ex, vx

70def lda_c_vwn(rho):
71    """
72    Wrapper function for `lda_c_vwn_` providing the LDA correlation energy and potential using VWN parametrization.
73
74    Parameters
75    ----------
76    rho : ndarray
77        Electron density array (spin-unpolarized).
78
79    Returns
80    -------
81    ec : ndarray
82        Correlation energy density.
83
84    vc : ndarray
85        Correlation potential.
86
87    Notes
88    -----
89    This function is equivalent to `lda_c_vwn_` but kept as a public-facing interface in PyFock.
90    """
91    ec, vc = lda_c_vwn_(rho)
92    return ec, vc

145def lda_c_vwn_cupy(rho):
146    """
147    Safe wrapper for the GPU-based VWN LDA correlation functional.
148
149    This function calls `lda_c_vwn_cupy_`, and replaces any NaNs in the result with zeros 
150    to ensure numerical stability in downstream calculations.
151
152    Parameters
153    ----------
154    rho : cupy.ndarray
155        Electron density array on GPU.
156
157    Returns
158    -------
159    ec : cupy.ndarray
160        Correlation energy density with NaNs replaced by 0.
161
162    vc : cupy.ndarray
163        Correlation potential with NaNs replaced by 0.
164
165    Notes
166    -----
167    This is a numerically safe version of `lda_c_vwn_cupy_` intended for production use.
168    """
169    ec, vc = lda_c_vwn_cupy_(rho)
170    vc[cp.isnan(vc)] = 0
171    ec[cp.isnan(ec)] = 0
172    return ec, vc

72def lda_c_pw(rho):
73    """
74    Numerically stable wrapper for the Perdew-Wang (PW92) LDA correlation functional.
75
76    This calls `lda_c_pw_`, then replaces any NaNs in the result with zeros to ensure 
77    downstream stability in SCF or post-processing steps.
78
79    Parameters
80    ----------
81    rho : ndarray
82        Electron density array (non-negative, spin-unpolarized).
83
84    Returns
85    -------
86    ec : ndarray
87        Correlation energy density with NaNs replaced by 0.
88
89    vc : ndarray
90        Correlation potential with NaNs replaced by 0.
91
92    Notes
93    -----
94    This is the CPU-based version. Use `lda_c_pw_cupy` for GPU acceleration.
95    """
96    ec, vc = lda_c_pw_(rho)
97    vc[np.isnan(vc)] = 0
98    ec[np.isnan(ec)] = 0
99    return ec, vc

158def lda_c_pw_cupy(rho):
159    """
160    Numerically stable GPU wrapper for the Perdew-Wang LDA correlation functional.
161
162    This function calls `lda_c_pw_cupy_` and replaces any NaNs in the resulting arrays
163    with zeros to ensure stability during molecular dynamics or SCF procedures.
164
165    Parameters
166    ----------
167    rho : cupy.ndarray
168        Electron density array on GPU (non-negative, spin-unpolarized).
169
170    Returns
171    -------
172    ec : cupy.ndarray
173        Correlation energy density with NaNs replaced by 0.
174
175    vc : cupy.ndarray
176        Correlation potential with NaNs replaced by 0.
177
178    Notes
179    -----
180    Use this function for production GPU workflows where numerical robustness is critical.
181    """
182    ec, vc = lda_c_pw_cupy_(rho)
183    vc[cp.isnan(vc)] = 0
184    ec[cp.isnan(ec)] = 0
185    return ec, vc

47def lda_c_pw_mod(rho):
48    ec, vc = lda_c_pw_mod_(rho)
49    vc[np.isnan(vc)] = 0
50    ec[np.isnan(ec)] = 0
51    return ec, vc

83def lda_c_pw_mod_cupy(rho):
84    ec, vc = lda_c_pw_mod_cupy_(rho)
85    vc[cp.isnan(vc)] = 0
86    ec[cp.isnan(ec)] = 0
87    return ec, vc

19def gga_x_pbe(rho, sigma):
20    """
21    Compute the restricted PBE (Perdew–Burke–Ernzerhof) exchange energy density and potential
22    using NumPy arrays for electron density and its gradient.
23
24    Adapted from
25    ----------
26    Eminus project:
27    https://github.com/wangenau/eminus/blob/main/eminus/xc/gga_x_pbe.py
28    Licensed under the Apache License, Version 2.0.
29
30    Reference
31    ----------
32    J. P. Perdew, K. Burke, and M. Ernzerhof,
33    "Generalized Gradient Approximation Made Simple",
34    Phys. Rev. Lett. 77, 3865 (1996).
35    https://doi.org/10.1103/PhysRevLett.77.3865
36
37    Parameters
38    ----------
39    rho : ndarray
40        Electron density array.
41    sigma : ndarray
42        Gradient of the electron density, defined as ∇ρ·∇ρ.
43
44    Returns
45    -------
46    ex : ndarray
47        Exchange energy density.
48    vx : ndarray
49        Functional derivative of the exchange energy with respect to density.
50    vsigma : ndarray
51        Functional derivative of the exchange energy with respect to the density gradient term σ.
52    """
53
54    mu = 0.2195149727645171 # Functional parameter
55
56    # rho_cutoff = 1e-12  # define rho_cutoff constant
57    rho = np.maximum(rho, 1e-12)
58
59    ex, vx = lda_x(rho)
60    gex, gvx, vsigmax = pbe_x_temp(rho, sigma)
61
62    ex += gex/rho
63    vx += gvx
64    vsigma = 0.5*vsigmax
65
66    vsigma[np.isnan(vsigma)] = 0
67    vx[np.isnan(vx)] = 0
68    ex[np.isnan(ex)] = 0
69
70    return ex, vx, vsigma

200def gga_x_pbe_cupy(rho, sigma):
201    """
202    Compute the restricted PBE exchange energy density and potential
203    using CuPy for GPU acceleration.
204
205    Adapted from
206    ----------
207    Eminus project:
208    https://github.com/wangenau/eminus/blob/main/eminus/xc/gga_x_pbe.py
209    Licensed under the Apache License, Version 2.0.
210
211    Reference
212    ----------
213    J. P. Perdew, K. Burke, and M. Ernzerhof,
214    "Generalized Gradient Approximation Made Simple",
215    Phys. Rev. Lett. 77, 3865 (1996).
216    https://doi.org/10.1103/PhysRevLett.77.3865
217
218    Parameters
219    ----------
220    rho : cp.ndarray
221        Electron density array (CuPy).
222    sigma : cp.ndarray
223        Gradient of the electron density, defined as ∇ρ·∇ρ (CuPy).
224
225    Returns
226    -------
227    ex : cp.ndarray
228        Exchange energy density.
229    vx : cp.ndarray
230        Functional derivative of the exchange energy with respect to density.
231    vsigma : cp.ndarray
232        Functional derivative of the exchange energy with respect to the density gradient term σ.
233    """
234
235    mu = 0.2195149727645171 # Functional parameter
236
237    # rho_cutoff = 1e-12  # define rho_cutoff constant
238    rho = cp.maximum(rho, 1e-12)
239
240    ex, vx = lda_x(rho)
241    gex, gvx, vsigmax = pbe_x_temp_cupy(rho, sigma)
242
243    ex += gex/rho
244    vx += gvx
245    vsigma = 0.5*vsigmax
246
247    vsigma[cp.isnan(vsigma)] = 0
248    vx[cp.isnan(vx)] = 0
249    ex[cp.isnan(ex)] = 0
250
251    return ex, vx, vsigma

64def gga_c_pbe(rho, sigma):
65    # Taken from: https://github.com/wangenau/eminus/blob/main/eminus/xc/gga_c_pbe.py
66    # Perdew-Burke-Ernzerhof parametrization of the correlation functional (spin-paired).
67    # Corresponds to the functional with the label GGA_C_PBE and ID 130 in Libxc.
68    # Reference: Phys. Rev. Lett. 78, 1396.
69    
70    ec, vc, vsigma = gga_c_pbe_(rho, sigma)
71    
72    vsigma[np.isnan(vsigma)] = 0
73    vc[np.isnan(vc)] = 0
74    ec[np.isnan(ec)] = 0
75
76    return ec, vc, vsigma

122def gga_c_pbe_cupy(rho, sigma):
123    # Taken from: https://github.com/wangenau/eminus/blob/main/eminus/xc/gga_c_pbe.py
124    # Perdew-Burke-Ernzerhof parametrization of the correlation functional (spin-paired).
125    # Corresponds to the functional with the label GGA_C_PBE and ID 130 in Libxc.
126    # Reference: Phys. Rev. Lett. 78, 1396.
127    
128    ec, vc, vsigma = gga_c_pbe_cupy_(rho, sigma)
129    
130    vsigma[cp.isnan(vsigma)] = 0
131    vc[cp.isnan(vc)] = 0
132    ec[cp.isnan(ec)] = 0
133
134    return ec, vc, vsigma

77def gga_x_b88(rho, sigma):
78    # Corresponds to 106 id in Libxc
79    ex = gga_x_b88_e(rho, sigma)
80    vx, vsigma = gga_x_b88_v(rho, sigma)
81    return ex, vx, vsigma

152def gga_x_b88_cupy(rho, sigma):
153    # Corresponds to 106 id in Libxc
154    ex = gga_x_b88_e_cupy(rho, sigma)
155    vx, vsigma = gga_x_b88_v_cupy(rho, sigma)
156    vsigma[cp.isnan(vsigma)] = 0
157    vx[cp.isnan(vx)] = 0
158    ex[cp.isnan(ex)] = 0
159    return ex, vx, vsigma

77def gga_c_lyp(rho, sigma):
78    # Corresponds to 131 id in Libxc
79    ec = gga_c_lyp_e(rho, sigma)
80    vc, vsigma = gga_c_lyp_v(rho, sigma)
81    return ec, vc, vsigma

145def gga_c_lyp_cupy(rho, sigma):
146    # Corresponds to 131 id in Libxc
147    ec = gga_c_lyp_e_cupy(rho, sigma)
148    vc, vsigma = gga_c_lyp_v_cupy(rho, sigma)
149    vsigma[cp.isnan(vsigma)] = 0
150    vc[cp.isnan(vc)] = 0
151    ec[cp.isnan(ec)] = 0
152    return ec, vc, vsigma

15def check_implemented(funcid):
16    """
17    Check if the given functional ID is implemented in PyFock.
18
19    Parameters
20    ----------
21    funcid : int
22        LibXC-style functional ID (e.g., 1 for LDA_X, 101 for GGA_X_PBE).
23
24    Raises
25    ------
26    SystemExit
27        If the functional ID is not supported by PyFock, prints an error and exits.
28
29    Notes
30    -----
31    Implemented functional IDs in PyFock:
32        - 1   : LDA_X
33        - 7   : LDA_C_VWN
34        - 12  : LDA_C_PW
35        - 13  : LDA_C_PW_MOD
36        - 101 : GGA_X_PBE
37        - 106 : GGA_X_B88
38        - 131 : GGA_C_LYP
39    For unsupported functionals, you can use LibXC directly.
40    """
41    if funcid not in [1, 7, 12, 13, 101, 106, 131]:
42        print('ERROR: The specified functional is not implemented in PyFock. You need to use LibXC to calculate the functional values.')
43        exit()

 45def func_compute(funcid, rho, sigma=None, use_gpu=True):
 46    """
 47    Compute exchange-correlation energy and potential using the specified functional.
 48
 49    Parameters
 50    ----------
 51    funcid : int
 52        Identifier for the XC functional. Matches LibXC IDs:
 53            - 1   : LDA_X
 54            - 7   : LDA_C_VWN
 55            - 12  : LDA_C_PW
 56            - 13  : LDA_C_PW_MOD
 57            - 101 : GGA_X_PBE
 58            - 106 : GGA_X_B88
 59            - 130 : GGA_C_PBE
 60            - 131 : GGA_C_LYP
 61
 62    rho : ndarray
 63        Electron density array. Should be shape-compatible with the expected input of the XC functional.
 64
 65    sigma : ndarray, optional
 66        Gradient of the density (∇ρ·∇ρ), required for GGA functionals.
 67
 68    use_gpu : bool, default=True
 69        Whether to use CuPy (GPU) versions of the functionals. If False, falls back to CPU (NumPy) versions.
 70
 71    Returns
 72    -------
 73    energy_density : ndarray
 74        The exchange-correlation energy density at each grid point.
 75
 76    potential : ndarray or tuple
 77        The exchange-correlation potential, and for GGA, possibly its derivative with respect to sigma.
 78
 79    Raises
 80    ------
 81    SystemExit
 82        If the functional is not implemented in PyFock.
 83
 84    Notes
 85    -----
 86    For unsupported functionals (i.e., other than those listed above), use LibXC.
 87    """
 88    if use_gpu:
 89        if funcid==1:
 90            return lda_x_cupy(rho)
 91        elif funcid==7:
 92            return lda_c_vwn_cupy(rho)
 93        elif funcid==12:
 94            return lda_c_pw_cupy(rho)
 95        elif funcid==13:
 96            return lda_c_pw_mod_cupy(rho)
 97        elif funcid==101:
 98            return gga_x_pbe_cupy(rho, sigma)
 99        elif funcid==106:
100            return gga_x_b88_cupy(rho, sigma)
101        elif funcid==130:
102            return gga_c_pbe_cupy(rho, sigma)
103        elif funcid==131:
104            return gga_c_lyp_cupy(rho, sigma)
105        else:
106            print('The specified functional is not implemented in PyFock. You need to use LibXC to calculate the functional values.')
107            exit()
108    else:
109        if funcid==1:
110            return lda_x(rho)
111        elif funcid==7:
112            return lda_c_vwn(rho)
113        elif funcid==12:
114            return lda_c_pw(rho)
115        elif funcid==13:
116            return lda_c_pw_mod(rho)
117        elif funcid==101:
118            return gga_x_pbe(rho, sigma)
119        elif funcid==106:
120            return gga_x_b88(rho, sigma)
121        elif funcid==130:
122            return gga_c_pbe(rho, sigma)
123        elif funcid==131:
124            return gga_c_lyp(rho, sigma)
125        else:
126            print('The specified functional is not implemented in PyFock. You need to use LibXC to calculate the functional values.')
127            exit()