add spin for abacus/stru (#751)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced data retrieval to include magnetic moment information in
atomic coordinates.
- Introduced a new function for transforming magnetic moment data into
Cartesian coordinates.
- Added structured input format for iron simulations, detailing atomic
and magnetic properties.

- **Bug Fixes**
- Improved handling of magnetic data registration in various format
classes.

- **Tests**
- Added a new test method to validate the reading of spin data from the
structured input file.
- Updated atomic positions in tests to include magnetic moment
information.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: root <pxlxingliang>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: pxlxingliang

dpdata/abacus/md.py

@@ -167,7 +167,7 @@ def get_frame(fname):
     with open_file(geometry_path_in) as fp:
         geometry_inlines = fp.read().split("\n")
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords, move = get_coords(
+    atom_names, natoms, types, coords, move, magmom = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
     # This coords is not to be used.

dpdata/abacus/relax.py

@@ -183,7 +183,7 @@ def get_frame(fname):
     with open_file(geometry_path_in) as fp:
         geometry_inlines = fp.read().split("\n")
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coord_tmp, move = get_coords(
+    atom_names, natoms, types, coord_tmp, move, magmom = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
 

dpdata/abacus/scf.py

@@ -258,6 +258,51 @@ def parse_stru_pos(pos_line):
     return pos, move, velocity, magmom, angle1, angle2, constrain, lambda1
 
 
+def get_atom_mag_cartesian(atommag, angle1, angle2):
+    """Transform atommag, angle1, angle2 to magmom in cartesian coordinates.
+
+    Parameters
+    ----------
+    atommag : float/list of float/None
+        Atom magnetic moment.
+    angle1 : float/None
+        value of angle1.
+    angle2 : float/None
+        value of angle2.
+    ABACUS support defining mag, angle1, angle2 at the same time.
+    angle1 is the angle between z-axis and real spin (in degrees).
+    angle2 is the angle between x-axis and real spin projection in xy-plane (in degrees).
+    If only mag is defined, then transfer it to magmom directly.
+    And if mag, angle1, angle2 are defined, then mag is only the norm of magmom, and the direction is defined by angle1 and angle2.
+    """
+    if atommag is None:
+        return None
+    if not (isinstance(atommag, list) or isinstance(atommag, float)):
+        raise RuntimeError(f"Invalid atommag: {atommag}")
+
+    if angle1 is None and angle2 is None:
+        if isinstance(atommag, list):
+            return atommag
+        else:
+            return [0, 0, atommag]
+    else:
+        a1 = 0
+        a2 = 0
+        if angle1 is not None:
+            a1 = angle1
+        if angle2 is not None:
+            a2 = angle2
+        if isinstance(atommag, list):
+            mag_norm = np.linalg.norm(atommag)
+        else:
+            mag_norm = atommag
+        return [
+            mag_norm * np.sin(np.radians(a1)) * np.cos(np.radians(a2)),
+            mag_norm * np.sin(np.radians(a1)) * np.sin(np.radians(a2)),
+            mag_norm * np.cos(np.radians(a1)),
+        ]
+
+
 def get_coords(celldm, cell, geometry_inlines, inlines=None):
     coords_lines = get_stru_block(geometry_inlines, "ATOMIC_POSITIONS")
     # assuming that ATOMIC_POSITIONS is at the bottom of the STRU file
@@ -268,16 +313,15 @@ def get_coords(celldm, cell, geometry_inlines, inlines=None):
     coords = []  # coordinations of atoms
     move = []  # move flag of each atom
     velocity = []  # velocity of each atom
-    mag = []  # magnetic moment of each atom
-    angle1 = []  # angle1 of each atom
-    angle2 = []  # angle2 of each atom
+    mags = []  # magnetic moment of each atom
     sc = []  # spin constraint flag of each atom
     lambda_ = []  # lambda of each atom
 
     ntype = get_nele_from_stru(geometry_inlines)
     line_idx = 1  # starting line of first element
     for it in range(ntype):
         atom_names.append(coords_lines[line_idx].split()[0])
+        atom_type_mag = float(coords_lines[line_idx + 1].split()[0])
         line_idx += 2
         atom_numbs.append(int(coords_lines[line_idx].split()[0]))
         line_idx += 1
@@ -302,17 +346,20 @@ def get_coords(celldm, cell, geometry_inlines, inlines=None):
             if imove is not None:
                 move.append(imove)
             velocity.append(ivelocity)
-            mag.append(imagmom)
-            angle1.append(iangle1)
-            angle2.append(iangle2)
             sc.append(iconstrain)
             lambda_.append(ilambda1)
 
+            # calculate the magnetic moment in cartesian coordinates
+            mag = get_atom_mag_cartesian(imagmom, iangle1, iangle2)
+            if mag is None:
+                mag = [0, 0, atom_type_mag]
+            mags.append(mag)
+
             line_idx += 1
     coords = np.array(coords)  # need transformation!!!
     atom_types = np.array(atom_types)
     move = np.array(move, dtype=bool)
-    return atom_names, atom_numbs, atom_types, coords, move
+    return atom_names, atom_numbs, atom_types, coords, move, mags
 
 
 def get_energy(outlines):
@@ -479,9 +526,12 @@ def get_frame(fname):
         outlines = fp.read().split("\n")
 
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords, move = get_coords(
-        celldm, cell, geometry_inlines, inlines
+    atom_names, natoms, types, coords, move, magmom = (
+        get_coords(  # here the magmom is the initial magnetic moment in STRU
+            celldm, cell, geometry_inlines, inlines
+        )
     )
+
     magmom, magforce = get_mag_force(outlines)
     if len(magmom) > 0:
         magmom = magmom[-1:]
@@ -565,7 +615,7 @@ def get_frame_from_stru(fname):
     nele = get_nele_from_stru(geometry_inlines)
     inlines = [f"ntype {nele}"]
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords, move = get_coords(
+    atom_names, natoms, types, coords, move, magmom = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
     data = {}
@@ -575,6 +625,7 @@ def get_frame_from_stru(fname):
     data["cells"] = cell[np.newaxis, :, :]
     data["coords"] = coords[np.newaxis, :, :]
     data["orig"] = np.zeros(3)
+    data["spins"] = np.array([magmom])
     if len(move) > 0:
         data["move"] = move[np.newaxis, :, :]
 

dpdata/plugins/abacus.py

@@ -20,7 +20,9 @@
 @Format.register("stru")
 class AbacusSTRUFormat(Format):
     def from_system(self, file_name, **kwargs):
-        return dpdata.abacus.scf.get_frame_from_stru(file_name)
+        data = dpdata.abacus.scf.get_frame_from_stru(file_name)
+        register_mag_data(data)
+        return data
 
     def to_system(self, data, file_name: FileType, frame_idx=0, **kwargs):
         """Dump the system into ABACUS STRU format file.
@@ -55,6 +57,7 @@ def register_mag_data(data):
             required=False,
             deepmd_name="spin",
         )
+        dpdata.System.register_data_type(dt)
         dpdata.LabeledSystem.register_data_type(dt)
     if "force_mags" in data:
         dt = DataType(
@@ -64,6 +67,7 @@ def register_mag_data(data):
             required=False,
             deepmd_name="force_mag",
         )
+        dpdata.System.register_data_type(dt)
         dpdata.LabeledSystem.register_data_type(dt)
 
 

tests/abacus.scf/stru_test

@@ -22,11 +22,11 @@ Cartesian    # Cartesian(Unit is LATTICE_CONSTANT)
 C
 0.0
 1
-5.192682633809 4.557725978258 4.436846615358 1 1 1
+5.192682633809 4.557725978258 4.436846615358 1 1 1 mag 0.000000000000 0.000000000000 0.000000000000
 H
 0.0
 4
-5.416431453540 4.011298860305 3.511161492417 0 0 0
-4.131588222365 4.706745191323 4.431136645083 1 0 1
-5.630930319126 5.521640894956 4.450356541303 1 0 1
-5.499851012568 4.003388899277 5.342621842622 0 1 1
+5.416431453540 4.011298860305 3.511161492417 0 0 0 mag 0.000000000000 0.000000000000 0.000000000000
+4.131588222365 4.706745191323 4.431136645083 1 0 1 mag 0.000000000000 0.000000000000 0.000000000000
+5.630930319126 5.521640894956 4.450356541303 1 0 1 mag 0.000000000000 0.000000000000 0.000000000000
+5.499851012568 4.003388899277 5.342621842622 0 1 1 mag 0.000000000000 0.000000000000 0.000000000000

tests/abacus.spin/STRU.spin

@@ -0,0 +1,24 @@
+ATOMIC_SPECIES
+Fe 55.845 Fe.upf
+
+NUMERICAL_ORBITAL
+Fe_gga_10au_200.0Ry_4s2p2d1f.orb
+
+LATTICE_CONSTANT
+1.880277359
+LATTICE_VECTORS
+    2.8274254848     0.0000000000     0.0000000000 #latvec1
+    0.0000000000     2.8274254848     0.0000000000 #latvec2
+    0.0000000000     0.0000000000     2.8274254848 #latvec3
+
+ATOMIC_POSITIONS
+Direct
+
+Fe #label
+1 #magnetism
+4 #number of atoms
+    0.0000000000     0.000000000     0.000000000 mag 0 0 2
+    0.1000000000     0.1000000000     0.1000000000 mag 3
+    0.2000000000     0.2000000000     0.2000000000 mag 3 angle1 90
+    0.3000000000     0.3000000000     0.3000000000 mag 3 4 0 angle1 90 angle2 90
+

tests/test_abacus_spin.py

@@ -155,3 +155,19 @@ def test_md(self):
         np.testing.assert_almost_equal(
             data["force_mags"], sys2.data["force_mags"], decimal=8
         )
+
+    def test_read_stru_spin(self):
+        mysys = dpdata.System("abacus.spin/STRU.spin", fmt="abacus/stru")
+        self.assertTrue("spins" in mysys.data)
+        print(mysys.data["spins"])
+
+        """
+    0.0000000000     0.000000000     0.000000000 mag 0 0 2
+    0.1000000000     0.1000000000     0.1000000000 mag 3
+    0.2000000000     0.2000000000     0.2000000000 mag 3 angle1 90
+    0.3000000000     0.3000000000     0.3000000000 mag 3 4 0 angle1 90 angle2 90
+        """
+        np.testing.assert_almost_equal(mysys.data["spins"][0][0], [0, 0, 2], decimal=8)
+        np.testing.assert_almost_equal(mysys.data["spins"][0][1], [0, 0, 3], decimal=8)
+        np.testing.assert_almost_equal(mysys.data["spins"][0][2], [3, 0, 0], decimal=8)
+        np.testing.assert_almost_equal(mysys.data["spins"][0][3], [0, 5, 0], decimal=8)