def e_xc_wb97mv_unpolarized(rho,
sigma,
tau,
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 unpolarized case.
10.1063/1.4952647
Args:
rho: Float numpy array with shape (num_grids,), the electron density.
sigma: Float numpy array with shape (num_grids,), the norm square of
density gradient.
tau: Float numpy array with shape (num_grids,), the 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.
"""
x = gga.get_reduced_density_gradient(rho, sigma)
t = get_mgga_t(rho, tau, polarized=False)
e_x_lda = lda.e_x_lda_unpolarized(rho)
e_c_lda_ss, e_c_lda_os = lda.decomposed_e_c_lda_unpolarized(rho)
f_x = f_b97m(
x,
t,
power_series=power_series_x,
gamma=gamma_x,
polarized=False)
f_c_ss = f_b97m(
x,
t,
power_series=power_series_ss,
gamma=gamma_ss,
polarized=False)
f_c_os = f_b97m(
x,
t,
power_series=power_series_os,
gamma=gamma_os,
polarized=False)
return (rsh.f_rsh(rho, omega=omega, polarized=False) * e_x_lda * f_x
+ e_c_lda_ss * f_c_ss + e_c_lda_os * f_c_os)
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 test_get_signature_of_lda(self, rho_index):
rho = xc_functionals.XCFunctional._signature_features['rho'][rho_index]
expected_e_xc = (lda.e_x_lda_unpolarized(2 * rho) # 2: spin factor
+ lda.e_c_lda_unpolarized(2 * rho))
mock_random_state = mock.Mock()
mock_random_state.choice = mock.Mock(side_effect=[[rho_index]])
e_xc = xc_functionals.lda_functional.get_signature(
{}, {}, {},
num_feature_samples=1,
signature='e_xc',
random_state=mock_random_state)
np.testing.assert_allclose(e_xc, expected_e_xc)
def e_x_unpolarized(rho, sigma, **kwargs):
"""Evaluates exchange energy density for spin unpolarized case.
Args:
rho: Float numpy array with shape (num_grids,), the electron density.
sigma: Float numpy array with shape (num_grids,), the norm square of
density gradient.
**kwargs: Dict, extra arguments for f_x.
Returns:
Float numpy array with shape (num_grids,), the exchange energy density.
"""
x = get_reduced_density_gradient(rho, sigma)
return lda.e_x_lda_unpolarized(rho) * f_x(x, **kwargs)
def _make_evaluator(num_mols, num_targets, eval_mode):
"""Makes an evaluator instance for testing."""
# B97 nonlinear parameters
gamma_x, gamma_css, gamma_cos = 0.004, 0.2, 0.006
num_grids_for_mols = [np.random.randint(5, 10) for _ in range(num_mols)]
grids_indices = [0] + list(np.cumsum(np.array(num_grids_for_mols)))
num_grids_all = np.sum(num_grids_for_mols)
rho = np.random.rand(num_grids_all)
sigma = np.random.rand(num_grids_all)
x = gga.get_reduced_density_gradient(rho, sigma)
e_lda_x = lda.e_x_lda_unpolarized(rho)
e_lda_css, e_lda_cos = lda.decomposed_e_c_lda_unpolarized(rho)
u_x = gga.u_b97(x, gamma=gamma_x)
u_css = gga.u_b97(x, gamma=gamma_css)
u_cos = gga.u_b97(x, gamma=gamma_cos)
rho_weights = np.random.rand(num_grids_all)
formula_matrix = np.random.rand(num_targets, num_mols)
sample_weights = np.random.rand(num_targets)
expected_exc_weighted = rho_weights * (gga.e_x_b97_unpolarized(
rho, sigma) + gga.e_c_b97_unpolarized(rho, sigma))
targets = formula_matrix @ jnp.array([
jnp.sum(expected_exc_weighted[grids_indices[i]:grids_indices[i + 1]])
for i in range(num_mols)
])
return evaluators.Evaluator(num_grids_for_mols=num_grids_for_mols,
rho_weights=rho_weights,
formula_matrix=formula_matrix,
targets=targets,
sample_weights=sample_weights,
e_lda_x=e_lda_x,
e_lda_css=e_lda_css,
e_lda_cos=e_lda_cos,
features_x={'u': u_x},
features_css={'u': u_css},
features_cos={'u': u_cos},
eval_mode=eval_mode)
def _parse_rho_and_derivs_unpolarized(rho_and_derivs, omega):
"""Computes common features and lda energies for spin unpolarized case."""
if rho_and_derivs.ndim != 2 or rho_and_derivs.shape[0] != 6:
raise ValueError(f'Wrong shape for rho_and_derivs. Expected (6, *), '
f'got {rho_and_derivs.shape}')
rho_s, grad_x_s, grad_y_s, grad_z_s, _, tau_s = rho_and_derivs
# first derivatives
sigma_ss = grad_x_s**2 + grad_y_s**2 + grad_z_s**2
x_s = gga.get_reduced_density_gradient(rho_s, sigma_ss, use_jax=False)
x2_s = x_s**2
u_s = gga.u_b97(x_s, gamma=GAMMA_X_B97, polarized=True)
del x_s
# second derivatives
tau_s = tau_s.copy()
t_s = mgga.get_mgga_t(rho_s, tau_s, polarized=True)
w_s = (t_s - 1) / (t_s + 1)
del t_s, tau_s
# LDA energy densities
e_lda_x = lda.e_x_lda_unpolarized(2 * rho_s)
if omega > 1e-8:
e_lda_x *= rsh.f_rsh(rho_s, omega=omega, polarized=True, use_jax=False)
e_lda_css, e_lda_cos = lda.decomposed_e_c_lda_unpolarized(2 * rho_s,
use_jax=False)
return {
'rho': rho_s,
'x2': x2_s,
'w': w_s,
'u': u_s,
}, {
'x': e_lda_x,
'css': e_lda_css,
'cos': e_lda_cos,
}
def xc_fun_unpolarized(rho, sigma, tau):
"""Evaluates XC energy density for spin unpolarized case.
Args:
rho: Float numpy array with shape (num_grids,), the electron density.
sigma: Float numpy array with shape (num_grids,), the norm square of
density gradient.
tau: Float numpy array with shape (num_grids,), the kinetic energy
density.
Returns:
Float numpy array with shape (num_grids,), the XC energy density.
"""
rho_s = 0.5 * rho
x_s = gga.get_reduced_density_gradient(rho_s,
0.25 * sigma,
use_jax=True)
x2_s = x_s**2
t_s = mgga.get_mgga_t(rho_s, 0.5 * tau, polarized=True)
w_s = (t_s - 1) / (t_s + 1)
e_lda_x = lda.e_x_lda_unpolarized(rho) * rsh.f_rsh(
rho, omega=omega, polarized=False, use_jax=True)
e_lda_css, e_lda_cos = lda.decomposed_e_c_lda_unpolarized(
rho, use_jax=True)
return self.eval_exc(**parameters,
e_lda_x=e_lda_x,
e_lda_css=e_lda_css,
e_lda_cos=e_lda_cos,
features={
'rho': rho_s,
'x2': x2_s,
'w': w_s
},
use_jax=True)
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))
class XCFunctional:
"""Exchange-correlation functional.
The exchange-correlation energy density e_xc is assumed to take the following
form:
e_xc(features) = e_x^LDA(rho) * f_x(features)
+ e_css^LDA(rho) * f_css(features)
+ e_cos^LDA(rho) * f_cos(features)
where e_x^LDA, e_css^LDA and e_cos are LDA exchange, same-spin correlation and
opposite-spin correlation energy densities, respectively; f_x, f_css and f_cos
are exchange, same-spin correlation and opposite-spin correlation enhancement
factors, respectively; features may include density, derivatives of density,
and/or other custom features.
"""
# representative values of common features, used to evaluate fingerprints
_signature_features = {
'rho': np.array([10., 1., 1.e-1, 1.e-2, 1.e-3, 1.e-4, 1.e-5]),
'x2': np.array([10., 100., 1000.]),
'w': np.array([-1., -0.5, 0., 0.5]),
'u': np.array([0.0, 0.5, 0.9])
}
# 2: spin factor because rho as a feature is defined as rho_sigma
_signature_e_lda_x = lda.e_x_lda_unpolarized(2 *
_signature_features['rho'])
_signature_e_lda_css, _signature_e_lda_cos = (
lda.decomposed_e_c_lda_unpolarized(2 * _signature_features['rho'],
use_jax=False))
def __init__(self, f_x=None, f_css=None, f_cos=None):
"""Initialize an XC functional.
Args:
f_x: Instance of enhancement_factors.EnhancementFactor class, the
exchange enhancement factor. If not specified, empty enhancement factor
will be used.
f_css: Instance of enhancement_factors.EnhancementFactor class, the
same-spin correlation enhancement factor. If not specified, empty
enhancement factor will be used.
f_cos: Instance of enhancement_factors.EnhancementFactor class, the
opposite-spin correlation enhancement factor. If not specified, empty
enhancement factor will be used.
"""
self.f_x = f_x or enhancement_factors.f_empty
self.f_css = f_css or enhancement_factors.f_empty
self.f_cos = f_cos or enhancement_factors.f_empty
self.feature_names = sorted(
set(self.f_x.feature_names + self.f_css.feature_names +
self.f_cos.feature_names))
self.num_parameters = (self.f_x.num_parameters +
self.f_css.num_parameters +
self.f_cos.num_parameters)
self.num_used_parameters = (self.f_x.num_used_parameters +
self.f_css.num_used_parameters +
self.f_cos.num_used_parameters)
@property
def parameters_spec(self):
"""Parameter specification of the functional."""
# NOTE(htm) parameters_spec cannot be stored as a regular data attribute of
# XCFunctional as it prevents making deepcopies of XCFunctional instances
return jax.tree_flatten({
'parameters_x': {
parameter_name: 0.
for parameter_name in sorted(self.f_x.parameter_names)
},
'parameters_css': {
parameter_name: 0.
for parameter_name in sorted(self.f_css.parameter_names)
},
'parameters_cos': {
parameter_name: 0.
for parameter_name in sorted(self.f_cos.parameter_names)
}
})[1]
def eval_exc(self,
parameters_x,
parameters_css,
parameters_cos,
e_lda_x,
e_lda_css,
e_lda_cos,
features=None,
features_x=None,
features_css=None,
features_cos=None,
use_jax=True):
"""Evalutates exchange-correlation energy density on grids.
Args:
parameters_x: Dict {parameter_name: parameter_value}, the parameters
for the exchange enhacement factor.
parameters_css: Dict {parameter_name: parameter_value}, the parameters
for the same-spin correlation enhacement factor.
parameters_cos: Dict {parameter_name: parameter_value}, the parameters
for the opposite-spin correlation enhacement factor.
e_lda_x: Float numpy array of shape (num_grids_all,), the LDA exchange
energy density on grids.
e_lda_css: Float numpy array of shape (num_grids_all,), the LDA
same-spin correlation energy density on grids.
e_lda_cos: Float numpy array of shape (num_grids_all,), the LDA
opposite-spin correlation energy density on grids.
features: Dict {feature_name: feature_value}, the features for evaluating
enhancement factors. feature_value's are float numpy array with shape
(num_grids_all,).
features_x: Dict {feature_name: feature_value}, if present, overrides
features for evaluating the exchange enhancement factor.
features_css: Dict {feature_name: feature_value}, if present, overrides
features for evaluating the same-spin correlation enhancement factor.
features_cos: Dict {feature_name: feature_value}, if present, overrides
features for evaluating the opposite-spin correlation enhancement
factor.
use_jax: Boolean, if True, use jax.numpy for calculations, otherwise use
numpy.
Returns:
Float numpy array with shape (num_grids_all,), the exchange-correlation
energy density on grids.
"""
features_x = features_x if features_x is not None else features
features_css = features_css if features_css is not None else features
features_cos = features_cos if features_cos is not None else features
return (
e_lda_x * self.f_x.eval(features_x, parameters_x, use_jax=use_jax)
+ e_lda_css *
self.f_css.eval(features_css, parameters_css, use_jax=use_jax) +
e_lda_cos *
self.f_cos.eval(features_cos, parameters_cos, use_jax=use_jax))
def eval_penalty(self, penalty_per_parameter):
"""Evaluates penalty of the functional based on number of used parameters.
Args:
penalty_per_parameter: Float, the penalty value per used parameter.
Returns:
Float, the penalty value.
"""
return penalty_per_parameter * self.num_used_parameters
def to_dict(self, parameters=None):
"""Converts the exchange-correlation functional to a dictionary.
Args:
parameters: Dict, the parameters of the functional. If present, parameters
will be included in the resulting dict. This flag is intended to serve
as a convenient option when storing a functional form together with a
set of parameters; the XCFunctional instance itself is independent of
parameters.
Returns:
Dict, the dictionary representation of exchange-correlation functional.
"""
functional_dict = {
'f_x': self.f_x.to_dict(),
'f_css': self.f_css.to_dict(),
'f_cos': self.f_cos.to_dict(),
}
if parameters is not None:
functional_dict['parameters'] = copy.deepcopy(parameters)
return functional_dict
@staticmethod
def from_dict(dictionary):
"""Loads an exchange-correlation functional from a dictionary.
Args:
dictionary: Dict, the dictionary representation of exchange-correlation
functional.
Returns:
Instance of XCFunctional, the loaded exchange-correlation functional.
"""
return XCFunctional(
f_x=enhancement_factors.EnhancementFactor.from_dict(
dictionary['f_x']),
f_css=enhancement_factors.EnhancementFactor.from_dict(
dictionary['f_css']),
f_cos=enhancement_factors.EnhancementFactor.from_dict(
dictionary['f_cos']),
)
def make_isomorphic_copy(self,
feature_names_x=None,
feature_names_css=None,
feature_names_cos=None,
num_shared_parameters=None,
num_variables=None):
"""Makes an isomorphic copy of the XCFunctional instance.
Args:
feature_names_x: List of strings, if present, specifies the features for
exchange enhancement factor.
feature_names_css: List of strings, if present, specifies the features for
same-spin correlation enhancement factor.
feature_names_cos: List of strings, if present, specifies the features for
opposite-spin correlation enhancement factor.
num_shared_parameters: Integer or sequence of 3 integers, if present,
specifies the number of shared parameters for each enhancement factor.
num_variables: Integer or sequence of 3 integers, if present, specifies
the number of variables for each enhancement factor.
Returns:
XCFunctional instance, the isomorphic copy.
"""
if num_shared_parameters is None or isinstance(num_shared_parameters,
int):
num_shared_parameters_x = num_shared_parameters
num_shared_parameters_css = num_shared_parameters
num_shared_parameters_cos = num_shared_parameters
else:
(num_shared_parameters_x, num_shared_parameters_css,
num_shared_parameters_cos) = num_shared_parameters
if num_variables is None or isinstance(num_variables, int):
num_variables_x = num_variables
num_variables_css = num_variables
num_variables_cos = num_variables
else:
num_variables_x, num_variables_css, num_variables_cos = num_variables
return XCFunctional(
f_x=self.f_x.make_isomorphic_copy(
feature_names=feature_names_x,
num_shared_parameters=num_shared_parameters_x,
num_variables=num_variables_x),
f_css=self.f_css.make_isomorphic_copy(
feature_names=feature_names_css,
num_shared_parameters=num_shared_parameters_css,
num_variables=num_variables_css),
f_cos=self.f_cos.make_isomorphic_copy(
feature_names=feature_names_cos,
num_shared_parameters=num_shared_parameters_cos,
num_variables=num_variables_cos),
)
def make_xc_fun_unpolarized(self, omega, **parameters):
"""Instantiates the functional form for spin unpolarized calculations.
Args:
omega: Float, the range separation parameter.
**parameters: Dict, the parameters for the functional form.
Returns:
Function, the resulting function that evaluates exchange correlation
energy density from rho, sigma and tau.
"""
def xc_fun_unpolarized(rho, sigma, tau):
"""Evaluates XC energy density for spin unpolarized case.
Args:
rho: Float numpy array with shape (num_grids,), the electron density.
sigma: Float numpy array with shape (num_grids,), the norm square of
density gradient.
tau: Float numpy array with shape (num_grids,), the kinetic energy
density.
Returns:
Float numpy array with shape (num_grids,), the XC energy density.
"""
rho_s = 0.5 * rho
x_s = gga.get_reduced_density_gradient(rho_s,
0.25 * sigma,
use_jax=True)
x2_s = x_s**2
t_s = mgga.get_mgga_t(rho_s, 0.5 * tau, polarized=True)
w_s = (t_s - 1) / (t_s + 1)
e_lda_x = lda.e_x_lda_unpolarized(rho) * rsh.f_rsh(
rho, omega=omega, polarized=False, use_jax=True)
e_lda_css, e_lda_cos = lda.decomposed_e_c_lda_unpolarized(
rho, use_jax=True)
return self.eval_exc(**parameters,
e_lda_x=e_lda_x,
e_lda_css=e_lda_css,
e_lda_cos=e_lda_cos,
features={
'rho': rho_s,
'x2': x2_s,
'w': w_s
},
use_jax=True)
return xc_fun_unpolarized
def make_xc_fun_polarized(self, omega, **parameters):
"""Instantiates the functional form for spin polarized calculations.
Args:
omega: Float, the range separation parameter.
**parameters: Dict, the parameters for the functional form.
Returns:
Function, the resulting function that evaluates exchange correlation
energy density from rho, sigma and tau.
"""
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))
return xc_fun_polarized
def make_eval_xc(self, omega, **parameters):
return xc.make_eval_xc_mgga(
xc_fun_unpolarized=self.make_xc_fun_unpolarized(
omega, **parameters),
xc_fun_polarized=self.make_xc_fun_polarized(omega, **parameters))
def __eq__(self, other):
return all([
self.f_x == other.f_x, self.f_css == other.f_css,
self.f_cos == other.f_cos
])
def __str__(self):
return json.dumps(self.to_dict(), indent=2)
def __repr__(self):
return self.__str__()
def get_signature(self,
parameters_x,
parameters_css,
parameters_cos,
num_feature_samples=10,
signature='e_xc',
random_state=0):
"""Computes a signature vector of the functional form with given parameters.
A random sample of features will be draw from self._signature_features,
which includes representative values of features in real DFT calculations.
The signature vector is defined as the exchange-correlation energy density
(e_xc) or exchange-correlation enhancement factor (F_xc) on the random
sample of features.
Args:
parameters_x: Dict {parameter_name: parameter_value}, the parameters
for the exchange enhacement factor.
parameters_css: Dict {parameter_name: parameter_value}, the parameters
for the same-spin correlation enhacement factor.
parameters_cos: Dict {parameter_name: parameter_value}, the parameters
for the opposite-spin correlation enhacement factor.
num_feature_samples: Integer, number of feature samples.
signature: String, the signature values to be evaluated. Possible values
are 'e_xc' for exchange-correlation energy density and 'F_xc' for
exchange-correlation enhancement factor.
random_state: Integer or instance of np.random.RandomState, the random
state used for the calculation.
Returns:
Float numpy array with shape (num_feature_samples,), the signature vector.
Raises:
ValueError, if feature names not in self._signature_features.
"""
for feature_name in self.feature_names:
if feature_name not in self._signature_features:
raise ValueError('Evaluating signature is not supported for '
f'feature {feature_name}')
if isinstance(random_state, int):
random_state = np.random.RandomState(random_state)
# randomly sample rho and corresponding lda energies
rho_sample_indices = random_state.choice(len(
self._signature_features['rho']),
size=num_feature_samples)
rho_sample = self._signature_features['rho'][rho_sample_indices]
e_lda_x_sample = self._signature_e_lda_x[rho_sample_indices]
e_lda_css_sample = self._signature_e_lda_css[rho_sample_indices]
e_lda_cos_sample = self._signature_e_lda_cos[rho_sample_indices]
# randomly sample other features
feature_samples = {'rho': rho_sample}
for feature_name in self.feature_names:
if feature_name == 'rho':
continue
feature_samples[feature_name] = random_state.choice(
self._signature_features[feature_name],
size=num_feature_samples)
e_xc_sample = self.eval_exc(parameters_x=parameters_x,
parameters_css=parameters_css,
parameters_cos=parameters_cos,
e_lda_x=e_lda_x_sample,
e_lda_css=e_lda_css_sample,
e_lda_cos=e_lda_cos_sample,
features=feature_samples)
if signature == 'e_xc':
return e_xc_sample
elif signature == 'F_xc':
return e_xc_sample / (e_lda_x_sample + e_lda_css_sample +
e_lda_cos_sample)
else:
raise ValueError(f'Unrecongnized signature flag {signature}')
def get_fingerprint(self,
num_feature_samples=10,
num_parameter_samples=10,
num_decimals=5):
"""Gets a fingerprint for the functional.
Fingerprint is evaluated as the MD5 hash value of functional singatures on
a random sample of feature and parameter values. Signatures will be
converted to strings with specified number of decimals.
Args:
num_feature_samples: Integer, number of samples of features.
num_parameter_samples: Integer, number of samples of parameters.
num_decimals: Integer, number of decimals when converting signatures to
strings.
Returns:
String, the fingerprint of the functional.
"""
format_string = f'{{:.{num_decimals}f}}'
# fix random seed to have consistent sampling behavior when running the
# code on different distributed workers
random_state = np.random.RandomState(0)
parameter_samples = random_state.rand(num_parameter_samples,
self.num_parameters)
signatures = []
for parameters in parameter_samples:
signatures.extend(
self.get_signature(**jax.tree_unflatten(
self.parameters_spec, parameters),
num_feature_samples=num_feature_samples,
random_state=random_state,
signature='e_xc'))
signature_string = ','.join(map(format_string.format, signatures))
return hashlib.md5(signature_string.encode('utf-8')).hexdigest()
def test_e_x_lda_erf_unpolarized_not_use_jax(self, rho, expected_e_x):
np.testing.assert_allclose(
lda.e_x_lda_unpolarized(rho) *
rsh.f_rsh(rho, omega=0.3, use_jax=False), expected_e_x)
def e_x_lda_erf(rho):
return lda.e_x_lda_unpolarized(rho) * rsh.f_rsh(
rho, omega=0.3, use_jax=True)