diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 9b87d4234f..7b82a8d187 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -279,6 +279,22 @@ def forward( rmse_ae.detach(), find_atom_ener ) + if self.has_v and "virial" in model_pred and "virial" in label: + find_virial = label.get("find_virial", 0.0) + pref_v = pref_v * find_virial + diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9) + l2_virial_loss = torch.mean(torch.square(diff_v)) + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + l2_virial_loss.detach(), find_virial + ) + loss += atom_norm * (pref_v * l2_virial_loss) + rmse_v = l2_virial_loss.sqrt() * atom_norm + more_loss["rmse_v"] = self.display_if_exist(rmse_v.detach(), find_virial) + if mae: + mae_v = torch.mean(torch.abs(diff_v)) * atom_norm + more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial) + if not self.inference: more_loss["rmse"] = torch.sqrt(loss.detach()) return model_pred, loss, more_loss diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index c958a62bf6..5676a0397a 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -70,14 +70,14 @@ def __init__( self, *args: Any, # underscore to prevent conflict with normal inputs - atomic_model_: T_AtomicModel | None = None, + atomic_model_: T_AtomicModel | None = None, # type: ignore **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) if atomic_model_ is not None: - self.atomic_model: T_AtomicModel = atomic_model_ + self.atomic_model: T_AtomicModel = atomic_model_ # type: ignore else: - self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs) + self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs) # type: ignore self.precision_dict = PRECISION_DICT self.reverse_precision_dict = RESERVED_PRECISION_DICT self.global_pt_float_precision = GLOBAL_PT_FLOAT_PRECISION @@ -138,6 +138,7 @@ def forward_common( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + coord_corr_for_virial: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -156,6 +157,9 @@ def forward_common( atomic parameter. nf x nloc x nda do_atomic_virial If calculate the atomic virial. + coord_corr_for_virial + The coordinates correction of the atoms for virial. + shape: nf x (nloc x 3) Returns ------- @@ -183,6 +187,14 @@ def forward_common( mixed_types=True, box=bb, ) + if coord_corr_for_virial is not None: + coord_corr_for_virial = coord_corr_for_virial.to(cc.dtype) + extended_coord_corr = torch.gather( + coord_corr_for_virial, 1, mapping.unsqueeze(-1).expand(-1, -1, 3) + ) + else: + extended_coord_corr = None + model_predict_lower = self.forward_common_lower( extended_coord, extended_atype, @@ -191,6 +203,7 @@ def forward_common( do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap, + extended_coord_corr=extended_coord_corr, ) model_predict = communicate_extended_output( model_predict_lower, @@ -247,6 +260,7 @@ def forward_common_lower( do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, + extended_coord_corr: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping @@ -273,6 +287,8 @@ def forward_common_lower( The data needed for communication for parallel inference. extra_nlist_sort whether to forcibly sort the nlist. + extended_coord_corr + coordinates correction for virial in extended region. nf x (nall x 3) Returns ------- @@ -305,6 +321,7 @@ def forward_common_lower( do_atomic_virial=do_atomic_virial, create_graph=self.training, mask=atomic_ret["mask"] if "mask" in atomic_ret else None, + extended_coord_corr=extended_coord_corr, ) model_predict = self.output_type_cast(model_predict, input_prec) return model_predict diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 0eb7014416..f08fba756b 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -53,17 +53,31 @@ def __init__( def process_spin_input( self, coord: torch.Tensor, atype: torch.Tensor, spin: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Generate virtual coordinates and types, concat into the input.""" + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Generate virtual coordinates and types, concat into the input. + + Returns + ------- + coord_spin : torch.Tensor + Concatenated coordinates with shape (nframes, 2*nloc, 3). + atype_spin : torch.Tensor + Concatenated atom types with shape (nframes, 2*nloc). + coord_corr : torch.Tensor + Coordinate correction for virial with shape (nframes, 2*nloc, 3). + """ nframes, nloc = atype.shape coord = coord.reshape(nframes, nloc, 3) spin = spin.reshape(nframes, nloc, 3) atype_spin = torch.concat([atype, atype + self.ntypes_real], dim=-1) - virtual_coord = coord + spin * (self.virtual_scale_mask.to(atype.device))[ - atype - ].reshape([nframes, nloc, 1]) + # spin_dist = s_i * \mu_i + spin_dist = spin * (self.virtual_scale_mask.to(atype.device))[atype].reshape( + [nframes, nloc, 1] + ) + virtual_coord = coord + spin_dist coord_spin = torch.concat([coord, virtual_coord], dim=-2) - return coord_spin, atype_spin + # for spin virial corr + coord_corr = torch.concat([torch.zeros_like(coord), -spin_dist], dim=-2) + return coord_spin, atype_spin, coord_corr def process_spin_input_lower( self, @@ -72,10 +86,28 @@ def process_spin_input_lower( extended_spin: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor + ]: """ Add `extended_spin` into `extended_coord` to generate virtual atoms, and extend `nlist` and `mapping`. - Note that the final `extended_coord_updated` with shape [nframes, nall + nall, 3] has the following order: + + Returns + ------- + extended_coord_updated : torch.Tensor + Updated coordinates with virtual atoms, shape (nframes, 2*nall, 3). + extended_atype_updated : torch.Tensor + Updated atom types with virtual atoms, shape (nframes, 2*nall). + nlist_updated : torch.Tensor + Updated neighbor list including virtual atoms. + mapping_updated : torch.Tensor or None + Updated mapping indices, or None if input mapping is None. + extended_coord_corr : torch.Tensor + Coordinate correction for virial with shape (nframes, 2*nall, 3). + + Notes + ----- + The final `extended_coord_updated` with shape [nframes, nall + nall, 3] has the following order: - [:, :nloc]: original nloc real atoms. - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms. - [:, nloc + nloc: nloc + nall]: ghost real atoms. @@ -83,13 +115,18 @@ def process_spin_input_lower( """ nframes, nall = extended_coord.shape[:2] nloc = nlist.shape[1] - virtual_extended_coord = extended_coord + extended_spin * ( + extended_spin_dist = extended_spin * ( self.virtual_scale_mask.to(extended_atype.device) )[extended_atype].reshape([nframes, nall, 1]) + virtual_extended_coord = extended_coord + extended_spin_dist virtual_extended_atype = extended_atype + self.ntypes_real extended_coord_updated = concat_switch_virtual( extended_coord, virtual_extended_coord, nloc ) + # for spin virial corr + extended_coord_corr = concat_switch_virtual( + torch.zeros_like(extended_coord), -extended_spin_dist, nloc + ) extended_atype_updated = concat_switch_virtual( extended_atype, virtual_extended_atype, nloc ) @@ -105,6 +142,7 @@ def process_spin_input_lower( extended_atype_updated, nlist_updated, mapping_updated, + extended_coord_corr, ) def process_spin_output( @@ -376,7 +414,7 @@ def spin_sampled_func() -> list[dict[str, Any]]: sampled = sampled_func() spin_sampled = [] for sys in sampled: - coord_updated, atype_updated = self.process_spin_input( + coord_updated, atype_updated, _ = self.process_spin_input( sys["coord"], sys["atype"], sys["spin"] ) tmp_dict = { @@ -407,7 +445,9 @@ def forward_common( do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: nframes, nloc = atype.shape - coord_updated, atype_updated = self.process_spin_input(coord, atype, spin) + coord_updated, atype_updated, coord_corr_for_virial = self.process_spin_input( + coord, atype, spin + ) if aparam is not None: aparam = self.expand_aparam(aparam, nloc * 2) model_ret = self.backbone_model.forward_common( @@ -417,6 +457,7 @@ def forward_common( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + coord_corr_for_virial=coord_corr_for_virial, ) model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: @@ -439,7 +480,7 @@ def forward_common( ) = self.process_spin_output( atype, model_ret[f"{var_name}_derv_c"], - add_mag=False, + add_mag=True, virtual_scale=False, ) return model_ret @@ -463,6 +504,7 @@ def forward_common_lower( extended_atype_updated, nlist_updated, mapping_updated, + extended_coord_corr_for_virial, ) = self.process_spin_input_lower( extended_coord, extended_atype, extended_spin, nlist, mapping=mapping ) @@ -478,6 +520,7 @@ def forward_common_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=extra_nlist_sort, + extended_coord_corr=extended_coord_corr_for_virial, ) model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: @@ -503,7 +546,7 @@ def forward_common_lower( extended_atype, model_ret[f"{var_name}_derv_c"], nloc, - add_mag=False, + add_mag=True, virtual_scale=False, ) return model_ret @@ -550,6 +593,11 @@ def translated_output_def(self) -> dict[str, Any]: output_def["force"].squeeze(-2) output_def["force_mag"] = deepcopy(out_def_data["energy_derv_r_mag"]) output_def["force_mag"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"]) + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"]) + output_def["atom_virial"].squeeze(-3) return output_def def forward( @@ -578,7 +626,10 @@ def forward( if self.backbone_model.do_grad_r("energy"): model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) model_predict["force_mag"] = model_ret["energy_derv_r_mag"].squeeze(-2) - # not support virial by far + if self.backbone_model.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3) return model_predict @torch.jit.export @@ -615,5 +666,10 @@ def forward_lower( model_predict["extended_force_mag"] = model_ret[ "energy_derv_r_mag" ].squeeze(-2) - # not support virial by far + if self.backbone_model.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -3 + ) return model_predict diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index 839a028fcf..bc8f652b47 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -156,6 +156,7 @@ def fit_output_to_model_output( do_atomic_virial: bool = False, create_graph: bool = True, mask: torch.Tensor | None = None, + extended_coord_corr: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Transform the output of the fitting network to the model output. @@ -192,6 +193,12 @@ def fit_output_to_model_output( model_ret[kk_derv_r] = dr if vdef.c_differentiable: assert dc is not None + if extended_coord_corr is not None: + dc_corr = ( + dr.squeeze(-2).unsqueeze(-1) + @ extended_coord_corr.unsqueeze(-2) + ).view(list(dc.shape[:-2]) + [1, 9]) # noqa: RUF005 + dc = dc + dc_corr model_ret[kk_derv_c] = dc model_ret[kk_derv_c + "_redu"] = torch.sum( model_ret[kk_derv_c].to(redu_prec), dim=1 diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index afa62403e7..4f88238fa1 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -2602,9 +2602,9 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { for (int j = 0; j < natoms * 3; j++) { force_mag[i][j] = force_mag_flat[i * natoms * 3 + j]; } - // for (int j = 0; j < 9; j++) { - // virial[i][j] = virial_flat[i * 9 + j]; - // } + for (int j = 0; j < 9; j++) { + virial[i][j] = virial_flat[i * 9 + j]; + } } }; /** @@ -2705,9 +2705,9 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { for (int j = 0; j < natoms * 3; j++) { force_mag[i][j] = force_mag_flat[i * natoms * 3 + j]; } - // for (int j = 0; j < 9; j++) { - // virial[i][j] = virial_flat[i * 9 + j]; - // } + for (int j = 0; j < 9; j++) { + virial[i][j] = virial_flat[i * 9 + j]; + } for (int j = 0; j < natoms; j++) { atom_energy[i][j] = atom_energy_flat[i * natoms + j]; } diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 4a0cff1520..3acb28a002 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -862,11 +862,11 @@ void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, flatten_vector(fm_flat, fm); std::copy(fm_flat.begin(), fm_flat.end(), force_mag); } - // if (virial) { - // std::vector v_flat; - // flatten_vector(v_flat, v); - // std::copy(v_flat.begin(), v_flat.end(), virial); - // } + if (virial) { + std::vector v_flat; + flatten_vector(v_flat, v); + std::copy(v_flat.begin(), v_flat.end(), virial); + } if (atomic_energy) { std::vector ae_flat; flatten_vector(ae_flat, ae); diff --git a/source/api_cc/src/DeepSpinPT.cc b/source/api_cc/src/DeepSpinPT.cc index 8ccf2fd383..cadf772d9a 100644 --- a/source/api_cc/src/DeepSpinPT.cc +++ b/source/api_cc/src/DeepSpinPT.cc @@ -251,8 +251,7 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, c10::IValue energy_ = outputs.at("energy"); c10::IValue force_ = outputs.at("extended_force"); c10::IValue force_mag_ = outputs.at("extended_force_mag"); - // spin model not suported yet - // c10::IValue virial_ = outputs.at("virial"); + bool has_virial = outputs.contains("virial"); torch::Tensor flat_energy_ = energy_.toTensor().view({-1}); torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU); ener.assign(cpu_energy_.data_ptr(), @@ -267,11 +266,16 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, dforce_mag.assign( cpu_force_mag_.data_ptr(), cpu_force_mag_.data_ptr() + cpu_force_mag_.numel()); - // spin model not suported yet - // torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); - // torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); - // virial.assign(cpu_virial_.data_ptr(), - // cpu_virial_.data_ptr() + cpu_virial_.numel()); + + if (has_virial) { + c10::IValue virial_ = outputs.at("virial"); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + } else { + virial.clear(); + } // bkw map force.resize(static_cast(nframes) * fwd_map.size() * 3); @@ -281,8 +285,6 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, select_map(force_mag, dforce_mag, bkw_map, 3, nframes, fwd_map.size(), nall_real); if (atomic) { - // spin model not suported yet - // c10::IValue atom_virial_ = outputs.at("extended_virial"); c10::IValue atom_energy_ = outputs.at("atom_energy"); torch::Tensor flat_atom_energy_ = atom_energy_.toTensor().view({-1}).to(floatType); @@ -292,19 +294,23 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, datom_energy.assign( cpu_atom_energy_.data_ptr(), cpu_atom_energy_.data_ptr() + cpu_atom_energy_.numel()); - // spin model not suported yet - // torch::Tensor flat_atom_virial_ = - // atom_virial_.toTensor().view({-1}).to(floatType); - // torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); - // datom_virial.assign( - // cpu_atom_virial_.data_ptr(), - // cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); atom_energy.resize(static_cast(nframes) * fwd_map.size()); - // atom_virial.resize(static_cast(nframes) * fwd_map.size() * 9); select_map(atom_energy, datom_energy, bkw_map, 1, nframes, fwd_map.size(), nall_real); - // select_map(atom_virial, datom_virial, bkw_map, 9, nframes, - // fwd_map.size(), nall_real); + if (outputs.contains("extended_virial")) { + c10::IValue atom_virial_ = outputs.at("extended_virial"); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + datom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); + atom_virial.resize(static_cast(nframes) * fwd_map.size() * 9); + select_map(atom_virial, datom_virial, bkw_map, 9, nframes, + fwd_map.size(), nall_real); + } else { + atom_virial.clear(); + } } } template void DeepSpinPT::compute>( @@ -415,8 +421,7 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, c10::IValue energy_ = outputs.at("energy"); c10::IValue force_ = outputs.at("force"); c10::IValue force_mag_ = outputs.at("force_mag"); - // spin model not suported yet - // c10::IValue virial_ = outputs.at("virial"); + bool has_virial = outputs.contains("virial"); torch::Tensor flat_energy_ = energy_.toTensor().view({-1}); torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU); ener.assign(cpu_energy_.data_ptr(), @@ -431,13 +436,16 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, force_mag.assign( cpu_force_mag_.data_ptr(), cpu_force_mag_.data_ptr() + cpu_force_mag_.numel()); - // spin model not suported yet - // torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); - // torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); - // virial.assign(cpu_virial_.data_ptr(), - // cpu_virial_.data_ptr() + cpu_virial_.numel()); + if (has_virial) { + c10::IValue virial_ = outputs.at("virial"); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + } else { + virial.clear(); + } if (atomic) { - // c10::IValue atom_virial_ = outputs.at("atom_virial"); c10::IValue atom_energy_ = outputs.at("atom_energy"); torch::Tensor flat_atom_energy_ = atom_energy_.toTensor().view({-1}).to(floatType); @@ -445,12 +453,17 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, atom_energy.assign( cpu_atom_energy_.data_ptr(), cpu_atom_energy_.data_ptr() + cpu_atom_energy_.numel()); - // torch::Tensor flat_atom_virial_ = - // atom_virial_.toTensor().view({-1}).to(floatType); - // torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); - // atom_virial.assign( - // cpu_atom_virial_.data_ptr(), - // cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); + if (outputs.contains("atom_virial")) { + c10::IValue atom_virial_ = outputs.at("atom_virial"); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + atom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); + } else { + atom_virial.clear(); + } } } diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt index 8ec3361d3a..a3e7d067f7 100644 --- a/source/api_cc/tests/CMakeLists.txt +++ b/source/api_cc/tests/CMakeLists.txt @@ -31,6 +31,8 @@ add_test( NAME runUnitTest_cc COMMAND runUnitTests_cc WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) +set_tests_properties(runUnitTest_cc PROPERTIES ENVIRONMENT + "TF_DISABLE_MLIR_BRIDGE=1") set_target_properties(runUnitTests_cc PROPERTIES INSTALL_RPATH "$ORIGIN/../lib") target_compile_definitions(runUnitTests_cc PUBLIC ${prec_def}) install(TARGETS runUnitTests_cc DESTINATION bin/) diff --git a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc b/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc index 8b569dd707..26a92a165a 100644 --- a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc +++ b/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc @@ -43,33 +43,34 @@ class TestInferDeepSpinDpaPt : public ::testing::Test { // atype = np.array([0, 1, 1, 0, 1, 1]) // box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]).reshape(1, -1) // dp = DeepPot("deeppot_dpa_spin.pth") - // e, f, _, ae, _, fm, _ = dp.eval(coord, box, atype, atomic=True, spin=spin) + // e, f, v, ae, av, fm, _ = dp.eval(coord, box, atype, atomic=True, spin=spin) // np.set_printoptions(precision=16) - // print(f"{e.ravel()=} {f.ravel()=} {fm.ravel()=} {ae.ravel()=}") + // print(f"{e.ravel()=} {f.ravel()=} {v.ravel()=} {fm.ravel()=} + // {ae.ravel()=}") std::vector expected_e = {-5.835211567762678, -5.071189078159807, - -5.044361601406714, -5.582324154346981, + -5.044361601406714, -5.582324154346982, -5.059906899269188, -5.074135576182056}; std::vector expected_f = { - -0.0619881702551019, 0.0646720543680939, 0.2137632336140025, - 0.037800173877136, -0.096327623008356, -0.1531911892384847, - -0.112204927558682, 0.0299145670766557, -0.0589474826303666, - 0.2278904556868233, 0.0382061907026398, 0.0888060647788163, - -0.0078898845686437, 0.0019385598635839, -0.0791616129664364, - -0.083607647181527, -0.0384037490026167, -0.0112690135575317}; + -0.0619881702551019, 0.0646720543680956, 0.2137632336140034, + 0.0378001738771361, -0.0963276230083563, -0.153191189238485, + -0.112204927558682, 0.0299145670766558, -0.0589474826303669, + 0.2278904556868226, 0.0382061907026395, 0.0888060647788163, + -0.0078898845686436, 0.0019385598635835, -0.079161612966436, + -0.0836076471815266, -0.0384037490026167, -0.0112690135575319}; std::vector expected_fm = { - -3.0778301386623275, - -1.3135930534661662, - -0.8332043979367366, + -3.077830138662336, + -1.3135930534661686, + -0.8332043979367388, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - -0.5452347545527696, - -0.2051506559632127, - -0.4908015055951312, + -0.5452347545527724, + -0.2051506559632141, + -0.4908015055951336, 0.0, 0.0, 0.0, @@ -77,10 +78,28 @@ class TestInferDeepSpinDpaPt : public ::testing::Test { 0.0, 0.0, }; + std::vector expected_tot_v = { + -0.02421162406248903, -0.007399464728803163, 0.012770296772824489, + -0.005129031187923775, 0.001381624482046452, 0.00660380864594899, + 0.013956096416149198, 0.006788089581579212, -0.0102462390797404}; + std::vector expected_atom_v = { + 1.56445561e-03, -2.09940556e-04, -2.38887855e-03, 1.97370249e-03, + 4.43423788e-03, -1.87644028e-03, 2.28936857e-05, -1.60139439e-03, + 2.42612568e-03, -8.18033202e-03, -3.63076402e-03, 3.29578144e-03, + -3.77806007e-03, -5.67947162e-04, 3.71520105e-03, 3.68597629e-03, + 3.58375798e-03, -3.53992730e-03, -1.25979731e-02, -5.71674421e-04, + 4.05411698e-03, 9.59636201e-05, 2.98503616e-05, 1.41985921e-03, + 3.86508016e-03, 1.60940945e-03, -4.82848144e-03, 9.92060149e-03, + 1.31159078e-03, -1.08861454e-03, -1.33848341e-04, 2.24873882e-03, + -3.35090985e-03, -2.53510394e-03, -2.99970681e-03, 4.43709623e-03, + -1.06689446e-02, -2.83469144e-03, 5.46665681e-03, -1.93195329e-03, + -1.85479946e-03, 2.29490155e-03, 4.39982122e-03, 1.49850013e-03, + -1.63963632e-03, -4.24943142e-03, -1.46398507e-03, 3.43123463e-03, + -1.35483560e-03, -2.90845595e-03, 4.40119696e-03, 4.51742900e-03, + 4.69752322e-03, -7.10141593e-03}; int natoms; double expected_tot_e; - // std::vector expected_tot_v; deepmd::DeepSpin dp; @@ -93,18 +112,12 @@ class TestInferDeepSpinDpaPt : public ::testing::Test { natoms = expected_e.size(); EXPECT_EQ(natoms * 3, expected_f.size()); EXPECT_EQ(natoms * 3, expected_fm.size()); - // EXPECT_EQ(natoms * 9, expected_v.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); expected_tot_e = 0.; - // expected_tot_v.resize(9); - // std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.); for (int ii = 0; ii < natoms; ++ii) { expected_tot_e += expected_e[ii]; } - // for (int ii = 0; ii < natoms; ++ii) { - // for (int dd = 0; dd < 9; ++dd) { - // expected_tot_v[dd] += expected_v[ii * 9 + dd]; - // } - // } }; void TearDown() override {}; @@ -121,10 +134,9 @@ TYPED_TEST(TestInferDeepSpinDpaPt, cpu_build_nlist) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial; @@ -132,16 +144,18 @@ TYPED_TEST(TestInferDeepSpinDpaPt, cpu_build_nlist) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); - EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (virial.empty()) { + return; + } + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } } TYPED_TEST(TestInferDeepSpinDpaPt, cpu_build_nlist_atomic) { @@ -153,10 +167,10 @@ TYPED_TEST(TestInferDeepSpinDpaPt, cpu_build_nlist_atomic) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial, atom_ener, atom_vir; @@ -165,24 +179,29 @@ TYPED_TEST(TestInferDeepSpinDpaPt, cpu_build_nlist_atomic) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); EXPECT_EQ(atom_ener.size(), natoms); - // EXPECT_EQ(atom_vir.size(), natoms * 9); EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (!virial.empty()) { + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + } for (int ii = 0; ii < natoms; ++ii) { EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); } - // for (int ii = 0; ii < natoms * 9; ++ii) { - // EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); - // } + if (atom_vir.empty()) { + return; + } + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } } template @@ -210,43 +229,59 @@ class TestInferDeepSpinDpaPtNopbc : public ::testing::Test { // atype = np.array([0, 1, 1, 0, 1, 1]) // box = None // dp = DeepPot("deeppot_dpa_spin.pth") - // e, f, _, ae, _, fm, _ = dp.eval(coord, box, atype, atomic=True, + // e, f, v, ae, av, fm, _ = dp.eval(coord, box, atype, atomic=True, // spin=spin) // np.set_printoptions(precision=16) - // print(f"{e.ravel()=} {f.ravel()=} {fm.ravel()=} {ae.ravel()=}") + // print(f"{e.ravel()=} {f.ravel()=} {v.ravel()=} {fm.ravel()=} + // {ae.ravel()=}") - std::vector expected_e = {-5.921669893870771, -5.1676693791758685, + std::vector expected_e = {-5.921669893870772, -5.167669379175869, -5.205933794558385, -5.58688965168251, -5.080322972018686, -5.08213772482076}; std::vector expected_f = { - -0.2929142244191496, 0.0801070990501456, 0.148216178514704, - 0.2929142244191503, -0.0801070990501454, -0.1482161785147037, - -0.2094984819251435, 0.0241594118950041, -0.0215199116994508, - 0.3068843038300324, -0.001620530344866, 0.1508093841389746, - -0.0122719879278721, 0.0186341247897136, -0.1137104245023705, - -0.0851138339770169, -0.0411730063398516, -0.0155790479371533}; - std::vector expected_fm = {-1.5298530476860008, - 0.0071315024546899, - 0.0650492472558729, - 0., - 0., - 0., - 0., - 0., - 0., - -0.6212052813442365, - -0.2290265978320395, - -0.5101405083352206, - 0., - 0., - 0., - 0., - 0., - 0.}; + -0.2929142244191496, 0.0801070990501461, 0.1482161785147047, + 0.2929142244191509, -0.080107099050146, -0.1482161785147047, + -0.2094984819251434, 0.024159411895004, -0.0215199116994508, + 0.3068843038300326, -0.0016205303448664, 0.1508093841389744, + -0.0122719879278721, 0.0186341247897136, -0.1137104245023706, + -0.0851138339770167, -0.0411730063398516, -0.0155790479371534}; + std::vector expected_fm = {-1.5298530476859904, + 0.00713150245469, + 0.0650492472558721, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.6212052813442372, + -0.2290265978320397, + -0.5101405083352208, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0}; + std::vector expected_tot_v = { + -0.006855934251658762, -0.00491299773599024, 0.001790727101442579, + -0.004700307257262127, -0.001921876254115465, 0.00522933533951051, + 0.002220184013052337, 0.005305572188769711, -0.0015617293208553679}; + std::vector expected_atom_v = { + 0.00223849, -0.00093775, -0.00169399, -0.00058709, 0.00024594, + 0.00044429, -0.00108942, 0.00045638, 0.00082442, 0.00331061, + -0.00138688, -0.00250533, -0.00159569, 0.00066847, 0.00120755, + -0.00286535, 0.00120035, 0.00216838, -0.00726106, 0.00110669, + -0.0011729, 0.00151598, -0.00037897, 0.0004289, -0.00116555, + 0.00041409, -0.00051682, 0.00734061, 0.00124985, -0.00139847, + 0.00025622, 0.00238308, -0.00372649, -0.00155401, -0.00312998, + 0.00476667, -0.00909511, -0.00319323, 0.00524226, -0.00252116, + -0.00193343, 0.00250048, 0.00436096, 0.00167675, -0.00186338, + -0.00338948, -0.00175168, 0.00331916, -0.00176855, -0.00290696, + 0.0043746, 0.00453355, 0.00468798, -0.00694099}; int natoms; double expected_tot_e; - // std::vector expected_tot_v; deepmd::DeepSpin dp; @@ -259,18 +294,12 @@ class TestInferDeepSpinDpaPtNopbc : public ::testing::Test { natoms = expected_e.size(); EXPECT_EQ(natoms * 3, expected_f.size()); EXPECT_EQ(natoms * 3, expected_fm.size()); - // EXPECT_EQ(natoms * 9, expected_v.size()); + EXPECT_EQ(9, expected_tot_v.size()); + EXPECT_EQ(natoms * 9, expected_atom_v.size()); expected_tot_e = 0.; - // expected_tot_v.resize(9); - // std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.); for (int ii = 0; ii < natoms; ++ii) { expected_tot_e += expected_e[ii]; } - // for (int ii = 0; ii < natoms; ++ii) { - // for (int dd = 0; dd < 9; ++dd) { - // expected_tot_v[dd] += expected_v[ii * 9 + dd]; - // } - // } }; void TearDown() override {}; @@ -287,10 +316,9 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_build_nlist) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial; @@ -298,16 +326,18 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_build_nlist) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); - EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (virial.empty()) { + return; + } + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } } TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_build_nlist_atomic) { @@ -319,10 +349,10 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_build_nlist_atomic) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial, atom_ener, atom_vir; @@ -331,24 +361,29 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_build_nlist_atomic) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); EXPECT_EQ(atom_ener.size(), natoms); - // EXPECT_EQ(atom_vir.size(), natoms * 9); EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (!virial.empty()) { + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + } for (int ii = 0; ii < natoms; ++ii) { EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); } - // for (int ii = 0; ii < natoms * 9; ++ii) { - // EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); - // } + if (atom_vir.empty()) { + return; + } + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } } TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist) { @@ -360,10 +395,9 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial; @@ -380,16 +414,18 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); - EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (virial.empty()) { + return; + } + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } } TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_atomic) { @@ -401,10 +437,10 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_atomic) { std::vector& expected_e = this->expected_e; std::vector& expected_f = this->expected_f; std::vector& expected_fm = this->expected_fm; - // std::vector& expected_v = this->expected_v; + std::vector& expected_tot_v = this->expected_tot_v; + std::vector& expected_atom_v = this->expected_atom_v; int& natoms = this->natoms; double& expected_tot_e = this->expected_tot_e; - // std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepSpin& dp = this->dp; double ener; std::vector force, force_mag, virial, atom_ener, atom_vir; @@ -421,22 +457,27 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_atomic) { EXPECT_EQ(force.size(), natoms * 3); EXPECT_EQ(force_mag.size(), natoms * 3); - // EXPECT_EQ(virial.size(), 9); EXPECT_EQ(atom_ener.size(), natoms); - // EXPECT_EQ(atom_vir.size(), natoms * 9); EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); for (int ii = 0; ii < natoms * 3; ++ii) { EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); EXPECT_LT(fabs(force_mag[ii] - expected_fm[ii]), EPSILON); } - // for (int ii = 0; ii < 3 * 3; ++ii) { - // EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); - // } + if (!virial.empty()) { + EXPECT_EQ(virial.size(), 9); + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + } for (int ii = 0; ii < natoms; ++ii) { EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); } - // for (int ii = 0; ii < natoms * 9; ++ii) { - // EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); - // } + if (atom_vir.empty()) { + return; + } + EXPECT_EQ(atom_vir.size(), natoms * 9); + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_atom_v[ii]), EPSILON); + } } diff --git a/source/tests/pt/model/test_autodiff.py b/source/tests/pt/model/test_autodiff.py index 31e06af751..104b838d01 100644 --- a/source/tests/pt/model/test_autodiff.py +++ b/source/tests/pt/model/test_autodiff.py @@ -57,6 +57,39 @@ def stretch_box(old_coord, old_box, new_box): return ncoord.reshape(old_coord.shape) +def finite_difference_cell_energy( + energy_func, + cell: np.ndarray, + delta: float = 1e-4, +) -> np.ndarray: + """ + Compute cell energy derivatives via central finite differences. + + Parameters + ---------- + energy_func : callable + Function that returns a scalar energy for a given cell. + cell : np.ndarray + Cell matrix with shape (3, 3). + delta : float + Perturbation size in unitless. + + Returns + ------- + np.ndarray + Energy derivatives dE/dh with shape (3, 3). + """ + deriv = np.zeros_like(cell) + for ii in range(3): + for jj in range(3): + perturb = np.zeros_like(cell) + perturb[ii, jj] = delta + ep = energy_func(cell + perturb) + em = energy_func(cell - perturb) + deriv[ii, jj] = (ep - em) / (2.0 * delta) + return deriv + + class ForceTest: def test( self, @@ -141,11 +174,17 @@ def test( cell = (cell) + 5.0 * torch.eye(3, device="cpu") coord = torch.rand([natoms, 3], dtype=dtype, device="cpu", generator=generator) coord = torch.matmul(coord, cell) + spin = torch.rand([natoms, 3], dtype=dtype, device="cpu", generator=generator) atype = torch.IntTensor([0, 0, 0, 1, 1]) # assumes input to be numpy tensor coord = coord.numpy() + spin = spin.numpy() cell = cell.numpy() - test_keys = ["energy", "force", "virial"] + test_spin = getattr(self, "test_spin", False) + if not test_spin: + test_keys = ["energy", "force", "virial"] + else: + test_keys = ["energy", "force", "force_mag", "virial"] def np_infer( new_cell, @@ -157,6 +196,7 @@ def np_infer( ).unsqueeze(0), torch.tensor(new_cell, device="cpu").unsqueeze(0), atype, + spins=torch.tensor(spin, device=env.DEVICE).unsqueeze(0), ) # detach ret = {key: to_numpy_array(result[key].squeeze(0)) for key in test_keys} @@ -251,3 +291,57 @@ def setUp(self) -> None: self.type_split = False self.test_spin = True self.model = get_model(model_params).to(env.DEVICE) + + +class TestEnergyModelSpinSeAVirial(unittest.TestCase, VirialTest): + def setUp(self) -> None: + model_params = copy.deepcopy(model_spin) + self.type_split = False + self.test_spin = True + self.model = get_model(model_params).to(env.DEVICE) + + +class TestEnergyModelSpinSeAVirialShear(unittest.TestCase): + def setUp(self) -> None: + model_params = copy.deepcopy(model_spin) + self.model = get_model(model_params).to(env.DEVICE) + + def test(self) -> None: + places = 5 + delta = 1e-4 + natoms = 5 + generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED) + atype = np.array([0, 0, 0, 1, 1], dtype=np.int32) + + # === Step 1. Prepare inputs === + cell = torch.rand([3, 3], dtype=dtype, device="cpu", generator=generator) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu") + coord = torch.rand([natoms, 3], dtype=dtype, device="cpu", generator=generator) + coord = torch.matmul(coord, cell) + spin = torch.rand([natoms, 3], dtype=dtype, device="cpu", generator=generator) + coord = coord.numpy() + spin = spin.numpy() + cell = cell.numpy() + + # === Step 2. Define energy and virial evaluators === + def np_infer(new_cell): + coord_new = stretch_box(coord, cell, new_cell) + result = eval_model( + self.model, + coord_new.reshape(1, natoms, 3), + new_cell.reshape(1, 3, 3), + atype, + spins=spin.reshape(1, natoms, 3), + ) + return result + + def energy_func(new_cell): + return np_infer(new_cell)["energy"].reshape(-1)[0] + + # === Step 3. Finite-difference virial by shear perturbations === + dE_dh = finite_difference_cell_energy(energy_func, cell, delta=delta) + virial_fd = -(dE_dh.T @ cell).reshape(9) + + # === Step 4. Compare with analytic virial === + virial_ref = np_infer(cell)["virial"].reshape(9) + np.testing.assert_almost_equal(virial_fd, virial_ref, decimal=places) diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index ddea392f33..66bb1082a0 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -115,7 +115,7 @@ def test_input_output_process(self) -> None: nframes, nloc = self.coord.shape[:2] self.real_ntypes = self.model.spin.get_ntypes_real() # 1. test forward input process - coord_updated, atype_updated = self.model.process_spin_input( + coord_updated, atype_updated, _ = self.model.process_spin_input( self.coord, self.atype, self.spin ) # compare atypes of real and virtual atoms @@ -174,6 +174,7 @@ def test_input_output_process(self) -> None: extended_atype_updated, nlist_updated, mapping_updated, + _, ) = self.model.process_spin_input_lower( extended_coord, extended_atype, extended_spin, nlist, mapping=mapping ) diff --git a/source/tests/pt/model/test_nosel.py b/source/tests/pt/model/test_nosel.py new file mode 100644 index 0000000000..fe349231ea --- /dev/null +++ b/source/tests/pt/model/test_nosel.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor.dpa3 import ( + RepFlowArgs, +) +from deepmd.pt.model.descriptor import ( + DescrptDPA3, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, +) + +from ...seed import ( + GLOBAL_SEED, +) +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) +from .test_mlp import ( + get_tols, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class TestDescrptDPA3Nosel(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ) -> None: + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + ua, + rus, + ruri, + acr, + nme, + prec, + ect, + optim, + ) in itertools.product( + [True, False], # update_angle + ["res_residual"], # update_style + ["norm", "const"], # update_residual_init + [0, 1], # a_compress_rate + [1, 2], # n_multi_edge_message + ["float64"], # precision + [False], # use_econf_tebd + [True, False], # optim_update + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + if prec == "float64": + atol = 1e-8 # marginal GPU test cases... + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=10, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei, + a_compress_rate=acr, + n_multi_edge_message=nme, + axis_neuron=4, + update_angle=ua, + update_style=rus, + update_residual_init=ruri, + optim_update=optim, + smooth_edge_update=True, + sel_reduce_factor=1.0, # test consistent when sel_reduce_factor == 1.0 + ) + + # dpa3 new impl + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + repflow.use_dynamic_sel = True + + # dpa3 new impl + dd1 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repflows.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.repflows.stddev = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + rd0, _, _, _, _ = dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + torch.tensor(self.mapping, dtype=int, device=env.DEVICE), + ) + # serialization + dd1.repflows.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd1.repflows.stddev = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + rd1, _, _, _, _ = dd1( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + torch.tensor(self.mapping, dtype=int, device=env.DEVICE), + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + + # def test_jit( + # self, + # ) -> None: + # rng = np.random.default_rng(100) + # nf, nloc, nnei = self.nlist.shape + # davg = rng.normal(size=(self.nt, nnei, 4)) + # dstd = rng.normal(size=(self.nt, nnei, 4)) + # dstd = 0.1 + np.abs(dstd) + # + # for ( + # ua, + # rus, + # ruri, + # acr, + # nme, + # prec, + # ect, + # ) in itertools.product( + # [True, False], # update_angle + # ["res_residual"], # update_style + # ["norm", "const"], # update_residual_init + # [0, 1], # a_compress_rate + # [1, 2], # n_multi_edge_message + # ["float64"], # precision + # [False], # use_econf_tebd + # ): + # dtype = PRECISION_DICT[prec] + # rtol, atol = get_tols(prec) + # + # repflow = RepFlowArgs( + # n_dim=20, + # e_dim=10, + # a_dim=10, + # nlayers=3, + # e_rcut=self.rcut, + # e_rcut_smth=self.rcut_smth, + # e_sel=nnei, + # a_rcut=self.rcut - 0.1, + # a_rcut_smth=self.rcut_smth, + # a_sel=nnei - 1, + # a_compress_rate=acr, + # n_multi_edge_message=nme, + # axis_neuron=4, + # update_angle=ua, + # update_style=rus, + # update_residual_init=ruri, + # ) + # + # # dpa3 new impl + # dd0 = DescrptDPA3( + # self.nt, + # repflow=repflow, + # # kwargs for descriptor + # exclude_types=[], + # precision=prec, + # use_econf_tebd=ect, + # type_map=["O", "H"] if ect else None, + # seed=GLOBAL_SEED, + # ).to(env.DEVICE) + # + # dd0.repflows.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + # dd0.repflows.stddev = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + # model = torch.jit.script(dd0) diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 8ec0edc180..7728aea895 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -915,7 +915,10 @@ def ff_spin(_spin): fdf.reshape(-1, 3), rff.reshape(-1, 3), decimal=places ) - if not test_spin: + # this option can be removed after other backends support spin virial + test_spin_virial = getattr(self, "test_spin_virial", False) + + if not test_spin or test_spin_virial: def ff_cell(bb): input_dict = { @@ -925,6 +928,8 @@ def ff_cell(bb): "aparam": aparam, "fparam": fparam, } + if test_spin: + input_dict["spin"] = spin return module(**input_dict)["energy"] fdv = ( @@ -944,13 +949,12 @@ def ff_cell(bb): "aparam": aparam, "fparam": fparam, } + if test_spin: + input_dict["spin"] = spin rfv = module(**input_dict)["virial"] np.testing.assert_almost_equal( fdv.reshape(-1, 9), rfv.reshape(-1, 9), decimal=places ) - else: - # not support virial by far - pass @unittest.skipIf(TEST_DEVICE == "cpu" and CI, "Skip test on CPU.") def test_device_consistence(self) -> None: diff --git a/source/tests/universal/pt/model/test_model.py b/source/tests/universal/pt/model/test_model.py index 09111e988f..2542b1f6c0 100644 --- a/source/tests/universal/pt/model/test_model.py +++ b/source/tests/universal/pt/model/test_model.py @@ -744,6 +744,8 @@ def setUpClass(cls) -> None: cls.expected_sel_type = ft.get_sel_type() cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + # this option can be removed after other backends support spin virial + cls.test_spin_virial = True @classmethod def tearDownClass(cls) -> None: