diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 1a9b2c38d0..d1710a2182 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -171,12 +171,12 @@ def __init__( ) self.starter_learning_rate = starter_learning_rate - self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference - self.has_f = (start_pref_f != 0.0 and limit_pref_f != 0.0) or inference - self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference - self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference - self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference - self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0 + self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference + self.has_f = (start_pref_f != 0.0 or limit_pref_f != 0.0) or inference + self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference + self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference + self.has_pf = (start_pref_pf != 0.0 or limit_pref_pf != 0.0) or inference + self.has_gf = start_pref_gf != 0.0 or limit_pref_gf != 0.0 self.start_pref_e = start_pref_e self.limit_pref_e = limit_pref_e @@ -304,16 +304,11 @@ def forward( l1_ener_loss = F.l1_loss( energy_pred.reshape([-1]), energy_label.reshape([-1]), - reduction="sum", + reduction="mean", ) - loss += pref_e * l1_ener_loss + loss += atom_norm * (pref_e * l1_ener_loss) more_loss["mae_e"] = self.display_if_exist( - F.l1_loss( - energy_pred.reshape([-1]), - energy_label.reshape([-1]), - reduction="mean", - ).detach(), - find_energy, + l1_ener_loss.detach() * atom_norm, find_energy ) # more_loss['log_keys'].append('rmse_e') else: @@ -367,11 +362,10 @@ def forward( rmse_f.detach(), find_force ) elif self.loss_func == "mae": - l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean") + l1_force_loss = paddle.mean(paddle.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist( l1_force_loss.detach(), find_force ) - l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum() loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION) else: raise NotImplementedError( @@ -386,7 +380,8 @@ def forward( if self.has_pf and "atom_pref" in label: atom_pref = label["atom_pref"] find_atom_pref = label.get("find_atom_pref", 0.0) - pref_pf = pref_pf * find_atom_pref + effective_find_pf = find_force * find_atom_pref + pref_pf = pref_pf * effective_find_pf atom_pref_reshape = atom_pref.reshape([-1]) if self.loss_func == "mse": @@ -395,18 +390,18 @@ def forward( ).mean() if not self.inference: more_loss["l2_pref_force_loss"] = self.display_if_exist( - l2_pref_force_loss.detach(), find_atom_pref + l2_pref_force_loss.detach(), effective_find_pf ) loss += (pref_pf * l2_pref_force_loss).to(GLOBAL_PD_FLOAT_PRECISION) rmse_pf = l2_pref_force_loss.sqrt() more_loss["rmse_pf"] = self.display_if_exist( - rmse_pf.detach(), find_atom_pref + rmse_pf.detach(), effective_find_pf ) elif self.loss_func == "mae": l1_pref_force_loss = (paddle.abs(diff_f) * atom_pref_reshape).mean() loss += (pref_pf * l1_pref_force_loss).to(GLOBAL_PD_FLOAT_PRECISION) more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss.detach(), find_atom_pref + l1_pref_force_loss.detach(), effective_find_pf ) else: raise NotImplementedError( @@ -416,7 +411,8 @@ def forward( if self.has_gf and "drdq" in label: drdq = label["drdq"] find_drdq = label.get("find_drdq", 0.0) - pref_gf = pref_gf * find_drdq + effective_find_gf = find_force * find_drdq + pref_gf = pref_gf * effective_find_gf force_reshape_nframes = force_pred.reshape([-1, natoms * 3]) force_label_reshape_nframes = force_label.reshape([-1, natoms * 3]) drdq_reshape = drdq.reshape( @@ -440,34 +436,47 @@ def forward( l2_gen_force_loss = paddle.square(diff_gen_force).mean() if not self.inference: more_loss["l2_gen_force_loss"] = self.display_if_exist( - l2_gen_force_loss.detach(), find_drdq + l2_gen_force_loss.detach(), effective_find_gf ) loss += (pref_gf * l2_gen_force_loss).to(GLOBAL_PD_FLOAT_PRECISION) rmse_gf = l2_gen_force_loss.sqrt() more_loss["rmse_gf"] = self.display_if_exist( - rmse_gf.detach(), find_drdq + rmse_gf.detach(), effective_find_gf ) if self.has_v and "virial" in model_pred and "virial" in label: find_virial = label.get("find_virial", 0.0) pref_v = pref_v * find_virial diff_v = label["virial"] - model_pred["virial"].reshape([-1, 9]) - l2_virial_loss = paddle.mean(paddle.square(diff_v)) - if not self.inference: - more_loss["l2_virial_loss"] = self.display_if_exist( - l2_virial_loss.detach(), find_virial + if self.loss_func == "mse" or self.use_huber: + l2_virial_loss = paddle.mean(paddle.square(diff_v)) + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + l2_virial_loss.detach(), find_virial + ) + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * model_pred["virial"].reshape([-1]), + atom_norm * label["virial"].reshape([-1]), + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_loss + rmse_v = l2_virial_loss.sqrt() * atom_norm + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial + ) + elif self.loss_func == "mae": + l1_virial_loss = paddle.mean(paddle.abs(diff_v)) + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss.detach() * atom_norm, find_virial ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) else: - l_huber_loss = custom_huber_loss( - atom_norm * model_pred["virial"].reshape([-1]), - atom_norm * label["virial"].reshape([-1]), - delta=self._huber_delta_virial, + raise NotImplementedError( + f"Loss type {self.loss_func} is not implemented for virial loss." ) - loss += pref_v * l_huber_loss - rmse_v = l2_virial_loss.sqrt() * atom_norm - more_loss["rmse_v"] = self.display_if_exist(rmse_v.detach(), find_virial) if mae: mae_v = paddle.mean(paddle.abs(diff_v)) * atom_norm more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial) @@ -534,7 +543,7 @@ def label_requirement(self) -> list[DataRequirementItem]: high_prec=True, ) ) - if self.has_f: + if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf: label_requirement.append( DataRequirementItem( "force", diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index b5ba880c93..5bc610931d 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -84,10 +84,10 @@ def __init__( and self.limit_pref_acdf >= 0.0 ), "Can not assign negative weight to `pref` and `pref_atomic`" - self.has_dos = (start_pref_dos != 0.0 and limit_pref_dos != 0.0) or inference - self.has_cdf = (start_pref_cdf != 0.0 and limit_pref_cdf != 0.0) or inference - self.has_ados = (start_pref_ados != 0.0 and limit_pref_ados != 0.0) or inference - self.has_acdf = (start_pref_acdf != 0.0 and limit_pref_acdf != 0.0) or inference + self.has_dos = (start_pref_dos != 0.0 or limit_pref_dos != 0.0) or inference + self.has_cdf = (start_pref_cdf != 0.0 or limit_pref_cdf != 0.0) or inference + self.has_ados = (start_pref_ados != 0.0 or limit_pref_ados != 0.0) or inference + self.has_acdf = (start_pref_acdf != 0.0 or limit_pref_acdf != 0.0) or inference assert self.has_dos or self.has_cdf or self.has_ados or self.has_acdf, ( AssertionError( diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 522dceafda..107654bcfc 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -160,12 +160,12 @@ def __init__( self.loss_func = loss_func self.starter_learning_rate = starter_learning_rate - self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference - self.has_f = (start_pref_f != 0.0 and limit_pref_f != 0.0) or inference - self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference - self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference - self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference - self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0 + self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference + self.has_f = (start_pref_f != 0.0 or limit_pref_f != 0.0) or inference + self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference + self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference + self.has_pf = (start_pref_pf != 0.0 or limit_pref_pf != 0.0) or inference + self.has_gf = start_pref_gf != 0.0 or limit_pref_gf != 0.0 # Hessian labels scale quadratically with the atom count. Unlike the # linear-size labels above, do not request them merely because a mock # inference loss is collecting requirements for ``dp change-bias``. diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index f1ce451e20..bd7e7c9a7b 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -122,11 +122,11 @@ def __init__( self.loss_func = loss_func self.starter_learning_rate = starter_learning_rate - self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference - self.has_fr = (start_pref_fr != 0.0 and limit_pref_fr != 0.0) or inference - self.has_fm = (start_pref_fm != 0.0 and limit_pref_fm != 0.0) or inference - self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference - self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference + self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference + self.has_fr = (start_pref_fr != 0.0 or limit_pref_fr != 0.0) or inference + self.has_fm = (start_pref_fm != 0.0 or limit_pref_fm != 0.0) or inference + self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference + self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference self.start_pref_e = start_pref_e self.limit_pref_e = limit_pref_e @@ -597,6 +597,18 @@ def label_requirement(self) -> list[DataRequirementItem]: high_prec=False, ) ) + if self.enable_atom_ener_coeff: + label_requirement.append( + DataRequirementItem( + "atom_ener_coeff", + ndof=1, + atomic=True, + must=False, + high_prec=False, + default=1.0, + source_policy="default", + ) + ) return label_requirement def serialize(self) -> dict: diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 715640653f..0becc2035d 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -1428,6 +1428,8 @@ def _load_single_data( else: # For non-atomic data, shape should be [ndof] data = np.full([ndof], vv["default"], dtype=dtype) + if vv["repeat"] != 1 and vv.get("special_shape") != "frame_major": + data = np.repeat(data, vv["repeat"], axis=-1) return np.float32(0.0), data # Branch 2: Data loading @@ -1526,6 +1528,8 @@ def _load_single_data( else: data = data.reshape([ndof]) + if vv["repeat"] != 1: + data = np.repeat(data, vv["repeat"], axis=-1) # Atomic: return [natoms, ndof] or flattened hessian above # Non-atomic: return [ndof] return np.float32(1.0), data diff --git a/source/tests/common/test_single_frame_repeat.py b/source/tests/common/test_single_frame_repeat.py new file mode 100644 index 0000000000..84f2900c70 --- /dev/null +++ b/source/tests/common/test_single_frame_repeat.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +from deepmd.utils.data import ( + DeepmdData, +) + + +@pytest.mark.parametrize("present", [False, True]) +@pytest.mark.parametrize("nframes", [1, 2]) +@pytest.mark.parametrize("backend", ["torch", "paddle"]) +def test_single_frame_repeat_matches_batch(tmp_path, present, nframes, backend): + """Single-frame and batch loaders repeat atomic labels identically.""" + np.savetxt(tmp_path / "type.raw", [1, 0], fmt="%d") + set_path = tmp_path / "set.000" + set_path.mkdir() + np.save(set_path / "coord.npy", np.zeros((nframes, 6))) + np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (nframes, 1))) + if present: + np.save(set_path / "atom_pref.npy", np.tile([[2.0, 3.0]], (nframes, 1))) + data = DeepmdData(str(tmp_path)) + data.add("atom_pref", 1, atomic=True, repeat=3, default=1.0) + expected = data._load_set(set_path)["atom_pref"][0] + frame = getattr(data, f"get_item_{backend}")(0) + np.testing.assert_array_equal(frame["atom_pref"].reshape(-1), expected) + assert frame["find_atom_pref"] == float(present) diff --git a/source/tests/pd/test_loss_prefactor_schedule_paddle.py b/source/tests/pd/test_loss_prefactor_schedule_paddle.py new file mode 100644 index 0000000000..c28f759b2c --- /dev/null +++ b/source/tests/pd/test_loss_prefactor_schedule_paddle.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import paddle +import pytest + +from deepmd.pd.loss.ener import ( + EnergyStdLoss, +) + + +@pytest.mark.parametrize( + "start,limit", [(0.0, 1.0), (1.0, 0.0), (0.0, 0.0), (1.0, 1.0)] +) +@pytest.mark.parametrize( + "term,labels", + [ + ("e", {"energy"}), + ("f", {"force"}), + ("v", {"virial"}), + ("ae", {"atom_ener"}), + ("pf", {"force", "atom_pref"}), + ("gf", {"force", "drdq"}), + ], +) +def test_energy_schedule_label_requirements(term, labels, start, limit): + """Request the labels needed whenever a scheduled term is active.""" + loss = EnergyStdLoss( + **{f"start_pref_{term}": start, f"limit_pref_{term}": limit}, + numb_generalized_coord=1, + ) + required = {item.key for item in loss.label_requirement} + assert required == (labels if start or limit else set()) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("learning_rate", [1.0, 0.5, 0.0]) +def test_energy_schedule_contributes_to_loss_and_gradient(start, limit, learning_rate): + """Apply the interpolated energy weight to the loss and gradient.""" + loss_fn = EnergyStdLoss(start_pref_e=start, limit_pref_e=limit) + prediction = paddle.to_tensor([[2.0]], dtype="float64", stop_gradient=False) + label = {"energy": paddle.zeros_like(prediction), "find_energy": 1.0} + _, loss, _ = loss_fn({}, lambda: {"energy": prediction}, label, 2, learning_rate) + weight = limit + (start - limit) * learning_rate + np.testing.assert_allclose(loss.numpy(), 2.0 * weight) + (gradient,) = paddle.grad(loss, prediction) + np.testing.assert_allclose(gradient.numpy(), [[2.0 * weight]]) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("term", ["f", "pf", "gf"]) +def test_force_schedule_contributes_to_loss_and_gradient(term, start, limit): + """Apply force, prefactor-force and generalized-force schedules.""" + loss_fn = EnergyStdLoss( + **{f"start_pref_{term}": start, f"limit_pref_{term}": limit}, + numb_generalized_coord=1, + ) + prediction = paddle.ones((1, 1, 3), dtype="float64") + prediction.stop_gradient = False + label = { + "force": paddle.zeros_like(prediction), + "find_force": 1.0, + "atom_pref": paddle.ones_like(prediction), + "find_atom_pref": 1.0, + "drdq": paddle.to_tensor([[[1.0], [0.0], [0.0]]], dtype="float64"), + "find_drdq": 1.0, + } + _, loss, _ = loss_fn({}, lambda: {"force": prediction}, label, 1, 0.25) + weight = limit + (start - limit) * 0.25 + np.testing.assert_allclose(loss.numpy(), weight) + expected = [2.0, 0.0, 0.0] if term == "gf" else [2.0 / 3] * 3 + (gradient,) = paddle.grad(loss, prediction) + np.testing.assert_allclose(gradient.numpy(), np.asarray([[expected]]) * weight) + + +@pytest.mark.parametrize("term", ["pf", "gf"]) +def test_force_terms_ignore_missing_force_labels(term): + """Do not train force terms against a missing-label placeholder.""" + loss_fn = EnergyStdLoss( + **{f"start_pref_{term}": 1.0, f"limit_pref_{term}": 1.0}, + numb_generalized_coord=1, + ) + prediction = paddle.ones((1, 1, 3), dtype="float64") + prediction.stop_gradient = False + label = { + "force": paddle.zeros_like(prediction), + "find_force": 0.0, + "atom_pref": paddle.ones_like(prediction), + "find_atom_pref": 1.0, + "drdq": paddle.to_tensor([[[1.0], [0.0], [0.0]]], dtype="float64"), + "find_drdq": 1.0, + } + _, loss, metrics = loss_fn({}, lambda: {"force": prediction}, label, 1, 0.25) + np.testing.assert_allclose(loss.numpy(), 0.0) + (gradient,) = paddle.grad(loss, prediction) + np.testing.assert_allclose(gradient.numpy(), 0.0) + assert np.isnan(metrics[f"rmse_{term}"].numpy()) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("term,key,width", [("e", "energy", 1), ("v", "virial", 9)]) +def test_mae_extensive_terms_are_normalized(term, key, width, start, limit): + """Normalize MAE energy and virial over the batch and atoms.""" + loss_fn = EnergyStdLoss( + **{f"start_pref_{term}": start, f"limit_pref_{term}": limit}, loss_func="mae" + ) + prediction = paddle.to_tensor([[2.0] * width, [6.0] * width], dtype="float64") + prediction.stop_gradient = False + label = {key: paddle.zeros_like(prediction), f"find_{key}": 1.0} + _, loss, metrics = loss_fn({}, lambda: {key: prediction}, label, 2, 0.25) + weight = limit + (start - limit) * 0.25 + np.testing.assert_allclose(loss.numpy(), 2.0 * weight) + np.testing.assert_allclose(metrics[f"mae_{term}"].numpy(), 2.0) + (gradient,) = paddle.grad(loss, prediction) + np.testing.assert_allclose(gradient.numpy(), weight / (4 * width)) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +def test_mae_force_uses_relative_error(start, limit): + """Apply relative normalization before computing force MAE.""" + loss_fn = EnergyStdLoss( + start_pref_f=start, limit_pref_f=limit, relative_f=1.0, loss_func="mae" + ) + prediction = paddle.to_tensor([[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]], dtype="float64") + prediction.stop_gradient = False + label = { + "force": paddle.to_tensor( + [[[3.0, 4.0, 0.0], [0.0, 0.0, 0.0]]], dtype="float64" + ), + "find_force": 1.0, + } + _, loss, metrics = loss_fn({}, lambda: {"force": prediction}, label, 2, 0.25) + weight = limit + (start - limit) * 0.25 + np.testing.assert_allclose(loss.numpy(), 25 / 36 * weight) + np.testing.assert_allclose(metrics["mae_f"].numpy(), 25 / 36) + (gradient,) = paddle.grad(loss, prediction) + expected = [[[-1 / 36, -1 / 36, 0.0], [1 / 6, 1 / 6, 1 / 6]]] + np.testing.assert_allclose(gradient.numpy(), np.asarray(expected) * weight) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0), (1.0, 1.0)]) +def test_virial_huber_takes_precedence_over_mae(start, limit): + """Keep the configured Huber virial loss when MAE is selected.""" + loss_fn = EnergyStdLoss( + start_pref_v=start, + limit_pref_v=limit, + loss_func="mae", + use_huber=True, + huber_delta=1.0, + ) + prediction = paddle.full((1, 9), 4.0, dtype="float64") + prediction.stop_gradient = False + label = {"virial": paddle.zeros_like(prediction), "find_virial": 1.0} + _, loss, _ = loss_fn({}, lambda: {"virial": prediction}, label, 2, 0.25) + weight = limit + (start - limit) * 0.25 + np.testing.assert_allclose(loss.numpy(), 1.5 * weight) + (gradient,) = paddle.grad(loss, prediction) + np.testing.assert_allclose(gradient.numpy(), weight / 18) diff --git a/source/tests/pt/test_loss_prefactor_schedule.py b/source/tests/pt/test_loss_prefactor_schedule.py new file mode 100644 index 0000000000..de02655e27 --- /dev/null +++ b/source/tests/pt/test_loss_prefactor_schedule.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import pytest +import torch + +from deepmd.pt.loss.dos import ( + DOSLoss, +) +from deepmd.pt.loss.ener import ( + EnergyStdLoss, +) +from deepmd.pt.loss.ener_spin import ( + EnergySpinLoss, +) +from deepmd.pt.utils.env import ( + DEVICE, +) + + +@pytest.mark.parametrize( + "start,limit", [(0.0, 1.0), (1.0, 0.0), (0.0, 0.0), (1.0, 1.0)] +) +@pytest.mark.parametrize( + "term,labels", + [ + ("e", {"energy"}), + ("f", {"force"}), + ("v", {"virial"}), + ("ae", {"atom_ener"}), + ("pf", {"force", "atom_pref"}), + ("gf", {"force", "drdq"}), + ], +) +def test_energy_schedule_label_requirements(term, labels, start, limit): + """Request the labels needed whenever a scheduled term is active.""" + loss = EnergyStdLoss( + **{f"start_pref_{term}": start, f"limit_pref_{term}": limit}, + numb_generalized_coord=1, + ) + required = {item.key for item in loss.label_requirement} + assert required == (labels if start or limit else set()) + + +@pytest.mark.parametrize( + "start,limit", [(0.0, 1.0), (1.0, 0.0), (0.0, 0.0), (1.0, 1.0)] +) +@pytest.mark.parametrize( + "term,label", + [ + ("e", "energy"), + ("fr", "force"), + ("fm", "force_mag"), + ("v", "virial"), + ("ae", "atom_ener"), + ], +) +def test_spin_schedule_label_requirements(term, label, start, limit): + """Request spin labels for either nonzero endpoint.""" + loss = EnergySpinLoss(**{f"start_pref_{term}": start, f"limit_pref_{term}": limit}) + required = {item.key for item in loss.label_requirement} + assert required == ({label} if start or limit else set()) + + +@pytest.mark.parametrize( + "start,limit", [(0.0, 1.0), (1.0, 0.0), (0.0, 0.0), (1.0, 1.0)] +) +@pytest.mark.parametrize("term", ["dos", "cdf", "ados", "acdf"]) +def test_dos_schedule_label_requirements(term, start, limit): + """Request DOS labels for either nonzero endpoint.""" + prefs = { + f"{bound}_pref_{name}": 0.0 + for name in ("dos", "cdf", "ados", "acdf") + for bound in ("start", "limit") + } + anchor = "ados" if term in ("dos", "cdf") else "dos" + prefs[f"start_pref_{anchor}"] = prefs[f"limit_pref_{anchor}"] = 1.0 + prefs[f"start_pref_{term}"] = start + prefs[f"limit_pref_{term}"] = limit + loss = DOSLoss(starter_learning_rate=1.0, numb_dos=2, **prefs) + required = {item.key for item in loss.label_requirement} + label = "dos" if term in ("dos", "cdf") else "atom_dos" + assert (label in required) == bool(start or limit) + + +@pytest.mark.parametrize("loss_type", [EnergyStdLoss, EnergySpinLoss]) +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("learning_rate", [1.0, 0.5, 0.0]) +def test_energy_schedule_contributes_to_loss_and_gradient( + loss_type, start, limit, learning_rate +): + """Apply the interpolated energy weight to the loss and gradient.""" + loss_fn = loss_type(start_pref_e=start, limit_pref_e=limit) + prediction = torch.tensor( + [[2.0]], dtype=torch.float64, device=DEVICE, requires_grad=True + ) + label = {"energy": torch.zeros_like(prediction), "find_energy": 1.0} + _, loss, _ = loss_fn({}, lambda: {"energy": prediction}, label, 2, learning_rate) + weight = limit + (start - limit) * learning_rate + torch.testing.assert_close( + loss, torch.tensor(2.0 * weight, dtype=loss.dtype, device=DEVICE) + ) + (gradient,) = torch.autograd.grad(loss, prediction) + torch.testing.assert_close(gradient, torch.full_like(prediction, 2.0 * weight)) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +def test_spin_energy_coefficients_are_requested_and_used(start, limit): + """Load and apply the coefficients used by spin energy loss.""" + loss_fn = EnergySpinLoss( + start_pref_e=start, limit_pref_e=limit, enable_atom_ener_coeff=True + ) + requirements = {item.key: item for item in loss_fn.label_requirement} + coeff = requirements["atom_ener_coeff"] + assert coeff.atomic and coeff.default == 1.0 and coeff.source_policy == "default" + prediction = torch.tensor( + [[[2.0], [3.0]]], dtype=torch.float64, device=DEVICE, requires_grad=True + ) + label = { + "energy": torch.zeros((1, 1), dtype=prediction.dtype, device=DEVICE), + "find_energy": 1.0, + "atom_ener_coeff": torch.tensor( + [[2.0, 0.0]], dtype=prediction.dtype, device=DEVICE + ), + } + model_pred = {"energy": prediction.sum(dim=1), "atom_energy": prediction} + _, loss, _ = loss_fn({}, lambda: model_pred, label, 2, 0.25) + weight = limit + (start - limit) * 0.25 + torch.testing.assert_close( + loss, torch.tensor(8.0 * weight, dtype=loss.dtype, device=DEVICE) + ) + (gradient,) = torch.autograd.grad(loss, prediction) + torch.testing.assert_close( + gradient, + torch.tensor([[[8.0], [0.0]]], dtype=prediction.dtype, device=DEVICE) * weight, + ) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("term", ["f", "pf", "gf"]) +def test_force_schedule_contributes_to_loss_and_gradient(term, start, limit): + """Apply force, prefactor-force and generalized-force schedules.""" + loss_fn = EnergyStdLoss( + **{f"start_pref_{term}": start, f"limit_pref_{term}": limit}, + numb_generalized_coord=1, + ) + prediction = torch.ones( + (1, 1, 3), dtype=torch.float64, device=DEVICE, requires_grad=True + ) + label = { + "force": torch.zeros_like(prediction), + "find_force": 1.0, + "atom_pref": torch.ones_like(prediction), + "find_atom_pref": 1.0, + "drdq": torch.tensor( + [[[1.0], [0.0], [0.0]]], dtype=torch.float64, device=DEVICE + ), + "find_drdq": 1.0, + } + _, loss, _ = loss_fn({}, lambda: {"force": prediction}, label, 1, 0.25) + weight = limit + (start - limit) * 0.25 + torch.testing.assert_close( + loss, torch.tensor(weight, dtype=loss.dtype, device=DEVICE) + ) + expected = [2.0, 0.0, 0.0] if term == "gf" else [2.0 / 3] * 3 + expected_grad = ( + torch.tensor([[expected]], dtype=prediction.dtype, device=DEVICE) * weight + ) + (gradient,) = torch.autograd.grad(loss, prediction) + torch.testing.assert_close(gradient, expected_grad) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +def test_magnetic_force_schedule_contributes_to_loss_and_gradient(start, limit): + """Apply the scheduled magnetic-force weight to loss and gradient.""" + loss_fn = EnergySpinLoss(start_pref_fm=start, limit_pref_fm=limit) + prediction = torch.tensor( + [[[1.0, 2.0, 3.0]]], dtype=torch.float64, device=DEVICE, requires_grad=True + ) + mask = torch.ones((1, 1, 1), dtype=torch.bool, device=DEVICE) + label = {"force_mag": torch.zeros_like(prediction), "find_force_mag": 1.0} + _, loss, _ = loss_fn( + {}, lambda: {"force_mag": prediction, "mask_mag": mask}, label, 1, 0.25 + ) + weight = limit + (start - limit) * 0.25 + torch.testing.assert_close( + loss, torch.tensor((14.0 / 3) * weight, dtype=loss.dtype, device=DEVICE) + ) + (gradient,) = torch.autograd.grad(loss, prediction) + torch.testing.assert_close(gradient, prediction.detach() * (2.0 / 3) * weight) + + +@pytest.mark.parametrize("start,limit", [(0.0, 1.0), (1.0, 0.0)]) +@pytest.mark.parametrize("learning_rate", [1.0, 0.5, 0.0]) +@pytest.mark.parametrize("term", ["dos", "cdf", "ados", "acdf"]) +def test_dos_schedule_contributes_to_loss_and_gradient( + term, start, limit, learning_rate +): + """Apply the scheduled DOS or cumulative-DOS weight.""" + prefs = { + f"{bound}_pref_{name}": 0.0 + for name in ("dos", "cdf", "ados", "acdf") + for bound in ("start", "limit") + } + prefs[f"start_pref_{term}"] = start + prefs[f"limit_pref_{term}"] = limit + loss_fn = DOSLoss( + starter_learning_rate=1.0, + numb_dos=2, + **prefs, + ) + prediction = torch.tensor( + [[1.0, 2.0]], dtype=torch.float64, device=DEVICE, requires_grad=True + ) + key = "atom_dos" if term in ("ados", "acdf") else "dos" + label = {key: torch.zeros_like(prediction), f"find_{key}": 1.0} + _, loss, _ = loss_fn({}, lambda: {key: prediction}, label, 1, learning_rate) + weight = limit + (start - limit) * learning_rate + cumulative = term in ("cdf", "acdf") + torch.testing.assert_close( + loss, + torch.tensor( + (5.0 if cumulative else 2.5) * weight, dtype=loss.dtype, device=DEVICE + ), + ) + (gradient,) = torch.autograd.grad(loss, prediction) + expected_grad = [[4.0, 3.0]] if cumulative else [[1.0, 2.0]] + torch.testing.assert_close( + gradient, + torch.tensor(expected_grad, dtype=prediction.dtype, device=DEVICE) * weight, + )