add the prepare for spin case

fpop/abacus.py

@@ -10,7 +10,7 @@
     Set,
     Dict,
     Optional,
-    Union,
+    Union
 )
 import numpy as np
 from dflow.python import (
@@ -314,7 +314,8 @@ def __init__(
             kpt_file: Optional[Union[str,Path]] = None,
             orb_files: Optional[Dict[str, Union[str,Path]]] = None,
             deepks_descriptor: Optional[Union[str,Path]] = None,
-            deepks_model: Optional[Union[str,Path]] = None
+            deepks_model: Optional[Union[str,Path]] = None,
+            constrain_elements: Optional[List[str]] = None,
     ):     
         """The input information of an ABACUS job except for STRU.
 
@@ -337,6 +338,9 @@ def __init__(
             The deepks descriptor file, by default None.
         deepks_model : str, optional
             The deepks model file, by default None.
+        constrain_elements : List[str], optional
+            The elements that need to constrain the magnetic moment, by default None. 
+            For the constrained elements, the flag "sc 1 1 1" will be added in STRU.
         """        
         self.input_file = input_file
         self._input = AbacusInputs.read_inputf(self.input_file)
@@ -347,7 +351,8 @@ def __init__(
         self._orb_files = {} if orb_files == None else self._read_dict_file(orb_files)
         self._deepks_descriptor = None if deepks_descriptor == None else (os.path.split(deepks_descriptor)[1], Path(deepks_descriptor).read_text())
         self._deepks_model = None if deepks_model == None else (os.path.split(deepks_model)[1], Path(deepks_model).read_bytes())
-
+        self._constrin_elements = constrain_elements
+
     def _read_dict_file(self,input_dict,out_dict=None):
         # input_dict is a dict whose value is a file.
         # The filename and context will make up a tuple, which is 
@@ -394,6 +399,9 @@ def get_deepks_descriptor(self):
 
     def get_deepks_model(self):
         return self._deepks_model
+
+    def get_constrain_elements(self):
+        return self._constrin_elements
 
     @staticmethod
     def read_inputf(inputf: Union[str,Path]) -> dict:
@@ -446,21 +454,15 @@ def write_pporb(self,element_list : List[str]):
         List[List]
             a list of the list of pp files, and orbital files 
         """  
-        need_orb = False
-        if self._input.get("basis_type","pw").lower() in ["lcao","lcao_in_pw"]:
-            need_orb = True
         pp,orb = [],[]
         for ielement in element_list:
             if ielement in self._pp_files:
                 Path(self._pp_files[ielement][0]).write_text(self._pp_files[ielement][1])
                 pp.append(self._pp_files[ielement][0])
-            if need_orb and ielement in self._orb_files:
+            if ielement in self._orb_files:
                 Path(self._orb_files[ielement][0]).write_text(self._orb_files[ielement][1]) 
                 orb.append(self._orb_files[ielement][0])
 
-        if not orb: 
-            orb = None
-
         return [pp,orb]     
 
     def write_deepks(self):
@@ -517,7 +519,7 @@ class PrepAbacus(PrepFp):
     def prep_task(
             self,
             conf_frame,
-            abacus_inputs: AbacusInputs,
+            inputs: AbacusInputs,
             prepare_image_config: Optional[Dict] = None,
             optional_input: Optional[Dict] = None,
             optional_artifact: Optional[Dict] = None,
@@ -539,13 +541,27 @@ def prep_task(
         """
 
         element_list = conf_frame['atom_names']
-        pp, orb = abacus_inputs.write_pporb(element_list)
-        dpks = abacus_inputs.write_deepks()
-        mass = abacus_inputs.get_mass(element_list)
-        conf_frame.to('abacus/stru', 'STRU', pp_file=pp,numerical_orbital=orb,numerical_descriptor=dpks,mass=mass)
+        pp, orb = inputs.write_pporb(element_list)
+        dpks = inputs.write_deepks()
+        mass = inputs.get_mass(element_list)
+
+        # if conf_frame has the spins, then we will use it as the initial mag and write it to STRU
+        mag = conf_frame.data.get("spins",None)
+        if mag is not None:
+            mag = mag[0] # spins is the mag of several frames, here we only use the first frame
 
-        abacus_inputs.write_input("INPUT")
-        abacus_inputs.write_kpt("KPT")
+        # if the constrain_elements is set, we will set the related flag in STRU    
+        sc = None
+        c_eles = inputs.get_constrain_elements()
+        if c_eles:
+            atom_names = conf_frame.data["atom_names"]
+            atom_types = [atom_names[i] for i in conf_frame.data["atom_types"]]
+            sc = [None if i not in c_eles else [1,1,1] for i in atom_types]
+
+        conf_frame.to('abacus/stru', 'STRU', pp_file=pp,numerical_orbital=orb,numerical_descriptor=dpks,mass=mass,mag=mag,sc=sc)
+
+        inputs.write_input("INPUT")
+        inputs.write_kpt("KPT")
 
         if optional_artifact:
             for file_name, file_path in optional_artifact.items():
@@ -632,15 +648,17 @@ def run_task(
             command = "abacus"
         # run abacus
         command = " ".join([command, ">", log_name])
-        ret, out, err = run_command(command, raise_error=False, try_bash=True,)
+        ret, out, err = run_command(command, raise_error=False, shell=True,)
         if ret != 0:
             raise TransientError(
                 "abacus failed\n", "out msg", out, "\n", "err msg", err, "\n"
             )
-        if not self.check_run_success(log_name):
-            raise TransientError(
-                "abacus failed , we could not check the exact cause . Please check log file ."
-            )
+        #if not self.check_run_success(log_name):
+        #    raise TransientError(
+        #        "abacus failed , we could not check the exact cause . Please check log file ."
+        #    )
+        if os.path.isdir(backward_dir_name):
+            shutil.rmtree(backward_dir_name)
         os.makedirs(Path(backward_dir_name))
         shutil.copyfile(log_name,Path(backward_dir_name)/log_name)
         for ii in backward_list:

fpop/prep_fp.py

@@ -21,6 +21,7 @@
     Optional,
     Union,
 )
+import dpdata
 
 class PrepFp(OP, ABC):
     r"""Prepares the working directories for first-principles (FP) tasks.
@@ -84,7 +85,7 @@
         pass
 
     @OP.exec_sign_check
     def execute(
             self,
             ip : OPIO,
     ) -> OPIO:
@@ -116,7 +117,7 @@
             - `task_names`: (`List[str]`) The name of tasks. Will be used as the identities of the tasks. The names of different tasks are different.
             - `task_paths`: (`Artifact(List[Path])`) The parepared working paths of the tasks. Contains all input files needed to start the FP. The order fo the Paths should be consistent with `op["task_names"]`
         """
-        import dpdata
+
 
         inputs = ip['inputs']
         confs = ip['confs']
@@ -135,6 +136,7 @@
         #System
         counter = 0
         # loop over list of System
+        self._register_spin()
         for system in confs:
             ss = dpdata.System(system, fmt=conf_format, labeled=False)
             for ff in range(ss.get_nframes()):
@@ -148,6 +150,17 @@
             'task_paths' : task_paths,
         })
 
