def test_decomposed_e_c_lda_polarized_not_use_jax(self, rhoa, rhob,
expected_e_c_aa,
expected_e_c_bb,
expected_e_c_ab):
e_c_aa, e_c_bb, e_c_ab = lda.decomposed_e_c_lda_polarized(
rhoa, rhob, use_jax=False)
np.testing.assert_allclose(e_c_aa, expected_e_c_aa)
np.testing.assert_allclose(e_c_bb, expected_e_c_bb)
np.testing.assert_allclose(e_c_ab, expected_e_c_ab)
def e_c_b97_polarized(rhoa,
rhob,
sigma_aa,
sigma_ab,
sigma_bb,
power_series_ss=((0, 0.1737), (1, 2.3487), (2, -2.4868)),
power_series_os=((0, 0.9454), (1, 0.7471), (2, -4.5961)),
gamma_ss=0.2,
gamma_os=0.006):
"""Evaluates B97 correlation energy density for spin polarized case.
10.1063/1.475007 Eq. 4-6.
Args:
rhoa: Float numpy array with shape (num_grids,), the spin up
electron density.
rhob: Float numpy array with shape (num_grids,), the spin down
electron density.
sigma_aa: Float numpy array with shape (num_grids,), the norm square of
density gradient (aa component).
sigma_ab: Float numpy array with shape (num_grids,), the norm square of
density gradient (ab component).
sigma_bb: Float numpy array with shape (num_grids,), the norm square of
density gradient (bb component).
power_series_ss: List of tuples of (integer, float), defines the power
series expansion of u for same-spin correlation.
power_series_os: List of tuples of (integer, float), defines the power
series expansion of u for opposite-spin correlation.
gamma_ss: Float, parameter.
gamma_os: Float, parameter.
Returns:
Float numpy array with shape (num_grids,), B97 correlation energy density.
"""
del sigma_ab
e_c_lda_aa, e_c_lda_bb, e_c_lda_ab = lda.decomposed_e_c_lda_polarized(
rhoa, rhob)
xa = get_reduced_density_gradient(rhoa, sigma_aa)
xb = get_reduced_density_gradient(rhob, sigma_bb)
xave = jnp.sqrt(0.5 * (xa**2 + xb**2))
f_c_aa = f_b97(xa,
power_series=power_series_ss,
gamma=gamma_ss,
polarized=True)
f_c_bb = f_b97(xb,
power_series=power_series_ss,
gamma=gamma_ss,
polarized=True)
f_c_ab = f_b97(xave,
power_series=power_series_os,
gamma=gamma_os,
polarized=True)
return e_c_lda_aa * f_c_aa + e_c_lda_bb * f_c_bb + e_c_lda_ab * f_c_ab
def test_parse_ks_info(self, num_mols_unpolarized, num_mols_polarized):
num_grids_unpolarized = np.random.randint(5, size=num_mols_unpolarized)
num_grids_polarized = np.random.randint(5, size=num_mols_polarized)
ks_info_unpolarized = [{
'weights': np.random.rand(num_grids),
'rho': np.random.rand(6, num_grids)
} for num_grids in num_grids_unpolarized]
ks_info_polarized = [{
'weights': np.random.rand(num_grids),
'rho': np.random.rand(2, 6, num_grids)
} for num_grids in num_grids_polarized]
# test the XC energy of a functional with enhancement factors
# F_x = F_css = F_cos = Identity(rho_sigma)
# XC energy = sum weights * e_lda * rho_sigma
expected_xc_energy = 0.
for ks_info in ks_info_unpolarized:
rho = ks_info['rho'][0, :]
expected_xc_energy += np.sum(
ks_info['weights'] * (0.5 * rho) *
(lda.e_x_lda_unpolarized(rho) + lda.e_c_lda_unpolarized(rho)))
for ks_info in ks_info_polarized:
rho_a = ks_info['rho'][0, 0, :]
rho_b = ks_info['rho'][1, 0, :]
e_lda_x_a = 0.5 * lda.e_x_lda_unpolarized(2 * rho_a)
e_lda_x_b = 0.5 * lda.e_x_lda_unpolarized(2 * rho_b)
e_lda_css_a, e_lda_css_b, e_lda_cos = lda.decomposed_e_c_lda_polarized(
rho_a, rho_b)
expected_xc_energy += np.sum(
ks_info['weights'] *
(rho_a * (e_lda_x_a + e_lda_css_a) + rho_b *
(e_lda_x_b + e_lda_css_b) + 0.5 *
(rho_a + rho_b) * e_lda_cos))
results = evaluators.parse_ks_info(
ks_info_unpolarized=ks_info_unpolarized,
ks_info_polarized=ks_info_polarized,
feature_names_x=['rho'],
feature_names_css=[],
feature_names_cos=[],
omega=0.)
xc_energy = np.sum(
results['rho_weights'] * results['features_x']['rho'] *
(results['e_lda_x'] + results['e_lda_css'] + results['e_lda_cos']))
self.assertAlmostEqual(xc_energy, expected_xc_energy)
def _parse_rho_and_derivs_polarized(rho_and_derivs, omega):
"""Computes common features and lda energies for spin polarized case."""
num_grids = rho_and_derivs.shape[-1]
if rho_and_derivs.ndim != 3 or rho_and_derivs.shape[0:2] != (2, 6):
raise ValueError(
f'Wrong shape for rho_and_derivs. Expected (2, 6, *), '
f'got {rho_and_derivs.shape}')
rho_a, grad_x_a, grad_y_a, grad_z_a, _, tau_a = rho_and_derivs[0]
rho_b, grad_x_b, grad_y_b, grad_z_b, _, tau_b = rho_and_derivs[1]
# first derivatives
sigma_aa = grad_x_a**2 + grad_y_a**2 + grad_z_a**2
sigma_bb = grad_x_b**2 + grad_y_b**2 + grad_z_b**2
x_a = gga.get_reduced_density_gradient(rho_a, sigma_aa, use_jax=False)
x_b = gga.get_reduced_density_gradient(rho_b, sigma_bb, use_jax=False)
x2_a = x_a**2
x2_b = x_b**2
u_a = gga.u_b97(x_a, gamma=GAMMA_X_B97, polarized=True)
u_b = gga.u_b97(x_b, gamma=GAMMA_X_B97, polarized=True)
u_ab = gga.u_b97(np.sqrt(0.5 * (x2_a + x2_b)),
gamma=GAMMA_X_B97,
polarized=True)
del x_a, x_b
# second derivatives
t_a = mgga.get_mgga_t(rho_a, tau_a, polarized=True)
t_b = mgga.get_mgga_t(rho_b, tau_b, polarized=True)
t_ab = 0.5 * (t_a + t_b)
w_a = (t_a - 1) / (t_a + 1)
w_b = (t_b - 1) / (t_b + 1)
w_ab = (t_ab - 1) / (t_ab + 1)
del tau_a, tau_b, t_a, t_b
# LDA energy densities
e_lda_x_a = 0.5 * lda.e_x_lda_unpolarized(2 * rho_a)
e_lda_x_b = 0.5 * lda.e_x_lda_unpolarized(2 * rho_b)
if omega > 1e-8:
e_lda_x_a *= rsh.f_rsh(rho_a,
omega=omega,
polarized=True,
use_jax=False)
e_lda_x_b *= rsh.f_rsh(rho_b,
omega=omega,
polarized=True,
use_jax=False)
e_lda_css_a, e_lda_css_b, e_lda_cos = lda.decomposed_e_c_lda_polarized(
rho_a, rho_b, use_jax=False)
return {
'rho': combine_arrays(rho_a, rho_b, 0.5 * (rho_a + rho_b)),
'x2': combine_arrays(x2_a, x2_b, 0.5 * (x2_a + x2_b)),
'w': combine_arrays(w_a, w_b, w_ab),
'u': combine_arrays(u_a, u_b, u_ab),
}, {
'x': combine_arrays(e_lda_x_a, e_lda_x_b, np.zeros(num_grids)),
'css': combine_arrays(e_lda_css_a, e_lda_css_b, np.zeros(num_grids)),
'cos': combine_arrays(np.zeros(num_grids), np.zeros(num_grids),
e_lda_cos)
}
def e_xc_wb97mv_polarized(rhoa,
rhob,
sigma_aa,
sigma_ab,
sigma_bb,
tau_a,
tau_b,
power_series_x=WB97MV_PARAMS['power_series_x'],
power_series_ss=WB97MV_PARAMS['power_series_ss'],
power_series_os=WB97MV_PARAMS['power_series_os'],
gamma_x=WB97MV_PARAMS['gamma_x'],
gamma_ss=WB97MV_PARAMS['gamma_ss'],
gamma_os=WB97MV_PARAMS['gamma_os'],
omega=WB97MV_PARAMS['omega']):
"""Evaluates wB97M-V xc energy density for spin polarized case.
10.1063/1.4952647
Args:
rhoa: Float numpy array with shape (num_grids,), the spin up electron
density.
rhob: Float numpy array with shape (num_grids,), the spin down electron
density.
sigma_aa: Float numpy array with shape (num_grids,), the norm square of
density gradient (aa component).
sigma_ab: Float numpy array with shape (num_grids,), the norm square of
density gradient (ab component).
sigma_bb: Float numpy array with shape (num_grids,), the norm square of
density gradient (bb component).
tau_a: Float numpy array with shape (num_grids,), the spin up kinetic
energy density.
tau_b: Float numpy array with shape (num_grids,), the spin down kinetic
energy density.
power_series_x: List of tuples of (integer, integer, float), defines the
power series expansion of w and u for exchange.
power_series_ss: List of tuples of (integer, integer, float), defines the
power series expansion of w and u for same-spin correlation.
power_series_os: List of tuples of (integer, integer, float), defines the
power series expansion of w and u for opposite-spin correlation.
gamma_x: Float, parameter.
gamma_ss: Float, parameter.
gamma_os: Float, parameter.
omega: Float, parameter.
Returns:
Float numpy array with shape (num_grids,), wB97M-V xc energy density.
"""
del sigma_ab
xa = gga.get_reduced_density_gradient(rhoa, sigma_aa)
xb = gga.get_reduced_density_gradient(rhob, sigma_bb)
xave = jnp.sqrt(0.5 * (xa ** 2 + xb ** 2))
ta = get_mgga_t(rhoa, tau_a, polarized=True)
tb = get_mgga_t(rhob, tau_b, polarized=True)
tave = 0.5 * (ta + tb)
# use e_x_lda_unpolarized and spin scaling to compute e_x_lda_a and e_x_lda_b
# separately
e_x_lda_a = 0.5 * lda.e_x_lda_unpolarized(2 * rhoa)
e_x_lda_b = 0.5 * lda.e_x_lda_unpolarized(2 * rhob)
e_c_lda_aa, e_c_lda_bb, e_c_lda_ab = lda.decomposed_e_c_lda_polarized(
rhoa, rhob)
f_x_a = f_b97m(
xa,
ta,
power_series=power_series_x,
gamma=gamma_x,
polarized=True)
f_x_b = f_b97m(
xb,
tb,
power_series=power_series_x,
gamma=gamma_x,
polarized=True)
f_c_aa = f_b97m(
xa,
ta,
power_series=power_series_ss,
gamma=gamma_ss,
polarized=True)
f_c_bb = f_b97m(
xb,
tb,
power_series=power_series_ss,
gamma=gamma_ss,
polarized=True)
f_c_ab = f_b97m(
xave,
tave,
power_series=power_series_os,
gamma=gamma_os,
polarized=True)
return (rsh.f_rsh(rhoa, omega=omega, polarized=True) * e_x_lda_a * f_x_a
+ rsh.f_rsh(rhob, omega=omega, polarized=True) * e_x_lda_b * f_x_b
+ e_c_lda_aa * f_c_aa + e_c_lda_bb * f_c_bb + e_c_lda_ab * f_c_ab)
def xc_fun_polarized(rho_a, rho_b, sigma_aa, sigma_ab, sigma_bb, tau_a,
tau_b):
"""Evaluates XC energy density for spin polarized case.
Args:
rho_a: Float numpy array with shape (num_grids,), the spin up electron
density.
rho_b: Float numpy array with shape (num_grids,), the spin down electron
density.
sigma_aa: Float numpy array with shape (num_grids,), the norm square of
density gradient (aa component).
sigma_ab: Float numpy array with shape (num_grids,), the norm square of
density gradient (ab component).
sigma_bb: Float numpy array with shape (num_grids,), the norm square of
density gradient (bb component).
tau_a: Float numpy array with shape (num_grids,), the spin up kinetic
energy density.
tau_b: Float numpy array with shape (num_grids,), the spin down kinetic
energy density.
Returns:
Float numpy array with shape (num_grids,), the XC energy density.
"""
del sigma_ab
rho_ab = 0.5 * (rho_a + rho_b)
x_a = gga.get_reduced_density_gradient(rho_a,
sigma_aa,
use_jax=True)
x_b = gga.get_reduced_density_gradient(rho_b,
sigma_bb,
use_jax=True)
x2_a = x_a**2
x2_b = x_b**2
x2_ab = 0.5 * (x2_a + x2_b)
t_a = mgga.get_mgga_t(rho_a, tau_a, polarized=True)
t_b = mgga.get_mgga_t(rho_b, tau_b, polarized=True)
t_ab = 0.5 * (t_a + t_b)
w_a = (t_a - 1) / (t_a + 1)
w_b = (t_b - 1) / (t_b + 1)
w_ab = (t_ab - 1) / (t_ab + 1)
e_lda_x_a = 0.5 * lda.e_x_lda_unpolarized(2 * rho_a) * rsh.f_rsh(
rho_a, omega=omega, polarized=True, use_jax=True)
e_lda_x_b = 0.5 * lda.e_x_lda_unpolarized(2 * rho_b) * rsh.f_rsh(
rho_b, omega=omega, polarized=True, use_jax=True)
e_lda_css_a, e_lda_css_b, e_lda_cos = lda.decomposed_e_c_lda_polarized(
rho_a, rho_b, use_jax=True)
features_a = {'rho': rho_a, 'x2': x2_a, 'w': w_a}
features_b = {'rho': rho_b, 'x2': x2_b, 'w': w_b}
features_ab = {'rho': rho_ab, 'x2': x2_ab, 'w': w_ab}
return (
e_lda_x_a * self.f_x.eval(
features_a, parameters['parameters_x'], use_jax=True) +
e_lda_x_b * self.f_x.eval(
features_b, parameters['parameters_x'], use_jax=True) +
e_lda_css_a * self.f_css.eval(
features_a, parameters['parameters_css'], use_jax=True) +
e_lda_css_b * self.f_css.eval(
features_b, parameters['parameters_css'], use_jax=True) +
e_lda_cos * self.f_cos.eval(
features_ab, parameters['parameters_cos'], use_jax=True))