def test_transform_two_body_elements():
l = 10
u = np.random.random((l, l, l, l)) + 1j * np.random.random((l, l, l, l))
C = np.random.random((l, l)) + 1j * np.random.random((l, l))
u_transformed = np.einsum(
"ls, kr, jq, ip, ijkl -> pqrs",
C,
C,
C.conj(),
C.conj(),
u,
optimize=True,
)
np.testing.assert_allclose(
u_transformed,
BasisSet.transform_two_body_elements(u, C, np=np),
atol=1e-10,
)
_u = np.dot(u, C)
_u = np.tensordot(_u, C, axes=(2, 0)).transpose(0, 1, 3, 2)
_u = np.tensordot(_u, C.conj(), axes=(1, 0)).transpose(0, 3, 1, 2)
_u = np.tensordot(C.conj().T, _u, axes=(1, 0))
np.testing.assert_allclose(_u,
BasisSet.transform_two_body_elements(u, C, np))
np.testing.assert_allclose(
_u, BasisSet.transform_two_body_elements(u, C, np, C_tilde=C.conj().T))
def test_anti_symmetric_properties():
l_half = 10
u = np.random.random((l_half, l_half, l_half, l_half))
# Make u symmetric
u = u + u.transpose(1, 0, 3, 2)
u = BasisSet.anti_symmetrize_u(BasisSet.add_spin_two_body(u, np=np))
np.testing.assert_allclose(u, -u.transpose(0, 1, 3, 2), atol=1e-10)
np.testing.assert_allclose(u, -u.transpose(1, 0, 2, 3), atol=1e-10)
np.testing.assert_allclose(u, u.transpose(1, 0, 3, 2), atol=1e-10)
def test_tddw(get_tddw):
tddw = get_tddw
h_dw = get_double_well_one_body_elements(
tddw.l // 2,
tddw._basis_set.omega,
tddw._basis_set.mass,
tddw._basis_set.barrier_strength,
dtype=np.complex128,
axis=0,
)
epsilon, C_dw = np.linalg.eigh(h_dw)
C = BasisSet.add_spin_one_body(C_dw, np=np)
tddw.change_basis(C[:, : tddw.l])
dip = np.load(os.path.join("tests", "dat", "tddw_dipole_moment.npy"))
np.testing.assert_allclose(
np.abs(dip), np.abs(tddw.dipole_moment), atol=1e-10
)
h = np.load(os.path.join("tests", "dat", "tddw_h.npy"))
np.testing.assert_allclose(h, tddw.h, atol=1e-10)
u = np.load(os.path.join("tests", "dat", "tddw_u.npy"))
np.testing.assert_allclose(np.abs(u), np.abs(tddw.u), atol=1e-10)
spf = np.load(os.path.join("tests", "dat", "tddw_spf.npy"))
np.testing.assert_allclose(np.abs(spf), np.abs(tddw.spf), atol=1e-10)
def test_transform_one_body_elements():
l = 10
h = np.random.random((l, l)) + 1j * np.random.random((l, l))
C = np.random.random((l, l)) + 1j * np.random.random((l, l))
h_transformed = np.einsum("ip, jq, ij", C.conj(), C, h, optimize=True)
np.testing.assert_allclose(
h_transformed,
BasisSet.transform_one_body_elements(h, C, np=np),
atol=1e-10,
)
_h = np.dot(h, C)
_h = np.dot(C.conj().T, _h)
np.testing.assert_allclose(_h,
BasisSet.transform_one_body_elements(h, C, np))
np.testing.assert_allclose(
_h, BasisSet.transform_one_body_elements(h, C, np, C_tilde=C.conj().T))
def test_add_spin_one_body():
l_half = 10
h = np.random.random((l_half, l_half))
l = l_half * 2
h_spin = np.zeros((l, l))
for p in range(l):
for q in range(l):
h_spin[p, q] = spin_delta(p, q) * h[p // 2, q // 2]
np.testing.assert_allclose(h_spin,
BasisSet.add_spin_one_body(h, np=np),
atol=1e-10)
def test_anti_symmetrize_u():
l_half = 10
u = np.random.random((l_half, l_half, l_half, l_half))
# Make u symmetric
u = u + u.transpose(1, 0, 3, 2)
l = l_half * 2
u_spin = np.zeros((l, l, l, l))
for p in range(l):
for q in range(l):
for r in range(l):
for s in range(l):
u_spin[p, q, r, s] = (spin_delta(p, r) * spin_delta(q, s) *
u[p // 2, q // 2, r // 2, s // 2])
u_spin[p, q, r,
s] -= (spin_delta(p, s) * spin_delta(q, r) *
u[p // 2, q // 2, s // 2, r // 2])
np.testing.assert_allclose(
u_spin,
BasisSet.anti_symmetrize_u(BasisSet.add_spin_two_body(u, np=np)),
atol=1e-10,
)
def setup_basis_set(
n,
l,
s,
h,
u,
dim=3,
particle_charge=-1,
np=None,
includes_spin=False,
anti_symmetrized_u=False,
**kwargs,
):
if np is None:
import numpy as np
bs = BasisSet(
l,
dim=dim,
np=np,
includes_spin=includes_spin,
anti_symmetrized_u=anti_symmetrized_u,
)
bs.h = h
bs.u = u
bs.s = s
bs.particle_charge = particle_charge
if "position" in kwargs.keys():
bs.position = kwargs["position"]
if "momentum" in kwargs.keys():
bs.momentum = kwargs["momentum"]
if "nuclear_repulsion_energy" in kwargs.keys():
bs.nuclear_repulsion_energy = kwargs["nuclear_repulsion_energy"]
return bs
def test_spin_two_body():
l_half = 10
u = np.random.random((l_half, l_half, l_half, l_half))
l = l_half * 2
u_spin = np.zeros((l, l, l, l))
for p in range(l):
for q in range(l):
for r in range(l):
for s in range(l):
u_spin[p, q, r, s] = (spin_delta(p, r) * spin_delta(q, s) *
u[p // 2, q // 2, r // 2, s // 2])
np.testing.assert_allclose(u_spin,
BasisSet.add_spin_two_body(u, np=np),
atol=1e-10)
def test_change_of_basis():
# We try to construct a two-dimensional double-well system from a
# two-dimensional harmonic oscillator system by using the change-of-basis
# function with C found from diagonalizing the double-well one-body
# Hamiltonian.
n = 2
l = 12
omega = 1
mass = 1
barrier_strength = 3
radius = 10
num_grid_points = 401
axis = 0
tdho = GeneralOrbitalSystem(
n,
TwoDimensionalHarmonicOscillator(
l, radius, num_grid_points, omega=omega
),
)
h_dw = get_double_well_one_body_elements(
l, omega, mass, barrier_strength, dtype=np.complex128, axis=axis
)
epsilon, C_dw = np.linalg.eigh(h_dw)
C = BasisSet.add_spin_one_body(C_dw, np=np)
tdho.change_basis(C)
tddw = GeneralOrbitalSystem(
n,
TwoDimensionalDoubleWell(
l,
radius,
num_grid_points,
omega=omega,
mass=mass,
barrier_strength=barrier_strength,
axis=axis,
),
)
tddw.change_basis(C)
np.testing.assert_allclose(tdho.u, tddw.u, atol=1e-7)
np.testing.assert_allclose(tdho.spf, tddw.spf, atol=1e-7)
def test_add_spin_spf():
spf = (np.arange(15) + 1).reshape(3, 5).T
n = 3
n_a = 2
n_b = n - n_a
l = 2 * spf.shape[0]
assert l == 10
m_a = l // 2 - n_a
assert m_a == 3
m_b = l // 2 - n_b
assert m_b == 4
new_spf = BasisSet.add_spin_spf(spf, np)
# Occupied spin-up
np.testing.assert_allclose(spf[0], new_spf[0])
np.testing.assert_allclose(spf[1], new_spf[2])
# Occupied spin-down
np.testing.assert_allclose(spf[0], new_spf[1])
# Virtual spin-up
np.testing.assert_allclose(spf[2], new_spf[4])
np.testing.assert_allclose(spf[3], new_spf[6])
np.testing.assert_allclose(spf[4], new_spf[8])
# Virtual spin-down
np.testing.assert_allclose(spf[1], new_spf[3])
np.testing.assert_allclose(spf[2], new_spf[5])
np.testing.assert_allclose(spf[3], new_spf[7])
np.testing.assert_allclose(spf[4], new_spf[9])
def construct_pyscf_system_rhf(
molecule,
basis="cc-pvdz",
add_spin=True,
anti_symmetrize=True,
np=None,
verbose=False,
charge=0,
cart=False,
dip=None,
**kwargs,
):
"""Convenience function setting up a closed-shell atom or a molecule from
PySCF as a ``QuantumSystem`` in RHF-basis using PySCF's RHF-solver.
Parameters
----------
molecule : str
String describing the atom or molecule. This gets passed to PySCF which
means that we support all the same string options as PySCF.
basis : str
String describing the basis set. PySCF determines which options are
available.
add_spin : bool
Whether or not to return a ``SpatialOrbitalSystem`` (``False``) or a
``GeneralOrbitalSystem`` (``True``). Default is ``True``.
anti_symmetrize : bool
Whether or not to anti-symmetrize the two-body elements in a
``GeneralOrbitalSystem``. This only applies if ``add_spin = True``.
Default is ``True``.
np : module
Array- and linear algebra module.
dip : list
[position, momentum] f.ex. from Dalton
Returns
-------
SpatialOrbitalSystem, GeneralOrbitalSystem
Depending on the choice of ``add_spin`` we return a
``SpatialOrbitalSystem`` (``add_spin = False``), or a
``GeneralOrbitalSystem`` (``add_spin = True``).
See Also
-------
PySCF
Example
-------
>>> # Set up the Beryllium atom centered at (0, 0, 0)
>>> system = construct_pyscf_system_rhf(
... "be 0 0 0", basis="cc-pVDZ", add_spin=False
... ) # doctest.ELLIPSIS
converged SCF energy = -14.5723...
>>> # Compare the number of occupied basis functions
>>> system.n == 4 // 2
True
>>> gos = system.construct_general_orbital_system()
>>> gos.n == 4
True
>>> system = construct_pyscf_system_rhf(
... "be 0 0 0", basis="cc-pVDZ"
... ) # doctest.ELLIPSIS
converged SCF energy = -14.5723...
>>> system.n == gos.n
True
"""
import pyscf
if np is None:
import numpy as np
# Build molecule in AO-basis
mol = pyscf.gto.Mole()
mol.unit = "bohr"
mol.charge = charge
mol.cart = cart
mol.build(atom=molecule, basis=basis, **kwargs)
nuclear_repulsion_energy = mol.energy_nuc()
n = mol.nelectron
assert (
n %
2 == 0), "We require closed shell, with an even number of particles"
l = mol.nao
hf = pyscf.scf.RHF(mol)
hf_energy = hf.kernel()
if not hf.converged:
warnings.warn("RHF calculation did not converge")
if verbose:
print(f"RHF energy: {hf.e_tot}")
charges = mol.atom_charges()
coords = mol.atom_coords()
nuc_charge_center = np.einsum("z,zx->x", charges, coords) / charges.sum()
mol.set_common_orig_(nuc_charge_center)
C = np.asarray(hf.mo_coeff)
h = pyscf.scf.hf.get_hcore(mol)
s = mol.intor_symmetric("int1e_ovlp")
u = mol.intor("int2e").reshape(l, l, l, l).transpose(0, 2, 1, 3)
position = mol.intor("int1e_r").reshape(3, l, l)
momentum = -1j * mol.intor("int1e_ipovlp").reshape(3, l, l)
if dip is not None:
diplen = dip[0]
arr_eq0 = np.allclose(diplen[0], position[0])
arr_eq1 = np.allclose(diplen[1], position[1])
arr_eq2 = np.allclose(diplen[2], position[2])
if not (arr_eq0 and arr_eq1 and arr_eq2):
print('WARNING:')
print('dipole_array_pyscf != dipole_array_dalton')
bs = BasisSet(l, dim=3, np=np)
bs.h = h
bs.s = s
bs.u = u
bs.nuclear_repulsion_energy = nuclear_repulsion_energy
bs.particle_charge = -1
bs.position = position
bs.momentum = momentum
if dip is not None:
bs.momentum = dip[1]
bs.change_module(np=np)
system = SpatialOrbitalSystem(n, bs)
system.change_basis(C)
return (system.construct_general_orbital_system(
anti_symmetrize=anti_symmetrize) if add_spin else system)
def construct_pyscf_system_ao(
molecule,
basis="cc-pvdz",
add_spin=True,
anti_symmetrize=True,
np=None,
**kwargs,
):
"""Convenience function setting up an atom or a molecule from PySCF as a
``QuantumSystem``.
Parameters
----------
molecule : str
String describing the atom or molecule. This gets passed to PySCF which
means that we support all the same string options as PySCF.
basis : str
String describing the basis set. PySCF determines which options are
available.
add_spin : bool
Whether or not to return a ``SpatialOrbitalSystem`` (``False``) or a
``GeneralOrbitalSystem`` (``True``). Default is ``True``.
anti_symmetrize : bool
Whether or not to anti-symmetrize the two-body elements in a
``GeneralOrbitalSystem``. This only applies if ``add_spin = True``.
Default is ``True``.
np : module
Array- and linear algebra module.
Returns
-------
SpatialOrbitalSystem, GeneralOrbitalSystem
Depending on the choice of ``add_spin`` we return a
``SpatialOrbitalSystem`` (``add_spin = False``), or a
``GeneralOrbitalSystem`` (``add_spin = True``).
See Also
-------
PySCF
Example
-------
>>> # Set up the Beryllium atom centered at (0, 0, 0)
>>> system = construct_pyscf_system_ao(
... "be 0 0 0", basis="cc-pVDZ", add_spin=False
... )
>>> # Compare the number of occupied basis functions
>>> system.n == 4 // 2
True
>>> gos = system.construct_general_orbital_system()
>>> gos.n == 4
True
"""
import pyscf
if np is None:
import numpy as np
mol = pyscf.gto.Mole()
mol.unit = "bohr"
mol.build(atom=molecule, basis=basis, **kwargs)
nuclear_repulsion_energy = mol.energy_nuc()
n = mol.nelectron
l = mol.nao
# n_a = (mol.nelectron + mol.spin) // 2
# n_b = n_a - mol.spin
# assert n_b == n - n_a
h = pyscf.scf.hf.get_hcore(mol)
s = mol.intor_symmetric("int1e_ovlp")
u = mol.intor("int2e").reshape(l, l, l, l).transpose(0, 2, 1, 3)
position = mol.intor("int1e_r").reshape(3, l, l)
bs = BasisSet(l, dim=3, np=np)
bs.h = h
bs.s = s
bs.u = u
bs.nuclear_repulsion_energy = nuclear_repulsion_energy
bs.particle_charge = -1
bs.position = position
bs.change_module(np=np)
system = SpatialOrbitalSystem(n, bs)
return (system.construct_general_orbital_system(
anti_symmetrize=anti_symmetrize) if add_spin else system)
def test_antisymmetric_two_body_elements(u):
l = len(u)
_u = BasisSet.anti_symmetrize_u(
BasisSet.add_spin_two_body(get_coulomb_elements(l // 2), np=np))
np.testing.assert_allclose(u, _u, atol=1e-6, rtol=1e-6)
def test_antisymmetric_one_body_elements(h):
l = len(h)
_h = BasisSet.add_spin_one_body(get_one_body_elements(l // 2), np=np)
np.testing.assert_allclose(h, _h, atol=1e-6, rtol=1e-6)
