Atomic data

Atomic data#

Atom definitions, element properties, electronic structure (SCF), orbital evaluation, electron initialization, and pretraining utilities.

Elements and atoms#

class jaqmc.utils.atomic.elements.Element(symbol, atomic_number, period, group, spin=None)[source]#

Chemical element.

symbol[source]#: official symbol of element.

atomic_number[source]#: atomic number of element.

period[source]#: period to which the element belongs.

spin[source]#: overrides default ground-state spin-configuration based on the element’s group (main groups only).

property unpaired_electron: int[source]#

Canonical spin configuration (via Hund’s rules) of neutral atom.

Returns:: Number of unpaired electrons (as required by PySCF) in the neutral atom’s ground state.
Raises:: NotImplementedError – If the element is a transition metal and the spin value is not set at initialization.

Examples

>>> from jaqmc.utils.atomic.elements import SYMBOLS
>>> SYMBOLS['H'].unpaired_electron
1
>>> SYMBOLS['C'].unpaired_electron
2
>>> SYMBOLS['Ne'].unpaired_electron
0

class jaqmc.utils.atomic.atom.Atom(symbol, coords, atomic_number=-1, charge=-1)[source]#

A single atom in a molecular or periodic system.

Parameters:

symbol (str) – Chemical element symbol (e.g. "H", "Fe").
coords (list[float]) – 3D Cartesian coordinates in Bohr.
atomic_number (int, default: -1) – Nuclear charge. Derived from symbol if not given.
charge (float, default: -1) – Effective charge seen by valence electrons. Equals atomic_number by default; reduced when using ECPs.

class jaqmc.utils.atomic.atomic_system.AtomicSystemConfig(*, atoms, basis='sto-3g', ecp=None, electron_spins, fixed_spins_per_atom=None, electron_init_width=1.0)[source]#

Configuration for an atomic system.

Parameters:

atoms (list[Atom]) – List of atoms in the system.
electron_spins (tuple[int, int]) – Tuple of two integers representing the number of up and down electrons.
basis (str | dict[str, str], default: 'sto-3g') – The basis set for Hartree-Fock pretrain. Can be a string (e.g., “sto-3g”, “ccecpccpvdz”) or a dict mapping element symbols to basis names (e.g., {“Fe”: “ccecpccpvdz”, “O”: “cc-pvdz”}).
ecp (str | dict[str, str] | None, default: None) – Effective core potential specification. Can be None (no ECP), a string (same ECP for all atoms, e.g., “ccecp”), or a dict mapping element symbols to ECP names (e.g., {“Fe”: “ccecp”}). Elements not in the dict use all-electron treatment.
fixed_spins_per_atom (list[tuple[int, int]] | None, default: None) – Optional list of fixed spin configurations per atom.
electron_init_width (float, default: 1.0) – Width of the Gaussian distribution for initializing electron positions.

Self-consistent field (SCF)#

class jaqmc.utils.atomic.scf.MolecularSCF(molecule=None, nelectrons=None, basis='cc-pVTZ', ecp=None, core_electrons=None, pyscf_mol=None, restricted=True)[source]#

Helper class for running Hartree-Fock (self-consistent field) with pyscf.

molecule[source]#: list of system.Atom objects giving the atoms in the molecule and their positions.

nelectrons[source]#: Tuple with number of alpha electrons and beta electrons.

basis[source]#: Basis set to use, best specified with the relevant string for a built-in basis set in pyscf. A user-defined basis set can be used (advanced). See https://sunqm.github.io/pyscf/gto.html#input-basis for more details.

pyscf_mol[source]#: the PySCF ‘Molecule’. If this is passed to the init, the molecule, nelectrons, and basis will not be used, and the calculations will be performed on the existing pyscf_mol

restricted[source]#: If true, use the restricted Hartree-Fock method, otherwise use the unrestricted Hartree-Fock method.

mean_field[source]#: the actual UHF object.

mo_coeff[source]#: The molecular orbital coefficients computed by Hartree-Fock.

