adding_thresholding.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The text.latex.preview rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The mathtext.fallback_to_cm rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: Support for setting the 'mathtext.fallback_to_cm' rcParam is deprecated since 3.3 and will be removed two minor releases later; use 'mathtext.fallback : 'cm' instead.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The validate_bool_maybe_none function was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The savefig.jpeg_quality rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The keymap.all_axes rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The animation.avconv_path rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n",
      "In /home/pozdn/.local/lib/python3.6/site-packages/matplotlib/mpl-data/stylelib/_classic_test.mplstyle: \n",
      "The animation.avconv_args rcparam was deprecated in Matplotlib 3.3 and will be removed two minor releases later.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import ase.io\n",
    "import tqdm\n",
    "from pytorch_prototype.nice_blocks import *\n",
    "from pytorch_prototype.utilities import *\n",
    "from matplotlib import pyplot as plt\n",
    "from torch import nn\n",
    "from nice.rascal_coefficients import process_structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "LAMBDA_MAX = 4\n",
    "clebsch = ClebschGordan(LAMBDA_MAX)\n",
    "HYPERS = {\n",
    "    'interaction_cutoff': 6.3,\n",
    "    'max_radial': 4,\n",
    "    'max_angular': LAMBDA_MAX,\n",
    "    'gaussian_sigma_type': 'Constant',\n",
    "    'gaussian_sigma_constant': 0.05,\n",
    "    'cutoff_smooth_width': 0.3,\n",
    "    'radial_basis': 'GTO'\n",
    "}\n",
    "MAGIC_NUMBER = 10000\n",
    "BATCH_SIZE = 500\n",
    "structures = process_structures(ase.io.read('methane.extxyz', index='0:10'))\n",
    "all_species = get_all_species(structures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split(coefficients):\n",
    "    lambd_max = max(coefficients.keys())\n",
    "    even, odd = {}, {}\n",
    "    for lambd in range(lambd_max + 1):\n",
    "        if lambd % 2 == 0:\n",
    "            even[str(lambd)] = coefficients[lambd]\n",
    "        else:\n",
    "            odd[str(lambd)] = coefficients[lambd]\n",
    "    return even, odd\n",
    "\n",
    "coefficients_all = get_coefs(structures, HYPERS, all_species)\n",
    "for key in coefficients_all.keys():\n",
    "    coefficients_all[key] *= MAGIC_NUMBER\n",
    "central = get_central_species(structures)\n",
    "mask = central == 6\n",
    "\n",
    "for key in coefficients_all.keys():\n",
    "    coefficients_all[key] = coefficients_all[key][mask]\n",
    " \n",
    "even, odd = split(coefficients_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "even: \n",
      "0 torch.Size([10, 8, 1])\n",
      "2 torch.Size([10, 8, 5])\n",
      "4 torch.Size([10, 8, 9])\n",
      "odd: \n",
      "1 torch.Size([10, 8, 3])\n",
      "3 torch.Size([10, 8, 7])\n"
     ]
    }
   ],
   "source": [
    "print(\"even: \")\n",
    "for key in even.keys():\n",
    "    print(key, even[key].shape)\n",
    "print(\"odd: \")\n",
    "for key in odd.keys():\n",
    "    print(key, odd[key].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 torch.Size([10, 192, 1])\n",
      "1 torch.Size([10, 128, 3])\n",
      "2 torch.Size([10, 384, 5])\n",
      "3 torch.Size([10, 256, 7])\n",
      "4 torch.Size([10, 384, 9])\n"
     ]
    }
   ],
   "source": [
    "block = ClebschCombining(clebsch.precomputed_, LAMBDA_MAX)\n",
    "result = block.forward(even, even)\n",
    "for key in result.keys():\n",
    "    print(key, result[key].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = {}\n",
    "some = [np.array([0, 2]), np.array([0, 2])]\n",
    "for l1 in range(LAMBDA_MAX + 1):\n",
    "    for l2 in range(LAMBDA_MAX + 1):\n",
    "        for lambd in range(abs(int(l1) - int(l2)), min(int(l1) + int(l2), LAMBDA_MAX) + 1):\n",
    "            key = f'{l1}_{l2}_{lambd}'\n",
    "            #print(key)\n",
    "            task[key] = some"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 torch.Size([10, 6, 1])\n",
      "1 torch.Size([10, 4, 3])\n",
      "2 torch.Size([10, 12, 5])\n",
      "3 torch.Size([10, 8, 7])\n",
      "4 torch.Size([10, 12, 9])\n"
     ]
    }
   ],
   "source": [
    "block = ClebschCombining(clebsch.precomputed_, LAMBDA_MAX, task = task)\n",
    "result = block.forward(even, even)\n",
    "for key in result.keys():\n",
    "    print(key, result[key].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Experiment(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Experiment, self).__init__()\n",
    "        self.register_buffer('asdf', torch.FloatTensor([1, 2, 3]))\n",
    "        self.register_buffer('q', torch.IntTensor([1, 2, 3]))\n",
    "    def forward(self):\n",
    "        print(self.asdf)\n",
    "        print(self.q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1., 2., 3.])\n",
      "tensor([1, 2, 3], dtype=torch.int32)\n"
     ]
    }
   ],
   "source": [
    "block = Experiment()\n",
    "block()"
   ]
  },
  {
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}