From 5912d0a1c2d86333478aea3399488a3cb11051ab Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Fri, 16 Jul 2021 17:43:44 -0400
Subject: [PATCH] add support for setting global energy for fix external to
 python module

---
 python/lammps/core.py                        | 17 +++++++++++++++++
 unittest/c-library/test_library_external.cpp | 13 ++++++++-----
 unittest/python/python-fix-external.py       |  5 +++++
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/python/lammps/core.py b/python/lammps/core.py
index 5079828ba8..44cd51bb33 100644
--- a/python/lammps/core.py
+++ b/python/lammps/core.py
@@ -1791,6 +1791,23 @@ class lammps(object):
 
   # -------------------------------------------------------------------------
 
+  def fix_external_set_energy_global(self, fix_id, eng):
+    """Get access to that array with per-atom forces of a fix external instance with a given fix ID.
+
+    This is a wrapper around the :cpp:func:`lammps_fix_external_get_force` function
+    of the C-library interface.
+
+    :param fix_id:  Fix-ID of a fix external instance
+    :type: string
+    :param eng:     potential energy to be added by fix external
+    :type: float
+    """
+
+    with ExceptionCheck(self):
+      return self.lib.lammps_fix_external_set_energy_global(self.lmp, fix_id.encode(), eng)
+
+  # -------------------------------------------------------------------------
+
   def get_neighlist(self, idx):
     """Returns an instance of :class:`NeighList` which wraps access to the neighbor list with the given index
 
diff --git a/unittest/c-library/test_library_external.cpp b/unittest/c-library/test_library_external.cpp
index f6f126c2b8..7c53daaef7 100644
--- a/unittest/c-library/test_library_external.cpp
+++ b/unittest/c-library/test_library_external.cpp
@@ -28,8 +28,11 @@ static void callback_one(void *handle, step_t timestep, int nlocal, tag_t *, dou
 {
     for (int i = 0; i < nlocal; ++i)
         f[i][0] = f[i][1] = f[i][2] = (double)timestep;
-    lammps_fix_external_set_energy_global(handle, "ext", 1.0);
-    double v[6] = {1.0,1.0,1.0,0.0,0.0,0.0 };
+    if (timestep < 10)
+        lammps_fix_external_set_energy_global(handle, "ext", 0.0);
+    else
+        lammps_fix_external_set_energy_global(handle, "ext", 1.0);
+    double v[6] = {1.0, 1.0, 1.0, 0.0, 0.0, 0.0};
     lammps_fix_external_set_virial_global(handle, "ext", v);
 }
 }
@@ -65,10 +68,10 @@ TEST(lammps_external, callback)
     ::testing::internal::CaptureStdout();
     lammps_set_fix_external_callback(handle, "ext", &callback_one, handle);
     lammps_command(handle, "run 10 post no");
-    double temp = lammps_get_thermo(handle, "temp");
-    double pe = lammps_get_thermo(handle, "pe");
+    double temp  = lammps_get_thermo(handle, "temp");
+    double pe    = lammps_get_thermo(handle, "pe");
     double press = lammps_get_thermo(handle, "press");
-    output      = ::testing::internal::GetCapturedStdout();
+    output       = ::testing::internal::GetCapturedStdout();
     if (verbose) std::cout << output;
     EXPECT_DOUBLE_EQ(temp, 1.0 / 30.0);
     EXPECT_DOUBLE_EQ(pe, 1.0 / 8.0);
diff --git a/unittest/python/python-fix-external.py b/unittest/python/python-fix-external.py
index 4f589bb5f6..02fe805626 100644
--- a/unittest/python/python-fix-external.py
+++ b/unittest/python/python-fix-external.py
@@ -8,6 +8,10 @@ def callback_one(lmp, ntimestep, nlocal, tag, x, f):
         f[i][0] = float(ntimestep)
         f[i][1] = float(ntimestep)
         f[i][2] = float(ntimestep)
+    if ntimestep < 10:
+        lmp.fix_external_set_energy_global("ext",0.5)
+    else:
+        lmp.fix_external_set_energy_global("ext",1.0)
 
 class PythonExternal(unittest.TestCase):
     def testExternalCallback(self):
@@ -35,6 +39,7 @@ class PythonExternal(unittest.TestCase):
         lmp.set_fix_external_callback("ext",callback_one,lmp)
         lmp.command("run 10 post no")
         self.assertAlmostEqual(lmp.get_thermo("temp"),1.0/30.0,14)
+        self.assertAlmostEqual(lmp.get_thermo("pe"),1.0/8.0,14)
 
     def testExternalArray(self):
         """Test fix external from Python with pf/array"""