eval_mos(positions)[source]#

Evaluates the Hartree-Fock single-particle orbitals at a set of points.

Parameters:

positions (Array) – numpy array of shape (N, 3) of the positions in space at which to evaluate the Hartree-Fock orbitals.

Return type:

tuple[Array, Array]

Returns:

Pair of numpy float64 arrays of shape (N, M) (deriv=False) or (4, N, M) (deriv=True), where 2M is the number of Hartree-Fock orbitals. The (i-th, j-th) element in the first (second) array gives the value of the j-th alpha (beta) Hartree-Fock orbital at the i-th electron position in positions. For restricted (RHF, ROHF) calculations, the two arrays will be identical. If deriv=True, the first index contains [value, x derivative, y derivative, z derivative].

Raises:

RuntimeError – If Hartree-Fock calculation has not been performed using run.
NotImplementedError – If Hartree-Fock calculation used Cartesian Gaussian-type orbitals as the underlying basis set.

eval_orbitals(pos, nspins)[source]#

Evaluates SCF orbitals at a set of positions.

Parameters:

pos (Array) – an array of electron positions to evaluate the orbitals at, of shape (…, nelec, 3), where the leading dimensions are arbitrary, nelec is the number of electrons and the spin up electrons are ordered before the spin down electrons.
nspins (tuple[int, int]) – tuple with number of spin up and spin down electrons.

Return type:

tuple[Array, Array]

Returns:

tuple with matrices of orbitals for spin up and spin down electrons, with the same leading dimensions as in pos.

Raises:

ValueError – If input is not a NumPy or JAX array.

eval_slater(pos, nspins)[source]#

Evaluate Slater determinant sign and log|det|.

Parameters:

pos (Array) – An array of electron positions to evaluate the orbitals at.
nspins (tuple[int, int]) – Tuple with number of spin up and spin down electrons.

Return type:

tuple[Array, Array]

Returns:

Tuple with sign and log absolute value of Slater determinant.

run(dm0=None)[source]#

Runs the Hartree-Fock calculation.

Parameters:: dm0 (ndarray | None, default: None) – Optional density matrix to initialize the calculation.
Returns:: A pyscf scf object (i.e. pyscf.scf.rhf.RHF, pyscf.scf.uhf.UHF or pyscf.scf.rohf.ROHF depending on the spin and restricted settings).

class jaqmc.utils.atomic.scf.PeriodicSCF(atoms=None, nelectrons=None, basis='cc-pVTZ', lattice_vectors=None, kpts=None, ecp=None, core_electrons=None, pyscf_cell=None, restricted=True, rcut=None)[source]#

Helper class for running periodic Hartree-Fock with k-point sampling.

This class handles the additional complexity of periodic systems compared to MolecularSCF:

K-point sampling: Electrons occupy Bloch states at different k-points in the Brillouin zone. The kpts parameter specifies the k-point mesh.
Bloch phase corrections: When electron coordinates are wrapped into the primitive cell, a phase factor \(e^{i\mathbf{k}\cdot\mathbf{R}}\) must be applied (see eval_mos() for details).
Variable occupancy: Different k-points may have different numbers of occupied orbitals due to band filling. This requires pre-extracting MO coefficients via _extract_mo_coeffs and padding for array ops.
K-point assignment: The klist parameter in eval_orbitals() specifies which k-point’s orbitals to use for each electron, enabling twist-averaged boundary conditions.

Parameters:

atoms (Sequence[Atom] | None, default: None) – List of Atom objects giving atom types and positions.
nelectrons (tuple[int, int] | None, default: None) – Tuple of (n_alpha, n_beta) electrons per primitive cell.
basis (str | Mapping[str, str] | None, default: 'cc-pVTZ') – Basis set string (e.g., "cc-pVTZ").
lattice_vectors (Array | ndarray | None, default: None) – Primitive cell lattice vectors, shape (3, 3).
kpts (Array | ndarray | None, default: None) – K-points for sampling, shape (nk, 3). If None, uses Gamma point only.
ecp (str | Mapping[str, str] | None, default: None) – Effective core potentials mapping.
core_electrons (Mapping[str, int] | None, default: None) – Core electrons mapping.
pyscf_cell (Cell | None, default: None) – Pre-built PySCF Cell. If provided, atoms/basis/lattice are ignored.
restricted (bool, default: True) – Use restricted (True) or unrestricted (False) HF.
rcut (float | None, default: None) – Cutoff radius for lattice sum. If None, uses PySCF estimate.

