nuts_and_bolts_of_e3nn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nuts and Bolts of `e3nn`\n",
    "\n",
    "## tutorial by: Tess E. Smidt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Everything in `e3nn` starts with a \"representation list\"\n",
    "This defines the data types of your spherical tensor. At every step in the network, there needs to be an `Rs` that identifies the entries in your spherical tensor.\n",
    "\n",
    "Let's start with a simple example where we are starting off with a benezene molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import e3nn\n",
    "torch.set_default_dtype(torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rs = [(2, 0)] # Two (2) scalar (L=0) channels: hydrogen and carbon\n",
    "\n",
    "# 3D coordinates of the atoms of the molecule\n",
    "C_geo = torch.tensor(\n",
    "    [[ 0.     ,  1.40272,  0.     ],\n",
    "     [-1.21479,  0.70136,  0.     ],\n",
    "     [-1.21479, -0.70136,  0.     ],\n",
    "     [ 0.     , -1.40272,  0.     ],\n",
    "     [ 1.21479, -0.70136,  0.     ],\n",
    "     [ 1.21479,  0.70136,  0.     ]]\n",
    ")\n",
    "H_geo = torch.tensor(\n",
    "    [[ 0.     ,  2.49029,  0.     ],\n",
    "     [-2.15666,  1.24515,  0.     ],\n",
    "     [-2.15666, -1.24515,  0.     ],\n",
    "     [ 0.     , -2.49029,  0.     ],\n",
    "     [ 2.15666, -1.24515,  0.     ],\n",
    "     [ 2.15666,  1.24515,  0.     ]]\n",
    ")\n",
    "geometry = torch.cat([C_geo, H_geo], axis=-2)\n",
    "\n",
    "# and features on each atom\n",
    "C_input = torch.tensor([[0., 1.] for i in range(C_geo.shape[-2])])\n",
    "H_input = torch.tensor([[1., 0.] for i in range(H_geo.shape[-2])])\n",
    "input = torch.cat([C_input, H_input])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a convolution\n",
    "A traditional convolution applies a filter function $W(\\vec{r_ij})$ to an input $I_j$ with an element-wise multiply $\\odot$ and sums for all $j$ neighboring $i$.\n",
    "\n",
    "$\\text{Traditional Convolution} = \\sum_{j \\in n_i} W(\\vec{r}) \\odot I_j$\n",
    "\n",
    "Our convolution differs in two ways:\n",
    "* $W(\\vec{r})$  is constrained to be seperable into a learned radial function $R(r)$ and spherical harmonics $Y_{lm}(\\hat{r})$\n",
    "* Our datatypes are geometric tensors so we replace the element-wise multiply with a tensor product and Clebsch-Gordon coefficients\n",
    "\n",
    "$\\text{E(3) Convolution} = \\sum_{j \\in n_{i}} R(r) Y_{lm}(\\hat{r}) \\otimes I_j$\n",
    "\n",
    "We will use 3 classes to build our convolutions an instance of a class based on `e3nn.radial.RadialModel`, a instance of `e3nn.kernel.Kernel` and an instance of `e3nn.point.operations.Convolution`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RadialModel\n",
    "We are going to use the `CosineBasisModel` function for our basis functions\n",
    "They look a bit like Gaussians, but they don't have long tails. The function is equal to $\\text{cos}^2(arg)$ modulo normalization constants. You can use any radial basis you like, just make sure it's well normalized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inside e3nn.radial.CosineBasisModel\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "max_radius = 3.0\n",
    "number_of_basis = 3\n",
    "radii = torch.linspace(0, max_radius, steps=number_of_basis)\n",
    "step = radii[1] - radii[0]\n",
    "basis = lambda x: x.div(step).add(1).relu().sub(2).neg().relu().add(1).mul(math.pi / 2).cos().pow(2)\n",
    "\n",
    "x = torch.linspace(-max_radius, max_radius, 1000)\n",
    "plt.plot(x, basis(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using functools.partial to set up classes\n",
    "[Mario Geiger](https://mariogeiger.ch/) is the [BDFL](https://en.wikipedia.org/wiki/Benevolent_dictator_for_life) of `e3nn`. `functools.partial` is Mario's favorite python function so we'll use it to define our radial function, kernel, convolution, and layers.\n",
    "\n",
    "We are going to define `RadialModel` by specifying every single argument of `CosineBasisModel` using `functools.partial` EXCEPT out_dim which will be passed later. The `CosineBasisModel` will then have a `radial_layer=2` layer fully connected neural network with hidden dimention `h=100` applied to the basis functions vectors.\n",
    "\n",
    "$$R_{c}(r_{ij}) = W_{ch} \\sigma(W_{hb} B_b(r_{ij}))$$\n",
    "\n",
    "where $b$ is the number of basis funcitons and $c$ is unintuitively the number of \"paths\" -- the number of valid combinations of $L_{input}$, $L_{filter}$, and $L_{output}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# e3nn has operations for point sets and for 3D images. We will be using points.\n",
    "import e3nn.point\n",
    "import e3nn.radial\n",
    "from functools import partial\n",
    "from e3nn.non_linearities import rescaled_act\n",
    "\n",
    "# We are going to define RadialModel by specifying every single argument\n",
    "# of CosineBasisModel EXCEPT out_dim which will be passed later\n",
    "radial_layers = 2\n",
    "sp = rescaled_act.Softplus(beta=5)\n",
    "RadialModel = partial(e3nn.radial.CosineBasisModel, max_radius=max_radius,\n",
    "                      number_of_basis=number_of_basis, h=100,\n",
    "                      L=radial_layers, act=sp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the kernel\n",
    "Now we define the kernel or filter function by passing `RadialModel` in yet another call to `partial`. This time we are specifying all arguements to `Kernel` EXCEPT `Rs_in` and `Rs_out`, which again will be passed later.\n",
    "\n",
    "One of the arguements to `Kernel` is `sh` which specifies which spherical harmonics to use. In the vast majority of cases, this should be set to the default `o3.spherical_harmonics_xyz`.\n",
    "\n",
    "### ...but just in case...\n",
    "If you need to backpropogate to your input geometry, then you must use `e3nn.o3.spherical_harmonics_xyz_backwardable`. This should only be used for low $L$ as (for the current implementation) gradients become unstable for high $L$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import e3nn.kernel\n",
    "K = partial(e3nn.kernel.Kernel, RadialModel=RadialModel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## And finally we define our convolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import e3nn.point.operations\n",
    "\n",
    "# If we wish to pass the convolution to a layer definition\n",
    "C = partial(e3nn.point.operations.Convolution, K)\n",
    "\n",
    "# Or alternatively, if we want to use the convolution directly,\n",
    "# we need to specify the `Rs` of the input and output\n",
    "Rs_in = [(2, 0)]\n",
    "Rs_out = [(4, 0), (4, 1), (4, 2)]\n",
    "convolution = e3nn.point.operations.Convolution(K(Rs_in, Rs_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nonlinearities\n",
    "To preserve equivariance, we need to use equivariant nonlinearities. One such example is `e3nn.non_linearities.gated_block.GatedBlock` which applies a nonlinearity by gating each $L > 0$ channel with an added scalar channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import e3nn.non_linearities as nl\n",
    "from e3nn.non_linearities import rescaled_act\n",
    "\n",
    "gated_block = nl.gated_block.GatedBlock(Rs_out, sp, rescaled_act.sigmoid)\n",
    "\n",
    "dimensionalities = [2 * L + 1 for mult, L in Rs_out for _ in range(mult)]\n",
    "gated_act = nl.GatedBlock(\n",
    "    dimensionalities, rescaled_act.sigmoid, rescaled_act.sigmoid)"
   ]
  }
 ],
 "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}