rename compute_free_energy to harmonic_free_energy and ensure grounds…

…tate energy is added when called from Harmonic

psiflow/free_energy/integration.py

@@ -11,19 +11,12 @@
 from psiflow.hamiltonians.hamiltonian import Hamiltonian, Zero
 from psiflow.sampling import SimulationOutput, Walker, sample
 from psiflow.sampling.walker import quench, randomize
-from psiflow.utils import multiply
+from psiflow.utils import compute_sum, multiply
 
 length = python_app(len, executors=["default_threads"])
 take_mean = python_app(np.mean, executors=["default_threads"])
 
 
-def _compute_sum(a, b):
-    return np.add(a, b)
-
-
-compute_sum = python_app(_compute_sum, executors=["default_threads"])
-
-
 @typeguard.typechecked
 def _integrate(x: np.ndarray, *args: float) -> np.ndarray:
     import scipy.integrate

psiflow/hamiltonians/_harmonic.py

@@ -13,7 +13,12 @@
 import psiflow
 from psiflow.geometry import Geometry, mass_weight
 from psiflow.hamiltonians.hamiltonian import Hamiltonian, evaluate_function
-from psiflow.utils import dump_json
+from psiflow.utils import compute_sum, dump_json, multiply
+
+
+@python_app(executors=["default_threads"])
+def get_energy(geometry: Geometry):
+    return geometry.energy
 
 
 @typeguard.typechecked
@@ -61,10 +66,6 @@ def serialize_calculator(self) -> DataFuture:
         def get_positions(geometry: Geometry):
             return geometry.per_atom.positions.copy().astype(float)
 
-        @python_app(executors=["default_threads"])
-        def get_energy(geometry: Geometry):
-            return geometry.energy
-
         return dump_json(
             hamiltonian=self.__class__.__name__,
             positions=get_positions(self.reference_geometry),
@@ -83,12 +84,15 @@ def compute_free_energy(
             self.hessian,
             self.reference_geometry,
         )
-        return compute_free_energy(
+        f_groundstate = get_energy(self.reference_geometry)
+        f_harmonic = harmonic_free_energy(
             frequencies,
             temperature=temperature,
             quantum=quantum,
             threshold=threshold,
         )
+        beta = 1 / (kB * temperature)
+        return compute_sum(multiply(f_groundstate, beta), f_harmonic)
 
     @staticmethod
     def deserialize_calculator(
@@ -150,7 +154,7 @@ def _compute_frequencies(hessian: np.ndarray, geometry: Geometry) -> np.ndarray:
 
 
 @typeguard.typechecked
-def _compute_free_energy(
+def _harmonic_free_energy(
     frequencies: Union[float, np.ndarray],
     temperature: float,
     quantum: bool = False,
@@ -176,4 +180,4 @@ def _compute_free_energy(
     return F
 
 
-compute_free_energy = python_app(_compute_free_energy, executors=["default_threads"])
+harmonic_free_energy = python_app(_harmonic_free_energy, executors=["default_threads"])

psiflow/utils.py

@@ -50,12 +50,11 @@ def setup_logger(module_name):
     return module_logger
 
 
-@typeguard.typechecked
-def _sum_integers(a: int, b: int) -> int:
-    return a + b
+def _compute_sum(a, b):
+    return np.add(a, b)
 
 
-sum_integers = python_app(_sum_integers, executors=["default_threads"])
+compute_sum = python_app(_compute_sum, executors=["default_threads"])
 
 
 @typeguard.typechecked

tests/test_sampling.py

@@ -313,7 +313,7 @@ def test_npt(dataset):
         walker.start.result().cell,
         output.state.result().cell,
     )
-    # volumea actually does change with barostat mode='nst' (?)
+    # volume actually does change with barostat mode='nst' (?)
     # assert np.allclose(
     #     np.linalg.det(walker.start.result().cell),
     #     np.linalg.det(output.state.result().cell),

tests/test_tools.py

@@ -2,7 +2,7 @@
 from ase.units import _c, second
 
 from psiflow.hamiltonians import EinsteinCrystal, get_mace_mp0
-from psiflow.hamiltonians._harmonic import compute_free_energy, compute_frequencies
+from psiflow.hamiltonians._harmonic import compute_frequencies, harmonic_free_energy
 from psiflow.tools import compute_harmonic, optimize, optimize_dataset
 
 
@@ -69,9 +69,9 @@ def test_dihydrogen(dataset_h2):
 
 def test_frequency_oscillator():
     for quantum in [True, False]:
-        f0 = compute_free_energy(1.0, 300, quantum=quantum).result()
-        f1 = compute_free_energy(1.1, 300, quantum=quantum).result()
+        f0 = harmonic_free_energy(1.0, 300, quantum=quantum).result()
+        f1 = harmonic_free_energy(1.1, 300, quantum=quantum).result()
         assert f1 > f0
 
-        f2 = compute_free_energy(1.0, 400, quantum=quantum).result()
+        f2 = harmonic_free_energy(1.0, 400, quantum=quantum).result()
         assert f0 > f2