4. The XC Functional Network

4.1. The RXCModel_GGA Class

class xcquinox.xc.RXCModel_GGA(xnet, cnet)

Bases: Module

__init__(xnet, cnet)

Initializes a combined X+C GGA model into one object that can be used in optimizing a functional.

Parameters:

• xnet (eqx.Module) – The exchange enhancement factor network

• cnet (eqx.Module) – The correlation enhancement factor network

__call__(inputs)

Calls the networks to generate an EPSILON value (rho* xc_energy_density), which libxc expects derivatives of.

Divide by rho to get the energy density itself.

This call calculates: RHO*( ( lda_x(RHO)*Fx(inputs) )+(pw92c_unpolarized(RHO)*Fc(inputs) )

Parameters:

inputs (array) – The expected inputs (here, [rho, sigma]) that will be vmapped against while calling the x/c enhancement factor networks.

Returns:

The epsilon values generated by the networks.

Return type:

array

4.2. The RXCModel_MGGA Class

class xcquinox.xc.RXCModel_MGGA(xnet, cnet)

Bases: Module

__init__(xnet, cnet)

Initializes a combined X+C MGGA model into one object that can be used in optimizing a functional.

Parameters:

• xnet (eqx.Module) – The exchange enhancement factor network

• cnet (eqx.Module) – The correlation enhancement factor network

__call__(inputs)

Calls the networks to generate an EPSILON value (rho* xc_energy_density), which libxc expects derivatives of.

Divide by rho to get the energy density itself.

This call calculates: RHO*( ( lda_x(RHO)*Fx(inputs) )+(pw92c_unpolarized(RHO)*Fc(inputs) )

Parameters:

inputs (array) – The expected inputs (here, [rho, sigma, lapl, tau]) that will be vmapped against while calling the x/c enhancement factor networks.

Returns:

The epsilon values generated by the networks.

Return type:

array

4.3. The LDA_X Class

class xcquinox.xc.LDA_X

Bases: Module

__init__()

__init__ Constructs an object whose forward pass computes the LDA exchange energy based on a given density input.

__call__(rho)

__call__ Computes the LDA exchange energy for a given value of the density.

\[E_x = -\frac{3}{4} \Big(\frac{3\rho}{\pi} \Big)^{1/3}\]

Parameters:

rho (float, broadcastable) – The value of the density

Returns:

The LDA exchange energy for the input value(s).

Return type:

float, broadcastable

4.4. The PW_C Class

class xcquinox.xc.PW_C

Bases: Module

__init__()

__init__ Constructs an object whose forward pass computes the UEG correlation energy, as parameterized by Perdew & Wang

DOI: 10.1103/PhysRevB.45.13244.

__call__(rs, zeta)

__call__ Forward pass, computing the correlation energy per electron for the given input values, rs and zeta, where

\[r_s = \Big[\frac{3}{4\pi (\rho_\uparrow+\rho_\downarrow)} \Big]^{1/3}\]

\[\zeta = \frac{\rho_\uparrow-\rho_\downarrow}{\rho_\uparrow+\rho_\downarrow}\]

Atomic units are assumed.

Parameters:

• rs (float, broadcastable) – The Wigner-Seiz radius corresponding to the given density value

• zeta (float, broadcastable) – The spin-polarization

4.5. The eXC Class

class xcquinox.xc.eXC(grid_models=[], heg_mult=True, pw_mult=True, level=1, exx_a=None, epsilon=1e-06, debug=False, nlstart_i=3, nlend_i=12, verbose=False)

Bases: Module

__init__(grid_models=[], heg_mult=True, pw_mult=True, level=1, exx_a=None, epsilon=1e-06, debug=False, nlstart_i=3, nlend_i=12, verbose=False)

__init__ Defines the XC functional

Constructed with two MLPs – one for the local exchange energy on the grid, the other for the local correlation energy.

Parameters:

• grid_models (list, optional) – list of eX (local exchange) or eC (local correlation). Defines the xc-models/enhancement factors, defaults to []

• heg_mult (bool, optional) – Use homoegeneous electron gas exchange (multiplicative if grid_models is not empty), defaults to True

• pw_mult (bool, optional) – Use homoegeneous electron gas correlation (Perdew & Wang), defaults to True

• level (int, optional) – Controls the number of density “descriptors” generated. 1: LDA, 2: GGA, 3:meta-GGA, 4: meta-GGA + electrostatic (nonlocal), defaults to 1

• exx_a (float, optional) – Exact exchange mixing parameter, defaults to None

• epsilon (float, optional) – Offset to avoid div/0 in calculations, defaults to 1e-8

• debug (bool, optional) – Controls printing of various stats throughout, defaults to False

• nlstart_i (int, optional) – If level > 3, this controls the number of CIDER Nonlocal parameters are selected, defaults to 3 as the first nonlocal CIDER descriptor

• nlend_i (int, optional) – If level > 3, this controls the number of CIDER Nonlocal parameters are selected, defaults to 12 as the last nonlocal CIDER descriptor

• verbose (bool, optional) – Flag to determine printing of a lot more extra information to the terminal, useful for debuggin, defaults to False

__call__(dm, ao_eval, grid_weights, mf=None, coor=None)

__call__ Forward call for the XC network to get the grid point e_xc

Generates the density-on-grid from the density matrix, atomic orbital evaluation, and the grid weights from a :pyscfad: calculation.

Parameters:

• dm (jax.Array) – Density matrix

• ao_eval (jax.Array) – Atomic orbitals evaluated on the grid

• grid_weights (jax.Array) – Grid weights associated to the grid on which the atomic orbitals are evaluated

• mf (pyscfad.dft.RKS kernel) – A pyscf(ad) converged calculation kernel if self.level > 3, used for building the CIDER nonlocal descriptors, defaults to None

• coor (jax.Array) – Grid coordinates associated to the grid on which the atomic orbitals are evaluated, passed into eval_grid_models, defaults to None

Returns:

Exc, exchange-correlation energy from integrating the network calls across the grid

Return type:

float

vprint(str)

Prints if verbose tag is set on the object

Parameters:

str (str) – The string to be printed

l_1(rho)

l_1 Level 1 (LDA-level) Descriptor – Creates dimensionless quantity from rho.

Equation 3 from the base paper.

\[x_0 = \rho^{1/3}\]

Parameters:

rho (jax.Array) – density

Returns:

Scaled density

Return type:

jax.Array

l_2(rho, gamma)

l_2 Level 2 (GGA-level) Descriptor – Reduced gradient density

Equation 5 from the base paper.

\[x_2=s=\frac{1}{2(3\pi^2)^{1/3}} \frac{|\nabla \rho|}{\rho^{4/3}}\]

Parameters:

• rho (jax.Array) – density

• gamma (jax.Array) – squared density gradient

Returns:

reduced density gradient s

Return type:

jax.Array

l_3(rho, gamma, tau)

l_3 Level 3 (MGGA-level) Descriptor – Reduced kinetic energy density

Equation 6 from the base paper.

\[\tau^W = \frac{|\nabla \rho|^2}{8\rho}, \tau^{unif} = \frac{3}{10} (3\pi^2)^{2/3}\rho^{5/3}.\]

Parameters:

• rho (jax.Array) – density

• gamma (jax.Array) – squared density gradient

• tau (jax.Array) – kinetic energy density

Returns:

reduced kinetic energy density

Return type:

jax.Array

l_4(rho, nl)

l_4 Level 4 (Non-local level) Descriptor – Unitless electrostatic potential

Parameters:

• rho (jax.Array) – density

• nl (jax.Array) – non-local values arising from density contractions

Returns:

the non-local descriptors

Return type:

jax.Array

nl_4(mf, dm, ao=None, gw=None, coor=None)

Level 4 (non-local level) descriptor generator – generates the CIDER nonlocal descriptors, given a density matrix and a converged kernel with associated mol and grids

Optional parameters not provided will use the values already associated with mf.grids; provide custom values for things like grid subsets.

Parameters:

• mf (pyscf(ad).dft.RKS kernel) – The converged calculation kernel, passed to RKSAnalyzer for descriptor generation

• dm (jax.Array) – Density matrix to use in nonlocal desciptor generation

• ao (jax.Array) – Atomic orbital values to use in nonlocal desciptor generation, defaults to None

• gw (jax.Array) – Grid weights to use in nonlocal desciptor generation, defaults to None

• coor (jax.Array) – Grid coordinates to use in nonlocal desciptor generation, defaults to None

Returns:

The non-local CIDER descriptors, from self.nlstart_i to self.nlend_i

Return type:

jax.Array

get_descriptors(rho0_a, rho0_b, gamma_a, gamma_b, gamma_ab, tau_a, tau_b, spin_scaling=False, mf=None, dm=None, ao=None, gw=None, coor=None)

get_descriptors Creates ‘ML-compatible’ descriptors from the electron density and its gradients, a & b correspond to spin channels

Parameters:

• rho0_a (jax.Array) – \(\rho\) in spin-channel a

• rho0_b (jax.Array) – \(\rho\) in spin-channel b

• gamma_a (jax.Array) – \(|\nabla \rho|^2\) in spin-channel b

• gamma_b (jax.Array) – \(|\nabla \rho|^2\) in spin-channel b

• gamma_ab (jax.Array) – \(|\nabla \rho|^2\), contracted from both spin channels

• tau_a (jax.Array) – KE density in spin-channel a

• tau_b (jax.Array) – KE density in spin-channel b

• spin_scaling (bool, optional) – Flag for spin-scaling, defaults to False

• mf (pyscf(ad).dft.RKS kernel) – The converged calculation kernel, passed to RKSAnalyzer for descriptor generation if self.level > 3, defaults to None

• dm (jax.Array) – Density matrix to use in nonlocal desciptor generation if self.level > 3, defaults to None

• ao (jax.Array) – Atomic orbital values to use in nonlocal desciptor generation if self.level > 3, defaults to None

• gw (jax.Array) – Grid weights to use in nonlocal desciptor generation if self.level > 3, defaults to None

• coor (jax.Array) – Grid coordinates to use in nonlocal desciptor generation if self.level > 3, defaults to None

Returns:

Array of the machine-learning descriptors on the grid

Return type:

jax.Array

eval_grid_models(rho, mf=None, dm=None, ao=None, gw=None, coor=None)

eval_grid_models Evaluates all models stored in self.grid_models along with HEG exchange and correlation

Parameters:

• rho (jax.Array) – List/array with [rho0_a,rho0_b,gamma_a,gamma_ab,gamma_b, dummy for laplacian, dummy for laplacian, tau_a, tau_b, non_loc_a, non_loc_b] Shape assumes, for instance, that rho0_a = rho[:, 0], etc.

• mf (pyscf(ad).dft.RKS kernel) – The converged calculation kernel, passed to RKSAnalyzer for nonlocal descriptor generation

• dm (jax.Array) – Density matrix to use in nonlocal desciptor generation

• ao (jax.Array) – Atomic orbitals to use in non-local descriptor generation, defaults to None

• gw (jax.Array) – Grid weights to use in non-local descr generation, defaults to None

• coor (jax.Array) – Grid coordinates to use in non-local descr generation, defaults to None

Returns:

The exchange-correlation energy density (on the grid)

Return type:

jax.Array