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)