diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py index 1c1fc246..d3d0e4a8 100644 --- a/dpdata/lammps/dump.py +++ b/dpdata/lammps/dump.py @@ -133,7 +133,7 @@ def safe_get_posi(lines, cell, orig=np.zeros(3), unwrap=False): def get_spintype(keys): - key_sp = ["sp", "spx", "spy", "spz"] + key_sp = ["sp", "spx", "spy", "spz"] key_csp = ["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"] lmp_sp_type = [key_sp, key_csp] for k in range(2): @@ -147,7 +147,7 @@ def safe_get_spin_force(lines): sp_type = get_spintype(keys) assert sp_type is not None, "Dump file does not contain spin!" id_idx = keys.index("id") - 2 - sp = keys.index(sp_type[0]) - 2 + sp = keys.index(sp_type[0]) - 2 spx = keys.index(sp_type[1]) - 2 spy = keys.index(sp_type[2]) - 2 spz = keys.index(sp_type[3]) - 2 @@ -254,10 +254,10 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False): system["coords"] = [safe_get_posi(lines, cell, np.array(orig), unwrap)] contain_spin = False blk, head = _get_block(lines, "ATOMS") - if 'sp' in head: + if "sp" in head: contain_spin = True spin_force = safe_get_spin_force(lines) - system['spins'] = [spin_force[:, :1] * spin_force[:, 1:4]] + system["spins"] = [spin_force[:, :1] * spin_force[:, 1:4]] for ii in range(1, len(array_lines)): bounds, tilt = get_dumpbox(array_lines[ii]) orig, cell = dumpbox2box(bounds, tilt) @@ -271,7 +271,10 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False): safe_get_posi(array_lines[ii], cell, np.array(orig), unwrap)[idx] ) if contain_spin: - system['spins'].append(safe_get_spin_force(array_lines[ii])[:, :1] * safe_get_spin_force(array_lines[ii])[:, 1:4]) + system["spins"].append( + safe_get_spin_force(array_lines[ii])[:, :1] + * safe_get_spin_force(array_lines[ii])[:, 1:4] + ) if contain_spin: system["spins"] = np.array(system["spins"]) system["cells"] = np.array(system["cells"]) diff --git a/dpdata/lammps/lmp.py b/dpdata/lammps/lmp.py index 4f8c1f8f..e78443c3 100644 --- a/dpdata/lammps/lmp.py +++ b/dpdata/lammps/lmp.py @@ -126,6 +126,7 @@ def get_posi(lines): posis.append([float(jj) for jj in ii.split()[2:5]]) return np.array(posis) + def get_spins(lines): atom_lines = get_atoms(lines) if len(atom_lines[0].split()) < 8: @@ -138,6 +139,7 @@ def get_spins(lines): spins = np.array(spins_ori) * np.array(spins_norm) return spins + def get_lmpbox(lines): box_info = [] tilt = np.zeros(3) @@ -163,8 +165,8 @@ def system_data(lines, type_map=None, type_idx_zero=True): system = {} system["atom_numbs"] = get_natoms_vec(lines) spins = get_spins(lines) - if spins is not None: - system['spins'] = np.array([spins]) + if spins is not None: + system["spins"] = np.array([spins]) system["atom_names"] = [] if type_map is None: for ii in range(len(system["atom_numbs"])): @@ -230,9 +232,9 @@ def from_system_data(system, f_idx=0): + ptr_float_fmt + "\n" ) - if 'spins' in system.keys(): + if "spins" in system.keys(): coord_fmt = ( - coord_fmt.strip('\n') + coord_fmt.strip("\n") + " " + ptr_float_fmt + " " @@ -243,9 +245,9 @@ def from_system_data(system, f_idx=0): + ptr_float_fmt + "\n" ) - spins_norm = np.linalg.norm(system['spins'][f_idx], axis=1) + spins_norm = np.linalg.norm(system["spins"][f_idx], axis=1) for ii in range(natoms): - if 'spins' in system.keys(): + if "spins" in system.keys(): if spins_norm[ii] != 0: ret += coord_fmt % ( ii + 1, @@ -253,9 +255,9 @@ def from_system_data(system, f_idx=0): system["coords"][f_idx][ii][0] - system["orig"][0], system["coords"][f_idx][ii][1] - system["orig"][1], system["coords"][f_idx][ii][2] - system["orig"][2], - system['spins'][f_idx][ii][0]/spins_norm[ii], - system['spins'][f_idx][ii][1]/spins_norm[ii], - system['spins'][f_idx][ii][2]/spins_norm[ii], + system["spins"][f_idx][ii][0] / spins_norm[ii], + system["spins"][f_idx][ii][1] / spins_norm[ii], + system["spins"][f_idx][ii][2] / spins_norm[ii], spins_norm[ii], ) else: @@ -265,9 +267,9 @@ def from_system_data(system, f_idx=0): system["coords"][f_idx][ii][0] - system["orig"][0], system["coords"][f_idx][ii][1] - system["orig"][1], system["coords"][f_idx][ii][2] - system["orig"][2], - system['spins'][f_idx][ii][0], - system['spins'][f_idx][ii][1], - system['spins'][f_idx][ii][2]+1, + system["spins"][f_idx][ii][0], + system["spins"][f_idx][ii][1], + system["spins"][f_idx][ii][2] + 1, spins_norm[ii], ) else: diff --git a/dpdata/plugins/vasp_deltaspin.py b/dpdata/plugins/vasp_deltaspin.py index 1638dd63..2a9d74df 100644 --- a/dpdata/plugins/vasp_deltaspin.py +++ b/dpdata/plugins/vasp_deltaspin.py @@ -1,6 +1,9 @@ -import numpy as np +from __future__ import annotations + import re +import numpy as np + import dpdata.vasp_deltaspin.outcar import dpdata.vasp_deltaspin.poscar from dpdata.format import Format @@ -14,7 +17,7 @@ class VASPPoscarFormat(Format): def from_system(self, file_name, **kwargs): with open(file_name) as fp: lines = [line.rstrip("\n") for line in fp] - with open(file_name[:-6] + 'INCAR') as fp: + with open(file_name[:-6] + "INCAR") as fp: lines_incar = [line.rstrip("\n") for line in fp] data = dpdata.vasp_deltaspin.poscar.to_system_data(lines, lines_incar) data = uniq_atom_names(data) @@ -38,10 +41,12 @@ def to_system(self, data, file_name, frame_idx=0, **kwargs): with open(file_name, "w") as fp: fp.write(w_str) - with open(file_name[:-6] + 'INCAR') as fp: + with open(file_name[:-6] + "INCAR") as fp: tmp_incar = fp.read() - res_incar = re.sub(r'MAGMOM[\s\S]*?\n\nM_CONST[\s\S]*?\n\n', m_str, tmp_incar, re.S) - with open(file_name[:-6] + 'INCAR', 'w') as fp: + res_incar = re.sub( + r"MAGMOM[\s\S]*?\n\nM_CONST[\s\S]*?\n\n", m_str, tmp_incar, re.S + ) + with open(file_name[:-6] + "INCAR", "w") as fp: fp.write(res_incar) @@ -78,10 +83,10 @@ def from_labeled_system( data["atom_types"], data["cells"], data["coords"], - data['spins'], + data["spins"], data["energies"], data["forces"], - data['mag_forces'], + data["mag_forces"], tmp_virial, ) = dpdata.vasp_deltaspin.outcar.get_frames( file_name, @@ -100,5 +105,3 @@ def from_labeled_system( data["virials"][ii] *= v_pref * vol data = uniq_atom_names(data) return data - - diff --git a/dpdata/system.py b/dpdata/system.py index 6174458d..2f4414f3 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -1215,7 +1215,9 @@ class LabeledSystem(System): DataType( "forces", np.ndarray, (Axis.NFRAMES, Axis.NATOMS, 3), deepmd_name="force" ), - DataType("mag_forces", np.ndarray, (Axis.NFRAMES, Axis.NATOMS, 3), required=False), + DataType( + "mag_forces", np.ndarray, (Axis.NFRAMES, Axis.NATOMS, 3), required=False + ), DataType( "virials", np.ndarray, diff --git a/dpdata/vasp_deltaspin/outcar.py b/dpdata/vasp_deltaspin/outcar.py index 822e1ef1..192106b1 100644 --- a/dpdata/vasp_deltaspin/outcar.py +++ b/dpdata/vasp_deltaspin/outcar.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import re import warnings @@ -93,11 +95,13 @@ def get_frames(fname, begin=0, step=1, ml=False, convergence_check=True): blk = get_outcar_block(fp, ml) cc += 1 - with open(fname[:-6] + 'OSZICAR') as f: + with open(fname[:-6] + "OSZICAR") as f: oszicar_blk = f.read() - mag_blk = re.findall(r'Magnetic Force(.*?)\n \n', oszicar_blk, re.S)[-1].split('\n') + mag_blk = re.findall(r"Magnetic Force(.*?)\n \n", oszicar_blk, re.S)[-1].split( + "\n" + ) mag_force = [list(map(float, ss.split()[-3:])) for ss in mag_blk[1:]] - spin_blk = re.findall(r'MW_current(.*?)\n ion', oszicar_blk, re.S)[-1].split('\n') + spin_blk = re.findall(r"MW_current(.*?)\n ion", oszicar_blk, re.S)[-1].split("\n") spin = [list(map(float, ss.split()[1:4])) for ss in spin_blk[1:]] all_mag_forces.append(mag_force) all_spins.append(spin) diff --git a/dpdata/vasp_deltaspin/poscar.py b/dpdata/vasp_deltaspin/poscar.py index 932d32c3..7c7905fb 100644 --- a/dpdata/vasp_deltaspin/poscar.py +++ b/dpdata/vasp_deltaspin/poscar.py @@ -1,4 +1,5 @@ #!/usr/bin/python3 +from __future__ import annotations import numpy as np @@ -35,16 +36,16 @@ def _to_system_data_lower(lines, lines_incar, cartesian=True): system["coords"] = np.array(system["coords"]) for idx, incar_param in enumerate(lines_incar): - if 'MAGMOM' in incar_param: + if "MAGMOM" in incar_param: start_index = idx - elif 'M_CONSTR' in incar_param: + elif "M_CONSTR" in incar_param: end_index = idx - system['spins'] = [lines_incar[start_index].replace('=', '').strip().split()[1:4]] - for idx in range(start_index+1, end_index-1): - system['spins'].append(lines_incar[idx].strip().split()[:3]) - system['spins'] = np.array([system['spins']]).astype('float64') - count = np.sum(np.linalg.norm(system['spins'][0], axis=1) > 0) - spins_idx = np.where(np.cumsum(system['atom_numbs']) <= count) + system["spins"] = [lines_incar[start_index].replace("=", "").strip().split()[1:4]] + for idx in range(start_index + 1, end_index - 1): + system["spins"].append(lines_incar[idx].strip().split()[:3]) + system["spins"] = np.array([system["spins"]]).astype("float64") + count = np.sum(np.linalg.norm(system["spins"][0], axis=1) > 0) + spins_idx = np.where(np.cumsum(system["atom_numbs"]) <= count) return system @@ -100,16 +101,20 @@ def from_system_data(system, f_idx=0, skip_zeros=True): magmom_incar = "MAGMOM = " mconstr_incar = "M_CONSTR = " - for idx, tmp_spin in enumerate(system['spins'][f_idx]): + for idx, tmp_spin in enumerate(system["spins"][f_idx]): if idx == 0: - magmom_incar += f"{tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" - mconstr_incar += f"{tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" - elif idx == len(system['spins'][f_idx]) - 1: - magmom_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \n" + magmom_incar += ( + f"{tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" + ) + mconstr_incar += ( + f"{tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" + ) + elif idx == len(system["spins"][f_idx]) - 1: + magmom_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \n" mconstr_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \n" else: - magmom_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" + magmom_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" mconstr_incar += f" {tmp_spin[0]:10.5f} {tmp_spin[1]:10.5f} {tmp_spin[2]:10.5f} \ \n" - magmom_incar += '\n' + mconstr_incar + '\n' + magmom_incar += "\n" + mconstr_incar + "\n" return ret, magmom_incar