test(pt_expt): add .pt2 (AOTInductor) unit tests and bug fixes

deepmd/pt_expt/infer/deep_eval.py

@@ -621,6 +621,22 @@ def _eval_model(
                 dtype=torch.float64,
                 device=DEVICE,
             )
+        elif self._is_pt2 and self.get_dim_fparam() > 0:
+            # .pt2 models are compiled with fparam as a required input.
+            # When the user omits fparam, fill with default values from metadata.
+            default_fp = self.metadata.get("default_fparam")
+            if default_fp is not None:
+                fparam_t = (
+                    torch.tensor(default_fp, dtype=torch.float64, device=DEVICE)
+                    .unsqueeze(0)
+                    .expand(nframes, -1)
+                    .contiguous()
+                )
+            else:
+                raise ValueError(
+                    f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) "
+                    "but was not provided, and no default_fparam is stored in the model."
+                )
         else:
             fparam_t = None
 

deepmd/pt_expt/utils/serialization.py

@@ -195,6 +195,7 @@ def _collect_metadata(model: torch.nn.Module) -> dict:
         "mixed_types": model.mixed_types(),
         "sel_type": model.get_sel_type(),
         "has_default_fparam": model.has_default_fparam(),
+        "default_fparam": model.get_default_fparam(),
         "fitting_output_defs": fitting_output_defs,
     }
 

source/api_cc/include/DeepPotPTExpt.h

@@ -202,6 +202,7 @@ class DeepPotPTExpt : public DeepPotBackend {
   int daparam;
   bool aparam_nall;
   bool has_default_fparam_;
+  std::vector<double> default_fparam_;
   double rcut;
   int gpu_id;
   bool gpu_enabled;

source/api_cc/src/DeepPotPTExpt.cc

@@ -636,6 +636,13 @@ void DeepPotPTExpt::init(const std::string& model,
   } else {
     has_default_fparam_ = false;
   }
+  default_fparam_.clear();
+  if (has_default_fparam_ && metadata.obj_val.count("default_fparam") &&
+      metadata["default_fparam"].type == JsonValue::Array) {
+    for (const auto& v : metadata["default_fparam"].as_array()) {
+      default_fparam_.push_back(v.as_double());
+    }
+  }
 
   type_map.clear();
   for (const auto& v : metadata["type_map"].as_array()) {
@@ -818,6 +825,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
                          valuetype_options)
             .to(torch::kFloat64)
             .to(device);
+  } else if (!default_fparam_.empty()) {
+    fparam_tensor =
+        torch::from_blob(const_cast<double*>(default_fparam_.data()),
+                         {1, static_cast<std::int64_t>(default_fparam_.size())},
+                         torch::TensorOptions().dtype(torch::kFloat64))
+            .clone()
+            .to(device);
+  } else if (dfparam > 0) {
+    throw deepmd::deepmd_exception(
+        "fparam is required for this model (dim_fparam=" +
+        std::to_string(dfparam) +
+        ") but was not provided, and no default_fparam is stored.");
   } else {
     fparam_tensor = torch::zeros({0}, options).to(device);
   }
@@ -982,6 +1001,15 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
       min_z = std::min(min_z, coord_d[ii * 3 + 2]);
       max_z = std::max(max_z, coord_d[ii * 3 + 2]);
     }
+    // Shift coords so minimum is at rcut (ensures all atoms are in [0, L))
+    double shift_x = rcut - min_x;
+    double shift_y = rcut - min_y;
+    double shift_z = rcut - min_z;
+    for (int ii = 0; ii < natoms; ++ii) {
+      coord_d[ii * 3 + 0] += shift_x;
+      coord_d[ii * 3 + 1] += shift_y;
+      coord_d[ii * 3 + 2] += shift_z;
+    }
     box_d.resize(9, 0.0);
     box_d[0] = (max_x - min_x) + 2.0 * rcut;
     box_d[4] = (max_y - min_y) + 2.0 * rcut;
@@ -1052,6 +1080,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener,
                          valuetype_options)
             .to(torch::kFloat64)
             .to(device);
+  } else if (!default_fparam_.empty()) {
+    fparam_tensor =
+        torch::from_blob(const_cast<double*>(default_fparam_.data()),
+                         {1, static_cast<std::int64_t>(default_fparam_.size())},
+                         torch::TensorOptions().dtype(torch::kFloat64))
+            .clone()
+            .to(device);
+  } else if (dfparam > 0) {
+    throw deepmd::deepmd_exception(
+        "fparam is required for this model (dim_fparam=" +
+        std::to_string(dfparam) +
+        ") but was not provided, and no default_fparam is stored.");
   } else {
     fparam_tensor = torch::zeros({0}, options).to(device);
   }

