support spin for deepmd format

dpdata/abacus/md.py

@@ -220,7 +220,7 @@ def get_frame(fname):
     if len(magmom) > 0:
         data["spins"] = magmom
     if len(magforce) > 0:
-        data["mag_forces"] = magforce
+        data["force_mags"] = magforce
     if len(move) > 0:
         data["move"] = move
 

dpdata/abacus/relax.py

@@ -217,7 +217,7 @@ def get_frame(fname):
     if len(magmom) > 0:
         data["spins"] = magmom
     if len(magforce) > 0:
-        data["mag_forces"] = magforce
+        data["force_mags"] = magforce
     if len(move) > 0:
         data["move"] = move
 

dpdata/abacus/scf.py

@@ -511,7 +511,7 @@ def get_frame(fname):
     if len(magmom) > 0:
         data["spins"] = magmom
     if len(magforce) > 0:
-        data["mag_forces"] = magforce
+        data["force_mags"] = magforce
     if len(move) > 0:
         data["move"] = move[np.newaxis, :, :]
     # print("atom_names = ", data['atom_names'])

dpdata/plugins/abacus.py

@@ -56,9 +56,9 @@ def register_mag_data(data):
             deepmd_name="spin",
         )
         dpdata.LabeledSystem.register_data_type(dt)
-    if "mag_forces" in data:
+    if "force_mags" in data:
         dt = DataType(
-            "mag_forces",
+            "force_mags",
             np.ndarray,
             (Axis.NFRAMES, Axis.NATOMS, 3),
             required=False,

dpdata/plugins/deepmd.py

@@ -10,17 +10,41 @@
 import dpdata.deepmd.hdf5
 import dpdata.deepmd.mixed
 import dpdata.deepmd.raw
+from dpdata.data_type import Axis, DataType
 from dpdata.driver import Driver
 from dpdata.format import Format
 
 if TYPE_CHECKING:
     import h5py
 
 
+def register_spin():
+    dt = DataType(
+        "spins",
+        np.ndarray,
+        (Axis.NFRAMES, Axis.NATOMS, 3),
+        required=False,
+        deepmd_name="spin",
+    )
+    dpdata.System.register_data_type(dt)
+    dpdata.LabeledSystem.register_data_type(dt)
+
+    dt = DataType(
+        "force_mags",
+        np.ndarray,
+        (Axis.NFRAMES, Axis.NATOMS, 3),
+        required=False,
+        deepmd_name="force_mag",
+    )
+    dpdata.System.register_data_type(dt)
+    dpdata.LabeledSystem.register_data_type(dt)
+
+
 @Format.register("deepmd")
 @Format.register("deepmd/raw")
 class DeePMDRawFormat(Format):
     def from_system(self, file_name, type_map=None, **kwargs):
+        register_spin()
         return dpdata.deepmd.raw.to_system_data(
             file_name, type_map=type_map, labels=False
         )
@@ -30,6 +54,7 @@ def to_system(self, data, file_name, **kwargs):
         dpdata.deepmd.raw.dump(file_name, data)
 
     def from_labeled_system(self, file_name, type_map=None, **kwargs):
+        register_spin()
         return dpdata.deepmd.raw.to_system_data(
             file_name, type_map=type_map, labels=True
         )
@@ -41,6 +66,7 @@ def from_labeled_system(self, file_name, type_map=None, **kwargs):
 @Format.register("deepmd/comp")
 class DeePMDCompFormat(Format):
     def from_system(self, file_name, type_map=None, **kwargs):
+        register_spin()
         return dpdata.deepmd.comp.to_system_data(
             file_name, type_map=type_map, labels=False
         )
@@ -69,6 +95,7 @@ def to_system(self, data, file_name, set_size=5000, prec=np.float64, **kwargs):
         dpdata.deepmd.comp.dump(file_name, data, set_size=set_size, comp_prec=prec)
 
     def from_labeled_system(self, file_name, type_map=None, **kwargs):
+        register_spin()
         return dpdata.deepmd.comp.to_system_data(
             file_name, type_map=type_map, labels=True
         )
@@ -149,6 +176,7 @@ def mix_system(self, *system, type_map, **kwargs):
         return dpdata.deepmd.mixed.mix_system(*system, type_map=type_map, **kwargs)
 
     def from_multi_systems(self, directory, **kwargs):
+        register_spin()
         sys_dir = []
         for root, dirs, files in os.walk(directory):
             if (
@@ -204,6 +232,8 @@ def _from_system(
         """
         import h5py
 
+        register_spin()
+
         if isinstance(file_name, (h5py.Group, h5py.File)):
             return dpdata.deepmd.hdf5.to_system_data(
                 file_name, "", type_map=type_map, labels=labels

tests/test_abacus_spin.py

@@ -31,7 +31,7 @@ def test_scf(self):
             decimal=8,
         )
         np.testing.assert_almost_equal(
-            data["mag_forces"][0],
+            data["force_mags"][0],
             [
                 [-0.0000175013, -0.0000418680, -0.3669618965],
                 [-0.0000161517, -0.0000195198, -0.3669821632],
@@ -47,7 +47,7 @@ def test_scf(self):
         sys2 = dpdata.LabeledSystem(self.dump_path, fmt="deepmd/npy")
         np.testing.assert_almost_equal(data["spins"], sys2.data["spins"], decimal=8)
         np.testing.assert_almost_equal(
-            data["mag_forces"], sys2.data["mag_forces"], decimal=8
+            data["force_mags"], sys2.data["force_mags"], decimal=8
         )
 
     def test_relax(self):
@@ -84,9 +84,9 @@ def test_relax(self):
             ]
         )
         self.assertEqual(len(data["spins"]), 3)
-        self.assertEqual(len(data["mag_forces"]), 3)
+        self.assertEqual(len(data["force_mags"]), 3)
         np.testing.assert_almost_equal(data["spins"], spins_ref, decimal=8)
-        np.testing.assert_almost_equal(data["mag_forces"], magforces_ref, decimal=8)
+        np.testing.assert_almost_equal(data["force_mags"], magforces_ref, decimal=8)
 
         # dump to deepmd-npy
         mysys.to(file_name=self.dump_path, fmt="deepmd/npy")
@@ -96,7 +96,7 @@ def test_relax(self):
         sys2 = dpdata.LabeledSystem(self.dump_path, fmt="deepmd/npy")
         np.testing.assert_almost_equal(data["spins"], sys2.data["spins"], decimal=8)
         np.testing.assert_almost_equal(
-            data["mag_forces"], sys2.data["mag_forces"], decimal=8
+            data["force_mags"], sys2.data["force_mags"], decimal=8
         )
 
     def test_md(self):
@@ -141,9 +141,9 @@ def test_md(self):
             ]
         )
         self.assertEqual(len(data["spins"]), 4)
-        self.assertEqual(len(data["mag_forces"]), 4)
+        self.assertEqual(len(data["force_mags"]), 4)
         np.testing.assert_almost_equal(data["spins"], spins_ref, decimal=8)
-        np.testing.assert_almost_equal(data["mag_forces"], magforces_ref, decimal=8)
+        np.testing.assert_almost_equal(data["force_mags"], magforces_ref, decimal=8)
 
         # dump to deepmd-npy
         mysys.to(file_name=self.dump_path, fmt="deepmd/npy")
@@ -153,5 +153,5 @@ def test_md(self):
         sys2 = dpdata.LabeledSystem(self.dump_path, fmt="deepmd/npy")
         np.testing.assert_almost_equal(data["spins"], sys2.data["spins"], decimal=8)
         np.testing.assert_almost_equal(
-            data["mag_forces"], sys2.data["mag_forces"], decimal=8
+            data["force_mags"], sys2.data["force_mags"], decimal=8
         )

tests/test_deepmd_spin.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import os
+import shutil
+import unittest
+
+from context import dpdata
+
+
+class TestDeepmdReadSpinNPY(unittest.TestCase):
+    def setUp(self):
+        self.tmp_save_path = "tmp.deepmd.spin/dump-tmp"
+
+    def tearDown(self):
+        if os.path.exists(self.tmp_save_path):
+            shutil.rmtree(self.tmp_save_path)
+
+    def check_Fe16(self, system):
+        self.assertTrue("spins" in system.data)
+        self.assertTrue("force_mags" in system.data)
+        self.assertEqual(system.data["spins"].shape, (2, 16, 3))
+        self.assertEqual(system.data["force_mags"].shape, (2, 16, 3))
+
+    def test_read_spin_npy(self):
+        system = dpdata.LabeledSystem("tmp.deepmd.spin/Fe16-npy", fmt="deepmd/npy")
+        self.check_Fe16(system)
+
+        system.to("deepmd/npy", self.tmp_save_path)
+        self.assertTrue(
+            os.path.isfile(os.path.join(self.tmp_save_path, "set.000/spin.npy"))
+        )
+        self.assertTrue(
+            os.path.isfile(os.path.join(self.tmp_save_path, "set.000/force_mag.npy"))
+        )
+
+    def test_read_spin_raw(self):
+        system = dpdata.LabeledSystem("tmp.deepmd.spin/Fe16-raw", fmt="deepmd/raw")
+        self.check_Fe16(system)
+
+        system.to("deepmd/raw", self.tmp_save_path)
+        self.assertTrue(os.path.isfile(os.path.join(self.tmp_save_path, "spin.raw")))
+        self.assertTrue(
+            os.path.isfile(os.path.join(self.tmp_save_path, "force_mag.raw"))
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/tmp.deepmd.spin/Fe16-npy/type.raw

@@ -0,0 +1,16 @@
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0

tests/tmp.deepmd.spin/Fe16-npy/type_map.raw

@@ -0,0 +1 @@
+Fe

tests/tmp.deepmd.spin/Fe16-raw/box.raw

@@ -0,0 +1,2 @@
+4.852406756199998838e+00 0.000000000000000000e+00 0.000000000000000000e+00 2.308035870000000114e-02 4.868205980599999982e+00 0.000000000000000000e+00 2.678336190000000019e-02 -1.273453610000000076e-02 7.575130651200000287e+00
+4.852406756199998838e+00 0.000000000000000000e+00 0.000000000000000000e+00 2.308035870000000114e-02 4.868205980599999982e+00 0.000000000000000000e+00 2.678336190000000019e-02 -1.273453610000000076e-02 7.575130651200000287e+00

tests/tmp.deepmd.spin/Fe16-raw/coord.raw

@@ -0,0 +1,2 @@
+7.265197249999999496e-02 4.801689380600000057e+00 7.557116530000000054e+00 2.502718440599999816e+00 4.848670014499999681e+00 7.528592681200000136e+00 6.829729569999999639e-02 2.503290553899999882e+00 7.542453851200000337e+00 2.387861248100000111e+00 2.392921626100000143e+00 1.227113879999999942e-02 1.248625939999999934e+00 1.235871750000000047e+00 1.962114260000000110e+00 3.579928339999999931e+00 1.308363060000000022e+00 1.944950500000000027e+00 1.261952210000000019e+00 3.670674079999999950e+00 1.937314930000000102e+00 3.648521740000000069e+00 3.517678329999999853e+00 1.842547099999999993e+00 7.917075379999999329e-02 4.015369999999999877e-03 3.783381850000000046e+00 2.364675399999999872e+00 3.114458000000000157e-02 3.819137730000000008e+00 4.877458919999999587e+00 2.472410029999999814e+00 3.767453989999999919e+00 2.441658079999999842e+00 2.449273470000000064e+00 3.787900969999999923e+00 1.207819179999999992e+00 1.185759330000000000e+00 5.700741540000000107e+00 3.672023149999999792e+00 1.266887440000000087e+00 5.690759779999999601e+00 1.225844649999999980e+00 3.688144059999999946e+00 5.677029669999999584e+00 3.729146230000000006e+00 3.577686089999999819e+00 5.653438930000000084e+00
+7.262864250000000688e-02 4.800096520600000360e+00 7.556544780000000294e+00 2.501984930599999934e+00 4.849790654499999576e+00 7.529256091199999723e+00 6.993106570000000299e-02 2.504862153899999999e+00 7.542524761200000150e+00 2.386465468099999931e+00 2.393284846099999807e+00 1.243044880000000046e-02 1.250970219999999911e+00 1.235900269999999912e+00 1.961877150000000070e+00 3.578310469999999910e+00 1.309157640000000011e+00 1.944905080000000064e+00 1.263246410000000042e+00 3.669684970000000046e+00 1.937208040000000020e+00 3.646978709999999957e+00 3.518230940000000029e+00 1.843518699999999955e+00 8.087652380000000252e-02 2.412619999999999869e-03 3.782949939999999955e+00 2.363124420000000114e+00 3.159122000000000319e-02 3.818357489999999910e+00 2.538490380000000060e-02 2.473290130000000087e+00 3.767864050000000020e+00 2.441273620000000033e+00 2.447853410000000007e+00 3.785945680000000202e+00 1.206648970000000043e+00 1.183960899999999983e+00 5.700463349999999707e+00 3.673867459999999863e+00 1.268293499999999963e+00 5.690915719999999567e+00 1.223096860000000063e+00 3.690037439999999780e+00 5.677984109999999696e+00 3.731158119999999911e+00 3.576031940000000020e+00 5.654460040000000021e+00

tests/tmp.deepmd.spin/Fe16-raw/energy.raw

@@ -0,0 +1,2 @@
+-5.454623401600000216e+04
+-5.454620955500000127e+04

tests/tmp.deepmd.spin/Fe16-raw/force.raw

@@ -0,0 +1,2 @@
+-1.106855285000000050e-01 1.000727386199999902e+00 2.883449123999999864e-01 -1.051338108999999993e-01 -2.162821185999999896e-01 1.511491688999999916e-01 -6.572733814000000230e-01 -1.100767982000000034e+00 1.898751854000000072e-01 7.363416830000000246e-01 2.437835339999999962e-01 -6.857924790000000104e-02 -6.825158524000000115e-01 -7.689053999999999334e-02 -4.883741223999999992e-01 9.032410482999999601e-01 -1.615540451700000046e+00 -3.581084193999999887e-01 -4.005809030999999965e-01 5.476169250000000044e-02 -2.907007447000000133e-01 3.671111684000000208e-01 1.718198793700000060e+00 3.327191424000000008e-01 -1.116415134899999995e+00 2.404410170999999874e-01 1.111270488000000028e-01 1.152163319400000052e+00 -2.274326597999999922e-01 -1.102415283000000024e-01 -2.244907613999999885e-01 -3.542942425999999756e-01 2.837310079999999859e-02 2.335184547999999971e-01 2.195733411000000135e-01 3.262157500999999793e-01 4.219752305000000203e-01 5.337641029999999898e-01 -1.105175815000000034e-01 -4.199524229000000108e-01 -8.862139848000000208e-01 -2.222332230000000075e-02 6.769434261000000230e-01 -5.677320351000000320e-01 -1.064142056000000058e-01 -7.742465350999999663e-01 1.033904147000000107e+00 1.273548630999999931e-01
+-1.255811501999999913e-01 1.024865058300000076e+00 3.048681815000000017e-01 -9.745471190000000450e-02 -2.231689940999999899e-01 1.499777952999999941e-01 -6.818135133999999598e-01 -1.119233886199999972e+00 2.084442448000000114e-01 7.586097389999999496e-01 2.483870280000000097e-01 -7.030565599999999415e-02 -7.146105834000000545e-01 -8.766770979999999325e-02 -4.958102687000000008e-01 9.329233330999999740e-01 -1.625196258799999960e+00 -3.685665191999999979e-01 -4.263045145000000091e-01 6.338842760000000520e-02 -2.973771901000000217e-01 3.982496473000000092e-01 1.715933820300000079e+00 3.187726661999999989e-01 -1.134394543700000035e+00 2.659650671000000077e-01 1.065910659000000038e-01 1.171039510700000008e+00 -2.463114885000000087e-01 -9.709863879999999781e-02 -2.318063747999999924e-01 -3.792864573999999989e-01 2.052522660000000165e-02 2.441585086999999987e-01 2.434691978000000090e-01 3.556763707000000219e-01 4.498362350999999815e-01 5.612586234999999979e-01 -1.060362118999999964e-01 -4.527059999999999973e-01 -8.982924847000000046e-01 -2.215639380000000019e-02 7.197915827999999960e-01 -5.931029316000000495e-01 -1.250889162999999915e-01 -8.099371649000000062e-01 1.048992988399999993e+00 1.175842438999999962e-01

tests/tmp.deepmd.spin/Fe16-raw/force_mag.raw

@@ -0,0 +1,2 @@
+-8.534410200000000155e-03 -1.003950009999999941e-02 -1.198012942999999997e-01 1.857438689999999989e-02 7.876794299999999999e-03 -1.158623106999999985e-01 5.790393100000000050e-03 2.902189410000000080e-02 -1.471329148999999947e-01 -1.627021049999999996e-02 -1.243863620000000038e-02 -1.431188442000000116e-01 1.983349000000000136e-03 -1.628727090000000113e-02 1.338718134000000115e-01 -1.675990100000000066e-02 2.088854290000000113e-02 1.289023514000000037e-01 7.555144599999999953e-03 2.097859640000000053e-02 1.455209802000000108e-01 -1.285195899999999917e-03 -5.159665590000000329e-02 1.251078151999999999e-01 1.242615849999999939e-02 5.394605999999999490e-04 -1.211362175000000041e-01 -1.565749859999999993e-02 4.241893000000000052e-03 -1.189902988999999994e-01 -8.991798299999999300e-03 8.944375299999999737e-03 -1.409907068000000041e-01 9.810309000000000130e-04 1.193021600000000043e-03 -1.438089848999999920e-01 -5.407705300000000166e-03 -1.619149939999999996e-02 1.402259667999999904e-01 2.398664899999999878e-03 1.248675119999999999e-02 1.249879150000000050e-01 -1.821928399999999922e-03 1.193867860000000052e-02 1.449831044000000069e-01 1.506308529999999916e-02 -2.018797589999999850e-02 1.371671825999999927e-01
+-8.948930500000000510e-03 -1.040653769999999936e-02 -1.205046651000000052e-01 1.876058129999999893e-02 8.277887200000000334e-03 -1.160853092000000003e-01 6.415928999999999903e-03 2.929689569999999882e-02 -1.481234593999999904e-01 -1.695458160000000089e-02 -1.269999410000000051e-02 -1.433629164000000122e-01 3.804538599999999875e-03 -1.687982010000000158e-02 1.344380966000000011e-01 -1.629795430000000045e-02 2.116517249999999919e-02 1.286579196000000036e-01 8.172215400000000171e-03 2.095696290000000073e-02 1.456517300999999875e-01 -1.081445799999999945e-03 -5.074005670000000107e-02 1.255218261999999918e-01 1.252434070000000080e-02 8.360950000000000239e-04 -1.213422064999999939e-01 -1.578363149999999909e-02 4.818495100000000364e-03 -1.191078504000000066e-01 -8.358033399999999224e-03 8.972258800000000589e-03 -1.413513788999999909e-01 1.125360099999999980e-03 1.306719399999999942e-03 -1.443008279999999921e-01 -5.158918000000000184e-03 -1.691780309999999996e-02 1.409441345999999928e-01 2.871428599999999859e-03 1.337089040000000000e-02 1.251760003999999871e-01 -1.211985000000000059e-03 1.288838999999999957e-02 1.463661454000000084e-01 1.788364250000000141e-02 -2.094214189999999906e-02 1.372846507000000105e-01

tests/tmp.deepmd.spin/Fe16-raw/spin.raw

@@ -0,0 +1,2 @@
+5.563920879999999930e-02 8.026123020000000552e-02 2.258392821700000219e+00 -2.267161628999999901e-01 -8.533636129999999653e-02 2.278170776100000072e+00 -6.958726640000000019e-02 -2.977513280999999901e-01 2.152763507299999990e+00 1.632430896000000020e-01 1.433680660999999967e-01 2.177797364699999960e+00 2.457801589999999980e-02 2.197890640000000062e-01 -2.175622943500000073e+00 2.034624560000000137e-01 -1.423908618999999987e-01 -2.227208061699999853e+00 -6.516286079999999981e-02 -2.002624396000000084e-01 -2.182219526199999926e+00 2.907600839999999881e-02 6.091792472000000513e-01 -2.112683871999999852e+00 -1.055395694000000051e-01 -4.771315269999999720e-02 2.195893971999999916e+00 1.742744213000000097e-01 -8.640444440000000226e-02 2.171807686100000190e+00 1.224203283000000059e-01 -1.428901210999999893e-01 2.129426829199999816e+00 2.808756769999999972e-02 -5.308771360000000272e-02 2.136427688399999969e+00 3.904636110000000299e-02 1.628157172999999958e-01 -2.162790827599999854e+00 -4.121785079999999712e-02 -1.434962392999999947e-01 -2.246669910100000056e+00 -1.821095639999999877e-02 -1.583821810000000108e-01 -2.097885662099999937e+00 -1.871916124000000126e-01 2.405559167000000109e-01 -2.183222719799999822e+00
+6.349788520000000658e-02 8.477241620000000588e-02 2.256907346999999842e+00 -2.246460507999999967e-01 -8.900857460000000110e-02 2.278921277699999859e+00 -7.359221209999999846e-02 -2.977365738000000150e-01 2.152393609499999805e+00 1.737881296000000075e-01 1.472407542999999985e-01 2.176163003700000154e+00 1.668498000000000030e-03 2.266339256000000080e-01 -2.172927181499999971e+00 1.942038004999999956e-01 -1.459203109999999970e-01 -2.226649295900000158e+00 -7.325918500000000444e-02 -2.010864508000000106e-01 -2.177766917800000090e+00 2.292739540000000159e-02 5.983890829000000355e-01 -2.116732264199999936e+00 -1.073844010000000043e-01 -5.292913860000000165e-02 2.194751144900000117e+00 1.736157278000000082e-01 -9.463951419999999670e-02 2.170929507399999903e+00 1.114323799999999975e-01 -1.447396732000000019e-01 2.129971659400000217e+00 2.287651989999999955e-02 -5.641734639999999856e-02 2.136625202099999843e+00 3.100547039999999865e-02 1.706691644999999979e-01 -2.161995885000000062e+00 -5.254033629999999705e-02 -1.529427707999999886e-01 -2.248103780000000107e+00 -2.916460049999999860e-02 -1.674864367999999992e-01 -2.095804000599999828e+00 -2.216720739999999967e-01 2.463171910999999992e-01 -2.178967231200000132e+00

tests/tmp.deepmd.spin/Fe16-raw/type.raw

@@ -0,0 +1,16 @@
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0

tests/tmp.deepmd.spin/Fe16-raw/type_map.raw

@@ -0,0 +1 @@
+Fe

tests/tmp.deepmd.spin/Fe16-raw/virial.raw

@@ -0,0 +1,2 @@
+7.557628675103746474e+00 -3.674192989340970272e-01 -7.797686352528614417e-01 -3.674192989340970272e-01 7.717710405205966850e+00 3.814578682575011093e-01 -7.797686352528614417e-01 3.814578682575011093e-01 7.217849924461329891e+00
+7.708199193940652449e+00 -3.671245997549927309e-01 -7.796781473872188561e-01 -3.671245997549927309e-01 7.854600132463434470e+00 3.768168653463239703e-01 -7.796781473872188561e-01 3.768168653463239703e-01 7.297379300880924902e+00