From b341ae74e99917ceb51ef87d8b8c379f1ce7c2dc Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Fri, 11 Oct 2024 10:04:18 +0800
Subject: [PATCH 01/15] add spin for lammps

---
 dpdata/lammps/dump.py     | 51 +++++++++++++++++++++++++++
 dpdata/lammps/lmp.py      | 72 +++++++++++++++++++++++++++++++++++----
 dpdata/plugins/lammps.py  | 22 ++++++++++--
 tests/lammps/spin.lmp     | 12 +++++++
 tests/lammps/traj.dump    | 52 ++++++++++++++++++++++++++++
 tests/test_lammps_spin.py | 54 +++++++++++++++++++++++++++++
 6 files changed, 254 insertions(+), 9 deletions(-)
 create mode 100644 tests/lammps/spin.lmp
 create mode 100644 tests/lammps/traj.dump
 create mode 100644 tests/test_lammps_spin.py

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 72fee27df..42aad00ea 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,6 +196,42 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
+def get_spin(lines):
+    blk, head = _get_block(lines, "ATOMS")
+    heads = head.split()
+    """
+    the spin info is stored in c_spin[1], c_spin[2], c_spin[3], c_spin[4] or sp, spx, spy, spz, which is the spin norm and the spin vector
+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    
+    """
+    key1 = ["sp", "spx", "spy", "spz"]
+    key2 = ["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"]
+    
+    # check if head contains spin info
+    if all(i in heads for i in key1):
+        key = key1
+    elif all(i in heads for i in key2):
+        key = key2
+    else:
+        return None
+    idx_id = heads.index("id") - 2
+    idx_sp = heads.index(key[0]) - 2
+    idx_spx = heads.index(key[1]) - 2
+    idx_spy = heads.index(key[2]) - 2
+    idx_spz = heads.index(key[3]) - 2
+    
+    norm = []
+    vec = []
+    id = []
+    for ii in blk:
+        words = ii.split()
+        norm.append([float(words[idx_sp])])
+        vec.append([float(words[idx_spx]), float(words[idx_spy]), float(words[idx_spz])])
+        id.append(int(words[idx_id]))
+    
+    spin = np.array(norm) * np.array(vec)
+    id, spin = zip(*sorted(zip(id, spin)))
+    return np.array(spin)
 
 def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
     array_lines = split_traj(lines)
@@ -216,6 +252,11 @@ 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 = get_spin(lines)
+    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 +269,16 @@ 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])
+            if spin is not None:
+                system["spins"].append(spin[idx])
+            else:
+                print(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
index 604b18d12..c7aa06862 100644
--- a/dpdata/lammps/lmp.py
+++ b/dpdata/lammps/lmp.py
@@ -126,6 +126,18 @@ 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:
+        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 = []
@@ -168,6 +180,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 +233,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
index 0327d444b..7faccb73f 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -4,13 +4,27 @@
 
 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
 
+import numpy as np
+
 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):
@@ -18,7 +32,9 @@ class LAMMPSLmpFormat(Format):
     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.
@@ -48,4 +64,6 @@ def from_system(
         self, file_name, type_map=None, begin=0, step=1, unwrap=False, **kwargs
     ):
         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)
+        register_spin(data)
+        return data
diff --git a/tests/lammps/spin.lmp b/tests/lammps/spin.lmp
new file mode 100644
index 000000000..d662cc2cd
--- /dev/null
+++ 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
new file mode 100644
index 000000000..e1682983a
--- /dev/null
+++ 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/test_lammps_spin.py b/tests/test_lammps_spin.py
new file mode 100644
index 000000000..ba6ca60c7
--- /dev/null
+++ b/tests/test_lammps_spin.py
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import os,shutil
+import unittest
+
+from context import dpdata
+import numpy as np
+
+class TestLmp(unittest.TestCase):
+    def test_dump_input(self):
+        tmp_system = dpdata.System(
+            os.path.join("poscars", "conf.lmp"), type_map=["O", "H"]
+        )
+        tmp_system.data["spins"] = [[[3,4,0],[0,4,3]],[]]
+        tmp_system.to("lammps/lmp", "tmp.lmp")
+        self.assertTrue(os.path.isfile("tmp.lmp"))
+        with open("tmp.lmp") as f:c = f.read()
+        
+        coord_ref = """     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"""
+        self.assertTrue(coord_ref in c)
+        if os.path.isfile("tmp.lmp"):os.remove("tmp.lmp")
+    
+    def test_read_input(self):    
+        # check if dpdata can read the spins
+        tmp_system = dpdata.System("lammps/spin.lmp", fmt="lammps/lmp", type_map=["O", "H"])
+        self.assertTrue((tmp_system.data["spins"][0] == [[3,4,0],[0,4,3]]).all())
+        
+        tmp_system.to(file_name="lammps/dump", fmt="deepmd/npy")
+        self.assertTrue(os.path.isfile("lammps/dump/set.000/spin.npy"))
+        
+        if os.path.isdir("lammps/dump"):shutil.rmtree("lammps/dump")
+        
+class TestDump(unittest.TestCase):
+    def test_read_dump_spin(self):
+        tmp_system = dpdata.System("lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"])
+        self.assertTrue(len(tmp_system.data["spins"]) == 2)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][0][0],[0,0,1.54706291],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][0][1],[0,0,1.54412869],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][0][-2],[0,0,1.65592002],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][0][-1],[0,0,0],decimal=8)
+        
+        np.testing.assert_almost_equal(tmp_system.data["spins"][1][0],[0.21021514724299958 , 1.0123821159859323 , -0.6159960941686954],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][1][1],[1.0057302798645609 , 0.568273899191638 , -0.2363447073875224],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][1][-2],[-0.28075943761984146 , -1.2845200151690905 , -0.0201237855118935],decimal=8)
+        np.testing.assert_almost_equal(tmp_system.data["spins"][1][-1],[0,0,0],decimal=8)
+        
+        tmp_system.to(file_name="lammps/dump", fmt="deepmd/npy")
+        self.assertTrue(os.path.isfile("lammps/dump/set.000/spin.npy"))
+        
+        if os.path.isdir("lammps/dump"):shutil.rmtree("lammps/dump")
+        
+        
+        
\ No newline at end of file