+    def _register_spin(self):
+        from dpdata.data_type import Axis, DataType
+        import numpy as np
+        dt = DataType(
+            "spins",
+            np.ndarray,
+            (Axis.NFRAMES, Axis.NATOMS, 3),
+            required=False,
+            deepmd_name="spin",
+        )
+        dpdata.System.register_data_type(dt)
 
     def _exec_one_frame(
             self,

fpop/run_fp.py

@@ -107,7 +107,7 @@
         pass
 
     @OP.exec_sign_check
     def execute(
         self,
         ip: OPIO,
     ) -> OPIO:
@@ -162,10 +162,18 @@
                 if not os.path.exists(ii):
                     raise FatalError(f"cannot file file/directory {ii}")
                 iname = ii.name
+                if os.path.isfile(iname):
+                    os.remove(iname)
+                elif os.path.isdir(iname):
+                    shutil.rmtree(iname)
                 Path(iname).symlink_to(ii)
             for ii in opt_input_files:
                 if os.path.exists(ii):
                     iname = ii.name
+                    if os.path.isfile(iname):
+                        os.remove(iname)
+                    elif os.path.isdir(iname):
+                        shutil.rmtree(iname)
                     Path(iname).symlink_to(ii)
             backward_dir_name = self.run_task(backward_dir_name,log_name,backward_list,run_image_config,optional_input)
 

tests/spin/deepmd-spin/type.raw

@@ -0,0 +1 @@
+0

tests/spin/deepmd-spin/type_map.raw

@@ -0,0 +1 @@
+Na

tests/test_abacus_inputs.py

@@ -43,7 +43,7 @@ def test_writefile1(self):
         self.assertTrue(os.path.isfile("KPT"))
         self.assertTrue(os.path.isfile("H.upf"))
         self.assertFalse(os.path.isfile("O.upf"))
-        self.assertFalse(os.path.isfile("H.orb"))
+        self.assertTrue(os.path.isfile("H.orb"))
         self.assertFalse(os.path.isfile("O.orb"))
 
         self.assertEqual("INPUT_PARAMETERS", Path("INPUT").read_text().split("\n")[0])
@@ -52,6 +52,7 @@ def test_writefile1(self):
 
         self.assertEqual(Path("KPT").read_text(),"tkpt")
         self.assertEqual(Path("H.upf").read_text(),"tH.upf")
+        self.assertEqual(Path("H.orb").read_text(),"tH.orb")
 
     def test_writefile2(self):
         abacusinput = AbacusInputs(input_file="../INPUT",
@@ -67,8 +68,8 @@ def test_writefile2(self):
         self.assertTrue(os.path.isfile("KPT"))
         self.assertTrue(os.path.isfile("H.upf"))
         self.assertTrue(os.path.isfile("O.upf"))
-        self.assertFalse(os.path.isfile("H.orb"))
-        self.assertFalse(os.path.isfile("O.orb"))
+        self.assertTrue(os.path.isfile("H.orb"))
+        self.assertTrue(os.path.isfile("O.orb"))
 
     def test_writefile3(self):
         abacusinput = AbacusInputs(input_file="../INPUT",

tests/test_prep_abacus.py

@@ -38,10 +38,11 @@
         )
 from fpop.abacus import PrepAbacus,AbacusInputs
 from typing import List
-from constants import POSCAR_1_content,POSCAR_2_content,dump_conf_from_poscar
+from constants import POSCAR_1_content,POSCAR_2_content,STRU1_content,dump_conf_from_poscar
 upload_packages.append("../fpop")
 upload_packages.append("./context.py")
 
+
 class TestPrepAbacus(unittest.TestCase):
     '''
     deepmd/npy format named ["data.000","data.001"].
@@ -200,4 +201,81 @@ def testWithoutOptionalParam(self):
 
         self.assertEqual(tdirs, step.outputs.parameters['task_names'].value)
 
+class TestPrepAbacusSpin(unittest.TestCase):
+    '''
+    deepmd/npy format named ["data.000","data.001"].
+    no optional_input or optional_artifact.
+    '''
+    def setUp(self):
+        self.source_path = Path('abacustest')
+        self.word_path = Path('task.000')
+        self.source_path.mkdir(parents=True, exist_ok=True)
+        self.word_path.mkdir(parents=True, exist_ok=True)
+
+        stru_path = self.source_path / "STRU_tmp"
+        stru_path.write_text(STRU1_content)
+        self.confs = dpdata.System(str(stru_path), fmt="abacus/stru")
+        self.confs.data["spins"] = [[[1,2,3],[4,5,6],[7,8,9],[1,1,1],
+                                     [2,2,2],[3,3,3],[4,4,4],[5,5,5]]]
+
+        (self.source_path/"INPUT").write_text('INPUT_PARAMETERS\ncalculation scf\nbasis_type lcao\n')
+        (self.source_path/"KPT").write_text('here kpt')
+        (self.source_path/"As.upf").write_text('here As upf')
+        (self.source_path/"Ga.upf").write_text('here Ga upf')
+        (self.source_path/"As.orb").write_text('here As orb')
+        (self.source_path/"Ga.orb").write_text('here Ga orb')
+        (self.source_path/'optional_test').write_text('here test')
+
+        self.abacus_inputs = AbacusInputs(
+            input_file=self.source_path/"INPUT",
+            kpt_file=self.source_path/"KPT",
+            pp_files={"As":self.source_path/"As.upf","Ga":self.source_path/"Ga.upf"},
+            orb_files={"As":self.source_path/"As.orb","Ga":self.source_path/"Ga.orb"},
+            constrain_elements=["As"] # only constrain As
+        )
+
+    def tearDown(self):
+        if os.path.isdir(self.source_path):
+            shutil.rmtree(self.source_path)
+        if os.path.isdir(self.word_path):
+            shutil.rmtree(self.word_path)
+
+    def test_prepare_abacus_spin(self):
+        pabacus = PrepAbacus()
+        os.chdir(self.word_path)
+
+        pabacus.prep_task(self.confs, self.abacus_inputs)
+
+        self.assertTrue(os.path.isfile('INPUT'))
+        self.assertTrue(os.path.isfile('KPT'))
+        self.assertTrue(os.path.isfile('STRU'))
+        self.assertTrue(os.path.isfile('As.upf'))
+        self.assertTrue(os.path.isfile('Ga.upf'))
+        self.assertTrue(os.path.isfile('As.orb'))
+        self.assertTrue(os.path.isfile('Ga.orb'))
+        self.assertEqual(Path('INPUT').read_text().split()[0],"INPUT_PARAMETERS")      
+        self.assertEqual(Path('KPT').read_text(),'here kpt')
+        self.assertEqual(Path('As.upf').read_text(),'here As upf')
+        self.assertEqual(Path('Ga.upf').read_text(),'here Ga upf')
+        self.assertEqual(Path('As.orb').read_text(),'here As orb')
+        self.assertEqual(Path('Ga.orb').read_text(),'here Ga orb')
+
+        with open("STRU") as f : c = f.read()
+        #print(c)
+        ref_c = '''Ga
+0.0
+4
+0.000000000000 0.000000000000 0.000000000000 1 1 1 mag 1.000000000000 2.000000000000 3.000000000000
+0.000000000000 2.875074596069 2.875074596069 1 1 1 mag 4.000000000000 5.000000000000 6.000000000000
+2.875074596069 0.000000000000 2.875074596069 1 1 1 mag 7.000000000000 8.000000000000 9.000000000000
+2.875074596069 2.875074596069 0.000000000000 1 1 1 mag 1.000000000000 1.000000000000 1.000000000000
+As
+0.0
+4
+1.437537298035 1.437537298035 1.437537298035 1 1 1 mag 2.000000000000 2.000000000000 2.000000000000 sc 1 1 1
+1.437537298035 4.312611894104 4.312611894104 1 1 1 mag 3.000000000000 3.000000000000 3.000000000000 sc 1 1 1
+4.312611894104 1.437537298035 4.312611894104 1 1 1 mag 4.000000000000 4.000000000000 4.000000000000 sc 1 1 1
+4.312611894104 4.312611894104 1.437537298035 1 1 1 mag 5.000000000000 5.000000000000 5.000000000000 sc 1 1 1'''
+        self.assertTrue(ref_c in c)
+        os.chdir("../")