From 7c8139dd1a0a1403aa5140d6937d0cebb5337f21 Mon Sep 17 00:00:00 2001
From: BowenD-UCB <84425382+BowenD-UCB@users.noreply.github.com>
Date: Thu, 26 Oct 2023 01:30:59 -0700
Subject: [PATCH] Added function to set MD starting temperature

---
 chgnet/model/dynamics.py   | 12 ++++++++++++
 chgnet/utils/vasp_utils.py | 14 +++++++++-----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/chgnet/model/dynamics.py b/chgnet/model/dynamics.py
index fc0f4d9a..4d12f765 100644
--- a/chgnet/model/dynamics.py
+++ b/chgnet/model/dynamics.py
@@ -14,6 +14,7 @@
 from ase.md.npt import NPT
 from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen, NPTBerendsen
 from ase.md.nvtberendsen import NVTBerendsen
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
 from ase.md.verlet import VelocityVerlet
 from ase.optimize.bfgs import BFGS
 from ase.optimize.bfgslinesearch import BFGSLineSearch
@@ -355,6 +356,7 @@ def __init__(
         ensemble: str = "nvt",
         thermostat: str = "Berendsen_inhomogeneous",
         temperature: int = 300,
+        starting_temperature: int | None = None,
         timestep: float = 2.0,
         pressure: float = 1.01325e-4,
         taut: float | None = None,
@@ -382,6 +384,10 @@ def __init__(
                 Default = "Nose-Hoover"
             temperature (float): temperature for MD simulation, in K
                 Default = 300
+            starting_temperature (float): starting temperature of MD simulation, in K
+                if set as None, the MD starts with the momentum carried by ase.Atoms
+                if input is a pymatgen.core.Structure, the MD starts at 0K
+                Default = None
             timestep (float): time step in fs
                 Default = 2
             pressure (float): pressure in GPa
@@ -434,6 +440,12 @@ def __init__(
         if isinstance(atoms, (Structure, Molecule)):
             atoms = atoms.to_ase_atoms()
 
+        if starting_temperature is not None:
+            MaxwellBoltzmannDistribution(
+                atoms, temperature_K=starting_temperature, force_temp=True
+            )
+            Stationary(atoms)
+
         self.atoms = atoms
         if isinstance(model, CHGNetCalculator):
             self.atoms.calc = model
diff --git a/chgnet/utils/vasp_utils.py b/chgnet/utils/vasp_utils.py
index 3e521da9..6a091aef 100644
--- a/chgnet/utils/vasp_utils.py
+++ b/chgnet/utils/vasp_utils.py
@@ -175,6 +175,8 @@ def solve_charge_by_mag(
                     (4.2, 5): 2
                 ))
     """
+    out_structure = structure.copy()
+    out_structure.remove_oxidation_states()
     ox_list = []
     solved_ox = True
     default_ox = default_ox or {"Li": 1, "O": -2}
@@ -183,10 +185,10 @@ def solve_charge_by_mag(
     }
 
     mag = structure.site_properties.get(
-        "final_magmom", structure.site_properties.get("magmom")
+        "magmom", structure.site_properties.get("magmom")
     )
 
-    for idx, site in enumerate(structure):
+    for idx, site in enumerate(out_structure):
         assigned = False
         if site.species_string in ox_ranges:
             for (min_mag, max_mag), mag_ox in ox_ranges[site.species_string].items():
@@ -201,7 +203,9 @@ def solve_charge_by_mag(
             solved_ox = False
 
     if solved_ox:
-        print(ox_list)
-        structure.add_oxidation_state_by_site(ox_list)
-        return structure
+        total_charge = sum(ox_list)
+        print(f"Solvec oxidation state, total charge={total_charge}")
+        out_structure.add_oxidation_state_by_site(ox_list)
+        return out_structure
+    print("Failed to solve oxidation state")
     return None