diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index 0036ae1620..953b413dd1 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -130,15 +130,17 @@ def call( local_label = xp.reshape( label_dict["atom_dos"], (-1, natoms, self.numb_dos) ) - diff = xp.reshape(local_pred - local_label, (-1, self.numb_dos)) + diff3d = local_pred - local_label # [nf, natoms, numb_dos] if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss_dos = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] + nf = diff3d.shape[0] + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf] + l2_local_loss_dos = xp.mean(per_frame_sum / per_frame_dof) else: - l2_local_loss_dos = xp.mean(xp.square(diff)) + l2_local_loss_dos = xp.mean(xp.square(diff3d)) loss += pref_ados * l2_local_loss_dos more_loss["rmse_local_dos"] = self.display_if_exist( xp.sqrt(l2_local_loss_dos), find_local @@ -155,15 +157,17 @@ def call( xp.reshape(label_dict["atom_dos"], (-1, natoms, self.numb_dos)), axis=-1, ) - diff = xp.reshape(local_pred_cdf - local_label_cdf, (-1, self.numb_dos)) + diff3d = local_pred_cdf - local_label_cdf # [nf, natoms, numb_dos] if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss_cdf = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] + nf = diff3d.shape[0] + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf] + l2_local_loss_cdf = xp.mean(per_frame_sum / per_frame_dof) else: - l2_local_loss_cdf = xp.mean(xp.square(diff)) + l2_local_loss_cdf = xp.mean(xp.square(diff3d)) loss += pref_acdf * l2_local_loss_cdf more_loss["rmse_local_cdf"] = self.display_if_exist( xp.sqrt(l2_local_loss_cdf), find_local @@ -175,15 +179,14 @@ def call( global_pred = xp.reshape(model_dict["dos"], (-1, self.numb_dos)) global_label = xp.reshape(label_dict["dos"], (-1, self.numb_dos)) diff = global_pred - global_label + # idiom 3: global dos is already padding-invariant; plain mean suffices + l2_global_loss_dos = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss_dos = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss_dos = xp.mean(xp.square(diff)) loss += pref_dos * l2_global_loss_dos more_loss["rmse_global_dos"] = self.display_if_exist( xp.sqrt(l2_global_loss_dos) / atom_num, find_global @@ -199,15 +202,14 @@ def call( xp.reshape(label_dict["dos"], (-1, self.numb_dos)), axis=-1 ) diff = global_pred_cdf - global_label_cdf + # idiom 3: global cdf is already padding-invariant; plain mean suffices + l2_global_loss_cdf = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss_cdf = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss_cdf = xp.mean(xp.square(diff)) loss += pref_cdf * l2_global_loss_cdf more_loss["rmse_global_cdf"] = self.display_if_exist( xp.sqrt(l2_global_loss_cdf) / atom_num, find_global diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7515f19b9a..5ee4916762 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -223,6 +223,19 @@ def call( atom_pref, ) + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghosts. + if "mask" in model_dict: + maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] + real_natoms = xp.sum(maskf, axis=-1) # [nf] + inv = xp.reshape(1.0 / real_natoms, (-1,)) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). + maskf = None + if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as # E = \sum_i \nu_i E_i @@ -280,73 +293,194 @@ def call( if self.has_e: if self.loss_func == "mse": l2_ener_loss = xp.mean(xp.square(energy - energy_hat)) - if not self.use_huber: - loss += atom_norm_ener**norm_exp * (pref_e * l2_ener_loss) + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + se = xp.square(energy - energy_hat) # [nf, k] + per_frame = xp.mean(xp.reshape(se, (_nf, -1)), axis=-1) # [nf] + if not self.use_huber: + loss += pref_e * xp.mean(per_frame * inv**norm_exp) + else: + inv_col = xp.reshape(inv, (_nf, 1)) # [nf, 1] + l_huber_loss = custom_huber_loss( + inv_col * energy, + inv_col * energy_hat, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame * inv**2)), find_energy + ) else: - l_huber_loss = custom_huber_loss( - atom_norm_ener * energy, - atom_norm_ener * energy_hat, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += atom_norm_ener**norm_exp * (pref_e * l2_ener_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm_ener * energy, + atom_norm_ener * energy_hat, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(l2_ener_loss) * atom_norm_ener, find_energy ) - loss += pref_e * l_huber_loss - more_loss["rmse_e"] = self.display_if_exist( - xp.sqrt(l2_ener_loss) * atom_norm_ener, find_energy - ) elif self.loss_func == "mae": l1_ener_loss = xp.mean(xp.abs(energy - energy_hat)) - loss += atom_norm_ener * (pref_e * l1_ener_loss) - more_loss["mae_e"] = self.display_if_exist( - l1_ener_loss * atom_norm_ener, find_energy - ) + if maskf is not None: + abs_e = xp.abs(energy - energy_hat) # [nf, k] + per_frame_ae = xp.mean( + xp.reshape(abs_e, (_nf, -1)), axis=-1 + ) # [nf] + l1_ener_masked = xp.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked, find_energy + ) + else: + loss += atom_norm_ener * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss * atom_norm_ener, find_energy + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener + if maskf is not None: + abs_e = xp.abs(energy - energy_hat) + per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1) + mae_e = xp.mean(per_frame_ae * inv) + else: + mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener more_loss["mae_e"] = self.display_if_exist(mae_e, find_energy) mae_e_all = xp.mean(xp.abs(energy - energy_hat)) more_loss["mae_e_all"] = self.display_if_exist(mae_e_all, find_energy) if self.has_f: if self.loss_func == "mse": l2_force_loss = xp.mean(xp.square(diff_f)) - if not self.use_huber: - loss += pref_f * l2_force_loss + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3] + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + # Masked MSE computed for rmse_f display regardless of use_huber. + sq_f = xp.square(diff_f_3d) * maskf_col # [nf, nloc, 3] + _pfs = xp.sum(xp.reshape(sq_f, (_nf, -1)), axis=-1) # [nf] + _pfd = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_force_masked = xp.mean(_pfs / _pfd) + if not self.use_huber: + loss += pref_f * l2_force_masked + else: + if not self.f_use_norm: + abs_e = xp.abs(diff_f_3d) + quad = 0.5 * xp.square(diff_f_3d) + lin = self._huber_delta_force * ( + abs_e - 0.5 * self._huber_delta_force + ) + huber_elem = xp.where( + abs_e <= self._huber_delta_force, quad, lin + ) + huber_masked = huber_elem * maskf_col + else: + diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + norm_2d = xp.reshape( + xp.linalg.vector_norm( + xp.reshape(diff_3, (-1, 3)), axis=1 + ), + (_nf, _nloc), + ) + abs_n = norm_2d + quad_n = 0.5 * xp.square(norm_2d) + lin_n = self._huber_delta_force * ( + abs_n - 0.5 * self._huber_delta_force + ) + huber_n = xp.where( + abs_n <= self._huber_delta_force, quad_n, lin_n + ) + huber_masked = xp.reshape(huber_n * maskf, (_nf, _nloc, 1)) + per_frame_sum = xp.sum( + xp.reshape(huber_masked, (_nf, -1)), axis=-1 + ) + if not self.f_use_norm: + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + else: + per_frame_dof = xp.sum(maskf, axis=-1) + l_huber_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l_huber_masked + more_loss["rmse_f"] = self.display_if_exist( + xp.sqrt(l2_force_masked), find_force + ) else: + if not self.use_huber: + loss += pref_f * l2_force_loss + else: + if not self.f_use_norm: + l_huber_loss = custom_huber_loss( + xp.reshape(force, (-1,)), + xp.reshape(force_hat, (-1,)), + delta=self._huber_delta_force, + ) + else: + force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) + force_diff_norm = xp.reshape( + xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) + ) + l_huber_loss = custom_huber_loss( + force_diff_norm, + xp.zeros_like(force_diff_norm), + delta=self._huber_delta_force, + ) + loss += pref_f * l_huber_loss + more_loss["rmse_f"] = self.display_if_exist( + xp.sqrt(l2_force_loss), find_force + ) + elif self.loss_func == "mae": + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - xp.reshape(force, (-1,)), - xp.reshape(force_hat, (-1,)), - delta=self._huber_delta_force, + abs_f = xp.abs(diff_f_3d) * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + l1_force_masked = xp.mean(per_frame_sum / per_frame_dof) + else: + diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + norm_2d = xp.reshape( + xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), + (_nf, _nloc), ) + masked_norm = norm_2d * maskf + per_frame_sum = xp.sum(masked_norm, axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) + l1_force_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l1_force_masked + more_loss["mae_f"] = self.display_if_exist( + l1_force_masked, find_force + ) + else: + if not self.f_use_norm: + l1_force_loss = xp.mean(xp.abs(diff_f)) else: force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) - force_diff_norm = xp.reshape( - xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) - ) - l_huber_loss = custom_huber_loss( - force_diff_norm, - xp.zeros_like(force_diff_norm), - delta=self._huber_delta_force, + l1_force_loss = xp.mean( + xp.linalg.vector_norm(force_diff_3, axis=1) ) - loss += pref_f * l_huber_loss - more_loss["rmse_f"] = self.display_if_exist( - xp.sqrt(l2_force_loss), find_force - ) - elif self.loss_func == "mae": - if not self.f_use_norm: - l1_force_loss = xp.mean(xp.abs(diff_f)) - else: - force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) - l1_force_loss = xp.mean(xp.linalg.vector_norm(force_diff_3, axis=1)) - loss += pref_f * l1_force_loss - more_loss["mae_f"] = self.display_if_exist(l1_force_loss, find_force) + loss += pref_f * l1_force_loss + more_loss["mae_f"] = self.display_if_exist( + l1_force_loss, find_force + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for force loss." ) if mae: - mae_f = xp.mean(xp.abs(diff_f)) + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_f = xp.abs(diff_f_3d) * maskf_col + per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + mae_f = xp.mean(per_frame_sum / per_frame_dof) + else: + mae_f = xp.mean(xp.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist(mae_f, find_force) if self.has_v: virial_reshape = xp.reshape(virial, (-1,)) @@ -355,30 +489,70 @@ def call( l2_virial_loss = xp.mean( xp.square(virial_hat_reshape - virial_reshape), ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization. + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + se_v = xp.square(v_hat_2d - v2d) # [nf, 9] + per_frame_v = xp.mean(se_v, axis=-1) # [nf] + if not self.use_huber: + loss += pref_v * xp.mean(per_frame_v * inv**norm_exp) + else: + inv_col = xp.reshape(inv, (_nf, 1)) # [nf, 1] + l_huber_v = custom_huber_loss( + inv_col * v2d, + inv_col * v_hat_2d, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_v + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_v * inv**2)), find_virial + ) else: - l_huber_loss = custom_huber_loss( - atom_norm * virial_reshape, - atom_norm * virial_hat_reshape, - delta=self._huber_delta_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 * virial_reshape, + atom_norm * virial_hat_reshape, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_loss + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(l2_virial_loss) * atom_norm, find_virial ) - loss += pref_v * l_huber_loss - more_loss["rmse_v"] = self.display_if_exist( - xp.sqrt(l2_virial_loss) * atom_norm, find_virial - ) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss * atom_norm, find_virial - ) + if maskf is not None: + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + abs_v = xp.abs(v_hat_2d - v2d) # [nf, 9] + per_frame_v = xp.mean(abs_v, axis=-1) # [nf] + l1_virial_masked = xp.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked, find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss * atom_norm, find_virial + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." ) if mae: - mae_v = xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) * atom_norm + if maskf is not None: + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + abs_v = xp.abs(v_hat_2d - v2d) + per_frame_v = xp.mean(abs_v, axis=-1) + mae_v = xp.mean(per_frame_v * inv) + else: + mae_v = ( + xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) * atom_norm + ) more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) if self.has_ae: atom_ener_reshape = xp.reshape(atom_ener, (-1,)) @@ -387,26 +561,67 @@ def call( l2_atom_ener_loss = xp.mean( xp.square(atom_ener_hat_reshape - atom_ener_reshape), ) - if not self.use_huber: - loss += pref_ae * l2_atom_ener_loss + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) + ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + sq_ae = xp.square(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = xp.sum(sq_ae, axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] + l2_ae_masked = xp.mean(per_frame_sum / per_frame_dof) + if not self.use_huber: + loss += pref_ae * l2_ae_masked + else: + # Huber applied element-wise then masked-mean. + diff_ae = ae_hat_2d - ae_2d + abs_ae = xp.abs(diff_ae) + quad_ae = 0.5 * xp.square(diff_ae) + lin_ae = self._huber_delta_energy * ( + abs_ae - 0.5 * self._huber_delta_energy + ) + huber_ae = xp.where( + abs_ae <= self._huber_delta_energy, quad_ae, lin_ae + ) + huber_ae_masked = huber_ae * maskf + per_frame_sum_h = xp.sum(huber_ae_masked, axis=-1) + l_huber_ae_masked = xp.mean(per_frame_sum_h / per_frame_dof) + loss += pref_ae * l_huber_ae_masked + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_ae_masked), find_atom_ener + ) else: - l_huber_loss = custom_huber_loss( - atom_ener_reshape, - atom_ener_hat_reshape, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += pref_ae * l2_atom_ener_loss + else: + l_huber_loss = custom_huber_loss( + atom_ener_reshape, + atom_ener_hat_reshape, + delta=self._huber_delta_energy, + ) + loss += pref_ae * l_huber_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener ) - loss += pref_ae * l_huber_loss - more_loss["rmse_ae"] = self.display_if_exist( - xp.sqrt(l2_atom_ener_loss), find_atom_ener - ) elif self.loss_func == "mae": l1_atom_ener_loss = xp.mean( xp.abs(atom_ener_hat_reshape - atom_ener_reshape) ) - loss += pref_ae * l1_atom_ener_loss - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss, find_atom_ener - ) + if maskf is not None: + ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) + ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + abs_ae = xp.abs(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = xp.sum(abs_ae, axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] + l1_ae_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l1_ae_masked + more_loss["mae_ae"] = self.display_if_exist( + l1_ae_masked, find_atom_ener + ) + else: + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atomic energy loss." @@ -418,18 +633,47 @@ def call( l2_pref_force_loss = xp.mean( xp.multiply(xp.square(diff_f), atom_pref_reshape), ) - loss += pref_pf * l2_pref_force_loss - more_loss["rmse_pf"] = self.display_if_exist( - xp.sqrt(l2_pref_force_loss), find_atom_pref - ) + if maskf is not None: + # Idiom 1 with pref weight (ncomp=3). + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + sq_pf = xp.square(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum( + xp.reshape(sq_pf, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_pf_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_pf * l2_pf_masked + more_loss["rmse_pf"] = self.display_if_exist( + xp.sqrt(l2_pf_masked), find_atom_pref + ) + else: + loss += pref_pf * l2_pref_force_loss + more_loss["rmse_pf"] = self.display_if_exist( + xp.sqrt(l2_pref_force_loss), find_atom_pref + ) elif self.loss_func == "mae": l1_pref_force_loss = xp.mean( xp.multiply(xp.abs(diff_f), atom_pref_reshape) ) - loss += pref_pf * l1_pref_force_loss - more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss, find_atom_pref - ) + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_pf = xp.abs(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum(xp.reshape(abs_pf, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + l1_pf_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_pf * l1_pf_masked + more_loss["mae_pf"] = self.display_if_exist( + l1_pf_masked, find_atom_pref + ) + else: + loss += pref_pf * l1_pref_force_loss + more_loss["mae_pf"] = self.display_if_exist( + l1_pref_force_loss, find_atom_pref + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." @@ -437,22 +681,40 @@ def call( if self.has_gf: find_drdq = label_dict["find_drdq"] drdq = label_dict["drdq"] - force_reshape_nframes = xp.reshape(force, (-1, natoms * 3)) - force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms * 3)) - drdq_reshape = xp.reshape( - drdq, (-1, natoms * 3, self.numb_generalized_coord) - ) - # "bij,bi->bj" einsum replaced with array-API-compatible ops - gen_force_hat = xp.sum( - drdq_reshape * force_hat_reshape_nframes[:, :, None], axis=1 - ) - gen_force = xp.sum(drdq_reshape * force_reshape_nframes[:, :, None], axis=1) - diff_gen_force = gen_force_hat - gen_force - l2_gen_force_loss = xp.mean(xp.square(diff_gen_force)) pref_gf = find_drdq * ( self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * lr_ratio ) + if maskf is not None: + # Mask per-atom forces before projecting onto generalized coords + # so ghost atoms don't contribute to the generalized force. + force_3d = xp.reshape(force, (_nf, _nloc, 3)) + force_hat_3d = xp.reshape(force_hat, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + masked_f = force_3d * maskf_col # [nf, nloc, 3] + masked_f_hat = force_hat_3d * maskf_col # [nf, nloc, 3] + f_flat = xp.reshape(masked_f, (_nf, _nloc * 3)) + f_hat_flat = xp.reshape(masked_f_hat, (_nf, _nloc * 3)) + drdq_reshape = xp.reshape( + drdq, (_nf, _nloc * 3, self.numb_generalized_coord) + ) + gen_force = xp.sum(drdq_reshape * f_flat[:, :, None], axis=1) + gen_force_hat = xp.sum(drdq_reshape * f_hat_flat[:, :, None], axis=1) + else: + force_reshape_nframes = xp.reshape(force, (-1, natoms * 3)) + force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms * 3)) + drdq_reshape = xp.reshape( + drdq, (-1, natoms * 3, self.numb_generalized_coord) + ) + gen_force_hat = xp.sum( + drdq_reshape * force_hat_reshape_nframes[:, :, None], axis=1 + ) + gen_force = xp.sum( + drdq_reshape * force_reshape_nframes[:, :, None], axis=1 + ) + # "bij,bi->bj" einsum replaced with array-API-compatible ops + diff_gen_force = gen_force_hat - gen_force + l2_gen_force_loss = xp.mean(xp.square(diff_gen_force)) loss += pref_gf * l2_gen_force_loss more_loss["rmse_gf"] = self.display_if_exist( xp.sqrt(l2_gen_force_loss), find_drdq diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 6262015c84..e39e0a6a9e 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -132,6 +132,19 @@ def call( # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size norm_exp = 2 if self.intensive_ener_virial else 1 + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghost padding atoms. + if "mask" in model_dict: + maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] + real_natoms = xp.sum(maskf, axis=-1) # [nf] + inv = xp.reshape(1.0 / real_natoms, (-1,)) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). + maskf = None + if self.has_e: energy_pred = model_dict["energy"] energy_label = label_dict["energy"] @@ -143,20 +156,45 @@ def call( atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener_pred.shape) energy_pred = xp.sum(atom_ener_coeff * atom_ener_pred, axis=1) if self.loss_func == "mse": - l2_ener_loss = xp.mean(xp.square(energy_pred - energy_label)) - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) - more_loss["rmse_e"] = self.display_if_exist( - xp.sqrt(l2_ener_loss) * atom_norm, find_energy - ) + se_e = xp.square(energy_pred - energy_label) # [nf, k] + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + per_frame_e = xp.mean(xp.reshape(se_e, (_nf, -1)), axis=-1) # [nf] + loss += pref_e * xp.mean(per_frame_e * inv**norm_exp) + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_e * inv**2)), find_energy + ) + else: + l2_ener_loss = xp.mean(se_e) + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(l2_ener_loss) * atom_norm, find_energy + ) elif self.loss_func == "mae": - abs_diff_e = xp.abs(energy_pred - energy_label) - l1_ener_loss = xp.sum(abs_diff_e) - loss += pref_e * l1_ener_loss - more_loss["mae_e"] = self.display_if_exist( - xp.mean(abs_diff_e), find_energy - ) + l1_ener_loss = xp.mean(xp.abs(energy_pred - energy_label)) + if maskf is not None: + # Idiom 2 (extensive) with abs: per-frame normalization by real-atom count. + abs_e = xp.abs(energy_pred - energy_label) # [nf, k] + per_frame_ae = xp.mean( + xp.reshape(abs_e, (_nf, -1)), axis=-1 + ) # [nf] + l1_ener_masked = xp.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked, find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss * atom_norm, find_energy + ) if mae: - mae_e = xp.mean(xp.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = xp.abs(energy_pred - energy_label) + per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1) + mae_e = xp.mean(per_frame_ae * inv) + else: + mae_e = xp.mean(xp.abs(energy_pred - energy_label)) * atom_norm more_loss["mae_e"] = self.display_if_exist(mae_e, find_energy) mae_e_all = xp.mean(xp.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist(mae_e_all, find_energy) @@ -167,24 +205,57 @@ def call( force_pred = model_dict["force"] force_label = label_dict["force"] if self.loss_func == "mse": - diff_fr = force_label - force_pred - l2_force_real_loss = xp.mean(xp.square(diff_fr)) - loss += pref_fr * l2_force_real_loss - more_loss["rmse_fr"] = self.display_if_exist( - xp.sqrt(l2_force_real_loss), find_force - ) - if mae: - mae_fr = xp.mean(xp.abs(force_label - force_pred)) - more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) + diff_fr = force_label - force_pred # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + sq_fr = xp.square(diff_fr) * maskf_col # [nf, nloc, 3] + per_frame_fr_sum = xp.sum( + xp.reshape(sq_fr, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_fr_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_force_real_loss = xp.mean(per_frame_fr_sum / per_frame_fr_dof) + loss += pref_fr * l2_force_real_loss + more_loss["rmse_fr"] = self.display_if_exist( + xp.sqrt(l2_force_real_loss), find_force + ) + if mae: + abs_fr = xp.abs(force_label - force_pred) * maskf_col + per_frame_mae_sum = xp.sum( + xp.reshape(abs_fr, (_nf, -1)), axis=-1 + ) + mae_fr = xp.mean(per_frame_mae_sum / per_frame_fr_dof) + more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) + else: + l2_force_real_loss = xp.mean(xp.square(diff_fr)) + loss += pref_fr * l2_force_real_loss + more_loss["rmse_fr"] = self.display_if_exist( + xp.sqrt(l2_force_real_loss), find_force + ) + if mae: + mae_fr = xp.mean(xp.abs(force_label - force_pred)) + more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) elif self.loss_func == "mae": - abs_diff_fr = xp.abs(force_label - force_pred) - per_atom_fr = xp.sum(abs_diff_fr, axis=-1) # [nf, na] - per_frame_fr = xp.mean(per_atom_fr, axis=-1) # [nf] - l1_force_real_loss = xp.sum(per_frame_fr) # scalar - loss += pref_fr * l1_force_real_loss - more_loss["mae_fr"] = self.display_if_exist( - xp.mean(abs_diff_fr), find_force - ) + abs_diff_fr = xp.abs(force_label - force_pred) # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3) with abs. + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_fr = abs_diff_fr * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum( + xp.reshape(abs_fr, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l1_force_real_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_fr * l1_force_real_masked + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_masked, find_force + ) + else: + l1_force_real_loss = xp.mean(abs_diff_fr) + loss += pref_fr * l1_force_real_loss + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_loss, find_force + ) if self.has_fm: find_force_mag = label_dict.get("find_force_mag", 0.0) @@ -208,16 +279,14 @@ def call( more_loss["mae_fm"] = self.display_if_exist(mae_fm, find_force_mag) elif self.loss_func == "mae": abs_diff_fm = xp.abs(diff_fm) # [nf, na, 3], zeros for non-magnetic - per_atom_fm = xp.sum(abs_diff_fm, axis=-1) # [nf, na] - mask_2d = mask_float[:, :, 0] # [nf, na] - per_frame_sum_fm = xp.sum(per_atom_fm, axis=-1) # [nf] - per_frame_count_fm = xp.sum(mask_2d, axis=-1) # [nf] - l1_force_mag_loss = xp.sum( - per_frame_sum_fm / per_frame_count_fm - ) # scalar + # Mean over frames, magnetic atoms and xyz (same reduction as + # force_mag MSE, force_real MAE and the displayed mae_fm) so the + # loss is batch-size independent: a 2-frame batch equals the mean + # of the two single-frame losses. + l1_force_mag_loss = xp.sum(abs_diff_fm) / (n_valid * 3) loss += pref_fm * l1_force_mag_loss more_loss["mae_fm"] = self.display_if_exist( - xp.sum(abs_diff_fm) / (n_valid * 3), find_force_mag + l1_force_mag_loss, find_force_mag ) if self.has_ae: @@ -225,46 +294,96 @@ def call( pref_ae = pref_ae * find_atom_ener atom_ener = model_dict["atom_energy"] atom_ener_label = label_dict["atom_ener"] - atom_ener_reshape = xp.reshape(atom_ener, (-1,)) - atom_ener_label_reshape = xp.reshape(atom_ener_label, (-1,)) - if self.loss_func == "mse": - l2_atom_ener_loss = xp.mean( - xp.square(atom_ener_label_reshape - atom_ener_reshape) - ) - loss += pref_ae * l2_atom_ener_loss - more_loss["rmse_ae"] = self.display_if_exist( - xp.sqrt(l2_atom_ener_loss), find_atom_ener - ) - elif self.loss_func == "mae": - l1_atom_ener_loss = xp.mean( - xp.abs(atom_ener_label_reshape - atom_ener_reshape) - ) - loss += pref_ae * l1_atom_ener_loss - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss, find_atom_ener - ) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae = xp.reshape(atom_ener, (_nf, _nloc, 1)) + ae_label = xp.reshape(atom_ener_label, (_nf, _nloc, 1)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + if self.loss_func == "mse": + sq = xp.square(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = xp.sum(xp.reshape(sq, (_nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] (ncomp=1) + l2_atom_ener_loss = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l2_atom_ener_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener + ) + elif self.loss_func == "mae": + abs_diff = xp.abs(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = xp.sum( + xp.reshape(abs_diff, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] (ncomp=1) + l1_atom_ener_loss = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) + else: + atom_ener_reshape = xp.reshape(atom_ener, (-1,)) + atom_ener_label_reshape = xp.reshape(atom_ener_label, (-1,)) + if self.loss_func == "mse": + l2_atom_ener_loss = xp.mean( + xp.square(atom_ener_label_reshape - atom_ener_reshape) + ) + loss += pref_ae * l2_atom_ener_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener + ) + elif self.loss_func == "mae": + l1_atom_ener_loss = xp.mean( + xp.abs(atom_ener_label_reshape - atom_ener_reshape) + ) + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) if self.has_v: find_virial = label_dict.get("find_virial", 0.0) pref_v = pref_v * find_virial virial_pred = xp.reshape(model_dict["virial"], (-1, 9)) virial_label = label_dict["virial"] - diff_v = virial_label - virial_pred + diff_v = virial_label - virial_pred # [nf, 9] if self.loss_func == "mse": - l2_virial_loss = xp.mean(xp.square(diff_v)) - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) - more_loss["rmse_v"] = self.display_if_exist( - xp.sqrt(l2_virial_loss) * atom_norm, find_virial - ) - if mae: - mae_v = xp.mean(xp.abs(diff_v)) * atom_norm - more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization by real-atom count. + se_v = xp.square(diff_v) # [nf, 9] + per_frame_v = xp.mean(se_v, axis=-1) # [nf] + loss += pref_v * xp.mean(per_frame_v * inv**norm_exp) + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_v * inv**2)), find_virial + ) + if mae: + abs_v = xp.abs(diff_v) # [nf, 9] + per_frame_mae_v = xp.mean(abs_v, axis=-1) # [nf] + mae_v = xp.mean(per_frame_mae_v * inv) + more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) + else: + l2_virial_loss = xp.mean(xp.square(diff_v)) + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(l2_virial_loss) * atom_norm, find_virial + ) + if mae: + mae_v = xp.mean(xp.abs(diff_v)) * atom_norm + more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(diff_v)) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss * atom_norm, find_virial - ) + if maskf is not None: + # Idiom 2 (extensive, k=9) with abs: per-frame normalization by real-atom count. + abs_v = xp.abs(diff_v) # [nf, 9] + per_frame_v = xp.mean(abs_v, axis=-1) # [nf] + l1_virial_masked = xp.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked, find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss * atom_norm, find_virial + ) more_loss["rmse"] = xp.sqrt(loss) return loss, more_loss diff --git a/deepmd/dpmodel/loss/property.py b/deepmd/dpmodel/loss/property.py index 7d658ff925..ef649c9b2f 100644 --- a/deepmd/dpmodel/loss/property.py +++ b/deepmd/dpmodel/loss/property.py @@ -85,8 +85,17 @@ def call( # Normalize by natoms for extensive properties (without mutating input) if not self.intensive: - pred = pred / natoms - label = label / natoms + if "mask" in model_dict: + # Per-frame real atom count: shape [nf] → broadcast over [nf, task_dim]. + real_natoms = xp.reshape( + xp.astype(xp.sum(model_dict["mask"], axis=-1), pred.dtype), + (-1,) + (1,) * (pred.ndim - 1), + ) + pred = pred / real_natoms + label = label / real_natoms + else: + pred = pred / natoms + label = label / natoms # Get out_std and out_bias if self.out_std is not None: diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index e367457451..4b21175185 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -103,11 +103,14 @@ def call( diff = xp.reshape(local_pred - local_label, (-1, self.tensor_size)) diff = diff * atomic_weight if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss = xp.sum(xp.square(diff)) / (n_valid * self.tensor_size) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff.dtype) # [nf, natoms] + nf = local_pred.shape[0] + diff3d = xp.reshape(diff, (nf, natoms, self.tensor_size)) + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.tensor_size # [nf] + l2_local_loss = xp.mean(per_frame_sum / per_frame_dof) else: l2_local_loss = xp.mean(xp.square(diff)) loss += local_weight * l2_local_loss @@ -129,15 +132,14 @@ def call( label_dict[self.label_name], (-1, self.tensor_size) ) diff = global_pred - global_label + # idiom 3: global tensor is already padding-invariant; plain mean suffices + l2_global_loss = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss = xp.mean(xp.square(diff)) loss += global_weight * l2_global_loss more_loss[f"rmse_global_{self.tensor_name}"] = self.display_if_exist( xp.sqrt(l2_global_loss) / atom_num, find_global diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 64acae0c05..bf2d2b9710 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -122,7 +122,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_dos = ( @@ -149,12 +149,19 @@ def forward( local_tensor_label_dos = label["atom_dos"].reshape( [-1, natoms, self.numb_dos] ) - diff = (local_tensor_pred_dos - local_tensor_label_dos).reshape( - [-1, self.numb_dos] - ) + diff = ( + local_tensor_pred_dos - local_tensor_label_dos + ) # [nf, natoms, numb_dos] if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss_dos = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = diff.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + sq = torch.square(diff) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.numb_dos # [nf] + l2_local_loss_dos = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss_dos = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_local_dos_loss"] = self.display_if_exist( l2_local_loss_dos.detach(), find_local @@ -173,12 +180,19 @@ def forward( local_tensor_label_cdf = torch.cumsum( label["atom_dos"].reshape([-1, natoms, self.numb_dos]), dim=-1 ) - diff = (local_tensor_pred_cdf - local_tensor_label_cdf).reshape( - [-1, self.numb_dos] - ) + diff = ( + local_tensor_pred_cdf - local_tensor_label_cdf + ) # [nf, natoms, numb_dos] if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss_cdf = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = diff.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + sq = torch.square(diff) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.numb_dos # [nf] + l2_local_loss_cdf = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss_cdf = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_local_cdf_loss"] = self.display_if_exist( l2_local_loss_cdf.detach(), find_local @@ -194,15 +208,12 @@ def forward( global_tensor_pred_dos = model_pred["dos"].reshape([-1, self.numb_dos]) global_tensor_label_dos = label["dos"].reshape([-1, self.numb_dos]) diff = global_tensor_pred_dos - global_tensor_label_dos + # idiom 3: global dos is already padding-invariant; plain mean suffices + l2_global_loss_dos = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss_dos = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss_dos = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_global_dos_loss"] = self.display_if_exist( l2_global_loss_dos.detach(), find_global @@ -222,15 +233,12 @@ def forward( label["dos"].reshape([-1, self.numb_dos]), dim=-1 ) diff = global_tensor_pred_cdf - global_tensor_label_cdf + # idiom 3: global cdf is already padding-invariant; plain mean suffices + l2_global_loss_cdf = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss_cdf = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss_cdf = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_global_cdf_loss"] = self.display_if_exist( l2_global_loss_cdf.detach(), find_global diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 50d83a4ac9..93f57618d6 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -225,7 +225,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_f = self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * coef @@ -239,6 +239,18 @@ def forward( # more_loss['log_keys'] = [] # showed when validation on the fly # more_loss['test_keys'] = [] # showed when doing dp test atom_norm = 1.0 / natoms + + # Per-frame mask: recover real-atom count per frame when mask is provided. + if "mask" in model_pred: + maskf = model_pred["mask"] # [nf, nloc], float + real_natoms_f = torch.sum(maskf, dim=-1) # [nf] + inv = (1.0 / real_natoms_f).reshape(-1) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). + maskf = None # Normalization exponent controls loss scaling with system size: # - norm_exp=2 (intensive_ener_virial=True): loss uses 1/N² scaling, making it independent of system size # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size @@ -267,19 +279,38 @@ def forward( more_loss["l2_ener_loss"] = self.display_if_exist( l2_ener_loss.detach(), find_energy ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization. + se = torch.square(energy_pred - energy_label) # [nf, k] + per_frame = torch.mean(se.reshape(_nf, -1), dim=-1) # [nf] + if not self.use_huber: + loss += pref_e * torch.mean(per_frame * inv**norm_exp) + else: + inv_col = inv.reshape(_nf, 1) + l_huber_loss = custom_huber_loss( + inv_col * energy_pred, + inv_col * energy_label, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + rmse_e = torch.sqrt(torch.mean(per_frame * inv**2)) + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy + ) else: - l_huber_loss = custom_huber_loss( - atom_norm * energy_pred, - atom_norm * energy_label, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * energy_pred, + atom_norm * energy_label, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + rmse_e = l2_ener_loss.sqrt() * atom_norm + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy ) - loss += pref_e * l_huber_loss - rmse_e = l2_ener_loss.sqrt() * atom_norm - more_loss["rmse_e"] = self.display_if_exist( - rmse_e.detach(), find_energy - ) # more_loss['log_keys'].append('rmse_e') elif self.loss_func == "mae": l1_ener_loss = F.l1_loss( @@ -287,18 +318,34 @@ def forward( energy_label.reshape(-1), reduction="mean", ) - loss += atom_norm * (pref_e * l1_ener_loss) - more_loss["mae_e"] = self.display_if_exist( - l1_ener_loss.detach() * atom_norm, - find_energy, - ) + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + l1_ener_masked = torch.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked.detach(), find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss.detach() * atom_norm, + find_energy, + ) # more_loss['log_keys'].append('rmse_e') else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + mae_e = torch.mean(per_frame_ae * inv) + else: + mae_e = ( + torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + ) more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy) mae_e_all = torch.mean(torch.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist( @@ -330,56 +377,136 @@ def forward( more_loss["l2_force_loss"] = self.display_if_exist( l2_force_loss.detach(), find_force ) - if not self.use_huber: - loss += (pref_f * l2_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + # Masked MSE computed for rmse_f display regardless of use_huber. + sq_f = torch.square(diff_f_3d) * maskf_col + _pfs = sq_f.reshape(_nf, -1).sum(dim=-1) + _pfd = maskf.sum(dim=-1) * 3 + l2_f_masked = torch.mean(_pfs / _pfd) + if not self.use_huber: + loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) + else: + if not self.f_use_norm: + abs_e = torch.abs(diff_f_3d) + quad = 0.5 * torch.square(diff_f_3d) + lin = self._huber_delta_force * ( + abs_e - 0.5 * self._huber_delta_force + ) + huber_elem = torch.where( + abs_e <= self._huber_delta_force, quad, lin + ) + huber_masked = huber_elem * maskf_col + per_frame_dof = maskf.sum(dim=-1) * 3 + else: + diff_3 = (force_label - force_pred).reshape( + _nf, _nloc, 3 + ) + norm_2d = torch.linalg.vector_norm( + diff_3.reshape(-1, 3), ord=2, dim=1 + ).reshape(_nf, _nloc) + abs_n = norm_2d + quad_n = 0.5 * torch.square(norm_2d) + lin_n = self._huber_delta_force * ( + abs_n - 0.5 * self._huber_delta_force + ) + huber_n = torch.where( + abs_n <= self._huber_delta_force, quad_n, lin_n + ) + huber_masked = (huber_n * maskf).reshape(_nf, _nloc, 1) + per_frame_dof = maskf.sum(dim=-1) + per_frame_sum = huber_masked.reshape(_nf, -1).sum(dim=-1) + l_huber_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l_huber_masked + else: + if not self.use_huber: + loss += (pref_f * l2_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + else: + if not self.f_use_norm: + l_huber_loss = custom_huber_loss( + force_pred.reshape(-1), + force_label.reshape(-1), + delta=self._huber_delta_force, + ) + else: + force_diff_norm = torch.linalg.vector_norm( + (force_label - force_pred).reshape(-1, 3), + ord=2, + dim=1, + keepdim=True, + ) + l_huber_loss = custom_huber_loss( + force_diff_norm, + torch.zeros_like(force_diff_norm), + delta=self._huber_delta_force, + ) + loss += pref_f * l_huber_loss + rmse_f = ( + l2_f_masked.sqrt() + if maskf is not None + else l2_force_loss.sqrt() + ) + more_loss["rmse_f"] = self.display_if_exist( + rmse_f.detach(), find_force + ) + elif self.loss_func == "mae": + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + if not self.f_use_norm: + abs_f = torch.abs(diff_f_3d) * maskf_col + per_frame_sum = abs_f.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l1_f_masked = torch.mean(per_frame_sum / per_frame_dof) + else: + diff_3 = (force_label - force_pred).reshape(_nf, _nloc, 3) + norm_2d = torch.linalg.vector_norm( + diff_3.reshape(-1, 3), ord=2, dim=1 + ).reshape(_nf, _nloc) + masked_norm = norm_2d * maskf + per_frame_sum = masked_norm.sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) + l1_f_masked = torch.mean(per_frame_sum / per_frame_dof) + more_loss["mae_f"] = self.display_if_exist( + l1_f_masked.detach(), find_force + ) + loss += (pref_f * l1_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - force_pred.reshape(-1), + l1_force_loss = F.l1_loss( force_label.reshape(-1), - delta=self._huber_delta_force, + force_pred.reshape(-1), + reduction="mean", ) else: - force_diff_norm = torch.linalg.vector_norm( + l1_force_loss = torch.linalg.vector_norm( (force_label - force_pred).reshape(-1, 3), ord=2, dim=1, keepdim=True, - ) - l_huber_loss = custom_huber_loss( - force_diff_norm, - torch.zeros_like(force_diff_norm), - delta=self._huber_delta_force, - ) - loss += pref_f * l_huber_loss - rmse_f = l2_force_loss.sqrt() - more_loss["rmse_f"] = self.display_if_exist( - rmse_f.detach(), find_force - ) - elif self.loss_func == "mae": - if not self.f_use_norm: - l1_force_loss = F.l1_loss( - force_label.reshape(-1), - force_pred.reshape(-1), - reduction="mean", + ).mean() + more_loss["mae_f"] = self.display_if_exist( + l1_force_loss.detach(), find_force ) - else: - l1_force_loss = torch.linalg.vector_norm( - (force_label - force_pred).reshape(-1, 3), - ord=2, - dim=1, - keepdim=True, - ).mean() - more_loss["mae_f"] = self.display_if_exist( - l1_force_loss.detach(), find_force - ) - loss += (pref_f * l1_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) + loss += (pref_f * l1_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for force loss." ) if mae: - mae_f = torch.mean(torch.abs(diff_f)) + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_f = torch.abs(diff_f_3d) * maskf_col + per_frame_sum = abs_f.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + mae_f = torch.mean(per_frame_sum / per_frame_dof) + else: + mae_f = torch.mean(torch.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist( mae_f.detach(), find_force ) @@ -400,17 +527,49 @@ def forward( more_loss["l2_pref_force_loss"] = self.display_if_exist( l2_pref_force_loss.detach(), find_atom_pref ) - loss += (pref_pf * l2_pref_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_pf = l2_pref_force_loss.sqrt() - more_loss["rmse_pf"] = self.display_if_exist( - rmse_pf.detach(), find_atom_pref - ) + if maskf is not None: + # Idiom 1 with pref weight (ncomp=3). + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + pf_3d = atom_pref.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + sq_pf = torch.square(diff_f_3d) * pf_3d * maskf_col + per_frame_sum = sq_pf.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l2_pf_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_pf * l2_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_pf = l2_pf_masked.sqrt() + more_loss["rmse_pf"] = self.display_if_exist( + rmse_pf.detach(), find_atom_pref + ) + else: + loss += (pref_pf * l2_pref_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + rmse_pf = l2_pref_force_loss.sqrt() + more_loss["rmse_pf"] = self.display_if_exist( + rmse_pf.detach(), find_atom_pref + ) elif self.loss_func == "mae": l1_pref_force_loss = (torch.abs(diff_f) * atom_pref_reshape).mean() - loss += (pref_pf * l1_pref_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss.detach(), find_atom_pref - ) + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + pf_3d = atom_pref.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_pf = torch.abs(diff_f_3d) * pf_3d * maskf_col + per_frame_sum = abs_pf.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l1_pf_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_pf * l1_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_pf"] = self.display_if_exist( + l1_pf_masked.detach(), find_atom_pref + ) + else: + loss += (pref_pf * l1_pref_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + more_loss["mae_pf"] = self.display_if_exist( + l1_pref_force_loss.detach(), find_atom_pref + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." @@ -420,15 +579,35 @@ def forward( drdq = label["drdq"] find_drdq = label.get("find_drdq", 0.0) pref_gf = pref_gf * find_drdq - force_reshape_nframes = force_pred.reshape(-1, natoms * 3) - force_label_reshape_nframes = force_label.reshape(-1, natoms * 3) - drdq_reshape = drdq.reshape(-1, natoms * 3, self.numb_generalized_coord) - gen_force_label = torch.einsum( - "bij,bi->bj", drdq_reshape, force_label_reshape_nframes - ) - gen_force = torch.einsum( - "bij,bi->bj", drdq_reshape, force_reshape_nframes - ) + if maskf is not None: + # Mask per-atom forces before projecting onto generalized coords. + f_3d = force_pred.reshape(_nf, _nloc, 3) * maskf.reshape( + _nf, _nloc, 1 + ) + f_hat_3d = force_label.reshape(_nf, _nloc, 3) * maskf.reshape( + _nf, _nloc, 1 + ) + f_flat = f_3d.reshape(_nf, _nloc * 3) + f_hat_flat = f_hat_3d.reshape(_nf, _nloc * 3) + drdq_reshape = drdq.reshape( + _nf, _nloc * 3, self.numb_generalized_coord + ) + gen_force = torch.einsum("bij,bi->bj", drdq_reshape, f_flat) + gen_force_label = torch.einsum( + "bij,bi->bj", drdq_reshape, f_hat_flat + ) + else: + force_reshape_nframes = force_pred.reshape(-1, natoms * 3) + force_label_reshape_nframes = force_label.reshape(-1, natoms * 3) + drdq_reshape = drdq.reshape( + -1, natoms * 3, self.numb_generalized_coord + ) + gen_force_label = torch.einsum( + "bij,bi->bj", drdq_reshape, force_label_reshape_nframes + ) + gen_force = torch.einsum( + "bij,bi->bj", drdq_reshape, force_reshape_nframes + ) diff_gen_force = gen_force_label - gen_force l2_gen_force_loss = torch.square(diff_gen_force).mean() if not self.inference: @@ -444,43 +623,78 @@ def forward( 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) + v2d = model_pred["virial"].reshape(-1, 9) + v_hat_2d = label["virial"].reshape(-1, 9) + diff_v = v_hat_2d - v2d if self.loss_func == "mse": l2_virial_loss = torch.mean(torch.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) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization. + se_v = torch.square(diff_v) # [nf, 9] + per_frame_v = torch.mean(se_v, dim=-1) # [nf] + if not self.use_huber: + loss += pref_v * torch.mean(per_frame_v * inv**norm_exp) + else: + inv_col = inv.reshape(_nf, 1) + l_huber_v = custom_huber_loss( + inv_col * v2d, + inv_col * v_hat_2d, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_v + rmse_v = torch.sqrt(torch.mean(per_frame_v * inv**2)) + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial + ) 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, + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * v2d.reshape(-1), + atom_norm * v_hat_2d.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 ) - 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 = F.l1_loss( - label["virial"].reshape(-1), - model_pred["virial"].reshape(-1), + v_hat_2d.reshape(-1), + v2d.reshape(-1), reduction="mean", ) - 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 maskf is not None: + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_v = torch.mean(abs_v, dim=-1) # [nf] + l1_v_masked = torch.mean(per_frame_v * inv) + loss += pref_v * l1_v_masked + more_loss["mae_v"] = self.display_if_exist( + l1_v_masked.detach(), find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss.detach() * atom_norm, + find_virial, + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." ) if mae: - mae_v = torch.mean(torch.abs(diff_v)) * atom_norm + if maskf is not None: + abs_v = torch.abs(diff_v) + per_frame_v = torch.mean(abs_v, dim=-1) + mae_v = torch.mean(per_frame_v * inv) + else: + mae_v = torch.mean(torch.abs(diff_v)) * atom_norm more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial) if self.has_ae and "atom_energy" in model_pred and "atom_ener" in label: @@ -499,29 +713,71 @@ def forward( more_loss["l2_atom_ener_loss"] = self.display_if_exist( l2_atom_ener_loss.detach(), find_atom_ener ) - if not self.use_huber: - loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae_2d = atom_ener.reshape(_nf, _nloc) + ae_hat_2d = atom_ener_label.reshape(_nf, _nloc) + sq_ae = torch.square(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = sq_ae.sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] + l2_ae_masked = torch.mean(per_frame_sum / per_frame_dof) + if not self.use_huber: + loss += (pref_ae * l2_ae_masked).to(GLOBAL_PT_FLOAT_PRECISION) + else: + diff_ae = ae_hat_2d - ae_2d + abs_ae = torch.abs(diff_ae) + quad_ae = 0.5 * torch.square(diff_ae) + lin_ae = self._huber_delta_energy * ( + abs_ae - 0.5 * self._huber_delta_energy + ) + huber_ae = torch.where( + abs_ae <= self._huber_delta_energy, quad_ae, lin_ae + ) + huber_ae_m = huber_ae * maskf + l_huber_ae = torch.mean(huber_ae_m.sum(dim=-1) / per_frame_dof) + loss += pref_ae * l_huber_ae + rmse_ae = l2_ae_masked.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) else: - l_huber_loss = custom_huber_loss( - atom_ener_reshape, - atom_ener_label_reshape, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += (pref_ae * l2_atom_ener_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + else: + l_huber_loss = custom_huber_loss( + atom_ener_reshape, + atom_ener_label_reshape, + delta=self._huber_delta_energy, + ) + loss += pref_ae * l_huber_loss + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener ) - loss += pref_ae * l_huber_loss - rmse_ae = l2_atom_ener_loss.sqrt() - more_loss["rmse_ae"] = self.display_if_exist( - rmse_ae.detach(), find_atom_ener - ) elif self.loss_func == "mae": l1_atom_ener_loss = F.l1_loss( atom_ener_reshape, atom_ener_label_reshape, reduction="mean", ) - loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss.detach(), find_atom_ener - ) + if maskf is not None: + ae_2d = atom_ener.reshape(_nf, _nloc) + ae_hat_2d = atom_ener_label.reshape(_nf, _nloc) + abs_ae = torch.abs(ae_hat_2d - ae_2d) * maskf + per_frame_sum = abs_ae.sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) + l1_ae_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_ae * l1_ae_masked).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_ae_masked.detach(), find_atom_ener + ) + else: + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atomic energy loss." diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 66473cd3f4..84c229e1c2 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -26,7 +26,7 @@ def _masked_force_mag_tensors( label: dict[str, torch.Tensor], model_pred: dict[str, torch.Tensor], -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Collect magnetic-force labels and predictions on spin-active atoms. Parameters @@ -42,54 +42,11 @@ def _masked_force_mag_tensors( Reference magnetic forces with shape ``(n_mag, 3)``. pred_fm : torch.Tensor Predicted magnetic forces with shape ``(n_mag, 3)``. - mag_counts : torch.Tensor - Number of spin-active atoms in each frame, with shape ``(nframes,)``. """ atomic_mask = model_pred["mask_mag"].expand(-1, -1, 3) label_fm = label["force_mag"][atomic_mask].reshape(-1, 3) pred_fm = model_pred["force_mag"][atomic_mask].reshape(-1, 3) - mag_counts = model_pred["mask_mag"].sum(dim=(1, 2)).to(torch.int64) - return label_fm, pred_fm, mag_counts - - -def _mean_within_segments( - values: torch.Tensor, - segment_lengths: torch.Tensor, -) -> torch.Tensor: - """Reduce ``values`` to a per-segment mean over contiguous frame blocks. - - Parameters - ---------- - values : torch.Tensor - Values laid out in frame order, with shape ``(n_values,)``. - segment_lengths : torch.Tensor - Length of each segment, with shape ``(n_segments,)``. - - Returns - ------- - torch.Tensor - Per-segment means with shape ``(n_segments,)``. Empty segments return - zero so frame-weighted reductions remain finite. - """ - nsegments = segment_lengths.shape[0] - if values.numel() == 0: - return torch.zeros( - nsegments, - dtype=values.dtype, - device=values.device, - ) - - segment_ids = torch.repeat_interleave( - torch.arange(nsegments, device=values.device, dtype=torch.long), - segment_lengths, - ) - totals = torch.zeros(nsegments, dtype=values.dtype, device=values.device) - totals.scatter_add_(0, segment_ids, values) - - means = torch.zeros(nsegments, dtype=values.dtype, device=values.device) - nonempty = segment_lengths > 0 - means[nonempty] = totals[nonempty] / segment_lengths[nonempty].to(values.dtype) - return means + return label_fm, pred_fm class EnergySpinLoss(TaskLoss): @@ -212,7 +169,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_fr = self.limit_pref_fr + (self.start_pref_fr - self.limit_pref_fr) * coef @@ -228,6 +185,20 @@ def forward( # - norm_exp=2 (intensive_ener_virial=True): loss uses 1/N² scaling, making it independent of system size # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size norm_exp = 2 if self.intensive_ener_virial else 1 + + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghost padding atoms. + if "mask" in model_pred: + maskf = model_pred["mask"] # [nf, nloc], float + real_natoms_f = torch.sum(maskf, dim=-1) # [nf] + inv = (1.0 / real_natoms_f).reshape(-1) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). + maskf = None + if self.has_e and "energy" in model_pred and "energy" in label: energy_pred = model_pred["energy"] energy_label = label["energy"] @@ -247,39 +218,65 @@ def forward( find_energy = label.get("find_energy", 0.0) pref_e = pref_e * find_energy if self.loss_func == "mse": - l2_ener_loss = torch.mean(torch.square(energy_pred - energy_label)) - if not self.inference: - more_loss["l2_ener_loss"] = self.display_if_exist( - l2_ener_loss.detach(), find_energy + se_e = torch.square(energy_pred - energy_label) # [nf, k] + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + per_frame_e = torch.mean(se_e.reshape(_nf, -1), dim=-1) # [nf] + if not self.inference: + more_loss["l2_ener_loss"] = self.display_if_exist( + torch.mean(per_frame_e).detach(), find_energy + ) + loss += pref_e * torch.mean(per_frame_e * inv**norm_exp) + rmse_e = torch.sqrt(torch.mean(per_frame_e * inv**2)) + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy + ) + else: + l2_ener_loss = torch.mean(se_e) + if not self.inference: + more_loss["l2_ener_loss"] = self.display_if_exist( + l2_ener_loss.detach(), find_energy + ) + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + rmse_e = l2_ener_loss.sqrt() * atom_norm + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy ) - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) - rmse_e = l2_ener_loss.sqrt() * atom_norm - more_loss["rmse_e"] = self.display_if_exist( - rmse_e.detach(), find_energy - ) # more_loss['log_keys'].append('rmse_e') elif self.loss_func == "mae": l1_ener_loss = F.l1_loss( energy_pred.reshape(-1), energy_label.reshape(-1), - reduction="sum", - ) - loss += 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, + reduction="mean", ) + if maskf is not None: + # Idiom 2 (extensive) with abs: per-frame normalization by real-atom count. + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) # [nf] + l1_ener_masked = torch.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked.detach(), find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss.detach() * atom_norm, find_energy + ) # more_loss['log_keys'].append('rmse_e') else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + mae_e = torch.mean(per_frame_ae * inv) + else: + mae_e = ( + torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + ) more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy) mae_e_all = torch.mean(torch.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist( @@ -290,31 +287,69 @@ def forward( find_force_r = label.get("find_force", 0.0) pref_fr = pref_fr * find_force_r if self.loss_func == "mse": - diff_fr = label["force"] - model_pred["force"] - l2_force_real_loss = torch.mean(torch.square(diff_fr)) - if not self.inference: - more_loss["l2_force_r_loss"] = self.display_if_exist( - l2_force_real_loss.detach(), find_force_r + diff_fr = label["force"] - model_pred["force"] # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + maskf_col = maskf.reshape(_nf, _nloc, 1) # [nf, nloc, 1] + sq_fr = torch.square(diff_fr) * maskf_col # [nf, nloc, 3] + per_frame_fr_sum = sq_fr.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_fr_dof = maskf.sum(dim=-1) * 3 # [nf] + l2_force_real_loss = torch.mean(per_frame_fr_sum / per_frame_fr_dof) + if not self.inference: + more_loss["l2_force_r_loss"] = self.display_if_exist( + l2_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_fr = l2_force_real_loss.sqrt() + more_loss["rmse_fr"] = self.display_if_exist( + rmse_fr.detach(), find_force_r ) - loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_fr = l2_force_real_loss.sqrt() - more_loss["rmse_fr"] = self.display_if_exist( - rmse_fr.detach(), find_force_r - ) - if mae: - mae_fr = torch.mean(torch.abs(diff_fr)) - more_loss["mae_fr"] = self.display_if_exist( - mae_fr.detach(), find_force_r + if mae: + abs_fr = torch.abs(diff_fr) * maskf_col + per_frame_mae_sum = abs_fr.reshape(_nf, -1).sum(dim=-1) + mae_fr = torch.mean(per_frame_mae_sum / per_frame_fr_dof) + more_loss["mae_fr"] = self.display_if_exist( + mae_fr.detach(), find_force_r + ) + else: + l2_force_real_loss = torch.mean(torch.square(diff_fr)) + if not self.inference: + more_loss["l2_force_r_loss"] = self.display_if_exist( + l2_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_fr = l2_force_real_loss.sqrt() + more_loss["rmse_fr"] = self.display_if_exist( + rmse_fr.detach(), find_force_r ) + if mae: + mae_fr = torch.mean(torch.abs(diff_fr)) + more_loss["mae_fr"] = self.display_if_exist( + mae_fr.detach(), find_force_r + ) elif self.loss_func == "mae": - l1_force_real_loss = F.l1_loss( - label["force"], model_pred["force"], reduction="none" - ) - more_loss["mae_fr"] = self.display_if_exist( - l1_force_real_loss.mean().detach(), find_force_r - ) - l1_force_real_loss = l1_force_real_loss.sum(-1).mean(-1).sum() - loss += (pref_fr * l1_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + abs_diff_fr = torch.abs( + label["force"] - model_pred["force"] + ) # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3) with abs. + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_fr = abs_diff_fr * maskf_col + per_frame_sum = abs_fr.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * 3 # [nf] + l1_force_real_masked = torch.mean(per_frame_sum / per_frame_dof) + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_masked.detach(), find_force_r + ) + loss += (pref_fr * l1_force_real_masked).to( + GLOBAL_PT_FLOAT_PRECISION + ) + else: + l1_force_real_loss = torch.mean(abs_diff_fr) + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l1_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for real force loss." @@ -323,7 +358,7 @@ def forward( if self.has_fm and "force_mag" in model_pred and "force_mag" in label: find_force_m = label.get("find_force_mag", 0.0) pref_fm = pref_fm * find_force_m - label_fm, pred_fm, mag_counts = _masked_force_mag_tensors(label, model_pred) + label_fm, pred_fm = _masked_force_mag_tensors(label, model_pred) if self.loss_func == "mse": diff_fm = label_fm - pred_fm l2_force_mag_loss = torch.mean(torch.square(diff_fm)) @@ -348,15 +383,14 @@ def forward( mae_fm.detach(), find_force_m ) elif self.loss_func == "mae": - per_atom_l1 = F.l1_loss(label_fm, pred_fm, reduction="none").sum(-1) - # Report the per-component magnetic-force MAE (mean over the - # masked components), consistent with ``mae_fr``, the ``mse`` - # branch, and the backend-independent reference. The loss - # contribution below keeps the per-atom reduction (summed xyz). + # Mean over frames, magnetic atoms and xyz (same reduction as + # force_mag MSE, force_real MAE and the displayed mae_fm) so the + # loss is batch-size independent: a 2-frame batch equals the mean + # of the two single-frame losses. + l1_force_mag_loss = torch.mean(torch.abs(label_fm - pred_fm)) more_loss["mae_fm"] = self.display_if_exist( - torch.mean(torch.abs(label_fm - pred_fm)).detach(), find_force_m + l1_force_mag_loss.detach(), find_force_m ) - l1_force_mag_loss = _mean_within_segments(per_atom_l1, mag_counts).sum() loss += (pref_fm * torch.nan_to_num(l1_force_mag_loss)).to( GLOBAL_PT_FLOAT_PRECISION ) @@ -370,68 +404,132 @@ def forward( atom_ener_label = label["atom_ener"] find_atom_ener = label.get("find_atom_ener", 0.0) pref_ae = pref_ae * find_atom_ener - atom_ener_reshape = atom_ener.reshape(-1) - atom_ener_label_reshape = atom_ener_label.reshape(-1) - if self.loss_func == "mse": - l2_atom_ener_loss = torch.square( - atom_ener_label_reshape - atom_ener_reshape - ).mean() - if not self.inference: - more_loss["l2_atom_ener_loss"] = self.display_if_exist( - l2_atom_ener_loss.detach(), find_atom_ener + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae = atom_ener.reshape(_nf, _nloc, 1) + ae_label = atom_ener_label.reshape(_nf, _nloc, 1) + maskf_col = maskf.reshape(_nf, _nloc, 1) # [nf, nloc, 1] + if self.loss_func == "mse": + sq = torch.square(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = sq.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] (ncomp=1) + l2_atom_ener_loss = torch.mean(per_frame_sum / per_frame_dof) + if not self.inference: + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss.detach(), find_atom_ener + ) + loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) + elif self.loss_func == "mae": + abs_diff = torch.abs(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = abs_diff.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] (ncomp=1) + l1_atom_ener_loss = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) + else: + raise NotImplementedError( + f"Loss type {self.loss_func} is not implemented for atomic energy loss." ) - loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_ae = l2_atom_ener_loss.sqrt() - more_loss["rmse_ae"] = self.display_if_exist( - rmse_ae.detach(), find_atom_ener - ) - elif self.loss_func == "mae": - l1_atom_ener_loss = F.l1_loss( - atom_ener_reshape, - atom_ener_label_reshape, - reduction="mean", - ) - loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss.detach(), find_atom_ener - ) else: - raise NotImplementedError( - f"Loss type {self.loss_func} is not implemented for atomic energy loss." - ) + atom_ener_reshape = atom_ener.reshape(-1) + atom_ener_label_reshape = atom_ener_label.reshape(-1) + if self.loss_func == "mse": + l2_atom_ener_loss = torch.square( + atom_ener_label_reshape - atom_ener_reshape + ).mean() + if not self.inference: + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss.detach(), find_atom_ener + ) + loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) + elif self.loss_func == "mae": + l1_atom_ener_loss = F.l1_loss( + atom_ener_reshape, + atom_ener_label_reshape, + reduction="mean", + ) + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) + else: + raise NotImplementedError( + f"Loss type {self.loss_func} is not implemented for atomic energy loss." + ) 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) + diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9) # [nf, 9] if self.loss_func == "mse": - l2_virial_loss = torch.mean(torch.square(diff_v)) - if not self.inference: - more_loss["l2_virial_loss"] = self.display_if_exist( - l2_virial_loss.detach(), find_virial + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization by real-atom count. + se_v = torch.square(diff_v) # [nf, 9] + per_frame_v = torch.mean(se_v, dim=-1) # [nf] + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + torch.mean(per_frame_v).detach(), find_virial + ) + loss += pref_v * torch.mean(per_frame_v * inv**norm_exp) + rmse_v = torch.sqrt(torch.mean(per_frame_v * inv**2)) + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial ) - loss += atom_norm**norm_exp * (pref_v * l2_virial_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 = torch.mean(torch.abs(diff_v)) * atom_norm - more_loss["mae_v"] = self.display_if_exist( - mae_v.detach(), find_virial + if mae: + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_mae_v = torch.mean(abs_v, dim=-1) # [nf] + mae_v = torch.mean(per_frame_mae_v * inv) + more_loss["mae_v"] = self.display_if_exist( + mae_v.detach(), find_virial + ) + else: + l2_virial_loss = torch.mean(torch.square(diff_v)) + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + l2_virial_loss.detach(), find_virial + ) + loss += atom_norm**norm_exp * (pref_v * l2_virial_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 = torch.mean(torch.abs(diff_v)) * atom_norm + more_loss["mae_v"] = self.display_if_exist( + mae_v.detach(), find_virial + ) elif self.loss_func == "mae": l1_virial_loss = F.l1_loss( label["virial"].reshape(-1), model_pred["virial"].reshape(-1), reduction="mean", ) - 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 maskf is not None: + # Idiom 2 (extensive, k=9) with abs: per-frame normalization by real-atom count. + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_v = torch.mean(abs_v, dim=-1) # [nf] + l1_virial_masked = torch.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked.detach(), find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss.detach() * atom_norm, find_virial + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index 51eb7a9319..2f081d2abe 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -40,6 +40,36 @@ def label_requirement(self) -> list[DataRequirementItem]: """Return data label requirements needed for this loss calculation.""" pass + @staticmethod + def _inject_atom_mask( + model_pred: dict[str, torch.Tensor], + input_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Recover the per-atom mask from atype for mixed_type batches. + + The exported forward drops the model's per-atom mask, so reconstruct it + here from ``atype`` (training-only). Ghost atoms have ``atype < 0``. + An all-ones mask is produced for non-mixed batches, keeping the loss + bit-identical to the pre-fix behavior. + + Parameters + ---------- + model_pred : dict[str, torch.Tensor] + Model predictions (modified in-place). + input_dict : dict[str, torch.Tensor] + Model inputs; must contain ``"atype"`` for injection to occur. + + Returns + ------- + dict[str, torch.Tensor] + ``model_pred`` with ``"mask"`` added if not already present. + """ + if "mask" not in model_pred and "atype" in input_dict: + atype = input_dict["atype"] + ref = model_pred.get("energy", input_dict.get("coord", atype)) + model_pred["mask"] = (atype >= 0).to(ref.dtype) + return model_pred + @staticmethod def display_if_exist(loss: torch.Tensor, find_property: float) -> torch.Tensor: """Display NaN if labeled property is not found. diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index af799d1a89..df865b10f1 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -100,14 +100,24 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) var_name = self.var_name nbz = model_pred[var_name].shape[0] assert model_pred[var_name].shape == (nbz, self.task_dim) assert label[var_name].shape == (nbz, self.task_dim) if not self.intensive: - model_pred[var_name] = model_pred[var_name] / natoms - label[var_name] = label[var_name] / natoms + if "mask" in model_pred: + # Per-frame real atom count: shape [nf] → broadcast over [nf, task_dim]. + real_natoms = ( + torch.sum(model_pred["mask"], dim=-1) + .to(dtype=model_pred[var_name].dtype) + .reshape(-1, 1) + ) + model_pred[var_name] = model_pred[var_name] / real_natoms + label[var_name] = label[var_name] / real_natoms + else: + model_pred[var_name] = model_pred[var_name] / natoms + label[var_name] = label[var_name] / natoms if self.out_std is None: out_std = model.atomic_model.out_std[0][0] diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index f329b79b20..c46c2ec0dd 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -101,7 +101,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) del learning_rate, mae if self.enable_atomic_weight: @@ -129,8 +129,16 @@ def forward( ) diff = diff * atomic_weight if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = local_tensor_pred.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + diff3d = diff.reshape(nf, natoms, self.tensor_size) + sq = torch.square(diff3d) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.tensor_size # [nf] + l2_local_loss = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss = torch.mean(torch.square(diff)) if not self.inference: more_loss[f"l2_local_{self.tensor_name}_loss"] = self.display_if_exist( l2_local_loss.detach(), find_local @@ -152,15 +160,12 @@ def forward( ) global_tensor_label = label[self.label_name].reshape([-1, self.tensor_size]) diff = global_tensor_pred - global_tensor_label + # idiom 3: global tensor is already padding-invariant; plain mean suffices + l2_global_loss = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss = torch.mean(torch.square(diff)) if not self.inference: more_loss[f"l2_global_{self.tensor_name}_loss"] = self.display_if_exist( l2_global_loss.detach(), find_global diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py new file mode 100644 index 0000000000..1109c21131 --- /dev/null +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -0,0 +1,2240 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reusable grad-accumulation invariant harness for dpmodel loss tests. + +This module provides ``assert_grad_accum_invariant`` for Tasks 2-5 that +verify the loss on a padded multi-frame batch equals mean(per_frame_loss). + +The dpmodel losses accept numpy arrays (via the array_api_compat backend). + +Scope / follow-ups (mixed_type padding fix, PR #5738) +---------------------------------------------------- +- The TF backend loss is not covered here and still has the mixed_type + dilution behavior; tracked in deepmodeling/deepmd-kit#5760. +- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; + tracked in deepmodeling/deepmd-kit#5761. +- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean + reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real + MAE; the grad-accum invariant is asserted by the test below. + +Constants +--------- +NA = 3 # real atoms in the short frame +NB = 5 # real atoms in the full-width frame (NB == NP) +NP = 5 # padded width (nloc) +""" + +import numpy as np + +from deepmd.dpmodel.loss.dos import ( + DOSLoss, +) +from deepmd.dpmodel.loss.ener import ( + EnergyLoss, +) +from deepmd.dpmodel.loss.ener_spin import EnergySpinLoss as EnergySpinLossDPModel +from deepmd.dpmodel.loss.property import ( + PropertyLoss, +) +from deepmd.dpmodel.loss.tensor import ( + TensorLoss, +) + +# --------------------------------------------------------------------------- +# Constants used by the multi-frame test harness (Tasks 2-5) +# --------------------------------------------------------------------------- + +NA = 3 # real atoms in frame A +NB = 5 # real atoms in frame B (== NP, so frame B is fully real) +NP = 5 # padded width (nloc) + + +# --------------------------------------------------------------------------- +# Reusable harness (used by Tasks 2-5; imported from here) +# --------------------------------------------------------------------------- + + +def assert_grad_accum_invariant( + loss_fn, + make_batch_A, + make_batch_B, + make_padded_batch, + rtol: float = 1e-5, + atol: float = 1e-6, +) -> None: + """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + + The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms + padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as + processing each frame separately and averaging. + + Parameters + ---------- + loss_fn : callable + Signature ``(model_pred, label, natoms) -> float``. The function + receives numpy-dict inputs and returns a scalar float. + make_batch_A : callable + Returns ``(model_pred, label, natoms)`` for frame A alone (1 frame, NA atoms). + make_batch_B : callable + Returns ``(model_pred, label, natoms)`` for frame B alone (1 frame, NB atoms). + make_padded_batch : callable + Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch + (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + rtol : float + Relative tolerance for ``np.isclose``. + atol : float + Absolute tolerance for ``np.isclose``. + """ + pred_A, label_A, natoms_A = make_batch_A() + pred_B, label_B, natoms_B = make_batch_B() + pred_pad, label_pad, natoms_pad = make_padded_batch() + + loss_A = float(loss_fn(pred_A, label_A, natoms_A)) + loss_B = float(loss_fn(pred_B, label_B, natoms_B)) + ref = 0.5 * (loss_A + loss_B) + + loss_pad = float(loss_fn(pred_pad, label_pad, natoms_pad)) + + assert np.isclose(loss_pad, ref, rtol=rtol, atol=atol), ( + f"Grad-accum invariant violated: padded_loss={loss_pad:.8f}, " + f"ref={ref:.8f}, diff={abs(loss_pad - ref):.2e}" + ) + + +# --------------------------------------------------------------------------- +# Helpers for constructing test data +# --------------------------------------------------------------------------- + +RNG = np.random.default_rng(42) +NUMB_DOS = 4 +TENSOR_SIZE = 3 + + +def _rnd(*shape): + return RNG.standard_normal(shape).astype(np.float64) + + +# --------------------------------------------------------------------------- +# Task 2: DOSLoss -- atomic (ados / acdf) and global (dos / cdf) +# --------------------------------------------------------------------------- + + +class TestDOSLossAtomicGradAccum: + """Per-frame masked mean (idiom 1) for atomic dos / acdf terms.""" + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=0.0, + limit_pref_dos=0.0, + start_pref_cdf=0.0, + limit_pref_cdf=0.0, + start_pref_ados=1.0, + limit_pref_ados=1.0, + start_pref_acdf=1.0, + limit_pref_acdf=1.0, + ) + + def _atom_dos_A(self): + return _rnd(NA, NUMB_DOS) + + def _atom_dos_B(self): + return _rnd(NB, NUMB_DOS) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def _make_batch_A(self): + pred = _rnd(NA, NUMB_DOS) + label = _rnd(NA, NUMB_DOS) + model_pred = { + "atom_dos": pred, + "mask": np.ones((1, NA), dtype=np.float64), + } + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + return model_pred, label_dict, NA + + def _make_batch_B(self): + pred = _rnd(NB, NUMB_DOS) + label = _rnd(NB, NUMB_DOS) + model_pred = { + "atom_dos": pred, + "mask": np.ones((1, NB), dtype=np.float64), + } + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + return model_pred, label_dict, NB + + def _make_padded_batch(self, pred_A_data, label_A_data, pred_B_data, label_B_data): + # Frame A: NA real atoms padded to NP; ghost pred/label = 0 + pred_A_pad = np.zeros((NP, NUMB_DOS), dtype=np.float64) + pred_A_pad[:NA] = pred_A_data + label_A_pad = np.zeros((NP, NUMB_DOS), dtype=np.float64) + label_A_pad[:NA] = label_A_data + mask_A = np.array([[1.0] * NA + [0.0] * (NP - NA)], dtype=np.float64) + + # Frame B: NB == NP, all real + mask_B = np.ones((1, NB), dtype=np.float64) + + atom_dos_pad = np.concatenate( + [pred_A_pad, pred_B_data], axis=0 + ) # [NP+NB, ncomp] + atom_dos_label = np.concatenate([label_A_pad, label_B_data], axis=0) + mask_pad = np.concatenate([mask_A, mask_B], axis=0) # [2, NP] + model_pred = {"atom_dos": atom_dos_pad, "mask": mask_pad} + label_dict = {"atom_dos": atom_dos_label, "find_atom_dos": 1.0} + return model_pred, label_dict, NP + + def test_ados_grad_accum_invariant(self): + """Atomic dos per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd(NA, NUMB_DOS) + label_A = _rnd(NA, NUMB_DOS) + pred_B = _rnd(NB, NUMB_DOS) + label_B = _rnd(NB, NUMB_DOS) + + def make_A(): + return ( + {"atom_dos": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"atom_dos": label_A, "find_atom_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + {"atom_dos": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"atom_dos": label_B, "find_atom_dos": 1.0}, + NB, + ) + + def make_padded(): + return self._make_padded_batch(pred_A, label_A, pred_B, label_B) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_acdf_grad_accum_invariant(self): + """Atomic cdf per-frame masked mean meets the grad-accum invariant.""" + # acdf uses cumsum of atom_dos -- same data pipeline + self.test_ados_grad_accum_invariant() + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(NB, NUMB_DOS) + label = _rnd(NB, NUMB_DOS) + with_mask = { + "atom_dos": pred, + "mask": np.ones((1, NB), dtype=np.float64), + } + without_mask = {"atom_dos": pred} + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +class TestDOSLossGlobalGradAccum: + """Plain mean (idiom 3) for global dos / cdf terms.""" + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_dos_grad_accum_invariant(self): + """Global dos plain mean meets the grad-accum invariant.""" + pred_A = _rnd(1, NUMB_DOS) + label_A = _rnd(1, NUMB_DOS) + pred_B = _rnd(1, NUMB_DOS) + label_B = _rnd(1, NUMB_DOS) + + def make_A(): + return ( + {"dos": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"dos": label_A, "find_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + {"dos": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"dos": label_B, "find_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = np.concatenate([pred_A, pred_B], axis=0) # [2, ncomp] + label_pad = np.concatenate([label_A, label_B], axis=0) + mask_pad = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 + ) + return ( + {"dos": pred_pad, "mask": mask_pad}, + {"dos": label_pad, "find_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_cdf_grad_accum_invariant(self): + """Global cdf plain mean meets the grad-accum invariant.""" + # cdf uses cumsum of dos -- reuse the same logic + self.test_dos_grad_accum_invariant() + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(2, NUMB_DOS) + label = _rnd(2, NUMB_DOS) + with_mask = {"dos": pred, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {"dos": pred} + label_dict = {"dos": label, "find_dos": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +# --------------------------------------------------------------------------- +# Task 2: TensorLoss -- local and global tensor +# --------------------------------------------------------------------------- + + +class TestTensorLossLocalGradAccum: + """Per-frame masked mean (idiom 1) for local tensor term.""" + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=1.0, + pref=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_local_grad_accum_invariant(self): + """Local tensor per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd(NA, TENSOR_SIZE) + label_A = _rnd(NA, TENSOR_SIZE) + pred_B = _rnd(NB, TENSOR_SIZE) + label_B = _rnd(NB, TENSOR_SIZE) + + def make_A(): + return ( + {"dipole": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"atom_dipole": label_A, "find_atom_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + {"dipole": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"atom_dipole": label_B, "find_atom_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = np.zeros((NP, TENSOR_SIZE), dtype=np.float64) + pred_A_pad[:NA] = pred_A + label_A_pad = np.zeros((NP, TENSOR_SIZE), dtype=np.float64) + label_A_pad[:NA] = label_A + mask_A = np.array([[1.0] * NA + [0.0] * (NP - NA)], dtype=np.float64) + mask_B = np.ones((1, NB), dtype=np.float64) + dipole_pad = np.concatenate([pred_A_pad, pred_B], axis=0) + label_pad = np.concatenate([label_A_pad, label_B], axis=0) + mask_pad = np.concatenate([mask_A, mask_B], axis=0) + return ( + {"dipole": dipole_pad, "mask": mask_pad}, + {"atom_dipole": label_pad, "find_atom_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(NB, TENSOR_SIZE) + label = _rnd(NB, TENSOR_SIZE) + with_mask = {"dipole": pred, "mask": np.ones((1, NB), dtype=np.float64)} + without_mask = {"dipole": pred} + label_dict = {"atom_dipole": label, "find_atom_dipole": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +class TestTensorLossGlobalGradAccum: + """Plain mean (idiom 3) for global tensor term.""" + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_global_grad_accum_invariant(self): + """Global tensor plain mean meets the grad-accum invariant.""" + pred_A = _rnd(1, TENSOR_SIZE) + label_A = _rnd(1, TENSOR_SIZE) + pred_B = _rnd(1, TENSOR_SIZE) + label_B = _rnd(1, TENSOR_SIZE) + + def make_A(): + return ( + {"global_dipole": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"dipole": label_A, "find_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + {"global_dipole": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"dipole": label_B, "find_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = np.concatenate([pred_A, pred_B], axis=0) + label_pad = np.concatenate([label_A, label_B], axis=0) + mask_pad = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 + ) + return ( + {"global_dipole": pred_pad, "mask": mask_pad}, + {"dipole": label_pad, "find_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(2, TENSOR_SIZE) + label = _rnd(2, TENSOR_SIZE) + with_mask = {"global_dipole": pred, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {"global_dipole": pred} + label_dict = {"dipole": label, "find_dipole": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +# --------------------------------------------------------------------------- +# Task 3: EnergyLoss -- energy, force, virial, atom_ener, atom_pref +# --------------------------------------------------------------------------- + + +def _full_ener_dicts(nf, nloc, energy_pred, energy_label, mask=None, **overrides): + """Build complete model_pred and label_dict for EnergyLoss.call. + + EnergyLoss.call fetches energy/force/virial/atom_energy unconditionally, + so all keys must be present regardless of which term is under test. + """ + model_pred = { + "energy": energy_pred, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "atom_energy": np.zeros((nf, nloc, 1), dtype=np.float64), + } + label_dict = { + "energy": energy_label, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "atom_ener": np.zeros((nf, nloc, 1), dtype=np.float64), + "atom_pref": np.zeros((nf, nloc * 3), dtype=np.float64), + "find_energy": 1.0, + "find_force": 0.0, + "find_virial": 0.0, + "find_atom_ener": 0.0, + "find_atom_pref": 0.0, + } + if mask is not None: + model_pred["mask"] = mask + model_pred.update({k: v for k, v in overrides.items() if k in model_pred}) + label_dict.update({k: v for k, v in overrides.items() if k in label_dict}) + return model_pred, label_dict + + +def _padded_force(f_A, f_B): + """Stack force arrays for a 2-frame padded batch (NA<=NP, NB==NP).""" + f_A_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_pad[:NA] = f_A + return np.stack([f_A_pad, f_B], axis=0) # [2, NP, 3] + + +def _padded_atom(arr_A, arr_B, ncomp): + """Pad arr_A from [NA, ncomp] to [NP, ncomp] with zeros, stack with arr_B.""" + pad = np.zeros((NP, ncomp), dtype=np.float64) + pad[:NA] = arr_A + return np.stack([pad, arr_B], axis=0) # [2, NP, ncomp] + + +def _padded_atom_flat(arr_A, arr_B, ncomp): + """Pad then reshape to [2, NP*ncomp] (for atom_pref shape).""" + return _padded_atom(arr_A, arr_B, ncomp).reshape(2, NP * ncomp) + + +_MASK_PAD = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 +) # [2, NP] + + +class TestDPModelEnergyLossEnerGradAccum: + """Idiom 2 (extensive) for the energy (has_e) term in EnergyLoss. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, e_A, e_A_hat, mask=np.ones((1, NA), dtype=np.float64) + ) + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, e_B, e_B_hat, mask=np.ones((1, NB), dtype=np.float64) + ) + return p, l, NB + + def make_padded(): + e_pad = np.concatenate([e_A, e_B], axis=0) # [2, 1] + e_hat_pad = np.concatenate([e_A_hat, e_B_hat], axis=0) + p, l = _full_ener_dicts(2, NP, e_pad, e_hat_pad, mask=_MASK_PAD) + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", intensive=False), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_mse_intensive_grad_accum(self): + """Energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", intensive=True), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_mae_grad_accum(self): + """Energy MAE meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mae"), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_huber_grad_accum(self): + """Energy Huber meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", use_huber=True), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _rnd(1, 1) + e_hat = _rnd(1, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_ener_dicts( + 1, NP, e, e_hat, mask=np.ones((1, NP), dtype=np.float64) + ) + p_nomask, l_nomask = _full_ener_dicts(1, NP, e, e_hat) + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nomask, l_nomask, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossForceGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for the force (has_f) term. + + Covers: mse, mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + use_huber=use_huber, + f_use_norm=f_use_norm, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force MSE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss("mse"), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force MAE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss("mae"), f_A, f_A_hat, f_B, f_B_hat) + + def test_huber_grad_accum(self): + """Force Huber meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant( + self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["find_force"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["find_force"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for the virial (has_v) term. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["virial"] = v_A[None] # [1, 9] + l["virial"] = v_A_hat[None] + l["find_virial"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["virial"] = v_B[None] + l["virial"] = v_B_hat[None] + l["find_virial"] = 1.0 + return p, l, NB + + def make_padded(): + v_pad = np.stack([v_A, v_B], axis=0) # [2, 9] + v_hat_pad = np.stack([v_A_hat, v_B_hat], axis=0) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["virial"] = v_pad + l["virial"] = v_hat_pad + l["find_virial"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", intensive=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mae_grad_accum(self): + """Virial MAE meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant(self._make_loss("mae"), v_A, v_A_hat, v_B, v_B_hat) + + def test_huber_grad_accum(self): + """Virial Huber meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", use_huber=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _rnd(9) + v_hat = _rnd(9) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["virial"] = v[None] + l_mask["virial"] = v_hat[None] + l_mask["find_virial"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["virial"] = v[None] + l_nm["virial"] = v_hat[None] + l_nm["find_virial"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for the atom_ener (has_ae) term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["atom_energy"] = ae_A[None] # [1, NA, 1] + l["atom_ener"] = ae_A_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["atom_energy"] = ae_B[None] + l["atom_ener"] = ae_B_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NB + + def make_padded(): + ae_pad = _padded_atom(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom(ae_A_hat, ae_B_hat, 1) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["atom_energy"] = ae_pad + l["atom_ener"] = ae_hat_pad + l["find_atom_ener"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom energy MSE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Atom energy MAE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_huber_grad_accum(self): + """Atom energy Huber (use_huber=True) meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + loss_obj = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + use_huber=True, + ) + self._run_invariant(loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-energy loss as no mask.""" + ae = _rnd(NP, 1) + ae_hat = _rnd(NP, 1) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["atom_energy"] = ae[None] + l_mask["atom_ener"] = ae_hat[None] + l_mask["find_atom_ener"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["atom_energy"] = ae[None] + l_nm["atom_ener"] = ae_hat[None] + l_nm["find_atom_ener"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossAtomPrefGradAccum: + """Idiom 1 with pref weight (ncomp=3) for the atom_pref (has_pf) term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _loss_fn_pf_only(self, loss_obj, model_pred, label, natoms): + """Return the atom_pref contribution to the loss (subtract force loss).""" + # Both pref_f and pref_pf are 1.0 here, but we want ONLY the pf term. + # Use a loss with pref_pf=1 and pref_f=0 to isolate it. + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _make_pf_only_loss(self, loss_func="mse"): + """Loss with pref_f=0 and pref_pf=1 to isolate atom_pref term.""" + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B): + """Invariant for atom_pref using force diff weighted by pref.""" + + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["force"] = f_A[None] + l["force"] = f_A_hat[None] + l["atom_pref"] = pf_A.reshape(1, NA * 3) + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["atom_pref"] = pf_B.reshape(1, NB * 3) + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + pf_pad = _padded_atom_flat(pf_A, pf_B, 3) # [2, NP*3] + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["atom_pref"] = pf_pad + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: float(loss_obj.call(1.0, na, mp, lb)[0]), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-pref MSE meets the grad-accum invariant.""" + f_A, f_A_hat = _rnd(NA, 3), _rnd(NA, 3) + pf_A = np.abs(_rnd(NA, 3)) + 0.1 # positive pref + f_B, f_B_hat = _rnd(NB, 3), _rnd(NB, 3) + pf_B = np.abs(_rnd(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mse"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_mae_grad_accum(self): + """Atom-pref MAE meets the grad-accum invariant.""" + f_A, f_A_hat = _rnd(NA, 3), _rnd(NA, 3) + pf_A = np.abs(_rnd(NA, 3)) + 0.1 + f_B, f_B_hat = _rnd(NB, 3), _rnd(NB, 3) + pf_B = np.abs(_rnd(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mae"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-pref loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + pf = np.abs(_rnd(NP, 3)) + 0.1 + loss_obj = self._make_pf_only_loss("mse") + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["atom_pref"] = pf.reshape(1, NP * 3) + l_mask["find_force"] = 1.0 + l_mask["find_atom_pref"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["atom_pref"] = pf.reshape(1, NP * 3) + l_nm["find_force"] = 1.0 + l_nm["find_atom_pref"] = 1.0 + + loss_m = float(loss_obj.call(1.0, NP, p_mask, l_mask)[0]) + loss_nm = float(loss_obj.call(1.0, NP, p_nm, l_nm)[0]) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +# --------------------------------------------------------------------------- +# Task 4: PropertyLoss -- extensive (not intensive) property +# --------------------------------------------------------------------------- + +PROP_TASK_DIM = 2 +PROP_VAR = "test_prop" + + +class TestPropertyLossExtensiveGradAccum: + """Idiom 2 (per-frame real-natoms normalization) for extensive PropertyLoss. + + The _loss_fn wrapper divides by nf so that the per-frame average matches + the separate-frame reference (property loss uses sum, not mean, over frames). + """ + + def _make_loss(self, loss_func="mse"): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func=loss_func, + intensive=False, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + """Return per-frame-averaged loss (raw loss / nf).""" + nf = model_pred[PROP_VAR].shape[0] + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) / nf + + def _run_invariant(self, loss_obj, p_A, l_A, p_B, l_B): + def make_A(): + return ( + {PROP_VAR: p_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {PROP_VAR: l_A}, + NA, + ) + + def make_B(): + return ( + {PROP_VAR: p_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {PROP_VAR: l_B}, + NB, + ) + + def make_padded(): + return ( + { + PROP_VAR: np.concatenate([p_A, p_B], axis=0), # [2, task_dim] + "mask": _MASK_PAD, + }, + {PROP_VAR: np.concatenate([l_A, l_B], axis=0)}, + NP, + ) + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """MSE extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mse"), p_A, l_A, p_B, l_B) + + def test_mae_grad_accum(self): + """MAE extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mae"), p_A, l_A, p_B, l_B) + + def test_smooth_mae_grad_accum(self): + """smooth_mae extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("smooth_mae"), p_A, l_A, p_B, l_B) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same extensive-property loss as no mask.""" + p = _rnd(2, PROP_TASK_DIM) + l = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + # PropertyLoss.call rebinds local `label` (no in-place mutation) so l is safe to reuse. + with_mask = {PROP_VAR: p, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {PROP_VAR: p} + loss_m, _ = loss_obj.call(1.0, NB, with_mask, {PROP_VAR: l}) + loss_nm, _ = loss_obj.call(1.0, NB, without_mask, {PROP_VAR: l}) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + def test_torch_backend_matches_numpy(self): + """The array-API path must work with torch tensors (pt_expt runs this). + + A padded batch is fed once as numpy and once as torch; the extensive + per-frame normalization uses xp.sum(mask, axis=-1), which raises under + the array_api_compat torch namespace if axis is passed positionally. + """ + import pytest + + torch = pytest.importorskip("torch") + # importing deepmd.pt sets a cuda:9999999 default-device trap; pin cpu. + dev = "cpu" + p = _rnd(2, PROP_TASK_DIM) + lab = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + np_pred = {PROP_VAR: p, "mask": _MASK_PAD} + pt_pred = { + PROP_VAR: torch.tensor(p, device=dev), + "mask": torch.tensor(_MASK_PAD, device=dev), + } + loss_np, _ = loss_obj.call(1.0, NP, np_pred, {PROP_VAR: lab}) + loss_pt, _ = loss_obj.call( + 1.0, NP, pt_pred, {PROP_VAR: torch.tensor(lab, device=dev)} + ) + assert np.isclose(float(loss_np), float(loss_pt)), ( + f"torch path must match numpy: {float(loss_np)} vs {float(loss_pt)}" + ) + + +class TestPropertyLossIntensiveUnaffectedByMask: + """Intensive property loss must be unchanged whether or not mask is present.""" + + def _make_loss(self): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func="mse", + intensive=True, + ) + + def test_intensive_ignores_mask(self): + """Intensive property: padded-batch loss == unmasked-batch loss.""" + p = _rnd(2, PROP_TASK_DIM) + l = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss() + with_mask = {PROP_VAR: p, "mask": _MASK_PAD} + without_mask = {PROP_VAR: p} + loss_m, _ = loss_obj.call(1.0, NP, with_mask, {PROP_VAR: l}) + loss_nm, _ = loss_obj.call(1.0, NP, without_mask, {PROP_VAR: l}) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"intensive property must ignore mask: {float(loss_m)} vs {float(loss_nm)}" + ) + + +# --------------------------------------------------------------------------- +# Task 5: EnergySpinLoss -- energy (has_e), force_real (has_fr), virial (has_v) +# Leave force_mag / mask_mag (has_fm) COMPLETELY UNTOUCHED. +# --------------------------------------------------------------------------- + +# Spin-specific test constants: NM magnetic atoms per frame (same count in +# both frames so that the pt fancy-index .view(nframes,-1,3) stays valid). +_NM = 2 + +_MASK_MAG_A = np.zeros((1, NA, 1), dtype=bool) # [1, NA, 1] +_MASK_MAG_A[0, :_NM, 0] = True # first NM atoms of frame A are magnetic + +_MASK_MAG_B = np.zeros((1, NB, 1), dtype=bool) # [1, NB, 1] +_MASK_MAG_B[0, :_NM, 0] = True # first NM atoms of frame B are magnetic + +# Padded batch (nf=2, nloc=NP): frame A padding slots are not magnetic. +_MASK_MAG_PAD_SPIN = np.zeros((2, NP, 1), dtype=bool) # [2, NP, 1] +_MASK_MAG_PAD_SPIN[0, :_NM, 0] = True # frame A: first NM real atoms are magnetic +_MASK_MAG_PAD_SPIN[1, :_NM, 0] = True # frame B: first NM real atoms are magnetic + +_MASK_PAD_SPIN = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 +) # [2, NP] + + +def _full_spin_dicts( + nf, nloc, energy_pred, energy_label, mask_mag, mask=None, **overrides +): + """Build complete model_pred and label_dict for EnergySpinLoss.call. + + EnergySpinLoss.call accesses "energy" unconditionally (for xp namespace), + so it must always be present. All other keys are guarded by has_* flags. + """ + model_pred = { + "energy": energy_pred, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "force_mag": np.zeros((nf, nloc, 3), dtype=np.float64), + "mask_mag": mask_mag, # [nf, nloc, 1] + "virial": np.zeros((nf, 9), dtype=np.float64), + } + label_dict = { + "energy": energy_label, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "force_mag": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "find_energy": 1.0, + "find_force": 0.0, + "find_force_mag": 0.0, + "find_virial": 0.0, + } + if mask is not None: + model_pred["mask"] = mask + model_pred.update({k: v for k, v in overrides.items() if k in model_pred}) + label_dict.update({k: v for k, v in overrides.items() if k in label_dict}) + return model_pred, label_dict + + +class TestDPModelEnerSpinLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + intensive_ener_virial=intensive, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + e_A, + e_A_hat, + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + e_B, + e_B_hat, + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + return p, l, NB + + def make_padded(): + e_pad = np.concatenate([e_A, e_B], axis=0) # [2, 1] + e_hat_pad = np.concatenate([e_A_hat, e_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, NP, e_pad, e_hat_pad, _MASK_MAG_PAD_SPIN, mask=_MASK_PAD_SPIN + ) + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + + def test_mae_grad_accum(self): + """Spin energy MAE meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, e_A, e_A_hat, e_B, e_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _rnd(1, 1) + e_hat = _rnd(1, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, NP, e, e_hat, _MASK_MAG_B, mask=np.ones((1, NP), dtype=np.float64) + ) + p_nm, l_nm = _full_spin_dicts(1, NP, e, e_hat, _MASK_MAG_B) + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossForceRealGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for force_real in EnergySpinLoss. + + Covers: mse; plus non-mixed no-op. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + return p, l, NB + + def make_padded(): + f_A_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_pad[:NA] = f_A + f_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_hat_pad[:NA] = f_A_hat + f_pad = np.stack([f_A_pad, f_B], axis=0) # [2, NP, 3] + f_hat_pad = np.stack([f_A_hat_pad, f_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force_real MSE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force_real MAE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, f_A, f_A_hat, f_B, f_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force_real loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["find_force"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["find_force"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for virial in EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + intensive_ener_virial=intensive, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["virial"] = v_A[None] # [1, 9] + l["virial"] = v_A_hat[None] + l["find_virial"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["virial"] = v_B[None] + l["virial"] = v_B_hat[None] + l["find_virial"] = 1.0 + return p, l, NB + + def make_padded(): + v_pad = np.stack([v_A, v_B], axis=0) # [2, 9] + v_hat_pad = np.stack([v_A_hat, v_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["virial"] = v_pad + l["virial"] = v_hat_pad + l["find_virial"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss(intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + + def test_mae_grad_accum(self): + """Spin virial MAE meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, v_A, v_A_hat, v_B, v_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _rnd(9) + v_hat = _rnd(9) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["virial"] = v[None] + l_mask["virial"] = v_hat[None] + l_mask["find_virial"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["virial"] = v[None] + l_nm["virial"] = v_hat[None] + l_nm["find_virial"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossForceMagUnchanged: + """Guard: the padding mask must NOT affect the force_mag / mask_mag term. + + The force_mag path uses mask_mag (spin virtual-atom mask), which is a + completely separate concept from the padding mask model_dict["mask"]. + After the Task-5 implementation, presenting a padding mask must leave + the force_mag loss bit-identical. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_padding_mask_does_not_affect_force_mag(self): + """force_mag loss is bit-identical with and without padding mask.""" + fm = _rnd(2, NP, 3) + fm_hat = _rnd(2, NP, 3) + loss_obj = self._make_loss() + + def _run(with_mask): + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm, + "mask_mag": _MASK_MAG_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + if with_mask: + pred["mask"] = _MASK_PAD_SPIN + loss, _ = loss_obj.call(1.0, NP, pred, lbl) + return float(loss) + + loss_with = _run(True) + loss_without = _run(False) + assert np.isclose(loss_with, loss_without), ( + f"force_mag loss must be unchanged by padding mask: " + f"{loss_with} vs {loss_without}" + ) + + +# --------------------------------------------------------------------------- +# Part A: EnergySpinLoss atom_ener (has_ae) grad-accum invariant +# --------------------------------------------------------------------------- + + +class TestDPModelEnerSpinLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for atom_ener in EnergySpinLoss. + + RED before the Part-A has_ae mask fix; GREEN after. + """ + + def _make_loss(self, loss_func="mse"): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["atom_energy"] = ae_A[None] # [1, NA, 1] + l["atom_ener"] = ae_A_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["atom_energy"] = ae_B[None] + l["atom_ener"] = ae_B_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NB + + def make_padded(): + ae_pad = _padded_atom(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom(ae_A_hat, ae_B_hat, 1) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["atom_energy"] = ae_pad + l["atom_ener"] = ae_hat_pad + l["find_atom_ener"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Spin atom_ener MSE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Spin atom_ener MAE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom_ener spin loss as no mask.""" + ae = _rnd(NP, 1) + ae_hat = _rnd(NP, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["atom_energy"] = ae[None] + l_mask["atom_ener"] = ae_hat[None] + l_mask["find_atom_ener"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["atom_energy"] = ae[None] + l_nm["atom_ener"] = ae_hat[None] + l_nm["find_atom_ener"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +# --------------------------------------------------------------------------- +# gen_force (has_gf) grad-accum invariant +# Ghost-atom forces are masked before projection, so the invariant is met. +# --------------------------------------------------------------------------- + +_NGEN = 2 # number of generalized coordinates for tests + + +class TestDPModelEnergyLossGenForceGradAccum: + """gen_force (has_gf) excludes ghost atoms via force masking before projection. + + The projected gen_force = sum_i(drdq_ij * masked_f_i) is frame-decomposable, + so mean(square(diff)) over [nf, ngen] already satisfies the invariant. + Expected GREEN immediately (no fix needed for gen_force). + """ + + def _make_loss(self): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=_NGEN, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B): + def make_A(): + p, l = _full_ener_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + l["drdq"] = drdq_A[None] # [1, NA*3, NGEN] + l["find_drdq"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + l["drdq"] = drdq_B[None] # [1, NB*3, NGEN] + l["find_drdq"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + # drdq ghost-atom slots are zero (ghost forces also zero, so no contribution) + drdq_A_pad = np.zeros((NP * 3, _NGEN), dtype=np.float64) + drdq_A_pad[: NA * 3] = drdq_A + drdq_pad = np.stack([drdq_A_pad, drdq_B], axis=0) # [2, NP*3, NGEN] + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + l["drdq"] = drdq_pad + l["find_drdq"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """gen_force MSE meets the grad-accum invariant (GREEN: already correct).""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + drdq_A = _rnd(NA * 3, _NGEN) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + drdq_B = _rnd(NB * 3, _NGEN) + self._run_invariant( + self._make_loss(), f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B + ) + + +# --------------------------------------------------------------------------- +# force_mag MSE grad-accum invariant (NM equal across frames, ghost atoms non-magnetic) +# +# NOTE: force_mag MAE (dpmodel) uses a plain mean over frames/atoms/xyz (like +# force_mag MSE and force_real MAE), so it meets the grad-accum invariant. +# Ghost atoms are correctly excluded (mask_mag=0 there). +# --------------------------------------------------------------------------- + + +class TestDPModelEnerSpinLossForceMagMSEGradAccum: + """force_mag MSE meets the grad-accum invariant when ghost atoms are non-magnetic. + + With equal magnetic-atom counts (NM) across frames and ghost atoms having + mask_mag=0, the global normalization n_valid = NM*nf correctly decomposes + per-frame, giving loss_pad == 0.5*(loss_A + loss_B). + Expected GREEN immediately (no fix needed for MSE force_mag). + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def test_mse_grad_accum(self): + """force_mag MSE (NM equal per frame, ghost atoms non-magnetic) meets invariant.""" + fm_A = _rnd(NA, 3) + fm_A_hat = _rnd(NA, 3) + fm_B = _rnd(NB, 3) + fm_B_hat = _rnd(NB, 3) + + def make_A(): + # Frame A: NA=3 real atoms, first NM=2 are magnetic + fm_A_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_full[:_NM] = fm_A[:_NM] + fm_A_hat_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_hat_full[:_NM] = fm_A_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_A_full[None], # [1, NA, 3] + "mask_mag": _MASK_MAG_A, + } + lbl = { + "force_mag": fm_A_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NA + + def make_B(): + fm_B_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_full[:_NM] = fm_B[:_NM] + fm_B_hat_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_hat_full[:_NM] = fm_B_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_B_full[None], # [1, NB, 3] + "mask_mag": _MASK_MAG_B, + } + lbl = { + "force_mag": fm_B_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NB + + def make_padded(): + # Pad frame A force_mag to NP width; ghost slots are zero (non-magnetic) + fm_A_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_pad[:_NM] = fm_A[:_NM] + fm_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_hat_pad[:_NM] = fm_A_hat[:_NM] + fm_B_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_pad[:_NM] = fm_B[:_NM] + fm_B_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_hat_pad[:_NM] = fm_B_hat[:_NM] + fm_pad = np.stack([fm_A_pad, fm_B_pad], axis=0) # [2, NP, 3] + fm_hat_pad = np.stack([fm_A_hat_pad, fm_B_hat_pad], axis=0) + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm_pad, + "mask_mag": _MASK_MAG_PAD_SPIN, # ghost atoms have mask_mag=False + "mask": _MASK_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat_pad, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) + + +class TestDPModelEnerSpinLossForceMagMAEGradAccum: + """force_mag MAE (dpmodel) meets the grad-accum invariant. + + ``force_mag`` MAE now reduces with a plain mean over frames, magnetic atoms + and xyz (same reduction as force_mag MSE and force_real MAE), so a padded + ``[A+B]`` batch (nf=2) equals the mean of the two single-frame losses and the + loss is batch-size independent. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def test_mae_grad_accum(self): + """force_mag MAE meets the frame-average grad-accum invariant.""" + fm_A = _rnd(NA, 3) + fm_A_hat = _rnd(NA, 3) + fm_B = _rnd(NB, 3) + fm_B_hat = _rnd(NB, 3) + + def make_A(): + fm_A_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_full[:_NM] = fm_A[:_NM] + fm_A_hat_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_hat_full[:_NM] = fm_A_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_A_full[None], # [1, NA, 3] + "mask_mag": _MASK_MAG_A, + } + lbl = { + "force_mag": fm_A_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NA + + def make_B(): + fm_B_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_full[:_NM] = fm_B[:_NM] + fm_B_hat_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_hat_full[:_NM] = fm_B_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_B_full[None], # [1, NB, 3] + "mask_mag": _MASK_MAG_B, + } + lbl = { + "force_mag": fm_B_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NB + + def make_padded(): + fm_A_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_pad[:_NM] = fm_A[:_NM] + fm_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_hat_pad[:_NM] = fm_A_hat[:_NM] + fm_B_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_pad[:_NM] = fm_B[:_NM] + fm_B_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_hat_pad[:_NM] = fm_B_hat[:_NM] + fm_pad = np.stack([fm_A_pad, fm_B_pad], axis=0) # [2, NP, 3] + fm_hat_pad = np.stack([fm_A_hat_pad, fm_B_hat_pad], axis=0) + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm_pad, + "mask_mag": _MASK_MAG_PAD_SPIN, # ghost atoms have mask_mag=False + "mask": _MASK_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat_pad, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py new file mode 100644 index 0000000000..094bcfeadb --- /dev/null +++ b/source/tests/pt/test_loss_padding.py @@ -0,0 +1,2174 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for mixed_type loss padding-mask support in the pt backend. + +Task 1: verify that TaskLoss._inject_atom_mask correctly recovers the per-atom +mask from atype (ghost atoms have atype < 0) so that later tasks can exclude +them from loss reductions. + +Harness +------- +assert_grad_accum_invariant -- reusable by Tasks 2-5 to check the + grad-accumulation invariant: loss on a padded multi-frame batch must equal + mean_over_frames(per_frame_loss). + +Scope / follow-ups (mixed_type padding fix, PR #5738) +---------------------------------------------------- +- The TF backend loss is not covered here and still has the mixed_type + dilution behavior; tracked in deepmodeling/deepmd-kit#5760. +- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; + tracked in deepmodeling/deepmd-kit#5761. +- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean + reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real + MAE; the grad-accum invariant is asserted by the test below. +""" + +import numpy as np +import torch + +from deepmd.pt.loss.dos import ( + DOSLoss, +) +from deepmd.pt.loss.ener import ( + EnergyStdLoss, +) +from deepmd.pt.loss.ener_spin import EnergySpinLoss as EnergySpinLossPT +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.loss.property import ( + PropertyLoss, +) +from deepmd.pt.loss.tensor import ( + TensorLoss, +) +from deepmd.pt.utils import ( + env, +) + +# --------------------------------------------------------------------------- +# Constants used by the multi-frame test harness (Tasks 2-5) +# --------------------------------------------------------------------------- + +NA = 3 # real atoms in frame A +NB = 5 # real atoms in frame B (== NP, so frame B is fully real) +NP = 5 # padded width (nloc) + + +# --------------------------------------------------------------------------- +# Reusable harness (used by Tasks 2-5; imported from here) +# --------------------------------------------------------------------------- + + +def assert_grad_accum_invariant( + loss_fn, + make_batch_A, + make_batch_B, + make_padded_batch, + rtol: float = 1e-5, + atol: float = 1e-6, +) -> None: + """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + + The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms + padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as + processing each frame separately and averaging. + + Parameters + ---------- + loss_fn : callable + Signature ``(model_pred, label, natoms) -> scalar torch.Tensor``. + make_batch_A : callable + Returns ``(model_pred, label, natoms)`` for frame A alone (1 frame, NA atoms). + make_batch_B : callable + Returns ``(model_pred, label, natoms)`` for frame B alone (1 frame, NB atoms). + make_padded_batch : callable + Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch + (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + rtol : float + Relative tolerance for ``torch.isclose``. + atol : float + Absolute tolerance for ``torch.isclose``. + """ + pred_A, label_A, natoms_A = make_batch_A() + pred_B, label_B, natoms_B = make_batch_B() + pred_pad, label_pad, natoms_pad = make_padded_batch() + + loss_A = loss_fn(pred_A, label_A, natoms_A) + loss_B = loss_fn(pred_B, label_B, natoms_B) + ref = 0.5 * (loss_A + loss_B) + + loss_pad = loss_fn(pred_pad, label_pad, natoms_pad) + + assert torch.isclose(loss_pad, ref, rtol=rtol, atol=atol), ( + f"Grad-accum invariant violated: padded_loss={loss_pad.item():.8f}, " + f"ref={ref.item():.8f}, diff={abs(loss_pad.item() - ref.item()):.2e}" + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +RNG = np.random.default_rng(42) +NUMB_DOS = 4 +TENSOR_SIZE = 3 + + +def _rnd_t(*shape): + """Create a random CPU tensor (device must be explicit to avoid default-device trap).""" + return torch.tensor(RNG.standard_normal(shape), dtype=torch.float64, device="cpu") + + +class _MockModel: + """Callable that ignores inputs and returns a fixed model_pred dict. + + Used in pt loss tests to bypass the actual model forward pass while still + exercising the loss computation code inside DOSLoss.forward / TensorLoss.forward. + The mask is pre-populated in ``pred`` so ``_inject_atom_mask`` leaves it alone. + """ + + def __init__(self, pred: dict): + self._pred = pred + + def __call__(self, **kwargs): + return dict(self._pred) # shallow copy; _inject_atom_mask may mutate it + + +# --------------------------------------------------------------------------- +# Task 2: DOSLoss -- atomic (ados / acdf) and global (dos / cdf) +# --------------------------------------------------------------------------- + + +class TestPTDOSLossAtomicGradAccum: + """Per-frame masked mean (idiom 1) for atomic dos / acdf terms. + + _loss_fn calls the ACTUAL pt DOSLoss.forward() via a mock model so that + RED/GREEN transitions directly reflect changes to deepmd/pt/loss/dos.py. + """ + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=0.0, + limit_pref_dos=0.0, + start_pref_cdf=0.0, + limit_pref_cdf=0.0, + start_pref_ados=1.0, + limit_pref_ados=1.0, + start_pref_acdf=1.0, + limit_pref_acdf=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, # no atype; mask already in model_pred → not re-injected + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_ados_grad_accum_invariant(self): + """Atomic dos per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd_t(NA, NUMB_DOS) + label_A = _rnd_t(NA, NUMB_DOS) + pred_B = _rnd_t(NB, NUMB_DOS) + label_B = _rnd_t(NB, NUMB_DOS) + + def make_A(): + return ( + { + "atom_dos": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"atom_dos": label_A, "find_atom_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "atom_dos": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"atom_dos": label_B, "find_atom_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = torch.zeros(NP, NUMB_DOS, dtype=torch.float64, device="cpu") + pred_A_pad[:NA] = pred_A + label_A_pad = torch.zeros(NP, NUMB_DOS, dtype=torch.float64, device="cpu") + label_A_pad[:NA] = label_A + mask_A = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA)], dtype=torch.float64, device="cpu" + ) + mask_B = torch.ones(1, NB, dtype=torch.float64, device="cpu") + atom_dos_pad = torch.cat([pred_A_pad, pred_B], dim=0) + atom_dos_label = torch.cat([label_A_pad, label_B], dim=0) + mask_pad = torch.cat([mask_A, mask_B], dim=0) + return ( + {"atom_dos": atom_dos_pad, "mask": mask_pad}, + {"atom_dos": atom_dos_label, "find_atom_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(NB, NUMB_DOS) + label = _rnd_t(NB, NUMB_DOS) + with_mask = { + "atom_dos": pred, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"atom_dos": pred} + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTDOSLossGlobalGradAccum: + """Plain mean (idiom 3) for global dos / cdf terms. + + _loss_fn calls the ACTUAL pt DOSLoss.forward() via a mock model. + """ + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_dos_grad_accum_invariant(self): + """Global dos plain mean meets the grad-accum invariant.""" + pred_A = _rnd_t(1, NUMB_DOS) + label_A = _rnd_t(1, NUMB_DOS) + pred_B = _rnd_t(1, NUMB_DOS) + label_B = _rnd_t(1, NUMB_DOS) + + def make_A(): + return ( + { + "dos": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"dos": label_A, "find_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "dos": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"dos": label_B, "find_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = torch.cat([pred_A, pred_B], dim=0) + label_pad = torch.cat([label_A, label_B], dim=0) + mask_pad = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", + ) + return ( + {"dos": pred_pad, "mask": mask_pad}, + {"dos": label_pad, "find_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(2, NUMB_DOS) + label = _rnd_t(2, NUMB_DOS) + with_mask = { + "dos": pred, + "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"dos": pred} + label_dict = {"dos": label, "find_dos": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# Task 2: TensorLoss -- local and global tensor +# --------------------------------------------------------------------------- + + +class TestPTTensorLossLocalGradAccum: + """Per-frame masked mean (idiom 1) for local tensor term. + + _loss_fn calls the ACTUAL pt TensorLoss.forward() via a mock model. + """ + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=1.0, + pref=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_local_grad_accum_invariant(self): + """Local tensor per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd_t(NA, TENSOR_SIZE) + label_A = _rnd_t(NA, TENSOR_SIZE) + pred_B = _rnd_t(NB, TENSOR_SIZE) + label_B = _rnd_t(NB, TENSOR_SIZE) + + def make_A(): + return ( + { + "dipole": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"atom_dipole": label_A, "find_atom_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "dipole": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"atom_dipole": label_B, "find_atom_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = torch.zeros(NP, TENSOR_SIZE, dtype=torch.float64, device="cpu") + pred_A_pad[:NA] = pred_A + label_A_pad = torch.zeros( + NP, TENSOR_SIZE, dtype=torch.float64, device="cpu" + ) + label_A_pad[:NA] = label_A + mask_A = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA)], dtype=torch.float64, device="cpu" + ) + mask_B = torch.ones(1, NB, dtype=torch.float64, device="cpu") + dipole_pad = torch.cat([pred_A_pad, pred_B], dim=0) + label_pad = torch.cat([label_A_pad, label_B], dim=0) + mask_pad = torch.cat([mask_A, mask_B], dim=0) + return ( + {"dipole": dipole_pad, "mask": mask_pad}, + {"atom_dipole": label_pad, "find_atom_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(NB, TENSOR_SIZE) + label = _rnd_t(NB, TENSOR_SIZE) + with_mask = { + "dipole": pred, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"dipole": pred} + label_dict = {"atom_dipole": label, "find_atom_dipole": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTTensorLossGlobalGradAccum: + """Plain mean (idiom 3) for global tensor term. + + _loss_fn calls the ACTUAL pt TensorLoss.forward() via a mock model. + """ + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_global_grad_accum_invariant(self): + """Global tensor plain mean meets the grad-accum invariant.""" + pred_A = _rnd_t(1, TENSOR_SIZE) + label_A = _rnd_t(1, TENSOR_SIZE) + pred_B = _rnd_t(1, TENSOR_SIZE) + label_B = _rnd_t(1, TENSOR_SIZE) + + def make_A(): + return ( + { + "global_dipole": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"dipole": label_A, "find_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "global_dipole": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"dipole": label_B, "find_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = torch.cat([pred_A, pred_B], dim=0) + label_pad = torch.cat([label_A, label_B], dim=0) + mask_pad = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", + ) + return ( + {"global_dipole": pred_pad, "mask": mask_pad}, + {"dipole": label_pad, "find_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(2, TENSOR_SIZE) + label = _rnd_t(2, TENSOR_SIZE) + with_mask = { + "global_dipole": pred, + "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"global_dipole": pred} + label_dict = {"dipole": label, "find_dipole": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# Task 1: unit tests for TaskLoss._inject_atom_mask +# --------------------------------------------------------------------------- + + +class TestInjectAtomMask: + """Unit tests for TaskLoss._inject_atom_mask. + + NOTE: source/tests/pt/__init__.py sets torch.set_default_device("cuda:9999999") + to catch tests that forget to specify device. All tensor creations here + must pass device="cpu" explicitly. + """ + + def test_injects_mask_from_atype(self) -> None: + """When model_pred has no mask, recover it from atype (<0 = ghost).""" + # nf=2, nloc=5; frame 0 has ghost atoms at positions 3 and 4 + atype = torch.tensor([[0, 1, 0, -1, -1], [0, 1, 0, 1, 0]], device="cpu") + energy = torch.zeros(2, 1, dtype=torch.float64, device="cpu") + input_dict = {"atype": atype} + model_pred = {"energy": energy} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" in result, "mask must be injected" + assert result["mask"].shape == (2, 5), f"wrong shape {result['mask'].shape}" + expected = torch.tensor( + [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]], + dtype=torch.float64, + device="cpu", + ) + assert torch.allclose(result["mask"], expected), ( + f"mask values wrong: {result['mask']}" + ) + assert result["mask"].dtype == energy.dtype, ( + "mask dtype must match reference tensor dtype" + ) + + def test_mask_dtype_follows_ref(self) -> None: + """Mask dtype follows the reference tensor (energy or coord or atype).""" + atype = torch.tensor([[0, -1]], dtype=torch.long, device="cpu") + model_pred = {"energy": torch.zeros(1, 1, dtype=torch.float32, device="cpu")} + input_dict = {"atype": atype} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert result["mask"].dtype == torch.float32 + + def test_fallback_to_coord_when_no_energy(self) -> None: + """When energy absent, ref falls back to coord.""" + atype = torch.tensor([[0, -1]], dtype=torch.long, device="cpu") + coord = torch.zeros(1, 2, 3, dtype=torch.float64, device="cpu") + model_pred = {} + input_dict = {"atype": atype, "coord": coord} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" in result + assert result["mask"].dtype == torch.float64 + + def test_mask_not_overwritten_if_present(self) -> None: + """When mask is already in model_pred, leave it unchanged.""" + existing = torch.ones(2, 5, device="cpu") * 0.5 + model_pred = { + "mask": existing, + "energy": torch.zeros(2, 1, device="cpu"), + } + input_dict = {"atype": torch.zeros(2, 5, dtype=torch.long, device="cpu")} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert result["mask"] is existing, "existing mask must not be overwritten" + + def test_no_atype_no_injection(self) -> None: + """When input_dict has no atype, mask must not be injected.""" + model_pred = {"energy": torch.zeros(2, 1, device="cpu")} + input_dict = {} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" not in result, "mask must not appear when atype is absent" + + def test_all_real_atoms_gives_all_ones(self) -> None: + """Non-mixed batch (all atype >= 0) yields all-ones mask.""" + atype = torch.tensor([[0, 1, 2], [1, 2, 0]], device="cpu") + energy = torch.zeros(2, 1, dtype=torch.float32, device="cpu") + model_pred = {"energy": energy} + input_dict = {"atype": atype} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert torch.all(result["mask"] == 1.0), ( + "all-real atoms must give all-ones mask" + ) + + +# --------------------------------------------------------------------------- +# Task 3: EnergyStdLoss -- energy, force, virial, atom_ener, atom_pref +# --------------------------------------------------------------------------- + +# Re-use same constants from Task 2 harness (NA=3, NB=5, NP=5). + +_MASK_PAD_PT = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", +) # [2, NP] + + +def _t(*shape, val=None): + """Random float64 CPU tensor.""" + if val is not None: + return torch.full(shape, val, dtype=torch.float64, device="cpu") + return torch.tensor(RNG.standard_normal(shape), dtype=torch.float64, device="cpu") + + +def _padded_force_t(f_A, f_B): + """Stack force tensors for 2-frame padded batch.""" + pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + pad[:NA] = f_A + return torch.stack([pad, f_B], dim=0) # [2, NP, 3] + + +def _padded_atom_t(a_A, a_B, ncomp): + """Pad arr_A from [NA, ncomp] to [NP, ncomp], stack with arr_B.""" + pad = torch.zeros(NP, ncomp, dtype=torch.float64, device="cpu") + pad[:NA] = a_A + return torch.stack([pad, a_B], dim=0) # [2, NP, ncomp] + + +class _EnerLossMockModel: + """Callable returning a fixed model_pred dict (mask pre-populated).""" + + def __init__(self, pred: dict): + self._pred = pred + + def __call__(self, **kwargs): + return dict(self._pred) + + +def _ener_loss_fn(loss_obj, model_pred, label, natoms): + """Call EnergyStdLoss.forward via mock model; return scalar loss tensor.""" + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_EnerLossMockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + +class TestPTEnergyLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in EnergyStdLoss. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + mp = { + "energy": e_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_A_hat, "find_energy": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "energy": e_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_B_hat, "find_energy": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "energy": torch.cat([e_A, e_B], dim=0), + "mask": _MASK_PAD_PT, + } + lb = { + "energy": torch.cat([e_A_hat, e_B_hat], dim=0), + "find_energy": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", intensive=True), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mae_grad_accum(self): + """Energy MAE meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant(self._make_loss("mae"), e_A, e_A_hat, e_B, e_B_hat) + + def test_huber_grad_accum(self): + """Energy Huber meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", use_huber=True), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _t(1, 1) + e_hat = _t(1, 1) + loss_obj = self._make_loss() + + mp_mask = { + "energy": e, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"energy": e} + lb = {"energy": e_hat, "find_energy": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossForceGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for the force term. + + Covers: mse, mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func=loss_func, + use_huber=use_huber, + f_use_norm=f_use_norm, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_A_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_B_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "force": _padded_force_t(f_A, f_B), # [2, NP, 3] + "mask": _MASK_PAD_PT, + } + lb = { + "force": _padded_force_t(f_A_hat, f_B_hat), + "find_force": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force MSE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss("mse"), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force MAE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss("mae"), f_A, f_A_hat, f_B, f_B_hat) + + def test_huber_grad_accum(self): + """Force Huber meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant( + self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + loss_obj = self._make_loss() + + mp_mask = { + "force": f.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0)} + lb = {"force": f_hat.unsqueeze(0), "find_force": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for the virial term. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + mp = { + "virial": v_A.unsqueeze(0), # [1, 9] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_A_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "virial": v_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_B_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "virial": torch.stack([v_A, v_B], dim=0), # [2, 9] + "mask": _MASK_PAD_PT, + } + lb = { + "virial": torch.stack([v_A_hat, v_B_hat], dim=0), + "find_virial": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", intensive=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mae_grad_accum(self): + """Virial MAE meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant(self._make_loss("mae"), v_A, v_A_hat, v_B, v_B_hat) + + def test_huber_grad_accum(self): + """Virial Huber meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", use_huber=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _t(9) + v_hat = _t(9) + loss_obj = self._make_loss() + + mp_mask = { + "virial": v.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"virial": v.unsqueeze(0)} + lb = {"virial": v_hat.unsqueeze(0), "find_virial": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for the atom_ener term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + mp = { + "atom_energy": ae_A.unsqueeze(0), # [1, NA, 1] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_A_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "atom_energy": ae_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_B_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NB + + def make_padded(): + ae_pad = _padded_atom_t(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom_t(ae_A_hat, ae_B_hat, 1) + mp = {"atom_energy": ae_pad, "mask": _MASK_PAD_PT} + lb = {"atom_ener": ae_hat_pad, "find_atom_ener": 1.0} + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-energy MSE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Atom-energy MAE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_huber_grad_accum(self): + """Atom-energy Huber (use_huber=True) meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + loss_obj = EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + use_huber=True, + ) + self._run_invariant(loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-energy loss as no mask.""" + ae = _t(NP, 1) + ae_hat = _t(NP, 1) + loss_obj = self._make_loss() + + mp_mask = { + "atom_energy": ae.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"atom_energy": ae.unsqueeze(0)} + lb = {"atom_ener": ae_hat.unsqueeze(0), "find_atom_ener": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossAtomPrefGradAccum: + """Idiom 1 with pref weight (ncomp=3) for the atom_pref term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_pf_only_loss(self, loss_func="mse"): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_A_hat.unsqueeze(0), + "atom_pref": pf_A.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_B_hat.unsqueeze(0), + "atom_pref": pf_B.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NB + + def make_padded(): + f_pad = _padded_force_t(f_A, f_B) + f_hat_pad = _padded_force_t(f_A_hat, f_B_hat) + pf_pad = _padded_atom_t(pf_A, pf_B, 3) # [2, NP, 3] + mp = {"force": f_pad, "mask": _MASK_PAD_PT} + lb = { + "force": f_hat_pad, + "atom_pref": pf_pad, + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-pref MSE meets the grad-accum invariant.""" + f_A, f_A_hat = _t(NA, 3), _t(NA, 3) + pf_A = torch.abs(_t(NA, 3)) + 0.1 + f_B, f_B_hat = _t(NB, 3), _t(NB, 3) + pf_B = torch.abs(_t(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mse"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_mae_grad_accum(self): + """Atom-pref MAE meets the grad-accum invariant.""" + f_A, f_A_hat = _t(NA, 3), _t(NA, 3) + pf_A = torch.abs(_t(NA, 3)) + 0.1 + f_B, f_B_hat = _t(NB, 3), _t(NB, 3) + pf_B = torch.abs(_t(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mae"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-pref loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + pf = torch.abs(_t(NP, 3)) + 0.1 + loss_obj = self._make_pf_only_loss("mse") + + mp_mask = { + "force": f.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0)} + lb = { + "force": f_hat.unsqueeze(0), + "atom_pref": pf.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# Task 4: PropertyLoss -- extensive (not intensive) property +# --------------------------------------------------------------------------- + +PROP_TASK_DIM = 2 +PROP_VAR = "test_prop" + + +# Unlike the other losses, PropertyLoss.forward builds out_std/out_bias (and the +# loss accumulator) on env.DEVICE and combines them with the *multi-dim* +# label/pred tensors (and divides pred/label by real_natoms from the mask). A +# CPU-input / env.DEVICE-buffer mix therefore raises on GPU CI (a 0-dim CPU +# scalar is auto-promoted, but a [nf, task_dim] tensor is not), so the property +# inputs and their masks must live on env.DEVICE. See #5738. +def _rnd_t_prop(*shape): + """Random float64 tensor on env.DEVICE (PropertyLoss holds device-resident buffers).""" + return torch.tensor( + RNG.standard_normal(shape), dtype=torch.float64, device=env.DEVICE + ) + + +_MASK_PAD_PROP = _MASK_PAD_PT.to(env.DEVICE) + + +class TestPTPropertyLossExtensiveGradAccum: + """Idiom 2 (per-frame real-natoms normalization) for extensive pt PropertyLoss. + + _loss_fn wraps the raw loss with /nf so the per-frame average matches the + separate-frame reference (pt PropertyLoss uses reduction='sum' over frames). + The mask is pre-populated in model_pred so _inject_atom_mask leaves it alone. + """ + + def _make_loss(self, loss_func="mse"): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func=loss_func, + intensive=False, + # Provide explicit out_std/out_bias to avoid accessing model.atomic_model. + out_std=[1.0] * PROP_TASK_DIM, + out_bias=[0.0] * PROP_TASK_DIM, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + """Return per-frame-averaged loss (raw loss / nf).""" + nf = model_pred[PROP_VAR].shape[0] + _, loss, _ = loss_obj.forward( + input_dict={}, # no atype; mask already in model_pred + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss / nf + + def _run_invariant(self, loss_obj, p_A, l_A, p_B, l_B): + def make_A(): + return ( + { + PROP_VAR: p_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device=env.DEVICE), + }, + {PROP_VAR: l_A.clone()}, + NA, + ) + + def make_B(): + return ( + { + PROP_VAR: p_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device=env.DEVICE), + }, + {PROP_VAR: l_B.clone()}, + NB, + ) + + def make_padded(): + return ( + { + PROP_VAR: torch.cat([p_A, p_B], dim=0), # [2, task_dim] + "mask": _MASK_PAD_PROP, + }, + {PROP_VAR: torch.cat([l_A, l_B], dim=0)}, + NP, + ) + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """MSE extensive property meets the grad-accum invariant.""" + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mse"), p_A, l_A, p_B, l_B) + + def test_mae_grad_accum(self): + """MAE extensive property meets the grad-accum invariant.""" + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mae"), p_A, l_A, p_B, l_B) + + def test_smooth_mae_grad_accum(self): + """smooth_mae extensive property meets the grad-accum invariant.""" + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("smooth_mae"), p_A, l_A, p_B, l_B) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same extensive-property loss as no mask.""" + p = _rnd_t_prop(2, PROP_TASK_DIM) + l = _rnd_t_prop(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + mp_mask = { + PROP_VAR: p, + "mask": torch.ones(2, NB, dtype=torch.float64, device=env.DEVICE), + } + mp_nm = {PROP_VAR: p} + # pt forward mutates label[var_name]; use separate dicts for each call. + lb_m = {PROP_VAR: l.clone()} + lb_nm = {PROP_VAR: l.clone()} + loss_m = self._loss_fn(loss_obj, mp_mask, lb_m, NB) + loss_nm = self._loss_fn(loss_obj, mp_nm, lb_nm, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTPropertyLossIntensiveUnaffectedByMask: + """Intensive property loss must be unchanged whether or not mask is present.""" + + def _make_loss(self): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func="mse", + intensive=True, + out_std=[1.0] * PROP_TASK_DIM, + out_bias=[0.0] * PROP_TASK_DIM, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_intensive_ignores_mask(self): + """Intensive property: masked batch == unmasked batch.""" + p = _rnd_t_prop(2, PROP_TASK_DIM) + l = _rnd_t_prop(2, PROP_TASK_DIM) + loss_obj = self._make_loss() + mp_mask = {PROP_VAR: p, "mask": _MASK_PAD_PROP} + mp_nm = {PROP_VAR: p} + # Use separate label dicts since pt forward mutates label[var_name]. + lb_m = {PROP_VAR: l.clone()} + lb_nm = {PROP_VAR: l.clone()} + loss_m = self._loss_fn(loss_obj, mp_mask, lb_m, NP) + loss_nm = self._loss_fn(loss_obj, mp_nm, lb_nm, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"intensive property must ignore mask: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# Task 5: EnergySpinLoss (pt) -- energy (has_e), force_real (has_fr), +# virial (has_v). Leave force_mag / mask_mag (has_fm) COMPLETELY UNTOUCHED. +# --------------------------------------------------------------------------- + +# Spin-specific test constants: NM magnetic atoms per frame (same count in +# both frames so that .view(nframes,-1,3) in the pt force_mag path is valid). +_NM_PT = 2 + +_MASK_MAG_A_PT = torch.zeros(1, NA, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_A_PT[0, :_NM_PT, 0] = True # first NM_PT atoms of frame A magnetic + +_MASK_MAG_B_PT = torch.zeros(1, NB, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_B_PT[0, :_NM_PT, 0] = True # first NM_PT atoms of frame B magnetic + +_MASK_MAG_PAD_SPIN_PT = torch.zeros(2, NP, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_PAD_SPIN_PT[0, :_NM_PT, 0] = True # frame A +_MASK_MAG_PAD_SPIN_PT[1, :_NM_PT, 0] = True # frame B + +_MASK_PAD_SPIN_PT = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", +) # [2, NP] + + +def _spin_loss_fn(loss_obj, model_pred, label, natoms): + """Call EnergySpinLoss.forward via mock model; return scalar loss tensor.""" + _, loss, _ = loss_obj.forward( + input_dict={}, # mask already in model_pred; no re-injection needed + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + +class TestPTEnerSpinLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in pt EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + intensive_ener_virial=intensive, + ) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + mp = { + "energy": e_A, + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_A_hat, "find_energy": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "energy": e_B, + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_B_hat, "find_energy": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "energy": torch.cat([e_A, e_B], dim=0), # [2, 1] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "energy": torch.cat([e_A_hat, e_B_hat], dim=0), + "find_energy": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss(intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + + def test_mae_grad_accum(self): + """Spin energy MAE meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, e_A, e_A_hat, e_B, e_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _t(1, 1) + e_hat = _t(1, 1) + loss_obj = self._make_loss() + mp_mask = { + "energy": e, + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"energy": e, "mask_mag": _MASK_MAG_B_PT} + lb = {"energy": e_hat, "find_energy": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossForceRealGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for force_real in pt EnergySpinLoss. + + Covers: mse; plus non-mixed no-op. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_A_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_B_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "force": _padded_force_t(f_A, f_B), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force": _padded_force_t(f_A_hat, f_B_hat), + "find_force": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force_real MSE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force_real MAE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, f_A, f_A_hat, f_B, f_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force_real loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + loss_obj = self._make_loss() + mp_mask = { + "force": f.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0), "mask_mag": _MASK_MAG_B_PT} + lb = {"force": f_hat.unsqueeze(0), "find_force": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for virial in pt EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + intensive_ener_virial=intensive, + ) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + mp = { + "virial": v_A.unsqueeze(0), # [1, 9] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_A_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "virial": v_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_B_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "virial": torch.stack([v_A, v_B], dim=0), # [2, 9] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "virial": torch.stack([v_A_hat, v_B_hat], dim=0), + "find_virial": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss(intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + + def test_mae_grad_accum(self): + """Spin virial MAE meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, v_A, v_A_hat, v_B, v_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _t(9) + v_hat = _t(9) + loss_obj = self._make_loss() + mp_mask = { + "virial": v.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"virial": v.unsqueeze(0), "mask_mag": _MASK_MAG_B_PT} + lb = {"virial": v_hat.unsqueeze(0), "find_virial": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossForceMagUnchanged: + """Guard: the padding mask must NOT affect the force_mag / mask_mag term. + + The force_mag path uses mask_mag (spin virtual-atom mask), which is a + completely separate concept from the padding mask model_pred["mask"]. + After the Task-5 implementation, presenting a padding mask must leave + the force_mag loss bit-identical. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_padding_mask_does_not_affect_force_mag(self): + """force_mag loss is bit-identical with and without padding mask.""" + fm = _rnd_t(2, NP, 3) + fm_hat = _rnd_t(2, NP, 3) + loss_obj = self._make_loss() + + def _run(with_mask): + mp = { + "force_mag": fm, + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + } + lb = { + "force_mag": fm_hat, + "find_force_mag": 1.0, + } + if with_mask: + mp["mask"] = _MASK_PAD_SPIN_PT + return _spin_loss_fn(loss_obj, mp, lb, NP) + + loss_with = _run(True) + loss_without = _run(False) + assert torch.isclose(loss_with, loss_without), ( + f"force_mag loss must be unchanged by padding mask: " + f"{loss_with.item()} vs {loss_without.item()}" + ) + + +# --------------------------------------------------------------------------- +# Part A: EnergySpinLoss atom_ener (has_ae) grad-accum invariant (pt) +# --------------------------------------------------------------------------- + + +class TestPTEnerSpinLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for atom_ener in pt EnergySpinLoss. + + RED before the Part-A has_ae mask fix; GREEN after. + """ + + def _make_loss(self, loss_func="mse"): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + mp = { + "atom_energy": ae_A.unsqueeze(0), # [1, NA, 1] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_A_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "atom_energy": ae_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_B_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NB + + def make_padded(): + ae_pad = _padded_atom_t(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom_t(ae_A_hat, ae_B_hat, 1) + mp = { + "atom_energy": ae_pad, + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = {"atom_ener": ae_hat_pad, "find_atom_ener": 1.0} + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Spin atom_ener MSE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Spin atom_ener MAE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom_ener spin loss as no mask.""" + ae = _t(NP, 1) + ae_hat = _t(NP, 1) + loss_obj = self._make_loss() + mp_mask = { + "atom_energy": ae.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = { + "atom_energy": ae.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"atom_ener": ae_hat.unsqueeze(0), "find_atom_ener": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# gen_force (has_gf) grad-accum invariant (pt) +# --------------------------------------------------------------------------- + +_NGEN_PT = 2 # number of generalized coordinates for tests + + +class TestPTEnergyLossGenForceGradAccum: + """gen_force (has_gf) excludes ghost atoms via force masking before projection. + + Expected GREEN immediately (no fix needed for gen_force). + """ + + def _make_loss(self): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=_NGEN_PT, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_A_hat.unsqueeze(0), + "find_force": 1.0, + "drdq": drdq_A.unsqueeze(0), # [1, NA*3, NGEN] + "find_drdq": 1.0, + } + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_B_hat.unsqueeze(0), + "find_force": 1.0, + "drdq": drdq_B.unsqueeze(0), # [1, NB*3, NGEN] + "find_drdq": 1.0, + } + return mp, lb, NB + + def make_padded(): + f_pad = _padded_force_t(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force_t(f_A_hat, f_B_hat) + # drdq ghost-atom slots are zero (ghost forces also zero, no contribution) + drdq_A_pad = torch.zeros( + NP * 3, _NGEN_PT, dtype=torch.float64, device="cpu" + ) + drdq_A_pad[: NA * 3] = drdq_A + drdq_pad = torch.stack([drdq_A_pad, drdq_B], dim=0) # [2, NP*3, NGEN] + mp = {"force": f_pad, "mask": _MASK_PAD_PT} + lb = { + "force": f_hat_pad, + "find_force": 1.0, + "drdq": drdq_pad, + "find_drdq": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """gen_force MSE meets the grad-accum invariant (GREEN: already correct).""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + drdq_A = _t(NA * 3, _NGEN_PT) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + drdq_B = _t(NB * 3, _NGEN_PT) + self._run_invariant( + self._make_loss(), f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B + ) + + +# --------------------------------------------------------------------------- +# force_mag MSE grad-accum invariant (pt) +# +# NOTE: force_mag MAE (pt) uses .sum() over frames, so MAE force_mag FAILS +# the invariant with a 2x factor when frames=2. This is a pre-existing +# frame-normalization artifact independent of ghost-atom masking. Ghost atoms +# are correctly excluded (mask_mag=False there). MAE artifact reported as +# NEEDS_CONTEXT in the audit report. +# --------------------------------------------------------------------------- + + +class TestPTEnerSpinLossForceMagMSEGradAccum: + """force_mag MSE meets the grad-accum invariant when ghost atoms are non-magnetic. + + With equal magnetic-atom counts (NM) across frames and ghost atoms having + mask_mag=False, the fancy-index selection excludes padding and the mean() + over [nf, NM, 3] satisfies the invariant. + Expected GREEN immediately (no fix needed for MSE force_mag). + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_mse_grad_accum(self): + """force_mag MSE (NM equal per frame, ghost atoms non-magnetic) meets invariant.""" + # Only magnetic-atom slots (first NM_PT) have non-zero values; others zero. + fm_A = _rnd_t(_NM_PT, 3) + fm_A_hat = _rnd_t(_NM_PT, 3) + fm_B = _rnd_t(_NM_PT, 3) + fm_B_hat = _rnd_t(_NM_PT, 3) + + def make_A(): + fm_A_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_full[:_NM_PT] = fm_A + fm_A_hat_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_hat_full[:_NM_PT] = fm_A_hat + mp = { + "force_mag": fm_A_full.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + } + lb = {"force_mag": fm_A_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NA + + def make_B(): + fm_B_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_full[:_NM_PT] = fm_B + fm_B_hat_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_hat_full[:_NM_PT] = fm_B_hat + mp = { + "force_mag": fm_B_full.unsqueeze(0), # [1, NB, 3] + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"force_mag": fm_B_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NB + + def make_padded(): + fm_A_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_pad[:_NM_PT] = fm_A + fm_A_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_hat_pad[:_NM_PT] = fm_A_hat + fm_B_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_pad[:_NM_PT] = fm_B + fm_B_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_hat_pad[:_NM_PT] = fm_B_hat + mp = { + "force_mag": torch.stack([fm_A_pad, fm_B_pad], dim=0), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force_mag": torch.stack([fm_A_hat_pad, fm_B_hat_pad], dim=0), + "find_force_mag": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) + + +class TestPTEnerSpinLossForceMagMAEGradAccum: + """force_mag MAE (pt) meets the grad-accum invariant. + + ``force_mag`` MAE now reduces with a plain mean over frames, magnetic atoms + and xyz (same reduction as force_mag MSE and force_real MAE), so a padded + ``[A+B]`` batch (nf=2) equals the mean of the two single-frame losses and the + loss is batch-size independent. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + + def test_mae_grad_accum(self): + """force_mag MAE meets the frame-average grad-accum invariant.""" + fm_A = _rnd_t(_NM_PT, 3) + fm_A_hat = _rnd_t(_NM_PT, 3) + fm_B = _rnd_t(_NM_PT, 3) + fm_B_hat = _rnd_t(_NM_PT, 3) + + def make_A(): + fm_A_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_full[:_NM_PT] = fm_A + fm_A_hat_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_hat_full[:_NM_PT] = fm_A_hat + mp = { + "force_mag": fm_A_full.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + } + lb = {"force_mag": fm_A_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NA + + def make_B(): + fm_B_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_full[:_NM_PT] = fm_B + fm_B_hat_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_hat_full[:_NM_PT] = fm_B_hat + mp = { + "force_mag": fm_B_full.unsqueeze(0), # [1, NB, 3] + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"force_mag": fm_B_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NB + + def make_padded(): + fm_A_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_pad[:_NM_PT] = fm_A + fm_A_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_hat_pad[:_NM_PT] = fm_A_hat + fm_B_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_pad[:_NM_PT] = fm_B + fm_B_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_hat_pad[:_NM_PT] = fm_B_hat + mp = { + "force_mag": torch.stack([fm_A_pad, fm_B_pad], dim=0), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force_mag": torch.stack([fm_A_hat_pad, fm_B_hat_pad], dim=0), + "find_force_mag": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + )