3. Network Structure

3.1. The GGA_FxNet_s Class

class xcquinox.net.GGA_FxNet_s(depth: int, nodes: int, seed: int, lob_lim=1.804)

Bases: Module

S: Exchange enhancementt factor for GGA.

The input to the network is the reduced density gradient, s, and the output is the enhancement factor, Fx.

__init__(depth: int, nodes: int, seed: int, lob_lim=1.804)

Constructor for the exchange enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 1 – just the gradient information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be the reduced density gradient, :s:, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.2. The GGA_FcNet_s Class

class xcquinox.net.GGA_FcNet_s(depth: int, nodes: int, seed: int, lob_lim=2.0)

Bases: Module

S: Correlation enhancement factor for GGA.

The input to the network is the reduced density gradient, s, and the output is the enhancement factor, Fc.

__init__(depth: int, nodes: int, seed: int, lob_lim=2.0)

Constructor for the correlation enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 2 – both the density and gradient information is passed to the network.

The default Lieb-Oxford bound the outputs are wrapped here is set to 2.0, to enforce the non-negativity of the correlation energy.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 2

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be the reduced density gradient, :s:, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.3. The GGA_FxNet_sigma Class

class xcquinox.net.GGA_FxNet_sigma(depth: int, nodes: int, seed: int, lob_lim=1.804, lower_rho_cutoff=1e-12)

Bases: Module

S: Exchange enhancementt factor for GGA.

It takes rho and sigma (grad_rho²) as input and outputs the exchange enhancement factor, Fx. It transforms the input to the reduced density gradient, s, and then passes it through the network.