From a9416dcf53569c31494eeb21a65239cf3a93064c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Oct 2024 02:11:00 +0000
Subject: [PATCH 02/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/lammps/dump.py     | 18 +++++---
 dpdata/lammps/lmp.py      |  7 +--
 dpdata/plugins/lammps.py  |  5 ++-
 tests/test_lammps_spin.py | 90 ++++++++++++++++++++++++++-------------
 4 files changed, 79 insertions(+), 41 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 42aad00ea..dba10ea6d 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,17 +196,18 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
+
 def get_spin(lines):
     blk, head = _get_block(lines, "ATOMS")
     heads = head.split()
     """
     the spin info is stored in c_spin[1], c_spin[2], c_spin[3], c_spin[4] or sp, spx, spy, spz, which is the spin norm and the spin vector
 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    
+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
     """
     key1 = ["sp", "spx", "spy", "spz"]
     key2 = ["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"]
-    
+
     # check if head contains spin info
     if all(i in heads for i in key1):
         key = key1
@@ -219,20 +220,23 @@ def get_spin(lines):
     idx_spx = heads.index(key[1]) - 2
     idx_spy = heads.index(key[2]) - 2
     idx_spz = heads.index(key[3]) - 2
-    
+
     norm = []
     vec = []
     id = []
     for ii in blk:
         words = ii.split()
         norm.append([float(words[idx_sp])])
-        vec.append([float(words[idx_spx]), float(words[idx_spy]), float(words[idx_spz])])
+        vec.append(
+            [float(words[idx_spx]), float(words[idx_spy]), float(words[idx_spz])]
+        )
         id.append(int(words[idx_id]))
-    
+
     spin = np.array(norm) * np.array(vec)
     id, spin = zip(*sorted(zip(id, spin)))
     return np.array(spin)
 
+
 def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
     array_lines = split_traj(lines)
     lines = array_lines[0]
@@ -274,7 +278,9 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
             if spin is not None:
                 system["spins"].append(spin[idx])
             else:
-                print(f"Warning: spin info is not found in frame {ii}, remove spin info.")
+                print(
+                    f"Warning: spin info is not found in frame {ii}, remove spin info."
+                )
                 system.pop("spins")
                 has_spin = False
     if has_spin:
diff --git a/dpdata/lammps/lmp.py b/dpdata/lammps/lmp.py
index c7aa06862..046201662 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:
@@ -180,11 +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
 
 
@@ -233,7 +234,7 @@ def from_system_data(system, f_idx=0):
         + ptr_float_fmt
         + "\n"
     )
-    
+
     if "spins" in system:
         coord_fmt = (
             coord_fmt.strip("\n")
diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index 7faccb73f..0b1584ccb 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -2,14 +2,14 @@
 
 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
 
-import numpy as np
-
 if TYPE_CHECKING:
     from dpdata.utils import FileType
 
@@ -25,6 +25,7 @@ def register_spin(data):
         )
         dpdata.System.register_data_type(dt)
 
+
 @Format.register("lmp")
 @Format.register("lammps/lmp")
 class LAMMPSLmpFormat(Format):
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index ba6ca60c7..be4811df7 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -1,54 +1,84 @@
 from __future__ import annotations
 
-import os,shutil
+import os
+import shutil
 import unittest
 
-from context import dpdata
 import numpy as np
+from context import dpdata
+
 
 class TestLmp(unittest.TestCase):
     def test_dump_input(self):
         tmp_system = dpdata.System(
             os.path.join("poscars", "conf.lmp"), type_map=["O", "H"]
         )
-        tmp_system.data["spins"] = [[[3,4,0],[0,4,3]],[]]
+        tmp_system.data["spins"] = [[[3, 4, 0], [0, 4, 3]], []]
         tmp_system.to("lammps/lmp", "tmp.lmp")
         self.assertTrue(os.path.isfile("tmp.lmp"))
-        with open("tmp.lmp") as f:c = f.read()
-        
+        with open("tmp.lmp") as f:
+            c = f.read()
+
         coord_ref = """     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"""
         self.assertTrue(coord_ref in c)
-        if os.path.isfile("tmp.lmp"):os.remove("tmp.lmp")
-    
-    def test_read_input(self):    
+        if os.path.isfile("tmp.lmp"):
+            os.remove("tmp.lmp")
+
+    def test_read_input(self):
         # check if dpdata can read the spins
-        tmp_system = dpdata.System("lammps/spin.lmp", fmt="lammps/lmp", type_map=["O", "H"])
-        self.assertTrue((tmp_system.data["spins"][0] == [[3,4,0],[0,4,3]]).all())
-        
+        tmp_system = dpdata.System(
+            "lammps/spin.lmp", fmt="lammps/lmp", type_map=["O", "H"]
+        )
+        self.assertTrue((tmp_system.data["spins"][0] == [[3, 4, 0], [0, 4, 3]]).all())
+
         tmp_system.to(file_name="lammps/dump", fmt="deepmd/npy")
         self.assertTrue(os.path.isfile("lammps/dump/set.000/spin.npy"))
