Coverage for eminus/xc/utils.py: 84.39%
173 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-21 12:19 +0000
1# SPDX-FileCopyrightText: 2023 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""Utility functions for exchange-correlation functionals."""
5import inspect
6import sys
8from eminus import backend as xp
9from eminus import config
10from eminus.logger import log
11from eminus.utils import add_maybe_none
13from .gga_c_chachiyo import gga_c_chachiyo, gga_c_chachiyo_spin
14from .gga_c_pbe import gga_c_pbe, gga_c_pbe_spin
15from .gga_c_pbe_sol import gga_c_pbe_sol, gga_c_pbe_sol_spin
16from .gga_x_chachiyo import gga_x_chachiyo, gga_x_chachiyo_spin
17from .gga_x_pbe import gga_x_pbe, gga_x_pbe_spin
18from .gga_x_pbe_sol import gga_x_pbe_sol, gga_x_pbe_sol_spin
19from .lda_c_chachiyo import lda_c_chachiyo, lda_c_chachiyo_spin
20from .lda_c_chachiyo_mod import lda_c_chachiyo_mod, lda_c_chachiyo_mod_spin
21from .lda_c_pw import lda_c_pw, lda_c_pw_spin
22from .lda_c_pw_mod import lda_c_pw_mod, lda_c_pw_mod_spin
23from .lda_c_vwn import lda_c_vwn, lda_c_vwn_spin
24from .lda_x import lda_x, lda_x_spin
25from .lda_xc_corr_ksdt import lda_xc_corr_ksdt
26from .lda_xc_gdsmfb import lda_xc_gdsmfb, lda_xc_gdsmfb_spin
27from .lda_xc_ksdt import lda_xc_ksdt, lda_xc_ksdt_spin
30def get_xc(xc, n_spin, Nspin, dn_spin=None, tau=None, xc_params=None, dens_threshold=0):
31 """Handle and get exchange-correlation functionals.
33 Args:
34 xc: Exchange and correlation identifier.
35 n_spin: Real-space electronic densities per spin channel.
36 Nspin: Number of spin states.
38 Keyword Args:
39 dn_spin: Real-space gradient of densities per spin channel.
40 tau: Real-space kinetic energy densities per spin channel.
41 xc_params: Exchange-correlation functional parameters.
42 dens_threshold: Do not treat densities smaller than the threshold.
44 Returns:
45 Exchange-correlation energy density and potential.
46 """
47 if isinstance(xc, str):
48 xc = parse_functionals(xc)
49 f_exch, f_corr = xc
50 if xc_params is None:
51 xc_params = {}
53 # Only use non-zero values of the density
54 n = xp.sum(n_spin, axis=0)
55 nz_mask = xp.atleast_2d(xp.nonzero(n > dens_threshold)[0])
56 n_nz = n[nz_mask]
57 # Zeta is only needed for non-zero values of the density
58 zeta_nz = get_zeta(n_spin[:, nz_mask])
59 # dn_spin is only needed for non-zero values of the density
60 if dn_spin is not None:
61 dn_spin_nz = dn_spin[:, nz_mask[0], :]
62 else:
63 dn_spin_nz = None
65 def handle_functional(fxc):
66 """Calculate a given functional fxc, same for exchange and correlation."""
67 # Calculate with the libxc extra...
68 if ":" in fxc:
69 from eminus.extras.libxc import libxc_functional
71 fxc = fxc.split(":")[-1]
72 exc, vxc, vsigma, vtau = libxc_functional(fxc, n_spin, Nspin, dn_spin, tau, xc_params)
73 # ...or use an internal functional
74 else:
75 if Nspin == 2 and fxc != "mock_xc":
76 fxc += "_spin"
77 exc_nz, vxc_nz, vsigma_nz = IMPLEMENTED[fxc](
78 n_nz, zeta=zeta_nz, dn_spin=dn_spin_nz, Nspin=Nspin, **xc_params
79 )
80 # Map the non-zero values back to the right dimension
81 exc = xp.zeros_like(n)
82 exc[nz_mask] = exc_nz
83 vxc = xp.zeros_like(n_spin)
84 for s in range(Nspin):
85 vxc[s, nz_mask] = vxc_nz[s]
86 if vsigma_nz is not None:
87 vsigma = xp.zeros((len(vsigma_nz), len(exc)))
88 for i in range(len(vsigma)):
89 vsigma[i, nz_mask] = vsigma_nz[i]
90 else:
91 vsigma = None
92 # There are no internal meta-GGAs
93 vtau = None
94 return exc, vxc, vsigma, vtau
96 ex, vx, vsigmax, vtaux = handle_functional(f_exch) # Calculate the exchange part
97 ec, vc, vsigmac, vtauc = handle_functional(f_corr) # Calculate the correlation part
98 return ex + ec, vx + vc, add_maybe_none(vsigmax, vsigmac), add_maybe_none(vtaux, vtauc)
101def get_exc(xc, n_spin, Nspin, dn_spin=None, tau=None, xc_params=None, dens_threshold=0):
102 """Get the exchange-correlation energy density.
104 This is a convenience function to interface :func:`~eminus.xc.utils.get_xc`.
106 Args:
107 xc: Exchange and correlation identifier.
108 n_spin: Real-space electronic densities per spin channel.
109 Nspin: Number of spin states.
111 Keyword Args:
112 dn_spin: Real-space gradient of densities per spin channel.
113 tau: Real-space kinetic energy densities per spin channel.
114 xc_params: Exchange-correlation functional parameters.
115 dens_threshold: Do not treat densities smaller than the threshold.
117 Returns:
118 Exchange-correlation energy potential.
119 """
120 exc, _, _, _ = get_xc(xc, n_spin, Nspin, dn_spin, tau, xc_params, dens_threshold)
121 return exc
124def get_vxc(xc, n_spin, Nspin, dn_spin=None, tau=None, xc_params=None, dens_threshold=0):
125 """Get the exchange-correlation potential.
127 This is a convenience function to interface :func:`~eminus.xc.utils.get_xc`.
129 Args:
130 xc: Exchange and correlation identifier.
131 n_spin: Real-space electronic densities per spin channel.
132 Nspin: Number of spin states.
134 Keyword Args:
135 dn_spin: Real-space gradient of densities per spin channel.
136 tau: Real-space kinetic energy densities per spin channel.
137 xc_params: Exchange-correlation functional parameters.
138 dens_threshold: Do not treat densities smaller than the threshold.
140 Returns:
141 Exchange-correlation energy density.
142 """
143 _, vxc, vsigma, vtau = get_xc(xc, n_spin, Nspin, dn_spin, tau, xc_params, dens_threshold)
144 return vxc, vsigma, vtau
147def parse_functionals(xc):
148 """Parse exchange-correlation functional strings to the internal format.
150 Args:
151 xc: Exchange and correlation identifier, separated by a comma.
153 Returns:
154 Exchange and correlation string.
155 """
156 # Check for combined aliases
157 try:
158 # Remove underscores when looking up in the dictionary
159 xc_ = xc.replace("_", "")
160 xc = ALIAS[xc_]
161 except KeyError:
162 pass
164 # Parse functionals
165 functionals = []
166 for f in xc.split(","):
167 if ":" in f or f in IMPLEMENTED:
168 f_xc = f
169 elif not f:
170 f_xc = "mock_xc"
171 else:
172 try:
173 # Remove underscores when looking up in the dictionary
174 f_ = f.replace("_", "")
175 f_xc = XC_MAP[f_]
176 except KeyError:
177 log.exception(f'No functional found for "{f}".')
178 raise
179 functionals.append(f_xc)
181 # If only one or no functional has been parsed append with mock functionals
182 functionals.extend("mock_xc" for _ in range(2 - len(functionals)))
183 return functionals
186def parse_xc_type(xc):
187 """Parse functional strings to identify the corresponding functional type.
189 Args:
190 xc: Exchange and correlation identifier, separated by a comma.
192 Returns:
193 Functional type.
194 """
195 xc_type = []
196 for func in xc:
197 if ":" in func:
198 xc_id = func.split(":")[-1]
199 # Try to parse the functional using pylibxc at first
200 try:
201 family = parse_xc_libxc(xc_id)
202 # Otherwise parse it with PySCF
203 except (ImportError, AssertionError):
204 family = parse_xc_pyscf(xc_id)
206 if family == 1:
207 xc_type.append("lda")
208 elif family == 2:
209 xc_type.append("gga")
210 elif family == 4:
211 xc_type.append("meta-gga")
212 else:
213 msg = "Unsupported functional family."
214 raise NotImplementedError(msg)
215 # Fall back to internal xc functionals
216 elif "gga" in func:
217 xc_type.append("gga")
218 else:
219 xc_type.append("lda")
221 # When mixing functional types use the higher level of theory
222 if xc_type[0] != xc_type[1]:
223 log.warning("Detected mixing of different functional types.")
224 if "meta-gga" in xc_type:
225 return "meta-gga"
226 return "gga"
227 return xc_type[0]
230def parse_xc_libxc(xc_id):
231 """Parse functional type by its ID using pylibxc.
233 Args:
234 xc_id: Functional ID or identifier.
236 Returns:
237 Functional type.
238 """
239 if not config.use_pylibxc:
240 raise AssertionError
241 import pylibxc
243 if not xc_id.isdigit():
244 xc_id = pylibxc.util.xc_functional_get_number(xc_id)
246 func = pylibxc.LibXCFunctional(int(xc_id), 1)
247 if func._needs_laplacian:
248 msg = "meta-GGAs that need a laplacian are not supported."
249 raise NotImplementedError(msg)
250 return func.get_family()
253def parse_xc_pyscf(xc_id):
254 """Parse functional type by its ID using PySCF.
256 Args:
257 xc_id: Functional ID or identifier.
259 Returns:
260 Functional type.
261 """
262 from pyscf.dft.libxc import is_gga, is_lda, is_meta_gga, needs_laplacian, XC_CODES
264 if not xc_id.isdigit():
265 xc_id = XC_CODES[xc_id.upper()]
267 if needs_laplacian(int(xc_id)):
268 msg = "meta-GGAs that need a laplacian are not supported."
269 raise NotImplementedError(msg)
270 # Use the same values as in parse_xc_libxc
271 if is_lda(xc_id):
272 return 1
273 if is_gga(xc_id):
274 return 2
275 if is_meta_gga(xc_id):
276 return 4
277 return -1
280def get_xc_defaults(xc):
281 """Get the default parameters and values for a given set of functionals.
283 Args:
284 xc: Exchange and correlation identifier, separated by a comma.
286 Returns:
287 Default parameters and values.
288 """
289 if isinstance(xc, str):
290 xc = parse_functionals(xc)
292 # Names of special kewyword arguments that should not be used in xc_params
293 SPECIAL_NAMES = ["dn_spin", "Nspin"]
295 params = {}
296 for func in xc:
297 # If pylibxc functionals are used determine the default values from it
298 if ":" in func:
299 # This only works for pylibxc, not with PySCF
300 if not config.use_pylibxc or "pylibxc" not in sys.modules:
301 msg = "ext_params only work with pylibxc as the libxc backend, not with pyscf."
302 raise NotImplementedError(msg)
303 from pylibxc import LibXCFunctional
305 fxc = func.split(":")[-1]
306 try:
307 f_xc = LibXCFunctional(int(fxc), 1)
308 except ValueError:
309 f_xc = LibXCFunctional(fxc, 1)
310 fxc_params = dict(zip(f_xc.get_ext_param_names(), f_xc.get_ext_param_default_values()))
312 # Analyze the signature for implemented functionals
313 if func in IMPLEMENTED:
314 sig = inspect.signature(IMPLEMENTED[func])
315 fxc_params = {
316 param.name: param.default
317 for param in sig.parameters.values()
318 if param.default is not inspect.Parameter.empty
319 }
321 # Remove special names from the parsed parameters
322 for special in SPECIAL_NAMES:
323 if special in fxc_params:
324 del fxc_params[special]
326 # Append all parameters, warn if a parameter has been used before
327 for name in fxc_params:
328 if name in params:
329 log.warning(
330 f'External parameter "{name}" is used in the exchange AND correlation part. '
331 "It will be passed to both functionals if used in xc_params."
332 )
333 params[name] = fxc_params[name]
334 return params
337def get_zeta(n_spin):
338 """Calculate the relative spin polarization.
340 Args:
341 n_spin: Real-space electronic densities per spin channel.
343 Returns:
344 Relative spin polarization.
345 """
346 # If only one spin is given return an array of ones as if the density only is in one channel
347 if len(n_spin) == 1:
348 return xp.ones_like(n_spin[0])
349 return (n_spin[0] - n_spin[1]) / (n_spin[0] + n_spin[1])
352def mock_xc(n, Nspin=1, **kwargs):
353 """Mock exchange-correlation functional with no effect (spin-paired).
355 Args:
356 n: Real-space electronic density.
357 Nspin: Number of spin states.
359 Keyword Args:
360 **kwargs: Throwaway arguments.
362 Returns:
363 Mock exchange-correlation energy density and potential.
364 """
365 zeros = xp.zeros_like(n)
366 return zeros, xp.stack([zeros] * Nspin), None
369#: Map functional names with their respective implementation.
370IMPLEMENTED = {
371 i.__name__: i
372 for i in (
373 mock_xc,
374 gga_c_chachiyo,
375 gga_c_chachiyo_spin,
376 gga_c_pbe,
377 gga_c_pbe_spin,
378 gga_c_pbe_sol,
379 gga_c_pbe_sol_spin,
380 gga_x_chachiyo,
381 gga_x_chachiyo_spin,
382 gga_x_pbe,
383 gga_x_pbe_spin,
384 gga_x_pbe_sol,
385 gga_x_pbe_sol_spin,
386 lda_x,
387 lda_x_spin,
388 lda_c_pw,
389 lda_c_pw_spin,
390 lda_c_pw_mod,
391 lda_c_pw_mod_spin,
392 lda_c_vwn,
393 lda_c_vwn_spin,
394 lda_c_chachiyo,
395 lda_c_chachiyo_spin,
396 lda_c_chachiyo_mod,
397 lda_c_chachiyo_mod_spin,
398 lda_xc_corr_ksdt,
399 lda_xc_gdsmfb,
400 lda_xc_gdsmfb_spin,
401 lda_xc_ksdt,
402 lda_xc_ksdt_spin,
403 )
404}
406#: Map functional shorthands to the actual functional names.
407XC_MAP = {
408 # lda_x
409 "1": "lda_x",
410 "s": "lda_x",
411 "lda": "lda_x",
412 "slater": "lda_x",
413 # lda_c_vwn
414 "7": "lda_c_vwn",
415 "vwn": "lda_c_vwn",
416 "vwn5": "lda_c_vwn",
417 # lda_c_pw
418 "12": "lda_c_pw",
419 "pw": "lda_c_pw",
420 "pw92": "lda_c_pw",
421 # lda_c_pw_mod
422 "13": "lda_c_pw_mod",
423 "pwmod": "lda_c_pw_mod",
424 "pw92mod": "lda_c_pw_mod",
425 # gga_x_pbe
426 "101": "gga_x_pbe",
427 "pbex": "gga_x_pbe",
428 # gga_x_pbe_sol
429 "116": "gga_x_pbe_sol",
430 "pbesolx": "gga_x_pbe_sol",
431 # gga_c_pbe
432 "130": "gga_c_pbe",
433 "pbec": "gga_c_pbe",
434 # gga_c_pbe_sol
435 "133": "gga_c_pbe_sol",
436 "pbesolc": "gga_c_pbe_sol",
437 # lda_xc_ksdt
438 "259": "lda_xc_ksdt",
439 "ksdt": "lda_xc_ksdt",
440 # lda_c_chachiyo
441 "287": "lda_c_chachiyo",
442 "chachiyo": "lda_c_chachiyo",
443 # gga_x_chachiyo
444 "298": "gga_x_chachiyo",
445 "chachiyox": "gga_x_chachiyo",
446 # lda_c_chachiyo_mod
447 "307": "lda_c_chachiyo_mod",
448 "chachiyomod": "lda_c_chachiyo_mod",
449 # gga_c_chachiyo
450 "309": "gga_c_chachiyo",
451 "chachiyoc": "gga_c_chachiyo",
452 # lda_xc_corr_ksdt
453 "318": "lda_xc_corr_ksdt",
454 "corrksdt": "lda_xc_corr_ksdt",
455 # lda_xc_gdsmfb
456 "577": "lda_xc_gdsmfb",
457 "gdsmfb": "lda_xc_gdsmfb",
458}
460#: Dictionary of common functional aliases.
461ALIAS = {
462 "svwn": "slater,vwn5",
463 "svwn5": "slater,vwn5",
464 "spw92": "slater,pw92mod",
465 "pbe": "pbex,pbec",
466 "pbesol": "pbesolx,pbesolc",
467 "chachiyo": "chachiyox,chachiyoc",
468}