Tutorial_molecular_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",
    "    HessianLoss,\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_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).\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 molecular cases as the examples. There are two parts of this tutorial, the first one focuses on using the Hessian matrix as an extended data type to train the energy model/functional. The second part uses the eigenvalues of the Hessian matrix instead. The implications are the connections to the experimental probes for the vibrational modes such as IR / Raman / inelastic neutron scatterings.\n",
    "This enables the energy model training or fine-tuning with realistic experimental data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8f51379",
   "metadata": {},
   "source": [
    "## Energy model training with molecular 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 consdier the propane molecules as discussed in our work.\n",
    "The structure of a propane molecule is mainly defined by the two torsion angles $\\phi_1, \\phi_2$ for the two C-C bonds.\n",
    "Due to the symmetry, we only need to consider the range $[0,2\\pi/3]$ for $\\phi_i$.\n",
    "In the valid dataset for the forces (in deriving the force MAE), we have scanned the $\\phi_1, \\phi_2$ angles ($31 \\times 31$ grid). At each torsion angle combination, we also add a small Gaussian noise to the atomic coordinates.\n",
    "\n",
    "In the training dataset, we only use four randomly generated propane structures (randomized torsion angles and perturbed atomic coordinates).\n",
    "For the conventional training procedure, only the energy / force data is used, while the full molecular Hessians are added in the new training procedure to utilize the second order derivative information.\n",
    "Below, we give more details for setting up these trainings.\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": "71edbe14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-01-11-21:14-xenophobic-aridatha\n",
      "2024-01-11-21:14-xenophobic-aridatha\n",
      "nums check 256 2048 10\n",
      "Loaded 40 training configurations from 'data/propane/propane_force_N4.xyz'\n",
      "Loaded 9610 validation configurations from 'data/propane/propane_v2_mesh_PBE_force.xyz'\n",
      "Total number of configurations: train=40, valid=9610, test=0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████| 40/40 [00:00<00:00, 23250.02it/s]\n",
      "100%|████████████████████████████████████| 9610/9610 [00:00<00:00, 24722.09it/s]\n",
      "0it [00:00, ?it/s]\n",
      "Pad train: 100%|███████████████████████████| 4356/4356 [00:16<00:00, 256.28it/s]\n",
      "Pad valid: 100%|███████████████████████████| 1000/1000 [00:03<00:00, 272.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of configurations: train=4356, valid=1000, test=0\n",
      "Compute the average number of neighbors: 10.000\n",
      "Do not normalize the radial basis (avg_r_min=None)\n",
      "Computed average Atomic Energies using least squares: {1: -6.220972980547945, 6: -2.3328648677054784}\n",
      "Create NequIP (JAX-MD version) with parameters {'use_sc': True, 'graph_net_steps': 2, 'hidden_irreps': '24x0e+24x0o+12x1e+12x1o+8x2e+8x2o', 'nonlinearities': {'e': 'swish', 'o': 'tanh'}, 'r_max': 5.0, 'avg_num_neighbors': 10.0, 'avg_r_min': None, 'num_species': 100, 'radial_basis': <function bessel_basis at 0x7fb719c41080>, 'radial_envelope': <function soft_envelope at 0x7fb719c40fe0>}\n",
      "num_params: 105836\n"
     ]
    }
   ],
   "source": [
    "# Load the config files\n",
    "with open('data/propane/mol-propane.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",
    "#print(len(EF_train_loader.graphs),len(EF_valid_loader.graphs))\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",
    "#print(len(ph_train_loader.graphs),len(ph_valid_loader.graphs))\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",
    "H_loss_fn = HessianLoss()\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_hessian(lambda *x: model_fn(w, *x), g)\n",
    ")\n"
   ]
  },
  {
   "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": "2057e4cf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimizer num_params: 317510\n",
      "Started training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_train:   0%|                                         | 0/4 [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%|███████████████████████████| 4/4 [00:05<00:00,  1.35s/it, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=44.3532, mae_e_per_atom=750.6 meV, mae_f=522.2 meV/Å, mae_s=0.8 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|██████████████████▉| 1066/1068 [00:15<00:00, 68.81it/s, n=9603]"
     ]
    },
    {
     "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%|███████████████████| 1068/1068 [00:20<00:00, 52.43it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=49.4872, mae_e_per_atom=754.2 meV, mae_f=550.4 meV/Å, mae_s=0.8 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 0:   2%|▏           | 10/500 [00:07<04:14,  1.92it/s, loss=3.499]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiled function `update_fn` for args:\n",
      "Outout: loss= 39.634\n",
      "Compilation time: 7.050s, cache size: 1\n",
      "Compiled function `update_fn` for args:\n",
      "Outout: loss= 35.488\n",
      "Compilation time: 0.015s, cache size: 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 0: 100%|███████████| 500/500 [00:13<00:00, 38.16it/s, loss=0.062]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.92it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.0677, mae_e_per_atom=0.3 meV, mae_f=24.0 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.53it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=1.0084, mae_e_per_atom=1.1 meV, mae_f=67.1 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 1: 100%|███████████| 500/500 [00:06<00:00, 83.25it/s, loss=0.034]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.81it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.0389, mae_e_per_atom=0.6 meV, mae_f=18.5 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.32it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=1.0545, mae_e_per_atom=1.4 meV, mae_f=68.2 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 2: 100%|███████████| 500/500 [00:06<00:00, 82.01it/s, loss=0.024]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.20it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.0249, mae_e_per_atom=0.7 meV, mae_f=15.1 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.15it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=1.0137, mae_e_per_atom=1.5 meV, mae_f=66.9 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 3: 100%|███████████| 500/500 [00:06<00:00, 82.24it/s, loss=0.034]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.15it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.0136, mae_e_per_atom=0.8 meV, mae_f=11.3 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.95it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=0.9484, mae_e_per_atom=1.3 meV, mae_f=63.6 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 4: 100%|███████████| 500/500 [00:06<00:00, 82.11it/s, loss=0.009]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.66it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.0079, mae_e_per_atom=0.7 meV, mae_f=8.6 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.82it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=0.8525, mae_e_per_atom=1.2 meV, mae_f=59.0 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 5: 100%|███████████| 500/500 [00:06<00:00, 82.29it/s, loss=0.005]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.06it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.0051, mae_e_per_atom=0.5 meV, mae_f=6.8 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.84it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=0.7972, mae_e_per_atom=1.3 meV, mae_f=55.4 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 6: 100%|███████████| 500/500 [00:06<00:00, 81.76it/s, loss=0.007]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.67it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0036, mae_e_per_atom=0.3 meV, mae_f=5.7 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.73it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=0.7833, mae_e_per_atom=1.2 meV, mae_f=53.6 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 7: 100%|███████████| 500/500 [00:06<00:00, 81.45it/s, loss=0.004]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.92it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0028, mae_e_per_atom=0.3 meV, mae_f=5.0 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.94it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=0.7885, mae_e_per_atom=1.0 meV, mae_f=52.6 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 8: 100%|███████████| 500/500 [00:06<00:00, 81.24it/s, loss=0.005]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.24it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0023, mae_e_per_atom=0.3 meV, mae_f=4.4 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.66it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=0.8120, mae_e_per_atom=1.0 meV, mae_f=52.5 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 9: 100%|███████████| 500/500 [00:06<00:00, 80.99it/s, loss=0.002]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.04it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_train: loss=0.0019, mae_e_per_atom=0.2 meV, mae_f=4.0 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.43it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_valid: loss=0.8345, mae_e_per_atom=1.0 meV, mae_f=52.4 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 10: 100%|██████████| 500/500 [00:05<00:00, 85.00it/s, loss=0.002]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.82it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_train: loss=0.0016, mae_e_per_atom=0.1 meV, mae_f=3.7 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.92it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_valid: loss=0.8648, mae_e_per_atom=1.1 meV, mae_f=52.8 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 11: 100%|██████████| 500/500 [00:05<00:00, 84.68it/s, loss=0.002]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 63.89it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_train: loss=0.0014, mae_e_per_atom=0.1 meV, mae_f=3.4 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 69.02it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_valid: loss=0.9025, mae_e_per_atom=1.0 meV, mae_f=53.4 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 12: 100%|██████████| 500/500 [00:05<00:00, 84.26it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.95it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_train: loss=0.0012, mae_e_per_atom=0.2 meV, mae_f=3.2 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 69.02it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_valid: loss=0.9284, mae_e_per_atom=1.2 meV, mae_f=53.8 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 13: 100%|██████████| 500/500 [00:05<00:00, 83.88it/s, loss=0.002]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.86it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_train: loss=0.0011, mae_e_per_atom=0.1 meV, mae_f=3.0 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 69.05it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_valid: loss=0.9599, mae_e_per_atom=1.1 meV, mae_f=54.3 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 14: 100%|██████████| 500/500 [00:05<00:00, 84.13it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.76it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_train: loss=0.0009, mae_e_per_atom=0.2 meV, mae_f=2.8 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.55it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 15: eval_valid: loss=0.9881, mae_e_per_atom=1.2 meV, mae_f=54.8 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 15: 100%|██████████| 500/500 [00:05<00:00, 83.52it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 68.34it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 16: eval_train: loss=0.0008, mae_e_per_atom=0.2 meV, mae_f=2.6 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.27it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 16: eval_valid: loss=1.0164, mae_e_per_atom=1.2 meV, mae_f=55.3 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 16: 100%|██████████| 500/500 [00:05<00:00, 83.70it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.20it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 17: eval_train: loss=0.0007, mae_e_per_atom=0.1 meV, mae_f=2.5 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 69.14it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 17: eval_valid: loss=1.0375, mae_e_per_atom=1.0 meV, mae_f=55.6 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 17: 100%|██████████| 500/500 [00:05<00:00, 84.00it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 63.45it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 18: eval_train: loss=0.0007, mae_e_per_atom=0.1 meV, mae_f=2.4 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.53it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 18: eval_valid: loss=1.0589, mae_e_per_atom=1.0 meV, mae_f=56.1 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 18: 100%|██████████| 500/500 [00:05<00:00, 83.69it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.97it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 19: eval_train: loss=0.0006, mae_e_per_atom=0.2 meV, mae_f=2.2 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.68it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 19: eval_valid: loss=1.0751, mae_e_per_atom=1.1 meV, mae_f=56.3 meV/Å, mae_s=1.4 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train interval 19: 100%|██████████| 500/500 [00:06<00:00, 82.83it/s, loss=0.001]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.39it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 20: eval_train: loss=0.0005, mae_e_per_atom=0.1 meV, mae_f=2.1 meV/Å, mae_s=1.2 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 68.39it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 20: eval_valid: loss=1.0961, mae_e_per_atom=1.2 meV, mae_f=56.6 meV/Å, mae_s=1.4 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": "5e6004b5-f604-4611-affb-b164732969bb",
   "metadata": {},
   "source": [
    "### Training with molecular 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": 3,
   "id": "0ddf289b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimizer num_params: 317510\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_train:   0%|                                         | 0/4 [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%|███████████████████████████| 4/4 [00:05<00:00,  1.37s/it, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=44.3532, mae_e_per_atom=750.565 meV, mae_f=522.178 meV/Å, mae_s=0.802 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid:  99%|██████████████████▉| 1061/1068 [00:15<00:00, 67.28it/s, n=9603]"
     ]
    },
    {
     "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%|███████████████████| 1068/1068 [00:20<00:00, 51.81it/s, n=9610]\n",
      "/home/shiangfang/miniconda3/envs/phonax/lib/python3.11/site-packages/numpy/core/fromnumeric.py:3464: RuntimeWarning: Mean of empty slice.\n",
      "  return _methods._mean(a, axis=axis, dtype=dtype,\n",
      "/home/shiangfang/miniconda3/envs/phonax/lib/python3.11/site-packages/numpy/core/_methods.py:192: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  ret = ret.dtype.type(ret / rcount)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=49.4872, mae_e_per_atom=754.195 meV, mae_f=550.353 meV/Å, mae_s=0.799 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 0: 100%|████████████████| 500/500 [00:13<00:00, 37.81it/s, loss=0.0621]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.69it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.0678, mae_e_per_atom=0.289 meV, mae_f=24.034 meV/Å, mae_s=1.226 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.49it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=1.0080, mae_e_per_atom=1.182 meV, mae_f=67.117 meV/Å, mae_s=1.404 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 1: 100%|█████████████████| 500/500 [00:06<00:00, 80.98it/s, loss=0.034]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.25it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.0389, mae_e_per_atom=0.588 meV, mae_f=18.516 meV/Å, mae_s=1.198 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.72it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=1.0539, mae_e_per_atom=1.356 meV, mae_f=68.166 meV/Å, mae_s=1.410 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 2: 100%|████████████████| 500/500 [00:06<00:00, 80.74it/s, loss=0.0243]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.40it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.0248, mae_e_per_atom=0.674 meV, mae_f=15.119 meV/Å, mae_s=1.193 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.30it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=1.0134, mae_e_per_atom=1.485 meV, mae_f=66.930 meV/Å, mae_s=1.417 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 3: 100%|████████████████| 500/500 [00:06<00:00, 80.75it/s, loss=0.0317]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.18it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.0135, mae_e_per_atom=0.788 meV, mae_f=11.331 meV/Å, mae_s=1.212 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.28it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=0.9512, mae_e_per_atom=1.365 meV, mae_f=63.675 meV/Å, mae_s=1.423 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 4: 100%|████████████████| 500/500 [00:06<00:00, 80.18it/s, loss=0.0104]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.87it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.0079, mae_e_per_atom=0.644 meV, mae_f=8.594 meV/Å, mae_s=1.202 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.27it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=0.8542, mae_e_per_atom=1.302 meV, mae_f=59.068 meV/Å, mae_s=1.428 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 5: 100%|███████████████| 500/500 [00:06<00:00, 80.01it/s, loss=0.00522]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.29it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.0051, mae_e_per_atom=0.463 meV, mae_f=6.796 meV/Å, mae_s=1.204 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.26it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=0.7981, mae_e_per_atom=1.298 meV, mae_f=55.446 meV/Å, mae_s=1.429 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 6: 100%|███████████████| 500/500 [00:06<00:00, 80.21it/s, loss=0.00986]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.86it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0037, mae_e_per_atom=0.341 meV, mae_f=5.736 meV/Å, mae_s=1.213 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 66.88it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=0.7852, mae_e_per_atom=1.171 meV, mae_f=53.699 meV/Å, mae_s=1.430 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 7: 100%|███████████████| 500/500 [00:06<00:00, 79.83it/s, loss=0.00312]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.27it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0028, mae_e_per_atom=0.269 meV, mae_f=5.025 meV/Å, mae_s=1.214 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.02it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=0.7903, mae_e_per_atom=1.104 meV, mae_f=52.668 meV/Å, mae_s=1.431 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 8: 100%|███████████████| 500/500 [00:06<00:00, 80.20it/s, loss=0.00389]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 62.86it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0022, mae_e_per_atom=0.223 meV, mae_f=4.422 meV/Å, mae_s=1.220 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.12it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=0.8143, mae_e_per_atom=1.123 meV, mae_f=52.510 meV/Å, mae_s=1.431 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 9: 100%|███████████████| 500/500 [00:06<00:00, 79.45it/s, loss=0.00161]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.30it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_train: loss=0.0019, mae_e_per_atom=0.175 meV, mae_f=4.029 meV/Å, mae_s=1.208 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:15<00:00, 67.02it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 0: eval_valid: loss=0.8359, mae_e_per_atom=1.090 meV, mae_f=52.501 meV/Å, mae_s=1.431 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:06<00:00, 19.46it/s, n_graphs=4356\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 0: hessian train hessian MAE 0.245 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|██| 28/28 [00:00<00:00, 81.14it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 0: hessian valid hessian MAE 0.237 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 0: 100%|████████████████| 500/500 [00:28<00:00, 17.69it/s, loss=0.0461]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.67it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_train: loss=0.0178, mae_e_per_atom=0.453 meV, mae_f=11.894 meV/Å, mae_s=1.216 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.63it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 1: eval_valid: loss=0.1337, mae_e_per_atom=1.211 meV, mae_f=24.011 meV/Å, mae_s=1.455 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 145.00it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 1: hessian train hessian MAE 0.089 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 149.45it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 1: hessian valid hessian MAE 0.082 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 1: 100%|████████████████| 500/500 [00:10<00:00, 48.89it/s, loss=0.0856]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 66.69it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_train: loss=0.0120, mae_e_per_atom=0.486 meV, mae_f=10.137 meV/Å, mae_s=1.217 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 66.41it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 2: eval_valid: loss=0.0857, mae_e_per_atom=1.127 meV, mae_f=19.744 meV/Å, mae_s=1.464 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 148.58it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 2: hessian train hessian MAE 0.072 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 150.35it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 2: hessian valid hessian MAE 0.067 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 2: 100%|████████████████| 500/500 [00:10<00:00, 49.15it/s, loss=0.0385]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.35it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_train: loss=0.0095, mae_e_per_atom=0.827 meV, mae_f=9.045 meV/Å, mae_s=1.212 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 66.54it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 3: eval_valid: loss=0.0834, mae_e_per_atom=0.944 meV, mae_f=19.136 meV/Å, mae_s=1.463 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 145.46it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 3: hessian train hessian MAE 0.065 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 148.37it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 3: hessian valid hessian MAE 0.058 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 3: 100%|████████████████| 500/500 [00:10<00:00, 48.93it/s, loss=0.0283]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.87it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_train: loss=0.0078, mae_e_per_atom=0.376 meV, mae_f=8.265 meV/Å, mae_s=1.208 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.54it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 4: eval_valid: loss=0.0695, mae_e_per_atom=0.869 meV, mae_f=17.837 meV/Å, mae_s=1.465 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 147.54it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 4: hessian train hessian MAE 0.059 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 150.55it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 4: hessian valid hessian MAE 0.051 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 4: 100%|████████████████| 500/500 [00:10<00:00, 48.81it/s, loss=0.0441]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.15it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_train: loss=0.0063, mae_e_per_atom=0.323 meV, mae_f=7.471 meV/Å, mae_s=1.219 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 66.03it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 5: eval_valid: loss=0.0584, mae_e_per_atom=0.692 meV, mae_f=16.446 meV/Å, mae_s=1.469 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 147.20it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 5: hessian train hessian MAE 0.054 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 151.47it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 5: hessian valid hessian MAE 0.048 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 5: 100%|████████████████| 500/500 [00:10<00:00, 48.81it/s, loss=0.0245]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 62.35it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_train: loss=0.0056, mae_e_per_atom=0.414 meV, mae_f=7.032 meV/Å, mae_s=1.225 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.97it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 6: eval_valid: loss=0.0581, mae_e_per_atom=0.796 meV, mae_f=16.354 meV/Å, mae_s=1.469 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 146.07it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 6: hessian train hessian MAE 0.050 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 150.46it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 6: hessian valid hessian MAE 0.046 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 6: 100%|████████████████| 500/500 [00:10<00:00, 48.86it/s, loss=0.0266]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 65.19it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_train: loss=0.0051, mae_e_per_atom=0.269 meV, mae_f=6.784 meV/Å, mae_s=1.210 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.84it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 7: eval_valid: loss=0.0596, mae_e_per_atom=0.561 meV, mae_f=16.195 meV/Å, mae_s=1.469 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 148.64it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 7: hessian train hessian MAE 0.049 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 150.32it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 7: hessian valid hessian MAE 0.043 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 7: 100%|████████████████| 500/500 [00:10<00:00, 48.79it/s, loss=0.0544]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.53it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_train: loss=0.0046, mae_e_per_atom=0.228 meV, mae_f=6.429 meV/Å, mae_s=1.226 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.63it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 8: eval_valid: loss=0.0598, mae_e_per_atom=0.524 meV, mae_f=16.076 meV/Å, mae_s=1.470 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 146.68it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 8: hessian train hessian MAE 0.048 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 153.00it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 8: hessian valid hessian MAE 0.043 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 8: 100%|████████████████| 500/500 [00:10<00:00, 48.76it/s, loss=0.0222]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 63.93it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_train: loss=0.0044, mae_e_per_atom=0.443 meV, mae_f=6.279 meV/Å, mae_s=1.209 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.58it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 9: eval_valid: loss=0.0584, mae_e_per_atom=0.799 meV, mae_f=15.824 meV/Å, mae_s=1.470 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 148.50it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 9: hessian train hessian MAE 0.046 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 150.27it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 9: hessian valid hessian MAE 0.040 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 9: 100%|██████████████████| 500/500 [00:10<00:00, 48.69it/s, loss=0.03]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 64.58it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_train: loss=0.0042, mae_e_per_atom=0.223 meV, mae_f=6.129 meV/Å, mae_s=1.206 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.63it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 10: eval_valid: loss=0.0584, mae_e_per_atom=0.522 meV, mae_f=15.539 meV/Å, mae_s=1.470 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 147.54it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 10: hessian train hessian MAE 0.045 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 149.69it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 10: hessian valid hessian MAE 0.040 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 10: 100%|███████████████| 500/500 [00:10<00:00, 48.75it/s, loss=0.0292]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 63.11it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_train: loss=0.0042, mae_e_per_atom=0.192 meV, mae_f=6.115 meV/Å, mae_s=1.208 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.46it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 11: eval_valid: loss=0.0540, mae_e_per_atom=0.486 meV, mae_f=15.138 meV/Å, mae_s=1.471 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 144.92it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 11: hessian train hessian MAE 0.043 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 148.84it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 11: hessian valid hessian MAE 0.038 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 11: 100%|███████████████| 500/500 [00:10<00:00, 48.64it/s, loss=0.0296]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 67.56it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_train: loss=0.0036, mae_e_per_atom=0.229 meV, mae_f=5.592 meV/Å, mae_s=1.230 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.42it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 12: eval_valid: loss=0.0534, mae_e_per_atom=0.479 meV, mae_f=14.925 meV/Å, mae_s=1.470 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 145.08it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 12: hessian train hessian MAE 0.042 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 147.38it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 12: hessian valid hessian MAE 0.036 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 12: 100%|███████████████| 500/500 [00:10<00:00, 48.46it/s, loss=0.0399]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 62.58it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_train: loss=0.0037, mae_e_per_atom=0.294 meV, mae_f=5.690 meV/Å, mae_s=1.205 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.54it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 13: eval_valid: loss=0.0567, mae_e_per_atom=0.472 meV, mae_f=15.105 meV/Å, mae_s=1.471 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 148.11it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 13: hessian train hessian MAE 0.041 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 147.91it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 13: hessian valid hessian MAE 0.037 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 13: 100%|███████████████| 500/500 [00:10<00:00, 48.38it/s, loss=0.0184]\n",
      "eval_train: 100%|███████████████████████████| 4/4 [00:00<00:00, 63.77it/s, n=36]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_train: loss=0.0036, mae_e_per_atom=0.220 meV, mae_f=5.612 meV/Å, mae_s=1.229 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval_valid: 100%|███████████████████| 1068/1068 [00:16<00:00, 65.34it/s, n=9610]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Interval 14: eval_valid: loss=0.0514, mae_e_per_atom=0.462 meV, mae_f=14.522 meV/Å, mae_s=1.472 meV/Å³\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian train: 100%|█| 121/121 [00:00<00:00, 145.88it/s, n_graphs=435\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 14: hessian train hessian MAE 0.040 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating hessian valid: 100%|█| 28/28 [00:00<00:00, 149.89it/s, n_graphs=1000]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interval 14: hessian valid hessian MAE 0.037 eV/A2 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Interval 14: 100%|███████████████| 500/500 [00:10<00:00, 48.39it/s, loss=0.0194]\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 = config[\"training\"][\"max_num_intervals\"],\n",
    "    steps_per_interval = config[\"training\"][\"steps_per_interval\"],\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=config['training']['warmup_num_intervals'],\n",
    "    phtrain_num_intervals=config['training']['phtrain_num_intervals'],\n",
    ")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04c7d71",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "89ba6ae6",
   "metadata": {},
   "source": [
    "# Energy model fine-tuning with Hessian spectrum data\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f4408c-6256-41c5-b2b8-c32598be0090",
   "metadata": {},
   "source": [
    "Here in this section,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3825e5b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create NequIP (JAX-MD version) with parameters {'use_sc': True, 'graph_net_steps': 2, 'hidden_irreps': '96x0e+ 96x0o  + 64x1e +64x1o +64x2e +64x2o', 'nonlinearities': {'e': 'swish', 'o': 'tanh'}, 'r_max': 5.0, 'avg_num_neighbors': 53.2808396682792, 'avg_r_min': None, 'num_species': 100, 'radial_basis': <function bessel_basis at 0x7f77bd701440>, 'radial_envelope': <function soft_envelope at 0x7f77bd7013a0>}\n"
     ]
    }
   ],
   "source": [
    "model_fn, params, num_message_passing, r_max = NequIP_JAXMD_molecule_model(os.path.join(os.getcwd(), 'trained-models'))\n",
    "\n",
    "train_crystal = ase.io.read('data/molecules/benzene.vasp')\n",
    "train_crystal.pbc=False\n",
    "train_graph = atoms_to_ext_graph(train_crystal, r_max, num_message_passing=1)\n",
    "train_graph = jax.tree_util.tree_map(to_f32, train_graph)\n",
    "\n",
    "test_crystal = ase.io.read('data/molecules/toluene.vasp')\n",
    "test_crystal.pbc=False\n",
    "test_graph = atoms_to_ext_graph(test_crystal, r_max, num_message_passing=1)\n",
    "test_graph = jax.tree_util.tree_map(to_f32, test_graph)\n",
    "\n",
    "train_masses = ase.data.atomic_masses[train_graph.nodes.species][train_graph.nodes.mask_primitive]\n",
    "test_masses = ase.data.atomic_masses[test_graph.nodes.species][test_graph.nodes.mask_primitive]\n",
    "\n",
    "def read_gamma_modes(fname):\n",
    "    all_modes = []\n",
    "    with open(fname,'r') as f:\n",
    "        for line in f:\n",
    "            if \"Mode\" in line:\n",
    "                all_modes.append(float(line.strip().split(',')[-1].strip().split()[0]))\n",
    "    all_modes = jnp.array(all_modes)\n",
    "    return all_modes\n",
    "benzene_gamma = read_gamma_modes('data/molecules/benzene_gamma')\n",
    "toluene_gamma = read_gamma_modes('data/molecules/toluene_gamma')\n",
    "\n",
    "\n",
    "def loss_fn_train(spectra_in):\n",
    "    return jnp.mean(jnp.abs(spectra_in-benzene_gamma))\n",
    "def loss_fn_test(spectra_in):\n",
    "    return jnp.mean(jnp.abs(spectra_in-toluene_gamma))\n",
    "def mae_spectra(pred_spectra,target_spectra):\n",
    "    return jnp.mean(jnp.abs(pred_spectra-target_spectra))\n",
    "\n",
    "gradient_transform = optax.adam(5E-5)\n",
    "\n",
    "optimizer_state = gradient_transform.init(params)\n",
    "\n",
    "spectra_predictor = jax.jit(\n",
    "    lambda w: predict_gamma_spectra_filter(w,model_fn,train_graph,train_masses)  #predict_gamma_spectra(lambda *x: model_fn(w, *x), test_graph)\n",
    ")\n",
    "\n",
    "valid_modes_original = predict_gamma_spectra_filter(params,model_fn,test_graph,test_masses)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ec9db688",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Benzene spectra MAE:  88.0602 Toluene spectra MAE:  74.18877\n",
      "Benzene spectra MAE:  85.86992 Toluene spectra MAE:  72.87502\n",
      "Benzene spectra MAE:  83.41573 Toluene spectra MAE:  70.76455\n",
      "Benzene spectra MAE:  80.89786 Toluene spectra MAE:  68.3744\n",
      "Benzene spectra MAE:  78.783806 Toluene spectra MAE:  65.71441\n",
      "Benzene spectra MAE:  76.697945 Toluene spectra MAE:  63.42712\n",
      "Benzene spectra MAE:  74.64182 Toluene spectra MAE:  61.302\n",
      "Benzene spectra MAE:  72.42226 Toluene spectra MAE:  59.635105\n",
      "Benzene spectra MAE:  70.277 Toluene spectra MAE:  58.176296\n",
      "Benzene spectra MAE:  68.560936 Toluene spectra MAE:  56.826687\n",
      "Benzene spectra MAE:  67.40696 Toluene spectra MAE:  55.650597\n",
      "Benzene spectra MAE:  66.224625 Toluene spectra MAE:  54.374023\n",
      "Benzene spectra MAE:  65.042625 Toluene spectra MAE:  53.331573\n",
      "Benzene spectra MAE:  64.294586 Toluene spectra MAE:  52.429073\n",
      "Benzene spectra MAE:  63.428223 Toluene spectra MAE:  51.64859\n",
      "Benzene spectra MAE:  62.439117 Toluene spectra MAE:  50.93556\n",
      "Benzene spectra MAE:  61.380157 Toluene spectra MAE:  50.285206\n",
      "Benzene spectra MAE:  60.242756 Toluene spectra MAE:  49.714634\n",
      "Benzene spectra MAE:  59.00773 Toluene spectra MAE:  49.298084\n",
      "Benzene spectra MAE:  57.67382 Toluene spectra MAE:  48.98265\n",
      "Benzene spectra MAE:  56.34242 Toluene spectra MAE:  48.78248\n",
      "Benzene spectra MAE:  54.92903 Toluene spectra MAE:  48.201534\n",
      "Benzene spectra MAE:  54.11241 Toluene spectra MAE:  47.244576\n",
      "Benzene spectra MAE:  53.31077 Toluene spectra MAE:  46.41502\n",
      "Benzene spectra MAE:  52.535686 Toluene spectra MAE:  45.65375\n",
      "Benzene spectra MAE:  51.76509 Toluene spectra MAE:  44.780025\n",
      "Benzene spectra MAE:  51.059464 Toluene spectra MAE:  43.89693\n",
      "Benzene spectra MAE:  50.526943 Toluene spectra MAE:  43.039215\n",
      "Benzene spectra MAE:  49.868187 Toluene spectra MAE:  42.671513\n",
      "Benzene spectra MAE:  49.047348 Toluene spectra MAE:  42.705963\n",
      "Benzene spectra MAE:  48.388454 Toluene spectra MAE:  42.822678\n",
      "Benzene spectra MAE:  47.755253 Toluene spectra MAE:  42.826183\n",
      "Benzene spectra MAE:  47.11308 Toluene spectra MAE:  42.75355\n",
      "Benzene spectra MAE:  46.578133 Toluene spectra MAE:  42.536575\n",
      "Benzene spectra MAE:  45.97153 Toluene spectra MAE:  42.20912\n",
      "Benzene spectra MAE:  45.3208 Toluene spectra MAE:  41.718925\n",
      "Benzene spectra MAE:  44.62356 Toluene spectra MAE:  41.27737\n",
      "Benzene spectra MAE:  43.90184 Toluene spectra MAE:  40.77172\n",
      "Benzene spectra MAE:  43.357983 Toluene spectra MAE:  40.28957\n",
      "Benzene spectra MAE:  42.801712 Toluene spectra MAE:  39.884377\n",
      "Benzene spectra MAE:  42.176624 Toluene spectra MAE:  39.593513\n",
      "Benzene spectra MAE:  41.4599 Toluene spectra MAE:  39.470165\n",
      "Benzene spectra MAE:  40.977062 Toluene spectra MAE:  39.478485\n",
      "Benzene spectra MAE:  40.220894 Toluene spectra MAE:  38.993576\n",
      "Benzene spectra MAE:  39.593704 Toluene spectra MAE:  38.46907\n",
      "Benzene spectra MAE:  39.111748 Toluene spectra MAE:  37.988007\n",
      "Benzene spectra MAE:  38.576054 Toluene spectra MAE:  37.638927\n",
      "Benzene spectra MAE:  38.0359 Toluene spectra MAE:  37.373554\n",
      "Benzene spectra MAE:  37.517166 Toluene spectra MAE:  37.04998\n",
      "Benzene spectra MAE:  36.96138 Toluene spectra MAE:  36.91375\n",
      "Benzene spectra MAE:  36.416367 Toluene spectra MAE:  36.88364\n",
      "Benzene spectra MAE:  35.918514 Toluene spectra MAE:  36.79705\n",
      "Benzene spectra MAE:  35.373184 Toluene spectra MAE:  36.588226\n",
      "Benzene spectra MAE:  34.751953 Toluene spectra MAE:  36.306423\n",
      "Benzene spectra MAE:  34.543457 Toluene spectra MAE:  35.959435\n",
      "Benzene spectra MAE:  34.365246 Toluene spectra MAE:  35.641357\n",
      "Benzene spectra MAE:  34.15272 Toluene spectra MAE:  35.43132\n",
      "Benzene spectra MAE:  33.86564 Toluene spectra MAE:  35.3514\n",
      "Benzene spectra MAE:  33.594784 Toluene spectra MAE:  35.32126\n",
      "Benzene spectra MAE:  33.294716 Toluene spectra MAE:  35.289627\n",
      "Benzene spectra MAE:  33.067604 Toluene spectra MAE:  35.245087\n",
      "Benzene spectra MAE:  32.840755 Toluene spectra MAE:  35.162014\n",
      "Benzene spectra MAE:  32.60466 Toluene spectra MAE:  35.09957\n",
      "Benzene spectra MAE:  32.348286 Toluene spectra MAE:  35.015137\n",
      "Benzene spectra MAE:  32.086887 Toluene spectra MAE:  34.911873\n",
      "Benzene spectra MAE:  31.7987 Toluene spectra MAE:  34.755207\n",
      "Benzene spectra MAE:  31.650072 Toluene spectra MAE:  34.55738\n",
      "Benzene spectra MAE:  31.460268 Toluene spectra MAE:  34.4682\n",
      "Benzene spectra MAE:  31.345434 Toluene spectra MAE:  34.859818\n",
      "Benzene spectra MAE:  31.50645 Toluene spectra MAE:  35.06376\n",
      "Benzene spectra MAE:  30.83881 Toluene spectra MAE:  34.63838\n",
      "Benzene spectra MAE:  31.14459 Toluene spectra MAE:  34.044376\n",
      "Benzene spectra MAE:  31.190289 Toluene spectra MAE:  33.90618\n",
      "Benzene spectra MAE:  30.74882 Toluene spectra MAE:  33.952316\n",
      "Benzene spectra MAE:  30.071407 Toluene spectra MAE:  34.368855\n",
      "Benzene spectra MAE:  30.691984 Toluene spectra MAE:  34.874187\n",
      "Benzene spectra MAE:  30.319986 Toluene spectra MAE:  34.73557\n",
      "Benzene spectra MAE:  29.873255 Toluene spectra MAE:  34.125744\n",
      "Benzene spectra MAE:  30.04958 Toluene spectra MAE:  33.801178\n",
      "Benzene spectra MAE:  29.73535 Toluene spectra MAE:  33.850185\n",
      "Benzene spectra MAE:  29.213913 Toluene spectra MAE:  34.117146\n",
      "Benzene spectra MAE:  29.44047 Toluene spectra MAE:  34.447258\n",
      "Benzene spectra MAE:  28.987076 Toluene spectra MAE:  34.212685\n",
      "Benzene spectra MAE:  29.076721 Toluene spectra MAE:  33.66917\n",
      "Benzene spectra MAE:  29.371578 Toluene spectra MAE:  33.37683\n",
      "Benzene spectra MAE:  29.265734 Toluene spectra MAE:  33.294617\n",
      "Benzene spectra MAE:  28.833128 Toluene spectra MAE:  33.330544\n",
      "Benzene spectra MAE:  28.275734 Toluene spectra MAE:  33.716316\n",
      "Benzene spectra MAE:  28.750504 Toluene spectra MAE:  34.162663\n",
      "Benzene spectra MAE:  28.589674 Toluene spectra MAE:  34.027157\n",
      "Benzene spectra MAE:  27.89568 Toluene spectra MAE:  33.541416\n",
      "Benzene spectra MAE:  28.184675 Toluene spectra MAE:  32.92458\n",
      "Benzene spectra MAE:  28.285603 Toluene spectra MAE:  32.767822\n",
      "Benzene spectra MAE:  27.762516 Toluene spectra MAE:  32.903427\n",
      "Benzene spectra MAE:  27.420477 Toluene spectra MAE:  33.378418\n",
      "Benzene spectra MAE:  27.914518 Toluene spectra MAE:  33.849636\n",
      "Benzene spectra MAE:  27.781593 Toluene spectra MAE:  33.85552\n",
      "Benzene spectra MAE:  27.137058 Toluene spectra MAE:  33.391182\n",
      "Benzene spectra MAE:  27.008383 Toluene spectra MAE:  33.029728\n",
      "Benzene spectra MAE:  27.040482 Toluene spectra MAE:  32.747658\n",
      "Benzene spectra MAE:  26.850199 Toluene spectra MAE:  32.707413\n",
      "Benzene spectra MAE:  26.42115 Toluene spectra MAE:  32.963615\n",
      "Benzene spectra MAE:  26.514862 Toluene spectra MAE:  33.36097\n",
      "Benzene spectra MAE:  26.490715 Toluene spectra MAE:  33.365154\n",
      "Benzene spectra MAE:  25.874168 Toluene spectra MAE:  32.95974\n",
      "Benzene spectra MAE:  25.902884 Toluene spectra MAE:  32.763943\n",
      "Benzene spectra MAE:  25.758955 Toluene spectra MAE:  32.729824\n",
      "Benzene spectra MAE:  25.508781 Toluene spectra MAE:  32.885303\n",
      "Benzene spectra MAE:  25.516783 Toluene spectra MAE:  33.07348\n",
      "Benzene spectra MAE:  25.240335 Toluene spectra MAE:  32.95123\n",
      "Benzene spectra MAE:  25.14027 Toluene spectra MAE:  32.65084\n",
      "Benzene spectra MAE:  25.143724 Toluene spectra MAE:  32.529636\n",
      "Benzene spectra MAE:  24.910168 Toluene spectra MAE:  32.654675\n",
      "Benzene spectra MAE:  24.689533 Toluene spectra MAE:  32.875557\n",
      "Benzene spectra MAE:  24.993746 Toluene spectra MAE:  33.17764\n",
      "Benzene spectra MAE:  24.56168 Toluene spectra MAE:  33.025314\n",
      "Benzene spectra MAE:  24.46237 Toluene spectra MAE:  32.622448\n",
      "Benzene spectra MAE:  24.590292 Toluene spectra MAE:  32.485966\n",
      "Benzene spectra MAE:  24.19494 Toluene spectra MAE:  32.66921\n",
      "Benzene spectra MAE:  24.052763 Toluene spectra MAE:  33.091988\n",
      "Benzene spectra MAE:  24.051283 Toluene spectra MAE:  33.181564\n",
      "Benzene spectra MAE:  23.76751 Toluene spectra MAE:  32.82424\n",
      "Benzene spectra MAE:  23.963024 Toluene spectra MAE:  32.550407\n",
      "Benzene spectra MAE:  23.717115 Toluene spectra MAE:  32.655903\n",
      "Benzene spectra MAE:  23.327044 Toluene spectra MAE:  33.077213\n",
      "Benzene spectra MAE:  23.567688 Toluene spectra MAE:  33.292645\n",
      "Benzene spectra MAE:  23.021465 Toluene spectra MAE:  33.022728\n",
      "Benzene spectra MAE:  22.969887 Toluene spectra MAE:  32.92324\n",
      "Benzene spectra MAE:  22.803402 Toluene spectra MAE:  32.976295\n",
      "Benzene spectra MAE:  22.852484 Toluene spectra MAE:  33.190998\n",
      "Benzene spectra MAE:  22.538872 Toluene spectra MAE:  32.95671\n",
      "Benzene spectra MAE:  22.580694 Toluene spectra MAE:  32.834778\n",
      "Benzene spectra MAE:  22.282635 Toluene spectra MAE:  33.078594\n",
      "Benzene spectra MAE:  22.17306 Toluene spectra MAE:  33.099617\n",
      "Benzene spectra MAE:  22.164658 Toluene spectra MAE:  32.94915\n",
      "Benzene spectra MAE:  22.027166 Toluene spectra MAE:  33.103825\n",
      "Benzene spectra MAE:  21.951939 Toluene spectra MAE:  32.923347\n",
      "Benzene spectra MAE:  21.839136 Toluene spectra MAE:  32.99021\n",
      "Benzene spectra MAE:  21.986462 Toluene spectra MAE:  33.22127\n",
      "Benzene spectra MAE:  21.641745 Toluene spectra MAE:  32.971745\n",
      "Benzene spectra MAE:  21.833471 Toluene spectra MAE:  32.71587\n",
      "Benzene spectra MAE:  21.729244 Toluene spectra MAE:  32.71679\n",
      "Benzene spectra MAE:  21.539633 Toluene spectra MAE:  32.992558\n",
      "Benzene spectra MAE:  21.28754 Toluene spectra MAE:  32.94238\n",
      "Benzene spectra MAE:  21.510956 Toluene spectra MAE:  32.755543\n",
      "Benzene spectra MAE:  21.454062 Toluene spectra MAE:  32.774925\n",
      "Benzene spectra MAE:  21.248854 Toluene spectra MAE:  32.92399\n",
      "Benzene spectra MAE:  21.137321 Toluene spectra MAE:  32.731754\n",
      "Benzene spectra MAE:  21.099957 Toluene spectra MAE:  32.795452\n",
      "Benzene spectra MAE:  20.971474 Toluene spectra MAE:  32.7257\n",
      "Benzene spectra MAE:  21.173023 Toluene spectra MAE:  32.88916\n",
      "Benzene spectra MAE:  20.832087 Toluene spectra MAE:  32.644386\n",
      "Benzene spectra MAE:  20.753923 Toluene spectra MAE:  32.62411\n",
      "Benzene spectra MAE:  20.822262 Toluene spectra MAE:  32.733482\n",
      "Benzene spectra MAE:  20.759804 Toluene spectra MAE:  32.441753\n",
      "Benzene spectra MAE:  20.638983 Toluene spectra MAE:  32.39528\n",
      "Benzene spectra MAE:  20.607626 Toluene spectra MAE:  32.59787\n",
      "Benzene spectra MAE:  20.378778 Toluene spectra MAE:  32.45335\n",
      "Benzene spectra MAE:  20.339884 Toluene spectra MAE:  32.335835\n",
      "Benzene spectra MAE:  20.332407 Toluene spectra MAE:  32.472614\n",
      "Benzene spectra MAE:  20.239595 Toluene spectra MAE:  32.240837\n",
      "Benzene spectra MAE:  20.226246 Toluene spectra MAE:  32.24198\n",
      "Benzene spectra MAE:  20.328602 Toluene spectra MAE:  32.427715\n",
      "Benzene spectra MAE:  20.085533 Toluene spectra MAE:  32.191895\n",
      "Benzene spectra MAE:  19.986507 Toluene spectra MAE:  32.227684\n",
      "Benzene spectra MAE:  19.999424 Toluene spectra MAE:  32.112762\n",
      "Benzene spectra MAE:  20.19882 Toluene spectra MAE:  32.37067\n",
      "Benzene spectra MAE:  19.853447 Toluene spectra MAE:  32.180584\n",
      "Benzene spectra MAE:  20.228842 Toluene spectra MAE:  31.81662\n",
      "Benzene spectra MAE:  20.288412 Toluene spectra MAE:  31.74039\n",
      "Benzene spectra MAE:  20.079302 Toluene spectra MAE:  31.879717\n",
      "Benzene spectra MAE:  19.731169 Toluene spectra MAE:  32.159687\n",
      "Benzene spectra MAE:  19.758245 Toluene spectra MAE:  32.21581\n",
      "Benzene spectra MAE:  19.657818 Toluene spectra MAE:  31.934578\n",
      "Benzene spectra MAE:  19.97201 Toluene spectra MAE:  31.644135\n",
      "Benzene spectra MAE:  19.745089 Toluene spectra MAE:  31.70183\n",
      "Benzene spectra MAE:  19.59485 Toluene spectra MAE:  32.055466\n",
      "Benzene spectra MAE:  19.751677 Toluene spectra MAE:  32.21109\n",
      "Benzene spectra MAE:  19.352377 Toluene spectra MAE:  31.938822\n",
      "Benzene spectra MAE:  19.385544 Toluene spectra MAE:  31.747982\n",
      "Benzene spectra MAE:  19.181694 Toluene spectra MAE:  31.84684\n",
      "Benzene spectra MAE:  19.245802 Toluene spectra MAE:  31.949173\n",
      "Benzene spectra MAE:  19.164923 Toluene spectra MAE:  31.734468\n",
      "Benzene spectra MAE:  19.07091 Toluene spectra MAE:  31.771381\n",
      "Benzene spectra MAE:  19.007778 Toluene spectra MAE:  31.745762\n",
      "Benzene spectra MAE:  18.952854 Toluene spectra MAE:  31.819159\n",
      "Benzene spectra MAE:  19.003359 Toluene spectra MAE:  31.62601\n",
      "Benzene spectra MAE:  18.953907 Toluene spectra MAE:  31.634296\n",
      "Benzene spectra MAE:  18.79339 Toluene spectra MAE:  31.725079\n",
      "Benzene spectra MAE:  18.771622 Toluene spectra MAE:  31.551928\n",
      "Benzene spectra MAE:  18.74063 Toluene spectra MAE:  31.491978\n",
      "Benzene spectra MAE:  18.645863 Toluene spectra MAE:  31.525913\n",
      "Benzene spectra MAE:  18.594358 Toluene spectra MAE:  31.592505\n",
      "Benzene spectra MAE:  18.575079 Toluene spectra MAE:  31.547064\n",
      "Benzene spectra MAE:  18.556665 Toluene spectra MAE:  31.585669\n",
      "Benzene spectra MAE:  18.42619 Toluene spectra MAE:  31.428925\n",
      "Benzene spectra MAE:  18.707493 Toluene spectra MAE:  31.139349\n",
      "Benzene spectra MAE:  18.414614 Toluene spectra MAE:  31.268448\n",
      "Benzene spectra MAE:  18.592087 Toluene spectra MAE:  31.625345\n",
      "Benzene spectra MAE:  18.381907 Toluene spectra MAE:  31.498476\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def update_fn(\n",
    "    params, optimizer_state\n",
    "):\n",
    "    loss, grad = jax.value_and_grad(\n",
    "        lambda params: loss_fn_train(spectra_predictor(params))\n",
    "    )(params)\n",
    "\n",
    "    updates, optimizer_state = gradient_transform.update(\n",
    "        grad, optimizer_state, params\n",
    "    )\n",
    "    params = optax.apply_updates(params, updates)\n",
    "    return loss, params, optimizer_state\n",
    "\n",
    "spectra_mae = []\n",
    "\n",
    "for ind in range(200):\n",
    "    loss, params, optimizer_state = update_fn(params, optimizer_state)\n",
    "    train_modes = predict_gamma_spectra_filter(params,model_fn,train_graph,train_masses)\n",
    "    valid_modes = predict_gamma_spectra_filter(params,model_fn,test_graph,test_masses)\n",
    "    train_mae = mae_spectra(train_modes,benzene_gamma)\n",
    "    valid_mae = mae_spectra(valid_modes,toluene_gamma)\n",
    "\n",
    "    print(\n",
    "        'Benzene spectra MAE: ', \n",
    "        train_mae,\n",
    "        'Toluene spectra MAE: ',\n",
    "        valid_mae,\n",
    "    )\n",
    "    spectra_mae.append([train_mae, valid_mae])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "401f2bb6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec406737",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "174405d0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d56f11b-3d7d-46fd-b20a-b257fb8dacf7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3de72966-b35c-4263-92a6-8f8e145b3b8f",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}