5. Utility Structures and Functions

5.1. gen_grid_s

xcquinox.utils.gen_grid_s(npts=30000, start_stop_rho=(0.01, 5), start_stop_s=(0.01, 5), train_pct=0.8, ranseed=92017, sigma=False)[source]

Generates a grid of rho, grad_rho, and reduced_density_gradient values. Will generate ~sqrt(npts) sized rho and s arrays, used to generate a grad_rho array.

The indices of these arrays are then randomly sub-sampled to yield (1 - train_pct) percent of validation indices, which are excluded from the training meshgrid.

Parameters:

  • npts (int, optional) – Number of points to generate in each array, defaults to 30000. The function will get close, but not necessarily exactly the number of points, unless the desired value is a perfect square.

  • start_stop_rho (tuple, optional) – Bounds for the generated reduced_density_gradient values, defaults to (0.01, 5)

  • start_stop_s (tuple, optional) – Bounds for the generated rho values, defaults to (0.01, 5)

  • train_pct (float, optional) – Percentage of generated BASE values to use in training, defaults to 0.8

  • ranseed (int, optional) – Seed to set for training/validation index selection, defaults to 92017

  • sigma (bool, optional) – If flagged, will return array of SIGMA (= grad_rho**2) values instead of “grad_rho” values, defaults to False

Returns:

List of [(train_inds, val_inds), (rho_values, grad_rho_values, s_values), (trho_flat, tgrad_rho_flat, ts_flat), (vrho_flat, vgrad_rho_flat, vs_flat)]

Return type:

list

5.2. PBE_Fx

xcquinox.utils.PBE_Fx(rho, grad_rho, lower_rho_cutoff=1e-12)[source]

Given density and density gradient magnitude values, calculates the PBE Exchange Enhancement Factor as denoted in:

Equation 14 from PBE paper – DOI: 10.1103/PhysRevLett.77.3865

Parameters:

  • rho (float, broadcastable array) – the density value on the grid

  • grad_rho (float, broadcastable array) – the magnitude of the density gradient values on the grid

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

Returns:

The numerical value(s) of the exchange enhancement factor

Return type:

float, broadcastable array

5.3. PBE_Fc

xcquinox.utils.PBE_Fc(rho, grad_rho, lower_rho_cutoff=1e-12)[source]

Given density and density gradient magnitude values, calculates the PBE Correlation Enhancement Factor as denoted in:

CRITICALLY, this function is only valid for the non-spin-polarized case where zeta = 0

Equation 3 from PBE paper – DOI: 10.1103/PhysRevLett.77.3865

Parameters:

  • rho (float, broadcastable array) – the density value on the grid

  • grad_rho (float, broadcastable array) – the magnitude of the density gradient values on the grid

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

Returns:

The numerical value(s) of the correlation enhancement factor

Return type:

float, broadcastable array

5.4. pw91_correlation_energy_density

xcquinox.utils.pw91_correlation_energy_density(rho)[source]

Calculate the correlation energy density according to the PW91 functional.

Parameters: - rho: electron density (numpy array or scalar) - grad_rho: gradient of the electron density (numpy array or scalar)

Returns: - correlation energy density (numpy array or scalar)

5.5. lda_x

xcquinox.utils.lda_x(rho)[source]

The HEG exchange energy density.

Parameters:

rho (jax.numpy array) – Total electron density array on a grid.

Returns:

Exchange energy density array.

Return type:

jax.numpy array

5.6. pw92c

xcquinox.utils.pw92c(rho)[source]

Implements the Perdew-Wang ‘92 local correlation (beyond RPA) for the unpolarized case. Reference: J.P.Perdew & Y.Wang, PRB, 45, 13244 (1992)

Parameters:

rho (jax.numpy array) – Total electron density array on a grid.

Returns:

Correlation energy density and potential arrays.

Return type:

tuple of (jax.numpy array, jax.numpy array)

5.7. pw92c_unpolarized

xcquinox.utils.pw92c_unpolarized(rho)[source]

Implements the Perdew-Wang ‘92 local correlation (beyond RPA) for the unpolarized case. Reference: J.P.Perdew & Y.Wang, PRB, 45, 13244 (1992)

