unit test

Author: chiang-yuan

mlip_arena/tasks/md.py

@@ -337,16 +337,42 @@ def run(
             traj = Trajectory(traj_file, "w", atoms)
 
             if not np.isnan(t_schedule).any():
-                MaxwellBoltzmannDistribution(
-                    atoms=atoms,
-                    temperature_K=t_schedule[last_step],
-                    rng=np.random.default_rng(seed=velocity_seed),
-                )
+                if atoms.get_temperature() == 0.0:
+                    MaxwellBoltzmannDistribution(
+                        atoms=atoms,
+                        temperature_K=t_schedule[0],
+                        rng=np.random.default_rng(seed=velocity_seed),
+                    )
+                else:
+                    atoms.set_momenta(
+                        atoms.get_momenta() * np.sqrt(
+                            t_schedule[0] / atoms.get_temperature()
+                        )
+                    )
 
             if zero_linear_momentum:
                 Stationary(atoms)
             if zero_angular_momentum:
                 ZeroRotation(atoms)
+    else:
+        if not np.isnan(t_schedule).any():
+            if atoms.get_temperature() == 0.0:
+                MaxwellBoltzmannDistribution(
+                    atoms=atoms,
+                    temperature_K=t_schedule[0],
+                    rng=np.random.default_rng(seed=velocity_seed),
+                )
+            else:
+                atoms.set_momenta(
+                    atoms.get_momenta() * np.sqrt(
+                        t_schedule[0] / atoms.get_temperature()
+                    )
+                )
+
+        if zero_linear_momentum:
+            Stationary(atoms)
+        if zero_angular_momentum:
+            ZeroRotation(atoms)
 
     fraction_traceless = dynamics_kwargs.pop("fraction_traceless", 1.0)
 

mlip_arena/tasks/viscosity.py

@@ -254,6 +254,8 @@ def run(
         )
         atoms = result["atoms"]
 
+    # rescale temperature and volume
+
     # Second stage: NVE equilibration
 
     if nve_eq_time > 0:
@@ -355,9 +357,9 @@ def store_stress(dyn: MolecularDynamics):
         "viscosity": {
             "units": "mPa·s",
             "components": {
-                "pxy": viscosity_components[0] * 1e3,
-                "pxz": viscosity_components[1] * 1e3,
-                "pyz": viscosity_components[2] * 1e3,
+                "xy": viscosity_components[0] * 1e3,
+                "xz": viscosity_components[1] * 1e3,
+                "yz": viscosity_components[2] * 1e3,
             },
             "average": eta_mPa_s,
             "final": eta_mPa_s.mean(),

tests/test_viscosity.py

@@ -0,0 +1,29 @@
+import numpy as np
+from ase.build import bulk
+from ase.calculators.lj import LennardJones
+
+from mlip_arena.tasks.viscosity import run as VISCOSITY
+
+
+def test_task_run():
+    """Test the viscosity task."""
+    atoms = bulk("Cu", "fcc", a=3.6, cubic=True) * (2, 2, 2)
+
+    atoms.positions = atoms.positions + 0.5 * np.random.rand(*atoms.positions.shape)
+
+    calculator = LennardJones(
+        epsilon=0.01, sigma=(atoms.get_volume() / len(atoms)) ** (1 / 3)
+    )
+
+    result = VISCOSITY(
+        atoms,
+        calculator=calculator,
+        temperature=1250,
+        pressure=0,
+        npt_eq_time=100,
+        nve_eq_time=100,
+        prod_time=100,
+    )
+    assert isinstance(result, dict)
+    assert "viscosity" in result
+    assert isinstance(result["viscosity"]["average"], float)