diff --git a/fpop/abacus.py b/fpop/abacus.py index 81e87e2..696fef3 100644 --- a/fpop/abacus.py +++ b/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: diff --git a/fpop/prep_fp.py b/fpop/prep_fp.py index 4cf3dd5..d22e727 100644 --- a/fpop/prep_fp.py +++ b/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. @@ -116,7 +117,7 @@ def execute( - `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 @@ def execute( #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 @@ def execute( '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, diff --git a/fpop/run_fp.py b/fpop/run_fp.py index 87ac1c3..798d552 100644 --- a/fpop/run_fp.py +++ b/fpop/run_fp.py @@ -162,10 +162,18 @@ def execute( 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) diff --git a/tests/spin/deepmd-spin/set.000/box.npy b/tests/spin/deepmd-spin/set.000/box.npy new file mode 100644 index 0000000..7428e80 Binary files /dev/null and b/tests/spin/deepmd-spin/set.000/box.npy differ diff --git a/tests/spin/deepmd-spin/set.000/coord.npy b/tests/spin/deepmd-spin/set.000/coord.npy new file mode 100644 index 0000000..94c428e Binary files /dev/null and b/tests/spin/deepmd-spin/set.000/coord.npy differ diff --git a/tests/spin/deepmd-spin/set.000/spin.npy b/tests/spin/deepmd-spin/set.000/spin.npy new file mode 100644 index 0000000..96c97f5 Binary files /dev/null and b/tests/spin/deepmd-spin/set.000/spin.npy differ diff --git a/tests/spin/deepmd-spin/type.raw b/tests/spin/deepmd-spin/type.raw new file mode 100644 index 0000000..573541a --- /dev/null +++ b/tests/spin/deepmd-spin/type.raw @@ -0,0 +1 @@ +0 diff --git a/tests/spin/deepmd-spin/type_map.raw b/tests/spin/deepmd-spin/type_map.raw new file mode 100644 index 0000000..1f7b9ab --- /dev/null +++ b/tests/spin/deepmd-spin/type_map.raw @@ -0,0 +1 @@ +Na diff --git a/tests/test_abacus_inputs.py b/tests/test_abacus_inputs.py index 1a58399..be87c0d 100644 --- a/tests/test_abacus_inputs.py +++ b/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", diff --git a/tests/test_prep_abacus.py b/tests/test_prep_abacus.py index 7c69c86..708b705 100644 --- a/tests/test_prep_abacus.py +++ b/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("../") diff --git a/tests/test_prep_abacus_spin.py b/tests/test_prep_abacus_spin.py new file mode 100644 index 0000000..55fbc87 --- /dev/null +++ b/tests/test_prep_abacus_spin.py @@ -0,0 +1,84 @@ +import os +import unittest +from dflow.python import ( + OPIO, +) + +import shutil +from pathlib import Path +from fpop.abacus import PrepAbacus,AbacusInputs + + +class TestPrepAbacus(unittest.TestCase): + ''' + deepmd/npy format named ["data.000","data.001"]. + no optional_input or optional_artifact. + ''' + def setUp(self): + # config in spin/deepmd-spin is contants.POSCAR_1_content, and add spins = [[[1,2,3]]] + # here to test if prep_fp can read the spins data. + self.confs = [Path("spin/deepmd-spin")] + self.ntasks = len(self.confs) + self.type_map = ['Na'] + + self.source_path = Path('abacustest') + self.source_path.mkdir(parents=True, exist_ok=True) + (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/"Na.upf").write_text('here upf') + (self.source_path/"Na.orb").write_text('here 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={"Na":self.source_path/"Na.upf"}, + orb_files={"Na":self.source_path/"Na.orb"} + ) + + def tearDown(self): + for ii in range(self.ntasks): + work_path = Path("task.%06d"%ii) + if work_path.is_dir(): + shutil.rmtree(work_path) + if os.path.isdir(self.source_path): + shutil.rmtree(self.source_path) + + def checkfile(self): + tdir = "task.000000" + print(tdir) + self.assertTrue(Path(tdir).is_dir()) + self.assertTrue(os.path.isfile(Path(tdir)/'INPUT')) + self.assertTrue(os.path.isfile(Path(tdir)/'KPT')) + self.assertTrue(os.path.isfile(Path(tdir)/'STRU')) + self.assertTrue(os.path.isfile(Path(tdir)/'Na.upf')) + self.assertTrue(os.path.isfile(Path(tdir)/'Na.orb')) + self.assertEqual((Path(tdir)/'INPUT').read_text().split()[0],"INPUT_PARAMETERS") + self.assertEqual((Path(tdir)/'KPT').read_text(),'here kpt') + self.assertEqual((Path(tdir)/'Na.upf').read_text(),'here upf') + self.assertEqual((Path(tdir)/'Na.orb').read_text(),'here orb') + stru_c = (Path(tdir)/'STRU').read_text() + self.assertTrue("0.000000000000 0.000000000000 0.000000000000 1 1 1 mag 1.000000000000 2.000000000000 3.000000000000" in stru_c) + return [tdir] + + def test_prepare(self): + op = PrepAbacus() + out = op.execute( + OPIO( + { + "prep_image_config" : {}, + "optional_artifact" : {"TEST": (self.source_path/'optional_test').absolute()}, + "confs" : self.confs, + "inputs" : self.abacus_inputs, + "type_map" : self.type_map, + } + ) + ) + + tdirs = self.checkfile() + + self.assertEqual(tdirs, out['task_names']) + self.assertEqual(tdirs, [str(ii) for ii in out['task_paths']]) + + for ii in out['task_names']: + self.assertEqual(Path(Path(ii)/'TEST').read_text(), "here test")