__init__(depth: int, nodes: int, seed: int, lob_lim=1.804, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 1 – just the gradient information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.4. The GGA_FcNet_sigma Class

class xcquinox.net.GGA_FcNet_sigma(depth: int, nodes: int, seed: int, lob_lim=2.0, lower_rho_cutoff=1e-12)

Bases: Module

S: Correlation enhancement factor for GGA.

It takes rho and sigma (grad_rho²) as input and outputs the correlation enhancement factor, Fc. It transforms the input to the reduced density gradient, s, and then passes it through the network.

__init__(depth: int, nodes: int, seed: int, lob_lim=2.0, lower_rho_cutoff=1e-12)

Constructor for the correlation enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 2 – both the density and gradient information is passed to the network.

The default Lieb-Oxford bound the outputs are wrapped here is set to 2.0, to enforce the non-negativity of the correlation energy.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 2

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.5. The GGA_FxNet_sigma_UNC Class

class xcquinox.net.GGA_FxNet_sigma_UNC(depth: int, nodes: int, seed: int, lob_lim=1.804, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=1.804, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 1 – just the gradient information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.6. The GGA_FcNet_sigma_UNC Class

class xcquinox.net.GGA_FcNet_sigma_UNC(depth: int, nodes: int, seed: int, lob_lim=2.0, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=2.0, lower_rho_cutoff=1e-12)

Constructor for the correlation enhancement factor object, for the GGA case.

In a GGA XC functional, the relevant quantities are (rho, grad_rho). Here, the network’s input size is hard-coded to 2 – both the density and gradient information is passed to the network.

The default Lieb-Oxford bound the outputs are wrapped here is set to 2.0, to enforce the non-negativity of the correlation energy.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 2

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, array of size 2 in order (rho, gradient_descriptor)) – _description_

Returns:

The enhancement factor value

Return type:

float

3.7. The MGGA_FxNet_sigma Class

class xcquinox.net.MGGA_FxNet_sigma(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the MGGA case.

In a MGGA XC functional, the relevant quantities are (rho, grad_rho, laplacian_rho, tau=kinetic energy density). Here, the network’s input size is hard-coded to 2 – just the gradient and alpha (related to tau) information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, one-dimensional array of size 4 in order [rho, sigma, laplacian_rho, alpha]) – A one-dimensional list/array of inputs [rho, sigma, laplacian_rho, alpha]

Returns:

The enhancement factor value

Return type:

float

3.8. The MGGA_FxNet_sigma Class

class xcquinox.net.MGGA_FxNet_sigma(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the MGGA case.

In a MGGA XC functional, the relevant quantities are (rho, grad_rho, laplacian_rho, tau=kinetic energy density). Here, the network’s input size is hard-coded to 2 – just the gradient and alpha (related to tau) information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, one-dimensional array of size 4 in order [rho, sigma, laplacian_rho, alpha]) – A one-dimensional list/array of inputs [rho, sigma, laplacian_rho, alpha]

Returns:

The enhancement factor value

Return type:

float

3.9. The MGGA_FxNet_sigma_transform Class

class xcquinox.net.MGGA_FxNet_sigma_transform(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the MGGA case.

In a MGGA XC functional, the relevant quantities are (rho, grad_rho, laplacian_rho, tau=kinetic energy density). Here, the network’s input size is hard-coded to 2 – just the gradient and alpha (related to tau) information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

This network transforms function inputs [rho, sigma, lapl, tau] to the below inputs for the network: rho -> log(rho**1/3 + 1e-5) [not sent to network itself, here, due to contraints] sigma -> (1-exp(-s**2))*log(s+1) tau -> log((alpha+1)/2)

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, one-dimensional array of size 4 in order [rho, sigma, laplacian_rho, alpha]) – A one-dimensional list/array of inputs [rho, sigma, laplacian_rho, alpha]

Returns:

The enhancement factor value

Return type:

float

3.10. The MGGA_FxNet_sigma_transform Class

class xcquinox.net.MGGA_FxNet_sigma_transform(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Bases: Module

__init__(depth: int, nodes: int, seed: int, lob_lim=1.174, lower_rho_cutoff=1e-12)

Constructor for the exchange enhancement factor object, for the MGGA case.

In a MGGA XC functional, the relevant quantities are (rho, grad_rho, laplacian_rho, tau=kinetic energy density). Here, the network’s input size is hard-coded to 2 – just the gradient and alpha (related to tau) information is passed to the network, to guarantee that the energy yielded from this multiplicative factor behaves correctly under uniform scaling of the electron density and obeys the spin-scaling relation.

This network transforms function inputs [rho, sigma, lapl, tau] to the below inputs for the network: rho -> log(rho**1/3 + 1e-5) [not sent to network itself, here, due to contraints] sigma -> (1-exp(-s**2))*log(s+1) tau -> log((alpha+1)/2)

Parameters:

• depth (int) – Depth of the neural network

• nodes (int) – Number of nodes in each layer

• seed (int) – The random seed to initiate baseline weight values for the network

• lob_lim (float, optional) – The Lieb-Oxford bound to respect, defaults to 1.804

• lower_rho_cutoff (float, optional) – a cut-off to bypass potential division by zero in the division by rho, defaults to 1e-12

__call__(inputs)

The network’s forward pass, resulting in the enhancement factor associated to the input gradient descriptor.

NOTE: This forward pass is explicitly NOT vectorized – it expects one grid point worth of data, the (rho, gradient_descriptor) values at that point. This structure expects the :jax.vmap: call to be coded OUTSIDE of the network class.

NOTE: Here, the gradient_descriptor is assumed to be Libxc’s/PySCF’s internal variable for the density gradient – sigma (gradient squared in non-spin-polarized, gradient contracted with itself in spin-polarized). This is so that we have easy access to automatic derivatives with respect to sigma, thus can generate v_sigma and use in convergence testing. However, within the call sigma is translated to the reduced density gradient, :s:, which the network is still assumed to be parameterized by, and the call is structured in such a way to respect the UEG limits for when gradients vanish. Namely, when s = 0, Fx = 1, so the resulting e = Fx*e_heg = e_heg.

Parameters:

inputs (tuple, list, one-dimensional array of size 4 in order [rho, sigma, laplacian_rho, alpha]) – A one-dimensional list/array of inputs [rho, sigma, laplacian_rho, alpha]

Returns:

The enhancement factor value

Return type:

float

3.11. The LOB Class

class xcquinox.net.LOB(limit: float)

Bases: Module

__init__(limit: float)

Utility function to squash output to [-1, limit-1] inteval.

Parameters:

limit (float, optional) – The Lieb-Oxford bound value to impose, defaults to 1.804

__call__(x)

Method calling the actual mapping of the input to the desired bounded region.

Parameters:

x (float) – Energy value to map back into the bounded region.

Returns:

Energy value mapped into bounded region.

Return type:

float

3.12. The eX Class

class xcquinox.net.eX(n_input, n_hidden=16, depth=3, use=[], ueg_limit=False, lob=1.804, seed=92017, spin_scaling=True)

Bases: Module

__init__(n_input, n_hidden=16, depth=3, use=[], ueg_limit=False, lob=1.804, seed=92017, spin_scaling=True)

__init__ Local exchange model based on MLP.

Receives density descriptors in this order : [rho, s, alpha, nl], where the input may be truncated depending on XC-level of approximation.

The MLP generated is hard-coded to have one output value – the predicted exchange energy given a specific input from the grid.

Parameters:

• n_input (int) – Input dimensions (LDA: 1, GGA: 2, meta-GGA: 3, …)

• n_hidden (int, optional) – Number of hidden nodes (three hidden layers used by default), defaults to 16

• depth (int, optional) – Depth of the MLP, defaults to 3

• use (list, optional) – Only these indices are used as input to the model (can be used to omit density as input to enforce uniform density scaling). These indices are also used to enforce UEG where the assumed order is [s, alpha, …], defaults to []

• ueg_limit (bool, optional) – Flag to determine whether or not to enforce uniform homoegeneous electron gas limit, defaults to False

• lob (float, optional) – Enforce this value as local Lieb-Oxford bound (don’t enforce if set to 0), defaults to 1.804

• seed (int, optional) – Random seed used to generate initial weights and biases for the MLP, defaults to 92017

__call__(rho, **kwargs)

__call__ Forward pass for the exchange network.

Uses :jax.vmap: to vectorize evaluation of the MLP on the descriptors, assuming a shape [batch, *, n_input]

Parameters:

rho (jax.Array) – The descriptors to the MLP – transformed densities and gradients appropriate to the XC-level. This network will only use the dimensions specified in self.use.

Returns:

The exchange energy on the grid

Return type:

jax.Array

3.13. The eC Class

class xcquinox.net.eC(n_input=2, n_hidden=16, depth=3, use=[], ueg_limit=False, lob=2.0, seed=92017, spin_scaling=False)

Bases: Module

__init__(n_input=2, n_hidden=16, depth=3, use=[], ueg_limit=False, lob=2.0, seed=92017, spin_scaling=False)

__init__ Local correlation model based on MLP.

Receives density descriptors in this order : [rho, spinscale, s, alpha, nl], where the input may be truncated depending on XC-level of approximation

Parameters:

• n_input (int) – Input dimensions (LDA: 2, GGA: 3 , meta-GGA: 4), defaults to 2.

• n_hidden (int, optional) – Number of hidden nodes (three hidden layers used by default), defaults to 16

• depth (int, optional) – Depth of the MLP, defaults to 3

• use (list, optional) – Only these indices are used as input to the model. These indices are also used to enforce UEG where the assumed order is [s, alpha, …], defaults to []

• ueg_limit (bool, optional) – Flag to determine whether or not to enforce uniform homoegeneous electron gas limit, defaults to False

• lob (float, optional) – Enforce this value as local Lieb-Oxford bound (don’t enforce if set to 0), defaults to 2.0

• seed (int, optional) – Random seed used to generate initial weights and biases for the MLP, defaults to 92017

__call__(rho, **kwargs)

__call__ Forward pass for the correlation network.

Uses :jax.vmap: to vectorize evaluation of the MLP on the descriptors, assuming a shape [*, n_input]

Parameters:

rho (jax.Array) – The descriptors to the MLP – transformed densities and gradients appropriate to the XC-level. This network will only use the dimensions specified in self.use in determining the UEG limits.

Returns:

The exchange energy on the grid

Return type:

jax.Array