Set initial magnetic moments

pyiron_atomistics/atomistics/structure/atoms.py

@@ -2566,17 +2566,63 @@ def set_initial_magnetic_moments(self, magmoms=None):
         Set array of initial magnetic moments.
 
         Args:
-            magmoms (numpy.ndarray/list): List of magneric moments
+            magmoms (numpy.ndarray/list/dict/float/None): List, dict or single
+                value of magneric moments
+
+
+        Example I:
+
+        >>> from pyiron_atomistics import Project
+        >>> structure = Project('.').create.structure.bulk('Ni', cubic=True)
+        >>> structure[-1] = 'Fe'
+        >>> v_dict = {'Fe': 1, 'Ni': 2}
+        >>> structure.set_initial_magnetic_moments(v_dict)
+        >>> structure.get_initial_magnetic_moments()
+        array([2, 2, 2, 1])
+
+
+        Example II:
+
+        >>> from pyiron_atomistics import Project
+        >>> structure = Project('.').create.structure.bulk('Ni', cubic=True)
+        >>> structure[-1] = 'Fe'
+        >>> structure.set_initial_magnetic_moments(1)
+        >>> print(structure.get_initial_magnetic_moments())
+
+        Output:
+
+        >>> array([1, 1, 1, 1])
+
+        If you want to make it non-magnetic, set `None`:
+
+        >>> structure.set_initial_magnetic_moments(None)
+
+        If a list or an array of length 3 (e.g. [1, 2, 3]) is set and the
+        number of atoms is not 3, it will be considered as a non-collinear
+        magnetic moments, and the values will be distributed to all the
+        atoms (like in Example II)
         """
         # pyiron part
         if magmoms is not None:
+            if isinstance(magmoms, dict):
+                if set(self.get_species_symbols()) != set(magmoms.keys()):
+                    raise ValueError(
+                        "Elements in structure {} not equal to those in dict {}".format(
+                            set(self.get_chemical_symbols()), set(magmoms.keys())
+                        )
+                    )
+                magmoms = [magmoms[c] for c in self.get_chemical_symbols()]
+            elif not hasattr(magmoms, "__len__") or (
+                len(magmoms) == 3 and len(self) != 3
+            ):
+                magmoms = len(self) * [magmoms]
             if len(magmoms) != len(self):
                 raise ValueError("magmons can be collinear or non-collinear.")
             if "spin" not in self._tag_list._lists.keys():
                 self.add_tag(spin=None)
             for ind, spin in enumerate(magmoms):
-                self.spin[ind] = spin
-        self.spins = magmoms
+                self.spin[ind] = spin  # For self._tag_list.spin
+        self.spins = magmoms  # For self.array['initial_magmoms']
 
     def rotate(
         self, a=0.0, v=None, center=(0, 0, 0), rotate_cell=False, index_list=None

tests/atomistics/structure/test_atoms.py

@@ -520,6 +520,13 @@ def test_set_initial_magnetic_moments(self):
         basis.set_repeat(2)
         self.assertTrue(np.array_equal(basis.spins, np.hstack([0.0, 2.0] * 8)))
         self.assertRaises(ValueError, basis.set_initial_magnetic_moments, magmoms=[4] * (len(basis) - 1))
+        basis = Atoms(elements=3 * ["Ni"] + ["Fe"], positions=np.random.random((4, 3)), cell=np.eye(3))
+        basis.set_initial_magnetic_moments({'Fe': 1, 'Ni': 2})
+        self.assertEqual(basis.get_initial_magnetic_moments().tolist(), [2, 2, 2, 1])
+        self.assertRaises(ValueError, basis.set_initial_magnetic_moments, {'Fe': 1})
+        basis = Atoms(elements=3 * ["Ni"] + ["Fe"], positions=np.random.random((4, 3)), cell=np.eye(3))
+        basis.set_initial_magnetic_moments(1)
+        self.assertEqual(basis.get_initial_magnetic_moments().tolist(), [1, 1, 1, 1])
 
     def test_get_parent_basis(self):
         periodic_table = PeriodicTable()