fix: Filter pseudopotentials by system atom types for SIESTA

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

dpgen/generator/run.py

@@ -3720,7 +3720,10 @@ def make_fp_siesta(iter_index, jdata):
     for ii in fp_tasks:
         os.chdir(ii)
         sys_data = dpdata.System("POSCAR").data
-        ret = make_siesta_input(sys_data, fp_pp_files, fp_params)
+        pps = []
+        for iii in sys_data["atom_names"]:
+            pps.append(fp_pp_files[jdata["type_map"].index(iii)])
+        ret = make_siesta_input(sys_data, pps, fp_params)
         with open("input", "w") as fp:
             fp.write(ret)
         os.chdir(cwd)

dpgen/tools/relabel.py

@@ -81,11 +81,14 @@ def make_pwscf(tdir, fp_params, mass_map, fp_pp_path, fp_pp_files, user_input):
     os.chdir(cwd)
 
 
-def make_siesta(tdir, fp_params, fp_pp_path, fp_pp_files):
+def make_siesta(tdir, fp_params, fp_pp_path, fp_pp_files, type_map):
     cwd = os.getcwd()
     os.chdir(tdir)
     sys_data = dpdata.System("POSCAR").data
-    ret = make_siesta_input(sys_data, fp_pp_files, fp_params)
+    pps = []
+    for iii in sys_data["atom_names"]:
+        pps.append(fp_pp_files[type_map.index(iii)])
+    ret = make_siesta_input(sys_data, pps, fp_params)
     open("input", "w").write(ret)
     os.chdir(cwd)
 
@@ -149,7 +152,7 @@ def create_init_tasks(target_folder, param_file, output, fp_json, verbose=True):
                     ".", fp_params, mass_map, fp_pp_files, fp_pp_files, user_input
                 )
             elif fp_style == "siesta":
-                make_siesta(".", fp_params, fp_pp_files, fp_pp_files)
+                make_siesta(".", fp_params, fp_pp_files, fp_pp_files, type_map)
             os.chdir(cwd_)
 
 
@@ -170,6 +173,7 @@ def create_tasks(
     sys = jdata["sys_configs"]
     # fp settings
     mass_map = jdata["mass_map"]
+    type_map = jdata["type_map"]
     fp_style = fp_jdata["fp_style"]
     fp_pp_path = fp_jdata["fp_pp_path"]
     fp_pp_files = fp_jdata["fp_pp_files"]
@@ -284,7 +288,7 @@ def create_tasks(
                     ".", fp_params, mass_map, fp_pp_files, fp_pp_files, user_input
                 )
             elif fp_style == "siesta":
-                make_siesta(".", fp_params, mass_map, fp_pp_files, fp_pp_files)
+                make_siesta(".", fp_params, mass_map, fp_pp_files, type_map)
             os.chdir(cwd_)
     os.chdir(cwd)
 

tests/generator/test_make_fp.py

@@ -1480,5 +1480,66 @@ def test_make_fp_custom(self):
         shutil.rmtree("iter.000000")
 
 
+class TestMakeFPSIESTASubsetElements(unittest.TestCase):
+    """Test SIESTA with subset of elements (e.g., only C when C and H pseudopotentials are provided)."""
+
+    def test_make_fp_siesta_subset_elements(self):
+        setUpModule()
+        if os.path.isdir("iter.000000"):
+            shutil.rmtree("iter.000000")
+        with open(param_siesta_file) as fp:
+            jdata = json.load(fp)
+        with open(machine_file) as fp:
+            mdata = json.load(fp)
+        
+        # Create a system with only C atoms (type 0), even though type_map includes C, H, N
+        md_descript = []
+        nsys = 2
+        nmd = 3
+        n_frame = 10
+        for ii in range(nsys):
+            tmp = []
+            for jj in range(nmd):
+                tmp.append(np.arange(0, 0.29, 0.29 / 10))
+            md_descript.append(tmp)
+        
+        # Only C atoms (all type 0) - this is the key part of the test
+        atom_types = [0, 0, 0, 0, 0, 0]
+        type_map = jdata["type_map"]  # This is ["C", "H", "N"]
+        
+        _make_fake_md(0, md_descript, atom_types, type_map)
+        make_fp(0, jdata, {})
+        
+        # Verify that the input file was created successfully
+        fp_path = os.path.join("iter.000000", "02.fp")
+        tasks = glob.glob(os.path.join(fp_path, "task.*"))
+        self.assertGreater(len(tasks), 0, "No FP tasks were created")
+        
+        # Check that input files exist and contain only C species
+        for task in tasks:
+            input_file = os.path.join(task, "input")
+            self.assertTrue(os.path.isfile(input_file), f"Input file not found in {task}")
+            
+            with open(input_file) as fp:
+                content = fp.read()
+                # Verify the input contains only C in the species block
+                self.assertIn("%block Chemical_Species_label", content)
+                self.assertIn("C", content)
+                # The species block should only have 1 species (C), not 3 (C, H, N)
+                lines = content.split("\n")
+                in_species_block = False
+                species_count = 0
+                for line in lines:
+                    if "%block Chemical_Species_label" in line:
+                        in_species_block = True
+                    elif "%endblock Chemical_Species_label" in line:
+                        in_species_block = False
+                    elif in_species_block and line.strip():
+                        species_count += 1
+                self.assertEqual(species_count, 1, f"Expected 1 species (C), but found {species_count}")
+        
+        shutil.rmtree("iter.000000")
+
+
 if __name__ == "__main__":
     unittest.main()