-        
-        if os.path.isdir("lammps/dump"):shutil.rmtree("lammps/dump")
-        
+
+        if os.path.isdir("lammps/dump"):
+            shutil.rmtree("lammps/dump")
+
+
 class TestDump(unittest.TestCase):
     def test_read_dump_spin(self):
-        tmp_system = dpdata.System("lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"])
+        tmp_system = dpdata.System(
+            "lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"]
+        )
         self.assertTrue(len(tmp_system.data["spins"]) == 2)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][0][0],[0,0,1.54706291],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][0][1],[0,0,1.54412869],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][0][-2],[0,0,1.65592002],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][0][-1],[0,0,0],decimal=8)
-        
-        np.testing.assert_almost_equal(tmp_system.data["spins"][1][0],[0.21021514724299958 , 1.0123821159859323 , -0.6159960941686954],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][1][1],[1.0057302798645609 , 0.568273899191638 , -0.2363447073875224],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][1][-2],[-0.28075943761984146 , -1.2845200151690905 , -0.0201237855118935],decimal=8)
-        np.testing.assert_almost_equal(tmp_system.data["spins"][1][-1],[0,0,0],decimal=8)
-        
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][0][0], [0, 0, 1.54706291], decimal=8
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][0][1], [0, 0, 1.54412869], decimal=8
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][0][-2], [0, 0, 1.65592002], decimal=8
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][0][-1], [0, 0, 0], decimal=8
+        )
+
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][1][0],
+            [0.21021514724299958, 1.0123821159859323, -0.6159960941686954],
+            decimal=8,
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][1][1],
+            [1.0057302798645609, 0.568273899191638, -0.2363447073875224],
+            decimal=8,
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][1][-2],
+            [-0.28075943761984146, -1.2845200151690905, -0.0201237855118935],
+            decimal=8,
+        )
+        np.testing.assert_almost_equal(
+            tmp_system.data["spins"][1][-1], [0, 0, 0], decimal=8
+        )
+
         tmp_system.to(file_name="lammps/dump", fmt="deepmd/npy")
         self.assertTrue(os.path.isfile("lammps/dump/set.000/spin.npy"))
-        
-        if os.path.isdir("lammps/dump"):shutil.rmtree("lammps/dump")
-        
-        
-        
\ No newline at end of file
+
+        if os.path.isdir("lammps/dump"):
+            shutil.rmtree("lammps/dump")

From 589d8bfbee65b00e2b4301b72091892959ed75ea Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Tue, 15 Oct 2024 13:01:28 +0800
Subject: [PATCH 03/15] fix bug

---
 dpdata/lammps/dump.py     | 41 ++++++++++++++++++++++++++++-----------
 dpdata/plugins/lammps.py  |  4 ++--
 tests/lammps/in.lmp       |  2 ++
 tests/test_lammps_spin.py | 30 ++++++++++++++++++++--------
 4 files changed, 56 insertions(+), 21 deletions(-)
 create mode 100644 tests/lammps/in.lmp

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 42aad00ea..6c2dd9b48 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,22 +196,40 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
-def get_spin(lines):
-    blk, head = _get_block(lines, "ATOMS")
-    heads = head.split()
+def get_spin_keys(inputfile):
+    # raed input file and get the keys for spin info in dump
+    # 1. find the index of "compute X X X sp spx spy spz ..." in the input file
+    # 2. construct the keys for spin info: c_X[idx],...
+    if inputfile is not None and os.path.isfile(inputfile):
+        with open(inputfile) as f:
+            lines = f.readlines()
+        for line in lines:
+            ls = line.split()
+            if len(ls) > 7 and ls[0] == "compute" and "sp" in ls and "spx" in ls and "spy" in ls and "spz" in ls:
+                idx_sp = "c_" + ls[1] + "[" + str(ls.index("sp")-3) + "]"
+                idx_spx = "c_" + ls[1] + "[" + str(ls.index("spx")-3) + "]"
+                idx_spy = "c_" + ls[1] + "[" + str(ls.index("spy")-3) + "]"
+                idx_spz = "c_" + ls[1] + "[" + str(ls.index("spz")-3) + "]"
+                return [idx_sp, idx_spx, idx_spy, idx_spz]
+    return None
+
+def get_spin(lines,spin_keys):
     """
-    the spin info is stored in c_spin[1], c_spin[2], c_spin[3], c_spin[4] or sp, spx, spy, spz, which is the spin norm and the spin vector
+    the spin info is stored in sp, spx, spy, spz or spin_keys, which is the spin norm and the spin vector
+    spin_keys may be like: c_spin[1], c_spin[2], c_spin[3], c_spin[4]
 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    
     """
+    blk, head = _get_block(lines, "ATOMS")
+    heads = head.split()
+    
     key1 = ["sp", "spx", "spy", "spz"]
-    key2 = ["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"]
     
     # check if head contains spin info
     if all(i in heads for i in key1):
         key = key1
-    elif all(i in heads for i in key2):
-        key = key2
+    elif spin_keys is not None and all(i in heads for i in spin_keys):
+        key = spin_keys
     else:
         return None
     idx_id = heads.index("id") - 2
@@ -233,7 +251,7 @@ def get_spin(lines):
     id, spin = zip(*sorted(zip(id, spin)))
     return np.array(spin)
 
