def test_change_of_basis():
n = 2
l = 10
dim = 2
new_l = 2 * l - n
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = spas.construct_general_orbital_system()
C_spas = RandomBasisSet.get_random_elements((spas.l, new_l), np)
C_gos = RandomBasisSet.get_random_elements((gos.l, new_l), np)
h_spas = change_basis_h(spas.h.copy(), C_spas)
u_spas = change_basis_u(spas.u.copy(), C_spas)
h_gos = change_basis_h(gos.h.copy(), C_gos)
u_gos = change_basis_u(gos.u.copy(), C_gos)
spas.change_basis(C_spas)
gos.change_basis(C_gos)
assert spas.l == new_l
assert all([new_l == s for s in spas.h.shape])
assert all([new_l == s for s in spas.u.shape])
np.testing.assert_allclose(h_spas, spas.h, atol=1e-12, rtol=1e-12)
np.testing.assert_allclose(u_spas, spas.u, atol=1e-12, rtol=1e-12)
assert gos.l == new_l
assert all([new_l == s for s in gos.h.shape])
assert all([new_l == s for s in gos.u.shape])
np.testing.assert_allclose(h_gos, gos.h, atol=1e-12, rtol=1e-12)
np.testing.assert_allclose(u_gos, gos.u, atol=1e-12, rtol=1e-12)
def test_spin_up_down():
n = 4
l = 10
dim = 2
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = spas.construct_general_orbital_system()
a = gos._basis_set.a
b = gos._basis_set.b
sigma_up = 0.5 * (gos._basis_set.sigma_x + 1j * gos._basis_set.sigma_y)
sigma_down = 0.5 * (gos._basis_set.sigma_x - 1j * gos._basis_set.sigma_y)
np.testing.assert_allclose(a, sigma_up @ b)
np.testing.assert_allclose(np.zeros_like(a), sigma_up @ a)
np.testing.assert_allclose(np.zeros_like(b), sigma_down @ b)
np.testing.assert_allclose(b, sigma_down @ a)
np.testing.assert_allclose(a, sigma_up @ sigma_down @ a)
np.testing.assert_allclose(b, sigma_down @ sigma_up @ b)
S_up = np.kron(spas.s, sigma_up)
S_down = np.kron(spas.s, sigma_down)
np.testing.assert_allclose(S_up, gos.spin_x + 1j * gos.spin_y)
np.testing.assert_allclose(S_down, gos.spin_x - 1j * gos.spin_y)
def test_gos_spin_matrices():
n = 4
l = 10
dim = 2
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = spas.construct_general_orbital_system()
for spin, sigma in zip(
[gos.spin_x, gos.spin_y, gos.spin_z],
[
gos._basis_set.sigma_x,
gos._basis_set.sigma_y,
gos._basis_set.sigma_z,
],
):
spin_2 = np.zeros((gos.l, gos.l), dtype=np.complex128)
for i in range(gos.l):
a = i % 2
g = i // 2
for j in range(gos.l):
b = j % 2
d = j // 2
spin_2[i, j] = 0.5 * spas.s[g, d] * sigma[a, b]
np.testing.assert_allclose(spin, spin_2)
def test_setters():
n = 2
l = 10
dim = 3
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = spas.construct_general_orbital_system()
assert gos.l == 2 * spas.l
def test_no_operators():
n = 4
l = 10
dim = 3
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
assert not spas.has_one_body_time_evolution_operator
assert not gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
np.testing.assert_allclose(spas.h_t(10), spas.h)
np.testing.assert_allclose(spas.u_t(10), spas.u)
np.testing.assert_allclose(gos.h_t(10), gos.h)
np.testing.assert_allclose(gos.u_t(10), gos.u)
spas.set_time_evolution_operator(
[],
add_h_0=False,
add_u_0=False,
)
gos.set_time_evolution_operator([], add_h_0=False, add_u_0=False)
assert not spas.has_one_body_time_evolution_operator
assert not gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
np.testing.assert_allclose(spas.h_t(0), np.zeros_like(spas.h))
np.testing.assert_allclose(spas.u_t(0), np.zeros_like(spas.u))
np.testing.assert_allclose(gos.h_t(0), np.zeros_like(gos.h))
np.testing.assert_allclose(gos.u_t(0), np.zeros_like(gos.u))
def test_spin_squared():
n = 2
l = 2
dim = 2
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
spas = SpatialOrbitalSystem(
n, ODQD(l, 8, 1001, potential=ODQD.HOPotential(1)))
gos = spas.construct_general_orbital_system(a=[1, 0], b=[0, 1])
overlap = spas.s
overlap_sq = np.einsum("pr, qs -> pqrs", overlap, overlap)
a = gos._basis_set.a
b = gos._basis_set.b
aa = np.kron(a, a)
ab = np.kron(a, b)
ba = np.kron(b, a)
bb = np.kron(b, b)
triplets = [aa, 1 / np.sqrt(2) * (ab + ba), bb]
singlet = [1 / np.sqrt(2) * (ab - ba)]
# S^2 with alpha = [1, 0]^T and beta = [0, 1]^T
S_sq_spin = np.zeros((4, 4))
S_sq_spin[0, 0] = 2
S_sq_spin[3, 3] = 2
S_sq_spin[1, 1] = 1
S_sq_spin[2, 2] = 1
S_sq_spin[1, 2] = 1
S_sq_spin[2, 1] = 1
S_sq_spin = S_sq_spin.reshape(2, 2, 2, 2)
for trip in triplets:
# Check that the eigenvalue of all triplet states is 2
np.testing.assert_allclose(trip.T @ S_sq_spin.reshape(4, 4) @ trip, 2)
# Check that the eigenvalue of the singlet state is 0
np.testing.assert_allclose(
singlet[0].T @ S_sq_spin.reshape(4, 4) @ singlet[0], 0)
def get_odho_ao():
n = 2
l = 20
grid_length = 5
num_grid_points = 1001
omega = 1
odho = SpatialOrbitalSystem(
n,
ODQD(
l, grid_length, num_grid_points, potential=ODQD.HOPotential(omega)
),
)
return odho
def test_overlap_squared():
n = 4
l = 10
dim = 2
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
overlap = spas.s
overlap_sq = np.einsum("pr, qs -> pqrs", overlap, overlap)
for p in range(spas.l):
for q in range(spas.l):
for r in range(spas.l):
for s in range(spas.l):
np.testing.assert_allclose(overlap_sq[p, q, r, s],
overlap[p, r] * overlap[q, s])
def construct_custom_system(
n,
l,
s,
h,
u,
dim=3,
particle_charge=-1,
np=None,
includes_spin=False,
anti_symmetrized_u=False,
system_type="general",
**kwargs,
):
if np is None:
import numpy as np
bs = setup_basis_set(
n,
l,
s,
h,
u,
dim,
particle_charge,
np,
includes_spin,
anti_symmetrized_u,
**kwargs,
)
if system_type.lower() == "general":
system = GeneralOrbitalSystem(n, bs)
elif system_type.lower() == "spatial":
system = SpatialOrbitalSystem(n, bs)
else:
raise NotImplementedError(
f"System type: {system_type} is not supported!")
if "C" in kwargs.keys():
system.change_basis(kwargs["C"])
return system
def test_single_time_evolution_operator():
n = 4
l = 10
dim = 3
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
assert not spas.has_one_body_time_evolution_operator
assert not gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
np.testing.assert_allclose(spas.h_t(10), spas.h)
np.testing.assert_allclose(spas.u_t(10), spas.u)
np.testing.assert_allclose(gos.h_t(10), gos.h)
np.testing.assert_allclose(gos.u_t(10), gos.u)
spas.set_time_evolution_operator(CustomOneBodyOperator(2, spas.h),
add_h_0=False)
gos.set_time_evolution_operator(CustomOneBodyOperator(3, gos.h),
add_u_0=False)
assert spas.has_one_body_time_evolution_operator
assert gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
np.testing.assert_allclose(
spas.h_t(0),
spas.h * 2,
)
np.testing.assert_allclose(spas.u_t(0), spas.u)
np.testing.assert_allclose(
gos.h_t(0),
gos.h + gos.h * 3,
)
np.testing.assert_allclose(gos.u_t(0), np.zeros_like(gos.u))
def test_multiple_time_evolution_operators():
n = 4
l = 10
dim = 3
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
assert not spas.has_one_body_time_evolution_operator
assert not gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
spas.set_time_evolution_operator(
[
CustomOneBodyOperator(2, spas.h),
CustomOneBodyOperator(3, spas.s),
AdiabaticSwitching(2),
],
add_u_0=False,
)
gos.set_time_evolution_operator(
(
CustomOneBodyOperator(1, gos.h),
CustomOneBodyOperator(3, gos.s),
CustomOneBodyOperator(-2, gos.position[0]),
),
add_h_0=False,
)
assert spas.has_one_body_time_evolution_operator
assert gos.has_one_body_time_evolution_operator
assert spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
np.testing.assert_allclose(
spas.h_t(0),
spas.h + spas.h * 2 + spas.s * 3,
)
np.testing.assert_allclose(spas.u_t(0), 2 * spas.u)
np.testing.assert_allclose(
gos.h_t(0),
gos.h + gos.s * 3 - gos.position[0] * 2,
)
np.testing.assert_allclose(gos.u_t(0), gos.u)
def test_single_dipole_time_evolution_operator():
n = 4
l = 10
dim = 3
omega = 0.25
spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
field = lambda t: np.sin(omega * 2)
polarization = np.zeros(dim)
polarization[0] = 1
spas.set_time_evolution_operator(
DipoleFieldInteraction(
field,
polarization,
))
gos.set_time_evolution_operator(
DipoleFieldInteraction(
field,
polarization,
))
assert spas.has_one_body_time_evolution_operator
assert gos.has_one_body_time_evolution_operator
assert not spas.has_two_body_time_evolution_operator
assert not gos.has_two_body_time_evolution_operator
for t in [0, 0.1, 0.5, 1.3]:
np.testing.assert_allclose(
spas.h_t(t),
spas.h - field(t) * spas.dipole_moment[0],
)
np.testing.assert_allclose(
gos.h_t(t),
gos.h - field(t) * gos.dipole_moment[0],
)
np.testing.assert_allclose(spas.u_t(t), spas.u)
np.testing.assert_allclose(gos.u_t(t), gos.u)
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)