def test_frozen_list1(self):
mo_occ = [[np.array([1,1,0,0,0]),],
[np.array([1,1,1,0,0]),]]
mp = fake_ump(frozen=([1,],[1,]), mo_occ=mo_occ, nkpts=1)
nocc = get_nocc(mp)
nmo = get_nmo(mp)
self.assertAlmostEqual(nocc, (1,2))
self.assertAlmostEqual(nmo, (4,4))
nocc = get_nocc(mp, per_kpoint=True)
nmo = get_nmo(mp, per_kpoint=True)
self.assertAlmostEqual(nocc, ([1],[2]))
self.assertAlmostEqual(nmo, ([4],[4]))
def test_no_frozen(self):
mo_occ = [[np.array([1,1,0,0,0]),],
[np.array([1,1,1,0,0]),]]
mp = fake_ump(frozen=None, mo_occ=mo_occ, nkpts=1)
nocc = get_nocc(mp)
nmo = get_nmo(mp)
self.assertAlmostEqual(nocc, (2,3))
self.assertAlmostEqual(nmo, (5,5))
nocc = get_nocc(mp, per_kpoint=True)
nmo = get_nmo(mp, per_kpoint=True)
self.assertAlmostEqual(nocc, ([2],[3]))
self.assertAlmostEqual(nmo, ([5],[5]))
def test_frozen_kpt_list1(self):
frozen = ([[2], [
0,
]], [[0, 1, 3, 4], [0, 2]])
mo_occ = [[
np.array([1, 1, 1, 0, 0]),
np.array([1, 1, 0, 0, 0]),
], [
np.array([1, 0, 0, 0, 0]),
np.array([1, 1, 1, 0, 0]),
]]
mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=2)
nocc = get_nocc(mp)
nmo = get_nmo(mp)
self.assertAlmostEqual(nocc, (2, 1))
self.assertAlmostEqual(
nmo, (5, 3)) # Take the max occ and max vir from kpoints -
# spin0: 2 occ (from first kpoint), 3 vir (from second kpoint)
# spin1: 1 occ (from second kpoint),2 vir (from second kpoint)
nocc = get_nocc(mp, per_kpoint=True)
nmo = get_nmo(mp, per_kpoint=True)
self.assertAlmostEqual(nocc, ([2, 1], [0, 1]))
self.assertAlmostEqual(nmo, ([4, 4], [1, 3]))
# Using `O` and `V` as occupied and virtual and `-` as a padded/frozen orbital,
# we will have the following orbital list
#
# alpha spin | beta spin
# k1 : O O | - V V | - | - V
# k2 : O - | V V V | O | V V
split_idx = padding_k_idx(mp, kind='split')
outa = [[y for x in split_idx[0][idx] for y in x]
for idx in range(2)] # Flatten the list for ease-of-use
outb = [[y for x in split_idx[1][idx] for y in x] for idx in range(2)]
expecteda = [[0, 1, 0], [1, 2, 0, 1,
2]] # [occ_idx, vir_idx] for alpha
expectedb = [[0], [1, 0, 1]] # [occ_idx, vir_idx] for beta
self.assertAlmostEqual(outa, expecteda)
self.assertAlmostEqual(outb, expectedb)
outa, outb = padding_k_idx(mp, kind='joint')
outa = [y for x in outa for y in x]
outb = [y for x in outb for y in x]
expecteda = [0, 1, 3, 4, 0, 2, 3, 4] # [occa_idx_k1, occa_idx_k2]
expectedb = [2, 0, 1, 2]
self.assertAlmostEqual(outa, expecteda)
self.assertAlmostEqual(outb, expectedb)
def test_frozen_list4(self):
frozen = ([0,2], [0,1,3,4])
mo_occ = [[np.array([1,1,0,0,0]),],
[np.array([1,1,1,0,0]),]]
mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=1)
nocc = get_nocc(mp)
nmo = get_nmo(mp)
self.assertAlmostEqual(nocc, (1,1))
self.assertAlmostEqual(nmo, (3,1))
nocc = get_nocc(mp, per_kpoint=True)
nmo = get_nmo(mp, per_kpoint=True)
self.assertAlmostEqual(nocc, ([1],[1]))
self.assertAlmostEqual(nmo, ([3],[1]))
def test_frozen_kpt_list2(self):
frozen = ([[0, 1], [
0,
]], [[0], [0, 1]])
mo_occ = [[
np.array([1, 1, 1, 0, 0]),
np.array([1, 1, 0, 0, 0]),
], [
np.array([1, 1, 0, 0, 0]),
np.array([1, 1, 1, 0, 0]),
]]
mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=2)
nocc = get_nocc(mp)
nmo = get_nmo(mp)
self.assertAlmostEqual(nocc, (1, 1))
self.assertAlmostEqual(
nmo,
(4, 4)) # We deleted the first occ from each of the alpha/beta, so
# we have 1 occupied and 3 virtuals.
nocc = get_nocc(mp, per_kpoint=True)
nmo = get_nmo(mp, per_kpoint=True)
self.assertAlmostEqual(nocc, ([1, 1], [1, 1]))
self.assertAlmostEqual(nmo, ([3, 4], [4, 3]))
# Using `O` and `V` as occupied and virtual and `-` as a padded/frozen orbital,
# we will have the following orbital list
#
# alpha spin | beta spin
# k1 : O | - V V | O | V V V
# k2 : O | V V V | O | - V V
split_idx = padding_k_idx(mp, kind='split')
outa = [[y for x in split_idx[0][idx] for y in x]
for idx in range(2)] # Flatten the list for ease-of-use
outb = [[y for x in split_idx[1][idx] for y in x] for idx in range(2)]
expecteda = [[0, 0], [1, 2, 0, 1, 2]] # [occ_idx, vir_idx] for alpha
expectedb = [[0, 0], [0, 1, 2, 1, 2]] # [occ_idx, vir_idx] for beta
self.assertAlmostEqual(outa, expecteda)
self.assertAlmostEqual(outb, expectedb)
outa, outb = padding_k_idx(mp, kind='joint')
outa = [y for x in outa for y in x]
outb = [y for x in outb for y in x]
expecteda = [0, 2, 3, 0, 1, 2, 3] # [occa_idx_k1, occa_idx_k2]
expectedb = [0, 1, 2, 3, 0, 2, 3]
self.assertAlmostEqual(outa, expecteda)
self.assertAlmostEqual(outb, expectedb)
