add spin for lammps (#738)

pxlxingliang · pre-commit-ci[bot] · wanghan-iapcm · web-flow · commit 7853b7b526a5 · 2024-10-23T15:47:04.000-04:00
Reference #728 to add the spin information for lammps.  ## Summary by CodeRabbit - **New Features** - Introduced functions to extract and normalize spin data from LAMMPS dump files. - Added support for registering spin data types in the system. - Created new simulation input and output files for LAMMPS. - Added new compute and dump commands for calculating and outputting spin properties in LAMMPS. - **Tests** - Implemented unit tests for verifying the handling of spin configurations in LAMMPS files, including scenarios with zero spin values and incomplete data.  --------- Co-authored-by: root <pxlxingliang> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
@@ -197,7 +197,91 @@ def load_file(fname: FileType, begin=0, step=1):
                 buff.append(line)
 
 
-def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
+def get_spin_keys(inputfile):
+    """
+    Read input file and get the keys for spin info in dump.
+
+    Parameters
+    ----------
+    inputfile : str
+        Path to the input file.
+
+    Returns
+    -------
+    list or None
+        List of spin info keys if found, None otherwise.
+    """
+    if inputfile is None:
+        return None
+
+    if not os.path.isfile(inputfile):
+        warnings.warn(f"Input file {inputfile} not found.")
+        return None
+
+    with open(inputfile) as f:
+        for line in f.readlines():
+            ls = line.split()
+            if (
+                len(ls) > 7
+                and ls[0] == "compute"
+                and all(key in ls for key in ["sp", "spx", "spy", "spz"])
+            ):
+                compute_name = ls[1]
+                return [
+                    f"c_{compute_name}[{ls.index(key) - 3}]"
+                    for key in ["sp", "spx", "spy", "spz"]
+                ]
+
+    return None
+
+
+def get_spin(lines, spin_keys):
+    """
+    Get the spin info from the dump file.
+
+    Parameters
+    ----------
+    lines : list
+        The content of the dump file.
+    spin_keys : list
+        The keys for spin info in dump file.
+    the spin info is stored in sp, spx, spy, spz or spin_keys, which is the spin norm and the spin vector
+    1 1 0.00141160 5.64868599 0.01005602 1.54706291 0.00000000 0.00000000 1.00000000 -1.40772100 -2.03739417 -1522.64797384 -0.00397809 -0.00190426 -0.00743976
+    """
+    blk, head = _get_block(lines, "ATOMS")
+    heads = head.split()
+
+    if spin_keys is not None and all(i in heads for i in spin_keys):
+        key = spin_keys
+    else:
+        return None
+
+    try:
+        idx_id = heads.index("id") - 2
+        idx_sp, idx_spx, idx_spy, idx_spz = (heads.index(k) - 2 for k in key)
+
+        norm = []
+        vec = []
+        atom_ids = []
+        for line in blk:
+            words = line.split()
+            norm.append([float(words[idx_sp])])
+            vec.append(
+                [float(words[idx_spx]), float(words[idx_spy]), float(words[idx_spz])]
+            )
+            atom_ids.append(int(words[idx_id]))
+
+        spin = np.array(norm) * np.array(vec)
+        atom_ids, spin = zip(*sorted(zip(atom_ids, spin)))
+        return np.array(spin)
+    except (ValueError, IndexError) as e:
+        warnings.warn(f"Error processing spin data: {str(e)}")
+        return None
+
+
+def system_data(
+    lines, type_map=None, type_idx_zero=True, unwrap=False, input_file=None
+):
     array_lines = split_traj(lines)
     lines = array_lines[0]
     system = {}
@@ -216,6 +300,12 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
     system["cells"] = [np.array(cell)]
     system["atom_types"] = get_atype(lines, type_idx_zero=type_idx_zero)
     system["coords"] = [safe_get_posi(lines, cell, np.array(orig), unwrap)]
+    spin_keys = get_spin_keys(input_file)
+    spin = get_spin(lines, spin_keys)
+    has_spin = False
+    if spin is not None:
+        system["spins"] = [spin]
+        has_spin = True
     for ii in range(1, len(array_lines)):
         bounds, tilt = get_dumpbox(array_lines[ii])
         orig, cell = dumpbox2box(bounds, tilt)
@@ -228,6 +318,18 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
         system["coords"].append(
             safe_get_posi(array_lines[ii], cell, np.array(orig), unwrap)[idx]
         )
+        if has_spin:
+            spin = get_spin(array_lines[ii], spin_keys)
+            if spin is not None:
+                system["spins"].append(spin[idx])
+            else:
+                warnings.warn(
+                    f"Warning: spin info is not found in frame {ii}, remove spin info."
+                )
+                system.pop("spins")
+                has_spin = False
+    if has_spin:
+        system["spins"] = np.array(system["spins"])
     system["cells"] = np.array(system["cells"])
     system["coords"] = np.array(system["coords"])
     return system
diff --git a/dpdata/lammps/lmp.py b/dpdata/lammps/lmp.py
@@ -127,6 +127,19 @@ def get_posi(lines):
     return np.array(posis)
 
 
+def get_spins(lines):
+    atom_lines = get_atoms(lines)
+    if len(atom_lines[0].split()) < 8:
+        return None
+    spins_ori = []
+    spins_norm = []
+    for ii in atom_lines:
+        iis = ii.split()
+        spins_ori.append([float(jj) for jj in iis[5:8]])
+        spins_norm.append([float(iis[-1])])
+    return np.array(spins_ori) * np.array(spins_norm)
+
+
 def get_lmpbox(lines):
     box_info = []
     tilt = np.zeros(3)
@@ -168,6 +181,11 @@ def system_data(lines, type_map=None, type_idx_zero=True):
     system["coords"] = [get_posi(lines)]
     system["cells"] = np.array(system["cells"])
     system["coords"] = np.array(system["coords"])
+
+    spins = get_spins(lines)
+    if spins is not None:
+        system["spins"] = np.array([spins])
+
     return system
 
 
@@ -216,14 +234,55 @@ def from_system_data(system, f_idx=0):
         + ptr_float_fmt
         + "\n"
     )
-    for ii in range(natoms):
-        ret += coord_fmt % (
-            ii + 1,
-            system["atom_types"][ii] + 1,
-            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],
+
+    if "spins" in system:
+        coord_fmt = (
+            coord_fmt.strip("\n")
+            + " "
+            + ptr_float_fmt
+            + " "
+            + ptr_float_fmt
+            + " "
+            + ptr_float_fmt
+            + " "
+            + ptr_float_fmt
+            + "\n"
         )
+        spins_norm = np.linalg.norm(system["spins"][f_idx], axis=1)
+    for ii in range(natoms):
+        if "spins" in system:
+            if spins_norm[ii] != 0:
+                ret += coord_fmt % (
+                    ii + 1,
+                    system["atom_types"][ii] + 1,
+                    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],
+                    spins_norm[ii],
+                )
+            else:
+                ret += coord_fmt % (
+                    ii + 1,
+                    system["atom_types"][ii] + 1,
+                    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,
+                    spins_norm[ii],
+                )
+        else:
+            ret += coord_fmt % (
+                ii + 1,
+                system["atom_types"][ii] + 1,
+                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],
+            )
     return ret
 
 
diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
@@ -2,23 +2,40 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
+
 import dpdata.lammps.dump
 import dpdata.lammps.lmp
+from dpdata.data_type import Axis, DataType
 from dpdata.format import Format
 from dpdata.utils import open_file
 
 if TYPE_CHECKING:
     from dpdata.utils import FileType
 
 
+def register_spin(data):
+    if "spins" in data:
+        dt = DataType(
+            "spins",
+            np.ndarray,
+            (Axis.NFRAMES, Axis.NATOMS, 3),
+            required=False,
+            deepmd_name="spin",
+        )
+        dpdata.System.register_data_type(dt)
+
+
 @Format.register("lmp")
 @Format.register("lammps/lmp")
 class LAMMPSLmpFormat(Format):
     @Format.post("shift_orig_zero")
     def from_system(self, file_name: FileType, type_map=None, **kwargs):
         with open_file(file_name) as fp:
             lines = [line.rstrip("\n") for line in fp]
-        return dpdata.lammps.lmp.to_system_data(lines, type_map)
+        data = dpdata.lammps.lmp.to_system_data(lines, type_map)
+        register_spin(data)
+        return data
 
     def to_system(self, data, file_name: FileType, frame_idx=0, **kwargs):
         """Dump the system in lammps data format.
@@ -45,7 +62,40 @@ def to_system(self, data, file_name: FileType, frame_idx=0, **kwargs):
 class LAMMPSDumpFormat(Format):
     @Format.post("shift_orig_zero")
     def from_system(
-        self, file_name, type_map=None, begin=0, step=1, unwrap=False, **kwargs
+        self,
+        file_name: str,
+        type_map: list[str] = None,
+        begin: int = 0,
+        step: int = 1,
+        unwrap: bool = False,
+        input_file: str = None,
+        **kwargs,
     ):
+        """Read the data from a lammps dump file.
+
+        Parameters
+        ----------
+        file_name : str
+            The dump file name
+        type_map : List[str], optional
+            The atom type list
+        begin : int, optional
+            The begin step
+        step : int, optional
+            The step
+        unwrap : bool, optional
+            Whether to unwrap the coordinates
+        input_file : str, optional
+            The input file name
+
+        Returns
+        -------
+        dict
+            The system data
+        """
         lines = dpdata.lammps.dump.load_file(file_name, begin=begin, step=step)
-        return dpdata.lammps.dump.system_data(lines, type_map, unwrap=unwrap)
+        data = dpdata.lammps.dump.system_data(
+            lines, type_map, unwrap=unwrap, input_file=input_file
+        )
+        register_spin(data)
+        return data
diff --git a/tests/lammps/in.lmp b/tests/lammps/in.lmp
@@ -0,0 +1,2 @@
+compute 	spin all property/atom sp spx spy spz fmx fmy fmz fx fy fz
+dump            dpgen_dump all custom 10 traj.dump id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10]
diff --git a/tests/lammps/spin.lmp b/tests/lammps/spin.lmp
@@ -0,0 +1,12 @@
+
+2 atoms
+2 atom types
+   0.0000000000    2.5243712000 xlo xhi
+   0.0000000000    2.0430257000 ylo yhi
+   0.0000000000    2.2254033000 zlo zhi
+   1.2621856000    1.2874292000    0.7485898000 xy xz yz
+
+Atoms # atomic
+
+     1      1    0.0000000000    0.0000000000    0.0000000000    0.6000000000    0.8000000000    0.0000000000    5.0000000000
+     2      2    1.2621856000    0.7018028000    0.5513885000    0.0000000000    0.8000000000    0.6000000000    5.0000000000
diff --git a/tests/lammps/traj.dump b/tests/lammps/traj.dump
@@ -0,0 +1,52 @@
+ITEM: TIMESTEP
+0
+ITEM: NUMBER OF ATOMS
+17
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+-4.0080511965879438e-02 5.7039029418994556e+00 -5.9179115295410201e-03
+1.4436085788922526e-02 5.6674744441011660e+00 -1.1487414836883500e-02
+7.8239288740356017e-03 5.6734038274259646e+00 6.8277359008788905e-04
+ITEM: ATOMS id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10]
+1 1 0.00141160 5.64868599 0.01005602 1.54706291 0.00000000 0.00000000 1.00000000 -1.40772100 -2.03739417 -1522.64797384 -0.00397809 -0.00190426 -0.00743976
+2 1 5.65283939 5.57449025 2.84281508 1.54412869 0.00000000 0.00000000 1.00000000 7.75304092 6.48949619 -1512.84926162 -0.00637234 -0.00733168 0.00661107
+3 1 0.00066480 2.78022036 0.01010716 1.54612979 0.00000000 0.00000000 1.00000000 -0.93618575 1.92206111 -1520.80305011 -0.00316673 0.00177893 -0.00744575
+4 1 5.65233666 2.85374747 2.84289453 1.54439093 0.00000000 0.00000000 1.00000000 8.11012818 -6.49922039 -1514.31557088 -0.00569217 0.00741000 0.00640353
+5 1 2.82063515 5.64869321 0.01007552 1.54714250 0.00000000 0.00000000 1.00000000 2.49070852 -2.14456666 -1523.53038650 0.00478410 -0.00213962 -0.00751154
+6 1 2.89579803 5.57439179 2.84287630 1.54415032 0.00000000 0.00000000 1.00000000 -8.03062338 6.63950296 -1513.41291897 0.00440396 -0.00717185 0.00633657
+7 1 2.82151287 2.78010538 0.01016303 1.54619615 0.00000000 0.00000000 1.00000000 2.71859584 1.98482729 -1521.34149633 0.00533453 0.00194532 -0.00745901
+8 1 2.89637049 2.85377083 2.84297332 1.54440023 0.00000000 0.00000000 1.00000000 -7.76758760 -6.67134514 -1514.43304618 0.00505040 0.00743195 0.00630302
+9 1 1.41106492 1.38817482 1.72302072 1.18134529 0.00000000 0.00000000 1.00000000 0.27170165 -0.00426695 -444.22843899 0.00100237 -0.00002725 -0.03385965
+10 1 1.41105247 1.38807861 3.96314606 1.18153407 0.00000000 0.00000000 1.00000000 -0.07722674 0.01368756 -337.08703133 -0.00066982 0.00007487 0.07887183
+11 1 1.41105864 4.21395432 1.43987180 1.71989299 0.00000000 0.00000000 1.00000000 -0.01511106 0.00320081 -1653.34500916 0.00010421 0.00007248 0.00634401
+12 1 1.41104843 4.21387554 4.24576823 1.71989825 0.00000000 0.00000000 1.00000000 -0.71645898 0.05923960 -1640.68070568 -0.00117959 0.00006676 -0.01467806
+13 1 4.27433865 1.38779084 1.43977211 1.72010048 0.00000000 0.00000000 1.00000000 0.45899480 0.03956420 -1653.36356942 0.00051885 0.00002313 0.00911600
+14 1 4.27436799 1.38772964 4.24586490 1.72010133 0.00000000 0.00000000 1.00000000 0.38385331 0.07301994 -1642.06086017 -0.00002034 0.00010335 -0.01688908
+15 1 4.27435427 4.21452597 1.39359689 1.65590121 0.00000000 0.00000000 1.00000000 -0.01658773 -0.06159007 -1659.12744163 0.00006470 -0.00006420 -0.01342935
+16 1 4.27434583 4.21455469 4.29208004 1.65592002 0.00000000 0.00000000 1.00000000 -0.15590720 -0.03252166 -1654.84697132 -0.00066755 -0.00003915 -0.00482188
+17 2 1.41096761 1.38958048 0.01029027 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 0.00000000 0.00000000 0.00048351 -0.00022876 -0.00645195
+ITEM: TIMESTEP
+10
+ITEM: NUMBER OF ATOMS
+17
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+-4.0080511965879438e-02 5.7039029418994556e+00 -5.9179115295410201e-03
+1.4436085788922526e-02 5.6674744441011660e+00 -1.1487414836883500e-02
+7.8239288740356017e-03 5.6734038274259646e+00 6.8277359008788905e-04
+ITEM: ATOMS id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10]
+1 1 0.00037565 5.64900783 0.00994919 1.20356102 0.17466098 0.84115562 -0.51181127 -77.61471611 -389.41594243 234.29512368 0.00514290 -0.02481576 0.01063015
+2 1 5.65299480 5.57370279 2.84182058 1.17910451 0.85296110 0.48195380 -0.20044424 -311.75775120 -175.76677913 79.45225558 -0.01239308 -0.00914070 0.01933082
+3 1 0.00076668 2.78053566 0.01181481 1.20779106 0.33415542 0.49831517 -0.80001384 -163.88630094 -248.58823709 387.72415159 -0.01738465 -0.02878227 0.01503087
+4 1 5.65188602 2.85285383 2.84413423 1.20124335 -0.83536303 -0.20314716 0.51078356 399.86863784 90.34522869 -236.39221701 0.02327635 -0.00046572 -0.00138388
+5 1 2.82101290 5.64942265 0.01091135 1.34670883 -0.98528016 0.07078135 -0.15560530 902.73741755 -62.52279896 140.44423419 0.01500524 0.00581151 0.00468238
+6 1 2.89400594 5.57477971 2.84333235 1.25424131 -0.94587492 0.11487066 0.30352161 528.43507318 -60.32699018 -171.89948334 -0.00478280 0.00069273 -0.00496159
+7 1 2.82260306 2.78052696 0.00917962 1.17249564 -0.99589145 0.06282562 -0.06521619 374.56568243 -26.39431071 20.98877908 0.01464486 -0.01010131 -0.00993410
+8 1 2.89632273 2.85545549 2.84070353 1.24297017 -0.44008251 -0.42493729 0.79104721 240.05525392 236.02796206 -448.18443804 0.00137705 0.01258804 -0.01817420
+9 1 1.41117683 1.38867159 1.72266429 1.19059484 0.71251804 -0.69714805 -0.07938914 -309.93474514 293.96860716 19.98886311 -0.03871152 0.00854863 -0.02757569
+10 1 1.41176544 1.38873530 3.96470435 1.17564502 -0.51932602 -0.74875017 0.41191463 181.72443401 263.91689829 -132.94216896 0.00122847 0.01674701 0.02707109
+11 1 1.41085716 4.21342650 1.43850987 1.19874662 -0.51890828 0.82913822 0.20800000 237.52969259 -379.65100512 -93.16140268 0.01185986 -0.01872789 0.00032128
+12 1 1.41088045 4.21340876 4.24487134 1.20157661 -0.86390154 -0.04516556 -0.50163154 388.97171693 21.75492170 227.68580658 0.02074490 0.00756366 0.01937948
+13 1 4.27525485 1.38812593 1.43912039 1.23209806 0.55809649 0.81404794 0.16079259 -335.92026314 -484.87463129 -91.14464759 -0.03675759 -0.03549076 0.00310277
+14 1 4.27483864 1.38696457 4.24782541 1.18431742 0.00519166 -0.92210080 0.38691492 -4.73957478 407.09534135 -171.59043210 -0.00911750 0.04394272 -0.01683249
+15 1 4.27528588 4.21463764 1.39334117 1.17456490 -0.93713453 -0.09927163 0.33455046 397.32993706 40.92599847 -141.68618750 0.01918926 -0.00534149 -0.01906574
+16 1 4.27407834 4.21327842 4.29226033 1.31499905 -0.21350543 -0.97682201 -0.01530327 180.98908307 848.25344747 12.36402507 0.00492895 0.04383813 0.00955221
+17 2 1.40675897 1.38612182 0.01000617 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 0.00000000 0.00000000 0.00174929 -0.00686653 -0.01117336
diff --git a/tests/lammps/traj_partial_spin.dump b/tests/lammps/traj_partial_spin.dump
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+compute spin all property/atom sp spx spy spz fmx fmy fmz fx fy fz`
	`2`	`+dump dpgen_dump all custom 10 traj.dump id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10]`