diff --git a/deepmd/dpmodel/loss/__init__.py b/deepmd/dpmodel/loss/__init__.py index 6ceb116d85..115f1f6b03 100644 --- a/deepmd/dpmodel/loss/__init__.py +++ b/deepmd/dpmodel/loss/__init__.py @@ -1 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.loss.dos import ( + DOSLoss, +) +from deepmd.dpmodel.loss.ener import ( + EnergyLoss, +) +from deepmd.dpmodel.loss.ener_spin import ( + EnergySpinLoss, +) +from deepmd.dpmodel.loss.property import ( + PropertyLoss, +) +from deepmd.dpmodel.loss.tensor import ( + TensorLoss, +) + +__all__ = [ + "DOSLoss", + "EnergyLoss", + "EnergySpinLoss", + "PropertyLoss", + "TensorLoss", +] diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py new file mode 100644 index 0000000000..0036ae1620 --- /dev/null +++ b/deepmd/dpmodel/loss/dos.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.loss.loss import ( + Loss, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class DOSLoss(Loss): + r"""Loss on DOS (density of states) for both local and global predictions. + + Parameters + ---------- + starter_learning_rate : float + The learning rate at the start of the training. + numb_dos : int + The number of DOS components. + start_pref_dos : float + The prefactor of global DOS loss at the start of the training. + limit_pref_dos : float + The prefactor of global DOS loss at the end of the training. + start_pref_cdf : float + The prefactor of global CDF loss at the start of the training. + limit_pref_cdf : float + The prefactor of global CDF loss at the end of the training. + start_pref_ados : float + The prefactor of atomic DOS loss at the start of the training. + limit_pref_ados : float + The prefactor of atomic DOS loss at the end of the training. + start_pref_acdf : float + The prefactor of atomic CDF loss at the start of the training. + limit_pref_acdf : float + The prefactor of atomic CDF loss at the end of the training. + **kwargs + Other keyword arguments. + """ + + def __init__( + self, + starter_learning_rate: float, + numb_dos: int, + start_pref_dos: float = 1.00, + limit_pref_dos: float = 1.00, + start_pref_cdf: float = 1000, + limit_pref_cdf: float = 1.00, + start_pref_ados: float = 0.0, + limit_pref_ados: float = 0.0, + start_pref_acdf: float = 0.0, + limit_pref_acdf: float = 0.0, + **kwargs: Any, + ) -> None: + self.starter_learning_rate = starter_learning_rate + self.numb_dos = numb_dos + self.start_pref_dos = start_pref_dos + self.limit_pref_dos = limit_pref_dos + self.start_pref_cdf = start_pref_cdf + self.limit_pref_cdf = limit_pref_cdf + self.start_pref_ados = start_pref_ados + self.limit_pref_ados = limit_pref_ados + self.start_pref_acdf = start_pref_acdf + self.limit_pref_acdf = limit_pref_acdf + + assert ( + self.start_pref_dos >= 0.0 + and self.limit_pref_dos >= 0.0 + and self.start_pref_cdf >= 0.0 + and self.limit_pref_cdf >= 0.0 + and self.start_pref_ados >= 0.0 + and self.limit_pref_ados >= 0.0 + and self.start_pref_acdf >= 0.0 + and self.limit_pref_acdf >= 0.0 + ), "Can not assign negative weight to `pref` and `pref_atomic`" + + self.has_dos = start_pref_dos != 0.0 or limit_pref_dos != 0.0 + self.has_cdf = start_pref_cdf != 0.0 or limit_pref_cdf != 0.0 + self.has_ados = start_pref_ados != 0.0 or limit_pref_ados != 0.0 + self.has_acdf = start_pref_acdf != 0.0 or limit_pref_acdf != 0.0 + + assert self.has_dos or self.has_cdf or self.has_ados or self.has_acdf, ( + "Can not assign zero weight to all pref terms" + ) + + def call( + self, + learning_rate: float, + natoms: int, + model_dict: dict[str, Array], + label_dict: dict[str, Array], + mae: bool = False, + ) -> tuple[Array, dict[str, Array]]: + """Calculate loss from model results and labeled results.""" + # Get array namespace from any available tensor + first_key = next(iter(model_dict)) + xp = array_api_compat.array_namespace(model_dict[first_key]) + + coef = learning_rate / self.starter_learning_rate + pref_dos = ( + self.limit_pref_dos + (self.start_pref_dos - self.limit_pref_dos) * coef + ) + pref_cdf = ( + self.limit_pref_cdf + (self.start_pref_cdf - self.limit_pref_cdf) * coef + ) + pref_ados = ( + self.limit_pref_ados + (self.start_pref_ados - self.limit_pref_ados) * coef + ) + pref_acdf = ( + self.limit_pref_acdf + (self.start_pref_acdf - self.limit_pref_acdf) * coef + ) + + loss = 0 + more_loss = {} + + if self.has_ados and "atom_dos" in model_dict and "atom_dos" in label_dict: + find_local = label_dict.get("find_atom_dos", 0.0) + pref_ados = pref_ados * find_local + local_pred = xp.reshape(model_dict["atom_dos"], (-1, natoms, self.numb_dos)) + local_label = xp.reshape( + label_dict["atom_dos"], (-1, natoms, self.numb_dos) + ) + diff = xp.reshape(local_pred - local_label, (-1, self.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) + else: + l2_local_loss_dos = xp.mean(xp.square(diff)) + loss += pref_ados * l2_local_loss_dos + more_loss["rmse_local_dos"] = self.display_if_exist( + xp.sqrt(l2_local_loss_dos), find_local + ) + + if self.has_acdf and "atom_dos" in model_dict and "atom_dos" in label_dict: + find_local = label_dict.get("find_atom_dos", 0.0) + pref_acdf = pref_acdf * find_local + local_pred_cdf = xp.cumulative_sum( + xp.reshape(model_dict["atom_dos"], (-1, natoms, self.numb_dos)), + axis=-1, + ) + local_label_cdf = xp.cumulative_sum( + 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)) + 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) + else: + l2_local_loss_cdf = xp.mean(xp.square(diff)) + loss += pref_acdf * l2_local_loss_cdf + more_loss["rmse_local_cdf"] = self.display_if_exist( + xp.sqrt(l2_local_loss_cdf), find_local + ) + + if self.has_dos and "dos" in model_dict and "dos" in label_dict: + find_global = label_dict.get("find_dos", 0.0) + pref_dos = pref_dos * find_global + 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 + 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(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 + ) + + if self.has_cdf and "dos" in model_dict and "dos" in label_dict: + find_global = label_dict.get("find_dos", 0.0) + pref_cdf = pref_cdf * find_global + global_pred_cdf = xp.cumulative_sum( + xp.reshape(model_dict["dos"], (-1, self.numb_dos)), axis=-1 + ) + global_label_cdf = xp.cumulative_sum( + xp.reshape(label_dict["dos"], (-1, self.numb_dos)), axis=-1 + ) + diff = global_pred_cdf - global_label_cdf + 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(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 + ) + + more_loss["rmse"] = xp.sqrt(loss) + return loss, more_loss + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + if self.has_ados or self.has_acdf: + label_requirement.append( + DataRequirementItem( + "atom_dos", + ndof=self.numb_dos, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.has_dos or self.has_cdf: + label_requirement.append( + DataRequirementItem( + "dos", + ndof=self.numb_dos, + atomic=False, + must=False, + high_prec=False, + ) + ) + return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "DOSLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "numb_dos": self.numb_dos, + "start_pref_dos": self.start_pref_dos, + "limit_pref_dos": self.limit_pref_dos, + "start_pref_cdf": self.start_pref_cdf, + "limit_pref_cdf": self.limit_pref_cdf, + "start_pref_ados": self.start_pref_ados, + "limit_pref_ados": self.limit_pref_ados, + "start_pref_acdf": self.start_pref_acdf, + "limit_pref_acdf": self.limit_pref_acdf, + } + + @classmethod + def deserialize(cls, data: dict) -> "DOSLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 201a77fcdd..1a99183a79 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -166,6 +166,7 @@ def call( natoms: int, model_dict: dict[str, Array], label_dict: dict[str, Array], + mae: bool = False, ) -> tuple[Array, dict[str, Array]]: """Calculate loss from model results and labeled results.""" energy = model_dict["energy"] @@ -266,6 +267,11 @@ def call( 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 + 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)) @@ -304,6 +310,9 @@ def call( raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for force loss." ) + if mae: + 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,)) virial_hat_reshape = xp.reshape(virial_hat, (-1,)) @@ -333,6 +342,9 @@ def call( 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 + more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) if self.has_ae: atom_ener_reshape = xp.reshape(atom_ener, (-1,)) atom_ener_hat_reshape = xp.reshape(atom_ener_hat, (-1,)) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py new file mode 100644 index 0000000000..a13d626764 --- /dev/null +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.loss.loss import ( + Loss, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class EnergySpinLoss(Loss): + r"""Loss on energy, real force, magnetic force and virial for spin models. + + Parameters + ---------- + starter_learning_rate : float + The learning rate at the start of the training. + start_pref_e : float + The prefactor of energy loss at the start of the training. + limit_pref_e : float + The prefactor of energy loss at the end of the training. + start_pref_fr : float + The prefactor of real force loss at the start of the training. + limit_pref_fr : float + The prefactor of real force loss at the end of the training. + start_pref_fm : float + The prefactor of magnetic force loss at the start of the training. + limit_pref_fm : float + The prefactor of magnetic force loss at the end of the training. + start_pref_v : float + The prefactor of virial loss at the start of the training. + limit_pref_v : float + The prefactor of virial loss at the end of the training. + start_pref_ae : float + The prefactor of atomic energy loss at the start of the training. + limit_pref_ae : float + The prefactor of atomic energy loss at the end of the training. + enable_atom_ener_coeff : bool + if true, the energy will be computed as \sum_i c_i E_i + loss_func : str + Loss function type: 'mse' or 'mae'. + **kwargs + Other keyword arguments. + """ + + def __init__( + self, + starter_learning_rate: float = 1.0, + start_pref_e: float = 0.0, + limit_pref_e: float = 0.0, + start_pref_fr: float = 0.0, + limit_pref_fr: float = 0.0, + start_pref_fm: float = 0.0, + limit_pref_fm: float = 0.0, + start_pref_v: float = 0.0, + limit_pref_v: float = 0.0, + start_pref_ae: float = 0.0, + limit_pref_ae: float = 0.0, + enable_atom_ener_coeff: bool = False, + loss_func: str = "mse", + **kwargs: Any, + ) -> None: + valid_loss_funcs = ["mse", "mae"] + if loss_func not in valid_loss_funcs: + raise ValueError( + f"Invalid loss_func '{loss_func}'. Must be one of {valid_loss_funcs}." + ) + self.loss_func = loss_func + self.starter_learning_rate = starter_learning_rate + self.start_pref_e = start_pref_e + self.limit_pref_e = limit_pref_e + self.start_pref_fr = start_pref_fr + self.limit_pref_fr = limit_pref_fr + self.start_pref_fm = start_pref_fm + self.limit_pref_fm = limit_pref_fm + self.start_pref_v = start_pref_v + self.limit_pref_v = limit_pref_v + self.start_pref_ae = start_pref_ae + self.limit_pref_ae = limit_pref_ae + self.enable_atom_ener_coeff = enable_atom_ener_coeff + self.has_e = self.start_pref_e != 0.0 or self.limit_pref_e != 0.0 + self.has_fr = self.start_pref_fr != 0.0 or self.limit_pref_fr != 0.0 + self.has_fm = self.start_pref_fm != 0.0 or self.limit_pref_fm != 0.0 + self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0 + self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0 + + def call( + self, + learning_rate: float, + natoms: int, + model_dict: dict[str, Array], + label_dict: dict[str, Array], + mae: bool = False, + ) -> tuple[Array, dict[str, Array]]: + """Calculate loss from model results and labeled results.""" + energy = model_dict["energy"] + xp = array_api_compat.array_namespace(energy) + + 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 + pref_fm = self.limit_pref_fm + (self.start_pref_fm - self.limit_pref_fm) * coef + pref_v = self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * coef + pref_ae = self.limit_pref_ae + (self.start_pref_ae - self.limit_pref_ae) * coef + + loss = 0 + more_loss = {} + atom_norm = 1.0 / natoms + + if self.has_e: + energy_pred = model_dict["energy"] + energy_label = label_dict["energy"] + find_energy = label_dict.get("find_energy", 0.0) + pref_e = pref_e * find_energy + if self.enable_atom_ener_coeff and "atom_energy" in model_dict: + atom_ener_pred = model_dict["atom_energy"] + atom_ener_coeff = label_dict["atom_ener_coeff"] + 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 * (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 + ) + if mae: + 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) + + if self.has_fr: + find_force = label_dict.get("find_force", 0.0) + pref_fr = pref_fr * find_force + 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) + 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 + ) + + if self.has_fm: + find_force_mag = label_dict.get("find_force_mag", 0.0) + pref_fm = pref_fm * find_force_mag + force_mag_pred = model_dict["force_mag"] + force_mag_label = label_dict["force_mag"] + mask_mag = model_dict["mask_mag"] + # mask_mag: [nframes, natoms, 1], bool -> use mask multiplication + mask_float = xp.astype(mask_mag, force_mag_pred.dtype) + # zero out non-magnetic atoms + diff_fm = (force_mag_label - force_mag_pred) * mask_float + n_valid = xp.sum(mask_float) + if self.loss_func == "mse": + l2_force_mag_loss = xp.sum(xp.square(diff_fm)) / (n_valid * 3) + loss += pref_fm * l2_force_mag_loss + more_loss["rmse_fm"] = self.display_if_exist( + xp.sqrt(l2_force_mag_loss), find_force_mag + ) + if mae: + mae_fm = xp.sum(xp.abs(diff_fm)) / (n_valid * 3) + 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 + 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 + ) + + if self.has_ae: + find_atom_ener = label_dict.get("find_atom_ener", 0.0) + 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 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 + if self.loss_func == "mse": + l2_virial_loss = xp.mean(xp.square(diff_v)) + loss += atom_norm * (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 + ) + + more_loss["rmse"] = xp.sqrt(loss) + return loss, more_loss + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + if self.has_e: + label_requirement.append( + DataRequirementItem( + "energy", + ndof=1, + atomic=False, + must=False, + high_prec=True, + ) + ) + if self.has_fr: + label_requirement.append( + DataRequirementItem( + "force", + ndof=3, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.has_fm: + label_requirement.append( + DataRequirementItem( + "force_mag", + ndof=3, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.has_v: + label_requirement.append( + DataRequirementItem( + "virial", + ndof=9, + atomic=False, + must=False, + high_prec=False, + ) + ) + if self.has_ae: + label_requirement.append( + DataRequirementItem( + "atom_ener", + ndof=1, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.enable_atom_ener_coeff: + label_requirement.append( + DataRequirementItem( + "atom_ener_coeff", + ndof=1, + atomic=True, + must=False, + high_prec=False, + default=1.0, + ) + ) + return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "EnergySpinLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "start_pref_e": self.start_pref_e, + "limit_pref_e": self.limit_pref_e, + "start_pref_fr": self.start_pref_fr, + "limit_pref_fr": self.limit_pref_fr, + "start_pref_fm": self.start_pref_fm, + "limit_pref_fm": self.limit_pref_fm, + "start_pref_v": self.start_pref_v, + "limit_pref_v": self.limit_pref_v, + "start_pref_ae": self.start_pref_ae, + "limit_pref_ae": self.limit_pref_ae, + "enable_atom_ener_coeff": self.enable_atom_ener_coeff, + "loss_func": self.loss_func, + } + + @classmethod + def deserialize(cls, data: dict) -> "EnergySpinLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py index 05878deabc..7d4e052dca 100644 --- a/deepmd/dpmodel/loss/loss.py +++ b/deepmd/dpmodel/loss/loss.py @@ -28,6 +28,7 @@ def call( natoms: int, model_dict: dict[str, Array], label_dict: dict[str, Array], + mae: bool = False, ) -> tuple[Array, dict[str, Array]]: """Calculate loss from model results and labeled results. diff --git a/deepmd/dpmodel/loss/property.py b/deepmd/dpmodel/loss/property.py new file mode 100644 index 0000000000..7d658ff925 --- /dev/null +++ b/deepmd/dpmodel/loss/property.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.loss.loss import ( + Loss, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class PropertyLoss(Loss): + r"""Loss on property predictions. + + Parameters + ---------- + task_dim : int + The output dimension of property fitting net. + var_name : str + The property variable name. + loss_func : str + The loss function: "smooth_mae", "mae", "mse", "rmse", "mape". + metric : list[str] + The metrics to report. + beta : float + The 'beta' parameter in 'smooth_mae' loss. + out_bias : list or None + The bias for normalization. + out_std : list or None + The standard deviation for normalization. + intensive : bool + Whether the property is intensive. + **kwargs + Other keyword arguments. + """ + + def __init__( + self, + task_dim: int, + var_name: str, + loss_func: str = "smooth_mae", + metric: list[str] | None = None, + beta: float = 1.00, + out_bias: list | None = None, + out_std: list | None = None, + intensive: bool = False, + **kwargs: Any, + ) -> None: + if metric is None: + metric = ["mae"] + self.task_dim = task_dim + self.var_name = var_name + self.loss_func = loss_func + self.metric = metric + self.beta = beta + self.out_bias = out_bias + self.out_std = out_std + self.intensive = intensive + + def call( + self, + learning_rate: float, + natoms: int, + model_dict: dict[str, Array], + label_dict: dict[str, Array], + mae: bool = False, + ) -> tuple[Array, dict[str, Array]]: + """Calculate loss from model results and labeled results.""" + del learning_rate, mae + var_name = self.var_name + pred = model_dict[var_name] + xp = array_api_compat.array_namespace(pred) + dev = array_api_compat.device(pred) + label = label_dict[var_name] + + # Normalize by natoms for extensive properties (without mutating input) + if not self.intensive: + pred = pred / natoms + label = label / natoms + + # Get out_std and out_bias + if self.out_std is not None: + out_std = xp.asarray(self.out_std, dtype=pred.dtype, device=dev) + else: + out_std = xp.ones((self.task_dim,), dtype=pred.dtype, device=dev) + if self.out_bias is not None: + out_bias = xp.asarray(self.out_bias, dtype=pred.dtype, device=dev) + else: + out_bias = xp.zeros((self.task_dim,), dtype=pred.dtype, device=dev) + + loss = xp.zeros((), dtype=pred.dtype, device=dev) + more_loss = {} + + norm_pred = (pred - out_bias) / out_std + norm_label = (label - out_bias) / out_std + diff = norm_label - norm_pred + + if self.loss_func == "smooth_mae": + abs_diff = xp.abs(diff) + smooth_l1 = xp.where( + abs_diff < self.beta, + 0.5 * diff**2 / self.beta, + abs_diff - 0.5 * self.beta, + ) + loss = loss + xp.sum(smooth_l1) + elif self.loss_func == "mae": + loss = loss + xp.sum(xp.abs(diff)) + elif self.loss_func == "mse": + loss = loss + xp.sum(xp.square(diff)) + elif self.loss_func == "rmse": + loss = loss + xp.sqrt(xp.mean(xp.square(diff))) + elif self.loss_func == "mape": + loss = loss + xp.mean(xp.abs((label - pred) / (label + 1e-3))) + else: + raise RuntimeError(f"Unknown loss function : {self.loss_func}") + + # metrics (computed on un-normalized values) + if "smooth_mae" in self.metric: + abs_raw = xp.abs(label - pred) + more_loss["smooth_mae"] = xp.mean( + xp.where( + abs_raw < self.beta, + 0.5 * (label - pred) ** 2 / self.beta, + abs_raw - 0.5 * self.beta, + ) + ) + if "mae" in self.metric: + more_loss["mae"] = xp.mean(xp.abs(label - pred)) + if "mse" in self.metric: + more_loss["mse"] = xp.mean(xp.square(label - pred)) + if "rmse" in self.metric: + more_loss["rmse"] = xp.sqrt(xp.mean(xp.square(label - pred))) + if "mape" in self.metric: + more_loss["mape"] = xp.mean(xp.abs((label - pred) / (label + 1e-3))) + + return loss, more_loss + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + label_requirement.append( + DataRequirementItem( + self.var_name, + ndof=self.task_dim, + atomic=False, + must=True, + high_prec=True, + ) + ) + return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "PropertyLoss", + "@version": 1, + "task_dim": self.task_dim, + "var_name": self.var_name, + "loss_func": self.loss_func, + "metric": self.metric, + "beta": self.beta, + "out_bias": self.out_bias, + "out_std": self.out_std, + "intensive": self.intensive, + } + + @classmethod + def deserialize(cls, data: dict) -> "PropertyLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py new file mode 100644 index 0000000000..e367457451 --- /dev/null +++ b/deepmd/dpmodel/loss/tensor.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.loss.loss import ( + Loss, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class TensorLoss(Loss): + r"""Loss on local and global tensors (e.g. dipole, polarizability). + + Parameters + ---------- + tensor_name : str + The name of the tensor in model predictions. + tensor_size : int + The size (dimension) of the tensor. + label_name : str + The name of the tensor in labels. + pref_atomic : float + The prefactor of the weight of atomic (local) loss. + pref : float + The prefactor of the weight of global loss. + enable_atomic_weight : bool + If true, atomic weight will be used in the loss calculation. + **kwargs + Other keyword arguments. + """ + + def __init__( + self, + tensor_name: str, + tensor_size: int, + label_name: str, + pref_atomic: float = 0.0, + pref: float = 0.0, + enable_atomic_weight: bool = False, + **kwargs: Any, + ) -> None: + self.tensor_name = tensor_name + self.tensor_size = tensor_size + self.label_name = label_name + self.local_weight = pref_atomic + self.global_weight = pref + self.enable_atomic_weight = enable_atomic_weight + + assert self.local_weight >= 0.0 and self.global_weight >= 0.0, ( + "Can not assign negative weight to `pref` and `pref_atomic`" + ) + self.has_local_weight = self.local_weight > 0.0 + self.has_global_weight = self.global_weight > 0.0 + assert self.has_local_weight or self.has_global_weight, ( + "Can not assign zero weight both to `pref` and `pref_atomic`" + ) + + def call( + self, + learning_rate: float, + natoms: int, + model_dict: dict[str, Array], + label_dict: dict[str, Array], + mae: bool = False, + ) -> tuple[Array, dict[str, Array]]: + """Calculate loss from model results and labeled results.""" + del learning_rate, mae + first_key = next(iter(model_dict)) + xp = array_api_compat.array_namespace(model_dict[first_key]) + + if self.enable_atomic_weight: + atomic_weight = xp.reshape(label_dict["atom_weight"], (-1, 1)) + else: + atomic_weight = 1.0 + + loss = 0 + more_loss = {} + + if ( + self.has_local_weight + and self.tensor_name in model_dict + and "atom_" + self.label_name in label_dict + ): + find_local = label_dict.get("find_atom_" + self.label_name, 0.0) + local_weight = self.local_weight * find_local + local_pred = xp.reshape( + model_dict[self.tensor_name], (-1, natoms, self.tensor_size) + ) + local_label = xp.reshape( + label_dict["atom_" + self.label_name], (-1, natoms, self.tensor_size) + ) + 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) + else: + l2_local_loss = xp.mean(xp.square(diff)) + loss += local_weight * l2_local_loss + more_loss[f"rmse_local_{self.tensor_name}"] = self.display_if_exist( + xp.sqrt(l2_local_loss), find_local + ) + + if ( + self.has_global_weight + and "global_" + self.tensor_name in model_dict + and self.label_name in label_dict + ): + find_global = label_dict.get("find_" + self.label_name, 0.0) + global_weight = self.global_weight * find_global + global_pred = xp.reshape( + model_dict["global_" + self.tensor_name], (-1, self.tensor_size) + ) + global_label = xp.reshape( + label_dict[self.label_name], (-1, self.tensor_size) + ) + diff = global_pred - global_label + 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(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 + ) + + more_loss["rmse"] = xp.sqrt(loss) + return loss, more_loss + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + if self.has_local_weight: + label_requirement.append( + DataRequirementItem( + "atomic_" + self.label_name, + ndof=self.tensor_size, + atomic=True, + must=False, + high_prec=False, + ) + ) + if self.has_global_weight: + label_requirement.append( + DataRequirementItem( + self.label_name, + ndof=self.tensor_size, + atomic=False, + must=False, + high_prec=False, + ) + ) + if self.enable_atomic_weight: + label_requirement.append( + DataRequirementItem( + "atomic_weight", + ndof=1, + atomic=True, + must=False, + high_prec=False, + default=1.0, + ) + ) + return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "TensorLoss", + "@version": 1, + "tensor_name": self.tensor_name, + "tensor_size": self.tensor_size, + "label_name": self.label_name, + "pref_atomic": self.local_weight, + "pref": self.global_weight, + "enable_atomic_weight": self.enable_atomic_weight, + } + + @classmethod + def deserialize(cls, data: dict) -> "TensorLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index bc77f34437..64acae0c05 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -14,6 +14,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.version import ( + check_version_compatibility, +) class DOSLoss(TaskLoss): @@ -264,3 +267,28 @@ def label_requirement(self) -> list[DataRequirementItem]: ) ) return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "DOSLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "numb_dos": self.numb_dos, + "start_pref_dos": self.start_pref_dos, + "limit_pref_dos": self.limit_pref_dos, + "start_pref_cdf": self.start_pref_cdf, + "limit_pref_cdf": self.limit_pref_cdf, + "start_pref_ados": self.start_pref_ados, + "limit_pref_ados": self.limit_pref_ados, + "start_pref_acdf": self.start_pref_acdf, + "limit_pref_acdf": self.limit_pref_acdf, + } + + @classmethod + def deserialize(cls, data: dict) -> "DOSLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 347db5cad8..df9885109d 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -18,6 +18,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.version import ( + check_version_compatibility, +) class EnergySpinLoss(TaskLoss): @@ -405,3 +408,31 @@ def label_requirement(self) -> list[DataRequirementItem]: ) ) return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "EnergySpinLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "start_pref_e": self.start_pref_e, + "limit_pref_e": self.limit_pref_e, + "start_pref_fr": self.start_pref_fr, + "limit_pref_fr": self.limit_pref_fr, + "start_pref_fm": self.start_pref_fm, + "limit_pref_fm": self.limit_pref_fm, + "start_pref_v": self.start_pref_v, + "limit_pref_v": self.limit_pref_v, + "start_pref_ae": self.start_pref_ae, + "limit_pref_ae": self.limit_pref_ae, + "enable_atom_ener_coeff": self.enable_atom_ener_coeff, + "loss_func": self.loss_func, + } + + @classmethod + def deserialize(cls, data: dict) -> "EnergySpinLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 189bcb2a4a..af799d1a89 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -16,6 +16,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.version import ( + check_version_compatibility, +) log = logging.getLogger(__name__) @@ -219,3 +222,26 @@ def label_requirement(self) -> list[DataRequirementItem]: ) ) return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "PropertyLoss", + "@version": 1, + "task_dim": self.task_dim, + "var_name": self.var_name, + "loss_func": self.loss_func, + "metric": self.metric, + "beta": self.beta, + "out_bias": self.out_bias, + "out_std": self.out_std, + "intensive": self.intensive, + } + + @classmethod + def deserialize(cls, data: dict) -> "PropertyLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 625a9b30bc..f329b79b20 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -14,6 +14,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.version import ( + check_version_compatibility, +) class TensorLoss(TaskLoss): @@ -205,3 +208,24 @@ def label_requirement(self) -> list[DataRequirementItem]: ) ) return label_requirement + + def serialize(self) -> dict: + """Serialize the loss module.""" + return { + "@class": "TensorLoss", + "@version": 1, + "tensor_name": self.tensor_name, + "tensor_size": self.tensor_size, + "label_name": self.label_name, + "pref_atomic": self.local_weight, + "pref": self.global_weight, + "enable_atomic_weight": self.enable_atomic_weight, + } + + @classmethod + def deserialize(cls, data: dict) -> "TensorLoss": + """Deserialize the loss module.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) diff --git a/deepmd/pt_expt/loss/__init__.py b/deepmd/pt_expt/loss/__init__.py index 19f76a0cba..77350a8cb0 100644 --- a/deepmd/pt_expt/loss/__init__.py +++ b/deepmd/pt_expt/loss/__init__.py @@ -1,8 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.pt_expt.loss.dos import ( + DOSLoss, +) from deepmd.pt_expt.loss.ener import ( EnergyLoss, ) +from deepmd.pt_expt.loss.ener_spin import ( + EnergySpinLoss, +) +from deepmd.pt_expt.loss.property import ( + PropertyLoss, +) +from deepmd.pt_expt.loss.tensor import ( + TensorLoss, +) __all__ = [ + "DOSLoss", "EnergyLoss", + "EnergySpinLoss", + "PropertyLoss", + "TensorLoss", ] diff --git a/deepmd/pt_expt/loss/dos.py b/deepmd/pt_expt/loss/dos.py new file mode 100644 index 0000000000..ce454204a3 --- /dev/null +++ b/deepmd/pt_expt/loss/dos.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.loss.dos import ( + DOSLoss, +) + +__all__ = ["DOSLoss"] diff --git a/deepmd/pt_expt/loss/ener.py b/deepmd/pt_expt/loss/ener.py index e5bd220bd0..e204b2596e 100644 --- a/deepmd/pt_expt/loss/ener.py +++ b/deepmd/pt_expt/loss/ener.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.loss.ener import EnergyLoss as EnergyLossDP -from deepmd.pt_expt.common import ( - torch_module, +from deepmd.dpmodel.loss.ener import ( + EnergyLoss, ) - -@torch_module -class EnergyLoss(EnergyLossDP): - pass +__all__ = ["EnergyLoss"] diff --git a/deepmd/pt_expt/loss/ener_spin.py b/deepmd/pt_expt/loss/ener_spin.py new file mode 100644 index 0000000000..fd367ff6c9 --- /dev/null +++ b/deepmd/pt_expt/loss/ener_spin.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.loss.ener_spin import ( + EnergySpinLoss, +) + +__all__ = ["EnergySpinLoss"] diff --git a/deepmd/pt_expt/loss/property.py b/deepmd/pt_expt/loss/property.py new file mode 100644 index 0000000000..9118cc8533 --- /dev/null +++ b/deepmd/pt_expt/loss/property.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.loss.property import ( + PropertyLoss, +) + +__all__ = ["PropertyLoss"] diff --git a/deepmd/pt_expt/loss/tensor.py b/deepmd/pt_expt/loss/tensor.py new file mode 100644 index 0000000000..66b93d4c8f --- /dev/null +++ b/deepmd/pt_expt/loss/tensor.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.loss.tensor import ( + TensorLoss, +) + +__all__ = ["TensorLoss"] diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9a2ee0301d..8f32ca660c 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -38,7 +38,11 @@ format_training_message_per_task, ) from deepmd.pt_expt.loss import ( + DOSLoss, EnergyLoss, + EnergySpinLoss, + PropertyLoss, + TensorLoss, ) from deepmd.pt_expt.model import ( get_model, @@ -81,6 +85,32 @@ def get_loss( if loss_type == "ener": loss_params["starter_learning_rate"] = start_lr return EnergyLoss(**loss_params) + elif loss_type == "dos": + loss_params["starter_learning_rate"] = start_lr + loss_params["numb_dos"] = _model.model_output_def()["dos"].output_size + return DOSLoss(**loss_params) + elif loss_type == "ener_spin": + loss_params["starter_learning_rate"] = start_lr + return EnergySpinLoss(**loss_params) + elif loss_type == "tensor": + model_output_type = _model.model_output_type() + if "mask" in model_output_type: + model_output_type.pop(model_output_type.index("mask")) + tensor_name = model_output_type[0] + loss_params["tensor_size"] = _model.model_output_def()[tensor_name].output_size + loss_params["label_name"] = tensor_name + if tensor_name == "polarizability": + tensor_name = "polar" + loss_params["tensor_name"] = tensor_name + return TensorLoss(**loss_params) + elif loss_type == "property": + task_dim = _model.get_task_dim() + var_name = _model.get_var_name() + intensive = _model.get_intensive() + loss_params["task_dim"] = task_dim + loss_params["var_name"] = var_name + loss_params["intensive"] = intensive + return PropertyLoss(**loss_params) else: raise ValueError(f"Unsupported loss type for pt_expt: {loss_type}") diff --git a/source/tests/consistent/loss/test_dos.py b/source/tests/consistent/loss/test_dos.py new file mode 100644 index 0000000000..8ed91873ec --- /dev/null +++ b/source/tests/consistent/loss/test_dos.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.loss.dos import DOSLoss as DOSLossDP +from deepmd.utils.argcheck import ( + loss_dos, +) + +from ..common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT, + INSTALLED_PT_EXPT, + CommonTest, + parameterized, +) +from .common import ( + LossTest, +) + +if INSTALLED_PT: + from deepmd.pt.loss.dos import DOSLoss as DOSLossPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch +else: + DOSLossPT = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.loss.dos import DOSLoss as DOSLossPTExpt +else: + DOSLossPTExpt = None +if INSTALLED_JAX: + from deepmd.jax.env import ( + jnp, + ) +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict + + +@parameterized( + (1.0, 0.0), # pref_dos + (1.0, 0.0), # pref_ados +) +class TestDOS(CommonTest, LossTest, unittest.TestCase): + @property + def data(self) -> dict: + (pref_dos, pref_ados) = self.param + return { + "start_pref_dos": pref_dos, + "limit_pref_dos": pref_dos / 2 if pref_dos else 0.0, + "start_pref_cdf": 0.0, + "limit_pref_cdf": 0.0, + "start_pref_ados": pref_ados, + "limit_pref_ados": pref_ados / 2 if pref_ados else 0.0, + "start_pref_acdf": 0.0, + "limit_pref_acdf": 0.0, + } + + skip_tf = True + skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT + skip_jax = not INSTALLED_JAX + skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_pd = True + + dp_class = DOSLossDP + pt_class = DOSLossPT + pt_expt_class = DOSLossPTExpt + jax_class = DOSLossDP + array_api_strict_class = DOSLossDP + args = loss_dos() + + def setUp(self) -> None: + (pref_dos, pref_ados) = self.param + if pref_dos == 0.0 and pref_ados == 0.0: + self.skipTest("Both pref_dos and pref_ados are 0") + CommonTest.setUp(self) + self.learning_rate = 1e-3 + rng = np.random.default_rng(20250326) + self.nframes = 2 + self.natoms = 6 + self.numb_dos = 4 + self.predict = { + "dos": rng.random((self.nframes, self.numb_dos)), + "atom_dos": rng.random((self.nframes, self.natoms, self.numb_dos)), + } + self.label = { + "dos": rng.random((self.nframes, self.numb_dos)), + "atom_dos": rng.random((self.nframes, self.natoms, self.numb_dos)), + "find_dos": 1.0, + "find_atom_dos": 1.0, + } + + @property + def additional_data(self) -> dict: + return { + "starter_learning_rate": 1e-3, + "numb_dos": self.numb_dos, + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError + + def eval_pt(self, pt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + _, loss, more_loss = pt_obj( + {}, + lambda: predict, + label, + self.natoms, + self.learning_rate, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.learning_rate, + self.natoms, + self.predict, + self.label, + ) + + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_jax(self, jax_obj: Any) -> Any: + predict = {kk: jnp.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = jax_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: + predict = {kk: array_api_strict.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: array_api_strict.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = array_api_strict_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: + loss = ret[0] + result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} + if len(ret) > 1: + more_loss = ret[1] + for k in sorted(more_loss): + if k.startswith("rmse_") or k.startswith("mae_"): + result[k] = np.atleast_1d( + np.asarray(more_loss[k], dtype=np.float64) + ) + return result + + @property + def rtol(self) -> float: + return 1e-10 + + @property + def atol(self) -> float: + return 1e-10 diff --git a/source/tests/consistent/loss/test_ener.py b/source/tests/consistent/loss/test_ener.py index 9885177287..dcb3988173 100644 --- a/source/tests/consistent/loss/test_ener.py +++ b/source/tests/consistent/loss/test_ener.py @@ -19,6 +19,7 @@ INSTALLED_JAX, INSTALLED_PD, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, CommonTest, parameterized, @@ -52,6 +53,10 @@ from deepmd.jax.env import ( jnp, ) +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.loss.ener import EnergyLoss as EnerLossPTExpt +else: + EnerLossPTExpt = None if INSTALLED_ARRAY_API_STRICT: import array_api_strict @@ -61,11 +66,12 @@ (False, True), # enable_atom_ener_coeff ("mse", "mae"), # loss_func (False, True), # f_use_norm + (False, True), # mae (dp test extra MAE metrics) ) class TestEner(CommonTest, LossTest, unittest.TestCase): @property def data(self) -> dict: - (use_huber, enable_atom_ener_coeff, loss_func, f_use_norm) = self.param + (use_huber, enable_atom_ener_coeff, loss_func, f_use_norm, _mae) = self.param return { "start_pref_e": 0.02, "limit_pref_e": 1.0, @@ -85,36 +91,41 @@ def data(self) -> dict: @property def skip_tf(self) -> bool: - (use_huber, enable_atom_ener_coeff, loss_func, f_use_norm) = self.param + (_use_huber, _enable_atom_ener_coeff, loss_func, f_use_norm, _mae) = self.param # Skip TF for MAE loss tests (not implemented in TF backend) return CommonTest.skip_tf or loss_func == "mae" or f_use_norm @property def skip_pd(self) -> bool: - (use_huber, enable_atom_ener_coeff, loss_func, f_use_norm) = self.param + (_use_huber, _enable_atom_ener_coeff, loss_func, f_use_norm, _mae) = self.param # Skip Paddle for MAE loss tests (not implemented in Paddle backend) return not INSTALLED_PD or loss_func == "mae" or f_use_norm skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT tf_class = EnerLossTF dp_class = EnerLossDP pt_class = EnerLossPT + pt_expt_class = EnerLossPTExpt jax_class = EnerLossDP pd_class = EnerLossPD array_api_strict_class = EnerLossDP args = loss_ener() def setUp(self) -> None: - (use_huber, enable_atom_ener_coeff, loss_func, f_use_norm) = self.param + (use_huber, _enable_atom_ener_coeff, loss_func, f_use_norm, mae) = self.param # Skip invalid combinations if f_use_norm and not (use_huber or loss_func == "mae"): self.skipTest("f_use_norm requires either use_huber or loss_func='mae'") if use_huber and loss_func == "mae": self.skipTest("Cannot use both huber and mae loss_func at the same time") + if loss_func == "mae" and mae: + self.skipTest("mae=True with loss_func='mae' is redundant") CommonTest.setUp(self) + self.mae = mae self.learning_rate = 1e-3 rng = np.random.default_rng(20250105) self.nframes = 2 @@ -178,7 +189,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: for kk, vv in self.label.items() } - loss, more_loss = obj.build( + loss, _more_loss = obj.build( self.learning_rate, [self.natoms], predict, @@ -208,6 +219,7 @@ def eval_pt(self, pt_obj: Any) -> Any: label, self.natoms, self.learning_rate, + mae=self.mae, ) loss = torch_to_numpy(loss) more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} @@ -219,8 +231,25 @@ def eval_dp(self, dp_obj: Any) -> Any: self.natoms, self.predict_dpmodel_style, self.label, + mae=self.mae, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = { + kk: numpy_to_torch(vv) for kk, vv in self.predict_dpmodel_style.items() + } + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=self.mae, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + def eval_jax(self, jax_obj: Any) -> Any: predict = {kk: jnp.asarray(vv) for kk, vv in self.predict_dpmodel_style.items()} label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} @@ -230,6 +259,7 @@ def eval_jax(self, jax_obj: Any) -> Any: self.natoms, predict, label, + mae=self.mae, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -247,6 +277,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: self.natoms, predict, label, + mae=self.mae, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -266,6 +297,7 @@ def eval_pd(self, pd_obj: Any) -> Any: label, self.natoms, self.learning_rate, + mae=self.mae, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -320,6 +352,7 @@ def data(self) -> dict: skip_tf = CommonTest.skip_tf skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = not INSTALLED_PD @@ -327,6 +360,7 @@ def data(self) -> dict: tf_class = EnerLossTF dp_class = EnerLossDP pt_class = EnerLossPT + pt_expt_class = EnerLossPTExpt jax_class = EnerLossDP pd_class = EnerLossPD array_api_strict_class = EnerLossDP @@ -390,7 +424,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: for kk, vv in self.label.items() } - loss, more_loss = obj.build( + loss, _more_loss = obj.build( self.learning_rate, [self.natoms], predict, @@ -420,6 +454,7 @@ def eval_pt(self, pt_obj: Any) -> Any: label, self.natoms, self.learning_rate, + mae=True, ) loss = torch_to_numpy(loss) more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} @@ -431,8 +466,25 @@ def eval_dp(self, dp_obj: Any) -> Any: self.natoms, self.predict_dpmodel_style, self.label, + mae=True, ) + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = { + kk: numpy_to_torch(vv) for kk, vv in self.predict_dpmodel_style.items() + } + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=True, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + def eval_jax(self, jax_obj: Any) -> Any: predict = {kk: jnp.asarray(vv) for kk, vv in self.predict_dpmodel_style.items()} label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} @@ -442,6 +494,7 @@ def eval_jax(self, jax_obj: Any) -> Any: self.natoms, predict, label, + mae=True, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -459,6 +512,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: self.natoms, predict, label, + mae=True, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -478,6 +532,7 @@ def eval_pd(self, pd_obj: Any) -> Any: label, self.natoms, self.learning_rate, + mae=True, ) loss = to_numpy_array(loss) more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} @@ -489,7 +544,7 @@ def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: if len(ret) > 1: more_loss = ret[1] for k in sorted(more_loss): - if k.startswith("rmse_"): + if k.startswith("rmse_") or k.startswith("mae_"): result[k] = np.atleast_1d( np.asarray(more_loss[k], dtype=np.float64) ) diff --git a/source/tests/consistent/loss/test_ener_spin.py b/source/tests/consistent/loss/test_ener_spin.py new file mode 100644 index 0000000000..2e8734c109 --- /dev/null +++ b/source/tests/consistent/loss/test_ener_spin.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.loss.ener_spin import EnergySpinLoss as EnerSpinLossDP +from deepmd.utils.argcheck import ( + loss_ener_spin, +) + +from ..common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT, + INSTALLED_PT_EXPT, + CommonTest, + parameterized, +) +from .common import ( + LossTest, +) + +if INSTALLED_PT: + from deepmd.pt.loss.ener_spin import EnergySpinLoss as EnerSpinLossPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch +else: + EnerSpinLossPT = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.loss.ener_spin import EnergySpinLoss as EnerSpinLossPTExpt +else: + EnerSpinLossPTExpt = None +if INSTALLED_JAX: + from deepmd.jax.env import ( + jnp, + ) +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict + + +@parameterized( + ("mse", "mae"), # loss_func + (False, True), # mae (dp test extra MAE metrics) +) +class TestEnerSpin(CommonTest, LossTest, unittest.TestCase): + @property + def data(self) -> dict: + (loss_func, _mae) = self.param + return { + "start_pref_e": 0.02, + "limit_pref_e": 1.0, + "start_pref_fr": 1000.0, + "limit_pref_fr": 1.0, + "start_pref_fm": 1000.0, + "limit_pref_fm": 1.0, + "start_pref_v": 1.0, + "limit_pref_v": 1.0, + "start_pref_ae": 1.0, + "limit_pref_ae": 1.0, + "loss_func": loss_func, + } + + skip_tf = True + skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT + skip_jax = not INSTALLED_JAX + skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_pd = True + + dp_class = EnerSpinLossDP + pt_class = EnerSpinLossPT + pt_expt_class = EnerSpinLossPTExpt + jax_class = EnerSpinLossDP + array_api_strict_class = EnerSpinLossDP + args = loss_ener_spin() + + def setUp(self) -> None: + (loss_func, mae) = self.param + if loss_func == "mae" and mae: + self.skipTest("mae=True with loss_func='mae' is redundant") + CommonTest.setUp(self) + self.mae = mae + self.learning_rate = 1e-3 + rng = np.random.default_rng(20250326) + self.nframes = 2 + self.natoms = 6 + n_magnetic = 4 + mask_mag = np.zeros((self.nframes, self.natoms, 1), dtype=bool) + mask_mag[:, :n_magnetic, :] = True + self.predict = { + "energy": rng.random((self.nframes,)), + "force": rng.random((self.nframes, self.natoms, 3)), + "force_mag": rng.random((self.nframes, self.natoms, 3)), + "mask_mag": mask_mag, + "virial": rng.random((self.nframes, 9)), + "atom_energy": rng.random((self.nframes, self.natoms)), + } + self.label = { + "energy": rng.random((self.nframes,)), + "force": rng.random((self.nframes, self.natoms, 3)), + "force_mag": rng.random((self.nframes, self.natoms, 3)), + "virial": rng.random((self.nframes, 9)), + "atom_ener": rng.random((self.nframes, self.natoms)), + "find_energy": 1.0, + "find_force": 1.0, + "find_force_mag": 1.0, + "find_virial": 1.0, + "find_atom_ener": 1.0, + } + + @property + def additional_data(self) -> dict: + return { + "starter_learning_rate": 1e-3, + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError + + def eval_pt(self, pt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + _, loss, more_loss = pt_obj( + {}, + lambda: predict, + label, + self.natoms, + self.learning_rate, + mae=self.mae, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.learning_rate, + self.natoms, + self.predict, + self.label, + mae=self.mae, + ) + + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=self.mae, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_jax(self, jax_obj: Any) -> Any: + predict = {kk: jnp.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = jax_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=self.mae, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: + predict = {kk: array_api_strict.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: array_api_strict.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = array_api_strict_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=self.mae, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: + loss = ret[0] + result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} + if len(ret) > 1: + more_loss = ret[1] + for k in sorted(more_loss): + if k.startswith("rmse_") or k.startswith("mae_"): + result[k] = np.atleast_1d( + np.asarray(more_loss[k], dtype=np.float64) + ) + return result + + @property + def rtol(self) -> float: + return 1e-10 + + @property + def atol(self) -> float: + return 1e-10 diff --git a/source/tests/consistent/loss/test_property.py b/source/tests/consistent/loss/test_property.py new file mode 100644 index 0000000000..7750eb6dae --- /dev/null +++ b/source/tests/consistent/loss/test_property.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.loss.property import PropertyLoss as PropertyLossDP +from deepmd.utils.argcheck import ( + loss_property, +) + +from ..common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT, + INSTALLED_PT_EXPT, + CommonTest, + parameterized, +) +from .common import ( + LossTest, +) + +if INSTALLED_PT: + from deepmd.pt.loss.property import PropertyLoss as PropertyLossPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch +else: + PropertyLossPT = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.loss.property import PropertyLoss as PropertyLossPTExpt +else: + PropertyLossPTExpt = None +if INSTALLED_JAX: + from deepmd.jax.env import ( + jnp, + ) +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict + + +@parameterized( + ("smooth_mae", "mae", "mse", "rmse", "mape"), # loss_func +) +class TestProperty(CommonTest, LossTest, unittest.TestCase): + @property + def data(self) -> dict: + (loss_func,) = self.param + return { + "loss_func": loss_func, + } + + skip_tf = True + skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT + skip_jax = not INSTALLED_JAX + skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_pd = True + + dp_class = PropertyLossDP + pt_class = PropertyLossPT + pt_expt_class = PropertyLossPTExpt + jax_class = PropertyLossDP + array_api_strict_class = PropertyLossDP + args = loss_property() + + def setUp(self) -> None: + CommonTest.setUp(self) + self.learning_rate = 1e-3 + rng = np.random.default_rng(20250326) + self.nframes = 2 + self.natoms = 6 + self.task_dim = 5 + self.var_name = "foo" + self.predict = { + self.var_name: rng.random((self.nframes, self.task_dim)), + } + self.label = { + self.var_name: rng.random((self.nframes, self.task_dim)), + } + + @property + def additional_data(self) -> dict: + return { + "task_dim": self.task_dim, + "var_name": self.var_name, + "out_bias": [0.1, 0.5, 1.2, -0.1, -10.0], + "out_std": [8.0, 10.0, 0.001, -0.2, -10.0], + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError + + def eval_pt(self, pt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + _, loss, more_loss = pt_obj( + {}, + lambda: predict, + label, + self.natoms, + self.learning_rate, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.learning_rate, + self.natoms, + self.predict, + self.label, + ) + + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_jax(self, jax_obj: Any) -> Any: + predict = {kk: jnp.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = jax_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: + predict = {kk: array_api_strict.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: array_api_strict.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = array_api_strict_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: + loss = ret[0] + result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} + if len(ret) > 1: + more_loss = ret[1] + for k in sorted(more_loss): + if ( + k.startswith("rmse") + or k.startswith("mae") + or k.startswith("mse") + or k.startswith("mape") + or k.startswith("smooth_mae") + ): + result[k] = np.atleast_1d( + np.asarray(more_loss[k], dtype=np.float64) + ) + return result + + @property + def rtol(self) -> float: + return 1e-10 + + @property + def atol(self) -> float: + return 1e-10 diff --git a/source/tests/consistent/loss/test_tensor.py b/source/tests/consistent/loss/test_tensor.py new file mode 100644 index 0000000000..06feb908fa --- /dev/null +++ b/source/tests/consistent/loss/test_tensor.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.loss.tensor import TensorLoss as TensorLossDP +from deepmd.utils.argcheck import ( + loss_tensor, +) + +from ..common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT, + INSTALLED_PT_EXPT, + CommonTest, + parameterized, +) +from .common import ( + LossTest, +) + +if INSTALLED_PT: + from deepmd.pt.loss.tensor import TensorLoss as TensorLossPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch +else: + TensorLossPT = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.loss.tensor import TensorLoss as TensorLossPTExpt +else: + TensorLossPTExpt = None +if INSTALLED_JAX: + from deepmd.jax.env import ( + jnp, + ) +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict + + +@parameterized( + (1.0, 0.0), # pref + (1.0, 0.0), # pref_atomic +) +class TestTensor(CommonTest, LossTest, unittest.TestCase): + @property + def data(self) -> dict: + (pref, pref_atomic) = self.param + return { + "pref": pref, + "pref_atomic": pref_atomic, + } + + skip_tf = True + skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT + skip_jax = not INSTALLED_JAX + skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_pd = True + + dp_class = TensorLossDP + pt_class = TensorLossPT + pt_expt_class = TensorLossPTExpt + jax_class = TensorLossDP + array_api_strict_class = TensorLossDP + args = loss_tensor() + + def setUp(self) -> None: + (pref, pref_atomic) = self.param + if pref == 0.0 and pref_atomic == 0.0: + self.skipTest("Both pref and pref_atomic are 0") + CommonTest.setUp(self) + self.learning_rate = 1e-3 + rng = np.random.default_rng(20250326) + self.nframes = 2 + self.natoms = 6 + self.tensor_name = "test_tensor" + self.label_name = "test_tensor" + self.tensor_size = 3 + self.predict = { + self.tensor_name: rng.random((self.nframes, self.natoms, self.tensor_size)), + "global_" + self.tensor_name: rng.random((self.nframes, self.tensor_size)), + } + self.label = { + "atom_" + self.label_name: rng.random( + (self.nframes, self.natoms, self.tensor_size) + ), + self.label_name: rng.random((self.nframes, self.tensor_size)), + "find_atom_" + self.label_name: 1.0, + "find_" + self.label_name: 1.0, + } + + @property + def additional_data(self) -> dict: + return { + "tensor_name": self.tensor_name, + "tensor_size": self.tensor_size, + "label_name": self.label_name, + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + raise NotImplementedError + + def eval_pt(self, pt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + _, loss, more_loss = pt_obj( + {}, + lambda: predict, + label, + self.natoms, + self.learning_rate, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.learning_rate, + self.natoms, + self.predict, + self.label, + ) + + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_jax(self, jax_obj: Any) -> Any: + predict = {kk: jnp.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = jax_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: + predict = {kk: array_api_strict.asarray(vv) for kk, vv in self.predict.items()} + label = {kk: array_api_strict.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = array_api_strict_obj( + self.learning_rate, + self.natoms, + predict, + label, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: + loss = ret[0] + result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} + if len(ret) > 1: + more_loss = ret[1] + for k in sorted(more_loss): + if k.startswith("rmse_") or k.startswith("mae_"): + result[k] = np.atleast_1d( + np.asarray(more_loss[k], dtype=np.float64) + ) + return result + + @property + def rtol(self) -> float: + return 1e-10 + + @property + def atol(self) -> float: + return 1e-10 diff --git a/source/tests/pt_expt/loss/test_dos.py b/source/tests/pt_expt/loss/test_dos.py new file mode 100644 index 0000000000..c1f7f51dbf --- /dev/null +++ b/source/tests/pt_expt/loss/test_dos.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the pt_expt DOSLoss wrapper.""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.loss.dos import DOSLoss as DOSLossDP +from deepmd.pt_expt.loss.dos import ( + DOSLoss, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.env import ( + PRECISION_DICT, +) + +from ...pt.model.test_mlp import ( + get_tols, +) +from ...seed import ( + GLOBAL_SEED, +) + + +def _make_data( + rng: np.random.Generator, + nframes: int, + natoms: int, + numb_dos: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Build model prediction and label dicts as torch tensors.""" + model_pred = { + "dos": torch.tensor( + rng.random((nframes, numb_dos)), dtype=dtype, device=device + ), + "atom_dos": torch.tensor( + rng.random((nframes, natoms, numb_dos)), dtype=dtype, device=device + ), + } + label = { + "dos": torch.tensor( + rng.random((nframes, numb_dos)), dtype=dtype, device=device + ), + "atom_dos": torch.tensor( + rng.random((nframes, natoms, numb_dos)), dtype=dtype, device=device + ), + "find_dos": torch.tensor(1.0, dtype=dtype, device=device), + "find_atom_dos": torch.tensor(1.0, dtype=dtype, device=device), + } + return model_pred, label + + +class TestDOSLoss: + def setup_method(self) -> None: + self.device = env.DEVICE + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + @pytest.mark.parametrize( + "has_dos,has_ados", + [(True, True), (True, False), (False, True)], + ) # which loss terms are active + def test_consistency(self, prec, has_dos, has_ados) -> None: + """Construct -> forward -> serialize/deserialize -> forward -> compare. + + Also compare with dpmodel. + """ + rng = np.random.default_rng(GLOBAL_SEED) + nframes, natoms, numb_dos = 2, 6, 4 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = DOSLoss( + starter_learning_rate=1e-3, + numb_dos=numb_dos, + start_pref_dos=1.0 if has_dos else 0.0, + limit_pref_dos=0.5 if has_dos else 0.0, + start_pref_cdf=0.0, + limit_pref_cdf=0.0, + start_pref_ados=1.0 if has_ados else 0.0, + limit_pref_ados=0.5 if has_ados else 0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + model_pred, label = _make_data( + rng, nframes, natoms, numb_dos, dtype, self.device + ) + + # Forward + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + assert "rmse" in more0 + + # Serialize / deserialize round-trip + loss1 = DOSLoss.deserialize(loss0.serialize()) + l1, more1 = loss1(learning_rate, natoms, model_pred, label) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + l1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + for key in more0: + np.testing.assert_allclose( + more0[key].detach().cpu().numpy(), + more1[key].detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=f"key={key}", + ) + + # Compare with dpmodel (numpy) + dp_loss = DOSLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_cdf_terms(self, prec) -> None: + """Test with CDF loss terms enabled.""" + rng = np.random.default_rng(GLOBAL_SEED + 1) + nframes, natoms, numb_dos = 2, 6, 4 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = DOSLoss( + starter_learning_rate=1e-3, + numb_dos=numb_dos, + start_pref_dos=0.0, + limit_pref_dos=0.0, + start_pref_cdf=1.0, + limit_pref_cdf=0.5, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=1.0, + limit_pref_acdf=0.5, + ) + + model_pred, label = _make_data( + rng, nframes, natoms, numb_dos, dtype, self.device + ) + + l0, _more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + + # Compare with dpmodel + dp_loss = DOSLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (cdf terms)", + ) diff --git a/source/tests/pt_expt/loss/test_ener.py b/source/tests/pt_expt/loss/test_ener.py index 37d7d4c703..68d3abbc8f 100644 --- a/source/tests/pt_expt/loss/test_ener.py +++ b/source/tests/pt_expt/loss/test_ener.py @@ -98,7 +98,7 @@ def test_consistency(self, prec, use_huber) -> None: start_pref_pf=0.0 if use_huber else 1.0, limit_pref_pf=0.0 if use_huber else 1.0, use_huber=use_huber, - ).to(self.device) + ) model_pred, label = _make_data(rng, nframes, natoms, dtype, self.device) @@ -133,7 +133,7 @@ def test_consistency(self, prec, use_huber) -> None: k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v for k, v in label.items() } - l_dp, more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) + l_dp, _more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) np.testing.assert_allclose( l0.detach().cpu().numpy(), diff --git a/source/tests/pt_expt/loss/test_ener_spin.py b/source/tests/pt_expt/loss/test_ener_spin.py new file mode 100644 index 0000000000..c39f4d166b --- /dev/null +++ b/source/tests/pt_expt/loss/test_ener_spin.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the pt_expt EnergySpinLoss wrapper.""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.loss.ener_spin import EnergySpinLoss as EnergySpinLossDP +from deepmd.pt_expt.loss.ener_spin import ( + EnergySpinLoss, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.env import ( + PRECISION_DICT, +) + +from ...pt.model.test_mlp import ( + get_tols, +) +from ...seed import ( + GLOBAL_SEED, +) + + +def _make_data( + rng: np.random.Generator, + nframes: int, + natoms: int, + n_magnetic: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Build model prediction and label dicts as torch tensors.""" + # mask_mag: True for magnetic atoms, False otherwise + mask_mag = torch.zeros((nframes, natoms, 1), dtype=torch.bool, device=device) + mask_mag[:, :n_magnetic, :] = True + model_pred = { + "energy": torch.tensor(rng.random((nframes,)), dtype=dtype, device=device), + "force": torch.tensor( + rng.random((nframes, natoms, 3)), dtype=dtype, device=device + ), + "force_mag": torch.tensor( + rng.random((nframes, natoms, 3)), dtype=dtype, device=device + ), + "mask_mag": mask_mag, + "virial": torch.tensor(rng.random((nframes, 9)), dtype=dtype, device=device), + "atom_energy": torch.tensor( + rng.random((nframes, natoms)), dtype=dtype, device=device + ), + } + label = { + "energy": torch.tensor(rng.random((nframes,)), dtype=dtype, device=device), + "force": torch.tensor( + rng.random((nframes, natoms, 3)), dtype=dtype, device=device + ), + "force_mag": torch.tensor( + rng.random((nframes, natoms, 3)), dtype=dtype, device=device + ), + "virial": torch.tensor(rng.random((nframes, 9)), dtype=dtype, device=device), + "atom_ener": torch.tensor( + rng.random((nframes, natoms)), dtype=dtype, device=device + ), + "find_energy": torch.tensor(1.0, dtype=dtype, device=device), + "find_force": torch.tensor(1.0, dtype=dtype, device=device), + "find_force_mag": torch.tensor(1.0, dtype=dtype, device=device), + "find_virial": torch.tensor(1.0, dtype=dtype, device=device), + "find_atom_ener": torch.tensor(1.0, dtype=dtype, device=device), + } + return model_pred, label + + +class TestEnergySpinLoss: + def setup_method(self) -> None: + self.device = env.DEVICE + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + @pytest.mark.parametrize("loss_func", ["mse", "mae"]) # loss function + def test_consistency(self, prec, loss_func) -> None: + """Construct -> forward -> serialize/deserialize -> forward -> compare. + + Also compare with dpmodel. + """ + rng = np.random.default_rng(GLOBAL_SEED) + nframes, natoms, n_magnetic = 2, 6, 4 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + if prec in ["single", "float32"]: + atol = max(atol, 2e-4) # relax for float32 rounding across envs + learning_rate = 1e-3 + + loss0 = EnergySpinLoss( + starter_learning_rate=1e-3, + start_pref_e=0.02, + limit_pref_e=1.0, + start_pref_fr=1000.0, + limit_pref_fr=1.0, + start_pref_fm=1000.0, + limit_pref_fm=1.0, + start_pref_v=1.0, + limit_pref_v=1.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + model_pred, label = _make_data( + rng, nframes, natoms, n_magnetic, dtype, self.device + ) + + # Forward + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + assert "rmse" in more0 + + # Serialize / deserialize round-trip + loss1 = EnergySpinLoss.deserialize(loss0.serialize()) + l1, more1 = loss1(learning_rate, natoms, model_pred, label) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + l1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + for key in more0: + np.testing.assert_allclose( + more0[key].detach().cpu().numpy(), + more1[key].detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=f"key={key}", + ) + + # Compare with dpmodel (numpy) + dp_loss = EnergySpinLossDP.deserialize(loss0.serialize()) + model_pred_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in model_pred.items() + } + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_partial_mask(self, prec) -> None: + """Test with partial magnetic atoms (some atoms non-magnetic).""" + rng = np.random.default_rng(GLOBAL_SEED + 1) + nframes, natoms, n_magnetic = 2, 6, 2 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = EnergySpinLoss( + starter_learning_rate=1e-3, + start_pref_e=0.02, + limit_pref_e=1.0, + start_pref_fr=1000.0, + limit_pref_fr=1.0, + start_pref_fm=1000.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + ) + + model_pred, label = _make_data( + rng, nframes, natoms, n_magnetic, dtype, self.device + ) + + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + assert "rmse_fm" in more0 + + # Compare with dpmodel + dp_loss = EnergySpinLossDP.deserialize(loss0.serialize()) + model_pred_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in model_pred.items() + } + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (partial mask)", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_all_masked(self, prec) -> None: + """Test with all atoms magnetic.""" + rng = np.random.default_rng(GLOBAL_SEED + 2) + nframes, natoms, n_magnetic = 2, 6, 6 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = EnergySpinLoss( + starter_learning_rate=1e-3, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1000.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + ) + + model_pred, label = _make_data( + rng, nframes, natoms, n_magnetic, dtype, self.device + ) + + l0, _more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + + # Compare with dpmodel + dp_loss = EnergySpinLossDP.deserialize(loss0.serialize()) + model_pred_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in model_pred.items() + } + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (all masked)", + ) diff --git a/source/tests/pt_expt/loss/test_property.py b/source/tests/pt_expt/loss/test_property.py new file mode 100644 index 0000000000..ee47a24b92 --- /dev/null +++ b/source/tests/pt_expt/loss/test_property.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the pt_expt PropertyLoss wrapper.""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.loss.property import PropertyLoss as PropertyLossDP +from deepmd.pt_expt.loss.property import ( + PropertyLoss, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.env import ( + PRECISION_DICT, +) + +from ...pt.model.test_mlp import ( + get_tols, +) +from ...seed import ( + GLOBAL_SEED, +) + + +def _make_data( + rng: np.random.Generator, + nframes: int, + task_dim: int, + var_name: str, + dtype: torch.dtype, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Build model prediction and label dicts as torch tensors.""" + model_pred = { + var_name: torch.tensor( + rng.random((nframes, task_dim)), dtype=dtype, device=device + ), + } + label = { + var_name: torch.tensor( + rng.random((nframes, task_dim)), dtype=dtype, device=device + ), + } + return model_pred, label + + +class TestPropertyLoss: + def setup_method(self) -> None: + self.device = env.DEVICE + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + @pytest.mark.parametrize( + "loss_func", ["smooth_mae", "mae", "mse", "rmse", "mape"] + ) # loss function + def test_consistency(self, prec, loss_func) -> None: + """Construct -> forward -> serialize/deserialize -> forward -> compare. + + Also compare with dpmodel. + """ + rng = np.random.default_rng(GLOBAL_SEED) + nframes = 2 + task_dim = 5 + var_name = "foo" + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + natoms = 6 + + loss0 = PropertyLoss( + task_dim=task_dim, + var_name=var_name, + loss_func=loss_func, + metric=["mae"], + beta=1.0, + out_bias=[0.1, 0.5, 1.2, -0.1, -10.0], + out_std=[8.0, 10.0, 0.001, -0.2, -10.0], + intensive=False, + ) + + model_pred, label = _make_data( + rng, nframes, task_dim, var_name, dtype, self.device + ) + + # Forward + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + + # Serialize / deserialize round-trip + loss1 = PropertyLoss.deserialize(loss0.serialize()) + l1, more1 = loss1(learning_rate, natoms, model_pred, label) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + l1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + for key in more0: + np.testing.assert_allclose( + more0[key].detach().cpu().numpy(), + more1[key].detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=f"key={key}", + ) + + # Compare with dpmodel (numpy) + dp_loss = PropertyLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = {k: v.detach().cpu().numpy() for k, v in label.items()} + l_dp, _more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + # Use relative tolerance: extreme out_std values (e.g. 0.001) can produce + # large loss values where torch/numpy accumulation order differs at machine epsilon. + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=max(rtol, 1e-14 if prec == "float64" else 1e-5), + atol=atol, + err_msg="pt_expt vs dpmodel", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_intensive(self, prec) -> None: + """Test intensive property (no division by natoms).""" + rng = np.random.default_rng(GLOBAL_SEED + 1) + nframes = 2 + task_dim = 3 + var_name = "band_gap" + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + natoms = 6 + + loss0 = PropertyLoss( + task_dim=task_dim, + var_name=var_name, + loss_func="mse", + metric=["mae", "rmse"], + beta=1.0, + out_bias=None, + out_std=None, + intensive=True, + ) + + model_pred, label = _make_data( + rng, nframes, task_dim, var_name, dtype, self.device + ) + + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + assert "mae" in more0 + assert "rmse" in more0 + + # Compare with dpmodel + dp_loss = PropertyLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = {k: v.detach().cpu().numpy() for k, v in label.items()} + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (intensive)", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_no_out_bias_std(self, prec) -> None: + """Test with out_bias and out_std as None (identity normalization).""" + rng = np.random.default_rng(GLOBAL_SEED + 2) + nframes = 2 + task_dim = 3 + var_name = "prop" + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + natoms = 6 + + loss0 = PropertyLoss( + task_dim=task_dim, + var_name=var_name, + loss_func="mae", + out_bias=None, + out_std=None, + intensive=False, + ) + + model_pred, label = _make_data( + rng, nframes, task_dim, var_name, dtype, self.device + ) + + l0, _ = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + + # Compare with dpmodel + dp_loss = PropertyLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = {k: v.detach().cpu().numpy() for k, v in label.items()} + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (no bias/std)", + ) diff --git a/source/tests/pt_expt/loss/test_tensor.py b/source/tests/pt_expt/loss/test_tensor.py new file mode 100644 index 0000000000..c9eae2d780 --- /dev/null +++ b/source/tests/pt_expt/loss/test_tensor.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the pt_expt TensorLoss wrapper.""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.loss.tensor import TensorLoss as TensorLossDP +from deepmd.pt_expt.loss.tensor import ( + TensorLoss, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.env import ( + PRECISION_DICT, +) + +from ...pt.model.test_mlp import ( + get_tols, +) +from ...seed import ( + GLOBAL_SEED, +) + + +def _make_data( + rng: np.random.Generator, + nframes: int, + natoms: int, + tensor_name: str, + label_name: str, + tensor_size: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Build model prediction and label dicts as torch tensors.""" + model_pred = { + tensor_name: torch.tensor( + rng.random((nframes, natoms, tensor_size)), dtype=dtype, device=device + ), + "global_" + tensor_name: torch.tensor( + rng.random((nframes, tensor_size)), dtype=dtype, device=device + ), + } + label = { + "atom_" + label_name: torch.tensor( + rng.random((nframes, natoms, tensor_size)), dtype=dtype, device=device + ), + label_name: torch.tensor( + rng.random((nframes, tensor_size)), dtype=dtype, device=device + ), + "find_atom_" + label_name: torch.tensor(1.0, dtype=dtype, device=device), + "find_" + label_name: torch.tensor(1.0, dtype=dtype, device=device), + } + return model_pred, label + + +class TestTensorLoss: + def setup_method(self) -> None: + self.device = env.DEVICE + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + @pytest.mark.parametrize( + "has_local,has_global", + [(True, True), (True, False), (False, True)], + ) # which loss terms are active + def test_consistency(self, prec, has_local, has_global) -> None: + """Construct -> forward -> serialize/deserialize -> forward -> compare. + + Also compare with dpmodel. + """ + rng = np.random.default_rng(GLOBAL_SEED) + nframes, natoms = 2, 6 + tensor_name = "test_tensor" + label_name = "test_tensor" + tensor_size = 3 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = TensorLoss( + tensor_name=tensor_name, + tensor_size=tensor_size, + label_name=label_name, + pref_atomic=1.0 if has_local else 0.0, + pref=1.0 if has_global else 0.0, + ) + + model_pred, label = _make_data( + rng, + nframes, + natoms, + tensor_name, + label_name, + tensor_size, + dtype, + self.device, + ) + + # Forward + l0, more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + assert "rmse" in more0 + + # Serialize / deserialize round-trip + loss1 = TensorLoss.deserialize(loss0.serialize()) + l1, more1 = loss1(learning_rate, natoms, model_pred, label) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + l1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + for key in more0: + np.testing.assert_allclose( + more0[key].detach().cpu().numpy(), + more1[key].detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=f"key={key}", + ) + + # Compare with dpmodel (numpy) + dp_loss = TensorLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _more_dp = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel", + ) + + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision + def test_with_atomic_weight(self, prec) -> None: + """Test with atomic weight enabled.""" + rng = np.random.default_rng(GLOBAL_SEED + 1) + nframes, natoms = 2, 6 + tensor_name = "dipole" + label_name = "dipole" + tensor_size = 3 + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + learning_rate = 1e-3 + + loss0 = TensorLoss( + tensor_name=tensor_name, + tensor_size=tensor_size, + label_name=label_name, + pref_atomic=1.0, + pref=1.0, + enable_atomic_weight=True, + ) + + model_pred, label = _make_data( + rng, + nframes, + natoms, + tensor_name, + label_name, + tensor_size, + dtype, + self.device, + ) + label["atom_weight"] = torch.tensor( + rng.random((nframes, natoms)), dtype=dtype, device=self.device + ) + + l0, _more0 = loss0(learning_rate, natoms, model_pred, label) + assert l0.shape == () + + # Compare with dpmodel + dp_loss = TensorLossDP.deserialize(loss0.serialize()) + model_pred_np = {k: v.detach().cpu().numpy() for k, v in model_pred.items()} + label_np = { + k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v + for k, v in label.items() + } + l_dp, _ = dp_loss(learning_rate, natoms, model_pred_np, label_np) + + np.testing.assert_allclose( + l0.detach().cpu().numpy(), + np.array(l_dp), + rtol=rtol, + atol=atol, + err_msg="pt_expt vs dpmodel (atomic weight)", + )