-def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False):
+def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False, input_name=None):
     array_lines = split_traj(lines)
     lines = array_lines[0]
     system = {}
@@ -252,7 +270,8 @@ 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 = get_spin(lines)
+    spin_keys = get_spin_keys(input_name)
+    spin = get_spin(lines,spin_keys)
     has_spin = False
     if spin is not None:
         system["spins"] = [spin]
@@ -270,11 +289,11 @@ 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 has_spin:
-            spin = get_spin(array_lines[ii])
+            spin = get_spin(array_lines[ii],spin_keys)
             if spin is not None:
                 system["spins"].append(spin[idx])
             else:
-                print(f"Warning: spin info is not found in frame {ii}, remove spin info.")
+                warnings.warn(f"Warning: spin info is not found in frame {ii}, remove spin info.")
                 system.pop("spins")
                 has_spin = False
     if has_spin:
diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index 7faccb73f..658bc9108 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -61,9 +61,9 @@ 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, type_map=None, begin=0, step=1, unwrap=False, input_name=None, **kwargs
     ):
         lines = dpdata.lammps.dump.load_file(file_name, begin=begin, step=step)
-        data = dpdata.lammps.dump.system_data(lines, type_map, unwrap=unwrap)
+        data = dpdata.lammps.dump.system_data(lines, type_map, unwrap=unwrap,input_name=input_name)
         register_spin(data)
         return data
diff --git a/tests/lammps/in.lmp b/tests/lammps/in.lmp
new file mode 100644
index 000000000..02f655fa5
--- /dev/null
+++ 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]
\ No newline at end of file
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index ba6ca60c7..7a9e82515 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -7,19 +7,33 @@
 import numpy as np
 
 class TestLmp(unittest.TestCase):
-    def test_dump_input(self):
-        tmp_system = dpdata.System(
+    def setUp(self):
+        self.tmp_system = dpdata.System(
             os.path.join("poscars", "conf.lmp"), type_map=["O", "H"]
         )
-        tmp_system.data["spins"] = [[[3,4,0],[0,4,3]],[]]
-        tmp_system.to("lammps/lmp", "tmp.lmp")
-        self.assertTrue(os.path.isfile("tmp.lmp"))
-        with open("tmp.lmp") as f:c = f.read()
+        self.tmp_system.data["spins"] = [[[3,4,0],[0,4,3]]]
+        self.lmp_coord_name = "tmp.lmp"
+    
+    def tearDown(self):
+        pass #if os.path.isfile(self.lmp_coord_name):os.remove(self.lmp_coord_name)
+    
+    def test_dump_input(self):
+        self.tmp_system.to("lammps/lmp", self.lmp_coord_name)
+        self.assertTrue(os.path.isfile(self.lmp_coord_name))
+        with open(self.lmp_coord_name) as f:c = f.read()
         
         coord_ref = """     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"""
         self.assertTrue(coord_ref in c)
-        if os.path.isfile("tmp.lmp"):os.remove("tmp.lmp")
+    
+    def test_dump_input_zero_spin(self):
+        self.tmp_system.data["spins"] = [[[0,0,0],[0,0,0]]]
+        self.tmp_system.to("lammps/lmp", self.lmp_coord_name)
+        self.assertTrue(os.path.isfile(self.lmp_coord_name))
+        with open(self.lmp_coord_name) as f:c = f.read()
+        coord_ref = """     1      1    0.0000000000    0.0000000000    0.0000000000    0.0000000000    0.0000000000    1.0000000000    0.0000000000
+     2      2    1.2621856000    0.7018028000    0.5513885000    0.0000000000    0.0000000000    1.0000000000    0.0000000000"""
+        self.assertTrue(coord_ref in c)
     
     def test_read_input(self):    
         # check if dpdata can read the spins
@@ -33,7 +47,7 @@ def test_read_input(self):
         
 class TestDump(unittest.TestCase):
     def test_read_dump_spin(self):
-        tmp_system = dpdata.System("lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"])
+        tmp_system = dpdata.System("lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"], input_name="lammps/in.lmp")
         self.assertTrue(len(tmp_system.data["spins"]) == 2)
         np.testing.assert_almost_equal(tmp_system.data["spins"][0][0],[0,0,1.54706291],decimal=8)
         np.testing.assert_almost_equal(tmp_system.data["spins"][0][1],[0,0,1.54412869],decimal=8)

