{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qsymm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sympy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ham = (\"hbar^2 / (2 * m) * (k_x**2 + k_y**2 + k_z**2) * eye(2) +\" +\n",
    "        \"alpha * sigma_x * k_x + alpha * sigma_y * k_y + alpha * sigma_z * k_z\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "H = qsymm.Model(ham)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{alpha*k_z:\n",
      "[[ 1.+0.j  0.+0.j]\n",
      " [ 0.+0.j -1.+0.j]],\n",
      "\n",
      "hbar**2*k_x**2/m:\n",
      "[[0.5+0.j 0. +0.j]\n",
      " [0. +0.j 0.5+0.j]],\n",
      "\n",
      "hbar**2*k_y**2/m:\n",
      "[[0.5+0.j 0. +0.j]\n",
      " [0. +0.j 0.5+0.j]],\n",
      "\n",
      "hbar**2*k_z**2/m:\n",
      "[[0.5+0.j 0. +0.j]\n",
      " [0. +0.j 0.5+0.j]],\n",
      "\n",
      "alpha*k_x:\n",
      "[[0.+0.j 1.+0.j]\n",
      " [1.+0.j 0.+0.j]],\n",
      "\n",
      "alpha*k_y:\n",
      "[[0.+0.j 0.-1.j]\n",
      " [0.+1.j 0.+0.j]],\n",
      "\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(H)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}\\alpha k_{z} + \\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{y}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m} & \\alpha k_{x} - i \\alpha k_{y}\\\\\\alpha k_{x} + i \\alpha k_{y} & - \\alpha k_{z} + \\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{y}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[alpha*k_z + hbar**2*k_x**2/(2*m) + hbar**2*k_y**2/(2*m) + hbar**2*k_z**2/(2*m),                                                         alpha*k_x - I*alpha*k_y],\n",
       "[                                                       alpha*k_x + I*alpha*k_y, -alpha*k_z + hbar**2*k_x**2/(2*m) + hbar**2*k_y**2/(2*m) + hbar**2*k_z**2/(2*m)]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H.tosympy(nsimplify=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(k_x, k_y, k_z)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H.momenta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ham2D = (\"hbar^2 / (2 * m) * (k_x**2 + k_z**2) * eye(2) +\" +\n",
    "         \"alpha * sigma_x * k_x + alpha * sigma_y * k_z\")\n",
    "H2D = qsymm.Model(ham2D, momenta=['k_x', 'k_z'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}\\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m} & \\alpha k_{x} - i \\alpha k_{z}\\\\\\alpha k_{x} + i \\alpha k_{z} & \\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[hbar**2*k_x**2/(2*m) + hbar**2*k_z**2/(2*m),                     alpha*k_x - I*alpha*k_z],\n",
       "[                    alpha*k_x + I*alpha*k_z, hbar**2*k_x**2/(2*m) + hbar**2*k_z**2/(2*m)]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H2D.tosympy(nsimplify=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(k_x, k_z)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H2D.momenta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Identity in 3D\n",
    "E = qsymm.identity(3)\n",
    "# Inversion in 3D\n",
    "I = qsymm.inversion(3)\n",
    "# 4-fold rotation around the x-axis\n",
    "C4 = qsymm.rotation(1/4, [1, 0, 0])\n",
    "# 3-fold rotation around the [1, 1, 1] axis\n",
    "C3 = qsymm.rotation(1/3, [1, 1, 1])\n",
    "# Time reversal\n",
    "TR = qsymm.time_reversal(3)\n",
    "# Particle-hole\n",
    "PH = qsymm.particle_hole(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$R\\left(\\frac{\\pi}{2}, [1 0 0]\\right)$"
      ],
      "text/plain": [
       "R(π/2, [1 0 0])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$ \\mathcal{T}$"
      ],
      "text/plain": [
       "T"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "192\n"
     ]
    }
   ],
   "source": [
    "cubic_gens = {I, C4, C3, TR, PH}\n",
    "cubic_group = qsymm.groups.generate_group(cubic_gens)\n",
    "print(len(cubic_group))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$R\\left(\\pi, [1 1 0]\\right)$"
      ],
      "text/plain": [
       "R(π, [1 1 0])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C3 * C4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$R\\left(-\\frac{2\\pi}{3}, [1 1 1]\\right)$"
      ],
      "text/plain": [
       "R(-2π/3, [1 1 1])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C3**-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}- \\alpha k_{z} + \\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{y}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m} & - \\alpha k_{x} - i \\alpha k_{y}\\\\- \\alpha k_{x} + i \\alpha k_{y} & \\alpha k_{z} + \\frac{\\hbar^{2} k_{x}^{2}}{2 m} + \\frac{\\hbar^{2} k_{y}^{2}}{2 m} + \\frac{\\hbar^{2} k_{z}^{2}}{2 m}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[-alpha*k_z + hbar**2*k_x**2/(2*m) + hbar**2*k_y**2/(2*m) + hbar**2*k_z**2/(2*m),                                                       -alpha*k_x - I*alpha*k_y],\n",
       "[                                                       -alpha*k_x + I*alpha*k_y, alpha*k_z + hbar**2*k_x**2/(2*m) + hbar**2*k_y**2/(2*m) + hbar**2*k_z**2/(2*m)]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "H_with_TR = TR.apply(H)\n",
    "H_with_TR.tosympy(nsimplify=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "sz = qsymm.ContinuousGroupGenerator(None, np.array([[1, 0], [0, -1]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}0 & - 2 i \\alpha k_{x} - 2 \\alpha k_{y}\\\\2 i \\alpha k_{x} - 2 \\alpha k_{y} & 0\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[                          0, -2*I*alpha*k_x - 2*alpha*k_y],\n",
       "[2*I*alpha*k_x - 2*alpha*k_y,                            0]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sz.apply(H).tosympy(nsimplify=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48 0\n"
     ]
    }
   ],
   "source": [
    "discrete_symm, continuous_symm = qsymm.symmetries(H, cubic_group)\n",
    "print(len(discrete_symm), len(continuous_symm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}1 & 0\\\\0 & 1\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[1, 0],\n",
       "[0, 1]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}k_{z} & k_{x} - i k_{y}\\\\k_{x} + i k_{y} & - k_{z}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[        k_z, k_x - I*k_y],\n",
       "[k_x + I*k_y,        -k_z]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}k_{x}^{2} + k_{y}^{2} + k_{z}^{2} & 0\\\\0 & k_{x}^{2} + k_{y}^{2} + k_{z}^{2}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[k_x**2 + k_y**2 + k_z**2,                        0],\n",
       "[                       0, k_x**2 + k_y**2 + k_z**2]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "family = qsymm.continuum_hamiltonian(discrete_symm, dim=3, total_power=2, prettify=True)\n",
    "qsymm.display_family(family)"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}