Adding a function that could deal with the magnetic_force and spin when using the ABACUS's deltaspin module

dpdata/abacus/scf.py

@@ -157,6 +157,44 @@ def get_energy(outlines):
     return Etot, False
 
 
+def get_magnetic(outlines, natoms):
+    """
+    Extracts magnetic moments from the output lines.
+
+    Parameters
+    ----------
+    outlines (list of str) : The lines from the output file.
+    natoms (int) : The number of atoms.
+
+    Returns
+    -------
+    np.ndarray: An array of magnetic moments.
+    """
+    Mag = []
+
+    # 定义正则表达式以匹配所需的数字
+    mag_pattern = re.compile(
+        r"Total Magnetism on atom:.*\(([\d.\-]+),\s*([\d.\-]+),\s*([\d.\-]+)\)"
+    )
+
+    for line in outlines:
+        # 使用正则表达式匹配磁性信息
+        match = mag_pattern.search(line)
+        if match:
+            # 提取匹配的数字并转换为浮点数，然后将元组转换为列表
+            Mag.append([float(num) for num in match.groups()])
+            # 如果Mag的长度达到了natoms，说明已经收集完所有原子的磁矩
+            if len(Mag) == natoms:
+                break
+
+    # 如果没有找到任何磁矩信息，返回一个空的NumPy数组
+    if len(Mag) == 0:
+        return np.array([[]])
+    else:
+        # 将Mag转换为NumPy数组并返回
+        return np.array(Mag)
+
+
 def collect_force(outlines):
     force = []
     for i, line in enumerate(outlines):
@@ -193,6 +231,65 @@ def get_force(outlines, natoms):
         return np.array(force[-1])  # only return the last force
 
 
+def collect_mag_force(outlines):
+    """
+    Collects magnetic forces from the output lines.
+
+    Parameters
+    ----------
+    outlines : The lines from the output file.
+
+    Returns
+    -------
+    list: A list of magnetic forces.
+    """
+    mag_force = []
+    for i, line in enumerate(outlines):
+        if "Magnetic force (Ry/uB)" in line:
+            value_pattern = re.compile(
+                r"^\s*ATOM\s+(\d+)\s+[-+]?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]+([eE][-+]?[0-9]+)?\s+[-+]?[0-9]+([eE][-+]?[0-9]+)?\s*$"
+            )
+            j = i
+            # find the first line of force
+            noforce = False
+            while not value_pattern.match(outlines[j]):
+                j += 1
+                if (
+                    j >= i + 10
+                ):  # if can not find the first line of force in 10 lines, then stop
+                    warnings.warn("Warning: can not find the first line of force")
+                    noforce = True
+                    break
+            if noforce:
+                break
+
+            mag_force.append([])
+            while value_pattern.match(outlines[j]):
+                mag_force[-1].append([float(ii) for ii in outlines[j].split()[1:4]])
+                j += 1
+    return mag_force  # only return the last force
+
+
+def get_mag_force(outlines, natoms):
+    """
+    Extracts magnetic forces from the output lines.
+
+    Parameters
+    ----------
+    outlines (list of str) : The lines from the output file.
+    natoms (int) : The number of atoms.
+
+    Returns
+    -------
+    np.ndarray: An array of magnetic forces.
+    """
+    mag_force = collect_mag_force(outlines)
+    if len(mag_force) == 0:
+        return [[]]
+    else:
+        return np.array(mag_force[-1])  # only return the last force
+
+
 def collect_stress(outlines):
     stress = []
     for i, line in enumerate(outlines):
@@ -230,6 +327,34 @@ def get_stress(outlines):
         return np.array(stress[-1]) * kbar2evperang3  # only return the last stress
 
 
+def check_deltaspin(path_in):
+    """
+    Checks if the deltaspin feature is enabled in the input file.
+
+    Parameters
+    ----------
+    path_in : The path to the input file.
+
+    Returns
+    -------
+    bool: True if deltaspin is enabled, None otherwise.
+    """
+    try:
+        with open(path_in) as file:
+            for line in file:
+                # 移除行首尾的空白字符，然后以空格分割
+                parts = line.strip().split()
+                # 检查是否是sc_mag_switch参数行
+                if len(parts) >= 2 and parts[0] == "sc_mag_switch":
+                    # 检查参数值是否为1
+                    return True if parts[1] == "1" else None
+        # 文件已读完，没有找到sc_mag_switch参数，返回None
+        return None
+    except FileNotFoundError:
+        print(f"File not found: {path_in}")
+        return None
+
+
 def get_frame(fname):
     data = {
         "atom_names": [],
@@ -239,6 +364,11 @@ def get_frame(fname):
         "coords": [],
         "energies": [],
         "forces": [],
+        "spin": [],
+        "mag_forces": [],
+        "coords_deltaspin": [],
+        "force_deltaspin": [],
+        "deltaspin": [],
     }
 
     if isinstance(fname, str):
@@ -256,6 +386,7 @@ def get_frame(fname):
 
     geometry_path_in = get_geometry_in(fname, inlines)
     path_out = get_path_out(fname, inlines)
+    deltaspin = check_deltaspin(path_in)
     if not (CheckFile(geometry_path_in) and CheckFile(path_out)):
         return data
 
@@ -280,10 +411,36 @@ def get_frame(fname):
     if stress is not None:
         stress *= np.abs(np.linalg.det(cell))
 
+    if deltaspin is not None:
+        spin = get_magnetic(outlines, natoms)
+        sp_norm = 1.49
+        virtual_len = 0.3
+        mag_forces = get_mag_force(outlines, natoms)
+        coords_deltaspin = np.hstack((coords, coords - spin / sp_norm * virtual_len))
+        force_deltaspin = np.hstack((force, mag_forces))
+    else:
+        mag_forces = None
+        coords_deltaspin = None
+        force_deltaspin = None
+
+    # print(force_deltaspin)
+    # print(coords_deltaspin)
+    # print(spin)
+    # print(mag_forces)
+    # print(force)
+    # print(coords)
+
     data["cells"] = cell[np.newaxis, :, :]
     data["coords"] = coords[np.newaxis, :, :]
     data["energies"] = np.array(energy)[np.newaxis]
     data["forces"] = force[np.newaxis, :, :]
+    if deltaspin is not None:
+        data["mag_forces"] = mag_forces[np.newaxis, :, :]
+        data["spin"] = spin[np.newaxis, :, :]
+        data["coords_deltaspin"] = coords_deltaspin[np.newaxis, :, :]
+        data["force_deltaspin"] = force_deltaspin[np.newaxis, :, :]
+        data["deltaspin"] = deltaspin
+
     if stress is not None:
         data["virials"] = stress[np.newaxis, :, :]
     data["orig"] = np.zeros(3)
@@ -295,6 +452,10 @@ def get_frame(fname):
     # print("energy = ", data['energies'])
     # print("force = ", data['forces'])
     # print("virial = ", data['virials'])
+    # print("spin = ", data['spin'])
+    # print("mag_forces = ", data['mag_forces'])
+    # print("force_deltaspin = ", data['force_deltaspin'])
+    # print("coords_deltaspin = ", data['coords_deltaspin'])
     return data
 
 

dpdata/deepmd/comp.py

@@ -133,13 +133,24 @@ def dump(folder, data, set_size=5000, comp_prec=np.float32, remove_sets=True):
     nframes = data["cells"].shape[0]
     cells = np.reshape(data["cells"], [nframes, 9]).astype(comp_prec)
     coords = np.reshape(data["coords"], [nframes, -1]).astype(comp_prec)
+    if "coords_deltaspin" in data:
+        coords_deltaspin = np.reshape(data["coords_deltaspin"], [nframes, -1]).astype(
+            comp_prec
+        )
     eners = None
     forces = None
     virials = None
+    force_deltaspin = None
+    if "deltaspin" in data:
+        deltaspin = data["deltaspin"]
     if "energies" in data:
         eners = np.reshape(data["energies"], [nframes]).astype(comp_prec)
     if "forces" in data:
         forces = np.reshape(data["forces"], [nframes, -1]).astype(comp_prec)
+    if "force_deltaspin" in data:
+        force_deltaspin = np.reshape(data["force_deltaspin"], [nframes, -1]).astype(
+            comp_prec
+        )
     if "virials" in data:
         virials = np.reshape(data["virials"], [nframes, 9]).astype(comp_prec)
     if "atom_pref" in data:
@@ -154,11 +165,21 @@ def dump(folder, data, set_size=5000, comp_prec=np.float32, remove_sets=True):
         set_folder = os.path.join(folder, "set.%03d" % ii)
         os.makedirs(set_folder)
         np.save(os.path.join(set_folder, "box"), cells[set_stt:set_end])
-        np.save(os.path.join(set_folder, "coord"), coords[set_stt:set_end])
+        if deltaspin is not None:
+            np.save(
+                os.path.join(set_folder, "coord"), coords_deltaspin[set_stt:set_end]
+            )
+        else:
+            np.save(os.path.join(set_folder, "coord"), coords[set_stt:set_end])
         if eners is not None:
             np.save(os.path.join(set_folder, "energy"), eners[set_stt:set_end])
         if forces is not None:
-            np.save(os.path.join(set_folder, "force"), forces[set_stt:set_end])
+            if deltaspin is not None:
+                np.save(
+                    os.path.join(set_folder, "force"), force_deltaspin[set_stt:set_end]
+                )
+            else:
+                np.save(os.path.join(set_folder, "force"), forces[set_stt:set_end])
         if virials is not None:
             np.save(os.path.join(set_folder, "virial"), virials[set_stt:set_end])
         if "atom_pref" in data:
@@ -185,6 +206,9 @@ def dump(folder, data, set_size=5000, comp_prec=np.float32, remove_sets=True):
             "energies",
             "forces",
             "virials",
+            "force_deltaspin",
+            "coords_deltaspin",
+            "deltaspin",
         ):
             # skip as these data contains specific rules
             continue