def get_frozen_mask(mp):
'''Boolean mask for orbitals in k-point post-HF method.
Creates a boolean mask to remove frozen orbitals and keep other orbitals for post-HF
calculations.
Args:
mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
Returns:
moidx (list of :obj:`ndarray` of `bool`): Boolean mask of orbitals to include.
'''
moidx = [[np.ones(x.size, dtype=bool) for x in mp.mo_occ[s]]
for s in [0, 1]]
if mp.frozen is None:
pass
elif isinstance(mp.frozen, (int, np.integer)):
for spin in [0, 1]:
for idx in moidx[spin]:
idx[:mp.frozen] = False
elif (_is_arraylike(mp.frozen[0]) and
isinstance(mp.frozen[0][0],
(int, np.integer))): # case example: ([0, 4], [0, 5, 6])
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
[
_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt],
ikpt) for ikpt in range(mp.nkpts)
]
for ikpt, kpt_occ in enumerate(moidx[spin]):
kpt_occ[mp.frozen[spin]] = False
elif (_is_arraylike(mp.frozen[0]) and isinstance(
mp.frozen[0][0],
(list, np.ndarray))): # case example: ([[0,],[]], [[0,1],[4]])
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
nkpts = len(mp.frozen[spin])
if nkpts != mp.nkpts:
raise RuntimeError(
'Frozen list has a different number of k-points (length) than passed in'
'mean-field/correlated calculation. \n\nCalculation nkpts = %d, frozen'
'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
for spin in [0, 1]:
[
_frozen_sanity_check(mp.frozen[spin][ikpt],
mp.mo_occ[spin][ikpt], ikpt)
for ikpt in range(mp.nkpts)
]
for ikpt, kpt_occ in enumerate(moidx[spin]):
kpt_occ[mp.frozen[spin][ikpt]] = False
else:
raise NotImplementedError('No known conversion for frozen %s' %
mp.frozen)
return moidx
def get_nmo(mp, per_kpoint=False):
'''Number of orbitals for k-point calculations.
Number of orbitals for use in a calculation with k-points, taking into account
frozen orbitals.
Note:
If `per_kpoint` is False, then the number of orbitals here is equal to max(nocc) + max(nvir),
where each max is done over all k-points. Otherwise the number of orbitals is returned
as a list of number of orbitals at each k-point.
Args:
mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
per_kpoint (bool, optional): True returns the number of orbitals at each k-point.
For a description of False, see Note.
Returns:
nmo (int, list of int): Number of orbitals. For return type, see description of arg
`per_kpoint`.
'''
if mp._nmo is not None:
return mp._nmo
nmo = [0, 0]
if isinstance(mp.frozen, (int, np.integer)):
for spin in [0, 1]:
nmo[spin] = [
len(mp.mo_occ[spin][k]) - mp.frozen for k in range(mp.nkpts)
]
elif mp.frozen is None:
nmo = [[len(mp.mo_occ[0][k]) for k in range(mp.nkpts)],
[len(mp.mo_occ[1][k]) for k in range(mp.nkpts)]]
elif (_is_arraylike(mp.frozen[0]) and
isinstance(mp.frozen[0][0],
(int, np.integer))): # case example: ([0, 4], [0, 5, 6])
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
[
_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt],
ikpt) for ikpt in range(mp.nkpts)
]
nmo[spin] = [
len(mp.mo_occ[spin][ikpt]) - len(mp.frozen[spin])
for ikpt in range(mp.nkpts)
]
elif (_is_arraylike(mp.frozen[0]) and isinstance(
mp.frozen[0][0],
(list, np.ndarray))): # case example: ([[0,],[]], [[0,1],[4]])
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
nkpts = len(mp.frozen[spin])
if nkpts != mp.nkpts:
raise RuntimeError(
'Frozen list has a different number of k-points (length) than passed in'
'mean-field/correlated calculation. \n\nCalculation nkpts = %d, frozen'
'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
for spin in [0, 1]:
[
_frozen_sanity_check(mp.frozen[spin][ikpt],
mp.mo_occ[spin][ikpt], ikpt)
for ikpt in range(mp.nkpts)
]
nmo[spin] = [
len(mp.mo_occ[spin][ikpt]) - len(mp.frozen[spin][ikpt])
for ikpt in range(nkpts)
]
else:
raise NotImplementedError('No known conversion for frozen %s' %
mp.frozen)
for spin in [0, 1]:
assert all(np.array(nmo[spin]) > 0), (
'Must have a positive number of orbitals! (spin=%d)'
'\n\nnmo %s\nfrozen %s\nmo_occ %s' %
(spin, nmo, mp.frozen, mp.mo_occ))
nmoa, nmob = nmo
if not per_kpoint:
# Depending on whether there are more occupied bands, we want to make sure that
# nmo has enough room for max(nocc) + max(nvir) number of orbitals for occupied
# and virtual space
nocca, noccb = mp.get_nocc(per_kpoint=True)
nmoa = np.amax(nocca) + np.max(np.array(nmoa) - np.array(nocca))
nmob = np.amax(noccb) + np.max(np.array(nmob) - np.array(noccb))
return nmoa, nmob
def get_nocc(mp, per_kpoint=False):
'''Number of occupied orbitals for k-point calculations.
Number of occupied orbitals for use in a calculation with k-points, taking into
account frozen orbitals.
Args:
mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
per_kpoint (bool, optional): True returns the number of occupied
orbitals at each k-point. False gives the max of this list.
Returns:
nocc (int, list of int): Number of occupied orbitals. For return type, see description of arg
`per_kpoint`.
Notes:
For specifying frozen orbitals inside mp, the following options are accepted:
+=========================+========================================+===============================+
| Argument (Example) | Argument Meaning | Example Meaning |
+=========================+========================================+===============================+
| int (1) | Freeze the same number of orbitals | Freeze one (lowest) orbital |
| | regardless of spin and/or kpt | for all kpts and spin cases |
+-------------------------+----------------------------------------+-------------------------------+
| 2-tuple of list of int | inner list: List of orbitals indices | Freeze the orbitals [0,4] for |
| ([0, 4], [0, 5, 6]) | to freeze at all kpts | spin0, and orbitals [0,5,6] |
| | outer list: Spin index | for spin1 at all kpts |
+-------------------------+----------------------------------------+-------------------------------+
| list(2) of list of list | inner list: list of orbital indices to | Freeze orbital 0 for spin0 at |
| ([[0,],[]], | freeze at each kpt for given spin | kpt0, and freeze orbital 0,1 |
| [[0,1],[4]]) | outer list: spin index | for spin1 at kpt0 and orbital |
| | | 4 at kpt1 |
+-------------------------+----------------------------------------+-------------------------------+
'''
for spin in [0, 1]:
for i, moocc in enumerate(mp.mo_occ[spin]):
if np.any(moocc % 1 != 0):
raise RuntimeError(
"Fractional occupation numbers encountered @ kp={:d}: {}. "
"This may have been caused by smearing of occupation numbers "
"in the mean-field calculation. If so, consider executing "
"mf.smearing_method = False; mf.mo_occ = mf.get_occ() prior "
"to calling this".format(i, moocc))
if mp._nocc is not None:
return mp._nocc
elif mp.frozen is None:
nocc = [[
np.count_nonzero(mp.mo_occ[0][k] > 0) for k in range(mp.nkpts)
], [np.count_nonzero(mp.mo_occ[1][k] > 0) for k in range(mp.nkpts)]]
elif isinstance(mp.frozen, (int, np.integer)):
nocc = [0] * 2
for spin in [0, 1]:
nocc[spin] = [
(np.count_nonzero(mp.mo_occ[spin][k] > 0) - mp.frozen)
for k in range(mp.nkpts)
]
elif (_is_arraylike(mp.frozen[0]) and
isinstance(mp.frozen[0][0],
(int, np.integer))): # case example: ([0, 4], [0, 5, 6])
nocc = [0] * 2
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
[
_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt],
ikpt) for ikpt in range(mp.nkpts)
]
nocc_spin = []
for ikpt in range(mp.nkpts):
max_occ_idx = np.max(np.where(mp.mo_occ[spin][ikpt] > 0))
frozen_nocc = np.sum(np.array(mp.frozen[spin]) <= max_occ_idx)
nocc_spin.append(
np.count_nonzero(mp.mo_occ[spin][ikpt]) - frozen_nocc)
nocc[spin] = nocc_spin
elif (_is_arraylike(mp.frozen[0]) and isinstance(
mp.frozen[0][0],
(list, np.ndarray))): # case example: ([[0,],[]], [[0,1],[4]])
assert (len(mp.frozen) == 2)
for spin in [0, 1]:
nkpts = len(mp.frozen[spin])
if nkpts != mp.nkpts:
raise RuntimeError(
'Frozen list has a different number of k-points (length) than passed in'
'mean-field/correlated calculation. \n\nCalculation nkpts = %d, frozen'
'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
nocc = [0] * 2
for spin in [0, 1]:
[
_frozen_sanity_check(frozen, mo_occ, ikpt) for ikpt, frozen,
mo_occ in zip(range(nkpts), mp.frozen[spin], mp.mo_occ[spin])
]
nocc_spin = []
for ikpt, frozen in enumerate(mp.frozen[spin]):
max_occ_idx = np.max(np.where(mp.mo_occ[spin][ikpt] > 0))
frozen_nocc = np.sum(np.array(frozen) <= max_occ_idx)
nocc_spin.append(
np.count_nonzero(mp.mo_occ[spin][ikpt]) - frozen_nocc)
nocc[spin] = nocc_spin
else:
raise NotImplementedError('No known conversion for frozen %s' %
mp.frozen)
for spin in [0, 1]:
assert any(np.array(nocc[spin]) > 0), (
'Must have occupied orbitals (spin=%d)! \n\nnocc %s\nfrozen %s\nmo_occ %s'
% (spin, nocc, mp.frozen, mp.mo_occ))
nocca, noccb = nocc
if not per_kpoint:
nocca = np.amax(nocca)
noccb = np.amax(noccb)
return nocca, noccb
