diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index 953b413dd1..fbf9332275 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -11,6 +11,10 @@ from deepmd.dpmodel.loss.loss import ( Loss, ) +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + masked_atom_num, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -134,11 +138,9 @@ def call( if "mask" in model_dict: # 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) + l2_local_loss_dos = masked_atom_mean( + xp.square(diff3d), maskf, self.numb_dos + ) else: l2_local_loss_dos = xp.mean(xp.square(diff3d)) loss += pref_ados * l2_local_loss_dos @@ -161,11 +163,9 @@ def call( if "mask" in model_dict: # 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) + l2_local_loss_cdf = masked_atom_mean( + xp.square(diff3d), maskf, self.numb_dos + ) else: l2_local_loss_cdf = xp.mean(xp.square(diff3d)) loss += pref_acdf * l2_local_loss_cdf @@ -181,12 +181,7 @@ def call( 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.mean( - xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) - ) - else: - atom_num = natoms + atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) 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 @@ -204,12 +199,7 @@ def call( 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.mean( - xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) - ) - else: - atom_num = natoms + atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) 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 5ee4916762..977eae4c07 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -11,6 +11,10 @@ from deepmd.dpmodel.loss.loss import ( Loss, ) +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + per_frame_component_mean, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -296,7 +300,7 @@ def call( 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] + per_frame = per_frame_component_mean(se) # [nf] if not self.use_huber: loss += pref_e * xp.mean(per_frame * inv**norm_exp) else: @@ -327,9 +331,7 @@ def call( l1_ener_loss = xp.mean(xp.abs(energy - energy_hat)) 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] + per_frame_ae = per_frame_component_mean(abs_e) # [nf] l1_ener_masked = xp.mean(per_frame_ae * inv) loss += pref_e * l1_ener_masked more_loss["mae_e"] = self.display_if_exist( @@ -346,8 +348,7 @@ def call( ) if mae: 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) + per_frame_ae = per_frame_component_mean(xp.abs(energy - energy_hat)) mae_e = xp.mean(per_frame_ae * inv) else: mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener @@ -362,10 +363,7 @@ def call( 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) + l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3) if not self.use_huber: loss += pref_f * l2_force_masked else: @@ -435,12 +433,8 @@ def call( 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: - 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) + l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) norm_2d = xp.reshape( @@ -474,11 +468,7 @@ def call( if mae: 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) + mae_f = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: mae_f = xp.mean(xp.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist(mae_f, find_force) @@ -494,7 +484,7 @@ def call( 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] + per_frame_v = per_frame_component_mean(se_v) # [nf] if not self.use_huber: loss += pref_v * xp.mean(per_frame_v * inv**norm_exp) else: @@ -526,8 +516,9 @@ def call( 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] + per_frame_v = per_frame_component_mean( + xp.abs(v_hat_2d - v2d) + ) # [nf] l1_virial_masked = xp.mean(per_frame_v * inv) loss += pref_v * l1_virial_masked more_loss["mae_v"] = self.display_if_exist( @@ -546,8 +537,7 @@ def call( 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) + per_frame_v = per_frame_component_mean(xp.abs(v_hat_2d - v2d)) mae_v = xp.mean(per_frame_v * inv) else: mae_v = ( @@ -565,10 +555,10 @@ def call( # 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) + l2_ae_masked = masked_atom_mean( + xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + ) if not self.use_huber: loss += pref_ae * l2_ae_masked else: @@ -609,10 +599,9 @@ def call( 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) + l1_ae_masked = masked_atom_mean( + xp.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + ) loss += pref_ae * l1_ae_masked more_loss["mae_ae"] = self.display_if_exist( l1_ae_masked, find_atom_ener @@ -637,13 +626,9 @@ def call( # 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) + l2_pf_masked = masked_atom_mean( + xp.square(diff_f_3d) * pf_3d, maskf, 3 + ) loss += pref_pf * l2_pf_masked more_loss["rmse_pf"] = self.display_if_exist( xp.sqrt(l2_pf_masked), find_atom_pref @@ -660,11 +645,7 @@ def call( 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) + l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3) loss += pref_pf * l1_pf_masked more_loss["mae_pf"] = self.display_if_exist( l1_pf_masked, find_atom_pref diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index e39e0a6a9e..c54288f49f 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -11,6 +11,10 @@ from deepmd.dpmodel.loss.loss import ( Loss, ) +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + per_frame_component_mean, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -159,7 +163,7 @@ def call( 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] + per_frame_e = per_frame_component_mean(se_e) # [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 @@ -174,9 +178,8 @@ def call( 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 + per_frame_ae = per_frame_component_mean( + xp.abs(energy_pred - energy_label) ) # [nf] l1_ener_masked = xp.mean(per_frame_ae * inv) loss += pref_e * l1_ener_masked @@ -190,8 +193,9 @@ def call( ) if mae: 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) + per_frame_ae = per_frame_component_mean( + xp.abs(energy_pred - energy_label) + ) mae_e = xp.mean(per_frame_ae * inv) else: mae_e = xp.mean(xp.abs(energy_pred - energy_label)) * atom_norm @@ -208,23 +212,15 @@ def call( 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) + l2_force_real_loss = masked_atom_mean(xp.square(diff_fr), maskf, 3) 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 = masked_atom_mean( + xp.abs(force_label - force_pred), maskf, 3 ) - 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)) @@ -239,13 +235,7 @@ def call( 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) + l1_force_real_masked = masked_atom_mean(abs_diff_fr, maskf, 3) loss += pref_fr * l1_force_real_masked more_loss["mae_fr"] = self.display_if_exist( l1_force_real_masked, find_force @@ -298,23 +288,18 @@ def call( # 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) + l2_atom_ener_loss = masked_atom_mean( + xp.square(ae_label - ae), maskf, 1 + ) 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) + l1_atom_ener_loss = masked_atom_mean( + xp.abs(ae_label - ae), maskf, 1 + ) loss += pref_ae * l1_atom_ener_loss more_loss["mae_ae"] = self.display_if_exist( l1_atom_ener_loss, find_atom_ener @@ -348,15 +333,15 @@ def call( if self.loss_func == "mse": 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] + per_frame_v = per_frame_component_mean(xp.square(diff_v)) # [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] + per_frame_mae_v = per_frame_component_mean( + xp.abs(diff_v) + ) # [nf] mae_v = xp.mean(per_frame_mae_v * inv) more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) else: @@ -372,8 +357,7 @@ def call( l1_virial_loss = xp.mean(xp.abs(diff_v)) 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] + per_frame_v = per_frame_component_mean(xp.abs(diff_v)) # [nf] l1_virial_masked = xp.mean(per_frame_v * inv) loss += pref_v * l1_virial_masked more_loss["mae_v"] = self.display_if_exist( diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py new file mode 100644 index 0000000000..068c11e2bc --- /dev/null +++ b/deepmd/dpmodel/loss/reduction.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared masked per-frame reduction idioms for the loss modules. + +These helpers factor out the three per-frame reduction patterns that the +mixed_type padding mask (PR #5738) introduced into every loss term (issue +#5768). They are written with ``array_api_compat`` so both the dpmodel +(numpy/jax/...) loss backend and the PyTorch loss backend can call them: the +PyTorch backend passes torch tensors and ``array_api_compat`` dispatches to the +torch namespace, preserving autograd and producing bit-identical results to the +previous hand-inlined torch code. + +Each helper implements ONLY the masked branch. Callers keep the original +non-masked expression in the ``else`` branch verbatim, so the "bit-identical +for non-mixed batches" guarantee from PR #5738 is preserved (defaulting the +mask to all-ones would change the reduction order at the ULP level and is +deliberately NOT done here). +""" + +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) + + +def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: + """Idiom 1: per-atom masked mean over ``ncomp`` components, averaged over frames. + + Parameters + ---------- + elem : Array + Non-negative per-element contribution of shape ``[nf, nloc, ncomp]`` + (already squared or abs'd, and pre-multiplied by any per-atom weight + such as ``atom_pref``). NOT yet multiplied by the mask. + maskf : Array + Per-atom real/ghost mask of shape ``[nf, nloc]`` (1.0 real, 0.0 ghost). + ncomp : int + Number of components per atom (force: 3, atom energy: 1, + dos: ``numb_dos``, tensor: ``tensor_size``). + + Returns + ------- + Array + ``mean_over_frames( sum(elem * mask) / (real_natoms * ncomp) )``. + An all-padding frame (zero real atoms) contributes a neutral ``0`` + instead of ``0/0 = NaN``. + """ + xp = array_api_compat.array_namespace(elem, maskf) + nf = elem.shape[0] + masked = elem * maskf[:, :, None] + per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * ncomp + # An all-padding frame has zero real atoms, so ``per_frame_dof`` is 0 and + # the ratio would be 0/0 = NaN -- poisoning the frame mean and, under + # autograd, its gradient. Divide by a safe denominator and map those frames + # to a neutral per-frame value of 0. Frames with real atoms are untouched, + # preserving the bit-identical guarantee. + has_dof = per_frame_dof > 0 + safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof)) + per_frame = xp.where( + has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum) + ) + return xp.mean(per_frame) + + +def per_frame_component_mean(err: Array) -> Array: + """Idiom 2 primitive: per-frame mean over the flattened component axis. + + Parameters + ---------- + err : Array + Per-frame error term of shape ``[nf, k]`` (already squared or abs'd). + + Returns + ------- + Array + Shape ``[nf]``: the mean over components for each frame. Callers apply + the extensive ``inv**exp`` weighting for both the loss term and the + RMSE display (which use different exponents), so ``err`` is reduced + once here and reused. + """ + xp = array_api_compat.array_namespace(err) + nf = err.shape[0] + return xp.mean(xp.reshape(err, (nf, -1)), axis=-1) + + +def masked_atom_num(mask: Array | None, natoms: Any, dtype: Any) -> Any: + """Idiom 3 companion: the display-only divisor for already-reduced globals. + + The global loss itself is a plain ``mean`` regardless of masking (global + quantities are padding-invariant); only the reported RMSE is divided by an + atom count. This returns that divisor. + + Parameters + ---------- + mask : Array or None + Per-atom mask of shape ``[nf, nloc]``, or ``None`` when not mixed_type. + natoms + Fallback atom count used when ``mask`` is ``None``. + dtype + Target dtype for the summed atom count (each backend passes the dtype + it currently uses: the diff's dtype for dpmodel, float32 for pt). + + Returns + ------- + Array or int + ``mean_over_frames(astype(sum(mask, axis=-1), dtype))`` when ``mask`` + is given, else ``natoms``. + """ + if mask is None: + return natoms + xp = array_api_compat.array_namespace(mask) + return xp.mean(xp.astype(xp.sum(mask, axis=-1), dtype)) diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index 4b21175185..eeddf2ecfc 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -11,6 +11,10 @@ from deepmd.dpmodel.loss.loss import ( Loss, ) +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + masked_atom_num, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -105,12 +109,12 @@ def call( if "mask" in model_dict: # 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) + diff3d = xp.reshape( + diff, (local_pred.shape[0], natoms, self.tensor_size) + ) + l2_local_loss = masked_atom_mean( + xp.square(diff3d), maskf, self.tensor_size + ) else: l2_local_loss = xp.mean(xp.square(diff)) loss += local_weight * l2_local_loss @@ -134,12 +138,7 @@ def call( 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.mean( - xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) - ) - else: - atom_num = natoms + atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) 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 bf2d2b9710..7046f73687 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -5,6 +5,10 @@ import torch +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + masked_atom_num, +) from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -154,12 +158,10 @@ def forward( ) # [nf, natoms, numb_dos] if "mask" in model_pred: # 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() + l2_local_loss_dos = masked_atom_mean( + torch.square(diff), maskf, self.numb_dos + ) else: l2_local_loss_dos = torch.mean(torch.square(diff)) if not self.inference: @@ -185,12 +187,10 @@ def forward( ) # [nf, natoms, numb_dos] if "mask" in model_pred: # 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() + l2_local_loss_cdf = masked_atom_mean( + torch.square(diff), maskf, self.numb_dos + ) else: l2_local_loss_cdf = torch.mean(torch.square(diff)) if not self.inference: @@ -210,10 +210,7 @@ def forward( 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).float().mean() - else: - atom_num = natoms + atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) if not self.inference: more_loss["l2_global_dos_loss"] = self.display_if_exist( l2_global_loss_dos.detach(), find_global @@ -235,10 +232,7 @@ def forward( 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).float().mean() - else: - atom_num = natoms + atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) 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 93f57618d6..bff23e2ba2 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -6,6 +6,10 @@ import torch import torch.nn.functional as F +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + per_frame_component_mean, +) from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -282,7 +286,7 @@ def forward( 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] + per_frame = per_frame_component_mean(se) # [nf] if not self.use_huber: loss += pref_e * torch.mean(per_frame * inv**norm_exp) else: @@ -319,8 +323,9 @@ def forward( reduction="mean", ) 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) + per_frame_ae = per_frame_component_mean( + torch.abs(energy_pred - energy_label) + ) l1_ener_masked = torch.mean(per_frame_ae * inv) loss += pref_e * l1_ener_masked more_loss["mae_e"] = self.display_if_exist( @@ -339,8 +344,9 @@ def forward( ) if mae: 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) + per_frame_ae = per_frame_component_mean( + torch.abs(energy_pred - energy_label) + ) mae_e = torch.mean(per_frame_ae * inv) else: mae_e = ( @@ -382,10 +388,9 @@ def forward( 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) + l2_f_masked = masked_atom_mean( + torch.square(diff_f_3d), maskf, 3 + ) if not self.use_huber: loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: @@ -456,12 +461,10 @@ def forward( 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) + l1_f_masked = masked_atom_mean( + torch.abs(diff_f_3d), maskf, 3 + ) else: diff_3 = (force_label - force_pred).reshape(_nf, _nloc, 3) norm_2d = torch.linalg.vector_norm( @@ -500,11 +503,7 @@ def forward( if mae: 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) + mae_f = masked_atom_mean(torch.abs(diff_f_3d), maskf, 3) else: mae_f = torch.mean(torch.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist( @@ -531,11 +530,9 @@ def forward( # 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) + l2_pf_masked = masked_atom_mean( + torch.square(diff_f_3d) * pf_3d, maskf, 3 + ) 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( @@ -554,11 +551,9 @@ def forward( 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) + l1_pf_masked = masked_atom_mean( + torch.abs(diff_f_3d) * pf_3d, maskf, 3 + ) 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 @@ -635,7 +630,7 @@ def forward( 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] + per_frame_v = per_frame_component_mean(se_v) # [nf] if not self.use_huber: loss += pref_v * torch.mean(per_frame_v * inv**norm_exp) else: @@ -671,8 +666,7 @@ def forward( reduction="mean", ) if maskf is not None: - abs_v = torch.abs(diff_v) # [nf, 9] - per_frame_v = torch.mean(abs_v, dim=-1) # [nf] + per_frame_v = per_frame_component_mean(torch.abs(diff_v)) # [nf] l1_v_masked = torch.mean(per_frame_v * inv) loss += pref_v * l1_v_masked more_loss["mae_v"] = self.display_if_exist( @@ -690,8 +684,7 @@ def forward( ) if mae: if maskf is not None: - abs_v = torch.abs(diff_v) - per_frame_v = torch.mean(abs_v, dim=-1) + per_frame_v = per_frame_component_mean(torch.abs(diff_v)) mae_v = torch.mean(per_frame_v * inv) else: mae_v = torch.mean(torch.abs(diff_v)) * atom_norm @@ -717,10 +710,10 @@ def forward( # 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) + per_frame_dof = maskf.sum(dim=-1) # [nf], kept for huber branch + l2_ae_masked = masked_atom_mean( + torch.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + ) if not self.use_huber: loss += (pref_ae * l2_ae_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: @@ -765,10 +758,9 @@ def forward( 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) + l1_ae_masked = masked_atom_mean( + torch.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1 + ) 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 diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 84c229e1c2..2cf8518ba5 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -6,6 +6,10 @@ import torch import torch.nn.functional as F +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + per_frame_component_mean, +) from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -221,7 +225,7 @@ def forward( 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] + per_frame_e = per_frame_component_mean(se_e) # [nf] if not self.inference: more_loss["l2_ener_loss"] = self.display_if_exist( torch.mean(per_frame_e).detach(), find_energy @@ -251,8 +255,9 @@ def forward( ) 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] + per_frame_ae = per_frame_component_mean( + torch.abs(energy_pred - energy_label) + ) # [nf] l1_ener_masked = torch.mean(per_frame_ae * inv) loss += pref_e * l1_ener_masked more_loss["mae_e"] = self.display_if_exist( @@ -270,8 +275,9 @@ def forward( ) if mae: 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) + per_frame_ae = per_frame_component_mean( + torch.abs(energy_pred - energy_label) + ) mae_e = torch.mean(per_frame_ae * inv) else: mae_e = ( @@ -290,11 +296,9 @@ def forward( 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) + l2_force_real_loss = masked_atom_mean( + torch.square(diff_fr), maskf, 3 + ) if not self.inference: more_loss["l2_force_r_loss"] = self.display_if_exist( l2_force_real_loss.detach(), find_force_r @@ -305,9 +309,7 @@ def forward( rmse_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) + mae_fr = masked_atom_mean(torch.abs(diff_fr), maskf, 3) more_loss["mae_fr"] = self.display_if_exist( mae_fr.detach(), find_force_r ) @@ -333,11 +335,7 @@ def forward( ) # [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) + l1_force_real_masked = masked_atom_mean(abs_diff_fr, maskf, 3) more_loss["mae_fr"] = self.display_if_exist( l1_force_real_masked.detach(), find_force_r ) @@ -409,12 +407,10 @@ def forward( # 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) + l2_atom_ener_loss = masked_atom_mean( + torch.square(ae_label - ae), maskf, 1 + ) if not self.inference: more_loss["l2_atom_ener_loss"] = self.display_if_exist( l2_atom_ener_loss.detach(), find_atom_ener @@ -425,10 +421,9 @@ def forward( 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) + l1_atom_ener_loss = masked_atom_mean( + torch.abs(ae_label - ae), maskf, 1 + ) 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 @@ -476,8 +471,7 @@ def forward( if self.loss_func == "mse": 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] + per_frame_v = per_frame_component_mean(torch.square(diff_v)) # [nf] if not self.inference: more_loss["l2_virial_loss"] = self.display_if_exist( torch.mean(per_frame_v).detach(), find_virial @@ -488,8 +482,9 @@ def forward( rmse_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] + per_frame_mae_v = per_frame_component_mean( + torch.abs(diff_v) + ) # [nf] mae_v = torch.mean(per_frame_mae_v * inv) more_loss["mae_v"] = self.display_if_exist( mae_v.detach(), find_virial @@ -518,8 +513,7 @@ def forward( ) 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] + per_frame_v = per_frame_component_mean(torch.abs(diff_v)) # [nf] l1_virial_masked = torch.mean(per_frame_v * inv) loss += pref_v * l1_virial_masked more_loss["mae_v"] = self.display_if_exist( diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index c46c2ec0dd..660cbd0337 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -5,6 +5,10 @@ import torch +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + masked_atom_num, +) from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -130,13 +134,13 @@ def forward( diff = diff * atomic_weight if "mask" in model_pred: # 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() + diff3d = diff.reshape( + local_tensor_pred.shape[0], natoms, self.tensor_size + ) + l2_local_loss = masked_atom_mean( + torch.square(diff3d), maskf, self.tensor_size + ) else: l2_local_loss = torch.mean(torch.square(diff)) if not self.inference: @@ -162,10 +166,7 @@ def forward( 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).float().mean() - else: - atom_num = natoms + atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) 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_reduction.py b/source/tests/common/dpmodel/test_loss_reduction.py new file mode 100644 index 0000000000..fecb0156c6 --- /dev/null +++ b/source/tests/common/dpmodel/test_loss_reduction.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the shared masked-reduction idioms (issue #5768).""" + +import numpy as np +import pytest + +from deepmd.dpmodel.loss.reduction import ( + masked_atom_mean, + masked_atom_num, + per_frame_component_mean, +) + +torch = pytest.importorskip("torch") + + +class TestMaskedAtomMean: + """Idiom 1: per-atom masked mean over ncomp components, averaged over frames.""" + + def _ref(self, elem, maskf, ncomp): + # reference reduction, numpy + nf = elem.shape[0] + masked = elem * maskf[:, :, None] + pfs = masked.reshape(nf, -1).sum(axis=-1) + pfd = maskf.sum(axis=-1) * ncomp + return (pfs / pfd).mean() + + @pytest.mark.parametrize("ncomp", [1, 3]) # atom-energy (1) and force (3) + def test_numpy_matches_reference(self, ncomp) -> None: + rng = np.random.default_rng(0) + elem = rng.random((2, 4, ncomp)) + maskf = np.array([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]]) + got = masked_atom_mean(elem, maskf, ncomp) + np.testing.assert_allclose(got, self._ref(elem, maskf, ncomp), rtol=0, atol=0) + + def test_torch_autograd_and_bit_identical(self) -> None: + elem_np = np.random.default_rng(1).random((2, 4, 3)) + maskf_np = np.array([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]]) + elem = torch.tensor(elem_np, requires_grad=True) + maskf = torch.tensor(maskf_np) + out = masked_atom_mean(elem, maskf, 3) + out.backward() + assert elem.grad is not None + # bit-identical to torch-native inline form + en = torch.tensor(elem_np) + mn = torch.tensor(maskf_np) + pfs = (en * mn[:, :, None]).reshape(2, -1).sum(dim=-1) + pfd = mn.sum(dim=-1) * 3 + ref = torch.mean(pfs / pfd) + assert out.item() == ref.item() + + def test_all_padding_frame_is_not_nan(self) -> None: + # a frame with zero real atoms has per_frame_dof == 0; the ratio must + # not become 0/0 = NaN (an independent finiteness invariant -- the + # reference formula shares the bug, so equality checks cannot catch it) + elem = np.random.default_rng(3).random((2, 4, 3)) + maskf = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) + got = masked_atom_mean(elem, maskf, 3) + assert np.isfinite(got) + # the all-padding frame contributes a neutral 0, so the result is the + # first frame's masked mean divided by the number of frames + pfs0 = (elem[0] * maskf[0][:, None]).reshape(-1).sum() + expected = (pfs0 / (maskf[0].sum() * 3)) / 2 + np.testing.assert_allclose(got, expected, rtol=0, atol=0) + + def test_all_padding_frame_torch_grad_is_not_nan(self) -> None: + # the guard must keep both the value and the autograd gradient finite + elem_np = np.random.default_rng(4).random((2, 4, 3)) + maskf_np = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) + elem = torch.tensor(elem_np, requires_grad=True) + maskf = torch.tensor(maskf_np) + out = masked_atom_mean(elem, maskf, 3) + out.backward() + assert torch.isfinite(out).item() + assert elem.grad is not None + assert torch.isfinite(elem.grad).all().item() + + +class TestPerFrameComponentMean: + """Idiom 2 primitive: per-frame mean over the flattened component axis.""" + + @pytest.mark.parametrize("k", [1, 9]) # energy (k=1) and virial (k=9) + def test_numpy_matches_reference(self, k) -> None: + rng = np.random.default_rng(2) + err = rng.random((3, k)) + got = per_frame_component_mean(err) + np.testing.assert_allclose( + got, err.reshape(3, -1).mean(axis=-1), rtol=0, atol=0 + ) + + def test_torch_bit_identical(self) -> None: + err_np = np.random.default_rng(3).random((3, 9)) + got = per_frame_component_mean(torch.tensor(err_np)) + ref = torch.mean(torch.tensor(err_np).reshape(3, -1), dim=-1) + assert torch.equal(got, ref) + + def test_torch_autograd(self) -> None: + err_np = np.random.default_rng(4).random((3, 9)) + err = torch.tensor(err_np, requires_grad=True) + out = per_frame_component_mean(err) + out.sum().backward() + assert err.grad is not None + + +class TestMaskedAtomNum: + """Idiom 3 companion: display-only divisor for already-reduced globals.""" + + def test_none_returns_natoms(self) -> None: + assert masked_atom_num(None, 17, np.float64) == 17 + + def test_numpy_mean_real_atoms(self) -> None: + mask = np.array([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) + got = masked_atom_num(mask, 3, np.float64) + np.testing.assert_allclose(got, np.mean(np.sum(mask, axis=-1)), rtol=0, atol=0) + + def test_torch_bit_identical_float32(self) -> None: + mask_np = np.array([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) + got = masked_atom_num(torch.tensor(mask_np), 3, torch.float32) + ref = torch.tensor(mask_np).sum(-1).float().mean() + assert got.item() == ref.item()