API #

class deepqmc.pretrain.gto.GTOBasis(*args, **kwargs)[source]#

Represent a GTO basis of a molecule.

classmethod from_pyscf(mol)[source]#

Create the input of the constructor from a pyscf molecule.

Parameters:: mol – a pyscf molecule with Cartesian Gaussian-type orbitals (mol.cart == True), such as the one returned by pyscf_from_hamil().

Loss functions#

class deepqmc.loss.LossFunction(*args, **kwargs)[source]#

Protocol for loss functions used during wave function training.

A LossFunction takes model parameters, an RNG key, and a batch of electron configurations and returns a scalar loss value together with auxiliary per-sample data.

class deepqmc.loss.LossFunctionFactory(*args, **kwargs)[source]#

Protocol for loss function factories.

A LossFunctionFactory constructs a LossFunction from a Hamiltonian and an ansatz, encapsulating the choice of objective (energy, overlap, spin, …) and any associated hyperparameters.

class deepqmc.loss.LossAndGradFunction(*args, **kwargs)[source]#

Protocol for combined loss-and-gradient functions.

A LossAndGradFunction has the same call signature as a LossFunction but additionally returns the gradient of the loss with respect to the model parameters. It is typically obtained by applying jax.value_and_grad() to a LossFunction.

deepqmc.loss.create_loss_fn(hamil, ansatz, clip_mask_fn, clip_mask_overlap_fn=None, alpha=None, spin_penalty=None, spin_penalty_type='squared', spin_penalty_states=None, scale_overlap_by=None, sort_states_by=None, min_gap_scale_factor=0.1, local_energy_batch_size=None)[source]#

Return type:: LossFunction

Energy loss#

deepqmc.loss.energy.compute_local_energy(rng, hamil, ansatz, params, phys_conf, batch_size=None)[source]#

Compute a batch of local energies.

Parameters:

rng (KeyArray) – rng key to use for the generation of the ECP quadratures.
hamil (MolecularHamiltonian) – the Hamiltonian of the system.
ansatz (ParametrizedWaveFunction) – the parametrized wave function.
params (Params) – the current parameters of the Ansatz.
phys_conf (PhysicalConfiguration) – a batch of input to the Ansatz.

Returns:

a tuple of local energy and: statistics.

Return type:

Tuple[Energy, Stats]

deepqmc.loss.energy.compute_mean_energy(local_energy, weight)[source]#

Compute the mean of a batch of local energies.

Parameters:

local_energy (Energy) – the batch of local energies.
weight (Weight) – the weight of each sample in the batch.

Returns:

a tuple of mean energy and: statistics.

Return type:

Tuple[Energy, Stats]

deepqmc.loss.energy.compute_mean_energy_tangent(local_energy, weight, log_psi_tangent, gradient_mask)[source]#

Compute the tangent of the mean energy with respect to the Ansatz parameters.

Parameters:

local_energy (Energy) – a batch of local energies.
weight (Weight) – the weights of each sample in the batch.
log_psi_tangent (Array) – the jvp of the WF values with respect to the Ansatz parameters.
gradient_mask (Array) – a boolean samplewise mask to apply to the gradients.

Returns:

the jvp of the mean energy with respect to the Ansatz parameters.

Return type:

Overlap loss#

class deepqmc.loss.overlap.OverlapGradientScaleFactory(*args, **kwargs)[source]#: Callable that computes the scaling factor of the overlap gradient.

deepqmc.loss.overlap.compute_mean_overlap(psi_ratio, weight)[source]#

Compute an estimate of the overlap matrix from WF ratios.

Parameters:

psi_ratio (Array) – the WF ratios \(\text{ratio}[i,\,j,\,:]=\frac{\Psi_i(r\sim\Psi^2_j)}{\Psi_j(r\sim\Psi^2_j)}\), shape: [n_wfs, n_wfs, elec_batch_size].
weight (Weight) – the sample weights, shape [n_wfs, elec_batch_size].

Returns:

tuple of the symmetric overlap matrix estimate, shaped [mol_batch_size, n_wfs, n_wfs], and overlap statistics.

Return type:

tuple[Array, Stats]

deepqmc.loss.overlap.compute_mean_overlap_tangent(psi_ratio, weight, log_psi_tangent, ratio_gradient_mask, overlap, scale_factory, data)[source]#

Compute the tangent of the overlap matrix with respect to the Ansatz parameters.

Parameters:

psi_ratio (Array) – the ratio of WF values.
weight (Weight) – the weight of each sample.
log_psi_tangent (Array) – the jvp of the WF values with respect to the parameters of the Ansatz.
ratio_gradient_mask (Array) – a samplewise boolean mask to apply to the gradients.
overlap (Array) – the overlap matrix estimate.
scale_factory (OverlapGradientScaleFactory) – function that computes the scaling factor of the overlap gradient.
data (DataDict) – input data passed to the scale_factory function.

Returns:

the jvp of the sum of the upper triangle of the overlap matrix with: respect to the Ansatz parameters.

Return type:

deepqmc.loss.overlap.compute_psi_ratio(ansatz, params, phys_conf)[source]#

Compute the ratio of all wave function for a batch of samples.

Parameters:

ansatz (Ansatz) – the ansatz object.
params (Params) – the current parameters of the Ansatz.
phys_conf (PhysicalConfiguration) – the input to the Ansatz, shape: [mol_batch_size, electronic_states, electron_batch_size, ...].

Returns:

the tuple of the WF ratios and overlap: statistics.

Return type:

tuple[Array, Stats]

deepqmc.loss.overlap.compute_single_sample_psi_ratios(psi, mean_log_psi)[source]#

Compute all possible ratios between the WFs for a single sample.

The mean of each WF value is subtracted before computing exponentials to avoid over/underflow.

Parameters:

psi (Psi) – all WF values for a single molecule and electron sample, shape: [electronic_states, electronic_states].
mean_log_psi (Array) – the mean log magnitude of the WFs, shape [electronic_states].

Returns:

the WF ratios: \(R[i,\,j]=\frac{\Psi_i(r\sim\Psi^2_j)}{\Psi_j(r\sim\Psi^2_j)}\), shape [electronic_states, electronic_states].

Return type:

deepqmc.loss.overlap.compute_wave_function_values(ansatz, params, phys_conf)[source]#

Compute the value of all WFs at samples drawn from all WFs.

Parameters:

ansatz (Ansatz) – the ansatz object.
params (Params) – PyTree of WF parameters, with leading axis over the different WFs. Shape: [n_wfs, ...].
phys_conf (PhysicalConfiguration) – input physical configuration samples, with leading axis over the different WFs the samples were drawn from. Shape: [n_wfs, elec_batch_size, ...]

Returns:

the WF values and auxiliary statistics. The WF values are \(\Psi[i, \, j, \, :] = \Psi_i({\bf r} \sim \Psi^2_j)\), shape: [n_wfs, n_wfs, elec_batch_size].

Return type:

tuple[Psi, Stats]

deepqmc.loss.overlap.no_scaling(data)[source]#

Return unit scaling factor.

Return type:: jax.Array

deepqmc.loss.overlap.scale_by_energy_gap(data, min_gap_scale_factor=0.1)[source]#

Scale the overlap gradient by the energy gap between the two states.

Return type:: jax.Array

deepqmc.loss.overlap.scale_by_energy_std(data, min_gap_scale_factor=0.01)[source]#

Scale the overlap gradient by the std. dev. of the states’ energies.

Return type:: jax.Array

deepqmc.loss.overlap.scale_by_max_gap_std(data, min_gap_scale_factor=0.1)[source]#

Scale the overlap gradient by the max of the energy gap and std. dev.

Return type:: jax.Array

deepqmc.loss.overlap.symmetrize_overlap_with_clipped_geometric_mean(x)[source]#

Symmetrize the overlap using the clipped geometric mean of it and its transpose.

Useful for computing an estimate of the overlap matrix using Monte Carlo samples from all WFs. Given input \(x_{ij}\) this function computes:

\[y_{ij}=\text{sign}(x_{ij}) \sqrt{\max(0, \, x_{ij} \cdot x_{ji})} \]

The product is clamped from below at zero only, to keep the square root defined when the signs of \(x_{ij}\) and \(x_{ji}\) differ (in which case \(y_{ij}\) is zero). Otherwise the two signs will agree, and we can use either one of them to compute the sign of \(y_{ij}\).

Parameters:: x (jax.Array) – the non-symmetric overlap values: \(x_{ij}=\frac1N\sum_{{\bf r}_j\sim\Psi^2_j}\frac{\Psi_i({\bf r}_j)} {\Psi_j({\bf r}_j)}\)
Return type:: jax.Array

Spin loss#

deepqmc.loss.spin.compute_mean_spin(spin_contriutions, weight, states=None)[source]#

Compute the mean of a batch of spin contributions.

Parameters:

spin_contriutions (Array) – the batch of local spin_contributions.
weight (Weight) – the weight of each sample in the batch.
states (list[int] | None) – (list[int] | None): list of state indices to compute spin for. If None, compute spin for all states.

Returns:

a tuple of spin expectation value and statistics.

Return type:

tuple[Array, Stats]

deepqmc.loss.spin.compute_mean_spin_raising_tangent(spin_raising_contributions, spin_raising_tangent, weight, log_psi_tangent, gradient_mask, states=None)[source]#

Compute the tangent of the spin raising operator with respect to the parameters.

Parameters:

spin_raising_contributions (Array) – a batch of spin raising contributions.
spin_raising_tangent (Array) – a batch of spin raising contribution tangents. This is the gradient of the local values of the spin raising contributions. Necessary, because the stochastic spin raising operator is not self-adjoint.
weight (Weight) – the weights of each sample in the batch.
log_psi_tangent (Array) – the jvp of the WF values with respect to the Ansatz parameters.
gradient_mask (Array) – a boolean samplewise mask to apply to the gradients.
states (list[int] | None) – (list[int] | None): list of state indices to compute spin for. If None, compute spin for all states.

Returns:

the jvp of the spin raising operator with respect to the: Ansatz parameters.

Return type:

deepqmc.loss.spin.compute_mean_spin_tangent(spin_contributions, weight, log_psi_tangent, gradient_mask, states=None)[source]#

Compute the tangent of the spin with respect to the Ansatz parameters.

Parameters:

spin_contributions (Array) – a batch of spin contributions.
weight (Weight) – the weights of each sample in the batch.
log_psi_tangent (Array) – the jvp of the WF values with respect to the Ansatz parameters.
gradient_mask (Array) – a boolean samplewise mask to apply to the gradients.
states (list[int] | None) – (list[int] | None): list of state indices to compute spin for. If None, compute spin for all states.

Returns:

the jvp of the spin with respect to the Ansatz parameters.

Return type:

deepqmc.loss.spin.compute_spin_contributions(hamil, ansatz, params, phys_conf, states=None)[source]#

Compute a batch of spin contributions.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
ansatz (Ansatz) – the Ansatz object.
params (Params) – the current parameters of the Ansatz.
phys_conf (PhysicalConfiguration) – a batch of input to the Ansatz.
states (list[int] | None) – (list[int] | None): list of state indices to compute spin for. If None, compute spin for all states.

Returns:

the samplewise contributions to spin expectation value.

Return type:

deepqmc.loss.spin.compute_spin_raising_contributions(rng, hamil, ansatz, phys_conf, params, batch_size=None, states=None)[source]#

Compute a batch of spin raising operator contributions.

Computes \(1 - \sum_{\alpha} \frac{\hat P_{\alpha\beta} \Psi}{\Psi}\) where a single \(\beta\) is sampled randomly from the spin down electrons.

Parameters:

rng (KeyArray) – a random key.
hamil (MolecularHamiltonian) – the Hamiltonian of the system.
ansatz (Ansatz) – the Ansatz object.
params (Params) – the current parameters of the Ansatz.
phys_conf (PhysicalConfiguration) – a batch of input to the Ansatz.
batch_size (int or None) – if specified, the batch size for the electron axis mapping. If None, use jax.vmap.
states (list[int] | None) – (list[int] | None): list of state indices to compute spin for. If None, compute spin for all states.

Returns:

the samplewise contributions to the stochastic spin raising: expectation value.

Return type:

Clipping#

deepqmc.loss.clip.clip_local_energy(clip_mask_fn, local_energy)[source]#

Apply a clipping function to the local energies.

The clipping function is twice vmapped: over the molecule batch, and electronic state dimensions.

Parameters:

