hydrogen_like_atom

Exactly solvable single-electron hydrogen-like atom.

HydrogenLikeWavefunction (Module) dataclass

Model which computes -decay_rate * r, with trainable decay_rate.

This model returns log|psi(x)|, where psi is the 1-s orbital exp(-decay_rate * r). This psi is an eigenfunction of the hydrogen-like Hamiltonian (one ion and one electron) when the decay rate is equal to 2 * nuclear_charge / (d - 1), with d > 1, where d is the dimension of the system.

Thus when the system is 3-d, then this model computes the ground-state wavefunction precisely when decay_rate = nuclear_charge.

Attributes:

Name	Type	Description
decay_rate	jnp.float32	initial decay rate in the model

__call__(self, x) special

Log of isotropic exponential decay. Computes -decay_rate * ||x||.

Parameters:

Name	Type	Description	Default
x	Array	single electron positions, of shape (..., 1, d), where d is the dimension of the system	required

Returns:

Type	Description
Array	log of exponential decay wavefunction, with shape x.shape[:-2]

Source code in vmcnet/examples/hydrogen_like_atom.py

@flax.linen.compact
def __call__(self, x: Array) -> Array:  # type: ignore[override]
    """Log of isotropic exponential decay. Computes -decay_rate * ||x||.

    Args:
        x (Array): single electron positions, of shape (..., 1, d), where d
            is the dimension of the system

    Returns:
        Array: log of exponential decay wavefunction, with shape x.shape[:-2]
    """
    r = jnp.linalg.norm(x, axis=-1)
    scaled_r = models.core.Dense(
        1,
        kernel_init=lambda key, shape, **kwargs: jnp.array(
            [[self.init_decay_rate]]
        ),
        use_bias=False,
    )(r)
    return -jnp.squeeze(scaled_r, axis=-1)

make_hydrogen_like_local_energy(log_psi_apply, charge, d=3)

Local energy calculation for the hydrogen-like atom in general dimension d.

Parameters:

Name	Type	Description	Default
log_psi_apply	Callable	a function which computes log|psi(x)| for single inputs x. It is okay for it to produce batch outputs on batches of x as long as it produces a single number for single x. Has the signature (params, single_x_in) -> log|psi(single_x_in)|	required
charge	jnp.float32	charge of the nucleus	required
d	int	Dimension of the system (number of coordinates that each ion and electron has). Defaults to 3.	3

Returns:

Type	Description
Callable	local energy function which computes -0.5 nabla^2 psi / psi + (Z / r), where psi is the wavefn, and Z is the charge of the nucleus. Has the signature (params, x) -> local energy array of shape (x.shape[0],)

Source code in vmcnet/examples/hydrogen_like_atom.py

def make_hydrogen_like_local_energy(
    log_psi_apply: Callable[[P, Array], Union[jnp.float32, Array]],
    charge: jnp.float32,
    d: int = 3,
) -> ModelApply[P]:
    """Local energy calculation for the hydrogen-like atom in general dimension d.

    Args:
        log_psi_apply (Callable): a function which computes log|psi(x)| for single
            inputs x. It is okay for it to produce batch outputs on batches of x as long
            as it produces a single number for single x. Has the signature
            (params, single_x_in) -> log|psi(single_x_in)|
        charge (jnp.float32): charge of the nucleus
        d (int, optional): Dimension of the system (number of coordinates that each
            ion and electron has). Defaults to 3.

    Returns:
        Callable: local energy function which computes -0.5 nabla^2 psi / psi + (Z / r),
        where psi is the wavefn, and Z is the charge of the nucleus. Has the signature
        (params, x) -> local energy array of shape (x.shape[0],)
    """
    ion_location = jnp.zeros((1, d))
    ion_charge = jnp.array([charge])

    kinetic_fn = physics.kinetic.create_continuous_kinetic_energy(log_psi_apply)
    potential_fn = physics.potential.create_electron_ion_coulomb_potential(
        ion_location, ion_charge
    )

    return physics.core.combine_local_energy_terms([kinetic_fn, potential_fn])