Repository: Quantum-Accelerators/quacc
Pull request #1816 (2024-03-01 23:41:38)

Stricter magmom handler and tests

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1816?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.41%. Comparing base [(`842e942`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/commit/842e942129d477387a3c9cdd63736e2b5b1e56e7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) to head [(`97b78dc`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1816?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).
> Report is 5 commits behind head on main.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1816   +/-   ##
=======================================
  Coverage   99.41%   99.41%           
=======================================
  Files          79       80    +1     
  Lines        3055     3101   +46     
=======================================
+ Hits         3037     3083   +46     
  Misses         18       18           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1816?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).


Patch:
diff --git a/src/quacc/recipes/espresso/_base.py b/src/quacc/recipes/espresso/_base.py
index b6588333..cef05711 100644
--- a/src/quacc/recipes/espresso/_base.py
+++ b/src/quacc/recipes/espresso/_base.py
@@ -81,7 +81,9 @@ def base_fn(
 
     final_atoms = run_calc(atoms, geom_file=geom_file, copy_files=copy_files)
 
-    return summarize_run(final_atoms, atoms, additional_fields=additional_fields)
+    return summarize_run(
+        final_atoms, atoms, move_magmoms=True, additional_fields=additional_fields
+    )
 
 
 def base_opt_fn(
@@ -152,7 +154,9 @@ def base_opt_fn(
 
     dyn = run_opt(atoms, relax_cell=relax_cell, copy_files=copy_files, **opt_flags)
 
-    return summarize_opt_run(dyn, additional_fields=additional_fields)
+    return summarize_opt_run(
+        dyn, move_magmoms=True, additional_fields=additional_fields
+    )
 
 
 def _prepare_atoms(
diff --git a/src/quacc/recipes/vasp/mp.py b/src/quacc/recipes/vasp/mp.py
index 1f9a6184..55575828 100644
--- a/src/quacc/recipes/vasp/mp.py
+++ b/src/quacc/recipes/vasp/mp.py
@@ -387,15 +387,17 @@ def mp_metagga_relax_flow(
     MPMetaGGARelaxFlowSchema
         Dictionary of results. See the type-hint for the data structure.
     """
-    (
-        mp_metagga_prerelax_job_,
-        mp_metagga_relax_job_,
-        mp_metagga_static_job_,
-    ) = customize_funcs(
-        ["mp_metagga_prerelax_job", "mp_metagga_relax_job", "mp_metagga_static_job"],
-        [mp_metagga_prerelax_job, mp_metagga_relax_job, mp_metagga_static_job],
-        parameters=job_params,
-        decorators=job_decorators,
+    (mp_metagga_prerelax_job_, mp_metagga_relax_job_, mp_metagga_static_job_) = (
+        customize_funcs(
+            [
+                "mp_metagga_prerelax_job",
+                "mp_metagga_relax_job",
+                "mp_metagga_static_job",
+            ],
+            [mp_metagga_prerelax_job, mp_metagga_relax_job, mp_metagga_static_job],
+            parameters=job_params,
+            decorators=job_decorators,
+        )
     )
 
     # Run the prerelax
diff --git a/src/quacc/schemas/ase.py b/src/quacc/schemas/ase.py
index 3cfb4d6c..c6ab70b8 100644
--- a/src/quacc/schemas/ase.py
+++ b/src/quacc/schemas/ase.py
@@ -12,7 +12,7 @@ from ase.vibrations.data import VibrationsData
 from quacc import SETTINGS, __version__
 from quacc.atoms.core import get_final_atoms_from_dyn
 from quacc.schemas.atoms import atoms_to_metadata
-from quacc.schemas.prep import prep_magmoms, prep_next_run
+from quacc.schemas.prep import prep_next_run
 from quacc.utils.dicts import clean_task_doc, recursive_dict_merge
 from quacc.utils.files import get_uri
 from quacc.wflow_tools.db import results_to_db
@@ -40,7 +40,7 @@ def summarize_run(
     final_atoms: Atoms,
     input_atoms: Atoms,
     charge_and_multiplicity: tuple[int, int] | None = None,
-    move_magmoms: bool = True,
+    move_magmoms: bool = False,
     additional_fields: dict[str, Any] | None = None,
     store: Store | bool | None = None,
 ) -> RunSchema:
@@ -103,9 +103,7 @@ def summarize_run(
 
     results = {"results": final_atoms.calc.results}
 
-    if move_magmoms:
-        final_atoms = prep_magmoms(final_atoms)
-    atoms_to_store = prep_next_run(final_atoms)
+    atoms_to_store = prep_next_run(final_atoms, move_magmoms=move_magmoms)
 
     if final_atoms:
         final_atoms_metadata = atoms_to_metadata(
@@ -130,7 +128,7 @@ def summarize_opt_run(
     trajectory: Trajectory | list[Atoms] = None,
     check_convergence: bool | None = None,
     charge_and_multiplicity: tuple[int, int] | None = None,
-    move_magmoms: bool = True,
+    move_magmoms: bool = False,
     additional_fields: dict[str, Any] | None = None,
     store: Store | bool | None = None,
 ) -> OptSchema:
diff --git a/src/quacc/schemas/cclib.py b/src/quacc/schemas/cclib.py
index c9a06bea..5a055e11 100644
--- a/src/quacc/schemas/cclib.py
+++ b/src/quacc/schemas/cclib.py
@@ -130,7 +130,6 @@ def cclib_summarize_run(
         final_atoms,
         input_atoms,
         charge_and_multiplicity=(attributes["charge"], attributes["mult"]),
-        move_magmoms=False,
         store=False,
     )
 
diff --git a/src/quacc/schemas/prep.py b/src/quacc/schemas/prep.py
index 6d36ab89..c1ce53e3 100644
--- a/src/quacc/schemas/prep.py
+++ b/src/quacc/schemas/prep.py
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def prep_next_run(atoms: Atoms) -> Atoms:
+def prep_next_run(atoms: Atoms, move_magmoms: bool = False) -> Atoms:
     """
     Prepares the Atoms object for a new run by stripping off the calculator and
     assigning a unique ID.
@@ -24,6 +24,9 @@ def prep_next_run(atoms: Atoms) -> Atoms:
     ----------
     atoms
         Atoms object
+    move_magmoms
+        Whether to move the magnetic moments from the calculator results to the
+        initial magnetic moments.
 
     Returns
     -------
@@ -32,6 +35,15 @@ def prep_next_run(atoms: Atoms) -> Atoms:
     """
     atoms = copy_atoms(atoms)
 
+    if (
+        move_magmoms
+        and hasattr(atoms, "calc")
+        and getattr(atoms.calc, "results", None) is not None
+    ):
+        atoms.set_initial_magnetic_moments(
+            atoms.calc.results.get("magmoms", [0.0] * len(atoms))
+        )
+
     # Clear off the calculator so we can run a new job. If we don't do this,
     # then something like atoms *= (2,2,2) still has a calculator attached,
     # which is a bit confusing.
@@ -50,42 +62,6 @@ def prep_next_run(atoms: Atoms) -> Atoms:
     return atoms
 
 
-def prep_magmoms(atoms: Atoms) -> Atoms:
-    """
-    Prepare the Atoms object for a new run by moving the converged magnetic
-    moments to the initial magnetic moments.
-
-    Parameters
-    ----------
-    atoms
-        Atoms object
-
-    Returns
-    -------
-    Atoms
-        Updated Atoms object.
-    """
-
-    atoms = copy_atoms(atoms)
-    if hasattr(atoms, "calc") and getattr(atoms.calc, "results", None) is not None:
-        # If there are initial magmoms set, then we should see what the final
-        # magmoms are. If they are present, move them to initial. If they are
-        # not present, it means the calculator doesn't support the "magmoms"
-        # property so we have to retain the initial magmoms given no further
-        # info.
-        if atoms.has("initial_magmoms"):
-            atoms.set_initial_magnetic_moments(
-                atoms.calc.results.get("magmoms", atoms.get_initial_magnetic_moments())
-            )
-        # If there are no initial magmoms set, just check the results and set
-        # everything to 0.0 if there is nothing there.
-        else:
-            atoms.set_initial_magnetic_moments(
-                atoms.calc.results.get("magmoms", [0.0] * len(atoms))
-            )
-    return atoms
-
-
 def set_magmoms(
     atoms: Atoms,
     elemental_mags_dict: dict[str, float] | None = None,


Test patch:
diff --git a/tests/core/calculators/vasp/test_vasp.py b/tests/core/calculators/vasp/test_vasp.py
index d231042b..9df8487f 100644
--- a/tests/core/calculators/vasp/test_vasp.py
+++ b/tests/core/calculators/vasp/test_vasp.py
@@ -14,7 +14,7 @@ from pymatgen.io.vasp.sets import MPRelaxSet, MPScanRelaxSet
 
 from quacc import SETTINGS
 from quacc.calculators.vasp import Vasp, presets
-from quacc.schemas.prep import prep_magmoms, prep_next_run
+from quacc.schemas.prep import prep_next_run
 
 FILE_DIR = Path(__file__).parent
 PSEUDO_DIR = FILE_DIR / "fake_pseudos"
@@ -315,28 +315,28 @@ def test_magmoms(atoms_mag, atoms_nomag, atoms_nospin):
 def test_prep_magmoms2(atoms_mag, atoms_nomag, atoms_nospin):
     atoms = deepcopy(atoms_mag)
     mags = atoms.get_magnetic_moments()
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
     assert atoms.has("initial_magmoms") is True
     assert atoms.get_initial_magnetic_moments().tolist() == mags.tolist()
 
     atoms = deepcopy(atoms_nomag)
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
     assert atoms.has("initial_magmoms") is True
     assert np.all(atoms.get_initial_magnetic_moments() == 0)
 
     atoms = deepcopy(atoms_nospin)
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
     assert atoms.has("initial_magmoms") is True
     assert np.all(atoms.get_initial_magnetic_moments() == 0)
 
     atoms = deepcopy(atoms_mag)
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     calc = Vasp(atoms, preset="BulkSet", mag_cutoff=10.0)
     atoms.calc = calc
     assert atoms.has("initial_magmoms") is True
@@ -346,7 +346,7 @@ def test_prep_magmoms2(atoms_mag, atoms_nomag, atoms_nospin):
     calc = Vasp(atoms)
     atoms.calc = calc
     atoms.calc.results = {"energy": -1.0, "magmoms": [0.0] * len(atoms)}
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     atoms *= (2, 2, 2)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
@@ -356,7 +356,7 @@ def test_prep_magmoms2(atoms_mag, atoms_nomag, atoms_nospin):
     calc = Vasp(atoms)
     atoms.calc = calc
     atoms.calc.results = {"energy": -1.0, "magmoms": [-0.02] * len(atoms)}
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     atoms *= (2, 2, 2)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
@@ -366,7 +366,7 @@ def test_prep_magmoms2(atoms_mag, atoms_nomag, atoms_nospin):
     calc = Vasp(atoms)
     atoms.calc = calc
     atoms.calc.results = {"energy": -1.0}
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     atoms *= (2, 2, 2)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
@@ -376,7 +376,7 @@ def test_prep_magmoms2(atoms_mag, atoms_nomag, atoms_nospin):
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
     atoms.calc.results = {"energy": -1.0, "magmoms": [3.14] * len(atoms)}
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     atoms *= (2, 2, 2)
     calc = Vasp(atoms, preset="BulkSet")
     atoms.calc = calc
diff --git a/tests/core/recipes/gaussian_recipes/conftest.py b/tests/core/recipes/gaussian_recipes/conftest.py
index 6e20070a..06413833 100644
--- a/tests/core/recipes/gaussian_recipes/conftest.py
+++ b/tests/core/recipes/gaussian_recipes/conftest.py
@@ -8,6 +8,7 @@ GAUSSIAN_DIR = Path(FILE_DIR, "gaussian_run")
 
 def mock_execute(self, *args, **kwargs):
     from shutil import copy
+
     copy(GAUSSIAN_DIR / "Gaussian.log.gz", Path(self.directory, "Gaussian.log.gz"))
 
 
@@ -21,6 +22,7 @@ def patch_execute(monkeypatch):
 def mock_read_results(self, *args, **kwargs):
     from ase.calculators.lj import LennardJones
     from ase.io import read
+
     atoms = read(Path(self.directory, "Gaussian.com"))
     atoms.calc = LennardJones()
     atoms.get_potential_energy()
diff --git a/tests/core/recipes/gaussian_recipes/test_gaussian_recipes.py b/tests/core/recipes/gaussian_recipes/test_gaussian_recipes.py
index cef6ed07..4f3df070 100644
--- a/tests/core/recipes/gaussian_recipes/test_gaussian_recipes.py
+++ b/tests/core/recipes/gaussian_recipes/test_gaussian_recipes.py
@@ -1,4 +1,5 @@
 from ase.build import molecule
+from numpy.testing import assert_array_equal
 
 from quacc.recipes.gaussian.core import relax_job, static_job
 
@@ -38,12 +39,14 @@ def test_static_job(tmp_path, monkeypatch):
         "2/9=2000",
     ]  # see ASE issue #660
     assert "opt" not in output["parameters"]
+    assert_array_equal(output["atoms"].get_initial_magnetic_moments(), [0, 0])
 
 
 def test_relax_job(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
 
     atoms = molecule("H2")
+    atoms.set_initial_magnetic_moments([0, 0])
 
     output = relax_job(atoms, 0, 1)
     assert output["natoms"] == len(atoms)
@@ -68,3 +71,23 @@ def test_relax_job(tmp_path, monkeypatch):
     assert output["parameters"]["basis"] == "def2svp"
     assert output["parameters"]["integral"] == "superfinegrid"
     assert output["parameters"]["ioplist"] == ["2/9=2000"]  # see ASE issue #660
+    assert_array_equal(output["atoms"].get_initial_magnetic_moments(), [0, 0])
+
+
+def test_relax_job_v2(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    atoms = molecule("H2")
+    atoms.set_initial_magnetic_moments([0, 3])
+
+    output = relax_job(atoms, 0, 3)
+    assert output["natoms"] == len(atoms)
+    assert output["parameters"]["charge"] == 0
+    assert output["parameters"]["mult"] == 3
+    assert output["parameters"]["opt"] == ""
+    assert output["parameters"]["xc"] == "wb97xd"
+    assert output["parameters"]["basis"] == "def2tzvp"
+    assert output["parameters"]["integral"] == "ultrafine"
+    assert "freq" not in output["parameters"]
+    assert "sp" not in output["parameters"]
+    assert_array_equal(output["atoms"].get_initial_magnetic_moments(), [0, 3])
diff --git a/tests/core/recipes/tblite_recipes/test_tblite_recipes.py b/tests/core/recipes/tblite_recipes/test_tblite_recipes.py
index 395c7c49..286e4e75 100644
--- a/tests/core/recipes/tblite_recipes/test_tblite_recipes.py
+++ b/tests/core/recipes/tblite_recipes/test_tblite_recipes.py
@@ -5,6 +5,7 @@ from copy import deepcopy
 
 import numpy as np
 from ase.build import bulk, molecule
+from numpy.testing import assert_array_equal
 
 from quacc import SETTINGS
 from quacc.recipes.tblite.core import freq_job, relax_job, static_job
@@ -20,6 +21,8 @@ def test_static_job_v1(tmp_path, monkeypatch):
     assert output["parameters"]["method"] == "GFN2-xTB"
     assert output["results"]["energy"] == pytest.approx(-137.96777594361672)
     assert np.array_equal(output["atoms"].get_positions(), atoms.get_positions())
+    assert_array_equal(output["atoms"].get_initial_magnetic_moments(), [0.0, 0.0, 0.0])
+    assert output["spin_multiplicity"] == 1
 
 
 def test_static_job_v2(tmp_path, monkeypatch):
@@ -30,6 +33,20 @@ def test_static_job_v2(tmp_path, monkeypatch):
     assert output["parameters"]["method"] == "GFN1-xTB"
     assert output["results"]["energy"] == pytest.approx(-156.96750578831137)
     assert np.array_equal(output["atoms"].get_positions(), atoms.get_positions())
+    assert output["spin_multiplicity"] == 1
+
+
+def test_static_job_v3(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    atoms = molecule("H")
+    atoms.set_initial_magnetic_moments([1.0])
+    output = static_job(atoms, method="GFN1-xTB")
+    assert output["parameters"]["method"] == "GFN1-xTB"
+    assert output["results"]["energy"] == pytest.approx(-10.92345239113973)
+    assert np.array_equal(output["atoms"].get_positions(), atoms.get_positions())
+    assert output["spin_multiplicity"] == 2
+    assert_array_equal(output["atoms"].get_initial_magnetic_moments(), [1.0])
 
 
 def test_relax_job(tmp_path, monkeypatch):
@@ -49,11 +66,13 @@ def test_relax_job_cell(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
 
     atoms = bulk("Cu")
+    atoms.set_initial_magnetic_moments([1.0])
     output = relax_job(atoms, method="GFN1-xTB", relax_cell=True)
     assert output["parameters"]["method"] == "GFN1-xTB"
     assert output["trajectory_results"][-1]["energy"] == pytest.approx(
         -130.4917433263215
     )
+    assert "spin_multiplicity" not in output
 
 
 def test_freq_job_v1(tmp_path, monkeypatch):
@@ -70,6 +89,7 @@ def test_freq_job_v1(tmp_path, monkeypatch):
     assert output["results"]["vib_freqs"][-1] == pytest.approx(3526.9940431752034)
     assert output["results"]["n_imag"] == 0
     assert output["results"]["imag_vib_freqs"] == []
+    assert output["spin_multiplicity"] == 1
 
     assert output["symmetry"]["point_group"] == "C2v"
     assert output["symmetry"]["rotation_number"] == 2
@@ -100,6 +120,7 @@ def test_freq_job_v2(tmp_path, monkeypatch):
     assert output["results"]["vib_freqs"] == []
     assert output["results"]["n_imag"] == 0
     assert output["results"]["imag_vib_freqs"] == []
+    assert output["spin_multiplicity"] == 2
 
     assert output["symmetry"]["linear"] is False
     assert output["symmetry"]["rotation_number"] == np.inf
@@ -114,6 +135,7 @@ def test_freq_job_v3(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
 
     atoms = molecule("CH3")
+    atoms.set_initial_magnetic_moments([1.0, 0.0, 0.0, 0.0])
     initial_atoms = deepcopy(atoms)
     output = freq_job(atoms, energy=-10.0, temperature=1000, pressure=20)
     assert output["atoms"] == initial_atoms
@@ -129,6 +151,7 @@ def test_freq_job_v3(tmp_path, monkeypatch):
     assert output["results"]["vib_energies"][-1] == pytest.approx(0.3880868821616259)
     assert output["results"]["n_imag"] == 0
     assert output["results"]["imag_vib_freqs"] == []
+    assert output["spin_multiplicity"] == 2
 
     assert output["parameters_thermo"]["temperature"] == 1000.0
     assert output["parameters_thermo"]["pressure"] == 20.0
@@ -143,6 +166,9 @@ def test_freq_job_v3(tmp_path, monkeypatch):
     assert output["results"]["gibbs_energy"] == pytest.approx(-11.100020872176652)
     assert "nid" in output
     assert "dir_name" in output
+    assert_array_equal(
+        output["atoms"].get_initial_magnetic_moments(), [1.0, 0.0, 0.0, 0.0]
+    )
 
 
 def test_unique_workdir(tmp_path, monkeypatch):
diff --git a/tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py b/tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py
index e3c937da..007777ab 100644
--- a/tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py
+++ b/tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py
@@ -525,6 +525,12 @@ def test_mp_metagga_relax_flow(tmp_path, monkeypatch):
 
     output = mp_metagga_relax_flow(atoms)
     assert output["static"]["nsites"] == len(atoms)
+    assert output["prerelax"]["parameters"]["gga"] == "ps"
+    assert output["prerelax"]["parameters"]["ismear"] == 0
+    assert output["prerelax"]["parameters"]["pp"] == "pbe"
+    assert output["prerelax"]["parameters"]["magmom"] == [0.6]
+    assert output["relax1"]["parameters"]["magmom"] == [0.0]
+    assert output["relax2"]["parameters"]["magmom"] == [0.0]
     assert output["relax2"]["parameters"]["metagga"].lower() == "r2scan"
     assert output["relax2"]["parameters"]["ediffg"] == -0.02
     assert output["relax2"]["parameters"]["encut"] == 680
@@ -532,13 +538,15 @@ def test_mp_metagga_relax_flow(tmp_path, monkeypatch):
     assert output["relax2"]["parameters"]["sigma"] == 0.05
     assert output["relax2"]["parameters"]["kspacing"] == 0.22
     assert output["relax2"]["parameters"]["pp"] == "pbe"
-    assert output["prerelax"]["parameters"]["gga"] == "ps"
-    assert output["prerelax"]["parameters"]["ismear"] == 0
-    assert output["prerelax"]["parameters"]["pp"] == "pbe"
 
     atoms = bulk("C")
+    atoms.set_initial_magnetic_moments([0.0, 0.0])
     output = mp_metagga_relax_flow(atoms)
-    assert output["static"]["nsites"] == len(atoms)
+    assert output["prerelax"]["parameters"]["ismear"] == 0
+    assert output["prerelax"]["parameters"]["pp"] == "pbe"
+    assert output["prerelax"]["parameters"]["magmom"] == [0.0, 0.0]
+    assert output["relax1"]["parameters"]["magmom"] == [0.0, 0.0]
+    assert output["relax2"]["parameters"]["magmom"] == [0.0, 0.0]
     assert output["relax2"]["parameters"]["metagga"].lower() == "r2scan"
     assert output["relax2"]["parameters"]["ediffg"] == -0.02
     assert output["relax2"]["parameters"]["encut"] == 680
@@ -547,17 +555,17 @@ def test_mp_metagga_relax_flow(tmp_path, monkeypatch):
         0.28329488761304206
     )
     assert output["relax2"]["parameters"]["pp"] == "pbe"
-    assert output["prerelax"]["parameters"]["ismear"] == 0
-    assert output["prerelax"]["parameters"]["pp"] == "pbe"
+    assert output["static"]["nsites"] == len(atoms)
 
     atoms = molecule("O2")
+    atoms.set_initial_magnetic_moments([1.0, 0.0])
     atoms.center(vacuum=10)
     atoms.pbc = True
     output = mp_metagga_relax_flow(atoms)
-    assert output["static"]["nsites"] == len(atoms)
-    assert output["static"]["parameters"]["ismear"] == -5
-    assert output["static"]["parameters"]["nsw"] == 0
-    assert output["static"]["parameters"]["algo"] == "fast"
+    assert output["prerelax"]["parameters"]["ismear"] == 0
+    assert output["prerelax"]["parameters"]["pp"] == "pbe"
+    assert output["prerelax"]["parameters"]["magmom"] == [1.0, 0.0]
+    assert output["relax1"]["parameters"]["magmom"] == [0.0, 0.0]
     assert output["relax2"]["parameters"]["metagga"].lower() == "r2scan"
     assert output["relax2"]["parameters"]["ediffg"] == -0.02
     assert output["relax2"]["parameters"]["encut"] == 680
@@ -566,13 +574,18 @@ def test_mp_metagga_relax_flow(tmp_path, monkeypatch):
         0.28329488761304206
     )
     assert output["relax2"]["parameters"]["pp"] == "pbe"
-    assert output["prerelax"]["parameters"]["ismear"] == 0
-    assert output["prerelax"]["parameters"]["pp"] == "pbe"
+    assert output["relax2"]["parameters"]["magmom"] == [0.0, 0.0]
+    assert output["static"]["nsites"] == len(atoms)
+    assert output["static"]["parameters"]["ismear"] == -5
+    assert output["static"]["parameters"]["nsw"] == 0
+    assert output["static"]["parameters"]["algo"] == "fast"
+    assert output["static"]["parameters"]["magmom"] == [0.0, 0.0]
 
 
 def test_mp_gga_relax_job():
     atoms = bulk("Ni") * (2, 1, 1)
     atoms[0].symbol = "O"
+    del atoms.arrays["initial_magmoms"]
     output = mp_gga_relax_job(atoms)
 
     assert output["nsites"] == len(atoms)
@@ -598,7 +611,7 @@ def test_mp_gga_relax_job():
         "lorbit": 11,
         "lreal": "auto",
         "lwave": False,
-        "magmom": [0.6, 0.6],
+        "magmom": [0.6, 5.0],
         "nelm": 100,
         "nsw": 99,
         "prec": "accurate",
@@ -612,6 +625,7 @@ def test_mp_gga_relax_job():
 def test_mp_gga_static_job():
     atoms = bulk("Ni") * (2, 1, 1)
     atoms[0].symbol = "O"
+    del atoms.arrays["initial_magmoms"]
     output = mp_gga_static_job(atoms)
     assert output["nsites"] == len(atoms)
     assert output["parameters"] == {
@@ -635,7 +649,7 @@ def test_mp_gga_static_job():
         "lorbit": 11,
         "lreal": False,
         "lwave": True,  # modified by us (sensible)
-        "magmom": [0.6, 0.6],
+        "magmom": [0.6, 5.0],
         "nelm": 100,
         "nsw": 0,
         "prec": "accurate",
@@ -648,8 +662,8 @@ def test_mp_gga_static_job():
 def test_mp_gga_relax_flow():
     atoms = bulk("Ni") * (2, 1, 1)
     atoms[0].symbol = "O"
+    del atoms.arrays["initial_magmoms"]
     output = mp_gga_relax_flow(atoms)
-    assert output["static"]["nsites"] == len(atoms)
     relax_params = {
         "algo": "fast",
         "ediff": 0.0001,
@@ -672,7 +686,7 @@ def test_mp_gga_relax_flow():
         "lorbit": 11,
         "lreal": "auto",
         "lwave": False,
-        "magmom": [0.6, 0.6],
+        "magmom": [0.6, 5.0],
         "nelm": 100,
         "nsw": 99,
         "prec": "accurate",
@@ -680,8 +694,11 @@ def test_mp_gga_relax_flow():
         "pp": "pbe",
         "setups": {"O": "", "Ni": "_pv"},
     }
+    relax2_params = relax_params.copy()
+    relax2_params["magmom"] = [0.0, 0.0]
+
     assert output["relax1"]["parameters"] == relax_params
-    assert output["relax2"]["parameters"] == relax_params
+    assert output["relax2"]["parameters"] == relax2_params
     assert output["static"]["parameters"] == {
         "algo": "fast",
         "ediff": 0.0001,
@@ -703,7 +720,7 @@ def test_mp_gga_relax_flow():
         "lorbit": 11,
         "lreal": False,
         "lwave": True,  # modified by us (sensible)
-        "magmom": [0.6, 0.6],
+        "magmom": [0.0, 0.0],
         "nelm": 100,
         "nsw": 0,
         "prec": "accurate",
@@ -716,6 +733,7 @@ def test_mp_gga_relax_flow():
 def test_mp_relax_flow_custom():
     atoms = bulk("Ni") * (2, 1, 1)
     atoms[0].symbol = "O"
+    del atoms.arrays["initial_magmoms"]
     output = mp_metagga_relax_flow(
         mp_gga_relax_flow(atoms, job_params={"mp_gga_relax_job": {"nsw": 0}})["static"][
             "atoms"
diff --git a/tests/core/schemas/test_ase.py b/tests/core/schemas/test_ase.py
index 88d49ca2..0b4033ca 100644
--- a/tests/core/schemas/test_ase.py
+++ b/tests/core/schemas/test_ase.py
@@ -65,7 +65,7 @@ def test_summarize_run4():
     atoms = read(os.path.join(RUN1, "OUTCAR.gz"))
     atoms.set_initial_magnetic_moments([3.14] * len(atoms))
     atoms.calc.results["magmoms"] = [2.0] * len(atoms)
-    results = summarize_run(atoms, initial_atoms)
+    results = summarize_run(atoms, initial_atoms, move_magmoms=True)
 
     assert atoms.calc is not None
     assert atoms.get_initial_magnetic_moments().tolist() == [3.14] * len(atoms)
diff --git a/tests/core/schemas/test_prep.py b/tests/core/schemas/test_prep.py
index ebca62a3..80f3b0ad 100644
--- a/tests/core/schemas/test_prep.py
+++ b/tests/core/schemas/test_prep.py
@@ -9,7 +9,7 @@ from ase.calculators.emt import EMT
 from ase.io import read
 
 from quacc.atoms.core import get_atoms_id
-from quacc.schemas.prep import prep_magmoms, prep_next_run
+from quacc.schemas.prep import prep_next_run
 
 
 @pytest.fixture()
@@ -55,16 +55,16 @@ def test_prep_next_run():
     atoms = bulk("Cu")
     atoms.calc = EMT()
     md5hash = "d4859270a1a67083343bec0ab783f774"
-    atoms = prep_next_run(atoms)
+    atoms = prep_next_run(atoms, move_magmoms=False)
     assert atoms.info.get("_id", None) == md5hash
     assert atoms.info.get("_old_ids", None) is None
     assert atoms.calc is None
-    atoms = prep_next_run(atoms)
+    atoms = prep_next_run(atoms, move_magmoms=False)
     assert atoms.info.get("_id", None) == md5hash
     assert atoms.info.get("_old_ids", None) == [md5hash]
     atoms[0].symbol = "Pt"
     new_md5hash = "52087d50a909572d58e01cfb49d4911b"
-    atoms = prep_next_run(atoms)
+    atoms = prep_next_run(atoms, move_magmoms=True)
     assert atoms.info.get("_old_ids", None) == [md5hash, md5hash]
     assert atoms.info.get("_id", None) == new_md5hash
 
@@ -72,16 +72,21 @@ def test_prep_next_run():
 def test_prep_magmoms(atoms_mag, atoms_nomag, atoms_nospin):
     atoms = deepcopy(atoms_mag)
     mags = atoms.get_magnetic_moments()
-    atoms = prep_magmoms(atoms)
+    atoms = prep_next_run(atoms, move_magmoms=True)
     assert atoms.has("initial_magmoms") is True
     assert atoms.get_initial_magnetic_moments().tolist() == mags.tolist()
 
     atoms = deepcopy(atoms_nomag)
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     assert atoms.has("initial_magmoms") is True
     assert np.all(atoms.get_initial_magnetic_moments() == 0)
 
     atoms = deepcopy(atoms_nospin)
-    atoms = prep_next_run(prep_magmoms(atoms))
+    atoms = prep_next_run(atoms, move_magmoms=True)
     assert atoms.has("initial_magmoms") is True
     assert np.all(atoms.get_initial_magnetic_moments() == 0)
+
+    atoms = deepcopy(atoms_mag)
+    mags = atoms.get_magnetic_moments()
+    atoms = prep_next_run(atoms, move_magmoms=False)
+    assert atoms.has("initial_magmoms") is False