{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build 2+3-body Mapped GP\n",
    "\n",
    "After the on-the-fly training is complete, we can play with the force field we obtained. \n",
    "We are going to do the following things:\n",
    "\n",
    "1. Parse the on-the-fly training trajectory to collect training data\n",
    "2. Reconstruct the GP model from the training trajectory\n",
    "3. Build up Mapped GP (MGP) for accelerated force field, and save coefficient file for LAMMPS\n",
    "4. Use LAMMPS to run fast simulation using MGP pair style\n",
    "\n",
    "## Parse OTF log file\n",
    "\n",
    "After the on-the-fly training is complete, we have a log file and can use the `otf_parser` module to parse the trajectory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from flare import otf_parser\n",
    "\n",
    "logdir = '../../../tests/test_files'\n",
    "file_name = f'{logdir}/AgI_snippet.out'\n",
    "hyp_no = 2 # use the hyperparameters from the 2nd training step\n",
    "otf_object = otf_parser.OtfAnalysis(file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct GP model from log file\n",
    "\n",
    "We can reconstruct GP model from the parsed log file (the on-the-fly training trajectory). Here we build up the GP model with 2+3 body kernel from the on-the-fly log file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp_model = otf_object.make_gp(hyp_no=hyp_no)\n",
    "gp_model.parallel = True\n",
    "gp_model.hyp_labels = ['sig2', 'ls2', 'sig3', 'ls3', 'noise']\n",
    "\n",
    "# write model to a binary file\n",
    "gp_model.write_model('AgI.gp', format='json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step `write_model` is to write this GP model into a binary file, \n",
    "so next time we can directly load the model from the pickle file as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flare.gp import GaussianProcess\n",
    "\n",
    "gp_model = GaussianProcess.from_file('AgI.gp.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Map the GP force field & Dump LAMMPS coefficient file\n",
    "\n",
    "To use the trained force field with accelerated version MGP, or in LAMMPS, we need to build MGP from GP model. \n",
    "Since 2-body and 3-body are both included, we need to set up the number of grid points for 2-body and 3-body in `grid_params`.\n",
    "We build up energy mapping, thus set `map_force=False`.\n",
    "See [MGP tutorial](https://flare.readthedocs.io/en/latest/tutorials/mgp.html) for more explanation of the MGP settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/yuxie/opt/anaconda3/lib/python3.8/site-packages/flare/mgp/mapxb.py:519: UserWarning: The minimal distance in training data is lower than the current lower bound, will reset lower bound to 2.129780094032889\n",
      "  warnings.warn(\n",
      "/Users/yuxie/opt/anaconda3/lib/python3.8/site-packages/flare/mgp/mapxb.py:519: UserWarning: The minimal distance in training data is lower than the current lower bound, will reset lower bound to 2.129780094032889\n",
      "  warnings.warn(\n",
      "/Users/yuxie/opt/anaconda3/lib/python3.8/site-packages/flare/mgp/mapxb.py:519: UserWarning: The minimal distance in training data is lower than the current lower bound, will reset lower bound to 2.129780094032889\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from flare.mgp import MappedGaussianProcess\n",
    "\n",
    "grid_params = {'twobody':   {'grid_num': [64]}, \n",
    "               'threebody': {'grid_num': [20, 20, 20]}}\n",
    "\n",
    "data = gp_model.training_statistics\n",
    "lammps_location = 'AgI_Molten'\n",
    "\n",
    "mgp_model = MappedGaussianProcess(grid_params, data['species'], \n",
    "    var_map=None, lmp_file_name='AgI_Molten', n_cpus=1)\n",
    "mgp_model.build_map(gp_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The coefficient file for LAMMPS mgp pair_style is automatically saved once the mapping is done. \n",
    "Saved as `lmp_file_name`. \n",
    "\n",
    "## Run LAMMPS with MGP pair style\n",
    "\n",
    "With the above coefficient file, we can run LAMMPS simulation with the mgp pair style. \n",
    "First download our mgp pair style files, compile your lammps executable with mgp pair style following our [instruction](https://flare.readthedocs.io/en/latest/installation/lammps.html) in the *Installation* section.\n",
    "\n",
    "1. One way to use it is running `lmp_executable < in.lammps > log.lammps` \n",
    "with the executable provided in our repository. \n",
    "When creating the input file, please note to set\n",
    "\n",
    "```\n",
    "newton off\n",
    "pair_style mgp\n",
    "pair_coeff * * <lmp_file_name> <chemical_symbols> yes/no yes/no\n",
    "```\n",
    "\n",
    "An example is using coefficient file `AgI_Molten.mgp` for AgI system, \n",
    "with two-body (the 1st `yes`) together with three-body (the 2nd `yes`).\n",
    "\n",
    "```\n",
    "pair_coeff * * AgI_Molten.mgp Ag I yes yes\n",
    "```\n",
    "\n",
    "2. Another way is to use the ASE LAMMPS interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/yuxie/opt/anaconda3/lib/python3.8/site-packages/ase/calculators/lammpsrun.py:191: UserWarning: You are using an old syntax to set 'parameters'.\n",
      "Please use LAMMPSRUN.set().\n",
      "  warnings.warn(self.legacy_warn_string.format(\"parameters\"))\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from flare.utils.element_coder import _Z_to_mass, _element_to_Z\n",
    "from flare.ase.calculator import FLARE_Calculator\n",
    "from ase.calculators.lammpsrun import LAMMPS\n",
    "\n",
    "from ase import Atoms\n",
    "\n",
    "# create test structure\n",
    "species = otf_object.gp_species_list[-1]\n",
    "positions = otf_object.position_list[-1]\n",
    "forces = otf_object.force_list[-1]\n",
    "otf_cell = otf_object.header['cell']\n",
    "structure = Atoms(symbols=species, cell=otf_cell, positions=positions)\n",
    "\n",
    "# get chemical symbols, masses etc.\n",
    "species = gp_model.training_statistics['species']\n",
    "specie_symbol_list = \" \".join(species)\n",
    "masses=[f\"{i} {_Z_to_mass[_element_to_Z[species[i]]]}\" for i in range(len(species))]\n",
    "\n",
    "# set up input params\n",
    "parameters = {'command': os.environ.get('lmp'), # set up executable for ASE\n",
    "              'newton': 'off',\n",
    "              'pair_style': 'mgp',\n",
    "              'pair_coeff': [f'* * {lammps_location + \".mgp\"} {specie_symbol_list} yes yes'],\n",
    "              'mass': masses}\n",
    "files = [lammps_location + \".mgp\"]\n",
    "\n",
    "# create ASE calc\n",
    "lmp_calc = LAMMPS(label=f'tmp_AgI', keep_tmp_files=True, tmp_dir='./tmp/',\n",
    "        parameters=parameters, files=files, specorder=species)\n",
    "\n",
    "structure.calc = lmp_calc\n",
    "\n",
    "# To compute energy, forces and stress\n",
    "# energy = structure.get_potential_energy()\n",
    "# forces = structure.get_forces()\n",
    "# stress = structure.get_stress()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. The third way to run LAMMPS is using our LAMMPS interface, please set the\n",
    "environment variable `$lmp` to the executable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flare import struc\n",
    "from flare.lammps import lammps_calculator\n",
    "\n",
    "# lmp coef file is automatically written now every time MGP is constructed\n",
    "\n",
    "# create test structure\n",
    "species = otf_object.gp_species_list[-1]\n",
    "positions = otf_object.position_list[-1]\n",
    "forces = otf_object.force_list[-1]\n",
    "otf_cell = otf_object.header['cell']\n",
    "structure = struc.Structure(otf_cell, species, positions)\n",
    "\n",
    "atom_types = [1, 2]\n",
    "atom_masses = [108, 127]\n",
    "atom_species = [1, 2] * 27\n",
    "\n",
    "# create data file\n",
    "data_file_name = 'tmp.data'\n",
    "data_text = lammps_calculator.lammps_dat(structure, atom_types,\n",
    "                                         atom_masses, atom_species)\n",
    "lammps_calculator.write_text(data_file_name, data_text)\n",
    "\n",
    "# create lammps input\n",
    "style_string = 'mgp'\n",
    "coeff_string = '* * {} Ag I yes yes'.format(lammps_location)\n",
    "lammps_executable = '$lmp'\n",
    "dump_file_name = 'tmp.dump'\n",
    "input_file_name = 'tmp.in'\n",
    "output_file_name = 'tmp.out'\n",
    "input_text = \\\n",
    "    lammps_calculator.generic_lammps_input(data_file_name, style_string,\n",
    "                                           coeff_string, dump_file_name)\n",
    "lammps_calculator.write_text(input_file_name, input_text)\n",
    "\n",
    "# To run lammps and get forces\n",
    "# lammps_calculator.run_lammps(lammps_executable, input_file_name,\n",
    "#                              output_file_name)\n",
    "\n",
    "# lammps_forces = lammps_calculator.lammps_parser(dump_file_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}