From 67450ce79af1bf7ae117ae3f8d0e5d321de23bdf Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 15 Oct 2024 05:03:31 +0000
Subject: [PATCH 04/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/lammps/dump.py     | 42 ++++++++++++++++++++++++++-------------
 dpdata/plugins/lammps.py  | 13 ++++++++++--
 tests/test_lammps_spin.py | 31 ++++++++++++++++++-----------
 3 files changed, 58 insertions(+), 28 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 9a44869fa..2d4f4ca25 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,6 +196,7 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
+
 def get_spin_keys(inputfile):
     # raed input file and get the keys for spin info in dump
     # 1. find the index of "compute X X X sp spx spy spz ..." in the input file
@@ -205,26 +206,34 @@ def get_spin_keys(inputfile):
             lines = f.readlines()
         for line in lines:
             ls = line.split()
-            if len(ls) > 7 and ls[0] == "compute" and "sp" in ls and "spx" in ls and "spy" in ls and "spz" in ls:
-                idx_sp = "c_" + ls[1] + "[" + str(ls.index("sp")-3) + "]"
-                idx_spx = "c_" + ls[1] + "[" + str(ls.index("spx")-3) + "]"
-                idx_spy = "c_" + ls[1] + "[" + str(ls.index("spy")-3) + "]"
-                idx_spz = "c_" + ls[1] + "[" + str(ls.index("spz")-3) + "]"
+            if (
+                len(ls) > 7
+                and ls[0] == "compute"
+                and "sp" in ls
+                and "spx" in ls
+                and "spy" in ls
+                and "spz" in ls
+            ):
+                idx_sp = "c_" + ls[1] + "[" + str(ls.index("sp") - 3) + "]"
+                idx_spx = "c_" + ls[1] + "[" + str(ls.index("spx") - 3) + "]"
+                idx_spy = "c_" + ls[1] + "[" + str(ls.index("spy") - 3) + "]"
+                idx_spz = "c_" + ls[1] + "[" + str(ls.index("spz") - 3) + "]"
                 return [idx_sp, idx_spx, idx_spy, idx_spz]
     return None
 
-def get_spin(lines,spin_keys):
+
+def get_spin(lines, spin_keys):
     """
     the spin info is stored in sp, spx, spy, spz or spin_keys, which is the spin norm and the spin vector
     spin_keys may be like: c_spin[1], c_spin[2], c_spin[3], c_spin[4]
-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
+    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
     """
     blk, head = _get_block(lines, "ATOMS")
     heads = head.split()
-    
+
     key1 = ["sp", "spx", "spy", "spz"]
-    
+
     # check if head contains spin info
     if all(i in heads for i in key1):
         key = key1
@@ -253,7 +262,10 @@ def get_spin(lines,spin_keys):
     id, spin = zip(*sorted(zip(id, spin)))
     return np.array(spin)
 
-def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False, input_name=None):
+
+def system_data(
+    lines, type_map=None, type_idx_zero=True, unwrap=False, input_name=None
+):
     array_lines = split_traj(lines)
     lines = array_lines[0]
     system = {}
@@ -273,7 +285,7 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False, input_na
     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_name)
-    spin = get_spin(lines,spin_keys)
+    spin = get_spin(lines, spin_keys)
     has_spin = False
     if spin is not None:
         system["spins"] = [spin]
@@ -291,11 +303,13 @@ def system_data(lines, type_map=None, type_idx_zero=True, unwrap=False, input_na
             safe_get_posi(array_lines[ii], cell, np.array(orig), unwrap)[idx]
         )
         if has_spin:
-            spin = get_spin(array_lines[ii],spin_keys)
+            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.")
+                warnings.warn(
+                    f"Warning: spin info is not found in frame {ii}, remove spin info."
+                )
                 system.pop("spins")
                 has_spin = False
     if has_spin:
diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index 1d75143f2..d44e1cafd 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -62,9 +62,18 @@ 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, input_name=None, **kwargs
+        self,
+        file_name,
+        type_map=None,
+        begin=0,
+        step=1,
+        unwrap=False,
+        input_name=None,
+        **kwargs,
     ):
         lines = dpdata.lammps.dump.load_file(file_name, begin=begin, step=step)
-        data = dpdata.lammps.dump.system_data(lines, type_map, unwrap=unwrap,input_name=input_name)
+        data = dpdata.lammps.dump.system_data(
+            lines, type_map, unwrap=unwrap, input_name=input_name
+        )
         register_spin(data)
         return data
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index 2a08de30a..b040f39bb 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -13,31 +13,33 @@ def setUp(self):
         self.tmp_system = dpdata.System(
             os.path.join("poscars", "conf.lmp"), type_map=["O", "H"]
         )
-        self.tmp_system.data["spins"] = [[[3,4,0],[0,4,3]]]
+        self.tmp_system.data["spins"] = [[[3, 4, 0], [0, 4, 3]]]
         self.lmp_coord_name = "tmp.lmp"
-    
+
     def tearDown(self):
-        pass #if os.path.isfile(self.lmp_coord_name):os.remove(self.lmp_coord_name)
-    
+        pass  # if os.path.isfile(self.lmp_coord_name):os.remove(self.lmp_coord_name)
+
     def test_dump_input(self):
         self.tmp_system.to("lammps/lmp", self.lmp_coord_name)
         self.assertTrue(os.path.isfile(self.lmp_coord_name))
-        with open(self.lmp_coord_name) as f:c = f.read()
-        
+        with open(self.lmp_coord_name) as f:
+            c = f.read()
+
         coord_ref = """     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"""
         self.assertTrue(coord_ref in c)
-    
+
     def test_dump_input_zero_spin(self):
-        self.tmp_system.data["spins"] = [[[0,0,0],[0,0,0]]]
+        self.tmp_system.data["spins"] = [[[0, 0, 0], [0, 0, 0]]]
         self.tmp_system.to("lammps/lmp", self.lmp_coord_name)
         self.assertTrue(os.path.isfile(self.lmp_coord_name))
-        with open(self.lmp_coord_name) as f:c = f.read()
+        with open(self.lmp_coord_name) as f:
+            c = f.read()
         coord_ref = """     1      1    0.0000000000    0.0000000000    0.0000000000    0.0000000000    0.0000000000    1.0000000000    0.0000000000
      2      2    1.2621856000    0.7018028000    0.5513885000    0.0000000000    0.0000000000    1.0000000000    0.0000000000"""
         self.assertTrue(coord_ref in c)
