perovskite_order_parameters_determine_irreps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Order parameters for octahedral tilting in perovskites\n",
    "In this notebook we will learn the irreps that describe octahedral tilting in perovskites, specifically those in space group $Pnma$ (62)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch_geometric as tg\n",
    "import e3nn\n",
    "from e3nn import rs, o3\n",
    "from e3nn.networks import GatedConvParityNetwork\n",
    "from e3nn.point.message_passing import Convolution\n",
    "import e3nn.point.data_helpers as dh\n",
    "from e3nn.tensor.spherical_tensor import SphericalTensor\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import pymatgen\n",
    "from pymatgen.core.structure import Structure\n",
    "from pymatgen.symmetry.analyzer import SpacegroupAnalyzer\n",
    "\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "device = \"cuda:2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import structures to use for training.\n",
    "We use structures based on the ICSD Structures 63157 and 16688 for space group $Pm\\bar{3}m$ (221) and $Pnma$ (62), respectively. We also make the following modifications:\n",
    "* transform the parent 221 structure to be in the 62 setting,\n",
    "  * ...so there's a one-to-one mapping of atoms between the two structures. \n",
    "* reshape the unit cells to have equal lattice parameters and replace atom types with site labels $A$, $B$, and $X$,\n",
    "  * ...while we can deal with changes in lattice size -- we didn't want to complicate this example.\n",
    "* rotate the 62 setting lattice to match the Cartesian coordinates of cubic setting lattice vectors,\n",
    "  * ...this makes it easier to interpret the order parameters.\n",
    "* and remove the $A$ site distortions for the $Pnma$ structures.\n",
    "  * ...so we can focus on only the irreps of the octahedral distortion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.055964 4.055964 0.000000\n",
      "-4.055964 4.055964 0.000000\n",
      "0.000000 0.000000 8.112000\n"
     ]
    }
   ],
   "source": [
    "struct_221_primitive = Structure.from_file('structures/221_orig.cif')\n",
    "struct_221 = Structure.from_file(\"structures/221_anon.cif\")\n",
    "struct_62 = Structure.from_file(filename=\"structures/62_anon_sorted.cif\")\n",
    "\n",
    "rlatt = torch.tensor(struct_221_primitive.lattice.reciprocal_lattice.abc)\n",
    "\n",
    "# Rotate space group 62 setting lattice to match Cartesian coordinates of cubic setting\n",
    "new_lattice = o3.rot_z(np.pi / 4.).T @ struct_221.lattice.matrix\n",
    "\n",
    "struct_221 = Structure(new_lattice, struct_221.species, struct_221.frac_coords)\n",
    "struct_62 = Structure(new_lattice, struct_62.species, struct_62.frac_coords)\n",
    "\n",
    "print(struct_221.lattice)\n",
    "lat = 4.055964"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "atom_types_list = ['A', 'B', 'X']\n",
    "atom_indices = [atom_types_list.index(str(s.species)[0]) for s in struct_221]  # Remove decorators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We set up the atom type and order parameter inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rs_order_param = [(1, 0, 1), (1, 0, -1), (1, 1, 1), (1, 1, -1), (1, 2, 1), (1, 2, -1)]\n",
    "A_input = torch.tensor([1., 0., 0.]).unsqueeze(0)\n",
    "B_input = torch.tensor([0., 1., 0.]).unsqueeze(0)\n",
    "X_input = torch.tensor([0., 0., 1.]).unsqueeze(0)\n",
    "atom_inputs = [A_input, B_input, X_input]\n",
    "Rs_in = [(3, 0, 1)] + Rs_order_param  # 3 atom types + order parameters\n",
    "N = len(struct_221)\n",
    "zeros = torch.zeros(N - 4, rs.dim(Rs_order_param))\n",
    "# order_param_input = torch.zeros(20, rs.dim(Rs_order_param))\n",
    "order_param_input = torch.zeros(4, rs.dim(Rs_order_param))\n",
    "# order_param_input.to(device)\n",
    "order_param_input.requires_grad = True\n",
    "\n",
    "lmax = 6\n",
    "\n",
    "Rs_out = [(1, 1, -1)]\n",
    "r_max = 3.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This function stitches together the fixed atom type inputs and the learnable order parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_features(order_param_input):\n",
    "    order_param = torch.cat([zeros, order_param_input], dim=0)\n",
    "\n",
    "    N = len(atom_indices)\n",
    "    all_atom_types = torch.cat([\n",
    "        atom_inputs[i] for i in atom_indices\n",
    "    ], dim=0)  # [N, atom_types]\n",
    "\n",
    "    return torch.cat([all_atom_types, order_param], dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We compute the displacements between the high and low symmetry structures, accounting for periodic boundary conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max distortion in Angstroms (62):  tensor(0.2565)\n"
     ]
    }
   ],
   "source": [
    "relative_vecs_62 = []\n",
    "for i, site in enumerate(struct_221):\n",
    "    sites = struct_62.get_sites_in_sphere(site.coords, 1., include_index=True, include_image=True)\n",
    "    sorted_sites = sorted(sites, key=lambda x: x[1])\n",
    "    r = sorted_sites[0][0].coords - site.coords\n",
    "    relative_vecs_62.append(r)\n",
    "relative_vecs_62 = np.stack(relative_vecs_62, axis=0)\n",
    "pos = torch.tensor(struct_221.cart_coords, dtype=torch.float64)\n",
    "relative_vecs_62 = torch.tensor(relative_vecs_62, dtype=torch.float64)\n",
    "print(\"max distortion in Angstroms (62): \", relative_vecs_62.norm(2, -1).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We define our model and optimizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_kwargs = {\n",
    "    'convolution': Convolution,\n",
    "    'Rs_in': Rs_in, 'Rs_out': Rs_out,\n",
    "    'mul': 6, 'lmax': lmax, 'layers': 3,\n",
    "    'max_radius': r_max,  # From dataset\n",
    "    'number_of_basis': 10,\n",
    "}\n",
    "\n",
    "model = GatedConvParityNetwork(**model_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = torch.optim.Adam(model.parameters(), 1e-3)\n",
    "order_opt = torch.optim.Adam([order_param_input], 1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, we overfit the model without symmetry breaking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('param_search_model_overfit.torch')['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Used to create overfit model\n",
    "# max_iter = 31\n",
    "\n",
    "# model.to(device)\n",
    "\n",
    "# for i in range(max_iter):\n",
    "#     features = create_features(order_param_input)\n",
    "#     data = dh.DataPeriodicNeighbors(features, Rs_in, pos, struct_221.lattice.matrix, r_max, y=relative_vecs_62)\n",
    "#     data.to(device)\n",
    "#     output = model(data.x, data.edge_index, data.edge_attr, n_norm=6)\n",
    "#     loss = ((output - data.y[:, [1, 2, 0]]).abs()).mean()\n",
    "#     if i % 10 == 0:\n",
    "#         print(\"model update: \", loss)\n",
    "#     opt.zero_grad()\n",
    "#     loss.backward()\n",
    "#     opt.step() \n",
    "\n",
    "# torch.save({'model': model.state_dict(), 'order_param_input': order_param_input},\n",
    "#            'param_search_model_overfit.torch')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now we train the model to recover the order parameters.\n",
    "We define a loss function on the order parameter input to encourage sparsity and penalize order parameters of higher order ($L$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def order_param_loss(order_param_input, Rs_order_param, scale=1e-2):\n",
    "    loss = scale * order_param_input.abs().mean()\n",
    "    for i, (m, L, p) in enumerate(Rs_order_param):\n",
    "        start = rs.dim(Rs_order_param[:i])\n",
    "        end = rs.dim(Rs_order_param[:i+1])\n",
    "        # Add a relatively severe penalty for higher L order parameters\n",
    "        loss = loss + scale * L ** 6 * order_param_input[:, start:end].abs().mean()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model update:  tensor(0.0676, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0676, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0629, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0623, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0180, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0244, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0170, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0176, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0106, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0104, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0105, device='cuda:2', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0110, grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "max_iter = 51\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "for i in range(max_iter):\n",
    "    order_param_input.to('cpu')\n",
    "    features = create_features(order_param_input)\n",
    "    data = dh.DataPeriodicNeighbors(features, Rs_in, pos, struct_221.lattice.matrix, r_max, y=relative_vecs_62)\n",
    "    data.to(device)\n",
    "    output = model(data.x, data.edge_index, data.edge_attr, n_norm=6)\n",
    "    loss = ((output - data.y[:, [1, 2, 0]]).abs()).mean()\n",
    "    if i % 10 == 0:\n",
    "        print(\"model update: \", loss)\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step() \n",
    "    \n",
    "    order_param_input.to('cpu')\n",
    "    features = create_features(order_param_input)\n",
    "    data = dh.DataPeriodicNeighbors(features, Rs_in, pos, struct_221.lattice.matrix, r_max, y=relative_vecs_62)\n",
    "    data.to(device)\n",
    "    output = model(data.x, data.edge_index, data.edge_attr, n_norm=6)\n",
    "    loss = ((output - data.y[:, [1, 2, 0]]).abs()).mean().cpu()\n",
    "    loss = loss + order_param_loss(order_param_input, Rs_order_param, scale=1e-2)\n",
    "    if i % 10 == 0:\n",
    "        print(\"order param update: \", loss)\n",
    "    order_opt.zero_grad()\n",
    "    loss.backward()\n",
    "    order_opt.step() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The largest contributions are for the `Rs = [(1, 1, 1)]` followed by `Rs = [(1, 2, 1)]`\n",
    "In this plot, the rows are different B sites and the columns are irreps components. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1, 0, 1), (1, 0, -1), (1, 1, 1), (1, 1, -1), (1, 2, 1), (1, 2, -1)]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Rs_order_param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fc40416da50>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWgAAABvCAYAAADBl1rZAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAQQklEQVR4nO3df7Bc5V3H8fcn94YEkpBAqPwoVEAiSp2BQgZRsYNCQ2C0oNYa2gG0lLRDmZFRxsHBUgZlhuK0KFrR8GOADvLDUiB2+N1Sa9VQQiZA+JmAcUhMQRJICDTk7r1f/zjPprub3bu75+y5d/fm85o5c3fPPs853z27++xznz3n+SoiMDOz/jNtsgMwM7Pm3ECbmfUpN9BmZn3KDbSZWZ9yA21m1qfcQJuZ9Sk30Ga2x5G0WNLLktZJuqzJ4zMk3Z0ef1LS4Wn9fElPSNou6e8b6pwg6blU53pJKhqnG2gz26NIGgK+AZwBHAOcI+mYhmIXAG9HxFHAdcBX0/odwJeBS5ts+gbgQmBBWhYXjdUNtJntaU4E1kXEaxGxE7gLOKuhzFnAben2t4BTJSki3ouIH5I11LtIOhjYNyJWRHb13+3A2UUDdQNtZnuaDwOv19zfkNY1LRMRFWArML/NNje02WbXhotuwMxsMp22aJ/YvHm0bt3qVTufp76Xuywilk1oYD3gBtrMBtrmt8b4t38/vG7d3Nmv7IiIhS2qbAQOq7l/aFrXrMwGScPAXGDzOGFsTNsZb5td8xCHmQ00BUzbqbqljaeABZKOkLQXsARY3lBmOXB+uv0p4HsxzsxyEbEJ2CbppHT2xnnAA3meTy33oM1ssAVMG+l8Vs6IqEi6GHgEGAJuiYjnJV0FrIyI5cDNwDclrQO2kDXiAEhaD+wL7CXpbGBRRLwAXATcCuwNPJSWQuTpRs1skJ1w7Mz4jwc/Urdu70PXPj3OEMfAcA/azAZbgEYmO4hyuIE2s8HW5RDHIHEDbWaDzT1oM7P+pACNFJ72oi+5gTazweYG2sysTwUw6gbazKz/hMA9aDOzPhSgkal5UbQbaDMbbAFU3EB3bN70feKgmfNy1Z02Lf/5jDOO2pK77uur8s8M+LZ25qr3i3Pznxs0Opb/DRlj+f8dnDZtLHfdIoIC/8IWOEW2SE6MvBfpSvkDjigQcKGq+WL+359s452d7xcbnwgR7kF37qCZ87jp+KW56s7ce0f7Qi0c9cA9ueteMvPq3HXv3Wt9rnr3nLop9z7fe3dW7ro/eX9m7rqzZr+fu26Rhmd0dCh/3Ur+usPTK7nrjuX8Eh0ezr/PSiX/R7rI65P3i/vc/7ytfaF23IM2M+tTIRjJ/yXczzr62mmXYNHMbDJFZVrdMlW0fSYdJlg0M5sc1R507dJG3qze6bE/T+tflnR6zfr1Kav3akkre/HUOhni2JVgMQVRTbD4Qi8CMDMrIoKues01nc5PkOUOfErS8jSnc9WurN6SlpBl9f6D1DldAnwUOAR4XNLPR0Q159ZvRMRbxZ9VppNn1UmCRTOzyREiKkN1Sxu5s3qn9XdFxAcR8d/AurS9UvRssEbSUkkrJa18ZyT/L/1mZl2Jrsegi2T1Hq9uAI9KelpSvtPYGnQyxNFJgkVSxtxlAL8w55CpOTmrmfWfEGO795oPaBgHnois3idHxEZJPwM8JumliPhBkQ120kDvSrBI1jAvAT5TZKdmZr3SYgz6rZKyeresGxHVv29Kuo9s6KNQA932f4HUva8mWHwRuCcini+yUzOz3ul6DLpIVu/lwJJ0lscRwALgR5JmSZoDIGkWsAhYU/SZdXShSkQ8CDxYdGdmZj0XNBviaF28QFbvVO4esrPYKsCXImJU0oHAfdnviAwD/xwRDxd9ar6S0MwGWoQY6/LilGadzoi4oub2DuD3W9S9Gri6Yd1rwLFdBdEBN9BmNtgCYnTqXD1Yq5QGetq0yD3pUZFJcQ6be2HuulvvvzZ33b9ZvDVXvadOzn8mzn7z38ldd86+7+WuOzKS/y1TZKa1oaHR9oVamDEz32yDUOz55lVkwqMiisxUmPdzW2j2vV3UybjzQHIP2swGWgSMuQdtZtaHmp8HPSW4gTazwRZ0/SPhoHADbWYDLRBjBX676mduoM1ssLkHbWbWnyJUKK1ZP3MDbWYDzkMcZmZ9KTzEYWbWp0KFLnDrZ26gzWyguQdtZta3pm4Pemp+7ZjZHiO71HuobmmnpKze424zDzfQZjbYQoxWptUt46nJ6n0GcAxwTsrWXWtXVm/gOrKs3jRk9V4M/IOkoQ632TU30GY20IJsNr3apY0ysnp3ss2ulTIGPTYmPtgxI1fdvNOUAry+9cbcdXl4bu6qeac5vf/Y/NM7Svnz8k6fkX/6zbGx/NNDFom5iA927JW77vD0Su66eZ9vkelGixzj3kz9OQm6TxrbLDP3LzfUr8vqLak2q/eKhrrVrN7tttk1/0hoZgMtAkZ3n250vKSxA8MNtJkNODE21tVobSlZvTvYZtc8Bm1mAy0CRitDdUsbPc/q3eE2u9a2By3pMOB24ECy8fhlEfG3RXdsZtYTXV5JWEZWb4Bm2yz61DoZ4qgAfxoRqyTNAZ6W9FhEvFB052ZmRVXP4uiqTo+zerfaZlFtG+iI2ARsSrfflfQi2a+WbqDNbNJ5utEkXU3zMeDJMoIxM+taQGVPTxoraTZwL3BJRGxr8vhSYCnAgTPyn1NsZtaN2NPn4pA0naxxviMivt2sTEQsi4iFEbFw3vR9ehmjmVlL1SGOLs7iGBidnMUhsl80X4yIr5cfkplZFwIqe3AP+teAc4HflLQ6LWeWHJeZWUeC7iZLGiSdnMXxQ2BAL9I3s6kupnAP2pd6m9lAi5DP4jAz60cBVKbQD4O1SmmgpSg0TWNeeaf9BPi9Dw7PXff1HZfnqrfxM7+de5873p+Zu+72rbNz1502lH+K1CKKTKM5Y2b+6VWLfPC7nMBnl+Hh/J+dIqeb5Y0XYGhoNFe9nkxB6x60mVl/Cl+oYmbWnwJRmUJnbtSams/KzPYYWQ9adUsRkvaX9Jiktenvfi3KnZ/KrJV0fs36EyQ9l5LHXp+uJUHSlZI2dnO6shtoMxtovW6ggcuA70bEAuC76X4dSfsDXyFLa3Ui8JWahvwG4EKyuaIXkCWXrbouIo5LS9uZ79xAm9nAGx1T3VJQbcLY24Czm5Q5HXgsIrZExNvAY8BiSQcD+0bEijTB/+0t6nfEDbSZDbQAKmOqWwo6ME2zDPBjsmQljZolnv1wWjY0WV91saRnJd3SauiklhtoMxtoEWJkrH4hZfWuWZbW1pH0uKQ1TZaz6rcdQfYd0As3AD8HHEc2x/7X2lXwWRxmNtACqOzehI6b1TsiTmv1mKQ3JB0cEZvSkMWbTYptBE6puX8o8P20/tCG9RvTPt+o2ceNwHdaxVDlHrSZDbQARqJ+Kag2Yez5wANNyjwCLJK0XxqqWAQ8koZGtkk6KZ29cV61fmrsq34HWNMuEPegzWygtehBF3ENcI+kC4D/AT4NIGkh8MWI+HxEbJH0l2TZvAGuiogt6fZFwK3A3sBDaQG4VtJxKeT1wBfaBeIG2swGWgAjvdxexGbg1CbrVwKfr7l/C3BLi3K/1GT9ud3G4gbazAZarxvofuIG2swGXFDp2YkW/cUNtJkNtABGGmfFmyLttbLT/Hq8Uen/yAbXmzkAeKvnOy3GMXWuH+NyTJ3px5iOjog5RTYwX0fG6fqrunV3xmefHu80u0FRSg86Ij7U6jFJK/vtwDmmzvVjXI6pM/0aU9FtBFCZoj1oD3GY2UALghEmJ5FE2dxAm9lAazoGPUVMRgO9bBL22Y5j6lw/xuWYOjMlY8pOs5uaPehSfiQ0M5so++rwWDh8Rd26JyoX+EdCM7PJNpV70KVNliRpsaSXU9qXZhkJZki6Oz3+pKTDy4ol7e8wSU9IekHS85L+uEmZUyRtrUlJc0WzbfU4rvUpPc7qZr9oK3N9Ok7PSjq+5HiOrnn+qyVtk3RJQ5kJOU5pztw3Ja2pWVcoHVFJMf21pJfS63OfpHkt6o77Wvc4po7SK7X7nPY4prtr4lkvaXWLul0dpwAq6WKVylS7aCUier4AQ8CrwJHAXsAzwDENZS4C/jHdXgLcXUYsNfs7GDg+3Z4DvNIkplOA75QZR5O41gMHjPP4mWSTrQg4CXhyAmMbIpuw/Gcn4zgBHweOB9bUrLsWuCzdvgz4apN6+wOvpb/7pdv7lRjTImA43f5qs5g6ea17HNOVwKUdvL7jfk57GVPD418DrujFcdpHH4kTpt9QtwAry35/TsRSVg/6RGBdRLwWETuBu8jSyNSqTSvzLeDUND1fKSJiU0SsSrffBV6kPtNBvzoLuD0yK4B5DdMWlulU4NWIaHXRUaki4gfAlobVudMRlRVTRDwaEZV0dwX18wGXrsVx6kQnn9Oex5Q+558G7uzJvsiGOGqXqaKsBrpVOpimZdKbeyswv6R46qThlI8BTzZ5+FckPSPpIUkfnYBwAnhU0tNqyPqQdHIsy7KE1h+iiT5OVUXSEU2Ez/HT6SUbtXute61deqXJOk6/DrwREWtbPN7VcQqCnYzWLUUUHUaTdLWk1yVtbyjf9bDuHjdhv6TZwL3AJRGxreHhVWT/zh8L/B1w/wSEdHJEHA+cAXxJ0scnYJ9tSdoL+CTwL00enozjtJvI/h/umwFHSZcDFeCOFkUm8rXuOr3SBDqH8XvPXR2nMWCnxuqWgopm9f7XtK7RBcDbEXEUcB3ZcNi4ymqgNwKH1dzflfalWRlJw8BcYHNJ8ZD2M52scb4jIr7d+HhEbIuI7en2g8B0SQeUGVNEVNPhvAncx+4vbCfHsgxnAKuiJk1P1WQcpxpvVId4NH46ogk9ZpL+EPgt4LPpi2M3HbzWPRMRb0TEaESMATe22NdkHKdh4HeBu1uV6fY4BcEHjNYtBRUaRosso/emJnW6HtYtq4F+Clgg6YjUE1tClkamVm1amU8B32v1xu6FdCBuBl6MiK+3KHNQ9YBJOpHs+JT2pSFplqQ51dtkPzY1psFZDpynzEnA1hYvfq+17OVM9HFqkDsdUVkBSVoM/BnwyYh4v0WZTl7rXsbUSXqlTj6nvXYa8FJEbGj2YJ7jNEawQ5W6paCyhtG6H9Yt69dHsrMPXiH7lfjytO4qsjcxwEyyf5/XAT8Cjizz11DgZLJ/h58FVqflTOCLZGlsAC4Gnif7NXsF8Kslx3Rk2tczab/V41Qbk4BvpOP4HLCwzJjSPmeRNbhza9ZN+HEi+4LYRDYf+wayfxHnk/3buRZ4HNg/lV0I3FRT93PpvbUO+KOSY1qXPnjV91X17KRDgAfHe61LjOmb6f3yLFmje3BjTOn+bp/TsmJK62+tvo9qyhY6TsDDwMqGZU3D/aUNdR5PZRqXs4B3Gsq+3WSflwJ/UXP/yzScNQNsb7i/Bji05v6rtDlbxVcSmpnVkPQycEr8NKv39yPi6IYy56QyX0j3/ymVu7OmzPaImF1z/xHgyoj4rzTU82PgQzFOI7zH/UhoZtZGWcNoXQ/ruoE2M6t3DfAJSWvJxsyvgSyrt6SbACLL4F3N6v0UNVm9JV0raQOwj6QNkq5M270ZmC9pHfAnNDk7pJGHOMzM+pR70GZmfcoNtJlZn3IDbWbWp9xAm5n1KTfQZmZ9yg20mVmfcgNtZtan3ECbmfWp/wcGomiqlwcaFwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "mappable = plt.imshow(order_param_input.detach().numpy(), cmap=\"plasma\", vmin=-0.01, vmax=0.01)\n",
    "plt.colorbar(mappable, shrink=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({'model': model.state_dict(), 'order_param_input': order_param_input},\n",
    "           'perovskite_order_parameters_determine_irreps.torch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.011, -0.007, -0.015],\n",
       "       [-0.015,  0.007,  0.011],\n",
       "       [ 0.015,  0.007, -0.011],\n",
       "       [-0.011, -0.007,  0.015]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "order_param_input.detach().numpy().round(3)[:, rs.dim(Rs_order_param[:2]): rs.dim(Rs_order_param[:2 + 1])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "order_param_input.to('cpu')\n",
    "features = create_features(order_param_input)\n",
    "data = dh.DataPeriodicNeighbors(features, Rs_in, pos, struct_221.lattice.matrix, r_max, y=relative_vecs_62)\n",
    "data.to(device)\n",
    "output = model(data.x, data.edge_index, data.edge_attr, n_norm=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Note, the network output does not perfect match the target output.\n",
    "This is because the octahedra in `structures/62_anon_sorted.cif` are slightly distorted. To capture this additional distortion, we would need to include order parameters for the $X$ sites as well as the $B$ sites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fc4240d0ad0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAADLCAYAAAC8opN3AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3de5hcVZnv8e+PQAjhEnJBcgG5SFDBwSDhNjNc5BqOSjgcQOQWHCAyjxxlkNEwiGAUAT0PnhnHI0SMBOUqXmgU5QmQIA6XSYMRDU5ME5HcQEkIAgFCut/zx17NVIrqququ2l29O7/P8+yna++99rtX7e6ufnuttddWRGBmZmbWKpu1ugJmZma2aXMyYmZmZi3lZMTMzMxaysmImZmZtZSTETMzM2spJyNmZmbWUk5GzMzMrKWcjJj1gqRnJB01GM8v6WxJvxqo8cxs8HIyYtaPJA1pdR3MzAYaJyNmdZL0PeCdwN2SXpH02bT9B5Kek/SSpF9K2rvkmBslfUvSPZJeBT4o6QOSfi3p5XTs7ZK+XHLMhyUtlLRW0sOS9ql2/gr1PE9Sh6Q1ktokjU/bd5UUkjYvKTtf0rmS3gtcBxycYq8tqf91kuam+j4oaZe+xjMzq8TJiFmdIuJM4FngIxGxTUR8Ne36OTAReAfwBHBz2aGnAVcC2wL/CfwYuBEYBdwK/M/ugpL2BWYDnwBGA9cDbZK2rHJ+So4/ArgKOAUYB/wJuK2O9/Z74HzgkRR7+5LdpwNfAsYACyu8v97GMzPbiJMRswZFxOyIeDki3gCuAN4vaURJkbsi4j8ioguYBGwO/FtEvBkRPyJLULpNB66PiMciojMi5gBvAAfVWZ3TgdkR8USqzyVkrRO7NvAWfxYRv0zxLk3xdm4gnpnZRpyMmDVA0hBJV0t6WtJfgWfSrjElxZaVvB4PrIiNn1BZun8X4DOpi2Zt6t7YOR1Xj/FkrSEARMQrwGpgQp3HV/JW/VK8Nb2oj5lZTU5GzHqn/DHXpwFTgaOAEcCuabt6OGYVMEFS6f7SVoZlwJURsX3JMjwibu3h/OVWkiU0WSWkrcm6e1YAr6bNw0vKj63y3t5WP0nbkHUvrWwgnpnZRpyMmPXO88DuJevbknWjrCb7o/yVGsc/AnQCF0jaXNJU4ICS/d8Gzpd0oDJbS/qQpG17OH+5W4GPS5okactUn8ci4pmI+AtZUnJGatH5B+BdZe9tJ0lDy2L+D0l/n7Z/CXg0IpY1EM/MbCNORsx65yrg86kL5WLgJrJukRXAU8Cj1Q6OiPXAicA5wFrgDOCnZAkNEdEOnAf8O/Ai0AGcXeX85fHvAy4DfkjWCvMu4NSSIucB/0yWPO0NPFyy7wFgEfCcpBdKtt8CXE7WPbNfqnMj8cysxSRNkbQ43Xk3o8L+iyQ9JelJSfd330WX9k2TtCQt05pSn427rs2sv0l6DLguIr7b6rqUk3QjsDwiPt/quphZc6T5jv4AHA0sBxYAH4uIp0rKfJCsVXWdpH8EDo+Ij0oaBbQDk8m6Yh8H9ouIFxupk1tGzPqZpMMkjU3dNNOAfYBftLpeZrbJOADoiIilqbX2NrKxb2+JiHkRsS6tPgrslF4fC8yNiDUpAZkLTGm0QpvXLmJmTfZu4A5ga2ApcFJErGptlcysKI46ZnisXt1ZtczCJ9YvAl4v2TQrImal1xPY+C6+5cCBVcKdQzafUk/HNnK3HuBkxKzfpQ+EWTULDgARcXar62BmG1u9uov5//HOqmW236rj9YiY3Oi5JJ1B1iVzWKOxqnE3jZmZWZEEqEtVlxpWsPGUAjulbRtR9lDOS4Hj06SHdR/bW05GzMzMCkSAOjerutSwAJgoabd06/2pQNtG58geTXE9WSLy55Jd9wLHSBopaSRwTNrWkKZ302w/dHiM32pE7YLWZ8PH5vfMsY4/jK1daADZY8/ncou97rl8HqdSxPvXav6f1UeLX87nIcb7jHsll7gAq58fmUvc0e9fmUtcgDc6RuUW22Dlay+xdv26vH5N3i5AXQ0cHrFB0gVkScQQskdILJI0E2iPiDbga8A2wA/SHI3PRsTxEbFG0pfIEhqAmRGxpoF3A+Rwa+9eI8bFTQc35bZj68EHLmmrXaiPTjziktxi5+FHD1yVW+wnrjo+l7gR/feZ1SxSPinUB+fn84/L8n95JJe4ADdde3Iucc98/vJc4gJ0TD0lt9gGZz0yh6deWtVvv9gfmDQsHnpg16plthm9+PFmjBnpLx7AamZmVjCNtIwMRE5GzMzMiiRAnUXs8O2ZkxEzM7MCEaANTkbMzMysVQJwN42ZmZm1kqpPwFo4TkbMzMyKpMFbewciJyNmZmYFM9gGsFadpk3S/pKelDRM0taSFkl6X39VzszMzMp0jxmpthRM1ZaRiFggqQ34MrAV8P2I+F2/1MzMzMwq2hTHjMwkm/b1deBTlQpImg5MBxg7bLumVc7MzMw2pgB1Fm8m52rqeVDeaLL56bcFhlUqEBGzImJyREweOXR4M+tnZmZm5QZZN009ycj1wGXAzcA1+VbHzMzMqgqgs8ZSg6QpkhZL6pA0o8L+QyU9IWmDpJPK9nVKWpiWpjwsrWo3jaSzgDcj4hZJQ4CHJR0REQ804+RmZmbWe+rqezdN+nv+TeBoYDmwQFJbRDxVUuxZ4Gzg4gohXouISX2uQAW1BrDeBNyUXncCBzbz5GZmZtZLATQ2ZuQAoCMilgJIug2YCryVjETEM2lfv3T61NNNY2ZmZgNJY2NGJgDLStaXp231GiapXdKjkk7oxXE98qRnZmZmRRKqp2VkjKT2kvVZETGrSTXYJSJWSNodeEDSbyPi6UYCOhkxMzMrkvpu7X0hIib3sG8FsHPJ+k5pW32nj1iRvi6VNB/YF2goGXE3jZmZWdF0qfpS3QJgoqTdJA0FTgXquitG0khJW6bXY4C/o2SsSV85GTEzMyuS7gGs1ZZqh0dsAC4A7gV+D9wREYskzZR0PLz1OJjlwMnA9ZIWpcPfC7RL+g0wD7i67C6cPnE3jZmZWdFEYzOwRsQ9wD1l275Q8noBWfdN+XEPA3/T0MkrcDJSQE9cdXxusT9/9KLahQaQPK/FBy5pylw+b3PiEZfkEjdPP3rgqlzifuXRfK5FfC6/R2j9y1cOziXuWdf4GaRWr7oGsBaKkxEzM7MiCYgGJj0biJyMmJmZFY1bRszMzKxlAugcXPefOBkxMzMrFDU8gHWgcTJiZmZWJI0/m2bAcTJiZmZWMB7AamZmZq3jMSNmZmbWWh4zYmZmZi0UATHIxoxUbedJ89RfWLJ+paRP518tMzMz61HnZtWXgqlV49nAWQCSNiN7st/3ywtJmi6pXVL7i+vXNb+WZmZmlgmIUNWlFklTJC2W1CFpRoX9h0p6QtIGSSeV7ZsmaUlapjXjLVXtpomIZyStlrQvsCPw64hYXaHcLGAWwF4jxkUzKmZmZmaVqKHWD0lDgG8CRwPLgQWS2sqevvsscDZwcdmxo4DLgclkQ2kfT8e+2OcKUd+YkRtShcaStZSYmZlZC9XT+lHFAUBHRCwFkHQbMBV4KxmJiGfSvq6yY48F5kbEmrR/LjAFuLWRCtWTWv04nWh/4N5GTmZmZmYN6p70rNpS3QRgWcn68rStHo0c26OaLSMRsV7SPGBtRHQ2ekIzMzNrTB2Tno2R1F6yPisNqRiQaiYjaeDqQcDJ+VfHzMzMqgoRXTU7Nl6IiMk97FsB7FyyvlPaVo8VwOFlx86v89ge1bq1dy+gA7g/IpY0ejIzMzNrgsZu7V0ATJS0m6ShZHfKttV55nuBYySNlDQSOIYmDOGodTfNU8DujZ7EzMzMmiMiW/p+fGyQdAFZEjEEmB0RiyTNBNojok3S/mRjRkcCH5H0xYjYOyLWSPoSWUIDMLN7MGsjPAOrmZlZoYhocGKziLgHuKds2xdKXi8g64KpdOxsmnx3rZMRMzOzovGzaczMzKxlgoZbRgYaJyNmZmYF0+CkZwOOk5EC2uOuO3KLvfOI83KJu+ylb+cSN6/6AjD/zFzCLlt3eS5x87wWfz3o9VziHnzQolzidkw9JZe4APMOfymXuEse3iuXuDYIhQbdU3udjJiZmRVI4G4aMzMza6VwN42ZmZm1VF0zsBaKkxEzM7Oiqf1smkJxMmJmZlYgEdDlZMTMzMxaymNGzMzMrGVCdPluGjMzM2ulwTaAtea7kXS+pIVp+aOkef1RMTMzM6ssuqovtUiaImmxpA5JMyrs31LS7Wn/Y5J2Tdt3lfRaSV5wXTPeT82WkYi4DrhO0hbAA8C1zTixmZmZ9V5EY7f2ShoCfBM4GlgOLJDUFhFPlRQ7B3gxIvaQdCpwDfDRtO/piJjU5wpU0Jt386/AAxFxd/kOSdMltUtqf3H9uubVzszMzN4mQlWXGg4AOiJiaUSsB24DppaVmQrMSa/vBI6UlNuo2bqSEUlnA7sAX6y0PyJmRcTkiJg8cujwJlbPzMzMNpKe2lttAcZ0NxKkZXpJhAnAspL15WkblcpExAbgJWB02rebpF9LelDSIc14SzW7aSTtB1wMHBJRT0+UmZmZ5amO1o8XImJyDqdeBbwzIlan/OAnkvaOiL82ErSelpELgFHAvDRY5YZGTmhmZmaNENFVfalhBbBzyfpOaVvFMpI2B0YAqyPijYhYDRARjwNPA3s2+o7qGcD68UZPYmZmZs0RNPygvAXAREm7kSUdpwKnlZVpA6YBjwAnkY0ZDUk7AGsiolPS7sBEYGkjlQHPM2JmZlYsacxInw+P2CDpAuBeYAgwOyIWSZoJtEdEG/Ad4HuSOoA1ZAkLwKHATElvAl3A+RGxpoF3AzgZMTMzKxjR1eCkZxFxD3BP2bYvlLx+HTi5wnE/BH7Y0MkrcDJiZmZWJJE9LG8wcTJiZmZWIEFj3TQDkZMRMzOzgmlwAOuA42TEzMysSKLxMSMDjaLJHU+S/gL8qc7iY4AXmlqB/GMXLW6esR03/9hFi5tn7KLFzTN20eLmGXsgxN0lInbIoQ4V/c3od8SPjjupapk9b/7W4zlNepaLpreM9OYbIqk9r4uVV+yixc0ztuPmH7tocfOMXbS4ecYuWtw8YxctbjNEQHS6m8bMzMxaaLB10zgZMTMzK5S6nsxbKK1ORmYVMHbR4uYZ23Hzj120uHnGLlrcPGMXLW6esYsWt3EBMchaRpo+gNXMzMzys/fIHeOOI8ofJbOx9/3o/27aA1jNrP9ICmBiRHS0ui5m1k8CugbZANbB1c5jNgBIekbSUa2uR7NIOlzS8oEaz2zTk40ZqbbUjCBNkbRYUoekGRX2bynp9rT/MUm7luy7JG1fLOnYZrwjJyNmBSTJrZpmm6ggu5um2lKNpCHAN4HjgL2Aj0naq6zYOcCLEbEH8HXgmnTsXmRP8N0bmAL8vxSvIU5GzCpIrRsXS3pS0kvpP4RhJfs/LGmhpLWSHpa0T9r+PeCdwN2SXpH0WUlzJH0m7Z8gKSR9Mq2/S9IaSZul9fPSfxxrJLVJGl9yzpD0SUlLgCUV6vz3kpZJOryH93S8pEWpzvMlvbcs9h4l6zdK+rKkrYGfA+PT+3lF0nhJV0i6M12XlyU9Ien9fY3Xu++O2SYuGktGgAOAjohYGhHrgduAqWVlpgJz0us7gSMlKW2/LSLeiIg/Ah0pXkOcjJj17BSyzH83YB/gbABJ+wKzgU8Ao4HrgTZJW0bEmcCzwEciYpuI+CrwIHB4inkYsBQ4tGT9oYjoknQEcFU67ziymYxvK6vTCcCBZP/NvEXSFOBW4H9FxPzyNyJpz7T/QmAHskeH3y1paLULEBGvkv33tDK9n20iYmXaPRX4ATAKuAX4iaQtGohnZnVpuJtmArCsZH152laxTERsAF4i+7yr59heczJi1rN/i4iVEbEGuBuYlLZPB66PiMciojMi5gBvAAf1EOdB4O9T68ehwFeBv0v7Dkv7AU4HZkfEExHxBnAJcHBpXy1wVUSsiYjXSradTJYQHRcR/9lDHT4K/Cwi5kbEm8D/AbYC/rb2ZejR4xFxZ4p3LTCMnq+BmTVRdG5WdQHGSGovWaa3us7VOBkx69lzJa/XAduk17sAn0ndHWslrQV2Bip2N0TE08CrZMnMIcBPgZWS3s3Gych4Sp7rFBGvAKvZ+L+O0v9Iul0I3BERv6vyXspjd6VYjfxH81ZdUrzl9HANzKx5IqinZeSFiJhcspTOm7KC7DOr205pG5XKpDFqI8g+j+o5ttecjJj13jLgyojYvmQZHhG3pv2VJu95EDgJGBoRK9L6NGAksDCVWUmW6ACQxleMZuNf9EqxTwZOkPTpKnUujy2yD5Tu2OuA4SXlx9Y4J5R8IKVWn53Sefoaz8zq1NWlqksNC4CJknZLXbWnAm1lZdrIPqMg++x6ILKJydqAU9PdNrsBE4GeWmTr5mTErPe+DZwv6UBltpb0IUnbpv3PA7uXHfMgcAHwy7Q+P63/KiI607ZbgY9LmiRpS+ArwGMR8UyN+qwEjgQ+LekfeyhzB/AhSUemcR2fIetaejjtXwicJmlIGn9yWMmxzwOjJY0oi7mfpBPTf00XpniPNhDPzOrS2JiRNAbkAuBe4PdkLauLJM2UdHwq9h2y39MO4CJgRjp2EdnnyVPAL4BPlnyG9ZmTEbNeioh24Dzg34EXyUaTn11S5Crg86kL5+K07UFgW/47GfkVWctB9zoRcR9wGfBDYBXwLrL/WOqp07NkCckMSedW2L8YOAP4Btlj0T9CNsh2fSry6bRtLdnYlZ+UHPtfZInS0vSeurti7iIbi/IicCZwYho/0td4ZlaHaPxuGiLinojYMyLeFRFXpm1fiIi29Pr1iDg5IvaIiAMiYmnJsVem494dET9vxnvydPBm1muSrgD2iIgzWl0Xs03Ne7YdHzfs+4mqZQ556ApPB29mZmZ58VN7zczMrIW6u2kGEycjZtZrEXFFq+tgtilzy4iZmZm1kNwyYmZmZi0UELXnEimUpicj2w8dHuO3av70AUW85yevH5XhO67NKTL8ccm4XOLm9f3bfeKqnCLDq89vn0vcvH4u8vwdyavOi19u+GGfFb1/nxdyiQuw6Mkdcom7175/ziUuwBsdo3KJm9fPXNH+zK587SXWrl/Xb9XOntpbtKtUXdOTkfFbjeCmg6fVLthLRewfk/L5Vf3AP/0sl7gAZxw3I5e4XblEhVu+cXVOkeHxaz+cS9y8fi7y/B3Jq84fnJ/PvGcPzb0xl7gAe084P5e48x7+Ri5xATqmnpJL3Lx+5vL6ecvLWY/MqV2omQI63U1jZmZmrRK+tdfMzMxazQNYzczMrHU8gNXMzMxaKQbhrb2D692YmZltAhp5am8tkkZJmitpSfo6sody01KZJZKmlWyfL2mxpIVpeUetczoZMTMzK5J0N021pUEzgPsjYiJwf1rfiKRRwOXAgcABwOVlScvpETEpLTXvW3cyYmZmViABdHWq6tKgqUD3/cpzgBMqlDkWmBsRayLiRWAuMKWvJ6yajEjaX9KTkoZJ2lrSIknv6+vJzMzMrFHVu2hSN80YSe0ly/RenGDHiOieUfI5YMcKZSYAy0rWl6dt3b6bumguk1QzO6o6gDUiFkhqA74MbAV8PyJ+VyuomZmZ5SPqm/TshYiY3NNOSfcBYyvsunTjc0Wo97PQnR4RKyRtC/wQOBO4qdoB9dxNMxNYALwOfKpSgZRxTQcYO2y73lTYzMzMeqnRQaoRcVRP+yQ9L2lcRKySNA6oNOZjBXB4yfpOwPwUe0X6+rKkW8jGlFRNRuoZMzIa2AbYFhhWqUBEzIqIyRExeeTQ4XWENDMzs75R3gNY24Duu2OmAXdVKHMvcIykkWng6jHAvZI2lzQGQNIWwIeBmj0q9dT4euAy4GbgmjrKm5mZWU4iai8Nuho4WtIS4Ki0jqTJkm7I6hBrgC+R9ZwsAGambVuSJSVPAgvJWlC+XeuEVbtpJJ0FvBkRt0gaAjws6YiIeKCv79DMzMwa09WZ382wEbEaOLLC9nbg3JL12cDssjKvAvv19py1BrDeROrniYhOsvuJzczMrEWyAayeDt7MzMxayE/tNTMzs5YJ5JYRMzMzay23jJiZmVnrBHQ2PuX7gOJkxMzMrEACt4yYmZlZK4XHjNRUxIyt99Put9bj1344t9gXHfPbXOLm9TOR57XY76Kf5hL3tOPe9jTuAe+Wn1+dS9zRD/xDLnHf3O7NXOICrN7s9Vzibv7KFrnELaKi/Q3p778ggW/tNTMzsxYrWsJWi5MRMzOzAskmPWt1LZrLyYiZmVnBdA2ylpH8Jrc3MzOzXHR1VV8aIWmUpLmSlqSvI3so9wtJayX9tGz7bpIek9Qh6XZJQ2ud08mImZlZgQRZy0i1pUEzgPsjYiJwf1qv5GvAmRW2XwN8PSL2AF4Ezql1QicjZmZmRRLQWWNp0FRgTno9BzihYjUi7gdeLt0mScARwJ21ji/lMSNmZmYFEojOfMeM7BgRq9Lr54Ade3HsaGBtRGxI68uBCbUOcjJiZmZWMHU0foyR1F6yPisiZnWvSLoPGFvhuEs3Ok9EqB8m43IyYmZmViBBXV0xL0TE5B5jRBzV0z5Jz0saFxGrJI0D/tyL6q0Gtpe0eWod2QlYUeugqmNGJM2UdGHJ+pWSPt2LSpmZmVmTdUX1pUFtwLT0ehpwV70HRkQA84CTenN8rQGss4GzACRtBpwKfL+8kKTpktolta9dv67eOpuZmVkvBdBZY2nQ1cDRkpYAR6V1JE2WdEN3IUkPAT8AjpS0XNKxadfngIskdZCNIflOrRNW7aaJiGckrZa0L9kAll9HxOoK5WYBswDeO2JcsR70YmZmVjB5/qFNf+ePrLC9HTi3ZP2QHo5fChzQm3PWM2bkBuBssoEus3sT3MzMzJqrzjEjhVLPPCM/BqYA+wP35lsdMzMzqyXnbpp+V7NlJCLWS5pHdt9wEd+jmZnZoBHAIHtOXu1kJA1cPQg4Of/qmJmZWS2DrWWg1q29ewEdZHPUL+mfKpmZmVlPultGqi1FU+tumqeA3fupLmZmZlZT0Jnr/TT9zzOwmpmZFczgSkWcjJiZmRVKAJ21HhdTsGzFyYiZmVmBBLChZqlicTJiZmZWMFG0po8aCpOM9MMTjJsuQrnEndh2ey5xAXYecV4ucZe99O1c4uZVXwAePC2XsMteuyyXuHlei78e/mouce884vlc4i45/qO5xAWYd/hLucT9/Wkn5hK3iIr2eZ/PJ33PsmfTFOsa1VLPDKxmZmY2gOR5a6+kUZLmSlqSvo7sodwvJK2V9NOy7TdK+qOkhWmZVOucTkbMzMwKpHsAa7WlQTPI5hebCNyf1iv5GnBmD/v+OSImpWVhrRM6GTEzMyuYnCc9mwrMSa/nACdUKhQR9wMvN346JyNmZmaFEmnSs2pLg3aMiFXp9XPAjn2IcaWkJyV9XdKWtQoXZgCrmZmZZeroihkjqb1kfVZEzOpekXQfMLbCcZeWrkREqPcjii8hS2KGArOAzwEzqx3gZMTMzKxAsmfT1MwPXoiIyT3GiDiqp32Snpc0LiJWSRoH/LlX9fvvVpU3JH0XuLjWMe6mMTMzK5icu2nagGnp9TTgrt4cnBIYJIlsvMnvah3jZMTMzKxAgqCrxtKgq4GjJS0BjkrrSJos6YbuQpIeAn4AHClpuaRj066bJf0W+C0wBvhyrRPW7KaRdD5wflodATwTER+s/z2ZmZlZMzXh9t0eRcRq4MgK29uBc0vWD+nh+CN6e86aLSMRcV1ETAL2B5YD1/b2JGZmZtYc3WNGcmwZ6Xe9GcD6r8ADEXF3+Q5J04HpAGOHbdekqpmZmVklm+R08JLOBnYBvlhpf0TMiojJETF5+6HDm1g9MzMzK7VJtoxI2o/stpxDIqIJE7uZmZlZ3zXljpkBpZ5umguAUcC87C4d2iPi3OqHmJmZWR5CsEGDq22gZjISER/vj4qYmZlZbVk3zeDiGVjNzMwKpnOQpSNORszMzAokCjpItRonI2ZmZgWzwS0jZmZm1ip1PiivUJyMmJmZFUgAG3KcDr4VnIyYmZkVSBC8Oci6aRTR3OxK0l+AP9VZfAzwQlMrkH/sosXNM7bj5h+7aHHzjF20uHnGLlrcPGMPhLi7RMQOOdShoqGb7Rzv2OKfqpZZsf4zj0fE5L7ElzQKuB3YFXgGOCUiXiwrMwn4FrAd0AlcGRG3p327AbcBo4HHgTMjYn21cza9ZaQ33xBJ7X29WK2KXbS4ecZ23PxjFy1unrGLFjfP2EWLm2fsosVthiB4U515nmIGcH9EXC1pRlr/XFmZdcBZEbFE0njgcUn3RsRa4Brg6xFxm6TrgHPIEpce1fVsGjMzMxsYguxBedWWBk0F5qTXc4AT3laHiD9ExJL0eiXwZ2AHZVO1HwHcWe34ch4zYmZmViCdseLe1W98bkyNYsMktZesz4qIWXWeYseIWJVePwfsWK2wpAOAocDTZF0zayNiQ9q9HJhQ64StTkbqvTADKXbR4uYZ23Hzj120uHnGLlrcPGMXLW6esYsWt2ERMaXRGJLuA8ZW2HVp2blC6vnWHUnjgO8B0yKiKz3Drvf1afYAVjMzMysuSYuBwyNiVUo25kfEuyuU2w6YD3wlIu5M2wT8BRgbERskHQxcERHHVjunx4yYmZlZqTZgWno9DbirvICkocCPgZu6ExHIWlKAecBJ1Y4v17JkRNIJkkLSe4oQN8/YzY4raWdJf0y3ZyFpZFrfdaDGLlrcItbZ1yL/uAWuc6ekhZJ+I+kJSX/baMw84+YZO886F8jVwNGSlgBHpXUkTZZ0QypzCnAocHa6XguV3e4L2Z03F0nqIBtD8p2aZ4yIlixk9zA/BHyxCHGLVmfgs2QDlgCuBy4Z6LGLFreIdfa18LXoIe4rJa+PBR4cyHGLWmcvVa57S04K2wArgD2BxQM9bhHrDGwBPAlcCCwCthjosYsWt4h19rXwteghbukf4JOBnwzkuEWts5cq170lJ4XTge+k1w8D+w3kuAWu87Fkt6QfncP3MJfYRYtbxDr7WvhaVIjZCSwE/gt4qYmfb7nELWqdvfS8tLlM6E0AAAFQSURBVGrMyMfIpoolff3YAI+bZ+w863wcsAp4XxNj5h27aHHzjF20uHnGLlrcPGPnEfe1iJgUEe8BpgA3qa/3aPZP3Dxj51ln60l/Zz/AKLJpZP9ENuf9MuBZ0m3GAy1uges8iawZ950p5rgmfg9ziV20uEWss6+Fr0UPcV8pW38eeMdAjVvUOnupct37/YQwHbi+bNuDwKEDMW4R6wwIeITUjAv8b+DmJn3/coldtLhFrLOvha9Fldil4yTeQ/aAuCEDNW5R6+ylynXv9xNm9x9PKdv2KeBbAzFuEetMluTcXrI+BHgCOKwJ1yKX2EWLW8Q6+1r4WlSJ3T1OYiHwG+BDjcbMM25R6+yl58UzsJqZmVlLeQZWMzMzayknI2ZmZtZSTkbMzMyspZyMmJmZWUs5GTEzM7OWcjJiZmZmLeVkxMzMzFrKyYiZmZm11P8HaOo28B3Wq5EAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 720x216 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(2, 1, figsize=(10, 3))\n",
    "bound = 0.2\n",
    "mappable = ax[0].imshow(data.y.cpu().detach().numpy().T, cmap=\"plasma\", vmax=bound, vmin=-bound)\n",
    "ax[1].imshow(output[:, [2, 0, 1]].cpu().detach().numpy().T, cmap=\"plasma\", vmax=bound, vmin=-bound)\n",
    "sites = [str(s.species)[0] for s in struct_62]\n",
    "\n",
    "ax[0].set_title('target output')\n",
    "ax[1].set_title('network output')\n",
    "\n",
    "ax[0].set_xticks(range(len(sites)))\n",
    "ax[0].set_xticklabels([''] * len(sites))\n",
    "ax[1].set_xticks(range(len(sites)))\n",
    "ax[1].set_xticklabels(sites)\n",
    "\n",
    "ax[0].set_yticks(range(3))\n",
    "ax[1].set_yticks(range(3))\n",
    "ax[0].set_yticklabels(['x', 'y', 'z'])\n",
    "ax[1].set_yticklabels(['x', 'y', 'z'])\n",
    "fig.colorbar(mappable, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_struct = Structure(new_lattice, struct_221.species,\n",
    "                       struct_221.cart_coords + output.cpu()[:, [2, 0, 1]].detach().numpy(),\n",
    "                       coords_are_cartesian=True)\n",
    "new_struct.to(filename='recovered_struct.cif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "62"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pga = SpacegroupAnalyzer(new_struct)\n",
    "pga.get_space_group_number()"
   ]
  }
 ],
 "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
}