Add formation energy comparison for MACE and CHGNet

Author: nedtaylor

example/python_pkg/visualisation/descriptor.ipynb

@@ -12,6 +12,7 @@
     "from ase import Atoms\n",
     "from raffle.generator import raffle_generator\n",
     "from mace.calculators import mace_mp\n",
+    "from chgnet.model import CHGNetCalculator\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd"
@@ -26,7 +27,9 @@
    "source": [
     "generator = raffle_generator()\n",
     "generator.distributions.set_history_len(10)\n",
-    "calc = mace_mp(model=\"medium\", dispersion=False, default_dtype=\"float32\", device='cpu')"
+    "mace = mace_mp(model=\"medium\", dispersion=False, default_dtype=\"float32\", device='cpu')\n",
+    "chgnet = CHGNetCalculator()\n",
+    "calc = mace"
    ]
   },
   {
@@ -232,10 +235,23 @@
    "outputs": [],
    "source": [
     "viability_dict = {}\n",
+    "ab_stack.calc = mace\n",
+    "ab_stack_energy_per_atom_mace = ab_stack.get_potential_energy() / len(ab_stack)\n",
+    "ab_stack.calc = chgnet\n",
+    "ab_stack_energy_per_atom_chgnet = ab_stack.get_potential_energy() / len(ab_stack)\n",
     "for struc_name, structure in structures_dict.items():\n",
+    "    structure.calc = mace\n",
+    "    total_energy_mace = structure.get_potential_energy()\n",
+    "    struc_energy_per_atom_mace = structure.get_potential_energy() / len(structure)\n",
+    "    struc_chgnet = structure.copy()\n",
+    "    struc_chgnet.calc = chgnet\n",
+    "    total_energy_chgnet = struc_chgnet.get_potential_energy()\n",
+    "    struc_energy_per_atom_chgnet = struc_chgnet.get_potential_energy() / len(struc_chgnet)\n",
     "    viability_dict[struc_name] = {\n",
-    "        \"total energy\": structure.get_potential_energy()/len(structure),\n",
-    "            \"formation energy\": structure.get_potential_energy()/len(structure) - ab_stack.get_potential_energy()/len(ab_stack),\n",
+    "        \"total energy (MACE)\": total_energy_mace,\n",
+    "        \"total energy (CHGNet)\": total_energy_chgnet,\n",
+    "        \"formation energy (MACE)\": struc_energy_per_atom_mace - ab_stack_energy_per_atom_mace,\n",
+    "        \"formation energy (CHGNet)\": struc_energy_per_atom_chgnet - ab_stack_energy_per_atom_chgnet,\n",
     "    }\n",
     "for key, value in n_body_dict.items():\n",
     "    generator.distributions.set_radius_distance_tol(value)\n",