-    
-    def test_read_input(self):    
+
+    def test_read_input(self):
         # check if dpdata can read the spins
         tmp_system = dpdata.System(
             "lammps/spin.lmp", fmt="lammps/lmp", type_map=["O", "H"]
@@ -53,7 +55,12 @@ def test_read_input(self):
 
 class TestDump(unittest.TestCase):
     def test_read_dump_spin(self):
-        tmp_system = dpdata.System("lammps/traj.dump", fmt="lammps/dump", type_map=["O", "H"], input_name="lammps/in.lmp")
+        tmp_system = dpdata.System(
+            "lammps/traj.dump",
+            fmt="lammps/dump",
+            type_map=["O", "H"],
+            input_name="lammps/in.lmp",
+        )
         self.assertTrue(len(tmp_system.data["spins"]) == 2)
         np.testing.assert_almost_equal(
             tmp_system.data["spins"][0][0], [0, 0, 1.54706291], decimal=8

From 204a8ac421fc3c403416df380b5e9caa7fce4d70 Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Tue, 15 Oct 2024 18:01:04 +0800
Subject: [PATCH 05/15] fix doc

---
 dpdata/lammps/dump.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 9a44869fa..d2087e6fd 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -215,6 +215,15 @@ def get_spin_keys(inputfile):
 
 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
     spin_keys may be like: c_spin[1], c_spin[2], c_spin[3], c_spin[4]
 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]

From ad5808a83920b738be7fcbfc6aeebd6cab87778a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 15 Oct 2024 10:02:50 +0000
Subject: [PATCH 06/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/lammps/dump.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 25824dab3..0d29828af 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -225,17 +225,14 @@ def get_spin_keys(inputfile):
 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
-    spin_keys may be like: c_spin[1], c_spin[2], c_spin[3], c_spin[4]
-    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
     """
     blk, head = _get_block(lines, "ATOMS")

From daa38ba96bb27d92379332660a1e72271cd9cafc Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Thu, 17 Oct 2024 11:29:11 +0800
Subject: [PATCH 07/15] improve codes

---
 dpdata/lammps/dump.py | 85 ++++++++++++++++++++++---------------------
 1 file changed, 43 insertions(+), 42 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 0d29828af..c82b01c72 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,29 +196,31 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
-
 def get_spin_keys(inputfile):
-    # raed input file and get the keys for spin info in dump
-    # 1. find the index of "compute X X X sp spx spy spz ..." in the input file
-    # 2. construct the keys for spin info: c_X[idx],...
-    if inputfile is not None and os.path.isfile(inputfile):
-        with open(inputfile) as f:
-            lines = f.readlines()
-        for line in lines:
+    """
+    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 or not os.path.isfile(inputfile):
+        return None
+
+    with open(inputfile) as f:
+        for line in f.readlines():
             ls = line.split()
-            if (
-                len(ls) > 7
-                and ls[0] == "compute"
-                and "sp" in ls
-                and "spx" in ls
-                and "spy" in ls
-                and "spz" in ls
-            ):
-                idx_sp = "c_" + ls[1] + "[" + str(ls.index("sp") - 3) + "]"
-                idx_spx = "c_" + ls[1] + "[" + str(ls.index("spx") - 3) + "]"
-                idx_spy = "c_" + ls[1] + "[" + str(ls.index("spy") - 3) + "]"
-                idx_spz = "c_" + ls[1] + "[" + str(ls.index("spz") - 3) + "]"
-                return [idx_sp, idx_spx, idx_spy, idx_spz]
+            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
 
 
@@ -235,39 +237,38 @@ def get_spin(lines, spin_keys):
     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()
 
     key1 = ["sp", "spx", "spy", "spz"]
 
-    # check if head contains spin info
     if all(i in heads for i in key1):
         key = key1
     elif spin_keys is not None and all(i in heads for i in spin_keys):
         key = spin_keys
     else:
         return None
-    idx_id = heads.index("id") - 2
-    idx_sp = heads.index(key[0]) - 2
-    idx_spx = heads.index(key[1]) - 2
-    idx_spy = heads.index(key[2]) - 2
-    idx_spz = heads.index(key[3]) - 2
-
-    norm = []
-    vec = []
-    id = []
-    for ii in blk:
-        words = ii.split()
-        norm.append([float(words[idx_sp])])
-        vec.append(
-            [float(words[idx_spx]), float(words[idx_spy]), float(words[idx_spz])]
-        )
-        id.append(int(words[idx_id]))
-
-    spin = np.array(norm) * np.array(vec)
-    id, spin = zip(*sorted(zip(id, spin)))
-    return np.array(spin)
 
+    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_name=None

From 866f14bfd26bc07c12c964a2ec776a456e8d86b2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 17 Oct 2024 03:29:24 +0000
Subject: [PATCH 08/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/lammps/dump.py | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index c82b01c72..d010367a7 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -196,17 +196,18 @@ def load_file(fname: FileType, begin=0, step=1):
             if cc >= begin and (cc - begin) % step == 0:
                 buff.append(line)
 
+
 def get_spin_keys(inputfile):
     """
     Read input file and get the keys for spin info in dump.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     inputfile : str
         Path to the input file.
 
-    Returns:
-    --------
+    Returns
+    -------
     list or None
         List of spin info keys if found, None otherwise.
     """
@@ -216,10 +217,16 @@ def get_spin_keys(inputfile):
     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"])):
+            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 [
+                    f"c_{compute_name}[{ls.index(key) - 3}]"
+                    for key in ["sp", "spx", "spy", "spz"]
+                ]
 
     return None
 
@@ -237,7 +244,6 @@ def get_spin(lines, spin_keys):
     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()
 
@@ -252,7 +258,7 @@ def get_spin(lines, spin_keys):
 
     try:
         idx_id = heads.index("id") - 2
-        idx_sp, idx_spx, idx_spy, idx_spz = [heads.index(k) - 2 for k in key]
+        idx_sp, idx_spx, idx_spy, idx_spz = (heads.index(k) - 2 for k in key)
 
         norm = []
         vec = []
@@ -260,7 +266,9 @@ def get_spin(lines, spin_keys):
         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])])