source/api_cc/tests/test_deeppot_a_fparam_aparam_ptexpt.cc

@@ -370,3 +370,89 @@ TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, has_default_fparam) {
   EXPECT_EQ(dp.dim_fparam(), 1);
   EXPECT_TRUE(dp.has_default_fparam());
 }
+
+// Eval without fparam should produce same result as eval with explicit default
+TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, eval_default_vs_explicit) {
+  using VALUETYPE = TypeParam;
+  deepmd::DeepPot& dp = this->dp;
+
+  std::vector<VALUETYPE> coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  00.25, 3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector<int> atype = {0, 0, 0, 0, 0, 0};
+  std::vector<VALUETYPE> box = {13., 0., 0., 0., 13., 0., 0., 0., 13.};
+  // The default fparam value from gen_fparam_aparam.py
+  std::vector<VALUETYPE> explicit_fparam = {0.25852028};
+  std::vector<VALUETYPE> aparam = {0.25852028, 0.25852028, 0.25852028,
+                                   0.25852028, 0.25852028, 0.25852028};
+
+  // Eval with explicit fparam
+  double e_explicit;
+  std::vector<VALUETYPE> f_explicit, v_explicit;
+  dp.compute(e_explicit, f_explicit, v_explicit, coord, atype, box,
+             explicit_fparam, aparam);
+
+  // Eval without fparam (should use default)
+  double e_default;
+  std::vector<VALUETYPE> f_default, v_default;
+  std::vector<VALUETYPE> empty_fparam;
+  dp.compute(e_default, f_default, v_default, coord, atype, box, empty_fparam,
+             aparam);
+
+  EXPECT_LT(fabs(e_explicit - e_default), EPSILON);
+  for (size_t ii = 0; ii < f_explicit.size(); ++ii) {
+    EXPECT_LT(fabs(f_explicit[ii] - f_default[ii]), EPSILON);
+  }
+  for (size_t ii = 0; ii < v_explicit.size(); ++ii) {
+    EXPECT_LT(fabs(v_explicit[ii] - v_default[ii]), EPSILON);
+  }
+}
+
+// Same test but with external nlist (LAMMPS path)
+TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, eval_default_vs_explicit_lmp) {
+  using VALUETYPE = TypeParam;
+  deepmd::DeepPot& dp = this->dp;
+
+  std::vector<VALUETYPE> coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  00.25, 3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector<int> atype = {0, 0, 0, 0, 0, 0};
+  std::vector<VALUETYPE> box = {13., 0., 0., 0., 13., 0., 0., 0., 13.};
+  std::vector<VALUETYPE> explicit_fparam = {0.25852028};
+  std::vector<VALUETYPE> aparam = {0.25852028, 0.25852028, 0.25852028,
+                                   0.25852028, 0.25852028, 0.25852028};
+
+  float rc = dp.cutoff();
+  int nloc = coord.size() / 3;
+  std::vector<VALUETYPE> coord_cpy;
+  std::vector<int> atype_cpy, mapping;
+  std::vector<std::vector<int> > nlist_data;
+  _build_nlist<VALUETYPE>(nlist_data, coord_cpy, atype_cpy, mapping, coord,
+                          atype, box, rc);
+  int nall = coord_cpy.size() / 3;
+  std::vector<int> ilist(nloc), numneigh(nloc);
+  std::vector<int*> firstneigh(nloc);
+  deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]);
+  convert_nlist(inlist, nlist_data);
+
+  // Eval with explicit fparam (external nlist)
+  double e_explicit;
+  std::vector<VALUETYPE> f_explicit, v_explicit;
+  dp.compute(e_explicit, f_explicit, v_explicit, coord_cpy, atype_cpy, box,
+             nall - nloc, inlist, 0, explicit_fparam, aparam);
+
+  // Eval without fparam (external nlist, should use default)
+  double e_default;
+  std::vector<VALUETYPE> f_default, v_default;
+  std::vector<VALUETYPE> empty_fparam;
+  dp.compute(e_default, f_default, v_default, coord_cpy, atype_cpy, box,
+             nall - nloc, inlist, 0, empty_fparam, aparam);
+
+  EXPECT_LT(fabs(e_explicit - e_default), EPSILON);
+  for (size_t ii = 0; ii < f_explicit.size(); ++ii) {
+    EXPECT_LT(fabs(f_explicit[ii] - f_default[ii]), EPSILON);
+  }
+  for (size_t ii = 0; ii < v_explicit.size(); ++ii) {
+    EXPECT_LT(fabs(v_explicit[ii] - v_default[ii]), EPSILON);
+  }
+}