From 876105d54f74237254d6678c4fe939f013ce5402 Mon Sep 17 00:00:00 2001
From: Yifan Li <yifanl0716@gmail.com>
Date: Sun, 17 Sep 2023 01:02:07 -0400
Subject: [PATCH] add deeptensor unittest for si units

---
 source/lmp/tests/test_deeptensor.py | 58 ++++++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 6 deletions(-)

diff --git a/source/lmp/tests/test_deeptensor.py b/source/lmp/tests/test_deeptensor.py
index 8c7d7d9640..cebe4e6cd7 100644
--- a/source/lmp/tests/test_deeptensor.py
+++ b/source/lmp/tests/test_deeptensor.py
@@ -6,6 +6,7 @@
     Path,
 )
 
+import constants
 import numpy as np
 import pytest
 from lammps import (
@@ -23,6 +24,7 @@
 pb_file2 = Path(__file__).parent / "deepdipole_new.pb"
 system_file = Path(__file__).parent.parent.parent / "tests"
 data_file = Path(__file__).parent / "data.lmp"
+data_file_si = Path(__file__).parent / "data.si"
 data_type_map_file = Path(__file__).parent / "data_type_map.lmp"
 
 # this is as the same as python and c++ tests, test_deepdipole.py
@@ -75,24 +77,49 @@ def setup_module():
     write_lmp_data(box, coord, type_OH, data_file)
     # TODO
     # write_lmp_data(box, coord, type_HO, data_type_map_file)
+    write_lmp_data(
+        box * constants.dist_metal2si,
+        coord * constants.dist_metal2si,
+        type_OH,
+        data_file_si,
+    )
 
 
 def teardown_module():
     os.remove(data_file)
     # os.remove(data_type_map_file)
+    os.remove(data_file_si)
 
 
-def _lammps(data_file) -> PyLammps:
+def _lammps(data_file, units="metal") -> PyLammps:
     lammps = PyLammps()
-    lammps.units("metal")
+    lammps.units(units)
     lammps.boundary("p p p")
     lammps.atom_style("atomic")
-    lammps.neighbor("2.0 bin")
+    if units == "metal" or units == "real":
+        lammps.neighbor("2.0 bin")
+    elif units == "si":
+        lammps.neighbor("2.0e-10 bin")
+    else:
+        raise ValueError("units should be metal, real, or si")
     lammps.neigh_modify("every 10 delay 0 check no")
     lammps.read_data(data_file.resolve())
-    lammps.mass("1 16")
-    lammps.mass("2 2")
-    lammps.timestep(0.0005)
+    if units == "metal" or units == "real":
+        lammps.mass("1 16")
+        lammps.mass("2 2")
+    elif units == "si":
+        lammps.mass("1 %.10e" % (16 * constants.mass_metal2si))
+        lammps.mass("2 %.10e" % (2 * constants.mass_metal2si))
+    else:
+        raise ValueError("units should be metal, real, or si")
+    if units == "metal":
+        lammps.timestep(0.0005)
+    elif units == "real":
+        lammps.timestep(0.5)
+    elif units == "si":
+        lammps.timestep(5e-16)
+    else:
+        raise ValueError("units should be metal, real, or si")
     lammps.fix("1 all nve")
     return lammps
 
@@ -109,6 +136,13 @@ def lammps():
 #    yield _lammps(data_file=data_type_map_file)
 
 
+@pytest.fixture
+def lammps_si():
+    lmp = _lammps(data_file=data_file_si, units="si")
+    yield lmp
+    lmp.close()
+
+
 def test_compute_deeptensor_atom(lammps):
     lammps.pair_style(f"deepmd {pb_file.resolve()}")
     lammps.pair_coeff("* *")
@@ -120,3 +154,15 @@ def test_compute_deeptensor_atom(lammps):
     assert np.array(lammps.variables["tensor"].value) == pytest.approx(
         expected_d[idx_map]
     )
+
+def test_compute_deeptensor_atom_si(lammps_si):
+    lammps_si.pair_style(f"deepmd {pb_file.resolve()}")
+    lammps_si.pair_coeff("* *")
+    lammps_si.compute(f"tensor all deeptensor/atom {pb_file2.resolve()}")
+    lammps_si.variable("tensor atom c_tensor[1]")
+    lammps_si.dump("1 all custom 1 dump id c_tensor[1]")
+    lammps_si.run(0)
+    idx_map = lammps_si.lmp.numpy.extract_atom("id") - 1
+    assert np.array(lammps_si.variables["tensor"].value) == pytest.approx(
+        expected_d[idx_map] * constants.dist_metal2si
+    )
\ No newline at end of file