+            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)
@@ -270,6 +278,7 @@ def get_spin(lines, spin_keys):
         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_name=None
 ):

From b1439dfd900c5f545c53164e0998b99dce63c963 Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Fri, 18 Oct 2024 15:50:55 +0800
Subject: [PATCH 09/15] modify input_name to input_file

---
 dpdata/lammps/dump.py     |  4 ++--
 dpdata/plugins/lammps.py  | 37 ++++++++++++++++++++++++++++++-------
 tests/test_lammps_spin.py |  2 +-
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index c82b01c72..97399b836 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -271,7 +271,7 @@ def get_spin(lines, spin_keys):
         return None
 
 def system_data(
-    lines, type_map=None, type_idx_zero=True, unwrap=False, input_name=None
+    lines, type_map=None, type_idx_zero=True, unwrap=False, input_file=None
 ):
     array_lines = split_traj(lines)
     lines = array_lines[0]
@@ -291,7 +291,7 @@ def system_data(
     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_name)
+    spin_keys = get_spin_keys(input_file)
     spin = get_spin(lines, spin_keys)
     has_spin = False
     if spin is not None:
diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index d44e1cafd..b07dc69fe 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -9,6 +9,7 @@
 from dpdata.data_type import Axis, DataType
 from dpdata.format import Format
 from dpdata.utils import open_file
+from typing import List, Dict
 
 if TYPE_CHECKING:
     from dpdata.utils import FileType
@@ -63,17 +64,39 @@ class LAMMPSDumpFormat(Format):
     @Format.post("shift_orig_zero")
     def from_system(
         self,
-        file_name,
-        type_map=None,
-        begin=0,
-        step=1,
-        unwrap=False,
-        input_name=None,
+        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)
         data = dpdata.lammps.dump.system_data(
-            lines, type_map, unwrap=unwrap, input_name=input_name
+            lines, type_map, unwrap=unwrap, input_file=input_file
         )
         register_spin(data)
         return data
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index b040f39bb..7e182f43f 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -59,7 +59,7 @@ def test_read_dump_spin(self):
             "lammps/traj.dump",
             fmt="lammps/dump",
             type_map=["O", "H"],
