basic-example-lode.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running this example requires [pyLODE](https://github.com/ceriottm/lode)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import ase.io\n",
    "\n",
    "import equistore.io\n",
    "\n",
    "from utils.pylode import PyLODESphericalExpansion\n",
    "from utils.models.soap import compute_power_spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_frames = 3\n",
    "frames = ase.io.read('data/BaTiO3_Training_set.xyz', f':{n_frames}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LODE calculations can take quite a while...\n",
    "# Therefore, it is useful to store results and only recalculate them if \n",
    "# needed!\n",
    "\n",
    "recalc = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "lode_hypers = dict(\n",
    "    smearing=3,\n",
    "    max_radial=2,\n",
    "    max_angular=3,\n",
    "    cutoff_radius=4.5,\n",
    "    potential_exponent=1,\n",
    "    radial_basis='GTO',\n",
    "    compute_gradients=True,\n",
    ")\n",
    "\n",
    "calculator = PyLODESphericalExpansion(lode_hypers)\n",
    "\n",
    "fname = \"precomputed.npz\"\n",
    "if recalc:\n",
    "    descriptor = calculator.compute(frames)\n",
    "    equistore.io.save(fname, descriptor, use_numpy=True)\n",
    "else:\n",
    "    descriptor = equistore.io.read(fname, use_numpy=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Lode descripor can also be precomputed using the `precompute_lode.py` from \n",
    "the script directory. \n",
    "\n",
    "You can either precompute the descriptor from the python \n",
    "interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scripts.precompute_lode import precompute_lode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "descriptor = precompute_lode(frames, lode_hypers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or from the command line with\n",
    "\n",
    "```bash\n",
    "scripts/precompute_lode.py -input_file \"data/BaTiO3_Training_set.xyz\" \\\n",
    "                           -index \":10\" \\\n",
    "                           -max_radial 2 \\\n",
    "                           -max_angular 0 \\\n",
    "                           -cutoff_radius 4.5 \\\n",
    "                           -smearing 3 \\\n",
    "                           -radial_basis \"GTO\" \\\n",
    "                           -subtract_center_contribution \\\n",
    "                           -compute_gradients \\\n",
    "                           -outfile \"precomputed_lode\"\n",
    "```\n",
    "\n",
    "The exact syntax is given on the help page\n",
    "\n",
    "```bash\n",
    "scripts/precompute_lode.py -h\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Application: Compute power spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example usage: Multiscale invariants (rho x V) combining short and long range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters for short range (SR) densities\n",
    "hypers_sr = dict(\n",
    "    smearing=3,\n",
    "    max_radial=2,\n",
    "    max_angular=1,\n",
    "    cutoff_radius=4.5,\n",
    "    potential_exponent=0, # 0 here means Gaussian/SR densities\n",
    "    radial_basis='GTO',\n",
    "    compute_gradients=True,\n",
    ")\n",
    "calculator_sr = PyLODESphericalExpansion(hypers_sr)\n",
    "descriptor_sr = calculator_sr.compute(frames)\n",
    "\n",
    "# Define hyperparameters for long range (LR) densities\n",
    "hypers_lr = dict(\n",
    "    smearing=3,\n",
    "    max_radial=1,\n",
    "    max_angular=1,\n",
    "    cutoff_radius=4.5,\n",
    "    potential_exponent=1, # 1 here means Coulombic/LR densities\n",
    "    radial_basis='monomial',\n",
    "    compute_gradients=True,\n",
    ")\n",
    "calculator_lr = PyLODESphericalExpansion(hypers_lr)\n",
    "descriptor_lr = calculator_sr.compute(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate SR x SR (rho x rho) invariants\n",
    "invariants_srsr = compute_power_spectrum(descriptor_sr)\n",
    "# Generate LR x LR (V x V) invariants\n",
    "invariants_lrlr = compute_power_spectrum(descriptor_lr)\n",
    "# Generate SR x LR (rho x V)\n",
    "invariants_multiscale = compute_power_spectrum(descriptor_sr, descriptor_lr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test that the coefficients obtained from the general method agree with those from the specialized version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'species_center': 8, 'species_neighbor_1': 8, 'species_neighbor_2': 8}\n"
     ]
    }
   ],
   "source": [
    "for k in invariants_srsr.keys.as_namedtuples(): \n",
    "    print(k.as_dict())\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Labels([( 8,  8), ( 8,  8), ( 8,  8), ( 8, 22), ( 8, 22), ( 8, 56),\n",
       "        (22,  8), (22,  8), (22,  8), (22, 22), (22, 22), (22, 56),\n",
       "        (56,  8), (56,  8), (56,  8), (56, 22), (56, 22), (56, 56)],\n",
       "       dtype={'names': ['species_center', 'species_neighbor_1'], 'formats': ['<i4', '<i4'], 'offsets': [0, 4], 'itemsize': 12})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invariants_srsr.keys[[\"species_center\", \"species_neighbor_1\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "invariants_srsr_2 = compute_power_spectrum(descriptor_sr, descriptor_sr)\n",
    "for key in invariants_srsr.keys:\n",
    "    # Get block from original code\n",
    "    block_ref = invariants_srsr.block(key)\n",
    "\n",
    "    # Compare this to invariants obtained from generalization\n",
    "    block_1 = invariants_srsr_2.block(key)\n",
    "\n",
    "    # Also compare with the invariants with swapped neighbor species\n",
    "    block_2 = invariants_srsr_2.block(key)\n",
    "\n",
    "    assert block_ref.values.shape == block_1.values.shape\n",
    "    assert block_ref.values.shape == block_2.values.shape\n",
    "    assert np.linalg.norm(block_ref.values - block_1.values) < 1e-15\n",
    "    assert np.linalg.norm(block_ref.values - block_2.values) < 1e-15"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.0 ('virtualenv': venv)",
   "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.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "9b1dab49005db54eeaf19407b770a132388229704ecb91b436b854691f0d31ef"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}