From b65e0d18bfe19f72192e46dde980b8b180346c4c Mon Sep 17 00:00:00 2001
From: Logan Ward <ward.logan.t@gmail.com>
Date: Wed, 15 Nov 2023 14:18:29 -0500
Subject: [PATCH] Compute heat capacities too

---
 jitterbug/compare.py  | 27 +++++++++++++++++++++++----
 pyproject.toml        |  1 +
 tests/test_compare.py |  3 +++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/jitterbug/compare.py b/jitterbug/compare.py
index 310ffeb..04f8e81 100644
--- a/jitterbug/compare.py
+++ b/jitterbug/compare.py
@@ -2,8 +2,10 @@
 from dataclasses import dataclass
 
 import ase
+from ase import units
 import numpy as np
 from ase.vibrations import VibrationsData
+from pmutt.statmech import StatMech, presets
 
 
 @dataclass
@@ -12,9 +14,15 @@ class HessianQuality:
 
     # Thermodynamics
     zpe: float
-    """Zero point energy (eV)"""
+    """Zero point energy (kcal/mol)"""
     zpe_error: float
     """Different between the ZPE and the target one"""
+    cp: list[float]
+    """Heat capacity as a function of temperature (units: kcal/mol/K)"""
+    cp_error: list[float]
+    """Difference between known and approximate heat capacity as a function of temperature (units: kcal/mol/K)"""
+    temps: list[float]
+    """Temperatures at which Cp was evaluated (units: K)"""
 
     # Vibrations
     vib_freqs: list[float]
@@ -47,11 +55,22 @@ def compare_hessians(atoms: ase.Atoms, known_hessian: np.ndarray, approx_hessian
     freq_error = np.subtract(approx_freqs[is_real], known_freqs[is_real])
     freq_mae = np.abs(freq_error).mean()
 
+    # Compare the enthalpy and heat capacity
+    known_harm = StatMech(vib_wavenumbers=np.real(known_freqs[is_real]), atoms=atoms, symmetrynumber=1, **presets['harmonic'])
+    approx_harm = StatMech(vib_wavenumbers=np.real(approx_freqs[is_real]), atoms=atoms, symmetrynumber=1, **presets['harmonic'])
+
+    temps = np.linspace(1., 373, 128)
+    known_cp = np.array([known_harm.get_Cp('kcal/mol/K', T=t) for t in temps])
+    approx_cp = np.array([approx_harm.get_Cp('kcal/mol/K', T=t) for t in temps])
+
     # Assemble into a result object
     return HessianQuality(
-        zpe=approx_vibs.get_zero_point_energy(),
-        zpe_error=(approx_vibs.get_zero_point_energy() - known_vibs.get_zero_point_energy()),
+        zpe=approx_vibs.get_zero_point_energy() * units.mol / units.kcal,
+        zpe_error=(approx_vibs.get_zero_point_energy() - known_vibs.get_zero_point_energy()) * units.mol / units.kcal,
         vib_freqs=np.real(approx_freqs[is_real]).tolist(),
         vib_errors=np.abs(freq_error),
-        vib_mae=freq_mae
+        vib_mae=freq_mae,
+        cp=approx_cp.tolist(),
+        cp_error=(known_cp - approx_cp).tolist(),
+        temps=temps.tolist()
     )
diff --git a/pyproject.toml b/pyproject.toml
index 0104d3c..738120b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,7 @@ dependencies = [
     "colmena==0.5.*",
     "parsl>=2023.04",
     "ase>3.22",
+    "pmutt>=1.4.9",
     "tqdm"
 ]
 
diff --git a/tests/test_compare.py b/tests/test_compare.py
index e2b471b..8bde1e1 100644
--- a/tests/test_compare.py
+++ b/tests/test_compare.py
@@ -2,6 +2,7 @@
 
 from ase.vibrations import VibrationsData
 from pytest import fixture
+import numpy as np
 
 from jitterbug.compare import compare_hessians
 
@@ -16,3 +17,5 @@ def example_hess() -> VibrationsData:
 def test_compare(example_hess):
     comp = compare_hessians(example_hess.get_atoms(), example_hess.get_hessian_2d(), example_hess.get_hessian_2d())
     assert comp.zpe_error == 0.
+    assert np.ndim(comp.cp_error) == 1
+    assert np.mean(comp.cp_error) == 0.