-            input_name="lammps/in.lmp",
+            input_file="lammps/in.lmp",
         )
         self.assertTrue(len(tmp_system.data["spins"]) == 2)
         np.testing.assert_almost_equal(

From 26f9875192eca0642e35de5b66210f37020c9b85 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 18 Oct 2024 07:51:28 +0000
Subject: [PATCH 10/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/plugins/lammps.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index b07dc69fe..d265c53e3 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 import numpy as np
 
@@ -9,7 +9,6 @@
 from dpdata.data_type import Axis, DataType
 from dpdata.format import Format
 from dpdata.utils import open_file
-from typing import List, Dict
 
 if TYPE_CHECKING:
     from dpdata.utils import FileType
@@ -64,12 +63,12 @@ class LAMMPSDumpFormat(Format):
     @Format.post("shift_orig_zero")
     def from_system(
         self,
-        file_name : str,
-        type_map: List[str] =None,
-        begin: int =0,
-        step: int =1,
-        unwrap: bool =False,
-        input_file: str =None,
+        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.
@@ -88,7 +87,7 @@ def from_system(
             Whether to unwrap the coordinates
         input_file : str, optional
             The input file name
-            
+
         Returns
         -------
         dict

From 596efafb3055a8d8af628e91b55545bdc87dee7c Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Fri, 18 Oct 2024 15:57:36 +0800
Subject: [PATCH 11/15] fix

---
 dpdata/plugins/lammps.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py
index d265c53e3..c7e5c7653 100644
--- a/dpdata/plugins/lammps.py
+++ b/dpdata/plugins/lammps.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -64,7 +64,7 @@ class LAMMPSDumpFormat(Format):
     def from_system(
         self,
         file_name: str,
-        type_map: List[str] = None,
+        type_map: list[str] = None,
         begin: int = 0,
         step: int = 1,
         unwrap: bool = False,

From 0a1bb76e77d099e62d0c1f6e7ed75c015b5f9419 Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Tue, 22 Oct 2024 10:09:05 +0800
Subject: [PATCH 12/15] add UT

---
 dpdata/lammps/dump.py               |  6 +---
 tests/lammps/traj_partial_spin.dump | 52 +++++++++++++++++++++++++++++
 tests/test_lammps_spin.py           | 48 ++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 5 deletions(-)
 create mode 100644 tests/lammps/traj_partial_spin.dump

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index b81ec5f18..6bfcf0305 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -247,11 +247,7 @@ def get_spin(lines, spin_keys):
     blk, head = _get_block(lines, "ATOMS")
     heads = head.split()
 
-    key1 = ["sp", "spx", "spy", "spz"]
-
-    if all(i in heads for i in key1):
-        key = key1
-    elif spin_keys is not None and all(i in heads for i in spin_keys):
+    if spin_keys is not None and all(i in heads for i in spin_keys):
         key = spin_keys
     else:
         return None
diff --git a/tests/lammps/traj_partial_spin.dump b/tests/lammps/traj_partial_spin.dump
new file mode 100644
index 000000000..5ee20855a
--- /dev/null
+++ b/tests/lammps/traj_partial_spin.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 
+1 1 0.00037565 5.64900783 0.00994919 
+2 1 5.65299480 5.57370279 2.84182058 
+3 1 0.00076668 2.78053566 0.01181481 
+4 1 5.65188602 2.85285383 2.84413423 
+5 1 2.82101290 5.64942265 0.01091135 
+6 1 2.89400594 5.57477971 2.84333235 
+7 1 2.82260306 2.78052696 0.00917962 
+8 1 2.89632273 2.85545549 2.84070353 
+9 1 1.41117683 1.38867159 1.72266429 
+10 1 1.41176544 1.38873530 3.96470435
+11 1 1.41085716 4.21342650 1.43850987
+12 1 1.41088045 4.21340876 4.24487134
+13 1 4.27525485 1.38812593 1.43912039
+14 1 4.27483864 1.38696457 4.24782541
+15 1 4.27528588 4.21463764 1.39334117
+16 1 4.27407834 4.21327842 4.29226033
+17 2 1.40675897 1.38612182 0.01000617
diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index 7e182f43f..5bdc0506f 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -6,6 +6,34 @@
 
 import numpy as np
 from context import dpdata
+from dpdata.lammps.dump import get_spin
+
+TRAJ_NO_ID = """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 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 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
+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
+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
+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
+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
+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
+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
+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
+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
+ 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
+ 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
+ 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
+ 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
+ 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
+ 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
+ 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
+ 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"""
 
 
 class TestLmp(unittest.TestCase):
@@ -99,3 +127,23 @@ def test_read_dump_spin(self):
 
         if os.path.isdir("lammps/dump"):
             shutil.rmtree("lammps/dump")
+        
+    def test_read_dump_partial_spin(self):
+        # test if dpdata can read the spins when the spin data is not complete
+        with self.assertWarns(UserWarning) as cm:
+            tmp_system = dpdata.System(
+                "lammps/traj_partial_spin.dump",
+                fmt="lammps/dump",
+                type_map=["O", "H"],
+                input_file="lammps/in.lmp",
+            )
+            self.assertTrue("spins" not in tmp_system.data)
+            
+        self.assertIn("Warning: spin info is not found in frame", str(cm.warning))
+    
+    def test_get_spin_failed(self):
+        with self.assertWarns(UserWarning) as cm:
+            spin = get_spin(TRAJ_NO_ID.split("\n"),["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"])
+            self.assertTrue(spin is None)
+            
+        self.assertIn("Error processing spin data:", str(cm.warning))
\ No newline at end of file

From 1c758bd939c3c10777b3584039cc25e7881e1c0c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 22 Oct 2024 02:09:20 +0000
Subject: [PATCH 13/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_lammps_spin.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py
index 5bdc0506f..d3d58920e 100644
--- a/tests/test_lammps_spin.py
+++ b/tests/test_lammps_spin.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 from context import dpdata
+
 from dpdata.lammps.dump import get_spin
 
 TRAJ_NO_ID = """ITEM: TIMESTEP
@@ -127,7 +128,7 @@ def test_read_dump_spin(self):
 
         if os.path.isdir("lammps/dump"):
             shutil.rmtree("lammps/dump")
-        
+
     def test_read_dump_partial_spin(self):
         # test if dpdata can read the spins when the spin data is not complete
         with self.assertWarns(UserWarning) as cm:
@@ -138,12 +139,15 @@ def test_read_dump_partial_spin(self):
                 input_file="lammps/in.lmp",
             )
             self.assertTrue("spins" not in tmp_system.data)
-            
+
         self.assertIn("Warning: spin info is not found in frame", str(cm.warning))
-    
+
     def test_get_spin_failed(self):
         with self.assertWarns(UserWarning) as cm:
-            spin = get_spin(TRAJ_NO_ID.split("\n"),["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"])
+            spin = get_spin(
+                TRAJ_NO_ID.split("\n"),
+                ["c_spin[1]", "c_spin[2]", "c_spin[3]", "c_spin[4]"],
+            )
             self.assertTrue(spin is None)
-            
-        self.assertIn("Error processing spin data:", str(cm.warning))
\ No newline at end of file
+
+        self.assertIn("Error processing spin data:", str(cm.warning))

From 78bab7e22c77d91e76d85a23bca06612f2ba63d6 Mon Sep 17 00:00:00 2001
From: root <pxlxingliang>
Date: Wed, 23 Oct 2024 17:33:52 +0800
Subject: [PATCH 14/15] throw warnings if inputfile is non-exist

---
 dpdata/lammps/dump.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index 6bfcf0305..a5dbc959e 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -211,9 +211,13 @@ def get_spin_keys(inputfile):
     list or None
         List of spin info keys if found, None otherwise.
     """
-    if inputfile is None or not os.path.isfile(inputfile):
+    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()

From ef386da401a6ff90b12f74a0853fffff9fd2eb9a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 23 Oct 2024 09:34:38 +0000
Subject: [PATCH 15/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 dpdata/lammps/dump.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py
index a5dbc959e..91f0c48dd 100644
--- a/dpdata/lammps/dump.py
+++ b/dpdata/lammps/dump.py
@@ -211,13 +211,13 @@ def get_spin_keys(inputfile):
     list or None
         List of spin info keys if found, None otherwise.
     """
-    if inputfile is None: 
+    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()