example_plot_with_radial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting with Radial Functions\n",
    "\n",
    "### using the `e3nn` repository\n",
    "\n",
    "### tutorial by: Tess E. Smidt (`blondegeek`)\n",
    "\n",
    "We are going to use class methods of the `FourierTensor` class to plot angular and radial Fourier transforms of geometry. This allows us to use this type of data as input or output to our models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "import numpy as np\n",
    "import torch\n",
    "import e3nn\n",
    "from e3nn.tensor.spherical_tensor import (\n",
    "    SphericalTensor,\n",
    "    projection\n",
    ")\n",
    "from e3nn.tensor.fourier_tensor import (\n",
    "    FourierTensor,\n",
    "    plot_on_grid\n",
    ")\n",
    "import e3nn.o3 as o3\n",
    "import e3nn.rs as rs\n",
    "\n",
    "import plotly\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import math\n",
    "\n",
    "torch.set_default_dtype(torch.float64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples of RadialModels for plotting\n",
    "Note, these differ from the models in `e3nn.radial` since they do not inherit from `torch.nn.Modules` and have NO LEARNED PARAMETERS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ConstantRadialModel():\n",
    "    def radial_function(r):\n",
    "        shape = r.shape\n",
    "        return torch.ones(list(shape) + [1])\n",
    "    return radial_function\n",
    "\n",
    "def FixedCosineRadialModel(max_radius, number_of_basis, min_radius=0.):\n",
    "    spacing = (max_radius - min_radius) / number_of_basis\n",
    "    radii = torch.linspace(min_radius, max_radius, number_of_basis)\n",
    "    step = radii[1] - radii[0]\n",
    "\n",
    "    def radial_function(r):\n",
    "        shape = r.shape\n",
    "        radial_shape = [1] * len(shape) + [number_of_basis]\n",
    "        centers = radii.reshape(*radial_shape)\n",
    "        return (r.unsqueeze(-1) - centers).div(step).add(1).relu().sub(2).neg().relu().add(1).mul(math.pi / 2).cos().pow(2)\n",
    "    \n",
    "    return radial_function\n",
    "\n",
    "def FixedGaussianRadialModel(max_radius, number_of_basis, min_radius=0.):\n",
    "    spacing = (max_radius - min_radius) / number_of_basis\n",
    "    radii = torch.linspace(min_radius, max_radius, number_of_basis)\n",
    "    gamma = 1. / spacing\n",
    "    \n",
    "    def radial_function(r):\n",
    "        shape = r.shape\n",
    "        radial_shape = [1] * len(shape) + [number_of_basis]\n",
    "        centers = radii.reshape(*radial_shape)\n",
    "        return torch.exp(-gamma * (r.unsqueeze(-1) - centers) ** 2)\n",
    "    \n",
    "    return radial_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up coordinates for tetrahedra and set `lmax`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tetra_coords = torch.tensor( # The easiest way to construct a tetrahedron is using opposite corners of a box\n",
    "    [[0., 0., 0.], [1., 1., 0.], [1., 0., 1.], [0., 1., 1.]]\n",
    ")\n",
    "tetra_coords -= tetra_coords.mean(-2)\n",
    "\n",
    "lmax = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create and plot spherical harmonic projection with radial functions\n",
    "(Also called an angular and radial Fourier Transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_radial = 3\n",
    "max_radius = 2.\n",
    "sphten = FourierTensor.from_geometry(tetra_coords, FixedCosineRadialModel(max_radius, n_radial), lmax)\n",
    "x, f = sphten.plot(5.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_max = float(f.abs().max()) * 0.5\n",
    "trace = go.Volume(\n",
    "    x=x[:,0], y=x[:,1], z=x[:,2], value=f,\n",
    "    isomin=-plot_max,\n",
    "    isomax=plot_max,\n",
    "    opacity=0.3, # needs to be small to see through all surfaces\n",
    "    surface_count=10, # needs to be a large number for good volume rendering\n",
    "    colorscale='RdBu'\n",
    ")\n",
    "go.Figure([trace])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create and plot spherical harmonic projection with magnitude as radius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sphten = SphericalTensor.from_geometry(tetra_coords, lmax)\n",
    "r, f = sphten.plot(res=100, relu=True)\n",
    "trace = go.Surface(x=r[..., 0], y=r[..., 1], z=r[..., 2], surfacecolor=f)\n",
    "go.Figure(trace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}