Tutorial.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# Tutorial to the GraNNField Simulation Package",
   "id": "2a22bfe29c6f350f"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 1. Data Extraction from AIMD",
   "id": "9f2e98d7329a1e04"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from ase.io import read, write\n",
    "from grannfield.data.data_from_files import extxyz_to_db"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "dir = 'path_to_aimd_data'",
   "id": "2d2be3e6cc0ec3f3"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Read AIMD trajectories\n",
    "atoms = read(dir+'/vasprun.xml', index=':')"
   ],
   "id": "372b20835ea82f32"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Create extxyz of AIMD trajectories\n",
    "write('aimd.extxyz', atoms, append=True, format='extxyz')"
   ],
   "id": "622e4cc1df01bc59"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Generate GraNNField database (.db) for model training\n",
    "extxyz_to_db('aimd.extxyz', 'database.db')"
   ],
   "id": "de7046727f4128a7"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 2. Model Training",
   "id": "dbf8e7a0597796e4"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import torch\n",
    "\n",
    "from torch.optim import Adam\n",
    "\n",
    "from grannfield.data.materials import MaterialsData, ConcatMaterialsData\n",
    "from grannfield.data.data import MaterialsLoader, create_subset, train_test_split\n",
    "\n",
    "from grannfield.learn.trainer import Trainer\n",
    "from grannfield.learn.hooks import CSVHook, ExponentialDecayHook, ReduceLROnPlateauHook\n",
    "\n",
    "from grannfield.utils.statistics import Normalizer"
   ],
   "id": "f9ee2049941f29a9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "dir = 'path_to_database_directory'",
   "id": "5c145bd0e9b5c709"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "dataset = ConcatMaterialsData(MaterialsData(os.path.join(dir, db)))",
   "id": "2e838644c1a45fa7"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "index_train = 'index_of_training'\n",
    "index_val = 'index_of_val'"
   ],
   "id": "b1dd95a4657978e2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "train = create_subset(dataset, index_train)\n",
    "val = create_subset(dataset, index_val)"
   ],
   "id": "e6e55c777c723105"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "train_loader = MaterialsLoader(train, batch_size=2, shuffle=True)\n",
    "val_loader = MaterialsLoader(val, batch_size=1)"
   ],
   "id": "d6c8e5602c67ca30"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from grannfield.models.escn.escn import eSCN\n",
    "\n",
    "model = eSCN(\n",
    "    max_neighbors = 500,\n",
    "    cutoff = 6.0,\n",
    "    use_pbc = True,\n",
    "    regress_forces = True,\n",
    "    offset_degree = 1\n",
    ")"
   ],
   "id": "dca8132370686d8d"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = 'cuda:0'\n",
    "else:\n",
    "    device = 'cpu'"
   ],
   "id": "b6811ff6b325bae8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "optimizer = Adam(model.parameters(), lr=1E-3, weight_decay=1E-6, betas=(0.9, 0.999), amsgrad=True)",
   "id": "14f7b290252e90eb"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from grannfield.learn.metrics import MeanAbsoluteError\n",
    "\n",
    "metrics = [MeanAbsoluteError(target='energy', model_output='energy'),\n",
    "           MeanAbsoluteError(target='forces', model_output='forces')]"
   ],
   "id": "9fba0456c04bef05"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "y_weight = 1\n",
    "dy_weight = 1000\n",
    "\n",
    "def loss(batch, result):\n",
    "    diff_energy = batch['energy']-result['energy']\n",
    "    err_sq_energy = torch.mean(diff_energy ** 2)\n",
    "\n",
    "    diff_forces = batch['forces'] - result['forces']\n",
    "    err_sq_forces = torch.mean(diff_forces ** 2)\n",
    "\n",
    "    err_sq = y_weight * err_sq_energy + dy_weight * err_sq_forces\n",
    "\n",
    "    return err_sq"
   ],
   "id": "577a41e30e045fa2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "chk = './'",
   "id": "89e9ff2f66ab437"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "hooks = [\n",
    "    CSVHook(log_path=chk, metrics=metrics),\n",
    "    ReduceLROnPlateauHook(\n",
    "        optimizer,\n",
    "        patience=5,\n",
    "        min_lr=1E-6,\n",
    "        factor=0.8,\n",
    "        stop_after_min=True\n",
    "    )\n",
    "]"
   ],
   "id": "8cf111ea6af1dd70"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "normalizers = torch.load('path_to_normalizers')",
   "id": "2bf5a4c000431901"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "trainer = Trainer(\n",
    "    model_path = chk,\n",
    "    model = model,\n",
    "    hooks = hooks,\n",
    "    loss_fn = loss,\n",
    "    optimizer = optimizer,\n",
    "    train_loader = train_loader,\n",
    "    checkpoint_interval = 5000,\n",
    "    validation_loader = val_loader,\n",
    "    validation_interval = len(train_loader),\n",
    "    normalizers = normalizers,\n",
    "    clip_grad_norm = None,\n",
    "    ema_decay = 0.999,\n",
    "    amp = False,\n",
    "    early_stopping_time = 3600*10,\n",
    "    device = device,\n",
    "    print_freq = 1000\n",
    ")"
   ],
   "id": "485894264e684aa3"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "n_epochs = 2000",
   "id": "755e38973ee0fb65"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "trainer.train(n_epochs=n_epochs)",
   "id": "b170aed29d24d2a8"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 3. Metrics Evaluation (Ionic Conductivity)",
   "id": "7597e1d4ed87fe3e"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from pymatgen.io.ase import AseAtomsAdaptor\n",
    "from pymatgen.analysis.diffusion.analyzer import DiffusionAnalyzer, get_arrhenius_plot, get_extrapolated_conductivity"
   ],
   "id": "138c2d32d74cdfc7"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "ase_to_pmg = AseAtomsAdaptor()",
   "id": "ef54b6dbd42ceafd"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "diffusing_species = 'Li'\n",
    "temperature = 300\n",
    "time_step = 2\n",
    "smoothed = False\n",
    "steps_to_ignore = 15000\n",
    "avg_nsteps = 1000\n",
    "step_skip = 1"
   ],
   "id": "d7fcf364e937f20d"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "diff_analyzer = DiffusionAnalyzer.from_structures('path_to_trajectories',\n",
    "                                                           specie=diffusing_species,\n",
    "                                                           temperature=temperatures,\n",
    "                                                           time_step=time_step,\n",
    "                                                           smoothed=smoothed,\n",
    "                                                           step_skip=step_skip,\n",
    "                                                           avg_nsteps=avg_nsteps)"
   ],
   "id": "4b89a1faf28c1e95"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "diffusivity = diff_analyzer.diffusivity",
   "id": "919679b0e150f11d"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 4. Molecular Dynamics (ASE)",
   "id": "11abad3b2791e4d4"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import shutil\n",
    "from datetime import datetime\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from pymatgen.io.ase import AseAtomsAdaptor\n",
    "import seaborn as sns\n",
    "\n",
    "from asap3.analysis.rdf import RadialDistributionFunction\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as mpatches\n",
    "from pymatgen.analysis.diffusion.analyzer import DiffusionAnalyzer, get_arrhenius_plot, get_extrapolated_conductivity\n",
    "\n",
    "from grannfield.data.materials import MaterialsData\n",
    "\n",
    "import time\n",
    "from ase import units\n",
    "from ase.io import read, write\n",
    "from ase.md.velocitydistribution import MaxwellBoltzmannDistribution\n",
    "from ase.md.velocitydistribution import Stationary, ZeroRotation\n",
    "\n",
    "from grannfield.utils.md import NoseHoover"
   ],
   "id": "fbdca509d4dc4fbd"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = 'cuda:1'\n",
    "else:\n",
    "    device = 'cpu'"
   ],
   "id": "3f791d3fcd694a10"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "chk = './'",
   "id": "1bbba71d0c3786b8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from grannfield.models.escn.escn import eSCN\n",
    "\n",
    "model = eSCN(\n",
    "    max_neighbors = 500,\n",
    "    cutoff = 6.0,\n",
    "    use_pbc = True,\n",
    "    regress_forces = True,\n",
    "    offset_degree = 5\n",
    ")"
   ],
   "id": "6eee6b4b1cf1de00"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "normalizers = 'path_to_normalizers'",
   "id": "b32825bd29f03cae"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "model.load_state_dict('path_to_model_checkpoint')",
   "id": "ee248136912cb8e6"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from grannfield.utils.calculator import GraNNFieldCalculator\n",
    "\n",
    "calculator = GraNNFieldCalculator(base_model=model, device=device, energy='energy',forces='forces', normalizers = normalizers, apply_constraint=False)"
   ],
   "id": "4df7a970dddd5d11"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "ase_to_pmg = AseAtomsAdaptor()",
   "id": "9533441b707ff33b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "def save_to_xyz(atoms, logdir, prefix=\"\"):\n",
    "    write(\n",
    "        filename=os.path.join(os.path.join(logdir), prefix + '.xyz'),\n",
    "        images=atoms,\n",
    "        format=\"extxyz\",\n",
    "        append=True,\n",
    "    )"
   ],
   "id": "120a4ae6488f0162"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "save_dir = 'path_for_saved_md_data'\n",
    "\n",
    "\n",
    "at = read('path_to_input_structure')\n",
    "at.set_calculator(calculator)\n",
    "\n",
    "temperature = 'target_temperature'\n",
    "timestep = 2 * units.fs\n",
    "n_steps = 25000\n",
    "\n",
    "MaxwellBoltzmannDistribution(atoms=at, temperature_K=temperature)\n",
    "\n",
    "ZeroRotation(at)\n",
    "Stationary(at)\n",
    "\n",
    "log_frequency = 100\n",
    "save_frequency = 10\n",
    "log_dir = save_dir\n",
    "prefix = 'prefix'\n",
    "\n",
    "print(prefix)\n",
    "\n",
    "save_to_xyz(at, logdir=log_dir, prefix=prefix)\n",
    "\n",
    "nvt_dyn = NoseHoover(\n",
    "atoms=at,\n",
    "timestep=timestep,\n",
    "temperature=temperature,\n",
    ")\n",
    "\n",
    "start = time.time()\n",
    "for step in range(1, n_steps):\n",
    "    nvt_dyn.run(steps=1)\n",
    "    end = time.time()\n",
    "    hours, rem = divmod(end-start, 3600)\n",
    "    minutes, seconds = divmod(rem, 60)\n",
    "\n",
    "    if not step % log_frequency:\n",
    "        print(\"GRANNFIELD_MD: Step \", step)\n",
    "        print(\"Elapsed Time [hh:mm:ss.ms]: {:0>2}:{:0>2}:{:05.2f}\".format(int(hours),int(minutes),seconds))\n",
    "        print(f\"Simulation Time [ps]: {timestep / units.fs * step / 1000}\")\n",
    "        print(f\"Temperature [K]: {at.get_temperature()}\")\n",
    "        print(f\"Total Energy [eV]: {at.get_potential_energy()}\")\n",
    "        print(f\"Forces Max [eV/Å]: {at.get_forces().max()}\\n\")\n",
    "\n",
    "    if not step % save_frequency:\n",
    "        save_to_xyz(at, logdir=log_dir, prefix=prefix)"
   ],
   "id": "3f70ea3ff97c0079"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 5. Molecular Dynamics ",
   "id": "33efc8f6a9cb41c1"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Example LAMMPS input file:\n",
    "\n",
    "units         metal\n",
    "boundary      p p p\n",
    "atom_style    atomic\n",
    "\n",
    "pair_style    grannfield/gpu .\n",
    "\n",
    "read_data     ./input.data\n",
    "\n",
    "pair_coeff    * * model_checkpoint.pth.tar Ge Li P S\n",
    "\n",
    "dump          1 all custom 1 output.lammpstrj id element x y z fx fy fz\n",
    "dump_modify   1 sort id element Ge Li P S\n",
    "\n",
    "thermo_style  custom step time cpu pe ke etotal temp\n",
    "thermo        1\n",
    "log           none\n",
    "\n",
    "velocity      all create 800 12345\n",
    "fix           1 all nvt temp 800 800 0.1\n",
    "timestep      2.0e-3\n",
    "run           25000\n"
   ],
   "id": "67edb6d6a7c9abfd"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}