Skip to content

Commit

Permalink
finish expanding utilities of feat_normalizer module
Browse files Browse the repository at this point in the history
  • Loading branch information
kylebystrom committed May 16, 2024
1 parent 6489801 commit 09fd8cd
Show file tree
Hide file tree
Showing 2 changed files with 209 additions and 20 deletions.
40 changes: 24 additions & 16 deletions ciderpress/dft/feat_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,11 +207,14 @@ def get_invariant_normalizer_from_exponent_params(power, a0, tau_mul):
return InhomogeneityNormalizer(const1, const2, power)


def get_normalizer_from_exponent_params(rho_pow, exp_pow, a0, tau_mul):
def get_normalizer_from_exponent_params(rho_pow, exp_pow, a0, tau_mul, gga=False):
power1 = rho_pow + 2.0 / 3 * exp_pow
power2 = exp_pow
tau_fac = tau_mul * 1.2 * (6 * np.pi**2) ** (2.0 / 3) / np.pi
B = np.pi / 2 ** (2.0 / 3) * (a0 - tau_fac)
if gga:
B = np.pi / 2 ** (2.0 / 3) * a0
else:
B = np.pi / 2 ** (2.0 / 3) * (a0 - tau_fac)
C = np.pi / 2 ** (2.0 / 3) * tau_fac # / CFC
const2 = C / B
const1 = B**exp_pow
Expand Down Expand Up @@ -277,15 +280,24 @@ def _get_rho_and_inh(self, X0T):
return rho_term, inh_term

def _get_drho_and_dinh(self, X0T, DX0T):
rho = X0T[0]
drho = DX0T[0]
grad_term = DX0T[1]
tau_term = DX0T[2]
grad = X0T[1]
dgrad = DX0T[1]
if self.slmode == "npa":
dinh = tau_term + 5.0 / 3 * grad_term
tau = X0T[2]
dtau = DX0T[2]
dinh = dtau + 5.0 / 3 * dgrad
elif self.slmode == "nst":
dinh = 0 # TODO
tau = X0T[2]
dtau = DX0T[2]
dinh = dtau / (CFC * rho ** (5.0 / 3))
dinh -= 5.0 / 3 * tau / (CFC * rho ** (8.0 / 3)) * drho
elif self.slmode == "np":
dinh = 5.0 / 3 * grad_term
dinh = 5.0 / 3 * dgrad
else:
dinh = dgrad / (8 * CFC * rho ** (8.0 / 3))
dinh -= 8.0 / 3 * grad / (8 * CFC * rho ** (11.0 / 3)) * drho
return drho, dinh

def get_normalized_feature_vector(self, X0T):
Expand All @@ -309,13 +321,7 @@ def get_derivative_of_normed_features(self, X0T, DX0T):
DX0TN = np.empty_like(DX0T)
rho, inh = self._get_rho_and_inh(X0T[None, :])
rho, inh = rho[0], inh[0]
drho = DX0T[0]
grad_term = DX0T[1]
tau_term = DX0T[2]
if self.slmode == "npa":
dinh = tau_term + 5.0 / 3 * grad_term
else:
dinh = tau_term
drho, dinh = self._get_drho_and_dinh(X0T, DX0T)
for i in range(self.nfeat):
if self[i] is not None:
DX0TN[i] = self[i].get_normed_feature_deriv(
Expand Down Expand Up @@ -353,11 +359,13 @@ def get_derivative_wrt_unnormed_features(self, X0T, df_dX0TN):
df_dX0T[:, 1] += 5.0 / 3 * dfdinh
df_dX0T[:, 2] += dfdinh
elif self.slmode == "nst":
df_dX0T[:, 2] += dfdinh
df_dX0T[:, 0] -= dfdinh * 5.0 / 3 * inh_term / rho_term
df_dX0T[:, 2] += dfdinh / (CFC * rho_term ** (5.0 / 3))
elif self.slmode == "np":
df_dX0T[:, 1] += 5.0 / 3 * dfdinh
else:
raise NotImplementedError
df_dX0T[:, 0] -= dfdinh * 8.0 / 3 * inh_term / rho_term
df_dX0T[:, 1] += dfdinh / (8 * CFC * rho_term ** (8.0 / 3))
for s in range(cond.shape[0]):
df_dX0T[s, ..., cond[s]] = 0.0
return df_dX0T
Expand Down
189 changes: 185 additions & 4 deletions ciderpress/dft/tests/test_feat_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ def setUpClass(cls):
assert rho_data.shape[0] == 5
cls.rho = rho_data
sigma = np.einsum("xg,xg->g", drho, drho)
s2 = get_s2(rho, sigma)
alpha = get_alpha(rho, sigma, tau)
cls.norm_list = FeatNormalizerList(
[
None,
Expand All @@ -63,15 +65,36 @@ def setUpClass(cls):
np.random.seed(72)
feat = np.random.normal(size=(cls.norm_list.nfeat - 3, N))
dfeat = np.random.normal(size=(cls.norm_list.nfeat, N))
cls.x = np.stack(
dfeat_np = np.concatenate([dfeat[:2], dfeat[3:]], axis=0)
dfeat_nst = dfeat.copy()
dfeat_ns = dfeat_np.copy()
C = 2 * (3 * np.pi**2) ** (1.0 / 3)
CF = 0.3 * (3 * np.pi**2) ** (2.0 / 3)
dsigma = rho ** (8.0 / 3) * dfeat[1]
dsigma += 8.0 / 3 * rho ** (5.0 / 3) * s2 * dfeat[0]
dsigma *= C * C
dtau = 5.0 / 3 * rho ** (2.0 / 3) * (alpha + 5.0 / 3 * s2) * dfeat[0]
dtau += 5.0 / 3 * rho ** (5.0 / 3) * dfeat[1]
dtau += rho ** (5.0 / 3) * dfeat[2]
dtau *= CF
dfeat_ns[1] = dsigma
dfeat_nst[1] = dsigma
dfeat_nst[2] = dtau
feat_list = [feat[i] for i in range(feat.shape[0])]
cls.x = cls.x_npa = np.stack(
[
rho,
get_s2(rho, sigma),
get_alpha(rho, sigma, tau),
s2,
alpha,
]
+ [feat[i] for i in range(feat.shape[0])]
+ feat_list
)[None, :, :]
cls.dx = dfeat
cls.dx_1e = cls.dx.copy()
cls.dx_1e[2] = 0
cls.dx_np = dfeat_np
cls.dx_ns = dfeat_ns
cls.dx_nst = dfeat_nst
cls.ref_norm_list = rho_normalizers.FeatNormalizerList(
[
rho_normalizers.ConstantNormalizer(3.14),
Expand All @@ -82,6 +105,50 @@ def setUpClass(cls):
),
]
)
cls.x_np = np.stack([rho, s2] + feat_list)[None, :, :]
cls.x_1e = np.stack([rho, s2, np.zeros_like(rho)] + feat_list)[None, :, :]
cls.x_nst = np.stack([rho, sigma, tau] + feat_list)[None, :, :]
cls.x_ns = np.stack([rho, sigma] + feat_list)[None, :, :]
cls.norm_list_np = FeatNormalizerList(
[
None,
None,
ConstantNormalizer(3.14),
DensityNormalizer(1.2, 0.75),
get_normalizer_from_exponent_params(0.0, 1.5, 1.0, 0.03125),
get_normalizer_from_exponent_params(
0.75, 1.5, 1.0 * 1.2, 0.03125 * 1.2
),
],
slmode="np",
)
cls.norm_list_nst = FeatNormalizerList(
[
None,
DensityNormalizer(0.25 / (3 * np.pi**2) ** (2.0 / 3), -8.0 / 3),
DensityNormalizer(10.0 / 3 / (3 * np.pi**2) ** (2.0 / 3), -5.0 / 3),
ConstantNormalizer(3.14),
DensityNormalizer(1.2, 0.75),
get_normalizer_from_exponent_params(0.0, 1.5, 1.0, 0.03125),
get_normalizer_from_exponent_params(
0.75, 1.5, 1.0 * 1.2, 0.03125 * 1.2
),
],
slmode="nst",
)
cls.norm_list_ns = FeatNormalizerList(
[
None,
DensityNormalizer(0.25 / (3 * np.pi**2) ** (2.0 / 3), -8.0 / 3),
ConstantNormalizer(3.14),
DensityNormalizer(1.2, 0.75),
get_normalizer_from_exponent_params(0.0, 1.5, 1.0, 0.03125),
get_normalizer_from_exponent_params(
0.75, 1.5, 1.0 * 1.2, 0.03125 * 1.2
),
],
slmode="ns",
)

def test_evaluate(self):
rho_data = self.rho.copy()
Expand Down Expand Up @@ -149,6 +216,120 @@ def test_evaluate(self):
for i in range(scaled_ref.shape[1]):
assert_almost_equal(scaled_feat[:, i], scaled_ref[:, i])

def test_evaluate_nst(self):
feat_ref = self.norm_list.get_normalized_feature_vector(self.x)
feat_ref[:, 2] += 5.0 / 3 * feat_ref[:, 1]
feat = self.norm_list_nst.get_normalized_feature_vector(self.x_nst)
dfeat_ref = self.norm_list.get_derivative_of_normed_features(self.x[0], self.dx)
dfeat_ref[2] += 5.0 / 3 * dfeat_ref[1]
dfeat = self.norm_list_nst.get_derivative_of_normed_features(
self.x_nst[0], self.dx_nst
)
featp = self.norm_list_nst.get_normalized_feature_vector(
self.x_nst + DELTA * self.dx_nst
)
featm = self.norm_list_nst.get_normalized_feature_vector(
self.x_nst - DELTA * self.dx_nst
)
dfeat_ref2 = ((featp - featm) / (2 * DELTA))[0]
for i in range(self.norm_list.nfeat):
assert_almost_equal(feat[:, i], feat_ref[:, i])
assert_almost_equal(dfeat[i], dfeat_ref[i])
assert_almost_equal(dfeat[i], dfeat_ref2[i])
v = 2 * feat
vfeat = self.norm_list_nst.get_derivative_wrt_unnormed_features(self.x_nst, v)
for i in range(self.norm_list.nfeat):
x = self.x_nst.copy()
x[:, i] += 0.5 * DELTA
featp = self.norm_list_nst.get_normalized_feature_vector(x)
ep = (featp**2).sum(axis=1)
x[:, i] -= DELTA
featm = self.norm_list_nst.get_normalized_feature_vector(x)
em = (featm**2).sum(axis=1)
assert_almost_equal(vfeat[:, i], (ep - em) / DELTA, 5)
usps = self.norm_list_nst.get_usps()
usps_ref = self.norm_list.get_usps()
muls = [0, 8, 5, 0, 0, 0]
for usp, usp_ref, mul in zip(usps, usps_ref, muls):
assert usp + mul == usp_ref, "{} {} {}".format(usp, usp_ref, mul)

def test_evaluate_np(self):
feat_ref = self.norm_list.get_normalized_feature_vector(self.x_1e)
assert_almost_equal(feat_ref[:, 2], 0)
feat_ref = feat_ref[:, [0, 1, 3, 4, 5, 6]]
feat = self.norm_list_np.get_normalized_feature_vector(self.x_np)
dfeat_ref = self.norm_list.get_derivative_of_normed_features(
self.x_1e[0], self.dx_1e
)
dfeat_ref = dfeat_ref[[0, 1, 3, 4, 5, 6]]
dfeat = self.norm_list_np.get_derivative_of_normed_features(
self.x_np[0], self.dx_np
)
featp = self.norm_list_np.get_normalized_feature_vector(
self.x_np + DELTA * self.dx_np
)
featm = self.norm_list_np.get_normalized_feature_vector(
self.x_np - DELTA * self.dx_np
)
dfeat_ref2 = ((featp - featm) / (2 * DELTA))[0]
for i in range(self.norm_list_np.nfeat):
assert_almost_equal(feat[:, i], feat_ref[:, i])
assert_almost_equal(dfeat[i], dfeat_ref[i])
assert_almost_equal(dfeat[i], dfeat_ref2[i])
v = 2 * feat
vfeat = self.norm_list_np.get_derivative_wrt_unnormed_features(self.x_np, v)
for i in range(self.norm_list_np.nfeat):
x = self.x_np.copy()
x[:, i] += 0.5 * DELTA
featp = self.norm_list_np.get_normalized_feature_vector(x)
ep = (featp**2).sum(axis=1)
x[:, i] -= DELTA
featm = self.norm_list_np.get_normalized_feature_vector(x)
em = (featm**2).sum(axis=1)
assert_almost_equal(vfeat[:, i], (ep - em) / DELTA, 5)
usps = self.norm_list_np.get_usps()
usps_ref = np.array(self.norm_list.get_usps())[[0, 1, 3, 4, 5, 6]]
muls = [0, 0, 0, 0, 0, 0]
for usp, usp_ref, mul in zip(usps, usps_ref, muls):
assert usp + mul == usp_ref, "{} {} {}".format(usp, usp_ref, mul)

def test_evaluate_ns(self):
feat_ref = self.norm_list_np.get_normalized_feature_vector(self.x_np)
feat = self.norm_list_ns.get_normalized_feature_vector(self.x_ns)
dfeat_ref = self.norm_list_np.get_derivative_of_normed_features(
self.x_np[0], self.dx_np
)
dfeat = self.norm_list_ns.get_derivative_of_normed_features(
self.x_ns[0], self.dx_ns
)
featp = self.norm_list_ns.get_normalized_feature_vector(
self.x_ns + DELTA * self.dx_ns
)
featm = self.norm_list_ns.get_normalized_feature_vector(
self.x_ns - DELTA * self.dx_ns
)
dfeat_ref2 = ((featp - featm) / (2 * DELTA))[0]
for i in range(self.norm_list_np.nfeat):
assert_almost_equal(feat[:, i], feat_ref[:, i])
assert_almost_equal(dfeat[i], dfeat_ref[i])
assert_almost_equal(dfeat[i], dfeat_ref2[i])
v = 2 * feat
vfeat = self.norm_list_ns.get_derivative_wrt_unnormed_features(self.x_ns, v)
for i in range(self.norm_list_np.nfeat):
x = self.x_ns.copy()
x[:, i] += 0.5 * DELTA
featp = self.norm_list_ns.get_normalized_feature_vector(x)
ep = (featp**2).sum(axis=1)
x[:, i] -= DELTA
featm = self.norm_list_ns.get_normalized_feature_vector(x)
em = (featm**2).sum(axis=1)
assert_almost_equal(vfeat[:, i], (ep - em) / DELTA, 5)
usps = self.norm_list_ns.get_usps()
usps_ref = self.norm_list_np.get_usps()
muls = [0, 8, 0, 0, 0]
for usp, usp_ref, mul in zip(usps, usps_ref, muls):
assert usp + mul == usp_ref, "{} {} {}".format(usp, usp_ref, mul)


if __name__ == "__main__":
unittest.main()

0 comments on commit 09fd8cd

Please sign in to comment.