clip_mask_fn (Callable[[Array], tuple[Energy, Array]]) – function taking as input an electron batch of local energies and returning a tuple of the clipped local energies and an identically shaped boolean mask array to be applied to the gradients.
local_energy (Array) – the electron batch of local energies, shape: [mol_batch_size, electronic_states, electron_batch_size // device_count].

Return type:

tuple[Energy, jax.Array]

deepqmc.loss.clip.clip_psi_ratio(clip_mask_fn, psi_ratio)[source]#

Apply a clipping function to the wave function ratios.

The clipping function is thrice vmapped: over the molecule batch, and the two electronic state dimensions of the wave function ratio array: \(\text{ratio}[i,\,j,\,:]=\frac{\Psi_i(r\sim\Psi^2_j)}{\Psi_j(r\sim\Psi^2_j)}\).

Parameters:

clip_mask_fn (Callable[[Array], tuple[Array, Array]]) – function taking as input an electron batch of ratios and returning a tuple of the clipped ratios and an identically shaped boolean mask array to be applied to the gradients.
psi_ratio (Array) – the electron batch of psi_ratios, shape: [mol_batch_size, electronic_states, electronic_states, electron_batch_size // device_count].

Returns:

the clipped WF ratios and gradient mask.

Return type:

deepqmc.loss.clip.median_clip_and_mask(x, clip_width, median_center, exclude_width=jax.numpy.inf)[source]#

Hard-clip values to a multiple of the mean absolute deviation from the center.

Parameters:

x (Array) – values to clip, shape [electron_batch_size].
clip_width (float) – number of mean absolute deviations (MADs) around the center within which values are kept; values outside this range are clipped to the boundary.
median_center (bool) – if True, use the median as the center; if False, use the mean.
exclude_width (float) – deviation threshold in MADs above which samples are excluded from gradient computation (gradient mask set to False). Defaults to jnp.inf (no exclusion).

Returns:

the clipped values and a boolean gradient: mask of the same shape, where False marks excluded outliers.

Return type:

deepqmc.loss.clip.median_log_squeeze_and_mask(x, clip_width=1.0, quantile=0.95, exclude_width=jax.numpy.inf)[source]#

Softly squeeze values toward the median using a log-squeeze function.

Values far from the median are continuously compressed rather than hard-clipped. The clipping scale is set to clip_width times the quantile-th quantile of absolute deviations from the median. Formally, the squeezed value is

\[\tilde{x}_i = \bar{x} + 2w\,\operatorname{log\_squeeze}\!\left(\frac{x_i - \bar{x}}{2w}\right), \]

where \(\bar{x}\) is the median and \(w = \texttt{clip\_width}\times Q_q(|x-\bar{x}|)\). For \(|x_i - \bar{x}| \ll w\) the squeezed value is close to \(x_i\); for large outliers it saturates at \(\bar{x} \pm 2w\).

Parameters:

x (Array) – values to squeeze, shape [electron_batch_size].
clip_width (float) – multiplier applied to the quantile to obtain the half-width \(w\) of the squeeze window.
quantile (float) – quantile of the absolute deviations used to set the natural scale of the distribution; default 0.95.
exclude_width (float) – deviation threshold in quantile units above which samples are excluded from gradient computation (gradient mask set to False). Defaults to jnp.inf (no exclusion).

Returns:

the squeezed values and a boolean gradient: mask of the same shape, where False marks excluded outliers.

Return type:

deepqmc.loss.clip.psi_ratio_clip_and_mask(psi_ratio, *, clip_width=10.0, exclude_width=jax.numpy.inf)[source]#

Clips WF ratios of a single batch of electron position samples.

Parameters:

psi_ratio (Array) – ratio of log WF values, shape [electron_batch_size]: \(\frac{\Psi_i({\bf r}_j)}{\Psi_j({\bf r}_j)}\).
clip_width (float) – clip width to use when clipping ratio.
exclude_width (float) – default: jnp.inf, deviation threshold above which outlier ratios are excluded from the overlap gradient computation.

Returns:

the clipped WF ratios and gradient mask.

Return type:

Sampling#

deepqmc.types.SamplerState = SamplerState#: Alias for dict. The state dict of any sampler, holding various data needed for MCMC sampling.

deepqmc.types.SamplerFactory = SamplerFactory#: Alias for Callable[[KeyArray, MolecularHamiltonian, Ansatz, list[Molecule], int, int], tuple[MoleculeIdxSampler, MultiNuclearGeometrySampler]]. A factory function that returns a tuple of a molecule index sampler and an electron and nuclei sampler.

Electron samplers#

class deepqmc.sampling.base.ElectronSampler(*args, **kwargs)[source]#

Protocol for ElectronSampler objects.

ElectronSampler objects implement Markov chain samplers for the electron positions. The samplers are assumed to implement a batch of walkers for a single electronic state on a single molecule and may be vmapped to fit the respective context they are used in. Electron samplers can be combined with chain().

init(rng, params, n, R)[source]#

Initializes the sampler state.

Parameters:

rng (KeyArray) – an rng key for the initialization of electron positions.
params (Params) – the parameters of the wave function that is being sampled.
n (int) – the number of walkers to propagate in parallel.
R (Array) – the nuclei positions of the molecular configuration.

Returns:

the sampler state holding electron positions and data about the sampler trajectory.

Return type:

sample(rng, state, params, R)[source]#

Propagates the sampler state.

Parameters:

rng (KeyArray) – an rng key for the proposal of electron positions.
state (SamplerState) – the state of the sampler from the previous step.
params (Params) – the parameters of the wave function that is being sampled.
R (Array) – the nuclei positions of the molecular configuration.

Returns:

the new sampler state, a physical configuration and statistics about the sampling trajectory.

Return type:

tuple[SamplerState, PhysicalConfiguration, Stats]

update(state, params, R)[source]#

Updates the sampler state.

The sampler state is updated to account for changes in the wave function due to a parameter update.

Parameters:

state (SamplerState) – the state of the sampler before parameter update.
params (Params) – the new parameters of the wave function.
R (Array) – the nuclei positions of the molecular configuration.

Returns:

the updated sampler state holding electron positions and data about the sampler trajectory.

Return type:

class deepqmc.sampling.MetropolisSampler(hamil, wf, *, sample_initializer, tau=1.0, target_acceptance=0.57, max_age=None)[source]#

Metropolis–Hastings Monte Carlo sampler.

The sample() method of this class returns electron coordinate samples from the distribution defined by the square of the sampled wave function.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the physical system.
wf (ParametrizedWaveFunction) – the wave function to sample.
sample_initializer (ElectronSampleInitializer) – (~deepqmc.sampling.electron_sample_initializers.ElectronSampleInitializer): callable that generates initial electron positions.
tau (float) – optional, the proposal step size scaling factor. Adjusted during every step if target_acceptance is specified.
target_acceptance (float) – optional, if specified the proposal step size will be scaled such that the ratio of accepted proposal steps approaches target_acceptance.
max_age (int) – optional, if specified the next proposed step will always be accepted for a walker that hasn’t moved in the last max_age steps.

class deepqmc.sampling.LangevinSampler(hamil, wf, *, sample_initializer, tau=1.0, target_acceptance=0.57, max_age=None)[source]#

Metropolis adjusted Langevin Monte Carlo sampler.

Derived from MetropolisSampler.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the physical system.
wf (ParametrizedWaveFunction) – the apply method of the haiku transformed ansatz object.
tau (float) – optional, the proposal step size scaling factor. Adjusted during every step if target_acceptance is specified.
target_acceptance (float) – optional, if specified the proposal step size will be scaled such that the ratio of accepted proposal steps approaches target_acceptance.
max_age (int) – optional, if specified the next proposed step will always be accepted for a walker that hasn’t moved in the last max_age steps.

class deepqmc.sampling.electron_samplers.OppositeSpinExchangeSampler(*, exchange_step_probability, up_logits_fn=None, down_logits_fn=None)[source]#

Add spin swapping steps into chained samplers.

This sampler proposes moves based on swapping the positions of a random pair of spin-up and spin-down electrons. This generally helps to equilibrate the spin of subsystems, when separated by a low probability region in space.

To control the frequency of spin swap proposals compared to regular proposals, this class performs an MCMC step with a spin swap proposal with probability exchange_step_probability, and a step with a normal proposal with probability 1 - exchange_step_probability. This leads to a well defined ratio between the two types of proposals when a large number of steps are considered, but can lead to surprising behavior with a single or a few number of sampling steps.

The sampler cannot be used as the last element of a sampler chain.

Parameters:

up_logits_fn (Callable) – function returning weights for spin-up elec swaps
down_logits_fn (Callable) – function returning weights for spin-down elec swaps

class deepqmc.sampling.DecorrSampler(*, length)[source]#

Insert decorrelating steps into chained samplers.

This sampler cannot be used as the last element of a sampler chain.

Parameters:: length (int) – the samples will be taken in every length MCMC step, that is, length \(-1\) decorrelating steps are inserted.

deepqmc.sampling.combine_samplers(samplers, hamil, wf)[source]#

Combine samplers to create more advanced sampling schemes.

Parameters:

samplers (list[ElectronSampler]) – one or more sampler instances to combine.
hamil (MolecularHamiltonian) – the molecular Hamiltonian.
wf (ParametrizedWaveFunction) – the wave function to sample.

Return type:

ElectronSampler

Electron sample initializers#

class deepqmc.sampling.electron_sample_initializers.ElectronSampleInitializer(*args, **kwargs)[source]#

Protocol for electron sample initializers.

These functions should take a nuclear configuration (charges and coordinates), along with the desired number of up and down spin electrons, and return an initial guess for the positions of these electrons. The returned positions will typically be used to initialize the walkers of an MCMC simulation, Consequently they should be not be too far from the equilibrium distribution of electrons for the given nuclear configuration. In other words, equilibrating the MCMC chains initialized with these electron positions should not take too long.

Parameters:

rng (KeyArray) – A random number generator seed.
charges (Array) – The atomic number of the nuclei.
ns_valence (Array) – The number of valence electrons for each atom. Without ECPs this equals the atomic number for each atom.
nuclear_coordinates (Array) – The nuclear coordinates.
n_up (int) – The number of spin-up electrons.
n_down (int) – The number of spin-down electrons.

class deepqmc.sampling.electron_sample_initializers.AtomCenteredDistribution(*args, **kwargs)[source]#: Protocol for electron distributions centered on nuclei.

class deepqmc.sampling.electron_sample_initializers.SingleGaussianDistribution(scale)[source]#

Distribution with a single Gaussian around each nucleus.

class deepqmc.sampling.electron_sample_initializers.ShellBasedDistribution(*args, **kwargs)[source]#

Distribution with an atomic shell structure.

sample_from_shell(rng, zeta, shape)[source]#

Draw samples from the distribution \(r \sim \exp(-2 \zeta ||r||)\).

Return type:: jax.Array

class deepqmc.sampling.electron_sample_initializers.AtomCenteredElectronInitializer(atom_centered_distribution)[source]#

Electron initializer that places electrons around nuclei.

Nuclei samplers#

class deepqmc.sampling.base.NucleiSampler(*args, **kwargs)[source]#

Protocol for nuclear geometry samplers.

NucleiSampler objects implement samplers for the nuclear coordinates, used during transferable training across multiple molecular geometries. The interface mirrors ElectronSampler but operates on nuclear positions rather than electron positions. Nuclei samplers are not using energy based accept and reject criteria.

init(nuc_coords)[source]#

Initialize the nuclear sampler state.

Parameters:: nuc_coords (Array) – initial nuclear coordinates of shape (n_nuc, 3).
Returns:: the initial sampler state.
Return type:: SamplerState

sample(rng, state)[source]#

Propose a new set of nuclear coordinates.

Parameters:

rng (KeyArray) – an rng key for the coordinate proposal.
state (SamplerState) – the current sampler state.

Returns:

the: updated sampler state, the proposed nuclear coordinates and sampling statistics.

Return type:

tuple[SamplerState, Array, Stats]

class deepqmc.sampling.base.ElectronWarp(*args, **kwargs)[source]#

Protocol for electron warp functions.

An ElectronWarp displaces the electron positions stored inside a sampler state in response to a change in nuclear geometry. Applying a warp before re-equilibrating the sampler avoids large acceptance-rate drops when nuclear coordinates move during optimization or potential energy exploration.

class deepqmc.sampling.nuclei_samplers.IdleNucleiSampler(charges)[source]#

Keeps track of nuclei without updating positions.

Parameters:: nuc_coords (Array) – initial coordinates of the sampled molecules

class deepqmc.sampling.nuclei_samplers.ConstraintNucleiSampler(charges, *, noise_fn=jax.random.normal, coordinate_transform=None, constraints=None)[source]#

Samples nuclear positions around a fixed geometry.

Parameters:

charges (Array) – the nuclear charges of the molecule (\(N_\text{nuc}\)).
noise_fn (Callable | list[Callable]) – a noise distribution (or per-coordinate list of distributions) to sample displacements from. Each callable must have the signature (rng, shape) -> ~jax.Array. Defaults to jax.random.normal().
coordinate_transform (InvertibleCoordinateTransform | None) – (~deepqmc.geom.coordinate_transform.InvertibleCoordinateTransform): optional, an invertible coordinate transform applied before adding noise. Defaults to a plain Cartesian transform.
constraints (list | None) – optional, a list of constraints of the form (idxs_at, idxs_set, fn).

class deepqmc.sampling.nuclei_samplers.PermutationNucleiSampler(charges)[source]#

Nuclei sampler that permutes nuclei with the same atomic number.

Parameters:: charges (Array) – the nuclear charges of the molecule (\(N_\text{nuc}\)).

class deepqmc.sampling.nuclei_samplers.ZMatrixSampler(charges, *, z_matrix_template)[source]#

Nuclei sampler sampling nuclei positions using a Z-matrix.

Parameters:

charges (Array) – the nuclear charges of the molecule (\(N_\text{nuc}\)).
z_matrix_template (StochasticZMatrixTemplate) – the template defining the Z-matrix connectivity and the noise distribution from which new internal coordinates are sampled.

deepqmc.sampling.nuclei_samplers.no_elec_warp(rng, R, dR, smpl_state)[source]#

Identity electron warp function.

Leaves the electron positions in smpl_state unchanged when the nuclei move.

Parameters:

rng (KeyArray) – unused, present for interface compatibility.
R (Array) – the new nuclear coordinates.
dR (Array) – the nuclear displacement, i.e. the difference between the new and the previous nuclear coordinates.
smpl_state (SamplerState) – the electron sampler state.

Return type:

deepqmc.sampling.nuclei_samplers.nn_elec_warp(rng, R, dR, smpl_state)[source]#

Nearest neighbor electron warp function.

Displaces each electron by the same displacement as its nearest nucleus, so that electrons remain attached to their nucleus as the nuclear geometry changes.

Parameters:

rng (KeyArray) – unused, present for interface compatibility.
R (Array) – the new nuclear coordinates.
dR (Array) – the nuclear displacement, i.e. the difference between the new and the previous nuclear coordinates.
smpl_state (SamplerState) – the electron sampler state.

Return type:

deepqmc.sampling.nuclei_samplers.fn_elec_warp(rng, R, dR, smpl_state, fn)[source]#

Electron warp function using a user-defined distance scaling function.

Displaces each electron by a weighted average of all nuclear displacements, with weights given by fn applied to the electron-nucleus distances.

Parameters:

rng (KeyArray) – unused, present for interface compatibility.
R (Array) – the new nuclear coordinates.
dR (Array) – the nuclear displacement, i.e. the difference between the new and the previous nuclear coordinates.
smpl_state (SamplerState) – the electron sampler state.
fn (Callable[[Array], Array]) – a function applied elementwise to the electron-nucleus distances to obtain the weights.

Return type:

tuple[MoleculeIdxSampler, MultiNuclearGeometrySampler]

Multi state and multi geometry samplers#

class deepqmc.sampling.MoleculeIdxSampler(rng, n_mols, batch_size, shuffle=False)[source]#

Sample molecule indexes for transferable training.

class deepqmc.sampling.combined_samplers.MultiElectronicStateSampler(sampler, n_state)[source]#

Sample from multiple electronic states in parallel.

This sampler applies vmap to an underlying ElectronSampler to sample from multiple electronic states in parallel.

Parameters:

sampler (ElectronSampler) – the electron sampler to use.
n_state (int) – the number of electronic states to sample from.

class deepqmc.sampling.MultiNuclearGeometrySampler(elec_sampler, nuc_sampler, warp_elec_fn, update_nuc_period, elec_equilibration_steps)[source]#

This sampler samples from multiple nuclear geometries in parallel.

Parameters:

elec_sampler (MultiElectronicStateSampler) – the electronic sampler to use
nuc_sampler (NucleiSampler) – the nuclei sampler to use.
warp_elec_fn (ElectronWarp) – the function that warps the electrons to the new nuclear geometry.
update_nuc_period (int) – optional, the number of steps between nuclear updates.
elec_equilibration_steps (int) – optional, the number of steps to equilibrate the electronic state after a nuclear update.

Setting up and equilibrating sampling#

deepqmc.sampling.initialize_sampling(rng, hamil, ansatz, mols, electronic_states, molecule_batch_size, *, elec_sampler, nuc_sampler=None, elec_warp_fn=None, update_nuc_period=None, elec_equilibration_steps=None)[source]#

Assemble the molecule-index and combined electron/nuclear samplers.

This is the function typically passed (as a SamplerFactory, partially applied with the sampler-specific keyword arguments) as the sampler_factory argument of train().

Parameters:

rng (KeyArray) – key used for PRNG.
hamil (MolecularHamiltonian) – the molecular Hamiltonian.
ansatz (Ansatz) – the wave function ansatz.
mols (list[Molecule]) – the molecules to sample from.
electronic_states (int) – the number of electronic states to sample.
molecule_batch_size (int) – the number of molecules to sample in each step.
elec_sampler (Callable) – a partially applied ElectronSampler, missing only the hamil and wf arguments, e.g. as created by combine_samplers().
nuc_sampler (Callable) – optional, a partially applied NucleiSampler, missing only the charges argument. Defaults to IdleNucleiSampler, i.e. fixed nuclear geometries.
elec_warp_fn (Callable) – optional, a ElectronWarp, used to move electrons along with the nuclei when the nuclear geometry is updated. Defaults to no_elec_warp().
update_nuc_period (int) – optional, the number of steps between nuclear geometry updates.
elec_equilibration_steps (int) – optional, the number of electron sampling steps to take between two nuclear geometry updates.

Returns:

tuple[~deepqmc.sampling.MoleculeIdxSampler, ~deepqmc.sampling.MultiNuclearGeometrySampler]: the molecule-index sampler and the combined electron/nuclear sampler.

Return type:

deepqmc.sampling.initialize_sampler_state(rng, sampler, params, electron_batch_size, nuc_coords)[source]#

Initialize the sampler state, split across the available devices.

Parameters:

rng (KeyArray) – key used for PRNG.
sampler (MultiNuclearGeometrySampler) – the sampler to initialize.
params (Params) – the wave function parameters.
electron_batch_size (int) – the total number of electron walkers to create, split evenly across all devices.
nuc_coords (Array) – the initial nuclear coordinates of the sampled molecule(s).

Returns:

the initialized, device-sharded sampler state.

Return type:

deepqmc.sampling.equilibrate(rng, params, molecule_idx_sampler, sampler, state, criterion, steps, *, block_size, n_blocks=5, allow_early_stopping=True)[source]#

Run MCMC sampling steps until the walkers have equilibrated.

A generator that repeatedly samples the wave function and, if allow_early_stopping is set, stops once criterion has stabilized: once block_size * n_blocks steps have been taken, the first and last block_size-sized blocks of criterion values are compared, and sampling stops as soon as the difference between their means is smaller than the smaller of their two standard deviations.

Parameters:

rng (KeyArray) – key used for PRNG.
params (Params) – the wave function parameters.
molecule_idx_sampler (MoleculeIdxSampler) – samples the indices of the molecules to consider in each step.
sampler (MultiNuclearGeometrySampler) – the sampler to equilibrate.
state (SamplerState) – the initial sampler state.
criterion (Callable) – a function of the sampled PhysicalConfiguration returning a scalar used to assess equilibration.
steps (Iterable[int]) – the step indices to (potentially) run.
block_size (int) – the number of steps in each of the two compared blocks.
n_blocks (int) – the number of blocks worth of samples to buffer before the equilibration criterion is first evaluated.
allow_early_stopping (bool) – if False, run through all of steps regardless of criterion.

Yields:

tuple – the current step, the updated sampler state, the sampled molecule indices, and the sampling statistics of that step.

Nuclear geometry#

This module implements the internal-coordinate machinery underlying nuclear geometry sampling: functions to compute bond lengths, angles and dihedral angles, coordinate transforms between Cartesian and internal coordinates, and Z-matrices. It is used e.g. by ConstraintNucleiSampler and ZMatrixSampler to sample or constrain nuclear positions in a coordinate system other than Cartesian.

deepqmc.geom.distance(i0, i1, coords)[source]#

Compute the Euclidean distance between two atoms.

Parameters:

i0 (int) – index of the first atom.
i1 (int) – index of the second atom.
coords (Array) – Cartesian coordinates of the atoms, of shape (n_atoms, 3).

Returns:

the distance between the atoms at indices i0 and i1.

Return type:

deepqmc.geom.angle(i0, i1, i2, coords)#

Compute the angle between three atoms.

This is the top-level, user-facing angle function. It has a well defined backward-mode gradient (as computed e.g. with jax.grad), and a well defined second derivative computed with forward-on-backward AD (as done e.g. by jax.hessian). If other combinations of jax differentation transformations are used, jax might raise and error, or the derivative might be ill-defined around zero.

Parameters:

i0 (int) – index of the first outer atom.
i1 (int) – index of the vertex atom, at which the angle is formed.
i2 (int) – index of the second outer atom.
coords (Array) – Cartesian coordinates of the atoms, of shape (n_atoms, 3).

Returns:

the angle between 0 and \(\pi\), in radians, spanned at i1 by the bonds to i0 and i2.

Return type:

deepqmc.geom.dihedral(i0, i1, i2, i3, coords)#

Compute the dihedral angle between four atoms.

This is the top-level, user-facing dihedral function. It has a well defined backward-mode gradient (as computed e.g. with jax.grad), and a well defined second derivative computed with forward-on-backward AD (as done e.g. by jax.hessian). If other combinations of jax differentation transformations are used, jax might raise and error, or the derivative might be ill-defined around zero.

Parameters:

i0 (int) – index of the atom defining, together with i1 and i2, the first of the two half-planes.
i1 (int) – index of the second atom, shared by both half-planes.
i2 (int) – index of the third atom, shared by both half-planes.
i3 (int) – index of the atom defining, together with i1 and i2, the second of the two half-planes.
coords (Array) – Cartesian coordinates of the atoms, of shape (n_atoms, 3).

Returns:

the signed dihedral angle around the i1-i2 bond, in radians, in the interval \((-\pi, \pi]\).

Return type:

StochasticZMatrixTemplate

Coordinate transforms#

A coordinate transform maps Cartesian nuclear coordinates to another coordinate representation. Invertible coordinate transforms can additionally map coordinates back to Cartesian space, which is required to use them as the coordinate_transform argument of ConstraintNucleiSampler.

class deepqmc.geom.coordinate_transform.CoordinateTransform(*args, **kwargs)[source]#

Protocol for coordinate transformations.

A CoordinateTransform maps Cartesian nuclear coordinates to another (possibly lower-dimensional) coordinate representation, e.g. a set of internal coordinates. It is used e.g. by ConstraintNucleiSampler to perform sampling steps in a coordinate system other than Cartesian.

from_cartesian(coords)[source]#

Transform Cartesian nuclear coordinates to this representation.

Parameters:: coords (Array) – Cartesian nuclear coordinates, of shape (n_nuc, 3).
Returns:: the coordinates in the target representation, of shape (len(self),).
Return type:: Array

class deepqmc.geom.coordinate_transform.InvertibleCoordinateTransform(*args, **kwargs)[source]#

Protocol for invertible coordinate transformations.

In addition to from_cartesian(), an InvertibleCoordinateTransform can also map coordinates back to Cartesian space. This is required to apply e.g. sampled noise in the transformed coordinate system, as done by ConstraintNucleiSampler.

to_cartesian(coords)[source]#

Transform coordinates in this representation back to Cartesian space.

Parameters:: coords (Array) – coordinates in this transform’s representation, of shape (len(self),).
Returns:: the corresponding Cartesian nuclear coordinates, of shape (n_nuc, 3).
Return type:: Array

class deepqmc.geom.coordinate_transform.CartesianCoordinateTransform(n_atoms)[source]#

Identity coordinate transform operating on flattened Cartesian coordinates.

This is the default coordinate_transform used by ConstraintNucleiSampler when none is specified, i.e. noise is added directly to the Cartesian nuclear coordinates.

Parameters:: n_atoms (int) – the number of atoms (nuclei) whose coordinates are transformed.

class deepqmc.geom.coordinate_transform.ZMatrixCoordinateTransform(zmatrix_template)[source]#

Invertible coordinate transform between Cartesian coordinates and a Z matrix.

The transformed coordinates are the flattened bond lengths, angles and dihedral angles described by zmatrix_template. This transform is typically passed as the coordinate_transform of a ConstraintNucleiSampler, to sample nuclear displacements in terms of bond lengths, angles and dihedrals rather than Cartesian coordinates.

Parameters:: zmatrix_template (ConcreteZMatrixTemplate) – the Z matrix template defining which atoms are connected by the bonds, angles and dihedrals making up the Z matrix.

class deepqmc.geom.coordinate_transform.RedundantInternalCoordinateTransform(internal_coordinates)[source]#

Coordinate transform to a (possibly redundant) set of internal coordinates.

This transform is not invertible: it merely evaluates a user-specified list of internal-coordinate functions, e.g. distance(), angle() or dihedral() partially applied to fixed atom indices, on the Cartesian nuclear coordinates. Unlike a Z matrix, the resulting coordinates need not have a one-to-one correspondence with the Cartesian coordinates and may be redundant.

Parameters:: internal_coordinates (Sequence[Callable[[...], jax.Array]]) – (~collections.abc.Sequence[~collections.abc.Callable[…, ~jax.Array]]): a sequence of functions, each called with the Cartesian nuclear coordinates as the keyword argument coords and returning a single scalar internal coordinate.

class deepqmc.geom.coordinate_transform.GeneralCoordinateTransform(n_coordinate, coordinate_transform_fn)[source]#

Coordinate transform wrapping an arbitrary user-defined transform function.

Parameters:

n_coordinate (int) – the number of coordinates returned by coordinate_transform_fn.
coordinate_transform_fn (Callable[[Array], Array]) – a function mapping Cartesian nuclear coordinates to a coordinate array of length n_coordinate.

class deepqmc.geom.coordinate_transform.SubsetCoordinateTransform(coordinate_idxs)[source]#

Coordinate transform selecting a subset of the flattened Cartesian coordinates.

Parameters:: coordinate_idxs (Sequence[int]) – indices into the flattened (n_nuc * 3) Cartesian coordinate array to select.

Z-matrices#

A Z-matrix specifies a molecular geometry in terms of bond lengths, bond angles and dihedral angles rather than Cartesian coordinates, each defined relative to previously placed atoms. ConcreteZMatrixTemplate defines the connectivity of such a Z-matrix and can be turned into a ZMatrixCoordinateTransform, while StochasticZMatrixTemplate additionally attaches a noise distribution to each entry, allowing new geometries to be sampled directly in internal coordinates (as used by ZMatrixSampler).

class deepqmc.geom.zmatrix.ConcreteZMatrixTemplate(line_templates)[source]#

Template for a concrete Z matrix representation.

The template includes the “connectivity” of the Z matrix, i.e. which atoms form bonds, angles, and dihedrals, but it doesn’t contain information about the values of these bond lengths, angles, and dihedrals.

clean_values(values)#

Wrap angle and dihedral values into their canonical ranges.

Bond angles are wrapped into \([0, \pi]\) and dihedral angles into \((-\pi, \pi]\), flipping the sign of the corresponding dihedral whenever an angle had to be reflected back into range. This keeps values produced e.g. by adding unconstrained Cartesian-like noise to a ZMatrixCoordinateTransform well-defined before they are turned back into a Z matrix.

Parameters:: values (Array) – the flattened bond length, angle and dihedral values to clean.
Returns:: values with the angle and dihedral entries wrapped into their canonical ranges.
Return type:: Array

concretize(values)#

Assign concrete values to the entries of this template.

Parameters:: values (Array) – the flattened bond length, angle and dihedral values, in the order in which they appear when iterating over the template’s lines.
Returns:: the resulting Z matrix (e.g. a ConcreteZMatrix or a StochasticZMatrix, depending on the concrete template type).
Return type:: Z

concretize_from_cartesian(cartesian)#

Assign concrete values to the entries of this template, computed from a Cartesian reference geometry.

Parameters:: cartesian (Array) – Cartesian nuclear coordinates, of shape (n_nuc, 3), used to evaluate the bond lengths, angles and dihedrals of this template.
Returns:: the resulting Z matrix, with values computed from cartesian.
Return type:: Z

classmethod from_simplified_config(line_templates)[source]#

Construct a template from a simplified, config-friendly specification.

This is the constructor typically used from Hydra configs, e.g. to build the zmatrix_template of a ZMatrixCoordinateTransform.

Parameters:

line_templates (Sequence[Any]) –

one entry per atom, in the same order as the nuclear charges. Each entry is either

a sequence (bond_atom_idx, angle_atom_idx, dihedral_atom_idx) of (up to three) atom indices, using None for entries that don’t apply (e.g. the first three atoms, which don’t need a full bond, angle and dihedral); or
a mapping with the key atom_idxs (the same sequence of indices as above) and, optionally, charge (the nuclear charge of the atom, only used for bookkeeping).

Returns:

the resulting Z matrix template.

Return type:

ConcreteZMatrixTemplate

zmatrix_constructor#: alias of ConcreteZMatrix

class deepqmc.geom.zmatrix.ConcreteZMatrix(lines)[source]#

A Z matrix with concrete (fixed) values for its bond lengths, angles and dihedrals.

Instances are typically obtained by calling concretize or concretize_from_cartesian on a ConcreteZMatrixTemplate, rather than constructed directly.

to_cartesian()[source]#

Convert the Z matrix to Cartesian nuclear coordinates.

Returns:: the Cartesian nuclear coordinates, of shape (n_nuc, 3).
Return type:: Array

property value#: Return the values of the Z matrix’s bond lengths, angles and dihedrals.

class deepqmc.geom.zmatrix.StochasticZMatrixTemplate(line_templates)[source]#

Template for a stochastic Z matrix.

The template includes the “connectivity” of the Z matrix, i.e. which atoms form bonds, angles, and dihedrals. Moreover, it includes a recipe for generating distributions for the bond lengths, angles, and dihedrals.

clean_values(values)#

Wrap angle and dihedral values into their canonical ranges.

Parameters:: values (Array) – the flattened bond length, angle and dihedral values to clean.
Returns:: values with the angle and dihedral entries wrapped into their canonical ranges.
Return type:: Array

concretize(values)#

Assign concrete values to the entries of this template.

Parameters:: values (Array) – the flattened bond length, angle and dihedral values, in the order in which they appear when iterating over the template’s lines.
Returns:: the resulting Z matrix (e.g. a ConcreteZMatrix or a StochasticZMatrix, depending on the concrete template type).
Return type:: Z

concretize_from_cartesian(cartesian)#

Assign concrete values to the entries of this template, computed from a Cartesian reference geometry.

Parameters:: cartesian (Array) – Cartesian nuclear coordinates, of shape (n_nuc, 3), used to evaluate the bond lengths, angles and dihedrals of this template.
Returns:: the resulting Z matrix, with values computed from cartesian.
Return type:: Z

classmethod from_simplified_config(lines)[source]#

Construct a template from a simplified, config-friendly specification.

This is the constructor typically used from Hydra configs, e.g. to build the z_matrix_template of a Z-matrix-based nuclei sampler (deepqmc.sampling.nuclei_samplers.ZMatrixSampler).

Parameters:

lines (Sequence[Mapping[str, Any]]) –

(~collections.abc.Sequence[~collections.abc.Mapping[str, typing.Any]]): one entry per atom, in the same order as the nuclear charges. Each entry is a mapping with the keys:

atom_idxs: a sequence (bond_atom_idx, angle_atom_idx, dihedral_atom_idx) of (up to three) atom indices, using None for entries that don’t apply (e.g. the first three atoms, which don’t need a full bond, angle and dihedral).
distribution_factories: a sequence of (up to three) DistributionFactory instances (or None), one per entry of atom_idxs, used to sample the corresponding bond length, angle or dihedral around the value found in the reference geometry.
charge (optional): the nuclear charge of the atom, only used for bookkeeping.

Returns:

the resulting Z matrix template.

Return type:

zmatrix_constructor#: alias of StochasticZMatrix

class deepqmc.geom.zmatrix.StochasticZMatrix(lines)[source]#

Stochastic Z matrix representation.

A Z matrix where the values of the bond lengths, angles, and dihedrals are defined by a distribution.

to_cartesian(rng)[source]#

Sample this Z matrix and convert it to Cartesian nuclear coordinates.

Parameters:: rng (KeyArray) – an rng key for sampling.
Returns:: the Cartesian nuclear coordinates, of shape (n_nuc, 3).
Return type:: Array

Distributions for stochastic Z-matrices#

These classes implement the DistributionFactory protocol, and are used to specify the noise distribution of the individual bond lengths, angles and dihedrals of a StochasticZMatrixTemplate.

class deepqmc.geom.zmatrix.stochastic.DistributionFactory(*args, **kwargs)[source]#

Protocol for distribution factories.

A DistributionFactory is called with the value of a bond length, angle or dihedral found in a reference geometry, and returns a sampler function for a (typically noisy) distribution over that coordinate. This is the extension point used to implement custom noise distributions for entries of a StochasticZMatrixTemplate.

class deepqmc.geom.zmatrix.UniformDistributionFactory(low, high)[source]#

Create uniform distributions over a fixed, absolute interval.

Note that the sampled values do not depend on loc, i.e. the reference value of the coordinate is ignored.

Parameters:

low (float) – the lower bound of the uniform distribution.
high (float) – the upper bound of the uniform distribution.

class deepqmc.geom.zmatrix.CenteredUniformDistributionFactory(low, high)[source]#

Create uniform distributions centered on loc.

Parameters:

low (float) – the offset below loc of the lower bound of the uniform distribution.
high (float) – the offset above loc of the upper bound of the uniform distribution.

class deepqmc.geom.zmatrix.RadiallyUniformDistributionFactory(low, high)[source]#

Create distributions over a fixed, absolute radial interval, with a probability density proportional to the sampled value rather than uniform in the value itself.

This samples r in [low, high] as the radius of a point picked uniformly at random inside an annulus between low and high, which is the appropriate measure e.g. for sampling bond lengths uniformly with respect to the enclosed area. The sampled values do not depend on loc.

Parameters:

low (float) – the lower bound of the sampled radius.
high (float) – the upper bound of the sampled radius.

class deepqmc.geom.zmatrix.CenteredRadiallyUniformDistributionFactory(low, high)[source]#

Like RadiallyUniformDistributionFactory, but centered on loc.

Samples r in [loc - low, loc + high], with a probability density proportional to r rather than uniform in r itself.

Parameters:

low (float) – the offset below loc of the lower bound of the sampled radius.
high (float) – the offset above loc of the upper bound of the sampled radius.

class deepqmc.geom.zmatrix.ClippedNormalDistributionFactory(scale, low=None, high=None)[source]#

Create normal distributions centered on loc, clipped to an absolute range.

Parameters:

scale (float) – the standard deviation of the normal distribution.
low (float | None) – optional, an absolute lower bound the samples are clipped to.
high (float | None) – optional, an absolute upper bound the samples are clipped to.

class deepqmc.geom.zmatrix.ClippedAsymmetricNormalDistributionFactory(low_scale, high_scale, low=None, high=None)[source]#

Create normal distributions centered on loc, with different standard deviations on either side of loc, clipped to an absolute range.

Parameters:

low_scale (float) – the standard deviation used for samples below loc.
high_scale (float) – the standard deviation used for samples above loc.
low (float | None) – optional, an absolute lower bound the samples are clipped to.
high (float | None) – optional, an absolute upper bound the samples are clipped to.

class deepqmc.geom.zmatrix.DeltaDistributionFactory(*args, **kwargs)[source]#

Create a degenerate “distribution” that deterministically returns loc.

Useful to keep a bond length, angle or dihedral fixed at its reference value while other entries of the same Z matrix are sampled stochastically.

Optimizers#

deepqmc.types.OptState = OptState#: Alias for Any. The state object of an optimizer, holding various data needed for optimization.

deepqmc.types.OptimizerFactory = OptimizerFactory#: Alias for Callable[[LossAndGradFunction], Optimizer]. A factory function that returns an Optimizer instance from a loss (and gradient) function.

class deepqmc.optimizer.Optimizer(loss_and_grad_fn)[source]#

Protocol for Optimizer objects.

init(rng, params, batch)[source]#

Initialize the optimizer state.

Parameters:

rng (KeyArray) – the RNG key used to initialize random components the of optimizer state.
params (Params) – the parameters of the wave function ansatz/ansatzes to be optimized during training.
batch (Batch) – a tuple containing a physical configuration, a set of sample weights and auxiliary data.

Returns:

the initial state of the optimizer

Return type:

OptState

step(rng, params, opt_state, batch)[source]#

Perform an optimization step.

Parameters:

rng (KeyArray) – the RNG key for the optimizer update.
params (Params) – the current parameters of the wave function ansatz/ansatzes.
opt_state (OptState) – the current state of the optimizer
batch (Batch) – a tuple containing a physical configuration, a set of sample weights and auxiliary data.

Returns:

tuple[~deepqmc.types.Params, ~deepqmc.types.OptState, ~deepqmc.types.Energy, ~jax.Array | None, ~deepqmc.types.Stats]: the new model parameters, an updated optimizer state, the energies obtained during the evaluation of the loss function, if applicable the wave function ratios obtained during the evaluation of the loss function and further statistics.

Return type:

tuple[Params, OptState, Energy, jax.Array | None, Stats]

class deepqmc.optimizer.NoOptimizer(loss_and_grad_fn)[source]#

Evaluation-only optimizer that freezes the wave function parameters.

Implements the Optimizer protocol without performing any parameter update. The loss function is still evaluated on each step so that energies and wave function statistics are collected, but gradients are discarded and the parameters are returned unchanged. Use this class to run inference with a trained ansatz.

Parameters:: loss_and_grad_fn (LossAndGradFunction) – callable that returns the loss, local energies, and gradients.

class deepqmc.optimizer.OptaxOptimizer(loss_and_grad_fn, *, optax_opt)[source]#

First-order optimizers bafrom the optax module.

Wraps any optax optimizer and handles device-parallel gradient averaging (pmean()) and parameter stacking automatically. Per-step statistics include opt/param_norm, opt/grad_norm, and opt/update_norm.

Parameters:

loss_and_grad_fn (LossAndGradFunction) – callable that returns the loss, local energies, and gradients.
optax_opt – an optax optimizer instance (e.g. optax.adam(learning_rate=1e-3)).

class deepqmc.optimizer.KFACOptimizer(loss_and_grad_fn, *, kfac)[source]#

Second-order optimizer using the KFAC method [Martens15].

Wraps the kfac_jax optimizer and wires up the multi-device infrastructure required by DeepQMC (pmap, pmap_axis_name, batch size extraction).

Parameters:

loss_and_grad_fn (LossAndGradFunction) – callable that returns the loss, local energies, and gradients; passed directly to kfac_jax as value_and_grad_func.
kfac – a partially-initialized kfac_jax optimizer constructor, i.e. a callable that accepts value_and_grad_func and related keyword arguments and returns the optimizer object.

deepqmc.optimizer.merge_states(params, merge_keys)[source]#

Average selected parameters across electronic states.

For each parameter key that contains at least one of the substrings in merge_keys, the parameter tensor is averaged along the state axis (axis 0) and the result is broadcast back so all states share the same values. Parameters whose keys do not match are left unchanged. This is used to enforce weight-sharing across electronic states during training.

Parameters:

params (Params) – parameter pytree; the outermost axis of each leaf is the electronic-state axis.
merge_keys (Optional[tuple[str, ...]]) – substrings used to select which parameter groups to merge; if None no parameters are merged.

Returns:

parameter pytree with the selected leaves: replaced by their state-averaged values.

Return type:

Params

deepqmc.kfacext.batch_size_extractor(batch)[source]#

Compute the batch size for KFAC.

KFAC requires a single batch dimension, we therefore flatten our batches resulting in batch dimensions that are a product of our various (molecule, electron) batch sizes. Note that each parameter receives gradients only from its samples, therefore the electronic state dimension is not included in this product.

Parameters:: batch (Batch) – a tuple containing a physical configuration, a set of sample weights and auxiliary data.
Returns:: the product of the molecule and electron batch size dimensions.
Return type:: int

Schedules#

Learning rate and damping schedules for use with KFACOptimizer or any optax optimizer.

deepqmc.utils.InverseSchedule(init_value, decay_rate)[source]#

Create a schedule that decays inversely proportional to the step count.

Returns a callable computing \(f(n) = \text{init\_value} / (1 + n / \text{decay\_rate})\). Commonly used as the learning rate or damping schedule of KFACOptimizer, e.g. via _target_: deepqmc.utils.InverseSchedule in a hydra config.

Parameters:

init_value (float) – the value of the schedule at step 0.
decay_rate (float) – the number of steps after which the value has decayed to half of init_value.

Returns:

a function mapping the step number to the current schedule value.

Return type:

Callable[[int], float]

deepqmc.utils.ConstantSchedule(value)[source]#

Create a schedule that returns the same value at every step.

Commonly used as the learning rate or damping schedule of KFACOptimizer, e.g. via _target_: deepqmc.utils.ConstantSchedule in a hydra config.

Parameters:: value (float) – the constant value of the schedule.
Returns:: a function mapping the step number to value.
Return type:: Callable[[int], float]

Wave functions#

class deepqmc.types.Psi(sign, log)[source]#

Represent wave function values.

The sign and log of the absolute value of the wave function are stored.

log: jax.Array#: Alias for field number 1

sign: jax.Array#: Alias for field number 0

deepqmc.types.WaveFunction = WaveFunction#: Alias for Callable[[PhysicalConfiguration], Psi]. A wave function that maps a Physical configuration to the (log) value of the wave function.

deepqmc.types.ParametrizedWaveFunction = ParametrizedWaveFunction#: Alias for Callable[[Params, PhysicalConfiguration], Psi]. A wave function that requires model parameters to be provided for its evaluation.

class deepqmc.types.Ansatz(*args, **kwargs)[source]#

Protocol for ansatz objects.

Ansatz objects represent a parametrized wave function Ansatz. New types of Ansatzes should implement this protocol to be compatible with the DeepQMC software suite. It is assumed that Ansatzes take as input a PhysicalConfiguration for a single sample of electron and nuclei configuration. To handle batches of samples, e.g. during training, the Ansatz is vmap-ed automatically by DeepQMC. The apply function of the Ansatz object is a ParametrizedWaveFunction().

apply(params, phys_conf, return_mos=False)[source]#

Evaluate the Ansatz.

Parameters:

params (Params) – the current parameters with which to evaluate the Ansatz.
phys_conf (PhysicalConfiguration) – a single sample on which to evaluate the Ansatz.
return_mos (bool) – whether to return the many-body orbitals instead of the wave function.

Returns:

the value of the wave function.

Return type:

Psi

init(rng, phys_conf)[source]#

Initialize the parameters of the Ansatz.

Parameters:

rng (KeyArray) – the RNG key used to generate the initial parameters.
phys_conf (PhysicalConfiguration) – a dummy input to the network of a single electron and nuclei configuration. The value of this can be anything, only its shape information is read.

Returns:

the initial parameters of the Ansatz.

Return type:

Params

class deepqmc.wf.NeuralNetworkWaveFunction(*args, **kwargs)[source]#

Implements the neural network wave function.

Configuration files to obtain the PauliNet [HermannNC20], FermiNet [PfauPRR20], Psiformer [Glehn23] and LapNet [Li24] architectures are provided. For a detailed description of the implemented architecture see [Schaetzle23].

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
omni_factory (Callable) – creates the omni net.
envelope (ExponentialEnvelopes) – the orbital envelopes.
backflow_op (Callable) – specifies how the backflow is applied to the orbitals.
n_determinants (int) – specifies the number of determinants
full_determinant (bool) – if False, the determinants are factorized into spin-up and spin-down parts.
cusp_electrons (Callable) – constructor of the electronic cusp module.
cusp_nuclei (Callable) – constructor of the nuclear cusp module.
backflow_transform (str) –
describes the backflow transformation. Possible values:
- 'mult': the backflow is a multiplicative factor
- 'add': the backflow is an additive term
- 'both': the backflow consist of a multiplicative factor and an
  additive term
conf_coeff (Callable) – returns a function that combines the determinants to obtain the WF value

Omni-net, envelopes and cusp corrections#

NeuralNetworkWaveFunction combines a GNN embedding with Jastrow, backflow, envelope and electronic/nuclear cusp factors, each configurable via a corresponding factory argument (omni_factory, envelope, cusp_electrons and cusp_nuclei).

class deepqmc.wf.omni.Backflow(n_orbitals, n_determinants, n_backflows, spin, multi_head=True, *, subnet_factory, name='Backflow')[source]#

The deep backflow factor.

Parameters:

n_orbitals (int) – the number of orbitals to compute backflow factors for.
n_determinants (int) – the number of determinants of the ansatz.
n_backflow (int) – the number of independent backflow factors for each orbital.
multi_head (bool) – if True, create separate MLPs for the n_backflow many backflows, otherwise use a single larger MLP for all.
name (str) – the name of this haiku module.

class deepqmc.wf.omni.Jastrow(*, sum_first, subnet_factory, name='Jastrow')[source]#

The deep Jastrow factor.

Parameters:

sum_first (bool) – if True, the electronic embeddings are summed before feeding them to the MLP. Otherwise the MLP is applied separately on each electron embedding, and the outputs are summed, yielding a (quasi) mean-field Jastrow factor.
name (str) – the name of this haiku module.

class deepqmc.wf.omni.NuclearGNNHead(*, one_particle_parameters)[source]#: A GNN head that predicts parameters from nucleus embeddings.

class deepqmc.wf.omni.OmniNet(hamil, n_orb_up, n_orb_down, n_determinants, n_backflows, *, embedding_dim, gnn_factory, jastrow_factory, backflow_factory, nuclear_gnn_head=None, use_attentive_stream=False)[source]#

Combine the GNN, the Jastrow and backflow MLPs.

A GNN is used to create embedding vectors for each electron, which are then fed into the Jastrow and/or backflow MLPs to produce the Jastrow–backflow part of deep QMC Ansatzes.

Parameters:

mol (Molecule) – the molecule to consider.
n_orb_up (int) – the number of spin-up orbitals in a single deterimant, to compute backflow factors for. This is equal to the number of spin-up electrons, except when full determinants are used, in which case it is equal to the total number of electrons.
n_orb_down (int) – the number of spin-down orbitals in a single deterimant, to compute backflow factors for. This is equal to the number of spin-down electrons, except when full determinants are used, in which case it is equal to the total number of electrons.
n_determinants (int) – the number of determinants to use.
n_backflows (int) – the number of independent backflow channels for each orbital, e.g. two channels are necessary if both additive and multiplicative backflows are used.
embedding_dim (int) – the length of the electron embedding vectors.
gnn_factory (Callable) – function that returns a GNN instance.
jastrow_factory (Callable) – function that returns a Jastrow instance.
backflow_factory (Callable) – function that returns a Backflow instance.
use_attentive_stream (bool) – if True, the GNN uses two streams (individual and attentive) as introduced in the LapNet architecture. In this case, the embedding dimension is doubled. After the GNN pass, one half of the features is discarded.

Envelopes#

class deepqmc.wf.env.ExponentialEnvelopes(hamil, n_determinants, *, isotropic, per_shell, per_orbital_exponent, spin_restricted, init_to_ones, softplus_zeta)[source]#: Create exponential envelopes centered on the nuclei.

class deepqmc.wf.env.SimplifiedNucleusDependentEnvelopes(hamil, n_determinants, *, n_envelope_per_nucleus, per_orbital_exponent, fixed_pi)[source]#: Envelopes with no trainable pi parameter.

Cusp corrections#

class deepqmc.wf.cusp.CuspAsymptotic(*, cusp_function, trainable_alpha)[source]#: Base class for nuclear and electronic cusps.

class deepqmc.wf.cusp.DeepQMCCusp[source]#

Compute the DeepQMC cusp factor.

Computes the factor: \(-\frac{\text{scale}}{\sum_{i<j}\alpha * (1 + \alpha r_{ij})}\), where \(r_{ij}\) are the electron-electron or electron-nuclei distances.

class deepqmc.wf.cusp.ElectronicCuspAsymptotic(*, same_scale, anti_scale, alpha=1.0, **kwargs)[source]#

Calculate a multiplicative factor, that implements the electronic cusps.

Parameters:

same_scale (float) – scaling factor to use for same spin electron cusps.
anti_scale (float) – scaling factor to use for anti spin electron cusps.
alpha (float) – default 1, the \(\alpha\) parameter in the above cusp equations.
trainable_alpha (bool) – whether the \(\alpha\) is trainable
cusp_function (Callable) – an instance of either DeepQMCCusp or PsiformerCusp.

class deepqmc.wf.cusp.NuclearCuspAsymptotic(nuclear_charges, *, alpha=1.0, **kwargs)[source]#

Calculate a multiplicative factor, that implements the nuclear cusps.

Parameters:

nuclear_charges (Array) – the array of nuclear charges of the molecule
alpha (float) – default 1, the \(\alpha\) parameter in the above cusp equations.
trainable_alpha (bool) – whether the \(\alpha\) is trainable
cusp_function (Callable) – an instance of either DeepQMCCusp or PsiformerCusp.

class deepqmc.wf.cusp.PsiformerCusp[source]#

Compute the Psiformer cusp factor.

Computes the factor: \(-\frac{\text{scale}}{\sum_{i<j}\alpha^2 * (\alpha + r_{ij})}\), where \(r_{ij}\) are the electron-electron or electron-nuclei distances.

Graph neural networks#

A graph neural network is the most important component of the neural network wave function ansatz. This module implements a general gnn framework, that can be configured to obtain a variety of different ansatzes.

Graphs#

This submodule implements the basic functionality for working with graphs.

deepqmc.gnn.graph.GraphEdgeBuilder(mask_self)[source]#

Create a function that builds graph edges.

Parameters:: mask_self (bool) – whether to mask edges that begin and end in the same node.

deepqmc.gnn.graph.GraphUpdate(aggregate_edges_for_nodes_fn, update_nodes_fn=None, update_edges_fn=None)[source]#

Create a function that updates a graph.

The update function is tailored to be used in GNNs.

Parameters:

aggregate_edges_for_nodes_fn (Callable) – function that aggregates edge features for each node, called before update_nodes_fn.
update_nodes_fn (Callable) – optional, function that updates the nodes.
update_edges_fn (Callable) – optional, function that updates the edges.

deepqmc.gnn.graph.MolecularGraphEdgeBuilder(n_nuc, n_up, n_down, edge_types, *, self_interaction)[source]#

Create a function that builds many types of molecular edges.

Parameters:

n_nuc (int) – number of nuclei.
n_up (int) – number of spin-up electrons.
n_down (int) – number of spin-down electrons.
edge_types (list[str]) –
list of edge type names to build. Possible names are:
- 'nn': nuclei->nuclei edges
- 'ne': nuclei->electrons edges
- 'en': electrons->nuclei edges
- 'same': edges between same-spin electrons
- 'anti': edges between opposite-spin electrons
- 'up': edges going from spin-up electrons to all electrons
- 'down': edges going from spin-down electrons to all electrons
self_interaction (bool) – whether edges between a particle and itself are considered

class deepqmc.gnn.edge_features.CombinedEdgeFeature(*, features)[source]#

Combine multiple edge features.

Parameters:: features (list[EdgeFeature]) – a list of edge feature objects to combine.

class deepqmc.gnn.edge_features.DifferenceEdgeFeature(*, log_rescale=False)[source]#

Return the difference vector as the edge features.

Parameters:: log_rescale (bool) – whether to rescale the features by \(\log(1 + d) / d\) where \(d\) is the length of the edge.

class deepqmc.gnn.edge_features.DistancePowerEdgeFeature(*, powers, eps=None, log_rescale=False)[source]#

Return powers of the distance as edge features.

Parameters:

powers (list[float]) – a list of powers to apply to the edge length.
eps (float | None) – a small value to add to the denominator when the power is negative.
log_rescale (bool) – whether to rescale the features by \(\log(1 + d) / d\) where \(d\) is the length of the edge.

class deepqmc.gnn.edge_features.EdgeFeature(*args, **kwargs)[source]#: Base class for all edge features.

class deepqmc.gnn.edge_features.GaussianEdgeFeature(*, n_gaussian, radius, offset)[source]#

Expand the distance in a Gaussian basis.

Parameters:

n_gaussian (int) – the number of gaussians to use, consequently the length of the feature vector
radius (float) – the radius within which to place gaussians
offset (bool) – whether to offset the position of the first Gaussian from zero.

Electron GNN#

This submodule provides the ElectronGNN architecture for defining neural network parametrized functions acting on graphs of electrons and nuclei.

class deepqmc.gnn.electron_gnn.ElectronEmbedding(n_nuc, n_up, n_down, embedding_dim, n_elec_types, elec_types, *, positional_embeddings, use_spin, project_to_embedding_dim)[source]#

Create initial embeddings for electrons.

Parameters:

n_nuc (int) – the number of nuclei.
n_up (int) – the number of spin up electrons.
n_down (int) – the number of spin down electrons.
embedding_dim (int) – the desired length of the embedding vectors.
n_elec_types (int) –
the number of electron types to differentiate. Usual values are:
- 1: treat all electrons as indistinguishable. Note that electrons
  with different spins can still become distinguishable during the later embedding update steps of the GNN.
- 2: treat spin up and spin down electrons as distinguishable already
  in the initial embeddings.
elec_types (Array) – an integer array with length equal to the number of electrons, with entries between 0 and n_elec_types. Specifies the type for each electron.
positional_embeddings (dict) – optional, if not None, a dict with edge types as keys, and edge features as values. Specifies the edge types and edge features to use when constructing the positional initial electron embeddings.
use_spin (bool) – only relevant if positional_embeddings is not False, if True, concatenate the spin of the given electron after the positional embedding features.
project_to_embedding_dim (bool) – only relevant if positional_embeddings is not False, if True, use a linear layer to project the initial embeddings to have length embedding_dim.

class deepqmc.gnn.electron_gnn.ElectronGNN(hamil, embedding_dim, *, n_interactions, edge_features, self_interaction, two_particle_stream_dim, nuclei_embedding, electron_embedding, layer_factory, ghost_coords=None)[source]#

A neural network acting on graphs defined by electrons and nuclei.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system on which the graph is defined.
embedding_dim (int) – the length of the electron embedding vectors.
n_interactions (int) – number of message passing interactions.
edge_features (dict) –
a dict of functions for each edge type, embedding the interparticle differences. Valid keys are:
- 'ne': for nucleus-electron edges
- 'nn': for nucleus-nucleus edges
- 'same': for same spin electron-electron edges
- 'anti': for opposite spin electron-electron edges
- 'up': for edges going from spin up electrons to all electrons
- 'down': for edges going from spin down electrons to all electrons
self_interaction (bool) – whether to consider edges where the sender and receiver electrons are the same.
two_particle_stream_dim (int) – the feature dimension of the two particle streams. Only active if deep_features are used.
nuclei_embedding (type[NucleiEmbedding]) – optional, the instance responsible for creating the initial nuclear embeddings. Set to None if nuclear embeddings are not needed.
electron_embedding (type[ElectronEmbedding]) – the instance that creates the initial electron embeddings.
layer_factory (type[ElectronGNNLayer]) – a callable that generates a layer of the GNN.
ghost_coords (Array) – optional, specifies the coordinates of one or more ghost atoms, useful for breaking spatial symmetries of the nuclear geometry.

edge_factory(phys_conf)[source]#: Compute all the graph edges used in the GNN.

class deepqmc.gnn.electron_gnn.ElectronGNNLayer(n_interactions, ilayer, n_nuc, n_up, n_down, embedding_dim, edge_types, self_interaction, node_data, two_particle_stream_dim, *, electron_residual, nucleus_residual, two_particle_residual, deep_features, update_features, update_rule, subnet_factory=None, subnet_factory_by_lbl=None)[source]#

The message passing layer of ElectronGNN.

Implements a message passing layer for the ElectronGNN architecture.

Parameters:

n_interactions (int) – the number of message passing interactions.
ilayer (int) – the index of this layer (0 <= ilayer < n_interactions).
n_nuc (int) – the number of nuclei.
n_up (int) – the number of spin up electrons.
n_down (int) – the number of spin down electrons.
embedding_dim (int) – the length of the electron embedding vectors.
edge_types (tuple[str]) – the types of edges to consider.
self_interaction (bool) – whether to consider edges where the sender and receiver electrons are the same.
node_data (dict[str, Any]) – a dictionary containing information about the nodes of the graph.
two_particle_stream_dim (int) – the feature dimension of the two particle streams.
electron_residual – whether a residual connection is used when updating the electron embeddings, either False, or an instance of ResidualConnection.
nucleus_residual – whether a residual connection is used when updating the nucleus embeddings, either False, or an instance of ResidualConnection.
two_particle_residual – whether a residual connection is used when updating the two particle embeddings, either False, or an instance of ResidualConnection.
deep_features – if False, the edge features are not updated throughout the GNN layers, if shared than in each layer a single MLP (u) is used to update all edge types, if separate then in each layer separate MLPs are used to update the different edge types.
update_features (list[UpdateFeature]) – a list of partially initialized update feature classes to use when computing the update features of the one particle embeddings. For more details see the documentation of update_features.
update_rule (str) –
how to combine the update features for the update of the one particle embeddings. Possible values:
- 'concatenate': run concatenated features through MLP
- 'featurewise': apply different MLP to each feature channel and sum
- 'featurewise_shared': apply the same MLP across feature channels
- 'sum': sum features before sending through an MLP
note that 'sum' and 'featurewise_shared' imply features of same size.
subnet_factory (Callable) – optional, a function that constructs the subnetworks of the GNN layer.
subnet_factory_by_lbl (dict) – optional, a dictionary of functions that construct subnetworks of the GNN layer. If both this and subnet_factory is specified, the specified values of subnet_factory_by_lbl will take precedence. If some keys are missing, the default value of subnet_factory will be used in their place. Possible keys are: (w, h, g or u).

class deepqmc.gnn.electron_gnn.NucleiEmbedding(n_up, n_down, charges, n_atom_types, *, embedding_dim, atom_type_embedding, subnet_type, edge_features)[source]#

Create initial embeddings for nuclei.

Parameters:

n_up (int) – the number of spin up electrons.
n_down (int) – the number of spin down electrons.
charges (Array) – the nuclear charges of the molecule.
n_atom_types (int) – the number of different atom types in the molecule.
embedding_dim (int) – the length of the output embedding vector
atom_type_embedding (bool) – if True, initial embeddings are the same for atoms of the same type (nuclear charge), otherwise they are different for all nuclei.
subnet_type (str) – the type of subnetwork to use for the embedding generation: - 'mlp': an MLP is used - 'embed': a haiku.Embed block is used
edge_features (EdgeFeature) – optional, the edge features to use when constructing the initial nuclear embeddings.

Update Features#

This submodule implements some common ways to compute update features for the node embeddings from the current node and edge embeddings. Instances of the below classes are callable, they take as input the current node and edge representations, and output a list of update features to be used for updating the node representations.

class deepqmc.gnn.update_features.ConvolutionElectronUpdateFeature(*args, edge_types, normalize, w_factory, h_factory, w_for_ne=True)[source]#

The convolution of node and edge embeddings as an update feature.

Returns the convolution of the node and edge embeddings for various edge types as separate update features.

Parameters:

n_up (int) – number of spin up electrons
n_down (int) – number of spin down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream
node_edge_mapping (NodeEdgeMapping) – mapping between the various node and edge types.
edge_types (list[str]) – list of edge types to sum over
normalize (bool) – whether to normalize the sum by the number of senders
w_factory (Callable) – factory function for the \(w\) matrix
h_factory (Callable) – factory function for the \(h\) matrix
w_for_ne (bool) – whether to use the \(w\) matrix for the \(ne\) edge type

class deepqmc.gnn.update_features.EdgeSumElectronUpdateFeature(*args, edge_types, normalize)[source]#

The (normalized) sum of the edge embeddings as an update feature.

Returns the (normalized) sum of the edge embeddings for various edge types as separate update features.

Parameters:

n_up (int) – number of spin up electrons
n_down (int) – number of spin down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream
node_edge_mapping (NodeEdgeMapping) – mapping between the various node and edge types.
edge_types (list[str]) – list of edge types to sum over
normalize (bool) – whether to normalize the sum by the number of senders

class deepqmc.gnn.update_features.NodeAttentionElectronUpdateFeature(*args, num_heads, mlp_factory, attention_residual, mlp_residual)[source]#

Create a single update feature by attenting over the nodes.

Returns the Psiformer update feature based on attention over the nodes.

Parameters:

n_up (int) – number of spin up electrons
n_down (int) – number of spin down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream
node_edge_mapping (NodeEdgeMapping) – mapping between the various node and edge types.
num_heads (int) – number of attention heads
mlp_factory (Type[MLP]) – factory function for the MLP
attention_residual (Optional[Residual]) – optional residual connection after the attention layer
mlp_residual (Optional[Residual]) – optional residual connection after the MLP layer

class deepqmc.gnn.update_features.NodeSumElectronUpdateFeature(*args, node_types, normalize)[source]#

The (normalized) sum of the node embeddings as an update feature.

Returns the (normalized) sum of the electron embeddings from the previous layer as a single update feature.

Parameters:

n_up (int) – number of spin up electrons
n_down (int) – number of spin down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream
node_edge_mapping (NodeEdgeMapping) – mapping between the various node and edge types.
node_types (list[str]) – list of node types to update
normalize (bool) – whether to normalize the sum by the number of nodes

class deepqmc.gnn.update_features.ResidualElectronUpdateFeature(n_up, n_down, two_particle_stream_dim, node_edge_mapping, is_last_layer=False)[source]#

Residual update feature.

Returns the unchanged electron embeddings from the previous layer as a single update feature.

class deepqmc.gnn.update_features.SparseDerivativeNodeAttentionElectronUpdateFeature(*args, num_heads, indiv_mlp_factory, attn_mlp_factory, attention_residual, attn_mlp_residual, indiv_mlp_residual)[source]#

LapNet-style dual-stream attention block as an electron update feature.

Implements one Transformer-like layer of the LapNet ansatz. The electron embedding is the concatenation of two streams of equal width, split along the feature axis:

the individual stream g — a per-electron embedding whose forward-Laplacian Jacobian stays sparse. Used as the queries and keys for attention.

the attentive stream h — a per-electron embedding whose Jacobian densifies once attention mixes contributions from all electrons. Used as the values for attention.

The attentive stream is updated by sparse multi-head attention (g``→Q/K, ``h``→V) followed by an MLP, both with optional residual connections. The individual stream is updated by its own MLP, except on the last layer (``is_last_layer), where it is zeroed out since only the attentive stream feeds the downstream determinant. The two updated streams are concatenated back into a single array on output.

Parameters:

n_up (int) – number of spin-up electrons
n_down (int) – number of spin-down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream, unused here
node_edge_mapping (NodeEdgeMapping) – node/edge mapping, unused here
num_heads (int) – number of attention heads. num_heads must divide the attentive embedding dimension; D_per_head = h.shape[-1] // num_heads.
indiv_mlp_factory (Type[MLP]) – factory for the MLP applied to the individual stream
attn_mlp_factory (Type[MLP]) – factory for the MLP applied to the attention output
attention_residual (Optional[ResidualConnection]) – optional residual connection wrapping the attention call
attn_mlp_residual (Optional[ResidualConnection]) – optional residual connection wrapping attn_mlp
indiv_mlp_residual (Optional[ResidualConnection]) – optional residual connection wrapping indiv_mlp

class deepqmc.gnn.update_features.UpdateFeature(n_up, n_down, two_particle_stream_dim, node_edge_mapping, is_last_layer=False)[source]#

Base class for all update features.

Parameters:

n_up (int) – number of spin up electrons
n_down (int) – number of spin down electrons
two_particle_stream_dim (int) – dimension of the two-particle stream
node_edge_mapping (NodeEdgeMapping) – mapping between the various node and edge types.

class deepqmc.gnn.utils.NodeEdgeMapping(edges, node_data=None)[source]#

A utility mapping between the various node and edge types.

For example it is often useful determine the sender/receiver node type of a given edge, or generate all the edge types with a given sender/receiver node.

Parameters:

edges (Sequence[str]) – all the edge types present in the graph
node_data (dict) – optional, data to store for the node types

Haiku#

Some additional neural network functionality is implemented in the package and documented here.

class deepqmc.hkext.GLU(out_dim, name=None, *, bias=True, layer_norm_before=True, activation=<PjitFunction of <function sigmoid>>, b_init=<function zeros>)[source]#

Gated Linear Unit.

Parameters:

out_dim (int) – the output dimension.
name (str) – optional, the name of the network.
bias (bool) – optional, whether to include a bias term.
layer_norm_before (bool) – optional, whether to apply layer normalization before the GLU operation.
activation (Callable) – default is sigmoid, the activation function.
b_init (Callable) – default is zeros, the initialization function for the bias term.

class deepqmc.hkext.Identity(*args, **kwargs)[source]#: Represent the identity operation.

class deepqmc.hkext.MLP(out_dim, name=None, *, hidden_layers, bias, last_linear, activation, init)[source]#

Represent a multilayer perceptron.

Parameters:

out_dim (int) – the output dimension.
name (str) – optional, the name of the network.
hidden_layers (tuple) – optional, either (‘log’, \(N_\text{layers}\)), in which case the network will have \(N_\text{layers}\) layers with logarithmically changing widths, or a tuple of ints specifying the width of each layer.
bias (bool | str) –
optional, specifies which layers should have a bias term. Possible values are
- True: all layers will have a bias term
- False: no layers will have a bias term
- 'not_last': all but the last layer will have a bias term
last_linear (bool) – optional, if True the activation function is not applied to the activation of the last layer.
activation (Callable) – optional, the activation function.
init (str | Callable) –
optional, specifies the initialization of the linear weights. Possible string values are:
- 'default': the default haiku initialization method is used.
- 'ferminet': the initialization method of the ferminet
  package is used.
- 'deeperwin': the initialization method of the deeperwin
  package is used.

class deepqmc.hkext.ResidualConnection(*, normalize)[source]#

Represent a residual connection between pytrees.

The residual connection is only added if inp and update have the same shape.

Parameters:: normalize (-) – if True the sum of inp and update is normalized with sqrt(2).

class deepqmc.hkext.SparseMultiHeadAttention(num_heads, key_size, w_init_scale=None, *, w_init=None, with_bias=True, b_init=None, value_size=None, model_size=None, name=None)[source]#

Drop-in subclass of hk.MultiHeadAttention.

Identical interface to the parent: just override __call__ to route the scaled-dot-product through sparse_attention, so folx picks up the wide-scope rule. Q/K/V projections and the output projection use Haiku’s standard hk.Linear (folx’s default rule handles those fine, since they’re per-electron Linears that preserve weak structure).

class deepqmc.hkext.SumPool(out_dim, name=None)[source]#

Represent a global sum pooling operation.

Parameters:

out_dim (int) – the output dimension.
name (str) – optional, the name of the network.

deepqmc.hkext.ssp(x)[source]#

Compute the shifted softplus activation function.

Computes the elementwise function \(\text{softplus}(x)=\log(1+\text{e}^x)+\log\frac{1}{2}\)

Return type:: Array

Observables#

class deepqmc.observable.ObservableMonitor(save_samples, period)[source]#

Base class for observable monitors evaluated during training or inference.

An ObservableMonitor encapsulates a physical observable (e.g. forces, spin) that is computed periodically from wave function samples. The lifecycle has two stages:

Construction — sets the sampling frequency and whether raw samples should be stored alongside the statistics.
Finalization — finalize() is called once the Hamiltonian and wave function are known; subclasses override it to build observable_fn. After finalization the monitor is ready to be called.

Subclasses must set the class attribute name and override finalize() to populate observable_fn. Set requires_energy to True when the observable function needs the local energies as an additional input.

Parameters:

save_samples (bool) – if True, the raw per-sample observable values are included in the returned statistics dictionary under the key '<name>/samples'.
period (int) – number of training steps between consecutive evaluations; must be at least 1.

finalize(hamil, wf)[source]#

Bind the monitor to a specific Hamiltonian and wave function.

Called once before training begins. The default implementation returns self unchanged; subclasses override this to construct observable_fn from hamil and wf.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the physical system.
wf (ParametrizedWaveFunction) – the parametrized wave function used during training.

Returns:

the finalized monitor (self).

Return type:

ObservableMonitor

Property monitors#

class deepqmc.observable.EnergyMonitor(save_samples, period)[source]#

Monitor the local energies during the calculation.

class deepqmc.observable.SpinMonitor(save_samples, period)[source]#

Monitor the total spin expectation value \(\langle S^2 \rangle\).

class deepqmc.observable.WaveFunctionMonitor(save_samples, period)[source]#

Monitor the wave function during the calculation.

class deepqmc.observable.PsiRatioMonitor(save_samples, period)[source]#

Monitor wave function ratios between electronic states.

class deepqmc.observable.ElectronPositionMonitor(save_samples, period)[source]#

Monitor the electron positions during training.

class deepqmc.observable.NuclearPositionMonitor(save_samples, period)[source]#

Monitor the nuclear positions during training.

class deepqmc.observable.OscillatorStrengthMonitor(save_samples, period)[source]#

Monitor oscillator strengths between electronic states.

Oscillator strength#

deepqmc.oscillator_strength.compute_oscillator_strength(local_energies, ratios, rs, local_energies_mask=None, ratios_mask=None)[source]#

Compute the oscillator strength and its error for a batch of samples.

Estimates the oscillator strength, transition dipole moment, and excitation energy between all pairs of electronic states from a batch of local energies and wave function ratios, together with their statistical errors. This is a batch postprocessing counterpart of OscillatorStrengthMonitor, useful for recomputing oscillator strengths from samples gathered after training, e.g. via read_and_convert_result().

Parameters:

local_energies (Array) – the electron batch of local energies, shape: [electronic_states, electron_batch_size].
ratios (Array) – the electron batch of wave function ratios, shape: [electronic_states, electronic_states, electron_batch_size].
rs (Array) – the electron batch of electron samples, shape: [electronic_states, electron_batch_size, n_electrons, 3].
local_energies_mask (Array) – optional, a boolean mask selecting the valid entries of local_energies.
ratios_mask (Array) – optional, a boolean mask selecting the valid entries of ratios.

Returns:

a tuple of (oscillator_strength, error), (transition_dipole_moment, error) and (excitation_energy, error) pairs, each array of shape [electronic_states, electronic_states].

Return type:

tuple[tuple[Array, Array], tuple[Array, Array], tuple[Array, Array]]

Force monitors#

class deepqmc.observable.BareForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor bare Hellmann-Feynman forces without variance reduction.

class deepqmc.observable.BareForceAntiMonitor(save_samples, period, coordinate_transform=None, cutoff=0.3)[source]#

Monitor bare Hellmann-Feynman forces with antithetic-sampling variance reduction.

class deepqmc.observable.ACZVForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZV estimator [Assaraf03].

class deepqmc.observable.ACZVZBForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZV-ZB estimator [Assaraf03].

class deepqmc.observable.ACZBForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZB estimator [Assaraf03].

class deepqmc.observable.ACZVQForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZVQ estimator [Assaraf03]; incompatible with ECPs.

class deepqmc.observable.ACZVQForceAntiMonitor(save_samples, period, coordinate_transform=None, cutoff=0.3)[source]#

Monitor HF forces using AC-ZVQ with antithetic sampling; incompatible with ECPs.

class deepqmc.observable.ACZVZBQForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZV-ZB-Q estimator [Assaraf03]; incompatible with ECPs.

class deepqmc.observable.ACZVQZBForceMonitor(save_samples, period, coordinate_transform=None)[source]#

Monitor HF forces using the AC-ZVQ-ZB hybrid estimator; incompatible with ECPs.

class deepqmc.observable.FiniteDifferenceForceMonitor(save_samples, period, h=0.001)[source]#

Monitor interatomic forces via a finite-difference scheme.

Force estimators#

The force monitors above wrap the following lower-level estimator functions, which construct the Hellmann-Feynman and finite-difference force evaluators used during training and evaluation. They can also be called directly for postprocessing analysis.

deepqmc.force.evaluate_hf_force_bare(hamil, wf, coordinate_transform=None)[source]#

Construct the bare estimator of the Hellmann-Feynman force.

The bare estimator is the direct gradient of the local energy with respect to the nuclear coordinates, without any variance-reduction terms. It therefore has the largest variance among the estimators implemented in this module, but is also the cheapest to evaluate. If the Hamiltonian uses a Gaussian-type effective core potential, the non-local ECP contribution to the force is added as well.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (rng, params, phys_conf) -> jax.Array that evaluates the bare force for a batch of samples.

Return type:

deepqmc.force.evaluate_hf_force_ac_zv(hamil, wf, coordinate_transform=None)[source]#

Construct the AC-ZV (zero-variance) Hellmann-Feynman force estimator.

Adds a zero-variance (ZV) correction term to evaluate_hf_force_bare(), computed via a JVP of the local energy through the nuclear coordinates [Tiihonen21]. This reduces the variance of the estimator compared to the bare estimator, at the cost of an additional local-energy-gradient evaluation.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (rng, params, phys_conf, e_loc=None, energy=None) -> jax.Array that evaluates the AC-ZV force for a batch of samples. If e_loc is not provided it is computed internally; energy is accepted for interface uniformity with the other estimators but is not used.

Return type:

deepqmc.force.evaluate_hf_force_ac_zvzb(hamil, wf, coordinate_transform=None)[source]#

Construct the AC-ZVZB (zero-variance zero-bias) Hellmann-Feynman force estimator.

Adds both the zero-variance (ZV) correction of evaluate_hf_force_ac_zv() and a zero-bias (ZB) correction to evaluate_hf_force_bare(). The ZB term corrects for the bias introduced by using finite Monte Carlo samples of a wave function that does not exactly satisfy the Schrödinger equation [Tiihonen21].

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (rng, params, phys_conf, e_loc, energy) -> jax.Array that evaluates the AC-ZVZB force for a batch of samples, given the local energies e_loc and the mean energy energy of the batch.

Return type:

deepqmc.force.evaluate_hf_force_ac_zb(hamil, wf, coordinate_transform=None)[source]#

Construct the AC-ZB (zero-bias) Hellmann-Feynman force estimator.

Adds only the zero-bias (ZB) correction term to evaluate_hf_force_bare(), without the zero-variance (ZV) term of evaluate_hf_force_ac_zv() [Tiihonen21]. Cheaper to evaluate than evaluate_hf_force_ac_zvzb(), but with a higher variance since it lacks the ZV correction.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (rng, params, phys_conf, e_loc, energy) -> jax.Array that evaluates the AC-ZB force for a batch of samples, given the local energies e_loc and the mean energy energy of the batch.

Return type:

deepqmc.force.evaluate_hf_force_ac_zvq(hamil, wf, coordinate_transform=None)[source]#

Construct the AC-ZVQ (zero-variance, closed-form) Hellmann-Feynman force estimator.

Combines the bare nuclear and electronic force with a zero-variance correction expressed in closed form through the auxiliary function Q, following [Assaraf03]. Unlike evaluate_hf_force_ac_zv(), this estimator does not require an rng key or an extra local-energy evaluation, but it is not compatible with effective core potentials.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (params, phys_conf) -> jax.Array that evaluates the AC-ZVQ force for a batch of samples.

Return type:

deepqmc.force.evaluate_hf_force_ac_zvzbq(hamil, wf, coordinate_transform=None)[source]#

Construct the AC-ZVZBQ (zero-variance zero-bias, closed-form) force estimator.

Adds a zero-bias (ZB) correction, expressed via the auxiliary function Q, to evaluate_hf_force_ac_zvq() [Assaraf03]. Like the ZVQ estimator, this does not require an rng key, but it is not compatible with effective core potentials.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (params, phys_conf, e_loc, energy) -> jax.Array that evaluates the AC-ZVZBQ force for a batch of samples, given the local energies e_loc and the mean energy energy of the batch.

Return type:

deepqmc.force.evaluate_hf_force_ac_zvqzb(hamil, wf, coordinate_transform=None)[source]#

Construct the hybrid AC-ZVQ + ZB Hellmann-Feynman force estimator.

Combines the closed-form zero-variance correction of evaluate_hf_force_ac_zvq() with a zero-bias correction computed via a general autodiff gradient of the log wave function with respect to the nuclear coordinates, instead of the closed-form Q function used in evaluate_hf_force_ac_zvzbq() [Assaraf03]. Not compatible with effective core potentials.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
coordinate_transform (InvertibleCoordinateTransform) – optional, the coordinate system in which the force is expressed. Defaults to Cartesian nuclear coordinates.

Returns:

a function of signature (params, phys_conf, e_loc, energy) -> jax.Array that evaluates the hybrid AC-ZVQ + ZB force for a batch of samples, given the local energies e_loc and the mean energy energy of the batch.

Return type:

deepqmc.force.evaluate_finite_difference_force(hamil, wf, step_size)[source]#

Construct a finite-difference estimator of the interatomic force.

Displaces each nuclear coordinate by \(\pm\) step_size (electron positions are co-displaced to follow the nearest nucleus) and estimates the force from the resulting change in the local energy, importance-weighted by the ratio of wave function values. Unlike the Hellmann-Feynman estimators in this module, this estimator does not require differentiating through the Hamiltonian, but its cost scales with the number of nuclear degrees of freedom, and its accuracy is limited by the finite step size.

Parameters:

hamil (MolecularHamiltonian) – the Hamiltonian of the system.
wf (ParametrizedWaveFunction) – the parametrized wave function.
step_size (float) – the finite-difference step size, in Cartesian nuclear coordinates.

Returns:

a function of signature (rng, params, phys_conf, e_loc, energy) -> jax.Array that evaluates the finite-difference force for a batch of samples, given the local energies e_loc. energy is accepted for interface uniformity with the other estimators but is not used.

Return type:

deepqmc.force.antithetic_wrapper(evaluate_force, wf, r_cut)[source]#

Wrap a force estimator with antithetic-sampling variance reduction.

For each sample, mirrors the electrons that lie within r_cut of their nearest nucleus through that nucleus, evaluates evaluate_force on both the original and the mirrored configuration, and returns their importance-weighted average. This reduces the variance of the force estimator without introducing additional bias. Only compatible with estimators that do not require the local energy or the mean energy as an input, e.g. evaluate_hf_force_bare() or evaluate_hf_force_ac_zvq().

Parameters:

evaluate_force (Callable) – a force estimator of signature (rng, params, phys_conf) -> jax.Array, e.g. as returned by evaluate_hf_force_bare().
wf (ParametrizedWaveFunction) – the parametrized wave function.
r_cut (float) – the cutoff radius around each nucleus within which electrons are mirrored.

Returns:

a function of signature (rng, params, phys_conf) -> jax.Array that evaluates the antithetic-sampling force estimate for a batch of samples.

Return type:

Multi-device execution#

DeepQMC parallelizes training and evaluation across the available GPUs, see Execution on multiple GPUs. The module below implements the low-level multi-device and multi-host primitives used throughout the package, and is useful when implementing custom ObservableMonitor, LossFunction, or other extensions that need to be aware of the underlying device parallelism.

deepqmc.parallel.PMAP_AXIS_NAME = 'device_axis'#: Alias for str. The default axis_name used for jax.pmap() calls and collective operations (e.g. jax.lax.pmean()) across the package.

deepqmc.parallel.align_rng_key_across_devices(rng)[source]#

Aligns rng keys on multiple devices.

Parameters:: rng – the same rng key stored on each single device.

deepqmc.parallel.all_device_max(x, axis_name='device_axis', **mean_kwargs)[source]#

Compute max across all devices.

Parameters:

x – the input data stored on multiple devices.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.all_device_mean(x, axis_name='device_axis', **mean_kwargs)[source]#

Compute mean across all devices.

Parameters:

x – the input data stored on multiple devices.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.all_device_median(x, axis_name='device_axis')[source]#

Compute median across all devices.

Parameters:

x – the input data stored on multiple devices.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.all_device_min(x, axis_name='device_axis', **mean_kwargs)[source]#

Compute min across all devices.

Parameters:

x – the input data stored on multiple devices.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.all_device_quantile(x, quantile, axis_name='device_axis')[source]#

Compute quantiles across all devices.

Parameters:

x – the input data stored on multiple devices.
quantile – probability for the quantiles to compute.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.all_device_std(x, axis_name='device_axis', **mean_kwargs)[source]#

Compute standard deviation across all devices.

Parameters:

x – the input data stored on multiple devices.
axis_name – optional, name of pmap-ed axis.

deepqmc.parallel.broadcast_to_devices(pytree)#

Broadcast an array stored on a single device to all devices.

The input array must already have the properly sized leading device axis (input.shape[0] == jax.device_count()). Useful for broadcasting data that differs across devices to the devices.

Return type:: T

deepqmc.parallel.gather_electrons_on_one_device(pytree, electron_batch_axis=3)[source]#

Gather electron sample type arrays on one device.

Many arrays (e.g. local energies, wave function values, etc.) are of the shape [n_device, ..., electron_batch_size / n_device, ...]. The total electron_batch_size many samples are stored across the devices. This function gathers arrays like these from the devices, and merges the electron batch axes to arrive at the output shape [..., electron_batch_size, ...]. The most common usecase involves arrays of shape [n_device, molecule_batch_size, electronic_states, electron_batch_size / n_device, ...] and hence the axis of the electron batch is 3. The electron_batch_axis argument can be used if the axis of the electron batch differs from the regular case.

Parameters:

pytree – a pytree of arrays all with shape: [n_device, ... , electron_batch_size / n_device, ...]
electron_batch_axis – the axis carrying the electron batch

Returns:

[..., electron_batch_size, ...].

Return type:

a pytree of arrays all with shape

deepqmc.parallel.local_slice()[source]#

Return a slice selecting the local devices from an array of all devices.

Return type:: slice

deepqmc.parallel.pexp_normalize_mean(x, axis_name='device_axis')[source]#: Compute the normalized-mean exponential of the input across many devices.

deepqmc.parallel.pmap(fn, axis_name='device_axis', **kwargs)[source]#: Alias of jax.pmap, with default axis_name value PMAP_AXIS_NAME for convenience.

deepqmc.parallel.pmap_all_gather(x)#

Gather data from all devices.

Includes it’s own pmap call inside.

deepqmc.parallel.pmap_pmean(x)#

Gather data using pmean from all devices.

Includes it’s own pmap call inside.

deepqmc.parallel.pmax(x, axis_name='device_axis', **kwargs)[source]#: Alias of jax.lax.pmax, with default axis_name value PMAP_AXIS_NAME for convenience.

deepqmc.parallel.pmean(x, axis_name='device_axis', **kwargs)[source]#: Alias of jax.lax.pmean, with default axis_name value PMAP_AXIS_NAME for convenience.

deepqmc.parallel.pmin(x, axis_name='device_axis', **kwargs)[source]#: Alias of jax.lax.pmin, with default axis_name value PMAP_AXIS_NAME for convenience.

deepqmc.parallel.replicate_on_devices(pytree, globally=False)[source]#

Replicate the input pytree on all devices.

Tiles the input arrays to add a leading device axis. The data will be the same across all devices. The effect is analogous to calling jnp.repeat(input[None], jax.device_count(), 0), except that it also works for pytrees, and the output array will be sharded across the devices. Useful for replicating the same data across all devices.

deepqmc.parallel.rng_iterator(rng)[source]#

Create an rng key iterator on each device.

Parameters:: rng (KeyArray) – rng key with a leading device axis, rng keys stored on each device.
Return type:: Generator[KeyArray, None, None]

deepqmc.parallel.scatter_electrons_to_devices(pytree)[source]#

Scatter electron sample type arrays across all devices.

Can be thought of as an inverse of gather_electrons_on_one_device().

Parameters:: pytree (T) – a pytree of arrays all with shape: [molecule_batch_size, electronic_states, electron_batch_size]
Returns:: [n_device, molecule_batch_size, electronic_states, electron_batch_size / n_device, ...]
Return type:: a pytree of arrays all with shape

deepqmc.parallel.select_one_device(pytree, idx=0)[source]#

Select one entry from the device axis.

Selects a single entry from the device axis, resulting in an array that is stored only on a single device. Useful for getting data that is identical across devices to a single device. Can be thought of as an inverse of replicate_on_devices().

Parameters:

pytree (T) – the input pytree of arrays.
idx – the index of the entry to select from the leading device axis.

Return type:

deepqmc.parallel.split_on_devices(rng, num)#

Call the jax.random.split function on each device.

Parameters:

rng – rng key with a leading device axis, rng keys stored on each device.
num (int) – the number of output keys on each device.

deepqmc.parallel.split_rng_key_to_devices(rng)[source]#

Create and place a separate rng key on each device.

Parameters:: rng – a simple rng key stored on a single device.

Logging#

class deepqmc.log.Checkpoint(step, path)[source]#

path: Path#: Alias for field number 1

step: int#: Alias for field number 0

class deepqmc.log.CheckpointStore(workdir, *, size=9223372036854775807, interval=1000)[source]#

Stores training checkpoints in the working directory.

Parameters:

workdir (str) – path where checkpoints are stored.
size (int) – maximum number of checkpoints stored at any time.
interval (int) – number of steps between two checkpoints.

class deepqmc.log.H5Logger(workdir, init_step, additional_keys_to_whitelist=None, aux_data=None, *, keys_whitelist=None)[source]#

Log selected training data to an HDF5 file.

Writes a result.h5 file in the working directory. Only entries whose key contains at least one of the whitelisted phrases are written; all others are silently dropped. The file is opened in SWMR mode so an external reader can access it while training is still running.

Parameters:

workdir (str) – directory in which result.h5 is created.
init_step (int) – initial training step; existing datasets in the file are resized to this length, enabling seamless resumption from a checkpoint.
additional_keys_to_whitelist (Optional[list[str]]) – extra key substrings to append to the default whitelist.
aux_data (Optional[dict]) – key-value pairs stored as HDF5 file attributes; typically used for static metadata such as molecular coordinates.
keys_whitelist (Optional[list[str]]) – overrides the default whitelist entirely when provided; additional_keys_to_whitelist is still appended on top.

update(single_device_data)[source]#

Write a new row of whitelisted entries to the HDF5 file.

Parameters:: single_device_data (Stats) – flat or nested statistics dictionary; entries whose flattened key contains a whitelisted phrase are appended to the corresponding dataset.

class deepqmc.log.MetricLogger(workdir, n_mol)[source]#

Protocol for a general MetricLogger.

update(step, single_device_stats, multi_device_stats, mol_idxs, prefix=None)[source]#

Update the MetricLogger with single and multi device stats.

Parameters:

step (int) – the step at which to add the new entries.
single_device_stats (Stats) – a dictionary containing the entries to add, that are on a single device.
multi_device_stats (Stats) – a dictionary containing the entries to add, that are stored over multiple devices.
mol_idxs (Array) – indices of molecules considered in the given step.
prefix (Optional[str]) – optional, an optional prefix to append to the keys.

class deepqmc.log.TensorboardMetricLogger(workdir, n_mol, *, max_queue=10)[source]#

An interface for writing metrics to Tensorboard.

update(step, single_device_stats, multi_device_stats, mol_idxs, prefix=None)[source]#

Update tensorboard writer with a dictionary of entries.

Parameters:

step (int) – the step at which to add the new entries.
single_device_stats (Stats) – a dictionary containing the entries to add, that are on a single device.
multi_device_stats (Stats) – a dictionary containing the entries to add, that are stored over multiple devices.
mol_idxs (Array) – indices of molecules considered in the given step.
prefix (Optional[str]) – an optional prefix to append to the stat keys.

Postprocessing#

The deepqmc.postprocess package collects utilities for analyzing a finished training or evaluation run: reinstantiating a trained ansatz from a checkpoint, reading logged observables from a workdir, and estimating their Monte Carlo sampling error.

Checkpoints and ansatzes#

deepqmc.postprocess.checkpoint_utils.load_parameters(chkpt_path, state=None)[source]#

Load ansatz parameters from a checkpoint file.

Parameters:

chkpt_path (Path) – path to a chkpt-*.pt file written by CheckpointStore.
state (int, optional) – if the ansatz has multiple electronic states, the index of the state to load parameters for. If None, the parameters for all electronic states are returned, with the electronic-state dimension preserved.

Returns:

the ansatz parameters stored in the checkpoint.

Return type:

Params

deepqmc.postprocess.checkpoint_utils.phys_conf_from_checkpoint(chkpt_path)[source]#

Load a PhysicalConfiguration from a checkpoint file.

Reconstructs the electron and nuclear sample positions stored in the sampler state of a checkpoint, broadcasting the nuclear positions to match the electronic-state and electron-batch dimensions of the electron positions. Assumes a single molecule, i.e. the returned mol_idx is all zeros.

Parameters:: chkpt_path (Path) – path to a chkpt-*.pt file written by CheckpointStore.
Returns:: the electron and nuclear positions stored in the checkpoint’s sampler state.
Return type:: PhysicalConfiguration

deepqmc.postprocess.ansatz_utils.instantiate_predefined_ansatz(ansatz_name, H)[source]#

Instantiate one of the predefined ansatzes.

The hydra configuration file ansatz_name.yaml must be present in the src/deepqmc/conf/ansatz directory.

Parameters:

ansatz_name (str) – the name of the predefined ansatz configuration, e.g. 'psiformer' or 'transpsiformer'.
H (MolecularHamiltonian) – the Hamiltonian of the system the ansatz is instantiated for.

Returns:

the instantiated wave function ansatz, with: uninitialized parameters.

Return type:

Ansatz

deepqmc.postprocess.ansatz_utils.instantiate_wf_from_checkpoint(ansatz_name, H, chkpt_path, state=0)[source]#

Instantiate a predefined ansatz and load its parameters from a checkpoint file.

Convenience function combining instantiate_predefined_ansatz() and load_parameters(), to obtain a ready-to-evaluate WaveFunction from a finished training run.

Parameters:

ansatz_name (str) – the name of the predefined ansatz configuration that was used for training, e.g. 'psiformer' or 'transpsiformer'.
H (MolecularHamiltonian) – the Hamiltonian of the system the ansatz was trained for.
chkpt_path (Path) – path to a chkpt-*.pt file written by CheckpointStore.
state (int) – the index of the electronic state for which to load parameters.

Returns:

the trained wave function, with its parameters: already bound, ready to be evaluated on a PhysicalConfiguration.

Return type:

WaveFunction

Reading workdirs#

deepqmc.postprocess.workdir.read_workdir(path, keys)[source]#

Read values of given keys from result.h5 files in a deepQMC workdir.

Automatically detects whether path contains a single-node training/ evaluation subdir, or multiple training_0, training_1, … / evaluation_0, evaluation_1, … multi-node subdirs, and concatenates the results of the latter along the electron batch axis.

Parameters:

path (Path) – the deepQMC workdir, e.g. as passed to hydra.run.dir.
keys (list[str]) – the names of the datasets to read from the result.h5 files, e.g. ['local_energy', 'mol_idxs'].

Returns:

a dictionary mapping each of the requested keys present in the result files to the corresponding array, and the iteration of the last checkpoint file found in the workdir, or None if no training/evaluation subdir or checkpoint was found.

Return type:

tuple[dict, Optional[int]]

deepqmc.postprocess.workdir.read_and_reshape_result(path, *keys, read_workdir=<function read_workdir>, gather_electrons=True)[source]#

Read and reshape results from a deepQMC workdir.

Use when retrieving all individual samples from a result file that contains results for a single molecule, and no mol_idxs entry, e.g. for a standard (non-transferable) training or evaluation run.

Parameters:

path (Path) – the deepQMC workdir, e.g. as passed to hydra.run.dir.
*keys (str) – the names of the datasets to read from the result.h5 files, e.g. 'local_energy'.
read_workdir (Callable) – optional, the function used to read the raw results from path, defaults to read_workdir().
gather_electrons (bool) – optional, whether to merge the per-device electron batch axis of the results via gather_electron_axis(). If False, only the results of the first device are kept.

Returns:

the array for the requested key, or a dictionary mapping each requested key to its array if more than one key was requested.

deepqmc.postprocess.workdir.read_and_convert_result(path, *keys, read_workdir=<function read_workdir>, gather_electrons=True)[source]#

Read and convert results from a deepQMC workdir to per molecule format.

Use when the result file contains results for multiple molecules, and the mol_idxs entry, e.g. for a transferable training or evaluation run. See also convert_to_per_molecule_format().

Parameters:

path (Path) – the deepQMC workdir, e.g. as passed to hydra.run.dir.
*keys (str) – the names of the datasets to read from the result.h5 files, e.g. 'local_energy'.
read_workdir (Callable) – optional, the function used to read the raw results from path, defaults to read_workdir().
gather_electrons (bool) – optional, whether to merge the per-device electron batch axis of the results via gather_electron_axis(). If False, only the results of the first device are kept.

Returns:

the array for the requested key, or a dictionary mapping each requested key to its array if more than one key was requested, rearranged into per molecule format.

deepqmc.postprocess.workdir.convert_to_per_molecule_format(raw_result, mol_idxs)[source]#

Convert results (local energies, psi values, etc.) to per molecule format.

Parameters:

raw_result (ndarray) – the result values in batched format used during training/evaluation, shape: [n_iter, molecule_batch_size, ...].
mol_idxs (ndarray) – the global dataset indices of the molecules considered in each iteration, shape: [n_iter, molecule_batch_size].

Returns:

the results rearranged into per molecule format, shape: [n_iter_per_molecule, n_molecules, ...].

Return type:

ndarray

deepqmc.postprocess.workdir.gather_electron_axis(pytree, electron_batch_axis=4)[source]#

Merge the per-device electron-sample axis of arrays read from a result file.

Local energies, wave function values and other per-sample quantities are stored in the result.h5 file of a deepQMC workdir with the same [n_iterations, n_device, ..., electron_batch_size / n_device, ...] layout they have during training, see gather_electrons_on_one_device(). This function merges the device axis (axis 1) into the electron batch axis, to arrive at the more convenient [n_iterations, ..., electron_batch_size, ...] layout.

Parameters:

pytree – a pytree of arrays all with shape [n_iterations, n_device, ..., electron_batch_size / n_device, ...].
electron_batch_axis (int) – optional, the axis carrying the electron batch before merging, e.g. 4 for arrays of shape [n_iterations, n_device, molecule_batch_size, electronic_states, electron_batch_size / n_device, ...], the most common case.

Returns:

a pytree of arrays all with shape [n_iterations, ..., electron_batch_size, ...].

deepqmc.postprocess.workdir.last_checkpoint_iteration(path)[source]#

Return the iteration of the last checkpoint file in a deepQMC subdir.

Return type:: int | None

deepqmc.postprocess.workdir.read_average_iteration_time(path, discard_first=10)[source]#

Read the average iteration time from a deepQMC workdir.

Needs the time entry in the result.h5 file. Computes the average wall-clock time per training/evaluation iteration, discarding the first few iterations (which include JIT compilation overhead) from the estimate.

Parameters:

path (Path) – the deepQMC workdir, e.g. as passed to hydra.run.dir.
discard_first (int) – optional, the number of initial iterations to discard from the estimate.

Returns:

the average wall-clock time per iteration, in seconds.

Return type:

float

Monte Carlo statistics#

deepqmc.postprocess.mc_utils.sampling_error(samples, walker_axis=-1, iteration_axis=0)[source]#

Estimate the statistical sampling error from a set of samples via blocking.

Splits the samples into blocks along iteration_axis (each block being one sampling iteration), averages within each block over walker_axis, and returns the standard error of the mean of the resulting block averages. This blocking approach reduces the impact of autocorrelation between consecutive Monte Carlo steps compared to a naive standard-error estimate over all samples.

Parameters:

samples (Array) – the array of samples, containing at least a walker and an iteration axis.
walker_axis (int) – optional, the axis indexing independent walkers, e.g. the electron batch axis.
iteration_axis (int) – optional, the axis indexing the sampling iterations, e.g. the training or evaluation steps.

Returns:

the estimated sampling error, with the walker_axis and iteration_axis reduced out.

Return type:

deepqmc.postprocess.mc_utils.clipped_mean_and_sampling_error(samples, lower=-inf, upper=inf, iteration_axis=0, walker_axis=1)[source]#

Compute the mean and sampling error of samples clipped to a given range.

Discards samples (and nan values) outside of [lower, upper] before computing the mean over both the walker and iteration axes, and estimates the sampling error from the standard deviation of the per-walker means. Useful for robustly estimating observables, e.g. local energies, in the presence of rare outlier samples.

Parameters:

samples (Array) – the array of samples, containing a walker and an iteration axis.
lower (float) – the lower bound of the range samples are clipped to.
upper (float) – the upper bound of the range samples are clipped to.
iteration_axis (int) – optional, the axis indexing the sampling iterations.
walker_axis (int) – optional, the axis indexing independent walkers.

Returns:

a tuple of the clipped mean, the sampling error of the mean, and a dictionary with the ratio of samples that were kept, i.e. not clipped, under the key 'kept_sample_ratio'.

Return type:

tuple[Array, Array, dict]

deepqmc.postprocess.mc_utils.clipped_batch_mean_and_std(samples, lower, upper, walker_axis=1)[source]#

Compute the batchwise mean and std. dev. of samples clipped to a given range.

Discards samples (and nan values) outside of [lower, upper] before computing the mean and standard deviation along walker_axis, independently for each remaining batch entry, e.g. each iteration.

Parameters:

samples (Array) – the array of samples.
lower (float) – the lower bound of the range samples are clipped to.
upper (float) – the upper bound of the range samples are clipped to.
walker_axis (int) – optional, the axis indexing independent walkers.

Returns:

a tuple of the clipped mean and standard deviation, with walker_axis reduced out.

Return type: