perovskite_order_parameters_spacegroup_74_from_62.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recovering output with intermediate symmetry by constraining order parameters\n",
    "In this example, we will show that by constraining the learnable order parameter, we can recover output with symmetry that is intermediate between input and target output symmetries."
   ]
  },
  {
   "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",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "device = \"cuda:1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymatgen\n",
    "from pymatgen.core.structure import Structure\n",
    "from pymatgen.symmetry.analyzer import SpacegroupAnalyzer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Same set up as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "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": [
    "### This time we only use a pseudovector order parameter according to our previous notebook's findings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rs_order_param = [(1, 1, 1)]  # 2020.08.25 ONLY NEED PSEUDOVECTOR (per atom)\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(2, rs.dim(Rs_order_param) - 1) # Remove z and same across pairs\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": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_features(order_param_input):\n",
    "    # Modify for z\n",
    "    new_order_param = torch.stack([order_param_input[:, 0].repeat(2),\n",
    "                                   torch.zeros(4), \n",
    "                                   order_param_input[:, 1].repeat(2)], dim=-1)\n",
    "    order_param = torch.cat([zeros, new_order_param], 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": "code",
   "execution_count": 7,
   "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": "code",
   "execution_count": 8,
   "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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = torch.optim.Adam(model.parameters(), 1e-2)\n",
    "order_opt = torch.optim.Adam([order_param_input], 1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('spacegroup_62_model_overfit.torch')['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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",
    "#            'spacegroup_62_model_overfit.torch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model update:  tensor(0.0676, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0676, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0423, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0379, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0424, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0294, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0449, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0476, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0318, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0353, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0320, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0299, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0332, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0335, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0295, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0303, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0298, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0298, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0317, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0312, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0302, device='cuda:1', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0312, grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "max_iter = 101\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "for i in range(max_iter):\n",
    "    # Update model\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",
    "    # Update input\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() + 1e-2 * order_param_input.abs().mean()\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": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({'model': model.state_dict(), 'order_param_input': order_param_input},\n",
    "           'spacegroup_74_from_62_model_1_e.torch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "order_param_input.to('cpu')\n",
    "features = create_features(order_param_input)\n",
    "# print(features)\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": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7efda501b750>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAADLCAYAAAC8opN3AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3de7hdVX3u8e8LGO6GhCgJFwEFtWhtkIjSFkSu8VgJxwOIFwgeJaVPOdWjtEK9owjYPnraU48QMTV44SJeiJbKg0CoLUKzwZQa2jRbRJIQUBJCQYSQvd/zx5yhK9u111o7a8299ty8n+eZT+ZlzN8ca+5k5bfHGHNM2SYiIiKiX7brdwUiIiLiuS3JSERERPRVkpGIiIjoqyQjERER0VdJRiIiIqKvkoxEREREXyUZiYiIiL5KMhIxBpLul3TcZLy+pLMk/eNEjRcRk1eSkYhxJGn7ftchImKiSTIS0SFJXwFeBHxX0hOS/qzc/w1JD0l6TNI/SHpFwzlflvQFSTdI+hXwBkmvlvRjSY+X514j6VMN5/yBpOWSNkq6XdKrWl2/ST3PljQoaYOkJZL2LvcfIMmSdmgou1TSeyT9FnAZcEQZe2ND/S+TdFNZ39sk7b+t8SIimkkyEtEh22cADwBvtr2b7c+Uh/4eOBh4IXA38LURp74duAjYHfhn4NvAl4HpwFXAf99SUNKhwCLgD4E9gcuBJZJ2bHF9Gs4/BrgYOA2YBfwcuLqDz/ZvwDnAj8rYezQcfgfwSWAGsLzJ5xtrvIiIrSQZieiS7UW2H7f9NPBx4HckTW0ocr3tf7I9DMwGdgD+2vYztr9FkaBssQC43PadtodsLwaeBl7XYXXeASyyfXdZnwsoWicO6OIj/p3tfyjjfaiMt18X8SIitpJkJKILkraXdImkn0r6T+D+8tCMhmKrG9b3BtZ66zdUNh7fH/hA2UWzseze2K88rxN7U7SGAGD7CWA9sE+H5zfzbP3KeBvGUJ+IiLaSjESMzcjXXL8dmAccB0wFDij3a5Rz1gH7SGo83tjKsBq4yPYeDcsutq8a5fojPUiR0BSVkHal6O5ZC/yq3L1LQ/mZLT7bb9RP0m4U3UsPdhEvImIrSUYixuZh4MUN27tTdKOsp/hP+dNtzv8RMAScK2kHSfOAwxuOfxE4R9JrVdhV0psk7T7K9Ue6CniXpNmSdizrc6ft+23/kiIpeWfZovM/gZeM+Gz7SpoyIuZ/k/T75f5PAnfYXt1FvIiIrSQZiRibi4EPl10o5wFXUnSLrAXuBe5odbLtTcBbgHcDG4F3At+jSGiwPQCcDfwN8CgwCJzV4voj4/8A+AjwTYpWmJcApzcUORv4U4rk6RXA7Q3HbgFWAA9JeqRh/9eBj1F0zxxW1rmbeBHRZ5LmSlpZPnl3fpPj75d0r6R7JN285Sm68th8SavKZX5P6rN113VEjDdJdwKX2f7bftdlJElfBtbY/nC/6xIRvVHOd/QfwPHAGmAZ8Dbb9zaUeQNFq+qTkv4IONr2WyVNBwaAORRdsXcBh9l+tJs6pWUkYpxJer2kmWU3zXzgVcD3+12viHjOOBwYtH1f2Vp7NcXYt2fZvtX2k+XmHcC+5fqJwE22N5QJyE3A3G4rtEP7IhHRYy8DrgV2Be4DTrG9rr9Vioi6OO6EXbx+/VDLMsvv3rQCeKph10LbC8v1fdj6Kb41wGtbhHs3xXxKo53bzdN6QJKRiHFXfiEsbFtwArB9Vr/rEBFbW79+mKX/9KKWZfbYefAp23O6vZakd1J0yby+21itpJsmIiKiTgwaVsuljbVsPaXAvuW+rah4KeeHgJPKSQ87PneskoxERETUiAANbddyaWMZcLCkA8tH708Hlmx1jeLVFJdTJCK/aDh0I3CCpGmSpgEnlPu60vNumj2m7OK9d57avmBss11mVvfOscH/mNm+0ARy0Esfqiz2kw9V8zqVOj6/1vb3rG208vFqXmL8qllPVBIXYP3D0yqJu+fvPFhJXICnB6dXFjvgwV8/xsZNT1b1z+Q3GTTcxen2ZknnUiQR21O8QmKFpAuBAdtLgL8AdgO+Uc7R+IDtk2xvkPRJioQG4ELbG7r4NEAFj/YeMnWWrzyiJ48dxyhefcGS9oW20VuOuaCy2FX41i0XVxb77otPqiSuPX7fWb0iVZNCvWFpNb+4rPnzH1USF+DKz55aSdwzHv5YJXEBBuedVlnsgDN/tJh7H1s3bv+wXz17J//wlgNaltltz5V39WLMyHjJANaIiIia6aZlZCJKMhIREVEnBg3VscN3dElGIiIiakSANicZiYiIiH4xkG6aiIiI6Ce1noC1dpKMRERE1EmXj/ZORElGIiIiamayDWBtOU2bpNdIukfSTpJ2lbRC0ivHq3IRERExwpYxI62WmmnZMmJ7maQlwKeAnYGv2v7JuNQsIiIimnoujhm5kGLa16eAP2lWQNICYAHAzJ2e37PKRURExNZk0FD9ZnJupZMX5e1JMT/97sBOzQrYXmh7ju0506bs0sv6RURExEiTrJumk2TkcuAjwNeAS6utTkRERLRkYKjN0oakuZJWShqUdH6T40dJulvSZkmnjDg2JGl5ufTkZWktu2kknQk8Y/vrkrYHbpd0jO1benHxiIiIGDsNb3s3Tfn/+eeB44E1wDJJS2zf21DsAeAs4LwmIX5te/Y2V6CJdgNYrwSuLNeHgNf28uIRERExRga6GzNyODBo+z4ASVcD84BnkxHb95fHxqXTp5NumoiIiJhIuhszsg+wumF7TbmvUztJGpB0h6STx3DeqDLpWURERJ1YnbSMzJA00LC90PbCHtVgf9trJb0YuEXSv9r+aTcBk4xERETUSWeP9j5ie84ox9YC+zVs71vu6+zy9tryz/skLQUOBbpKRtJNExERUTfDar20tgw4WNKBkqYApwMdPRUjaZqkHcv1GcDv0TDWZFslGYmIiKiTLQNYWy2tTrc3A+cCNwL/Blxre4WkCyWdBM++DmYNcCpwuaQV5em/BQxI+hfgVuCSEU/hbJN000RERNSNu5uB1fYNwA0j9n20YX0ZRffNyPNuB367q4s3kWSkhu6++KTKYn/4+BXtC00gVd6LV1/Qk7l8fsNbjrmgkrhV+tYtF1cS99N3VHMv/MHqXqH1558+opK4Z16ad5BGpzoawForSUYiIiLqxOAuJj2biJKMRERE1E1aRiIiIqJvDAxNrudPkoxERETUiroewDrRJBmJiIiok+7fTTPhJBmJiIiomQxgjYiIiP7JmJGIiIjor4wZiYiIiD6ywZNszEjLdp5ynvr3NWxfJOm91VcrIiIiRjW0XeulZtrVeBFwJoCk7Sje7PfVkYUkLZA0IGng0U1P9r6WERERUTDYarm0I2mupJWSBiWd3+T4UZLulrRZ0ikjjs2XtKpc5vfiI7XsprF9v6T1kg4F9gJ+bHt9k3ILgYUAh0yd5V5ULCIiIppRV60fkrYHPg8cD6wBlklaMuLtuw8AZwHnjTh3OvAxYA7FUNq7ynMf3eYK0dmYkSvKCs2kaCmJiIiIPuqk9aOFw4FB2/cBSLoamAc8m4zYvr88Njzi3BOBm2xvKI/fBMwFruqmQp2kVt8uL/Qa4MZuLhYRERFd2jLpWaultX2A1Q3ba8p9nejm3FG1bRmxvUnSrcBG20PdXjAiIiK608GkZzMkDTRsLyyHVExIbZORcuDq64BTq69OREREtGTh4bYdG4/YnjPKsbXAfg3b+5b7OrEWOHrEuUs7PHdU7R7tPQQYBG62varbi0VEREQPdPdo7zLgYEkHSppC8aTskg6vfCNwgqRpkqYBJ9CDIRztnqa5F3hxtxeJiIiI3rCLZdvP92ZJ51IkEdsDi2yvkHQhMGB7iaTXUIwZnQa8WdInbL/C9gZJn6RIaAAu3DKYtRuZgTUiIqJWhLuc2Mz2DcANI/Z9tGF9GUUXTLNzF9Hjp2uTjERERNRN3k0TERERfWO6bhmZaJKMRERE1EyXk55NOElGauig66+tLPZ+U8+uJO7qx75YSdyq6gvA0jMqCbv6yY9VErfKe/Gfr3uqkrhHvG5FJXEH551WSVyAW49+rJK4q24/pJK4MQlZk+6tvUlGIiIiasSkmyYiIiL6yemmiYiIiL7qaAbWWkkyEhERUTft301TK0lGIiIiasSG4SQjERER0VcZMxIRERF9YzGcp2kiIiKinybbANa2n0bSOZKWl8vPJN06HhWLiIiI5jzcemlH0lxJKyUNSjq/yfEdJV1THr9T0gHl/gMk/bohL7isF5+nbcuI7cuAyyQ9D7gF+GwvLhwRERFjZ3f3aK+k7YHPA8cDa4BlkpbYvreh2LuBR20fJOl04FLgreWxn9qevc0VaGIsn+avgFtsf3fkAUkLJA1IGnh005O9q11ERET8BlstlzYOBwZt32d7E3A1MG9EmXnA4nL9OuBYSZWNmu0oGZF0FrA/8Ilmx20vtD3H9pxpU3bpYfUiIiJiK+Vbe1stwIwtjQTlsqAhwj7A6obtNeU+mpWxvRl4DNizPHagpB9Luk3Skb34SG27aSQdBpwHHGl30hMVERERVeqg9eMR23MquPQ64EW215f5wXckvcL2f3YTtJOWkXOB6cCt5WCVK7q5YERERHRDeLj10sZaYL+G7X3LfU3LSNoBmAqst/207fUAtu8Cfgq8tNtP1MkA1nd1e5GIiIjoDdP1i/KWAQdLOpAi6TgdePuIMkuA+cCPgFMoxoxa0guADbaHJL0YOBi4r5vKQOYZiYiIqJdyzMg2n25vlnQucCOwPbDI9gpJFwIDtpcAXwK+ImkQ2ECRsAAcBVwo6RlgGDjH9oYuPg2QZCQiIqJmxHCXk57ZvgG4YcS+jzasPwWc2uS8bwLf7OriTSQZiYiIqBMXL8ubTJKMRERE1IjprptmIkoyEhERUTNdDmCdcJKMRERE1Im7HzMy0cg97niS9Evg5x0WnwE80tMKVB+7bnGrjJ241ceuW9wqY9ctbpWx6xa3ytgTIe7+tl9QQR2a+u09X+hvvfGUlmVe+rUv3FXRpGeV6HnLyFh+IJIGqrpZVcWuW9wqYydu9bHrFrfK2HWLW2XsusWtMnbd4vaCDR5KN01ERET00WTrpkkyEhERUSsdvZm3VvqdjCysYey6xa0yduJWH7tucauMXbe4VcauW9wqY9ctbvcMnmQtIz0fwBoRERHVecW0vXztMSNfJbO1V37r/zy3B7BGxPiRZOBg24P9rktEjBPD8CQbwDq52nkiJgBJ90s6rt/16BVJR0taM1HjRTz3FGNGWi1tI0hzJa2UNCjp/CbHd5R0TXn8TkkHNBy7oNy/UtKJvfhESUYiakhSWjUjnqNM8TRNq6UVSdsDnwfeCBwCvE3SISOKvRt41PZBwOeAS8tzD6F4g+8rgLnA/yvjdSXJSEQTZevGeZLukfRY+RvCTg3H/0DSckkbJd0u6VXl/q8ALwK+K+kJSX8mabGkD5TH95FkSX9cbr9E0gZJ25XbZ5e/cWyQtETS3g3XtKQ/lrQKWNWkzr8vabWko0f5TCdJWlHWeamk3xoR+6CG7S9L+pSkXYG/B/YuP88TkvaW9HFJ15X35XFJd0v6nW2NN7afTsRznLtLRoDDgUHb99neBFwNzBtRZh6wuFy/DjhWksr9V9t+2vbPgMEyXleSjESM7jSKzP9A4FXAWQCSDgUWAX8I7AlcDiyRtKPtM4AHgDfb3s32Z4DbgKPLmK8H7gOOatj+oe1hSccAF5fXnUUxk/HVI+p0MvBait9mniVpLnAV8D9sLx35QSS9tDz+PuAFFK8O/66kKa1ugO1fUfz29GD5eXaz/WB5eB7wDWA68HXgO5Ke10W8iOhI1900+wCrG7bXlPualrG9GXiM4vuuk3PHLMlIxOj+2vaDtjcA3wVml/sXAJfbvtP2kO3FwNPA60aJcxvw+2Xrx1HAZ4DfK4+9vjwO8A5gke27bT8NXAAc0dhXC1xse4PtXzfsO5UiIXqj7X8epQ5vBf7O9k22nwH+EtgZ+N32t2FUd9m+roz3WWAnRr8HEdFDHtqu5QLMkDTQsCzod51bSTISMbqHGtafBHYr1/cHPlB2d2yUtBHYD2ja3WD7p8CvKJKZI4HvAQ9KehlbJyN70/BeJ9tPAOvZ+reOxt9ItngfcK3tn7T4LCNjD5exuvmN5tm6lPHWMMo9iIjesemkZeQR23MalsZ5U9ZSfGdtsW+5j2ZlyjFqUym+jzo5d8ySjESM3WrgItt7NCy72L6qPN5s8p7bgFOAKbbXltvzgWnA8rLMgxSJDgDl+Io92foferPYpwInS3pvizqPjC2KL5QtsZ8EdmkoP7PNNaHhC6ls9dm3vM62xouIDg0Pq+XSxjLgYEkHll21pwNLRpRZQvEdBcV31y0uJiZbApxePm1zIHAwMFqLbMeSjESM3ReBcyS9VoVdJb1J0u7l8YeBF4845zbgXOAfyu2l5fY/2h4q910FvEvSbEk7Ap8G7rR9f5v6PAgcC7xX0h+NUuZa4E2Sji3HdXyAomvp9vL4cuDtkrYvx5+8vuHch4E9JU0dEfMwSW8pf2t6Xxnvji7iRURHuhszUo4BORe4Efg3ipbVFZIulHRSWexLFP9OB4H3A+eX566g+D65F/g+8McN32HbLMlIxBjZHgDOBv4GeJRiNPlZDUUuBj5cduGcV+67Ddid/0pG/pGi5WDLNrZ/AHwE+CawDngJxW8sndTpAYqE5HxJ72lyfCXwTuD/UrwW/c0Ug2w3lUXeW+7bSDF25TsN5/47RaJ0X/mZtnTFXE8xFuVR4AzgLeX4kW2NFxEdcPdP02D7Btsvtf0S2xeV+z5qe0m5/pTtU20fZPtw2/c1nHtRed7LbP99Lz5TpoOPiDGT9HHgINvv7HddIp5rXr773r7i0D9sWebIH34808FHREREVfLW3oiIiOijLd00k0mSkYgYM9sf73cdIp7L0jISERERfaS0jEREREQfGdx+LpFa6XkysseUXbz3zr2fPqDKZ36q+pFWVecq/wrWsc51U8fn16r6+a18vOuXfTb1qkN/WUlcgHt+/IJK4lZZ56cGp1cWO2Ddrx9j46Ynx+1rrnhr7+T6Vu15MrL3zlO58oj57QuOUZX9Y1I1/z1UVeeq6gv1rHPd1LGvt6qf3xuWVjPv2a23f7GSuAD7TT27krhV1nnVSW+tLHbAmXd8eXwvaBhKN01ERET0i/Nob0RERPRbBrBGRERE/2QAa0RERPSTJ+GjvZPr00RERDwHdPPW3nYkTZd0k6RV5Z/TRik3vyyzStL8hv1LJa2UtLxcXtjumklGIiIi6qR8mqbV0qXzgZttHwzcXG5vRdJ04GPAa4HDgY+NSFreYXt2ufyi3QWTjERERNSIgeEhtVy6NA9YXK4vBk5uUuZE4CbbG2w/CtwEzN3WC7ZMRiS9RtI9knaStKukFZJeua0Xi4iIiG617qIpu2lmSBpoWBaM4QJ72V5Xrj8E7NWkzD7A6obtNeW+Lf627KL5iKS22VHLAay2l0laAnwK2Bn4qu2ftAsaERER1XBnk549YnvOaAcl/QCY2eTQh7a+lq2xz3r4DttrJe0OfBM4A7iy1QmdPE1zIbAMeAr4k2YFyoxrAcDMnZ4/lgpHRETEGHU7SNX2caMdk/SwpFm210maBTQb87EWOLphe19gaRl7bfnn45K+TjGmpGUy0smYkT2B3YDdgZ2aFbC90PYc23OmTdmlg5ARERGxbVT1ANYlwJanY+YD1zcpcyNwgqRp5cDVE4AbJe0gaQaApOcBfwC07VHppMaXAx8BvgZc2kH5iIiIqIjdfunSJcDxklYBx5XbSJoj6YqiDt4AfJKi52QZcGG5b0eKpOQeYDlFC0rbFy+17KaRdCbwjO2vS9oeuF3SMbZv2dZPGBEREd0ZHqruYVjb64Fjm+wfAN7TsL0IWDSizK+Aw8Z6zXYDWK+k7OexPUTxPHFERET0STGANdPBR0RERB/lrb0RERHRN0ZpGYmIiIj+SstIRERE9I9hqPsp3yeUJCMRERE1YtIyEhEREf3kjBmJSW7sryCYvKr6zaOO93iy/Rb2XFPHv3NVmQx/l00e7Y2IiIg+mwxJVaMkIxERETVSTHrW71r0VpKRiIiImhmeZC0j1U1uHxEREZUYHm69dEPSdEk3SVpV/jltlHLfl7RR0vdG7D9Q0p2SBiVdI2lKu2smGYmIiKgRU7SMtFq6dD5ws+2DgZvL7Wb+Ajijyf5Lgc/ZPgh4FHh3uwsmGYmIiKgTw1CbpUvzgMXl+mLg5KbVsG8GHm/cJ0nAMcB17c5vlDEjERERNWLEULVjRvayva5cfwjYawzn7glstL253F4D7NPupCQjERERNdNB48cMSQMN2wttL9yyIekHwMwm531oq+vY1jhMVJNkJCIiokZMR10xj9ieM2oM+7jRjkl6WNIs2+skzQJ+MYbqrQf2kLRD2TqyL7C23Uktx4xIulDS+xq2L5L03jFUKiIiInps2K2XLi0B5pfr84HrOz3RtoFbgVPGcn67AayLgDMBJG0HnA58dWQhSQskDUgaeHTTk53WOSIiIsbIwFCbpUuXAMdLWgUcV24jaY6kK7YUkvRD4BvAsZLWSDqxPPRB4P2SBinGkHyp3QVbdtPYvl/SekmHUgxg+bHt9U3KLQQWAhwydVZeghAREVGhKv+jLf+fP7bJ/gHgPQ3bR45y/n3A4WO5ZidjRq4AzqIY6LJoLMEjIiKitzocM1Irncwz8m1gLvAa4MZqqxMRERHtVNxNM+7atozY3iTpVornhuv4GSMiIiYNA5PsPXntk5Fy4OrrgFOrr05ERES0M9laBto92nsIMEgxR/2q8alSREREjGZLy0irpW7aPU1zL/DicapLREREtGWGKn2eZvxlBtaIiIiamVypSJKRiIiIWjEw1O51MTXLVpKMRERE1IiBzW1L1UuSkYiIiJpx3Zo+2uh5MmLAVq/D1lJVb10+6PprK4kLsN/UsyuJu/qxL1YSt6r6VqmO96KqOt9y0lsribuqorgAtx79WCVxB+edVkncKlX1XT8Ob6yvteLdNJPrHnUyA2tERERMIFU+2itpuqSbJK0q/5w2SrnvS9oo6Xsj9n9Z0s8kLS+X2e2umWQkIiKiRrYMYG21dOl8ivnFDgZuLreb+QvgjFGO/ant2eWyvN0Fk4xERETUTMWTns0DFpfri4GTmxWyfTPwePeXSzISERFRKy4nPWu1dGkv2+vK9YeAvbYhxkWS7pH0OUk7tiucp2kiIiJqpoOumBmSBhq2F9peuGVD0g+AmU3O+1Djhm1r7COKL6BIYqYAC4EPAhe2OiHJSERERI0U76Zpmx88YnvOqDHs40Y7JulhSbNsr5M0C/jFmOr3X60qT0v6W+C8duekmyYiIqJmKu6mWQLML9fnA9eP5eQygUGSKMab/KTdOUlGIiIiasSY4TZLly4Bjpe0Cjiu3EbSHElXbCkk6YfAN4BjJa2RdGJ56GuS/hX4V2AG8Kl2F2zbTSPpHOCccnMqcL/tN3T+mSIiIqKXevD47qhsrweObbJ/AHhPw/aRo5x/zFiv2bZlxPZltmcDrwHWAJ8d60UiIiKiN7aMGamwZWTcjWUA618Bt9j+7sgDkhYACwBm7vT8HlUtIiIimnlOTgcv6Sxgf+ATzY7bXmh7ju05e0zZpYfVi4iIiEbPyZYRSYdRPJZzpO0eTOwWERER264nT8xMKJ1005wLTAduLZ7SYcD2e1qfEhEREVWwYLMmV9tA22TE9rvGoyIRERHRXtFNM7lkBtaIiIiaGZpk6UiSkYiIiBpxTQeptpJkJCIiomY2p2UkIiIi+qXDF+XVSpKRiIiIGjGwucLp4PshyUhERESNGPPMJOumkd3b7ErSL4Gfd1h8BvBITytQfey6xa0yduJWH7tucauMXbe4VcauW9wqY0+EuPvbfkEFdWhqynb7+YXP+98ty6zd9IG7bM/ZlviSpgPXAAcA9wOn2X50RJnZwBeA5wNDwEW2rymPHQhcDewJ3AWcYXtTq2v2vGVkLD8QSQPberP6FbtucauMnbjVx65b3Cpj1y1ulbHrFrfK2HWL2wvGPKOhKi9xPnCz7UsknV9uf3BEmSeBM22vkrQ3cJekG21vBC4FPmf7akmXAe+mSFxG1dG7aSIiImJiMMWL8lotXZoHLC7XFwMn/0Yd7P+wvapcfxD4BfACFVO1HwNc1+r8kTJmJCIiokaGvPbG9U9/cEabYjtJGmjYXmh7YYeX2Mv2unL9IWCvVoUlHQ5MAX5K0TWz0fbm8vAaYJ92F+x3MtLpjZlIsesWt8rYiVt97LrFrTJ23eJWGbtucauMXbe4XbM9t9sYkn4AzGxy6EMjrmVp9Ed3JM0CvgLMtz1cvsNu7PXp9QDWiIiIqC9JK4Gjba8rk42ltl/WpNzzgaXAp21fV+4T8Etgpu3Nko4APm77xFbXzJiRiIiIaLQEmF+uzweuH1lA0hTg28CVWxIRKFpSgFuBU1qdP1LfkhFJJ0uypJfXIW6VsXsdV9J+kn5WPp6FpGnl9gETNXbd4taxzrkX1cetcZ2HJC2X9C+S7pb0u93GrDJulbGrrHONXAIcL2kVcFy5jaQ5kq4oy5wGHAWcVd6v5Soe94XiyZv3SxqkGEPypbZXtN2XheIZ5h8Cn6hD3LrVGfgzigFLAJcDF0z02HWLW8c6517kXowS94mG9ROB2yZy3LrWOUuL+96Xi8JuwFrgpcDKiR63jnUGngfcA7wPWAE8b6LHrlvcOtY59yL3YpS4jf8Bnwp8ZyLHrWuds7S47325KLwD+FK5fjtw2ESOW+M6n0jxSPrxFfwMK4ldt7h1rHPuRe5Fk5hDwHLg34HHevj9VkncutY5y+hLv8aMvI1iqljKP982weNWGbvKOr8RWAe8socxq45dt7hVxq5b3Cpj1y1ulbGriPtr27NtvxyYC1ypbX1Gc3ziVhm7yjrHaMY7+wGmU0wj+3OKOe9XAw9QPmY80eLWuM6zKZpxX1TGnNXDn2ElsesWt451zr3IvRgl7hMjth8GXjhR49a1zlla3PdxvyAsAC4fsQA7n/8AAADBSURBVO824KiJGLeOdQYE/IiyGRf4X8DXevTzqyR23eLWsc65F7kXLWI3jpN4OcUL4rafqHHrWucsLe77uF+weP547oh9fwJ8YSLGrWOdKZKcaxq2twfuBl7fg3tRSey6xa1jnXMvci9axN4yTmI58C/Am7qNWWXcutY5y+hLZmCNiIiIvsoMrBEREdFXSUYiIiKir5KMRERERF8lGYmIiIi+SjISERERfZVkJCIiIvoqyUhERET0VZKRiIiI6Kv/D4J0HoI1Yjc7AAAAAElFTkSuQmCC\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": 16,
   "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='intermediate_struct_in_space_group_74.cif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "74"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pga = SpacegroupAnalyzer(new_struct)\n",
    "pga.get_space_group_number()"
   ]
  },
  {
   "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
}