eval_mos(positions)[source]#

Evaluate MOs at positions using all k-points.

For periodic systems, the molecular orbital at k-point \(\mathbf{k}\) is constructed from Bloch-summed atomic orbitals:

\[\psi_{n\mathbf{k}}(\mathbf{r}) = \sum_\mu C_{\mu n}^{\mathbf{k}} \phi_{\mu\mathbf{k}}(\mathbf{r})\]

where \(C_{\mu n}^{\mathbf{k}}\) are the MO coefficients and \(\phi_{\mu\mathbf{k}}\) are Bloch AOs.

When an electron coordinate \(\mathbf{r}\) lies outside the primitive cell, it is wrapped back: \(\mathbf{r} = \mathbf{r}' + \mathbf{R}\) where \(\mathbf{r}'\) is in the cell and \(\mathbf{R}\) is a lattice vector. Due to Bloch’s theorem, this introduces a phase factor:

\[\psi_{n\mathbf{k}}(\mathbf{r}) = e^{i\mathbf{k}\cdot\mathbf{R}} \psi_{n\mathbf{k}}(\mathbf{r}')\]

Parameters:: positions (Array) – Electron coordinates with shape (N, 3).
Return type:: tuple[Array, Array]
Returns:: Tuple of (alpha_mos, beta_mos), each with shape (nk, N, max_occ) where max_occ is the maximum number of occupied orbitals across k-points. Zero-padded for k-points with fewer occupied orbitals.
Raises:: RuntimeError – If Hartree-Fock calculation has not been performed.

eval_orbitals(pos, nspins)[source]#

Evaluate occupied orbital matrices for the Slater determinant.

This method produces dense orbital matrices by evaluating all electrons at all k-points and concatenating the results. The k-point structure is implicit in the MO coefficients from the periodic HF calculation.

For each spin, the orbital matrix has shape (n_spin, n_spin) where:

Row i corresponds to electron i
Column j corresponds to orbital j (ordered by k-point)

The resulting matrices can be used directly with jnp.linalg.slogdet to compute the Slater determinant. The block-diagonal structure (electrons at different k-points being orthogonal) emerges naturally from the physics.

Parameters:

pos (Array) – Electron positions with shape (..., nelec, 3).
nspins (tuple[int, int]) – Tuple of (n_alpha, n_beta) electrons.

Return type:

tuple[Array, Array]

Returns:

Tuple of (alpha_matrix, beta_matrix) orbital matrices with shapes (..., n_alpha, n_alpha) and (..., n_beta, n_beta).

Raises:

RuntimeError – If Hartree-Fock calculation has not been performed.

eval_slater(pos, nspins)[source]#

Evaluate Slater determinant sign and log|det|.

Parameters:

pos (Array) – An array of electron positions to evaluate the orbitals at.
nspins (tuple[int, int]) – Tuple with number of spin up and spin down electrons.

Return type:

Array

Returns:

Log value of Slater determinant (complex).

get_orbital_kpoints()[source]#

Get k-point assignment for each orbital in the wavefunction.

In periodic systems, each orbital is associated with a specific k-point. This method constructs an array where each row is the k-point vector for the corresponding orbital.

The orbitals are ordered as: all spin-up orbitals (concatenated across k-points), then all spin-down orbitals.

Return type:: Array
Returns:: Array of shape (n_orbitals, 3) where n_orbitals is the total number of occupied orbitals across all k-points and spins. Each row contains the k-point coordinates for that orbital.
Raises:: RuntimeError – If Hartree-Fock calculation has not been performed.

Example

For a system with 4 k-points and 1 occupied orbital per k-point per spin, the returned array has shape (8, 3) with k-points ordered as [k0, k1, k2, k3, k0, k1, k2, k3].

run(dm0=None)[source]#

Run the k-point HF calculation.

Parameters:: dm0 (ndarray | None, default: None) – Optional density matrix to initialize the calculation.
Returns:: The PySCF mean field object after SCF convergence.

Gaussian-type orbitals#

class jaqmc.utils.atomic.gto.AtomicOrbitalEvaluator(atom_list, basis_dict)[source]#

An atomic orbital evaluator using GTO basis in JAX.

atom_list[source]#: A list of tuples (Z, [x, y, z]) that specifies the atom types (Z) and positions (x,y,z) in Bohr for the Mol

basis_dict[source]#: A dictionary of cGTO specifications of the form {“Z”: [[L, [alpha_1, w_1], [alpha_2, w_2], …], …], …}

spec[source]#: The dict returned by _get_orbital_construction_dict().

classmethod from_pyscf(mol)[source]#

Initialize from a pyscf.gto.Mole.

Return type:: Self
Returns:: The initialized instance.

class jaqmc.utils.atomic.gto.PBCAtomicOrbitalEvaluator(atom_list, basis_dict, image_translation_vectors)[source]#

Evaluator for Bloch-summed Gaussian-type orbitals in periodic systems.

This class constructs crystal orbitals suitable for periodic boundary conditions by summing localized atomic orbitals over lattice translation vectors with k-point-dependent phase factors. The resulting orbitals satisfy Bloch’s theorem and can be used for k-point sampling in solid-state calculations.

The evaluator wraps an AtomicOrbitalEvaluator for the underlying GTO evaluation and handles the lattice summation internally.

eval_aos[source]#: The underlying AtomicOrbitalEvaluator for computing localized atomic orbitals.

lattice_vectors[source]#: Array of shape [L, 3] containing lattice translation vectors \(\mathbf{L}\) used in the Bloch sum, sorted by distance from the origin.

Example:

cell = pyscf.pbc.gto.Cell(atom="H 0 0 0", a=np.eye(3) * 4, basis="sto-3g")
cell.build()
kpts = cell.make_kpts([2, 2, 2])
evaluator = PBCAtomicOrbitalEvaluator.from_pyscf(cell)
aos = evaluator(coords, kpts=kpts)  # shape: [nk, ncoords, nao]

jaqmc.utils.atomic.gto.solid_harmonic(r, l_max)[source]#

Computes all solid harmonics r**ell Y_{ell, m} for all ell <= l_max.

Returns:: all solid harmonics r**ell Y_{ell, m} for all ell <= l_max.
Return type:: Array

jaqmc.utils.atomic.gto.cart2sph(r)[source]#

Convert Cartesian coordinates to spherical coordinates.

Parameters:: r (Array) – Cartesian coordinates with x, y, and z on the last axis.
Return type:: tuple[Array, Array, Array]
Returns:: Tuple (rho, phi, theta) where rho is the radial distance, phi is the azimuthal angle in the xy plane, and theta is the polar angle measured from the positive z axis.

Effective core potentials#

jaqmc.utils.atomic.ecp.get_valence_spin_config(symbol, ecp)[source]#

Get valence electron spin configuration for an element with an ECP.

Builds a temporary PySCF molecule to determine the number of valence electrons after core electrons are replaced by the ECP. Returns mol.nelec directly, which splits the valence electrons into (n_alpha, n_beta) based on the element’s ground-state spin.

Parameters:

symbol (str) – Element symbol (e.g., "Fe").
ecp (str | dict[str, str]) – ECP name (e.g., "ccecp") or per-element mapping.

Return type:

tuple[int, int]

Returns:

Tuple of (n_alpha, n_beta) valence electrons.

jaqmc.utils.atomic.ecp.get_core_electrons(atoms, ecp)[source]#

Get the number of core electrons removed by ECP for each element.

Parameters:

atoms (Sequence[Atom]) – List of Atom objects.
ecp (str | dict[str, str] | None) – ECP specification (name, per-element mapping, or None).

Return type:

dict[str, int]

Returns:

Mapping from element symbol to number of core electrons removed. Elements without core electrons removed are omitted.

Electron initialization#

jaqmc.utils.atomic.initialization.distribute_spins(rngs, atoms, total_spins)[source]#

Distributes total spins among atoms to match the target system spin.

The function initializes spins from atomic defaults and iteratively adjusts the distribution (transferring spins between channels) until the global total_spins constraint is met.

Parameters:

rngs (PRNGKey) – Random number generator key.
atoms (list[Atom]) – List of Atom objects defining the system.
total_spins (tuple[int, int]) – A tuple (n_alpha, n_beta) representing the total number of electrons in the system.

Return type:

list[tuple[int, int]]

Returns:

A list of tuples, where each tuple (n_alpha, n_beta) corresponds to the spin configuration of an atom in the atoms list.

jaqmc.utils.atomic.initialization.initialize_electrons_gaussian(rng, atom_coords, spins_per_atom, batch_size, init_width)[source]#

Initializes electron positions with Gaussian distribution around atoms.

Electrons are assigned to atoms based on the local spin configuration provided in spins_per_atom. The initial positions are centered on the atoms with added Gaussian noise.

Parameters:

rng (PRNGKey) – Random number generator key.
atom_coords (Array) – Array of atom coordinates with shape (natoms, 3).
spins_per_atom (list[tuple[int, int]]) – List of (n_alpha, n_beta) tuples for each atom, dictating how many electrons are initialized around each nucleus.
batch_size (int) – Number of walkers.
init_width (float) – tandard deviation (scale) of the Gaussian noise added to the atomic centers.

Return type:

Array

Returns:

Electron positions array of shape (batch_size, nelec, 3).

Pretraining#

jaqmc.utils.atomic.pretrain.make_pretrain_log_amplitude(log_psi_fn, scf_log_amplitude_fn, scf_fraction=0.0)[source]#

Creates a log amplitude function for pretraining sampling.

Creates a function that returns either the SCF ansatz, the neural network ansatz, or a weighted mixture of the two. This allows sampling from an SCF-biased distribution during pretraining.

Parameters:

log_psi_fn (WavefunctionEvaluate[TypeVar(DataT, bound= Data), Array]) – Neural network log amplitude function.
scf_log_amplitude_fn (Callable[[TypeVar(DataT, bound= Data)], Array]) – Function that takes data and returns the SCF log amplitude.
scf_fraction (float, default: 0.0) – Mixing fraction for SCF (0.0 = pure NN, 1.0 = pure SCF).

Return type:

WavefunctionEvaluate[TypeVar(DataT, bound= Data), Array]

Returns:

A log amplitude function for sampling.

Type Parameters:

DataT – Concrete Data subtype consumed by both input callables.

Raises:

ValueError – If scf_fraction is not between 0 and 1.

jaqmc.utils.atomic.pretrain.make_pretrain_loss(orbitals_fn, scf, nspins, full_det=False)[source]#

Returns a pretrain loss estimator matching NN orbitals to HF orbitals.

Used by both molecule and solid workflows to pretrain the wavefunction against Hartree-Fock orbitals from an SCF calculation.

Parameters:

orbitals_fn (NumericWavefunctionEvaluate) – Function to evaluate NN orbitals.
scf (MolecularSCF | PeriodicSCF) – SCF instance (MolecularSCF or PeriodicSCF).
nspins (tuple[int, int]) – Electron spin counts as (n_alpha, n_beta).
full_det (bool, default: False) – Whether to use full determinant.

Return type:

Estimator

Atomic data

Contents

Atomic data#

Elements and atoms#

Self-consistent field (SCF)#

Gaussian-type orbitals#

Effective core potentials#

Electron initialization#

Pretraining#