perovskite_order_parameters_with_explicit_k.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recovering order parameters with k-vectors\n",
    "This notebook provides a general sketch of how one would go about recovering order parameters with k-vectors. To learn the k-vectors is more involved but theoretically possible but beyond the scope of this paper."
   ]
  },
  {
   "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:3\""
   ]
  },
  {
   "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": [
    "### In this example, we simplify the order parameter recovery process a bit by stating which k-vectors we wish to recover order parameters for, in this case $M$ and $R$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rs_order_param = [(1, 1, 1)]  # Pseudovector\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",
    "\n",
    "R = torch.ones(3) * 0.5 # In L=1 basis\n",
    "M1 = torch.tensor([1., 0., 1.]) * 0.5 # In L=1 basis, xy\n",
    "kvecs = torch.stack([R, M1], dim=0)\n",
    "\n",
    "pseudo_order_param_input = torch.zeros(kvecs.shape[0], 3) \n",
    "# pseudo_order_param_input[0] = torch.tensor([0.4, 0, 0.4])\n",
    "pseudo_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 is the crucial part, we need to combine k-vector information with learnable order parameters.\n",
    "Note, that we fix an origin for the distortion with our choice of centering `perm_normalized_B_coords`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_features(\n",
    "    kvecs,\n",
    "    pseudo_order_param_input,\n",
    "    normalized_B_coords):\n",
    "    \n",
    "    perm_normalized_B_coords = normalized_B_coords[:, [1, 2, 0]]\n",
    "    \n",
    "    vol = 1.\n",
    "    order_param_input = torch.cos(2 * math.pi / vol * (kvecs.unsqueeze(0) * perm_normalized_B_coords.unsqueeze(1)).sum(-1)) # [r, k]\n",
    "    order_param_input = order_param_input.unsqueeze(-1) * pseudo_order_param_input.unsqueeze(0) # [r, k, 3]\n",
    "    order_param_input = order_param_input.sum(-2)  # sum over kvectors\n",
    "    \n",
    "    assert list(order_param_input.shape) == [4, 3]\n",
    "    \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": [
    "### Same as before, define displacements, model, and optimizers."
   ]
  },
  {
   "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": "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([pseudo_order_param_input], 1e-3)"
   ]
  },
  {
   "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('spacegroup_62_model_overfit.torch')['model'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define our origin for the distortion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalized_B_coords = (torch.tensor(struct_62.cart_coords[-4:]) - torch.tensor(struct_62.cart_coords[-4:-3])) / lat\n",
    "normalized_B_coords.requires_grad = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model update:  tensor(0.0676, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0676, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0581, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0563, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0197, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0199, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0130, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0123, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0102, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0104, grad_fn=<AddBackward0>)\n",
      "model update:  tensor(0.0106, device='cuda:3', grad_fn=<MeanBackward0>)\n",
      "order param update:  tensor(0.0106, 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(kvecs, \n",
    "                               pseudo_order_param_input, \n",
    "                               normalized_B_coords)\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(kvecs, \n",
    "                               pseudo_order_param_input, \n",
    "                               normalized_B_coords)\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 + 1e-2 * pseudo_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": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.5000, 0.5000, 0.5000],\n",
       "        [0.5000, 0.0000, 0.5000]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kvecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "rounded_pseudo = torch.tensor(pseudo_order_param_input.detach().numpy().round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Here we recover an order parameter analogous to $(a000bb)$ in the Howard 1998 notation.\n",
    "where $(a00)$ cooresponds to a $z$ value for $M1$ and (0bb) coorespond to $-x$ and $+y$ values (minus sign comes from different in convention for $b$ vs. $a$ axis of 62 space group structure)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f704c564cd0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXEAAADnCAYAAADoz2ZOAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAQLElEQVR4nO3dfaxlVXnH8e+PQd584dUCYYjQQDQjbWqKUNOYGHlxSGPHpNJAkzo2GtJEYk1rDK0RDEIjjanWaGMnOgnYVDSmhrFOOhl5+aMttUyRaNGSmYKGoaKdGcSoRZx7n/5x9tjDzZ2Zc+eeM2fW3t/PZOfuvfY6ez/n/PHMylpr75WqQpLUpuPmHYAk6ciZxCWpYSZxSWqYSVySGmYSl6SGHT/vACRp1q68+pTau3dhorqPPPz8tqpaP+OQpsYkLqn39u5d4IEHz5uo7mknPnHWjMOZKpO4pAEILK6ZdxAzYRKX1H8FWejnEKBJXFLvBchi5h3GTJjEJfVfQRbnHcRsmMQlDYNJXJIaVZCevuvPJC5pEOxOkaRWFWShn01xk7ikYbAlLkltGk0xtCUuSW0qbIlLUsucnSJJrSrI/nkHMRsmcUnD0NNF4U3ikgbBeeKS1CoHNiWpbQ5sSlLLbIlPT/LiOi6nz+PWg/HKl062nqBW58SL9s07hN575OHn91TVy1dzjRRkwfeJT81xOZ1TTnjXPG49GHe97tl5hzAIF93zhXmH0Hunn/TEd6dyIVviktQoBzYlqXEObEpSu/q6xmY/l3+WpHEFLGSy7TCSrE/yWJJdSW5a5vyJST7fnf9akgu68quS/HuSb3Z/3ziNr2YSlzQMixNuh5BkDfBJ4BpgHXB9knVLqr0DeKaqLgI+CtzRle8B3lxVvwJsBD676u+ESVzSENQKtkO7DNhVVY9X1fPA3cCGJXU2AHd2+18ErkiSqvp6Vf13V/4ocHKSE1fztcAkLmkQAosTbod2HvDk2PHurmzZOlW1H3gWOHNJnd8BHq6qnx3xV+o4sClpGGrigc2zkuwYO95UVZumFUaSVzPqYrl6GtcziUvqv1rRWwz3VNWlBzn3FHD+2PHarmy5OruTHA+cCuwFSLIW+BLwtqr6r4kjOgS7UyQNw3RmpzwEXJzkwiQnANcBW5bU2cJo4BLgrcB9VVVJTgO+AtxUVf88ra9lEpfUf8VU+sS7Pu4bgW3At4EvVNWjSW5N8ttdtc8AZybZBfwxcGAa4o3ARcDNSR7ptl9a7VezO0XSMEzeJ37oy1RtBbYuKbt5bP854NplPncbcNtUghhjEpc0DL47RZJalam1xI81JnFJ/VdQPX13iklc0jC4KIQkNaqwO0WSmmZ3iiS1yoFNSWrXgYd9esgkLmkQyoFNSWqY3SmS1Ci7UySpZQ5sSlLbbIlLUpuqRlsfmcQlDcNCP5dPMIlL6r+Csk9cklo10Ur2TTKJSxoGW+KS1C67UySpVYXLs0lSy8rZKZLUqPKJTUlqmn3iktQypxhKUsNsiUtSm6pcFEKSGhb7xA8myQLwze5aTwC/X1U/XO11JWmqeprEpzFx8n+r6teq6hJgH/CuKVxTkqanoBYz0daaac9+fxA4b8rXlKTVOzBX/HBbY6bWJ55kDXAF8JmDnL8BuAEgnDat20rSROwTP7iTkzzCqAX+bWD7cpWqahOwCWDNcWt7usaGpGNSpbezU6bWJw68Agj2iUs6xhSjlvgk2+EkWZ/ksSS7kty0zPkTk3y+O/+1JBd05WcmuT/Jj5N8YlrfbWp94lX1U+DdwJ8kceqipGPLYibbDqHrNv4kcA2wDrg+ybol1d4BPFNVFwEfBe7oyp8DPgC8d5pfa6oDm1X1deAbwPXTvK4krUpNrSV+GbCrqh6vqueBu4ENS+psAO7s9r8IXJEkVfWTqvonRsl8albdYq6qlyw5fvNqrylJUzf5wOZZSXaMHW/qxvRgNPb35Ni53cDlSz7/izpVtT/Js8CZwJ4VxzwBuz0kDcCKntjcU1WXzjKaaTKJS+q/6b075Sng/LHjtV3ZcnV2d+ODpwJ7p3Hz5fRzqQtJGjPF2SkPARcnuTDJCcB1wJYldbYAG7v9twL3VdXMplXbEpc0DFN42Kfr474R2AasATZX1aNJbgV2VNUWRg88fjbJLkavIrnuwOeTfAd4GXBCkrcAV1fVt1YTk0lc0gBM770oVbUV2Lqk7Oax/eeAaw/y2QumEsQYk7ik/isfu5ekptViP4cATeKSBqEW5x3BbJjEJfVf0eRrZidhEpfUe+XybJLUNpO4JLXMJC5JjSpYXHB2iiS1q6friZnEJQ2AA5uS1KwDL8DqI5O4pP4rpvbulGONSVzSIPjYvSQ1yz5xSWpXweyWZZgvk7ik3nNgU5Ja58CmJLXLlrgktarCorNTJKldtsQlqWUmcUlqU5XLs0lS0+xOkaSGmcQlqVnOTpGkdrnavSS1y8fuJalxJnFJapVTDCWpZQ5sSlKz+twn3s//miRpiapMtB1OkvVJHkuyK8lNy5w/Mcnnu/NfS3LB2Lk/7cofS/KmaXwvk7ik/qvpJPEka4BPAtcA64Drk6xbUu0dwDNVdRHwUeCO7rPrgOuAVwPrgb/urrcqc+lOeeVLF7jrdc/O49aD8WdffeW8QxiEP9/wu/MOYQDumMI1prbG5mXArqp6HCDJ3cAG4FtjdTYAH+z2vwh8Ikm68rur6mfAE0l2ddd7cDUB2RKXNAyLmWyDs5LsGNtuGLvKecCTY8e7uzKWq1NV+4FngTMn/OyKObApqfeqWMnslD1Vdeks45kmW+KSBqFqsu0wngLOHzte25UtWyfJ8cCpwN4JP7tiJnFJgzCl2SkPARcnuTDJCYwGKrcsqbMF2NjtvxW4r6qqK7+um71yIXAx8G+r/V52p0gagOkMbFbV/iQ3AtuANcDmqno0ya3AjqraAnwG+Gw3cLmPUaKnq/cFRoOg+4F3VdXCamMyiUsahGk97FNVW4GtS8puHtt/Drj2IJ+9Hbh9KoF0TOKSeq+qv09smsQlDcLiQj+HAE3ikgZgag/7HHNM4pL6z+4USWpXn99iaBKXNAgmcUlqlotCSFK7CmrRlrgkNck+cUlq3AQvt2qSSVzSICzaEpekRjlPXJLaVc5OkaS22RKXpFY5xVCS2mZLXJIaVb7FUJLaZhKXpFYVLDg7RZLa5GP3ktS4Wpx3BLNhEpc0AA5sSlK7ynenSFKzCnzsXpJaZneKJDUrdqdIUquqXBRCkprmC7AkqWH2iUtSo6pgwZa4JLWrry3xfk6clKQXGM1OmWRb1V2SM5JsT7Kz+3v6Qept7OrsTLJxrPz2JE8m+fGk9zSJS+q90QuwJttW6Sbg3qq6GLi3O36BJGcAtwCXA5cBt4wl+y93ZRMziUsahKpMtK3SBuDObv9O4C3L1HkTsL2q9lXVM8B2YP0oxvrXqvreSm5on7ik/itYWJg4QZ+VZMfY8aaq2jThZ88eS8JPA2cvU+c84Mmx491d2RExiUvqvRW+T3xPVV16sJNJvgqcs8yp97/gnlWVZOaPGB22OyVJJfnbsePjk/xPkn/ojl+V5MEkP0vy3lkGK0lHZnoDm1V1ZVVdssx2D/D9JOcCdH9/sMwlngLOHzte25UdkUn6xH8CXJLk5O74qiU33Ae8G/jIkQYhSTM14aDmFAY2twAHZptsBO5Zps424Ookp3cDmld3ZUdk0oHNrcBvdfvXA587cKKqflBVDwE/P9IgJGmWCo7KFEPgw8BVSXYCV3bHJLk0yacBqmof8CHgoW67tSsjyV8k2Q2ckmR3kg8e7oaT9onfDdzcdaH8KrAZeP1KvlmSG4AbAM456WUr+agkrdrReAFWVe0FrlimfAfwzrHjzYzy6NJ67wPet5J7TpTEq+obSS5g1ArfupIbjF1jE7AJYN2p5/b0fWKSjlU+dj/q6/kI8AbgzJlEI0kz4KtoRzYDP6yqbyZ5w4zikaSZGPyiEFW1G/j40vIk5wA7gJcBi0neA6yrqh9NLUpJWqXBtsSr6iXLlD0APNDtP81onqMkHbMGm8QlqXUHphj2kUlcUv8VLNgSl6Q2FaGwJS5JzVq0JS5J7eppDjeJS+q/0cDmvKOYDZO4pEFwYFOSGtbTHG4Sl9R/BSzOO4gZMYlLGgRb4pLUMFviktSo0ULJ845iNkzikgZhYd4BzIhJXFLvObApSY0ziUtSw3raJW4Sl9R/dqdIUtOK6mlb3CQuaRCcnSJJjbI7RZIaV5mwO6WxXheTuKRBsCUuSY2yO0WSGrfQWj/JhEziknqvwCmGktSyvnanHDfvACTpaKhMtq1GkjOSbE+ys/t7+kHqbezq7EyysSs7JclXkvxnkkeTfHiSe5rEJfXeaGCzJtpW6Sbg3qq6GLi3O36BJGcAtwCXA5cBt4wl+49U1auA1wC/meSaw93QJC5pEBYn3FZpA3Bnt38n8JZl6rwJ2F5V+6rqGWA7sL6qflpV9wNU1fPAw8Daw93QJC6p94piYcJtlc6uqu91+08DZy9T5zzgybHj3V3ZLyQ5DXgzo9b8ITmwKWkQVtBVclaSHWPHm6pq04GDJF8Fzlnmc+8fP6iqSiZ9TPT/JTke+Bzw8ap6/HD1TeKSBmEFg5Z7qurSg16n6sqDnUvy/STnVtX3kpwL/GCZak8Bbxg7Xgs8MHa8CdhZVR+bJFi7UyT13lEc2NwCbOz2NwL3LFNnG3B1ktO7Ac2ruzKS3AacCrxn0huaxCUNQk34b5U+DFyVZCdwZXdMkkuTfBqgqvYBHwIe6rZbq2pfkrWMumTWAQ8neSTJOw93Q7tTJA3C0XjYp6r2AlcsU74DeOfY8WZg85I6u4EVz1Q3iUvqvQOzU/rIJC5pEBZXPlGkCSZxSb13YGCzj+aSxL/9o6f3vHbbHd+dx71X4Sxgz7yD6LnmfuPt2+YdwYo19xsDr5jGRfqZwueUxKvq5fO472ok2XGouaNaPX/j2Rvyb2xLXJIaVcB+k7gktWoqc8CPSSbxyW06fBWtkr/x7A3yN3ZgU4y/AEez4W88e4P9jeMUQ0lqlqvdS1Lj7E6RpEaNHrvvZ1vcJC5pEPraEvdVtIeQ5NYk7xk7vj3JH80zpj5K8ofdazcfSfJEkvvnHVPfJHltkm8kOSnJi7vV1C+Zd1xH01F6n/hRZ0v80DYDfw98LMlxwHWMVqfWFFXVp4BPJXkRcB/wl3MOqXeq6qEkW4DbgJOBv62q/5hzWEeNUwwHqqq+k2RvktcwWvD06937gjUbfwXcV1VfnncgPXUro0UIngPePedYjrrFFb+puw0m8cP7NPB2Rgujbj50VR2pJG9n9KKjG+ccSp+dCbwEeBFwEvCT+YZz9NgSH7YvMWrBvAj4vTnH0ktJfh14L/D6qurnFIJjw98AHwAuBO5gQP9hFsXPnZ0yTFX1fDfQ9sOqWph3PD11I3AGcH8SgB1Vddi1BTW5JG8Dfl5Vf5dkDfAvSd5YVffNO7ajxZb4QHUDmr8BXDvvWPqqqv5g3jH0XVXdBdzV7S8Al883oqOvr0ncKYaHkGQdsAu4t6p2zjseSUemKBayONHWGlvih1BV3wJ+ed5xSFqdAhdKlqRWFfB8g63sSaSqn/87SdIBSf6R0fqik9hTVetnGc80mcQlqWEObEpSw0ziktQwk7gkNcwkLkkNM4lLUsP+Dz3vDsCw90NGAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "mappable = ax.imshow(rounded_pseudo, cmap='plasma')\n",
    "ax.set_yticks(range(len(kvecs)))\n",
    "ax.set_yticklabels([\"R\", \"M1\", \"M2\", \"M3\"])\n",
    "ax.set_xticks(range(3))\n",
    "ax.set_xticklabels(['y', 'z', 'x'])\n",
    "fig.colorbar(mappable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0270, -0.0000,  0.0270],\n",
       "        [-0.0000, -0.0210,  0.0000]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rounded_pseudo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({'model': model.state_dict(), \n",
    "            'kvecs': kvecs,\n",
    "            'pseudo_order_param_input': pseudo_order_param_input},\n",
    "           'spacegroup_62_kvec_model_L_1(1).torch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = create_features(kvecs, \n",
    "                           rounded_pseudo, \n",
    "                           normalized_B_coords)\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": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0270, -0.0210,  0.0270],\n",
       "        [ 0.0270,  0.0210, -0.0270],\n",
       "        [-0.0270,  0.0210,  0.0270],\n",
       "        [ 0.0270, -0.0210, -0.0270]], grad_fn=<SliceBackward>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features[-4:, -3:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f703c7f5690>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAADLCAYAAAC8opN3AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3de7hdVX3u8e8LGO6GhCgJFwEFtGhtkIjQFkGu8VgJhwOIFwgeJaVPOdWjtII3FEWg7aOnPfUIEanBCxcRJVoqDwKhtgjNBlNqaDFbRJIQUBKCIELI3u/5Y87Qle3aa62dtWbWnjvv53nmk3kZ8zfHmjtZ+e0xxhxTtomIiIjol636XYGIiIjYsiUZiYiIiL5KMhIRERF9lWQkIiIi+irJSERERPRVkpGIiIjoqyQjERER0VdJRiLGQNJDko6ZiNeXdKakfx6v8SJi4koyErEZSdq633WIiBhvkoxEdEjSV4CXAd+R9LSkvyj3f0PSo5KelPRPkl7dcM6XJX1B0k2Sfg28SdLrJP1I0lPluddK+nTDOX8kaYmktZLulPTaVtdvUs+zJA1KWiNpoaTdy/37SLKkbRrKLpL0Xkm/A1wGHFbGXttQ/8sk3VLW9w5Je29qvIiIZpKMRHTI9unAw8Bbbe9k+y/LQ/8I7A+8FLgX+NqIU98BXATsDPwr8C3gy8BU4Grgv28oKOkg4Ergj4FdgcuBhZK2bXF9Gs4/CrgYOBWYAfwcuKaDz/YfwNnAD8vYuzQcfifwKWAasKTJ5xtrvIiIjSQZieiS7SttP2X7OeATwO9JmtxQ5Ebb/2J7GJgJbAP8re3nbd9AkaBsMA+43PbdtodsLwCeAw7tsDrvBK60fW9Zn/MpWif26eIj/oPtfyrjfaSMt1cX8SIiNpJkJKILkraWdImkn0r6FfBQeWhaQ7HlDeu7Ayu98RsqG4/vDXyw7KJZW3Zv7FWe14ndKVpDALD9NLAa2KPD85t5oX5lvDVjqE9ERFtJRiLGZuRrrt8BzAGOASYD+5T7Nco5q4A9JDUeb2xlWA5cZHuXhmUH21ePcv2RHqFIaIpKSDtSdPesBH5d7t6hofz0Fp/tt+onaSeK7qVHuogXEbGRJCMRY/MY8PKG7Z0pulFWU/yn/Jk25/8QGALOkbSNpDnAIQ3HvwicLekNKuwo6S2Sdh7l+iNdDbxb0kxJ25b1udv2Q7Z/SZGUvKts0fmfwCtGfLY9JU0aEfO/SfrDcv+ngLtsL+8iXkTERpKMRIzNxcBHyy6Uc4GrKLpFVgL3A3e1Otn2OuAk4D3AWuBdwHcpEhpsDwBnAX8HPAEMAme2uP7I+N8HPgZ8k6IV5hXAaQ1FzgL+nCJ5ejVwZ8Ox24ClwKOSHm/Y/3XgAorumYPLOncTLyL6TNJsSQ+UT96d1+T4ByTdL+k+SbdueIquPDZX0rJymduT+mzcdR0Rm5uku4HLbP99v+sykqQvAytsf7TfdYmI3ijnO/oJcCywAlgMvN32/Q1l3kTRqvqMpD8BjrT9NklTgQFgFkVX7D3Awbaf6KZOaRmJ2MwkHSFpetlNMxd4LfC9ftcrIrYYhwCDth8sW2uvoRj79gLbt9t+pty8C9izXD8euMX2mjIBuQWY3W2FtmlfJCJ67JXAdcCOwIPAybZX9bdKEVEXxxy3g1evHmpZZsm965YCzzbsmm97frm+Bxs/xbcCeEOLcO+hmE9ptHO7eVoPSDISsdmVXwjz2xYcB2yf2e86RMTGVq8eZtG/vKxlmV22H3zW9qxuryXpXRRdMkd0G6uVdNNERETUiUHDarm0sZKNpxTYs9y3ERUv5fwIcEI56WHH545VkpGIiIgaEaChrVoubSwG9pe0b/no/WnAwo2uUbya4nKKROQXDYduBo6TNEXSFOC4cl9Xet5Ns8ukHbz79pPbF4xNtsP06t45NviT6e0LjSP7HfBoZbGfebSa16nU8fm1tr9nbaIHnqrmJcavnfF0JXEBVj82pZK4u/7eI5XEBXhucGplsQMe+c2TrF33TFX/TH6bQcNdnG6vl3QORRKxNcUrJJZKuhAYsL0Q+CtgJ+Ab5RyND9s+wfYaSZ+iSGgALrS9potPA1TwaO+Bk2f4qsN68thxjOJ15y9sX2gTnXTU+ZXFrsINt11cWex7Lz6hkrj25vvO6hWpmhTqTYuq+cVlxYd/WElcgKs+e0olcU9/7IJK4gIMzjm1stgBZ/xwAfc/uWqz/cN+3czt/IPb9mlZZqddH7inF2NGNpcMYI2IiKiZblpGxqMkIxEREXVi0FAdO3xHl2QkIiKiRgRofZKRiIiI6BcD6aaJiIiIflLrCVhrJ8lIREREnXT5aO94lGQkIiKiZibaANaW07RJer2k+yRtJ2lHSUslvWZzVS4iIiJG2DBmpNVSMy1bRmwvlrQQ+DSwPfBV2z/eLDWLiIiIprbEMSMXUkz7+izwZ80KSJoHzAOYvt2Le1a5iIiI2JgMGqrfTM6tdPKivF0p5qffGdiuWQHb823Psj1ryqQdelm/iIiIGGmCddN0koxcDnwM+BpwabXViYiIiJYMDLVZ2pA0W9IDkgYlndfk+Bsl3StpvaSTRxwbkrSkXHrysrSW3TSSzgCet/11SVsDd0o6yvZtvbh4REREjJ2GN72bpvz//PPAscAKYLGkhbbvbyj2MHAmcG6TEL+xPXOTK9BEuwGsVwFXletDwBt6efGIiIgYIwPdjRk5BBi0/SCApGuAOcALyYjth8pjm6XTp5NumoiIiBhPuhszsgewvGF7RbmvU9tJGpB0l6QTx3DeqDLpWURERJ1YnbSMTJM00LA93/b8HtVgb9srJb0cuE3Sv9v+aTcBk4xERETUSWeP9j5ue9Yox1YCezVs71nu6+zy9sryzwclLQIOArpKRtJNExERUTfDar20thjYX9K+kiYBpwEdPRUjaYqkbcv1acAf0DDWZFMlGYmIiKiTDQNYWy2tTrfXA+cANwP/AVxne6mkCyWdAC+8DmYFcApwuaSl5em/AwxI+jfgduCSEU/hbJJ000RERNSNu5uB1fZNwE0j9n28YX0xRffNyPPuBH63q4s3kWSkhu69+ITKYn/02KXtC40jVd6L153fk7l8fstJR51fSdwq3XDbxZXE/cxd1dwLf6i6V2h9+DOHVRL3jEvzDtLoVEcDWGslyUhERESdGNzFpGfjUZKRiIiIuknLSERERPSNgaGJ9fxJkpGIiIhaUdcDWMebJCMRERF10v27acadJCMRERE1kwGsERER0T8ZMxIRERH9lTEjERER0Uc2eIKNGWnZzlPOU//+hu2LJL2v+mpFRETEqIa2ar3UTLsaXwmcASBpK4o3+311ZCFJ8yQNSBp4Yt0zva9lREREFAy2Wi7tSJot6QFJg5LOa3L8jZLulbRe0skjjs2VtKxc5vbiI7XsprH9kKTVkg4CdgN+ZHt1k3LzgfkAB06e4V5ULCIiIppRV60fkrYGPg8cC6wAFktaOOLtuw8DZwLnjjh3KnABMItiKO095blPbHKF6GzMyBVlhaZTtJREREREH3XS+tHCIcCg7QcBJF0DzAFeSEZsP1QeGx5x7vHALbbXlMdvAWYDV3dToU5Sq2+VF3o9cHM3F4uIiIgubZj0rNXS2h7A8obtFeW+TnRz7qjatozYXifpdmCt7aFuLxgRERHd6WDSs2mSBhq255dDKsaltslIOXD1UOCU6qsTERERLVl4uG3HxuO2Z41ybCWwV8P2nuW+TqwEjhxx7qIOzx1Vu0d7DwQGgVttL+v2YhEREdED3T3auxjYX9K+kiZRPCm7sMMr3wwcJ2mKpCnAcfRgCEe7p2nuB17e7UUiIiKiN+xi2fTzvV7SORRJxNbAlbaXSroQGLC9UNLrKcaMTgHeKumTtl9te42kT1EkNAAXbhjM2o3MwBoREVErwl1ObGb7JuCmEfs+3rC+mKILptm5V9Ljp2uTjERERNRN3k0TERERfWO6bhkZb5KMRERE1EyXk56NO0lGami/G6+rLPZek8+qJO7yJ79YSdyq6gvAotMrCbv8mQsqiVvlvfjVoc9WEvewQ5dWEndwzqmVxAW4/cgnK4m77M4DK4kbE5A14d7am2QkIiKiRky6aSIiIqKfnG6aiIiI6KuOZmCtlSQjERERddP+3TS1kmQkIiKiRmwYTjISERERfZUxIxEREdE3FsN5miYiIiL6aaINYG37aSSdLWlJufxM0u2bo2IRERHRnIdbL+1Imi3pAUmDks5rcnxbSdeWx++WtE+5fx9Jv2nICy7rxedp2zJi+zLgMkkvAm4DPtuLC0dERMTY2d092itpa+DzwLHACmCxpIW2728o9h7gCdv7SToNuBR4W3nsp7ZnbnIFmhjLp/kb4Dbb3xl5QNI8SQOSBp5Y90zvahcRERG/xVbLpY1DgEHbD9peB1wDzBlRZg6woFy/HjhaUmWjZjtKRiSdCewNfLLZcdvzbc+yPWvKpB16WL2IiIjYSPnW3lYLMG1DI0G5zGuIsAewvGF7RbmPZmVsrweeBHYtj+0r6UeS7pB0eC8+UttuGkkHA+cCh9ud9ERFRERElTpo/Xjc9qwKLr0KeJnt1WV+8G1Jr7b9q26CdtIycg4wFbi9HKxyRTcXjIiIiG4ID7de2lgJ7NWwvWe5r2kZSdsAk4HVtp+zvRrA9j3AT4EDuv1EnQxgfXe3F4mIiIjeMF2/KG8xsL+kfSmSjtOAd4wosxCYC/wQOJlizKglvQRYY3tI0suB/YEHu6kMZJ6RiIiIeinHjGzy6fZ6SecANwNbA1faXirpQmDA9kLgS8BXJA0CaygSFoA3AhdKeh4YBs62vaaLTwMkGYmIiKgZMdzlpGe2bwJuGrHv4w3rzwKnNDnvm8A3u7p4E0lGIiIi6sTFy/ImkiQjERERNWK666YZj5KMRERE1EyXA1jHnSQjERERdeLux4yMN3KPO54k/RL4eYfFpwGP97QC1ceuW9wqYydu9bHrFrfK2HWLW2XsusWtMvZ4iLu37ZdUUIemfnfXl/qGN5/csswBX/vCPRVNelaJnreMjOUHImmgqptVVey6xa0yduJWH7tucauMXbe4VcauW9wqY9ctbi/Y4KF000REREQfTbRumiQjERERtdLRm3lrpd/JyPwaxq5b3CpjJ271sesWt8rYdYtbZey6xa0ydt3ids/gCdYy0vMBrBEREVGdV0/ZzdcdNfJVMht7zQ3/Z8sewBoRm48kA/vbHux3XSJiMzEMT7ABrBOrnSdiHJD0kKRj+l2PXpF0pKQV4zVexJanGDPSamkbQZot6QFJg5LOa3J8W0nXlsfvlrRPw7Hzy/0PSDq+F58oyUhEDUlKq2bEFsoUT9O0WlqRtDXweeDNwIHA2yUdOKLYe4AnbO8HfA64tDz3QIo3+L4amA38vzJeV5KMRDRRtm6cK+k+SU+WvyFs13D8jyQtkbRW0p2SXlvu/wrwMuA7kp6W9BeSFkj6YHl8D0mW9Kfl9iskrZG0Vbl9VvkbxxpJCyXt3nBNS/pTScuAZU3q/IeSlks6cpTPdIKkpWWdF0n6nRGx92vY/rKkT0vaEfhHYPfy8zwtaXdJn5B0fXlfnpJ0r6Tf29R4Y/vpRGzh3F0yAhwCDNp+0PY64Bpgzogyc4AF5fr1wNGSVO6/xvZztn8GDJbxupJkJGJ0p1Jk/vsCrwXOBJB0EHAl8MfArsDlwEJJ29o+HXgYeKvtnWz/JXAHcGQZ8wjgQeCNDds/sD0s6Sjg4vK6MyhmMr5mRJ1OBN5A8dvMCyTNBq4G/oftRSM/iKQDyuPvB15C8erw70ia1OoG2P41xW9Pj5SfZyfbj5SH5wDfAKYCXwe+LelFXcSLiI503U2zB7C8YXtFua9pGdvrgScpvu86OXfMkoxEjO5vbT9iew3wHWBmuX8ecLntu20P2V4APAccOkqcO4A/LFs/3gj8JfAH5bEjyuMA7wSutH2v7eeA84HDGvtqgYttr7H9m4Z9p1AkRG+2/a+j1OFtwD/YvsX288BfA9sDv9/+NozqHtvXl/E+C2zH6PcgInrIQ1u1XIBpkgYalnn9rnMrSUYiRvdow/ozwE7l+t7AB8vujrWS1gJ7AU27G2z/FPg1RTJzOPBd4BFJr2TjZGR3Gt7rZPtpYDUb/9bR+BvJBu8HrrP94xafZWTs4TJWN7/RvFCXMt4KRrkHEdE7Np20jDxue1bD0jhvykqK76wN9iz30axMOUZtMsX3USfnjlmSkYixWw5cZHuXhmUH21eXx5tN3nMHcDIwyfbKcnsuMAVYUpZ5hCLRAaAcX7ErG/9Dbxb7FOBESe9rUeeRsUXxhbIh9jPADg3lp7e5JjR8IZWtPnuW19nUeBHRoeFhtVzaWAzsL2nfsqv2NGDhiDILKb6joPjuus3FxGQLgdPKp232BfYHRmuR7ViSkYix+yJwtqQ3qLCjpLdI2rk8/hjw8hHn3AGcA/xTub2o3P5n20PlvquBd0uaKWlb4DPA3bYfalOfR4CjgfdJ+pNRylwHvEXS0eW4jg9SdC3dWR5fArxD0tbl+JMjGs59DNhV0uQRMQ+WdFL5W9P7y3h3dREvIjrS3ZiRcgzIOcDNwH9QtKwulXShpBPKYl+i+Hc6CHwAOK88dynF98n9wPeAP234DttkSUYixsj2AHAW8HfAExSjyc9sKHIx8NGyC+fcct8dwM78VzLyzxQtBxu2sf194GPAN4FVwCsofmPppE4PUyQk50l6b5PjDwDvAv4vxWvR30oxyHZdWeR95b61FGNXvt1w7n9SJEoPlp9pQ1fMjRRjUZ4ATgdOKsePbGq8iOiAu3+aBts32T7A9itsX1Tu+7jtheX6s7ZPsb2f7UNsP9hw7kXlea+0/Y+9+EyZDj4ixkzSJ4D9bL+r33WJ2NK8aufdfcVBf9yyzOE/+ESmg4+IiIiq5K29ERER0UcbumkmkiQjETFmtj/R7zpEbMnSMhIRERF9pLSMRERERB8Z3H4ukVrpeTKyy6QdvPv2vZ8+oMpnfqr6kVZV5x2nr60oMgz+ZHr7QuPIfgc82r7QJvr1o7tUFrtuqvo38sBTXb/ss6nXHvTLSuIC3Pejl1QSt8o6Pzs4tZK4dfzvsIrv5VW/eZK1657ZbLejeGtvHe/+6HqejOy+/WSuOmxu+4JjVGX/mFRN2lBVnQ/+8I2VxAU46ajzK4tdhRvmX1xZ7Hs+M/Illluuqv6NvGlRNfOe3X7nFyuJC7DX5LMqiVtlnZed8LZK4lb196JKVXwvn3HXl3sesyXDULppIiIiol+cR3sjIiKi3zKANSIiIvonA1gjIiKinzwBH+2dWJ8mIiJiC9DNW3vbkTRV0i2SlpV/Thml3NyyzDJJcxv2L5L0gKQl5fLSdtdMMhIREVEn5dM0rZYunQfcant/4NZyeyOSpgIXAG8ADgEuGJG0vNP2zHL5RbsLJhmJiIioEQPDQ2q5dGkOsKBcXwCc2KTM8cAtttfYfgK4BZi9qRdsmYxIer2k+yRtJ2lHSUslvWZTLxYRERHdat1FU3bTTJM00LDMG8MFdrO9qlx/FNitSZk9gOUN2yvKfRv8fdlF8zFJbbOjlgNYbS+WtBD4NLA98FXbP24XNCIiIqrhziY9e9z2rNEOSvo+0GzK7Y9sfC1bY5/d7p22V0raGfgmcDpwVasTOnma5kJgMfAs8GfNCpQZ1zyA6du9eCwVjoiIiDHqdpCq7WNGOybpMUkzbK+SNANoNuZjJXBkw/aewKIy9sryz6ckfZ1iTEnLZKSTMSO7AjsBOwPbNStge77tWbZnTZm0QwchIyIiYtOo6gGsC4ENT8fMBZq9g+Rm4DhJU8qBq8cBN0vaRtI0AEkvAv4IaNuj0kmNLwc+BnwNuLSD8hEREVERu/3SpUuAYyUtA44pt5E0S9IVRR28BvgURc/JYuDCct+2FEnJfcASihaUti9eatlNI+kM4HnbX5e0NXCnpKNs37apnzAiIiK6MzxU3cOwtlcDRzfZPwC8t2H7SuDKEWV+DRw81mu2G8B6FWU/j+0hiueJIyIiok+KAayZDj4iIiL6KG/tjYiIiL4xSstIRERE9FdaRiIiIqJ/DEPdT/k+riQZiYiIqBGTlpGIiIjoJ2fMSIwD9158QmWxP3rs0spiV6HKe3Hwh5tNOti9k446v5K4VbrhtourCbzo9GrixkbG/mqRztTxt/Mq7sXmvgsmj/ZGREREn9UxEWwlyUhERESNFJOe9bsWvZVkJCIiomaGJ1jLSHWT20dEREQlhodbL92QNFXSLZKWlX9OGaXc9yStlfTdEfv3lXS3pEFJ10qa1O6aSUYiIiJqxBQtI62WLp0H3Gp7f+DWcruZvwKajUK/FPic7f2AJ4D3tLtgkpGIiIg6MQy1Wbo0B1hQri8ATmxaDftW4KnGfZIEHAVc3+78RhkzEhERUSNGDFU7ZmQ326vK9UeB3cZw7q7AWtvry+0VwB7tTkoyEhERUTMdNH5MkzTQsD3f9vwNG5K+D0xvct5HNrqObVU1UU2DJCMRERE1Yjrqinnc9qxRY9jHjHZM0mOSZtheJWkG8IsxVG81sIukbcrWkT2Ble1OajlmRNKFkt7fsH2RpPeNoVIRERHRY8NuvXRpITC3XJ8LdDwdtW0DtwMnj+X8dgNYrwTOAJC0FXAa8NWRhSTNkzQgaeCJdc90WueIiIgYIwNDbZYuXQIcK2kZcEy5jaRZkq7YUEjSD4BvAEdLWiHp+PLQh4APSBqkGEPypXYXbNlNY/shSaslHUQxgOVHtlc3KTcfmA9w4OQZlfctRUREbMmq/I+2/H/+6Cb7B4D3NmwfPsr5DwKHjOWanYwZuQI4k2Kgy5VjCR4RERG91eGYkVrpZJ6RbwGzgdcDN1dbnYiIiGin4m6aza5ty4jtdZJup3huuI6fMSIiYsIwMMHek9c+GSkHrh4KnFJ9dSIiIqKdidYy0O7R3gOBQYo56pdtnipFRETEaDa0jLRa6qbd0zT3Ay/fTHWJiIiItsxQpc/TbH6ZgTUiIqJmJlYqkmQkIiKiVgwMtXtdTM2ylSQjERERNWJgfdtS9ZJkJCIiomZct6aPNnqejBiw1euwtVTVW5f3u/G6SuIC7DX5rEriLn/yi5XEraq+ACw6vZKwy5+5oJK4Vd6LXx36bCVxbzviV5XEXXbC2yqJC3D7kU9WEndwzqmVxIX6fSdvhjfW11rxbpqJdY86mYE1IiIixpEqH+2VNFXSLZKWlX9OGaXc9yStlfTdEfu/LOlnkpaUy8x210wyEhERUSMbBrC2Wrp0HsX8YvsDt5bbzfwVMFoT8p/bnlkuS9pdMMlIREREzVQ86dkcYEG5vgA4sVkh27cCT3V/uSQjERERteJy0rNWS5d2s72qXH8U2G0TYlwk6T5Jn5O0bbvCeZomIiKiZjroipkmaaBhe77t+Rs2JH0fmN7kvI80bti2xj6i+HyKJGYSMB/4EHBhqxOSjERERNRI8W6atvnB47ZnjRrDPma0Y5IekzTD9ipJM4BfjKl+/9Wq8pykvwfObXdOumkiIiJqpuJumoXA3HJ9LnDjWE4uExgkiWK8yY/bnZNkJCIiokaMGW6zdOkS4FhJy4Bjym0kzZJ0xYZCkn4AfAM4WtIKSceXh74m6d+BfwemAZ9ud8G23TSSzgbOLjcnAw/ZflPnnykiIiJ6qQeP747K9mrg6Cb7B4D3NmwfPsr5R431mm1bRmxfZnsm8HpgBfDZsV4kIiIiemPDmJEKW0Y2u7EMYP0b4Dbb3xl5QNI8YB7A9O1e3KOqRURERDNb5HTwks4E9gY+2ey47fm2Z9metcukHXpYvYiIiGi0RbaMSDqY4rGcw233YGK3iIiI2HQ9eWJmXOmkm+YcYCpwe/GUDgO239v6lIiIiKiCBes1sdoG2iYjtt+9OSoSERER7RXdNBNLZmCNiIiomaEJlo4kGYmIiKgR13SQaitJRiIiImpmfVpGIiIiol86fFFerSQZiYiIqBED6yucDr4fkoxERETUiDHPT7BuGtm9za4k/RL4eYfFpwGP97QC1ceuW9wqYydu9bHrFrfK2HWLW2XsusWtMvZ4iLu37ZdUUIemJm21l1/6ov/dsszKdR+8x/asTYkvaSpwLbAP8BBwqu0nRpSZCXwBeDEwBFxk+9ry2L7ANcCuwD3A6bbXtbpmz1tGxvIDkTSwqTerX7HrFrfK2Ilbfey6xa0ydt3iVhm7bnGrjF23uL1gzPMaqvIS5wG32r5E0nnl9odGlHkGOMP2Mkm7A/dIutn2WuBS4HO2r5F0GfAeisRlVB29myYiIiLGB1O8KK/V0qU5wIJyfQFw4m/Vwf6J7WXl+iPAL4CXqJiq/Sjg+lbnj5QxIxERETUy5JU3r37uQ9PaFNtO0kDD9nzb8zu8xG62V5XrjwK7tSos6RBgEvBTiq6ZtbbXl4dXAHu0u2C/k5FOb8x4il23uFXGTtzqY9ctbpWx6xa3yth1i1tl7LrF7Zrt2d3GkPR9YHqTQx8ZcS1Loz+6I2kG8BVgru3h8h12Y69PrwewRkRERH1JegA40vaqMtlYZPuVTcq9GFgEfMb29eU+Ab8EptteL+kw4BO2j291zYwZiYiIiEYLgbnl+lzgxpEFJE0CvgVctSERgaIlBbgdOLnV+SP1LRmRdKIkS3pVHeJWGbvXcSXtJeln5eNZSJpSbu8zXmPXLW4d65x7UX3cGtd5SNISSf8m6V5Jv99tzCrjVhm7yjrXyCXAsZKWAceU20iaJemKssypwBuBM8v7tUTF475QPHnzAUmDFGNIvtT2irb7slA8w/wD4JN1iFu3OgN/QTFgCeBy4PzxHrtucetY59yL3ItR4j7dsH48cMd4jlvXOmdpcd/7clHYCVgJHAA8MN7j1rHOwIuA+4D3A0uBF4332HWLW8c6517kXowSt/E/4FOAb4/nuHWtc5YW970vF4V3Al8q1+8EDh7PcWtc5+MpHkk/toKfYSWx6xa3jnXOvci9aBJzCFgC/CfwZA+/3yqJW9c6Zxl96deYkbdTTBVL+efbx3ncKmNXWec3A6uA1/QwZtWx6xa3yth1i1tl7LrFrTJ2FYlYmi4AAAElSURBVHF/Y3um7VcBs4GrtKnPaG6euFXGrrLOMZrNnf0AUymmkf05xZz3y4GHKR8zHm9xa1znmRTNuC8rY87o4c+wkth1i1vHOude5F6MEvfpEduPAS8dr3HrWucsLe77Zr8gzAMuH7HvDuCN4zFuHesMCPghZTMu8L+Ar/Xo51dJ7LrFrWOdcy9yL1rEbhwn8SqKF8RtPV7j1rXOWVrc981+weL549kj9v0Z8IXxGLeOdaZIcq5t2N4auBc4ogf3opLYdYtbxzrnXuRetIi9YZzEEuDfgLd0G7PKuHWtc5bRl8zAGhEREX2VGVgjIiKir5KMRERERF8lGYmIiIi+SjISERERfZVkJCIiIvoqyUhERET0VZKRiIiI6KskIxEREdFX/x+QaSz8t26MnQAAAABJRU5ErkJggg==\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": 20,
   "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_kvec.cif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "62"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pga = SpacegroupAnalyzer(new_struct, symprec=1e-3)\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
}