add ut for lammps atomic energy

Author: iProzd

source/lmp/tests/test_lammps_spin.py

@@ -150,6 +150,46 @@ def test_pair_deepmd(lammps):
     lammps.run(1)
 
 
+def test_pair_deepmd_virial(lammps):
+    lammps.pair_style(f"deepspin {pb_file.resolve()}")
+    lammps.pair_coeff("* *")
+    lammps.compute("peatom all pe/atom pair")
+    lammps.compute("pressure all pressure NULL pair")
+    lammps.compute("virial all centroid/stress/atom NULL pair")
+    lammps.variable("eatom atom c_peatom")
+    # for ii in range(9):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     lammps.variable(f"pressure{jj} equal c_pressure[{ii+1}]")
+    # for ii in range(9):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     lammps.variable(f"virial{jj} atom c_virial[{ii+1}]")
+    # lammps.dump(
+    #     "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)])
+    # )
+    lammps.run(0)
+    assert lammps.eval("pe") == pytest.approx(expected_e)
+    for ii in range(4):
+        assert lammps.atoms[ii].force == pytest.approx(
+            expected_f[lammps.atoms[ii].id - 1]
+        )
+    idx_map = lammps.lmp.numpy.extract_atom("id") - 1
+    assert np.array(lammps.variables["eatom"].value) == pytest.approx(
+        expected_ae[idx_map]
+    )
+    # vol = box[1] * box[3] * box[5]
+    # for ii in range(6):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     assert np.array(
+    #         lammps.variables[f"pressure{jj}"].value
+    #     ) / constants.nktv2p == pytest.approx(
+    #         -expected_v[idx_map, jj].sum(axis=0) / vol
+    #     )
+    # for ii in range(9):
+    #     assert np.array(
+    #         lammps.variables[f"virial{ii}"].value
+    #     ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii])
+
+
 def test_pair_deepmd_model_devi(lammps):
     lammps.pair_style(
         f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1"

source/lmp/tests/test_lammps_spin_pt.py

@@ -146,6 +146,46 @@ def test_pair_deepmd(lammps):
     lammps.run(1)
 
 
+def test_pair_deepmd_virial(lammps):
+    lammps.pair_style(f"deepspin {pb_file.resolve()}")
+    lammps.pair_coeff("* *")
+    lammps.compute("peatom all pe/atom pair")
+    lammps.compute("pressure all pressure NULL pair")
+    lammps.compute("virial all centroid/stress/atom NULL pair")
+    lammps.variable("eatom atom c_peatom")
+    # for ii in range(9):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     lammps.variable(f"pressure{jj} equal c_pressure[{ii+1}]")
+    # for ii in range(9):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     lammps.variable(f"virial{jj} atom c_virial[{ii+1}]")
+    # lammps.dump(
+    #     "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)])
+    # )
+    lammps.run(0)
+    assert lammps.eval("pe") == pytest.approx(expected_e)
+    for ii in range(4):
+        assert lammps.atoms[ii].force == pytest.approx(
+            expected_f[lammps.atoms[ii].id - 1]
+        )
+    idx_map = lammps.lmp.numpy.extract_atom("id") - 1
+    assert np.array(lammps.variables["eatom"].value) == pytest.approx(
+        expected_ae[idx_map]
+    )
+    # vol = box[1] * box[3] * box[5]
+    # for ii in range(6):
+    #     jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii]
+    #     assert np.array(
+    #         lammps.variables[f"pressure{jj}"].value
+    #     ) / constants.nktv2p == pytest.approx(
+    #         -expected_v[idx_map, jj].sum(axis=0) / vol
+    #     )
+    # for ii in range(9):
+    #     assert np.array(
+    #         lammps.variables[f"virial{ii}"].value
+    #     ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii])
+
+
 def test_pair_deepmd_model_devi(lammps):
     lammps.pair_style(
         f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1"