Parameters:

rho (jax.numpy array) – Total electron density array on a grid.

Returns:

Correlation energy density array.

Return type:

jax.numpy array

5.8. pad_array

xcquinox.utils.pad_array(arr, max_arr, shape=None, device=CpuDevice(id=0))[source]

Utility function designed to pad an array to a given maximum size, for use during jitted forward passes.

Since JAX jit compilations re-compile everytime a new array shape is passed, padding is a tactic to avoid this memory consumption.

Parameters:

  • arr (jax.Array) – The original array meant to be padded

  • max_arr (jax.Array) – The array whose size is to be emulated, or an empty value if manually specifying shape

  • shape (tuple of ints, optional) – The array shape you want to pad to, if known a priori, defaults to None

  • device (jax.Device, optional) – Where to place the padded arrays, defaults to jax.devices(‘cpu’)[0]

Returns:

Padded version of arr

Return type:

jax.Array

5.9. pad_array_list

xcquinox.utils.pad_array_list(arrlst, device=CpuDevice(id=0))[source]

Automatically finds maximum size from a list of arrays, and pads all arrays to that size

Parameters:

  • arrlst (list) – List of jax.Arrays with varying sizes

  • device (jax.Device, optional) – Where to place the newly padded arrays, defaults to jax.devices(‘cpu’)[0]

Returns:

List of the padded arrays

Return type:

list of jax.Array

5.10. ase_atoms_to_mol

xcquinox.utils.ase_atoms_to_mol(atoms, basis='6-311++G(3df,2pd)', charge=0, spin=None)[source]

Converts an ASE.Atoms object into a PySCF(ad).gto.Mol object

Parameters:

  • atoms (ASE.Atoms) – ASE.Atoms object of a single molecule/system

  • basis (str, optional) – Basis set to assign in PySCF(AD), defaults to ‘6-311++G(3df,2pd)’

  • charge (int, optional) – Global charge of molecule, defaults to 0

  • spin (int, optional) – Specify spin if desired. Defaults to None, which has PySCF guess spin based on electron number/occupation. Use with care.

Returns:

A string of the molecule’s name, and the Mole object for pyscfad to work with.

Return type:

(str, pyscfad.gto.Mole)

5.11. get_spin

xcquinox.utils.get_spin(at)[source]

Returns single-atom, ground-state spin configurations, since pyscf does not guess this correctly.

Relies on xcquinox.utils.spins_dict, a hard-coded dictionary populated with atom names and associated spins.

Parameters:

at (ase.Atoms) – ase.Atoms object, with an associated ase.Atoms.info dictionary optionally containing keys ‘spin’, ‘radical’, or ‘openshell’ and associated integer pair. If specified, returns that value. Otherwise, uses the value in spins_dict for single atoms.

Returns:

The guessed spin value

Return type:

int

5.12. get_dm_moe

xcquinox.utils.get_dm_moe(dm, eri, vxc_grad_func, mo_occ, hc, s, ogd, alpha0=0.7)[source]

Generates the DM, molecular orbital energies, and molecular orbital coefficients for a model

_extended_summary_

Parameters:

  • dm (jax.Array) – The initial DM starting point for the calculation.

  • eri (jax.Array) – The electron repulsion integrals to use in generating the effective potential

  • vxc_grad_func (function) – A function F(dm) = E, with which the derivative is taken to generate Vxc

  • mo_occ (jax.Array) – The molecular orbital occupation numbers for the molecule

  • hc (jax.Array) – The molecule’s core Hamiltonian

  • s (jax.Array) – The molecule’s overlap matrix

  • ogd (tuple/jax.Array) – The molecule’s original DM dimensions

  • alpha0 (float, optional) – The mixing parameter used for the one-shot DM creation, defaults to 0.7

Returns:

tuple of (density matrix, molecular orbital energies, molecular orbital coefficients)

Return type:

tuple of jax.Array

5.13. make_rdm1

class xcquinox.utils.make_rdm1[source]

Bases: Module

__init__()[source]

