Tutorial_crystal_hessian_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4ba0e8af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import sys\n",
    "import pickle\n",
    "\n",
    "from typing import Callable, Dict, List, Optional\n",
    "import haiku as hk\n",
    "import ase\n",
    "import ase.io\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import optax\n",
    "import yaml\n",
    "\n",
    "from phonax.datasets import (\n",
    "    datasets,\n",
    "    ph_datasets,\n",
    ")\n",
    "from phonax.optimizer import optimizer\n",
    "from phonax.energy_force_train import energy_force_train\n",
    "from phonax.loss import (\n",
    "    WeightedEnergyFrocesStressLoss,\n",
    "    crystalHessianLoss,\n",
    ")\n",
    "from phonax.nequip_model import NequIP_JAXMD_model\n",
    "\n",
    "from phonax.utils import (\n",
    "    create_directory_with_random_name,\n",
    ")\n",
    "\n",
    "from phonax.phonons_train import (\n",
    "    ph_evaluate,\n",
    "    hessian_train,\n",
    "    two_stage_hessian_train,\n",
    ")\n",
    "\n",
    "from phonax.predictors import (\n",
    "    predict_energy_forces_stress,\n",
    "    predict_crystal_hessian,\n",
    ")\n",
    "from phonax.phonons import (\n",
    "    atoms_to_ext_graph, \n",
    "    predict_gamma_spectra_filter\n",
    ")\n",
    "\n",
    "from phonax.data_utils import to_f32\n",
    "\n",
    "from phonax.trained_models import NequIP_JAXMD_molecule_model   \n",
    "\n",
    "jax.config.update(\"jax_debug_nans\", True)\n",
    "jax.config.update(\"jax_debug_infs\", True)\n",
    "np.set_printoptions(precision=3, suppress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbdbebf",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "Previously, we have been making the second derivative Hessian matrix predictions given an energy model, which is trained with zeroth and first order derivative data (i.e. energy and forces).\n",
    "Here in this tutorial, we will demonstrate an alternative view of the Hessian data, and use them as part of the training data (the Hessian matrix, or the eigenvalue spectrum) for the periodic crystal solids.\n",
    "\n",
    "As discussed in our work, this augments the training dataset in converging the energy models. Locally, the second derivative data improves the local curvatures of the energy model landscape.\n",
    "To demonstrate this here, we use periodic crystal cases as the examples, and focus on using the Hessian matrix as an extended data type to train the energy model/functional. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8f51379",
   "metadata": {},
   "source": [
    "## Energy model training with crystalline Hessian data\n",
    "\n",
    "This section focuses on using the Hessian matrix as an extended data type to train the energy model/functional for the interatomic potentials.\n",
    "The training procedure can be viewed as an extension to the conventional interatomic potential training using only the zeroth and first order training data for the energy landscape (i.e. energy and force).\n",
    "\n",
    "For the specific example, we use Si crystal (mp-149) example as discussed in our work.\n",
    "When trained with energy and force data for the energy model, we found the training can be improved by the augmented supercell geometries, which effectively create finer mesh grid in the momentum space sampling.\n",
    "However, in our training example here, we will only use the $1 \\times 1 \\times 1$ energy and force data, combined with the crystal Hessian data to achieve good phonon predictions, without supercell geometries for the energy and forces.\n",
    "\n",
    "Compared to the molecular cases, the derivation is more complicated for the periodic crystalline solids.\n",
    "This is due to the periodic boundary conditions which require the construction for the extended graphs in the computations.\n",
    "To describe the phonon states, one has to introduce crystal momentum $\\vec{k}$ as well from the Bloch theorem.\n",
    "In summary, the Hessian training involves the random selection for momentum $\\vec{k}$, and two (random) displacement patterns at this momentum. \n",
    "The one can project out the dynamical matrix element given these two displacement patterns. (In analogy to the molecular cases, one can also choose random atomic sites and two displacement vectors on these sites to perform the projection.)\n",
    "With enough projections (or slices) of the Hessians from the trained model, the energy model would converge to the groundtruth energy functional at all momentum $\\vec{k}$ and bewteen any states.\n",
    "\n",
    "One note about the random direction (or patterns) here. If the displacements are constrained to the x, y or z direcitons, the specific coordinate frame breaks the (rotational) equivariance in the training procedure, despite our energy model respects the equivarnace. The randomized projections restore the rotational equivariance.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4bf1adb",
   "metadata": {},
   "source": [
    "### Initialize the training procedures\n",
    "\n",
    "Here we load the training configurations and datasets, initialize the model, loss functions, and the predictors used for the next training steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ba810af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-17-20:48-nice-katrine\n",
      "2024-03-17-20:48-nice-katrine\n",
      "nums check 1500 9000 128\n",
      "Loaded 1000 training configurations from 'data/mp-149/mp-149-SC111_train.xyz'\n",
      "Loaded 200 validation configurations from 'data/mp-149/mp-149-SC333_valid.xyz'\n",
      "Total number of configurations: train=1000, valid=200, test=0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████| 1000/1000 [00:00<00:00, 16391.24it/s]\n",
      "100%|████████████████████████████████████████| 200/200 [00:00<00:00, 736.38it/s]\n",
      "0it [00:00, ?it/s]\n",
      "Pad train: 100%|███████████████████████████| 2000/2000 [00:06<00:00, 306.12it/s]\n",
      "Pad valid: 100%|█████████████████████████████| 100/100 [00:00<00:00, 305.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of configurations: train=2000, valid=100, test=0\n",
      "Compute the average number of neighbors: 27.992\n",
      "Do not normalize the radial basis (avg_r_min=None)\n",
      "Computed average Atomic Energies using least squares: {14: -5.531418045215}\n",
      "Create NequIP (JAX-MD version) with parameters {'use_sc': True, 'graph_net_steps': 2, 'hidden_irreps': '16x0e + 16x0o +12x1e + 12x1o + 8x2e +8x2o', 'nonlinearities': {'e': 'swish', 'o': 'tanh'}, 'r_max': 5.0, 'avg_num_neighbors': 27.992, 'avg_r_min': None, 'num_species': 100, 'radial_basis': <function bessel_basis at 0x7ff21a0435e0>, 'radial_envelope': <function soft_envelope at 0x7ff21a043670>}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-17 20:49:00.392487: W external/xla/xla/service/gpu/nvptx_compiler.cc:742] The NVIDIA driver's CUDA version is 12.0 which is older than the ptxas CUDA version (12.4.99). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_params: 58920\n"
     ]
    }
   ],
   "source": [
    "# Load the config files\n",
    "with open('data/mp-149/mp149-fh.yaml') as f:\n",
    "    config = yaml.load(f, Loader=yaml.FullLoader)\n",
    "        \n",
    "if 'save_dir_name' in config:\n",
    "    save_dir_name = config['save_dir_name']\n",
    "else:\n",
    "    save_dir_name = create_directory_with_random_name()\n",
    "print(save_dir_name)\n",
    "        \n",
    "# Save config\n",
    "with open(f\"{save_dir_name}/config.yaml\", \"w\") as f:\n",
    "    yaml.dump(config, f)\n",
    "    \n",
    "# Load the energy/force training/valid dataset\n",
    "EF_train_loader, EF_valid_loader, EF_test_loader, EF_r_max = datasets(\n",
    "    r_max = config[\"cutoff\"],\n",
    "    config_dataset = config[\"dataset\"],\n",
    ")\n",
    "\n",
    "# Load the Hessian training/valid dataset\n",
    "ph_train_loader, ph_valid_loader, ph_test_loader, ph_r_max, num_message_passing = ph_datasets(\n",
    "    r_max = config[\"cutoff\"],\n",
    "    config_dataset = config[\"ph_dataset\"],\n",
    "    num_message_passing = config['model']['num_layers'],\n",
    ")\n",
    "\n",
    "assert ph_r_max==EF_r_max\n",
    "r_max = EF_r_max\n",
    "\n",
    "# Initialize the NequIP model\n",
    "model_fn, params, num_message_passing = NequIP_JAXMD_model(\n",
    "    r_max=r_max,\n",
    "    atomic_energies_dict={},\n",
    "    train_graphs=EF_train_loader.graphs,\n",
    "    initialize_seed=config[\"model\"][\"seed\"],\n",
    "    num_species = config[\"model\"][\"num_species\"],\n",
    "    use_sc = True,\n",
    "    graph_net_steps = config[\"model\"][\"num_layers\"],\n",
    "    hidden_irreps = config[\"model\"][\"internal_irreps\"],\n",
    "    nonlinearities =  {'e': 'swish', 'o': 'tanh'},\n",
    "    save_dir_name = save_dir_name,\n",
    "    reload = config[\"initialization\"]['reload'] if 'reload' in config[\"initialization\"] else None,\n",
    ")\n",
    "    \n",
    "print(\"num_params:\", sum(p.size for p in jax.tree_util.tree_leaves(params)))\n",
    "\n",
    "# Define Loss functions for energy/force and Hessians\n",
    "EF_loss_fn = WeightedEnergyFrocesStressLoss(\n",
    "    energy_weight = config[\"training\"][\"energy_weight\"],\n",
    "    forces_weight = config[\"training\"][\"forces_weight\"],\n",
    "    stress_weight = config[\"training\"][\"stress_weight\"],\n",
    ")\n",
    "\n",
    "\n",
    "H_loss_fn = crystalHessianLoss()\n",
    "\n",
    "# Predictors for the energy/force/stress and Hessians\n",
    "EF_predictor = jax.jit(\n",
    "    lambda w, g: predict_energy_forces_stress(lambda *x: model_fn(w, *x), g)\n",
    ")\n",
    "\n",
    "H_predictor = jax.jit(\n",
    "    lambda w, g: predict_crystal_hessian(lambda *x: model_fn(w, *x), g)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5dfaf95",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e5d7b929-3e9a-450a-910b-b1db2641ee24",
   "metadata": {},
   "source": [
    "### Training with energy / force training data\n",
    "\n",
    "In this section, we demonstrate the energy model training by using only the energy and force training data, as is done in the conventional interatomic potential model training procedure.\n",
    "Note in this simple propane molecule example, we only use four molecular configurations in the training dataset.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f0ceaea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimizer num_params: 176762\n",
      "Started training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_train:   0%|                                         | 0/7 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiled function `model` for args:\n",
      "cache size: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_train: 100%|██████████████████████████| 7/7 [00:09<00:00,  1.29s/it, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=354.1589, mae_e_per_atom=2768.2 meV, mae_f=1523.9 meV/Å, mae_s=27.9 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid:   0%|                                        | 0/40 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiled function `model` for args:\n",
      "cache size: 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:08<00:00,  4.68it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=310.6482, mae_e_per_atom=2763.0 meV, mae_f=1373.5 meV/Å, mae_s=9.1 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 0:   2%|▏            | 8/500 [00:15<11:46,  1.44s/it, loss=4.978]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiled function `update_fn` for args:\n",
      "Outout: loss= 330.353\n",
      "Compilation time: 15.633s, cache size: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 0: 100%|███████████| 500/500 [00:23<00:00, 20.83it/s, loss=0.090]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 57.04it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.1209, mae_e_per_atom=51.2 meV, mae_f=24.3 meV/Å, mae_s=68.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 59.81it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=3.1920, mae_e_per_atom=14.5 meV, mae_f=165.5 meV/Å, mae_s=28.3 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 1: 100%|███████████| 500/500 [00:08<00:00, 61.77it/s, loss=0.096]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 60.56it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.0811, mae_e_per_atom=38.1 meV, mae_f=20.8 meV/Å, mae_s=62.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.73it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=2.3827, mae_e_per_atom=10.6 meV, mae_f=143.2 meV/Å, mae_s=19.9 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 2: 100%|███████████| 500/500 [00:08<00:00, 61.05it/s, loss=0.051]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 58.21it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.0578, mae_e_per_atom=28.3 meV, mae_f=18.4 meV/Å, mae_s=60.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 67.52it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=2.0508, mae_e_per_atom=11.3 meV, mae_f=132.9 meV/Å, mae_s=14.9 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 3: 100%|███████████| 500/500 [00:08<00:00, 61.28it/s, loss=0.050]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 60.07it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.0440, mae_e_per_atom=20.4 meV, mae_f=17.1 meV/Å, mae_s=58.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.87it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=1.9256, mae_e_per_atom=13.8 meV, mae_f=128.2 meV/Å, mae_s=13.3 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 4: 100%|███████████| 500/500 [00:08<00:00, 62.17it/s, loss=0.032]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 43.55it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.0367, mae_e_per_atom=14.9 meV, mae_f=16.2 meV/Å, mae_s=57.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 68.01it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=1.9194, mae_e_per_atom=15.9 meV, mae_f=127.9 meV/Å, mae_s=13.9 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 5: 100%|███████████| 500/500 [00:08<00:00, 61.08it/s, loss=0.039]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 58.88it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.0330, mae_e_per_atom=11.2 meV, mae_f=15.6 meV/Å, mae_s=57.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.76it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=1.9192, mae_e_per_atom=16.5 meV, mae_f=128.0 meV/Å, mae_s=15.1 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 6: 100%|███████████| 500/500 [00:08<00:00, 61.93it/s, loss=0.032]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 60.87it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0299, mae_e_per_atom=8.9 meV, mae_f=15.1 meV/Å, mae_s=56.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.03it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=1.9328, mae_e_per_atom=17.2 meV, mae_f=128.5 meV/Å, mae_s=16.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 7: 100%|███████████| 500/500 [00:08<00:00, 60.88it/s, loss=0.029]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 58.04it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0285, mae_e_per_atom=6.8 meV, mae_f=14.9 meV/Å, mae_s=57.1 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 66.62it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=1.9493, mae_e_per_atom=17.1 meV, mae_f=129.0 meV/Å, mae_s=17.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 8: 100%|███████████| 500/500 [00:08<00:00, 59.55it/s, loss=0.026]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 56.83it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0273, mae_e_per_atom=5.4 meV, mae_f=14.8 meV/Å, mae_s=57.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 60.40it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=1.9598, mae_e_per_atom=18.1 meV, mae_f=129.3 meV/Å, mae_s=18.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 9: 100%|███████████| 500/500 [00:08<00:00, 60.13it/s, loss=0.040]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 56.82it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_train: loss=0.0266, mae_e_per_atom=4.3 meV, mae_f=14.7 meV/Å, mae_s=57.1 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 60.23it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_valid: loss=1.9724, mae_e_per_atom=18.0 meV, mae_f=129.6 meV/Å, mae_s=18.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 10: 100%|██████████| 500/500 [00:08<00:00, 61.58it/s, loss=0.024]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 58.26it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_train: loss=0.0262, mae_e_per_atom=3.7 meV, mae_f=14.6 meV/Å, mae_s=57.1 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.63it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_valid: loss=1.9867, mae_e_per_atom=18.0 meV, mae_f=130.0 meV/Å, mae_s=19.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 11: 100%|██████████| 500/500 [00:08<00:00, 60.92it/s, loss=0.029]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 59.70it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_train: loss=0.0248, mae_e_per_atom=3.2 meV, mae_f=14.2 meV/Å, mae_s=56.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 66.29it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_valid: loss=2.0080, mae_e_per_atom=17.9 meV, mae_f=130.7 meV/Å, mae_s=19.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 12: 100%|██████████| 500/500 [00:08<00:00, 59.58it/s, loss=0.027]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.43it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_train: loss=0.0252, mae_e_per_atom=2.9 meV, mae_f=14.3 meV/Å, mae_s=57.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 59.40it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_valid: loss=2.0368, mae_e_per_atom=18.1 meV, mae_f=131.8 meV/Å, mae_s=19.6 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 13: 100%|██████████| 500/500 [00:08<00:00, 59.57it/s, loss=0.026]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 57.63it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_train: loss=0.0247, mae_e_per_atom=2.6 meV, mae_f=14.2 meV/Å, mae_s=57.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 60.09it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_valid: loss=2.0768, mae_e_per_atom=18.4 meV, mae_f=133.1 meV/Å, mae_s=19.8 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 14: 100%|██████████| 500/500 [00:08<00:00, 60.09it/s, loss=0.028]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 57.83it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_train: loss=0.0240, mae_e_per_atom=2.5 meV, mae_f=14.1 meV/Å, mae_s=57.0 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 65.10it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_valid: loss=2.1226, mae_e_per_atom=18.5 meV, mae_f=134.6 meV/Å, mae_s=20.0 meV/Å³\n",
      "Training complete\n"
     ]
    }
   ],
   "source": [
    "# force only training\n",
    "\n",
    "Fonly_params = params.copy()\n",
    "\n",
    "gradient_transform, steps_per_interval, max_num_intervals = optimizer(\n",
    "    lr = config[\"training\"][\"learning_rate\"],\n",
    "    max_num_intervals = config[\"training\"][\"max_num_intervals\"],\n",
    "    steps_per_interval = config[\"training\"][\"steps_per_interval\"],\n",
    "    # weight_decay = config[\"training\"][\"weight_decay\"],\n",
    ")\n",
    "optimizer_state = gradient_transform.init(params)\n",
    "print(\"optimizer num_params:\", sum(p.size for p in jax.tree_util.tree_leaves(optimizer_state)))\n",
    "    \n",
    "train_out = energy_force_train(\n",
    "    EF_predictor,\n",
    "    Fonly_params,\n",
    "    optimizer_state,\n",
    "    EF_train_loader,\n",
    "    EF_valid_loader,\n",
    "    EF_test_loader,\n",
    "    gradient_transform,\n",
    "    EF_loss_fn,\n",
    "    max_num_intervals,\n",
    "    steps_per_interval,\n",
    "    save_dir_name,\n",
    "    ema_decay = config[\"training\"][\"ema_decay\"],\n",
    "    patience = config[\"training\"][\"patience\"],\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2e78eb5",
   "metadata": {},
   "source": [
    "### Training with crystal Hessians\n",
    "\n",
    "In the next section, we show the energy model training with additional second order Hessian data beynd the force data at the first order derivative.\n",
    "Specifically, there are two tages in the training procedure.\n",
    "The first stage, or the warm-up stage, uses only the energy and force data to initialize the energy model training.\n",
    "After trained with several warm-up steps, the second training stage adds the Hessian data in computing the loss function.\n",
    "At the second stage, the update gradients come from both the energy / force training losses, and the Hessian losses, and we combine the two (weighted sum) to get the overall training gradient.\n",
    "\n",
    "From the training results below, one can see the valid force MAE further decreases once we have included the Hessian training data in the second stage.\n",
    "\n",
    "We note that, it remains an open question to find the optimal training procedure in adding the second order Hessian data. For example, one can perform the training without warm-up steps, or gradually ramp up the training weights associated with the Hessian losses.\n",
    "It would be relevant to examine the energy model converged for their predictions and stability under molecular dynamics simulations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "99a2cad3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimizer num_params: 176762\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.72it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=350.4000, mae_e_per_atom=2766.672 meV, mae_f=1520.624 meV/Å, mae_s=27.592 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.54it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=310.6482, mae_e_per_atom=2762.956 meV, mae_f=1373.478 meV/Å, mae_s=9.123 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 0: 100%|█████████████████| 100/100 [00:08<00:00, 11.55it/s, loss=0.334]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.36it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.3489, mae_e_per_atom=63.654 meV, mae_f=45.458 meV/Å, mae_s=76.216 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.58it/s, n=200]\n",
      "/home/shiangfang/miniconda3/envs/phonax-test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:3504: RuntimeWarning: Mean of empty slice.\n",
      "  return _methods._mean(a, axis=axis, dtype=dtype,\n",
      "/home/shiangfang/miniconda3/envs/phonax-test/lib/python3.9/site-packages/numpy/core/_methods.py:129: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  ret = ret.dtype.type(ret / rcount)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=6.2858, mae_e_per_atom=24.377 meV, mae_f=225.995 meV/Å, mae_s=38.006 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 1: 100%|█████████████████| 100/100 [00:01<00:00, 56.50it/s, loss=0.237]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.81it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.2397, mae_e_per_atom=60.819 meV, mae_f=37.990 meV/Å, mae_s=73.932 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.27it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=5.0234, mae_e_per_atom=20.326 meV, mae_f=205.108 meV/Å, mae_s=35.567 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 2: 100%|█████████████████| 100/100 [00:01<00:00, 56.37it/s, loss=0.201]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.44it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.1708, mae_e_per_atom=57.545 meV, mae_f=30.172 meV/Å, mae_s=72.316 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.99it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=4.1606, mae_e_per_atom=18.014 meV, mae_f=188.633 meV/Å, mae_s=33.427 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 3: 100%|█████████████████| 100/100 [00:01<00:00, 56.24it/s, loss=0.131]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 53.61it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.1324, mae_e_per_atom=54.698 meV, mae_f=25.123 meV/Å, mae_s=69.920 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.38it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=3.4726, mae_e_per_atom=15.964 meV, mae_f=172.589 meV/Å, mae_s=30.537 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 4: 100%|█████████████████| 100/100 [00:01<00:00, 57.77it/s, loss=0.116]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.45it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.1156, mae_e_per_atom=50.600 meV, mae_f=23.319 meV/Å, mae_s=67.989 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.87it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=3.1076, mae_e_per_atom=14.023 meV, mae_f=162.980 meV/Å, mae_s=27.840 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 5: 100%|████████████████| 100/100 [00:01<00:00, 56.08it/s, loss=0.0991]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.56it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.1044, mae_e_per_atom=47.462 meV, mae_f=22.521 meV/Å, mae_s=66.569 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.94it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=2.8853, mae_e_per_atom=13.012 meV, mae_f=156.960 meV/Å, mae_s=25.590 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 6: 100%|████████████████| 100/100 [00:01<00:00, 57.62it/s, loss=0.0894]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.45it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0983, mae_e_per_atom=44.340 meV, mae_f=22.062 meV/Å, mae_s=65.483 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.10it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=2.7210, mae_e_per_atom=12.173 meV, mae_f=152.400 meV/Å, mae_s=23.739 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 7: 100%|████████████████| 100/100 [00:01<00:00, 56.39it/s, loss=0.0808]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 52.94it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0914, mae_e_per_atom=41.924 meV, mae_f=21.627 meV/Å, mae_s=64.558 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.12it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=2.5975, mae_e_per_atom=11.521 meV, mae_f=148.789 meV/Å, mae_s=22.075 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 8: 100%|████████████████| 100/100 [00:01<00:00, 57.05it/s, loss=0.0754]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 52.38it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0856, mae_e_per_atom=39.410 meV, mae_f=21.086 meV/Å, mae_s=63.525 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.53it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=2.5067, mae_e_per_atom=10.993 meV, mae_f=146.045 meV/Å, mae_s=20.640 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 9: 100%|████████████████| 100/100 [00:01<00:00, 57.85it/s, loss=0.0805]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 40.51it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=0.0790, mae_e_per_atom=37.290 meV, mae_f=20.649 meV/Å, mae_s=62.700 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.58it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=2.4362, mae_e_per_atom=10.849 meV, mae_f=143.879 meV/Å, mae_s=19.314 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 126.35it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 0: hessian train hessian MAE 0.479 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 133.07it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 0: hessian valid hessian MAE 0.433 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 0: 100%|████████████████| 100/100 [00:18<00:00,  5.44it/s, loss=0.0448]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.23it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.0757, mae_e_per_atom=34.710 meV, mae_f=20.462 meV/Å, mae_s=61.519 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.33it/s, n=200]\n",
      "/home/shiangfang/miniconda3/envs/phonax-test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:3504: RuntimeWarning: Mean of empty slice.\n",
      "  return _methods._mean(a, axis=axis, dtype=dtype,\n",
      "/home/shiangfang/miniconda3/envs/phonax-test/lib/python3.9/site-packages/numpy/core/_methods.py:129: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  ret = ret.dtype.type(ret / rcount)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=2.3719, mae_e_per_atom=10.849 meV, mae_f=141.880 meV/Å, mae_s=18.027 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 126.46it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 1: hessian train hessian MAE 0.457 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 132.97it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 1: hessian valid hessian MAE 0.411 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 1: 100%|████████████████| 100/100 [00:02<00:00, 36.85it/s, loss=0.0518]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.35it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.0740, mae_e_per_atom=32.935 meV, mae_f=20.662 meV/Å, mae_s=61.145 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.57it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=2.3144, mae_e_per_atom=10.936 meV, mae_f=140.052 meV/Å, mae_s=16.943 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.06it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 2: hessian train hessian MAE 0.433 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 136.88it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 2: hessian valid hessian MAE 0.392 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 2: 100%|████████████████| 100/100 [00:02<00:00, 36.68it/s, loss=0.0451]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.99it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.0715, mae_e_per_atom=31.479 meV, mae_f=20.420 meV/Å, mae_s=60.747 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 62.99it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=2.2295, mae_e_per_atom=10.423 meV, mae_f=137.292 meV/Å, mae_s=16.206 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 126.48it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 3: hessian train hessian MAE 0.418 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 131.34it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 3: hessian valid hessian MAE 0.380 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 3: 100%|████████████████| 100/100 [00:02<00:00, 37.01it/s, loss=0.0383]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.34it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.0677, mae_e_per_atom=29.833 meV, mae_f=20.019 meV/Å, mae_s=60.252 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.81it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=2.1422, mae_e_per_atom=10.807 meV, mae_f=134.483 meV/Å, mae_s=15.545 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.51it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 4: hessian train hessian MAE 0.397 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 131.82it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 4: hessian valid hessian MAE 0.367 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 4: 100%|████████████████| 100/100 [00:02<00:00, 36.56it/s, loss=0.0463]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.25it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.0651, mae_e_per_atom=28.728 meV, mae_f=19.844 meV/Å, mae_s=59.665 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.69it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=2.0504, mae_e_per_atom=10.341 meV, mae_f=131.388 meV/Å, mae_s=15.051 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.54it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 5: hessian train hessian MAE 0.379 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 134.84it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 5: hessian valid hessian MAE 0.352 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 5: 100%|████████████████| 100/100 [00:02<00:00, 36.80it/s, loss=0.0365]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 56.47it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.0633, mae_e_per_atom=27.656 meV, mae_f=19.825 meV/Å, mae_s=59.839 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.93it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=1.9543, mae_e_per_atom=9.357 meV, mae_f=128.219 meV/Å, mae_s=14.608 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.42it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 6: hessian train hessian MAE 0.355 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 115.22it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 6: hessian valid hessian MAE 0.334 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 6: 100%|████████████████| 100/100 [00:02<00:00, 37.37it/s, loss=0.0383]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.68it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0605, mae_e_per_atom=26.427 meV, mae_f=19.536 meV/Å, mae_s=59.127 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.77it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=1.8760, mae_e_per_atom=9.807 meV, mae_f=125.472 meV/Å, mae_s=14.239 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 129.90it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 7: hessian train hessian MAE 0.335 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 136.03it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 7: hessian valid hessian MAE 0.318 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 7: 100%|████████████████| 100/100 [00:02<00:00, 36.73it/s, loss=0.0356]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.48it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0568, mae_e_per_atom=25.827 meV, mae_f=18.935 meV/Å, mae_s=58.800 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.59it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=1.8115, mae_e_per_atom=9.152 meV, mae_f=123.114 meV/Å, mae_s=13.962 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.87it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 8: hessian train hessian MAE 0.317 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 131.98it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 8: hessian valid hessian MAE 0.303 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 8: 100%|████████████████| 100/100 [00:02<00:00, 36.67it/s, loss=0.0306]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 56.01it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0568, mae_e_per_atom=25.080 meV, mae_f=19.158 meV/Å, mae_s=58.743 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 58.25it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=1.7514, mae_e_per_atom=8.762 meV, mae_f=120.722 meV/Å, mae_s=13.726 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.66it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 9: hessian train hessian MAE 0.300 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 115.79it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 9: hessian valid hessian MAE 0.287 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 9: 100%|████████████████| 100/100 [00:02<00:00, 37.36it/s, loss=0.0352]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.11it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_train: loss=0.0559, mae_e_per_atom=23.885 meV, mae_f=19.127 meV/Å, mae_s=58.606 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.16it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_valid: loss=1.7283, mae_e_per_atom=8.120 meV, mae_f=119.191 meV/Å, mae_s=13.502 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 129.01it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 10: hessian train hessian MAE 0.280 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 117.61it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 10: hessian valid hessian MAE 0.268 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 10: 100%|███████████████| 100/100 [00:02<00:00, 37.00it/s, loss=0.0266]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.38it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_train: loss=0.0512, mae_e_per_atom=22.784 meV, mae_f=18.466 meV/Å, mae_s=58.023 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.86it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_valid: loss=1.7035, mae_e_per_atom=7.555 meV, mae_f=117.914 meV/Å, mae_s=13.345 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.34it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 11: hessian train hessian MAE 0.266 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 118.22it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 11: hessian valid hessian MAE 0.257 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 11: 100%|███████████████| 100/100 [00:02<00:00, 37.36it/s, loss=0.0287]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.41it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_train: loss=0.0493, mae_e_per_atom=21.998 meV, mae_f=18.189 meV/Å, mae_s=57.843 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.20it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_valid: loss=1.6886, mae_e_per_atom=7.192 meV, mae_f=116.797 meV/Å, mae_s=13.161 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.77it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 12: hessian train hessian MAE 0.258 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 118.06it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 12: hessian valid hessian MAE 0.244 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 12: 100%|███████████████| 100/100 [00:02<00:00, 37.27it/s, loss=0.0287]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.42it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_train: loss=0.0476, mae_e_per_atom=21.239 meV, mae_f=17.903 meV/Å, mae_s=57.252 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.48it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_valid: loss=1.6668, mae_e_per_atom=6.643 meV, mae_f=115.970 meV/Å, mae_s=13.018 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.62it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 13: hessian train hessian MAE 0.252 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 116.46it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 13: hessian valid hessian MAE 0.236 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 13: 100%|███████████████| 100/100 [00:02<00:00, 37.32it/s, loss=0.0279]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 39.25it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_train: loss=0.0456, mae_e_per_atom=20.295 meV, mae_f=17.663 meV/Å, mae_s=57.203 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.79it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_valid: loss=1.6496, mae_e_per_atom=6.278 meV, mae_f=115.209 meV/Å, mae_s=12.963 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 125.37it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 14: hessian train hessian MAE 0.247 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 136.11it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 14: hessian valid hessian MAE 0.233 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 14: 100%|███████████████| 100/100 [00:02<00:00, 36.40it/s, loss=0.0251]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 57.01it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_train: loss=0.0445, mae_e_per_atom=19.449 meV, mae_f=17.601 meV/Å, mae_s=57.338 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 61.98it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_valid: loss=1.6303, mae_e_per_atom=5.856 meV, mae_f=114.142 meV/Å, mae_s=12.927 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 126.99it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 15: hessian train hessian MAE 0.245 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 133.30it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 15: hessian valid hessian MAE 0.228 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 15: 100%|███████████████| 100/100 [00:02<00:00, 36.79it/s, loss=0.0211]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 56.04it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 16: eval_train: loss=0.0432, mae_e_per_atom=18.529 meV, mae_f=17.457 meV/Å, mae_s=57.112 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.29it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 16: eval_valid: loss=1.6119, mae_e_per_atom=5.500 meV, mae_f=113.539 meV/Å, mae_s=12.937 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.04it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 16: hessian train hessian MAE 0.243 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 132.96it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 16: hessian valid hessian MAE 0.226 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 16: 100%|███████████████| 100/100 [00:02<00:00, 36.78it/s, loss=0.0261]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.04it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 17: eval_train: loss=0.0422, mae_e_per_atom=17.726 meV, mae_f=17.330 meV/Å, mae_s=57.210 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.93it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 17: eval_valid: loss=1.5932, mae_e_per_atom=5.649 meV, mae_f=112.603 meV/Å, mae_s=12.966 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.99it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 17: hessian train hessian MAE 0.242 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 134.87it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 17: hessian valid hessian MAE 0.224 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 17: 100%|███████████████| 100/100 [00:02<00:00, 36.24it/s, loss=0.0299]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.03it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 18: eval_train: loss=0.0410, mae_e_per_atom=17.018 meV, mae_f=17.037 meV/Å, mae_s=57.282 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.85it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 18: eval_valid: loss=1.5857, mae_e_per_atom=4.935 meV, mae_f=112.251 meV/Å, mae_s=13.025 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.46it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 18: hessian train hessian MAE 0.239 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 117.11it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 18: hessian valid hessian MAE 0.223 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 18: 100%|███████████████| 100/100 [00:02<00:00, 37.12it/s, loss=0.0274]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.91it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 19: eval_train: loss=0.0395, mae_e_per_atom=16.140 meV, mae_f=16.893 meV/Å, mae_s=57.456 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.96it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 19: eval_valid: loss=1.5846, mae_e_per_atom=5.065 meV, mae_f=112.129 meV/Å, mae_s=13.094 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.57it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 19: hessian train hessian MAE 0.238 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 116.97it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 19: hessian valid hessian MAE 0.223 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 19: 100%|███████████████| 100/100 [00:02<00:00, 37.37it/s, loss=0.0224]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 53.96it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 20: eval_train: loss=0.0406, mae_e_per_atom=15.505 meV, mae_f=17.174 meV/Å, mae_s=57.060 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.87it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 20: eval_valid: loss=1.5896, mae_e_per_atom=4.675 meV, mae_f=112.198 meV/Å, mae_s=13.170 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.58it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 20: hessian train hessian MAE 0.235 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 118.04it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 20: hessian valid hessian MAE 0.222 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 20: 100%|███████████████| 100/100 [00:02<00:00, 37.31it/s, loss=0.0267]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.82it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 21: eval_train: loss=0.0376, mae_e_per_atom=14.525 meV, mae_f=16.640 meV/Å, mae_s=56.267 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 56.36it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 21: eval_valid: loss=1.5729, mae_e_per_atom=4.689 meV, mae_f=111.393 meV/Å, mae_s=13.264 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.88it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 21: hessian train hessian MAE 0.233 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 118.12it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 21: hessian valid hessian MAE 0.217 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 21: 100%|███████████████| 100/100 [00:02<00:00, 37.33it/s, loss=0.0252]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 54.99it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 22: eval_train: loss=0.0381, mae_e_per_atom=14.368 meV, mae_f=16.729 meV/Å, mae_s=56.615 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.77it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 22: eval_valid: loss=1.5825, mae_e_per_atom=4.284 meV, mae_f=111.426 meV/Å, mae_s=13.348 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 129.60it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 22: hessian train hessian MAE 0.231 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 113.55it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 22: hessian valid hessian MAE 0.216 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 22: 100%|███████████████| 100/100 [00:02<00:00, 37.43it/s, loss=0.0193]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 55.68it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 23: eval_train: loss=0.0382, mae_e_per_atom=13.771 meV, mae_f=16.852 meV/Å, mae_s=56.329 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 57.64it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 23: eval_valid: loss=1.5838, mae_e_per_atom=4.199 meV, mae_f=111.130 meV/Å, mae_s=13.455 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 128.79it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 23: hessian train hessian MAE 0.229 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 117.03it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 23: hessian valid hessian MAE 0.210 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 23: 100%|███████████████| 100/100 [00:02<00:00, 37.33it/s, loss=0.0281]\n",
      "eval_train: 100%|██████████████████████████| 7/7 [00:00<00:00, 41.06it/s, n=889]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 24: eval_train: loss=0.0373, mae_e_per_atom=13.259 meV, mae_f=16.614 meV/Å, mae_s=56.476 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|████████████████████████| 40/40 [00:00<00:00, 63.57it/s, n=200]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 24: eval_valid: loss=1.5729, mae_e_per_atom=4.049 meV, mae_f=111.037 meV/Å, mae_s=13.589 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 1000/1000 [00:07<00:00, 127.20it/s, n_graphs=2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 24: hessian train hessian MAE 0.227 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 50/50 [00:00<00:00, 135.60it/s, n_graphs=100]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 24: hessian valid hessian MAE 0.210 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 24: 100%|███████████████| 100/100 [00:02<00:00, 37.38it/s, loss=0.0193]\n"
     ]
    }
   ],
   "source": [
    "FH_params = params.copy()\n",
    "\n",
    "# restart the optimizer state\n",
    "gradient_transform, steps_per_interval, max_num_intervals = optimizer(\n",
    "    lr = config[\"training\"][\"learning_rate\"],\n",
    "    max_num_intervals = 20, #config[\"training\"][\"max_num_intervals\"],\n",
    "    steps_per_interval = 100, # config[\"training\"][\"steps_per_interval\"],\n",
    "    # weight_decay = config[\"training\"][\"weight_decay\"],\n",
    ")\n",
    "optimizer_state = gradient_transform.init(params)\n",
    "print(\"optimizer num_params:\", sum(p.size for p in jax.tree_util.tree_leaves(optimizer_state)))\n",
    "\n",
    "two_stage_hessian_train(\n",
    "    energy_forces_stress_predictor=EF_predictor,\n",
    "    phonon_predictor=H_predictor,\n",
    "    params=FH_params,\n",
    "    gradient_transform=gradient_transform,\n",
    "    optimizer_state=optimizer_state,\n",
    "    steps_per_interval=steps_per_interval,\n",
    "    EF_loss_fn=EF_loss_fn,\n",
    "    H_loss_fn=H_loss_fn,\n",
    "    EF_train_loader=EF_train_loader,\n",
    "    EF_valid_loader=EF_valid_loader,\n",
    "    H_train_loader=ph_train_loader,\n",
    "    H_valid_loader=ph_valid_loader,\n",
    "    warmup_num_intervals=10 ,#config['training']['warmup_num_intervals'],\n",
    "    phtrain_num_intervals=25  ,#config['training']['phtrain_num_intervals'],\n",
    "    periodic_crystal=True,\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2057e4cf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04c7d71",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68156e4a-22f9-4956-9540-8c7e9bc1b2a9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}