fix: support electron temperature in lammps npt task group (#255)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced two optional parameters for enhanced electron temperature
configuration in the task group functionality.
- Added a new NPT simulation input template for improved flexibility in
testing configurations.
- Implemented a new test method to validate the correct integration of
the NPT simulation input.

- **Bug Fixes**
- Enhanced the robustness of testing capabilities for NPT simulations,
ensuring accurate validation of different parameter configurations.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Co-authored-by: Han Wang <[email protected]>

dpgen2/exploration/task/make_task_group_from_config.py

@@ -45,6 +45,8 @@ def npt_task_group_args():
     doc_use_clusters = "Calculate atomic model deviation"
     doc_relative_f_epsilon = "Calculate relative force model deviation"
     doc_relative_v_epsilon = "Calculate relative virial model deviation"
+    doc_ele_temp_f = "The electron temperature set by frame style"
+    doc_ele_temp_a = "The electron temperature set by atomistic style"
 
     return [
         Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]),
@@ -92,6 +94,20 @@ def npt_task_group_args():
             default=None,
             doc=doc_relative_v_epsilon,
         ),
+        Argument(
+            "ele_temp_f",
+            float,
+            optional=True,
+            default=None,
+            doc=doc_ele_temp_f,
+        ),
+        Argument(
+            "ele_temp_a",
+            float,
+            optional=True,
+            default=None,
+            doc=doc_ele_temp_a,
+        ),
     ]
 
 

tests/exploration/test_exploration_group.py

@@ -72,6 +72,43 @@
 """
 )
 
+in_template_npt_fparam = textwrap.dedent(
+    """variable        NSTEPS          equal 1000
+variable        THERMO_FREQ     equal 10
+variable        DUMP_FREQ       equal 10
+variable        TEMP            equal %f
+variable        ELE_TEMP        equal %f
+variable        PRES            equal %f
+variable        TAU_T           equal 0.100000
+variable        TAU_P           equal 0.500000
+
+units           metal
+boundary        p p p
+atom_style      atomic
+
+neighbor        1.0 bin
+
+box          tilt large
+if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data conf.lmp"
+change_box   all triclinic
+mass            1 10.000000
+mass            2 20.000000
+pair_style      deepmd model.000.pb model.001.pb model.002.pb  out_freq ${THERMO_FREQ} out_file model_devi.out fparam ${ELE_TEMP}
+pair_coeff      * *
+
+thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz
+thermo          ${THERMO_FREQ}
+dump            1 all custom ${DUMP_FREQ} traj.dump id type x y z fx fy fz
+restart         10000 dpgen.restart
+
+if "${restart} == 0" then "velocity        all create ${TEMP} 1111"
+fix             1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}
+
+timestep        0.001000
+run             ${NSTEPS} upto
+"""
+)
+
 in_template_nvt = textwrap.dedent(
     """variable        NSTEPS          equal 1000
 variable        THERMO_FREQ     equal 10
@@ -312,6 +349,75 @@ def test_npt(self, mock_random):
                 in_template_npt % (self.tt[j_idx], self.pp[k_idx]),
             )
 
+    @patch("dpgen2.exploration.task.lmp.lmp_input.random")
+    def test_npt_ele_temp(self, mock_random):
+        mock_random.randrange.return_value = 1110
+        self.confs = ["foo"]
+        self.tt = [100]
+        self.pp = [1]
+        self.numb_model = 3
+        self.mass_map = [10, 20]
+        ele_temp = 10.0
+
+        cpt_group = NPTTaskGroup()
+        cpt_group.set_md(
+            self.numb_model,
+            self.mass_map,
+            self.tt,
+            self.pp,
+            ele_temp_f=ele_temp,
+        )
+        cpt_group.set_conf(
+            self.confs,
+        )
+        task_group = cpt_group.make_task()
+
+        ngroup = len(task_group)
+        self.assertEqual(ngroup, len(self.confs) * len(self.tt) * len(self.pp))
+
+        for ii in range(ngroup):
+            i_idx = ii // (len(self.tt) * len(self.pp))
+            j_idx = (ii - len(self.tt) * len(self.pp) * i_idx) // len(self.pp)
+            k_idx = ii - len(self.tt) * len(self.pp) * i_idx - len(self.pp) * j_idx
+            self.assertEqual(
+                task_group[ii].files()[lmp_conf_name],
+                self.confs[i_idx],
+            )
+            self.assertEqual(
+                task_group[ii].files()[lmp_input_name],
+                in_template_npt_fparam % (self.tt[j_idx], ele_temp, self.pp[k_idx]),
+            )
+
+        cpt_group = NPTTaskGroup()
+        cpt_group.set_md(
+            self.numb_model,
+            self.mass_map,
+            self.tt,
+            self.pp,
+            ele_temp_a=ele_temp,
+        )
+        cpt_group.set_conf(
+            self.confs,
+        )
+        task_group = cpt_group.make_task()
+
+        ngroup = len(task_group)
+        self.assertEqual(ngroup, len(self.confs) * len(self.tt) * len(self.pp))
+
+        for ii in range(ngroup):
+            i_idx = ii // (len(self.tt) * len(self.pp))
+            j_idx = (ii - len(self.tt) * len(self.pp) * i_idx) // len(self.pp)
+            k_idx = ii - len(self.tt) * len(self.pp) * i_idx - len(self.pp) * j_idx
+            self.assertEqual(
+                task_group[ii].files()[lmp_conf_name],
+                self.confs[i_idx],
+            )
+            in_template_npt_aparam = in_template_npt_fparam.replace("fparam", "aparam")
+            self.assertEqual(
+                task_group[ii].files()[lmp_input_name],
+                in_template_npt_aparam % (self.tt[j_idx], ele_temp, self.pp[k_idx]),
+            )
+
     @patch("dpgen2.exploration.task.lmp.lmp_input.random")
     def test_nvt(self, mock_random):
         mock_random.randrange.return_value = 1110