A equinox.Module object whose forward pass constructs the 1-electron density matrix, given input molecular coefficients and occupations.

__call__(mo_coeff, mo_occ)[source]

Forward pass calculating one-particle reduced density matrix.

Parameters:

  • mo_coeff (jax.Array) – Molecular orbital coefficients

  • mo_occ (jax.Array) – Molecular orbital occupation numbers

Returns:

The RDM1

Return type:

jax.Array

5.14. get_rho

class xcquinox.utils.get_rho[source]

Bases: Module

__init__()[source]

A equinox.Module object whose forward pass constructs the density on the grid, given input density matrix and atomic orbital evaluations.

__call__(dm, ao_eval)[source]

Forward pass computing the density on the grid, projecting the density matrix onto the atomic orbitals

Parameters:

  • dm (jax.Array) – Density matrix

  • ao_eval (jax.Array) – Atomic orbitals evaluated on the grid, NOT higher order derivatives.

Returns:

Density on the grid

Return type:

jax.Array

5.15. energy_tot

class xcquinox.utils.energy_tot[source]

Bases: Module

__init__()[source]

A equniox.Module object whose forward pass constructs the total energy: (electron-electron + electron-ion; ion-ion not included)

__call__(dm, hcore, veff)[source]

Forward pass contraction to find total electron energy (e-e + e-ion)

Parameters:

  • dm (jax.Array) – Density matrix

  • hcore (jax.Array) – Core Hamiltonian

  • veff (jax.Array) – Effective potential for the electrons

Returns:

Total energy

Return type:

jax.Array/float

5.16. get_fock

class xcquinox.utils.get_fock[source]

Bases: Module

__init__()[source]

A equniox.Module object whose forward pass constructs the Fock matrix

__call__(hc, veff)[source]

Returns the Fock matrix, the sum of the core Hamiltonian and the effective potential (including xc-potential)

Parameters:

  • hc (jax.Array) – Core Hamiltonian

  • veff (jax.Array) – Effective Potential (including the xc-contribution)

Returns:

Fock matrix

Return type:

jax.Array

5.17. get_hcore

class xcquinox.utils.get_hcore[source]

Bases: Module

__init__()[source]

A equniox.Module object whose forward pass constructs the core Hamiltonian, given the nuclear potential and kinetic energy matrix

__call__(v, t)[source]

“Core” Hamiltionian forward pass, includes ion-electron and kinetic contributions

\[H_{core} = T + V_{nuc-elec}\]
Parameters:

  • v (jax.Array) – Electron-ion interaction energy

  • t (jax.Array) – Kinetic energy

Returns:

Hcore

Return type:

jax.Array

5.18. get_veff

class xcquinox.utils.get_veff[source]

Bases: Module

__init__()[source]

A equniox.Module object whose forward pass constructs the one-electron effective potential (not including local xc-potential, to be found with the network)

__call__(dm, eri)[source]

Forward pass, constructing the classical parts of the effective potential

Parameters:

  • dm (jax.Arary) – Density matrix

  • eri (jax.Array) – Electron repulsion integral tensor

Returns:

Effective potential

Return type:

jax.Array

5.19. eig

class xcquinox.utils.eig[source]

Bases: Module

__init__()[source]

A equniox.Module object whose forward pass solves the generalized eigenvalue problem for the Hamiltonian, using Cholesky decomposition

__call__(h, s_chol, ogdim)[source]

Solver for generalized eigenvalue problem, finding the molecular energies (eigenvalues) and orbital coefficients (eigenvectors)

Parameters:

  • h (jax.Array) – Hamiltionian

  • s_chol (jax.Array) – (Inverse) of the Cholesky decomposition of overlap matrix S

  • ogdim (tuple, iterable of some original shape) – The original dimensions of the input array, or their current shapes if not padded. JAX jit compilations require static arrays to avoid recompilation with each call, and backpropagating through this eigendecomposition with padded arrays results in NaNs. Arrays are masked to original dimensions to avoid this.

Returns:

The eigenvalues (molecular orbital enegies) and eigenvectors (molecular orbital coefficients)

Return type:

tuple of jax.Arrays