grogu/test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use numpy number of threads one\n",
    "# print(threadpool_info())\n",
    "# from threadpoolctl import threadpool_info, threadpool_limits\n",
    "# user_api = threadpool_info()[0][\"user_api\"]\n",
    "# threadpool_limits(limits=1, user_api=user_api)\n",
    "# print(threadpool_info())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.14.3\n",
      "1.24.4\n"
     ]
    }
   ],
   "source": [
    "from sys import getsizeof\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "import sisl\n",
    "from src.grogupy import *\n",
    "from mpi4py import MPI\n",
    "import warnings\n",
    "\n",
    "# runtime information\n",
    "times = dict()\n",
    "times[\"start_time\"] = timer()\n",
    "########################\n",
    "# it works if data is in downloads folder\n",
    "########################\n",
    "sisl.__version__\n",
    "\n",
    "try:\n",
    "    print(sisl.__version__)\n",
    "except:\n",
    "    print(\"sisl version unknown.\")\n",
    "\n",
    "try:\n",
    "    print(np.__version__)\n",
    "except:\n",
    "    print(\"numpy version unknown.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([[-8.20799665e+01,  1.42696825e-04, -8.30905002e-04],\n",
      "       [ 1.41821186e-04, -8.20626555e+01, -1.88623702e-01],\n",
      "       [ 6.97129733e-04,  1.88607418e-01, -8.92465137e+01]]), -84.46304519743876, array([[ 2.38307873e+00,  1.42259005e-04, -6.68876347e-05],\n",
      "       [ 1.42259005e-04,  2.40038974e+00, -8.14185600e-06],\n",
      "       [-6.68876347e-05, -8.14185600e-06, -4.78346847e+00]]), array([[ 0.00000000e+00,  4.37819604e-07, -7.64017368e-04],\n",
      "       [-4.37819604e-07,  0.00000000e+00, -1.88615560e-01],\n",
      "       [ 7.64017368e-04,  1.88615560e-01,  0.00000000e+00]]))\n",
      "[[-8.20799665e+01  1.42259005e-04 -6.68876347e-05]\n",
      " [ 1.42259005e-04 -8.20626555e+01 -8.14185600e-06]\n",
      " [-6.68876347e-05 -8.14185600e-06 -8.92465137e+01]] -84.46304519743875 [ 2.38307873e+00  2.40038974e+00  1.42259005e-04 -6.68876347e-05\n",
      " -8.14185600e-06] [ 1.88615560e-01 -7.64017368e-04 -4.37819604e-07]\n"
     ]
    },
    {
     "ename": "LinAlgError",
     "evalue": "Singular matrix",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mLinAlgError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[187], line 60\u001b[0m\n\u001b[1;32m     57\u001b[0m                 M \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mouter(ekkprime, ekprimek)\n\u001b[1;32m     58\u001b[0m                 V \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m E[j] \u001b[38;5;241m*\u001b[39m ekkprime\u001b[38;5;241m.\u001b[39mflatten()\n\u001b[0;32m---> 60\u001b[0m K \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinalg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43mM\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mV\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36msolve\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/Documents/oktatás/elte/phd/grogu_project/.venv/lib/python3.9/site-packages/numpy/linalg/linalg.py:386\u001b[0m, in \u001b[0;36msolve\u001b[0;34m(a, b)\u001b[0m\n\u001b[1;32m    384\u001b[0m signature \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDD->D\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m isComplexType(t) \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdd->d\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    385\u001b[0m extobj \u001b[38;5;241m=\u001b[39m get_linalg_error_extobj(_raise_linalgerror_singular)\n\u001b[0;32m--> 386\u001b[0m r \u001b[38;5;241m=\u001b[39m \u001b[43mgufunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msignature\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msignature\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextobj\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mextobj\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m wrap(r\u001b[38;5;241m.\u001b[39mastype(result_t, copy\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m))\n",
      "File \u001b[0;32m~/Documents/oktatás/elte/phd/grogu_project/.venv/lib/python3.9/site-packages/numpy/linalg/linalg.py:89\u001b[0m, in \u001b[0;36m_raise_linalgerror_singular\u001b[0;34m(err, flag)\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_raise_linalgerror_singular\u001b[39m(err, flag):\n\u001b[0;32m---> 89\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m LinAlgError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSingular matrix\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mLinAlgError\u001b[0m: Singular matrix"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "\n",
    "with open(\"./Fe3GeTe2_fdf_test.pickle\", \"rb\") as file:\n",
    "    out = pickle.load(file)\n",
    "\n",
    "ref_xcf = out[\"parameters\"][\"ref_xcf_orientations\"]\n",
    "energies = out[\"pairs\"][0][\"energies\"]\n",
    "\n",
    "\n",
    "def fit_energies(energies, ref_xcf):\n",
    "    M = np.zeros((9, 9))\n",
    "    V = np.zeros(9)\n",
    "\n",
    "    for i in range(len(ref_xcf)):\n",
    "        E = energies[i]\n",
    "        vw = ref_xcf[i][\"vw\"][:2]\n",
    "        o = ref_xcf[i][\"o\"]\n",
    "\n",
    "        vw = np.array([[vw[0], vw[0]], [vw[0], vw[1]], [vw[1], vw[0]], [vw[1], vw[1]]])\n",
    "\n",
    "        for j in range(len(vw)):\n",
    "            aa = np.cross(vw[j, 0], o)\n",
    "            bb = np.cross(vw[j, 1], o)\n",
    "            ekkprime = np.outer(aa, bb)\n",
    "            ekprimek = np.outer(bb, aa)\n",
    "            M += np.outer(ekkprime, ekprimek)\n",
    "            V += E[j] * ekkprime.flatten()\n",
    "\n",
    "    J = np.linalg.solve(M, V)\n",
    "    J = J.reshape(3, 3) * sisl.unit_convert(\"eV\", \"meV\")\n",
    "    J_s = 0.5 * (J + J.T)\n",
    "    D = 0.5 * (J - J.T)\n",
    "    J_iso = np.trace(J_s) / 3\n",
    "    J_S = J_s - np.eye(3) * J_iso\n",
    "\n",
    "    return J, J_iso, J_S, D\n",
    "\n",
    "\n",
    "print(fit_energies(energies, ref_xcf))\n",
    "print(\n",
    "    out[\"pairs\"][0][\"J\"],\n",
    "    out[\"pairs\"][0][\"J_iso\"],\n",
    "    out[\"pairs\"][0][\"J_S\"],\n",
    "    out[\"pairs\"][0][\"D\"],\n",
    ")\n",
    "\n",
    "\n",
    "energies = out[\"magnetic_entities\"][0][\"energies\"]\n",
    "M = np.zeros((9, 9))\n",
    "V = np.zeros(9)\n",
    "\n",
    "for i in range(len(ref_xcf)):\n",
    "    E = energies[i]\n",
    "    vw = ref_xcf[i][\"vw\"]\n",
    "    o = ref_xcf[i][\"o\"]\n",
    "\n",
    "    vw = np.array([[vw[0], vw[0]], [vw[1], vw[1]], [vw[2], vw[2]]])\n",
    "\n",
    "    for j in range(len(vw)):\n",
    "        aa = np.cross(vw[j, 0], o)\n",
    "        bb = np.cross(vw[j, 1], o)\n",
    "        ekkprime = np.outer(aa, bb)\n",
    "        ekprimek = np.outer(bb, aa)\n",
    "        M += np.outer(ekkprime, ekprimek)\n",
    "        V += E[j] * ekkprime.flatten()\n",
    "\n",
    "K = np.linalg.solve(M, V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matrix([[a1], [a2], [a3]])\n",
      "Matrix([[b1], [b2], [b3]])\n",
      "Matrix([[J11, J12, J13], [J21, J22, J23], [J31, J32, J33]])\n",
      "J11*a1*b1 + J12*a1*b2 + J13*a1*b3 + J21*a2*b1 + J22*a2*b2 + J23*a2*b3 + J31*a3*b1 + J32*a3*b2 + J33*a3*b3\n",
      "J11*a1*b1 + J12*a1*b2 + J13*a1*b3 + J21*a2*b1 + J22*a2*b2 + J23*a2*b3 + J31*a3*b1 + J32*a3*b2 + J33*a3*b3\n",
      "Matrix([[J11], [J12], [J13], [J21], [J22], [J23], [J31], [J32], [J33]])\n",
      "Matrix([[J11*a1*b1 + J12*a1*b2 + J13*a1*b3 + J21*a2*b1 + J22*a2*b2 + J23*a2*b3 + J31*a3*b1 + J32*a3*b2 + J33*a3*b3]])\n",
      "Matrix([[a1*b1], [a1*b2], [a1*b3], [a2*b1], [a2*b2], [a2*b3], [a3*b1], [a3*b2], [a3*b3]])\n"
     ]
    }
   ],
   "source": [
    "import sympy as sy\n",
    "\n",
    "a = sy.symbols(\"a1:4\")\n",
    "b = sy.symbols(\"b1:4\")\n",
    "J = sy.symbols(\"J(1:4)(1:4)\")\n",
    "\n",
    "\n",
    "aa = sy.Matrix(a)\n",
    "bb = sy.Matrix(b)\n",
    "\n",
    "print(aa)\n",
    "print(bb)\n",
    "\n",
    "JJ = sy.Matrix(J).reshape(3, 3)\n",
    "\n",
    "print(JJ)\n",
    "\n",
    "print((aa.T @ JJ @ bb)[0].expand())\n",
    "print(sy.trace(JJ @ (bb @ aa.T)))\n",
    "print(JJ.reshape(9, 1))\n",
    "print(JJ.reshape(1, 9) @ (aa @ bb.T).reshape(9, 1))\n",
    "print((aa @ bb.T).reshape(9, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'parameters': {'infile': '/Users/danielpozsar/Downloads/nojij/Fe3GeTe2/monolayer/soc/lat3_791/Fe3GeTe2.fdf',\n",
       "  'outfile': './Fe3GeTe2_fdf_test.pickle',\n",
       "  'scf_xcf_orientation': array([0., 0., 1.]),\n",
       "  'ref_xcf_orientations': [{'o': array([1., 0., 0.]),\n",
       "    'vw': array([[0.        , 1.        , 0.        ],\n",
       "           [0.        , 0.        , 1.        ],\n",
       "           [0.        , 0.70710678, 0.70710678]])},\n",
       "   {'o': array([0., 1., 0.]),\n",
       "    'vw': array([[1.        , 0.        , 0.        ],\n",
       "           [0.        , 0.        , 1.        ],\n",
       "           [0.70710678, 0.        , 0.70710678]])},\n",
       "   {'o': array([0., 0., 1.]),\n",
       "    'vw': array([[1.        , 0.        , 0.        ],\n",
       "           [0.        , 1.        , 0.        ],\n",
       "           [0.70710678, 0.70710678, 0.        ]])}],\n",
       "  'kset': 10,\n",
       "  'kdirs': 'xy',\n",
       "  'ebot': -12.906878959999998,\n",
       "  'eset': 600,\n",
       "  'esetp': 1000.0,\n",
       "  'parallel_solver_for_Gk': False,\n",
       "  'padawan_mode': True,\n",
       "  'parallel_size': 1,\n",
       "  'automatic_ebot': True,\n",
       "  'cell': array([[ 3.79100000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "         [-1.89550000e+00,  3.28310231e+00,  0.00000000e+00],\n",
       "         [ 1.25954923e-15,  2.18160327e-15,  2.05700000e+01]])},\n",
       " 'magnetic_entities': [{'atom': 3,\n",
       "   'l': [2],\n",
       "   'orbital_indices': array([41, 42, 43, 44, 45, 46, 47, 48, 49, 50], dtype=int32),\n",
       "   'spin_box_indices': array([ 82,  83,  84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,\n",
       "           95,  96,  97,  98,  99, 100, 101]),\n",
       "   'tags': ['[3]Fe([2])'],\n",
       "   'xyz': [array([-7.33915874e-06,  4.14927851e-06,  1.16575858e+01])],\n",
       "   'energies': array([[1.17772527, 1.17771974, 1.17772812],\n",
       "          [1.17764184, 1.177567  , 1.17761357],\n",
       "          [1.17892446, 1.17892335, 1.17892284]]),\n",
       "   'K': array([[-0.07483657,  0.00106563, -0.0091485 ],\n",
       "          [ 0.00106563, -0.00553059, -0.00561574],\n",
       "          [-0.0091485 , -0.00561574,  0.        ]]),\n",
       "   'K_consistency': 6.819575556149537e-05},\n",
       "  {'atom': 4,\n",
       "   'l': [2],\n",
       "   'orbital_indices': array([56, 57, 58, 59, 60, 61, 62, 63, 64, 65], dtype=int32),\n",
       "   'spin_box_indices': array([112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124,\n",
       "          125, 126, 127, 128, 129, 130, 131]),\n",
       "   'tags': ['[4]Fe([2])'],\n",
       "   'xyz': [array([-7.32698766e-06,  4.15827452e-06,  8.91242254e+00])],\n",
       "   'energies': array([[1.17774   , 1.17766249, 1.17769662],\n",
       "          [1.1776063 , 1.17760676, 1.17759665],\n",
       "          [1.17893438, 1.17893546, 1.17893594]]),\n",
       "   'K': array([[ 0.00045681, -0.00101949,  0.00987698],\n",
       "          [-0.00101949, -0.07750577,  0.00462617],\n",
       "          [ 0.00987698,  0.00462617,  0.        ]]),\n",
       "   'K_consistency': 7.687732837413641e-05},\n",
       "  {'atom': 5,\n",
       "   'l': [2],\n",
       "   'orbital_indices': array([71, 72, 73, 74, 75, 76, 77, 78, 79, 80], dtype=int32),\n",
       "   'spin_box_indices': array([142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154,\n",
       "          155, 156, 157, 158, 159, 160, 161]),\n",
       "   'tags': ['[5]Fe([2])'],\n",
       "   'xyz': [array([ 1.89546671,  1.09439132, 10.2850027 ])],\n",
       "   'energies': array([[0.82604465, 0.82600318, 0.82602942],\n",
       "          [0.82596082, 0.82597202, 0.82596098],\n",
       "          [0.83074663, 0.83074663, 0.83074663]]),\n",
       "   'K': array([[ 1.12053382e-02, -4.12902668e-06,  5.44290076e-03],\n",
       "          [-4.12902668e-06, -4.14783466e-02, -5.50288071e-03],\n",
       "          [ 5.44290076e-03, -5.50288071e-03,  0.00000000e+00]]),\n",
       "   'K_consistency': 5.2686617336705766e-05}],\n",
       " 'pairs': [{'ai': 0,\n",
       "   'aj': 1,\n",
       "   'Ruc': array([0, 0, 0]),\n",
       "   'dist': 2.745163300331324,\n",
       "   'tags': ['[3]Fe([2])', '[4]Fe([2])'],\n",
       "   'energies': array([[-8.92392323e-02,  1.88623702e-04, -1.88607418e-04,\n",
       "           -8.91422199e-02],\n",
       "          [-8.92537950e-02,  8.30905002e-07, -6.97129733e-07,\n",
       "           -8.91766857e-02],\n",
       "          [-7.49830910e-02, -1.42696825e-07, -1.41821186e-07,\n",
       "           -7.49832472e-02]]),\n",
       "   'J_iso': -84.46304519743875,\n",
       "   'J_S': array([ 2.38307873e+00,  2.40038974e+00,  1.42259005e-04, -6.68876347e-05,\n",
       "          -8.14185600e-06]),\n",
       "   'D': array([ 1.88615560e-01, -7.64017368e-04, -4.37819604e-07]),\n",
       "   'J': array([[-8.20799665e+01,  1.42259005e-04, -6.68876347e-05],\n",
       "          [ 1.42259005e-04, -8.20626555e+01, -8.14185600e-06],\n",
       "          [-6.68876347e-05, -8.14185600e-06, -8.92465137e+01]])},\n",
       "  {'ai': 0,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([0, 0, 0]),\n",
       "   'dist': 2.5835033632437767,\n",
       "   'tags': ['[3]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-0.04075836,  0.00119288, -0.00112921, -0.04080086],\n",
       "          [-0.04056911,  0.00207638, -0.00198342, -0.04056736],\n",
       "          [-0.04367583, -0.00022997, -0.00023025, -0.04340975]]),\n",
       "   'J_iso': -41.630212919435856,\n",
       "   'J_S': array([-0.35834154, -0.60813515,  0.23011103, -0.046484  , -0.03183377]),\n",
       "   'D': array([ 1.16104682e+00, -2.02989980e+00,  1.42005544e-04]),\n",
       "   'J': array([[-4.19885545e+01,  2.30111030e-01, -4.64839962e-02],\n",
       "          [ 2.30111030e-01, -4.22383481e+01, -3.18337650e-02],\n",
       "          [-4.64839962e-02, -3.18337650e-02, -4.06637362e+01]])},\n",
       "  {'ai': 1,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([0, 0, 0]),\n",
       "   'dist': 2.583501767937866,\n",
       "   'tags': ['[4]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-0.04077002, -0.00117329,  0.00113811, -0.04082399],\n",
       "          [-0.04057372, -0.00204387,  0.00197878, -0.0405759 ],\n",
       "          [-0.04367653, -0.00022997, -0.00023026, -0.04341045]]),\n",
       "   'J_iso': -41.63843353643032,\n",
       "   'J_S': array([-0.35473711, -0.61182503,  0.23011277,  0.03254538,  0.01758745]),\n",
       "   'D': array([-1.15570028e+00,  2.01132360e+00,  1.42811452e-04]),\n",
       "   'J': array([[-4.19931707e+01,  2.30112766e-01,  3.25453781e-02],\n",
       "          [ 2.30112766e-01, -4.22502586e+01,  1.75874450e-02],\n",
       "          [ 3.25453781e-02,  1.75874450e-02, -4.06718714e+01]])},\n",
       "  {'ai': 0,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-1, -1,  0]),\n",
       "   'dist': 2.5834973202859075,\n",
       "   'tags': ['[3]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-4.05164545e-02, -2.37526408e-03,  2.29370374e-03,\n",
       "           -4.04788963e-02],\n",
       "          [-4.09074301e-02, -2.05208851e-07, -8.74099273e-08,\n",
       "           -4.09801092e-02],\n",
       "          [-4.32782877e-02,  1.10167960e-07,  1.95301791e-07,\n",
       "           -4.38099217e-02]]),\n",
       "   'J_iso': -41.66184991462266,\n",
       "   'J_S': array([-7.33165557e-01, -2.16742080e-01, -1.52734875e-04,  1.46309389e-04,\n",
       "           4.07801676e-02]),\n",
       "   'D': array([-2.33448391e+00,  5.88994618e-05, -4.25669156e-05]),\n",
       "   'J': array([[-4.23950155e+01, -1.52734875e-04,  1.46309389e-04],\n",
       "          [-1.52734875e-04, -4.18785920e+01,  4.07801676e-02],\n",
       "          [ 1.46309389e-04,  4.07801676e-02, -4.07119423e+01]])},\n",
       "  {'ai': 1,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-1, -1,  0]),\n",
       "   'dist': 2.583495745338251,\n",
       "   'tags': ['[4]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-4.05171468e-02,  2.37529161e-03, -2.29373272e-03,\n",
       "           -4.04795842e-02],\n",
       "          [-4.08811902e-02,  1.57284973e-07, -4.45157439e-07,\n",
       "           -4.09421171e-02],\n",
       "          [-4.32789799e-02,  1.09351818e-07,  1.96420584e-07,\n",
       "           -4.38106181e-02]]),\n",
       "   'J_iso': -41.651606039885934,\n",
       "   'J_S': array([-7.24761554e-01, -2.27675997e-01, -1.52886201e-04,  1.43936233e-04,\n",
       "          -4.07794409e-02]),\n",
       "   'D': array([ 2.33451216e+00, -3.01221206e-04, -4.35343826e-05]),\n",
       "   'J': array([[-4.23763676e+01, -1.52886201e-04,  1.43936233e-04],\n",
       "          [-1.52886201e-04, -4.18792820e+01, -4.07794409e-02],\n",
       "          [ 1.43936233e-04, -4.07794409e-02, -4.06991685e+01]])},\n",
       "  {'ai': 0,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-1,  0,  0]),\n",
       "   'dist': 2.583541444641373,\n",
       "   'tags': ['[3]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-0.04075762,  0.00117311, -0.00113797, -0.04081104],\n",
       "          [-0.04055754, -0.00207635,  0.00198375, -0.04055565],\n",
       "          [-0.04366146,  0.00022983,  0.00023003, -0.04339595]]),\n",
       "   'J_iso': -41.62320943278901,\n",
       "   'J_S': array([-0.35258884, -0.61303902, -0.22992743,  0.04629825, -0.01757271]),\n",
       "   'D': array([ 1.15553939e+00,  2.03005138e+00, -1.00670740e-04]),\n",
       "   'J': array([[-4.19757983e+01, -2.29927428e-01,  4.62982483e-02],\n",
       "          [-2.29927428e-01, -4.22362485e+01, -1.75727100e-02],\n",
       "          [ 4.62982483e-02, -1.75727100e-02, -4.06575816e+01]])},\n",
       "  {'ai': 1,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-1,  0,  0]),\n",
       "   'dist': 2.5835398672184064,\n",
       "   'tags': ['[4]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-0.04074738, -0.00119276,  0.00112911, -0.04078932],\n",
       "          [-0.04056216,  0.00204384, -0.00197854, -0.04056432],\n",
       "          [-0.04366217,  0.00022983,  0.00023003, -0.04339665]]),\n",
       "   'J_iso': -41.62033567430795,\n",
       "   'J_S': array([-0.36015198, -0.6054113 , -0.22992935, -0.03265044,  0.03182307]),\n",
       "   'D': array([-1.16093344e+00, -2.01119076e+00, -9.97127859e-05]),\n",
       "   'J': array([[-4.19804877e+01, -2.29929354e-01, -3.26504440e-02],\n",
       "          [-2.29929354e-01, -4.22257470e+01,  3.18230690e-02],\n",
       "          [-3.26504440e-02,  3.18230690e-02, -4.06547724e+01]])},\n",
       "  {'ai': 1,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-2,  0,  0]),\n",
       "   'dist': 5.951322298958084,\n",
       "   'tags': ['[4]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-1.72611605e-03, -4.36527901e-05,  1.26576311e-04,\n",
       "           -1.82722454e-03],\n",
       "          [-1.80427815e-03, -2.63038710e-04,  2.20144926e-04,\n",
       "           -1.73941292e-03],\n",
       "          [-1.75263911e-03,  1.36519152e-03, -1.52871309e-03,\n",
       "           -1.52707483e-03]]),\n",
       "   'J_iso': -1.7294575999129687,\n",
       "   'J_S': array([ 0.09621372, -0.06047422,  0.08176078,  0.02144689, -0.04146176]),\n",
       "   'D': array([-0.08511455,  0.24159182,  1.4469523 ]),\n",
       "   'J': array([[-1.63324388,  0.08176078,  0.02144689],\n",
       "          [ 0.08176078, -1.78993182, -0.04146176],\n",
       "          [ 0.02144689, -0.04146176, -1.7651971 ]])},\n",
       "  {'ai': 1,\n",
       "   'aj': 2,\n",
       "   'Ruc': array([-3,  0,  0]),\n",
       "   'dist': 9.638732176310562,\n",
       "   'tags': ['[4]Fe([2])', '[5]Fe([2])'],\n",
       "   'energies': array([[-7.80160843e-04,  3.56600605e-05, -6.10780433e-05,\n",
       "           -7.42312414e-04],\n",
       "          [-7.94983595e-04, -5.77431729e-05,  5.08380963e-05,\n",
       "           -8.34625148e-04],\n",
       "          [ 2.17237059e-04, -1.43586609e-04,  1.74512092e-04,\n",
       "            2.13496975e-04]]),\n",
       "   'J_iso': -0.45355799416161025,\n",
       "   'J_S': array([ 0.14299391,  0.19102032, -0.01546274,  0.00345254,  0.01270899]),\n",
       "   'D': array([ 0.04836905,  0.05429063, -0.15904935]),\n",
       "   'J': array([[-0.31056409, -0.01546274,  0.00345254],\n",
       "          [-0.01546274, -0.26253768,  0.01270899],\n",
       "          [ 0.00345254,  0.01270899, -0.78757222]])}],\n",
       " 'runtime': {'start_time': 0.881972083,\n",
       "  'setup_time': 0.979693,\n",
       "  'H_and_XCF_time': 1.361691666,\n",
       "  'site_and_pair_dictionaries_time': 1.40248525,\n",
       "  'k_set_time': 1.432338791,\n",
       "  'reference_rotations_time': 1.68477325,\n",
       "  'green_function_inversion_time': 281.406242666,\n",
       "  'end_time': 281.949395875}}"
      ]
     },
     "execution_count": 151,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0., 0., 1.]), array([0., 1., 0.]))"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scf_xcf_orientation = out[\"parameters\"][\"scf_xcf_orientation\"]\n",
    "\n",
    "ref_xcf_orientations = out[\"parameters\"][\"ref_xcf_orientations\"]\n",
    "\n",
    "scf_xcf_orientation\n",
    "\n",
    "from grogupy import *\n",
    "from grogupy.utilities import *\n",
    "\n",
    "\n",
    "def generate_perpendicular(orientation, reference):\n",
    "    x = np.array([1, 0, 0])\n",
    "    y = np.array([0, 1, 0])\n",
    "    z = np.array([0, 0, 1])\n",
    "\n",
    "    R_to_reference = RotMa2b(z, reference)\n",
    "    x = R_to_reference @ x\n",
    "    y = R_to_reference @ y\n",
    "    z = R_to_reference @ z\n",
    "\n",
    "    R_to_orientation = RotMa2b(z, orientation)\n",
    "    x = R_to_orientation @ x\n",
    "    y = R_to_orientation @ y\n",
    "\n",
    "    return x, y\n",
    "\n",
    "\n",
    "generate_perpendicular(ref_xcf_orientations[0][\"o\"], scf_xcf_orientation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1. 0. 0.]\n",
      "[[0.         1.         0.        ]\n",
      " [0.         0.         1.        ]\n",
      " [0.         0.70710678 0.70710678]]\n",
      "(array([0., 0., 1.]), array([0., 1., 0.]))\n",
      "[0. 1. 0.]\n",
      "[[1.         0.         0.        ]\n",
      " [0.         0.         1.        ]\n",
      " [0.70710678 0.         0.70710678]]\n",
      "(array([-1.,  0.,  0.]), array([ 0.,  0., -1.]))\n",
      "[0. 0. 1.]\n",
      "[[1.         0.         0.        ]\n",
      " [0.         1.         0.        ]\n",
      " [0.70710678 0.70710678 0.        ]]\n",
      "(array([-1.,  0.,  0.]), array([0., 1., 0.]))\n"
     ]
    }
   ],
   "source": [
    "for ref in ref_xcf_orientations:\n",
    "    print(ref[\"o\"])\n",
    "    print(ref[\"vw\"])\n",
    "    print(generate_perpendicular(ref[\"o\"], scf_xcf_orientation))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [
    {
     "ename": "LinAlgError",
     "evalue": "Singular matrix",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mLinAlgError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[182], line 25\u001b[0m\n\u001b[1;32m     22\u001b[0m                 M \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mouter(ekkprime, ekprimek)\n\u001b[1;32m     23\u001b[0m                 V \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m E[j] \u001b[38;5;241m*\u001b[39m ekkprime\u001b[38;5;241m.\u001b[39mflatten()\n\u001b[0;32m---> 25\u001b[0m K \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinalg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43mM\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mV\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36msolve\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/Documents/oktatás/elte/phd/grogu_project/.venv/lib/python3.9/site-packages/numpy/linalg/linalg.py:386\u001b[0m, in \u001b[0;36msolve\u001b[0;34m(a, b)\u001b[0m\n\u001b[1;32m    384\u001b[0m signature \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDD->D\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m isComplexType(t) \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdd->d\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    385\u001b[0m extobj \u001b[38;5;241m=\u001b[39m get_linalg_error_extobj(_raise_linalgerror_singular)\n\u001b[0;32m--> 386\u001b[0m r \u001b[38;5;241m=\u001b[39m \u001b[43mgufunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msignature\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msignature\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextobj\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mextobj\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m wrap(r\u001b[38;5;241m.\u001b[39mastype(result_t, copy\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m))\n",
      "File \u001b[0;32m~/Documents/oktatás/elte/phd/grogu_project/.venv/lib/python3.9/site-packages/numpy/linalg/linalg.py:89\u001b[0m, in \u001b[0;36m_raise_linalgerror_singular\u001b[0;34m(err, flag)\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_raise_linalgerror_singular\u001b[39m(err, flag):\n\u001b[0;32m---> 89\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m LinAlgError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSingular matrix\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mLinAlgError\u001b[0m: Singular matrix"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "################################################################################\n",
    "#################################### INPUT #####################################\n",
    "################################################################################\n",
    "path = (\n",
    "    \"/Users/danielpozsar/Downloads/nojij/Fe3GeTe2/monolayer/soc/lat3_791/Fe3GeTe2.fdf\"\n",
    ")\n",
    "outfile = \"./Fe3GeTe2_notebook\"\n",
    "\n",
    "# this information needs to be given at the input!!\n",
    "scf_xcf_orientation = np.array([0, 0, 1])  # z\n",
    "# list of reference directions for around which we calculate the derivatives\n",
    "# o is the quantization axis, v and w are two axes perpendicular to it\n",
    "# at this moment the user has to supply o,v,w on the input.\n",
    "# we can have some default for this\n",
    "ref_xcf_orientations = [\n",
    "    dict(o=np.array([1, 0, 0]), vw=[np.array([0, 1, 0]), np.array([0, 0, 1])]),\n",
    "    dict(o=np.array([0, 1, 0]), vw=[np.array([1, 0, 0]), np.array([0, 0, 1])]),\n",
    "    dict(o=np.array([0, 0, 1]), vw=[np.array([1, 0, 0]), np.array([0, 1, 0])]),\n",
    "]\n",
    "magnetic_entities = [\n",
    "    dict(atom=3, l=2),\n",
    "    dict(atom=4, l=2),\n",
    "    dict(atom=5, l=2),\n",
    "]\n",
    "pairs = [\n",
    "    dict(ai=0, aj=1, Ruc=np.array([0, 0, 0])),\n",
    "    dict(ai=0, aj=2, Ruc=np.array([0, 0, 0])),\n",
    "    dict(ai=1, aj=2, Ruc=np.array([0, 0, 0])),\n",
    "    dict(ai=0, aj=2, Ruc=np.array([-1, -1, 0])),\n",
    "    dict(ai=1, aj=2, Ruc=np.array([-1, -1, 0])),\n",
    "    dict(ai=0, aj=2, Ruc=np.array([-1, 0, 0])),\n",
    "    dict(ai=1, aj=2, Ruc=np.array([-1, 0, 0])),\n",
    "    dict(ai=1, aj=2, Ruc=np.array([-2, 0, 0])),\n",
    "    dict(ai=1, aj=2, Ruc=np.array([-3, 0, 0])),\n",
    "]\n",
    "\n",
    "# Brilloun zone sampling and Green function contour integral\n",
    "kset = 100\n",
    "kdirs = \"xy\"\n",
    "ebot = -13\n",
    "eset = 300\n",
    "esetp = 1000\n",
    "################################################################################\n",
    "#################################### INPUT #####################################\n",
    "################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300\n",
      "100\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "eset_space = np.linspace(0, eset - 1, eset, dtype=int)\n",
    "kset_space = np.linspace(0, kset - 1, kset, dtype=int)\n",
    "orient_space = np.linspace(\n",
    "    0, len(ref_xcf_orientations) - 1, len(ref_xcf_orientations), dtype=int\n",
    ")\n",
    "print(len(kset_space))\n",
    "print(len(orient_space))\n",
    "print(len(eset_space))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "combinations = product(eset_space, eset_space, eset_space)\n",
    "asd = np.array_split(np.array(list(combinations)), 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================================================================================================\n",
      "SLURM job ID not found.\n",
      "Input file: \n",
      "/Users/danielpozsar/Downloads/nojij/Fe3GeTe2/monolayer/soc/lat3_791/Fe3GeTe2.fdf\n",
      "Output file: \n",
      "./Fe3GeTe2_notebook.pickle\n",
      "Number of nodes in the parallel cluster:  1\n",
      "================================================================================================================================================================\n",
      "Cell [Ang]: \n",
      "[[ 3.79100000e+00  0.00000000e+00  0.00000000e+00]\n",
      " [-1.89550000e+00  3.28310231e+00  0.00000000e+00]\n",
      " [ 1.25954923e-15  2.18160327e-15  2.05700000e+01]]\n",
      "================================================================================================================================================================\n",
      "DFT axis: \n",
      "[0 0 1]\n",
      "Quantization axis and perpendicular rotation directions:\n",
      "[1 0 0]  --»  [array([0, 1, 0]), array([0, 0, 1])]\n",
      "[0 1 0]  --»  [array([1, 0, 0]), array([0, 0, 1])]\n",
      "[0 0 1]  --»  [array([1, 0, 0]), array([0, 1, 0])]\n",
      "================================================================================================================================================================\n",
      "Parameters for the contour integral:\n",
      "Number of k points:  3\n",
      "k point directions:  xy\n",
      "Ebot:  -13\n",
      "Eset:  300\n",
      "Esetp:  1000\n",
      "================================================================================================================================================================\n",
      "Setup done. Elapsed time: 3.275603625 s\n",
      "================================================================================================================================================================\n"
     ]
    }
   ],
   "source": [
    "# MPI parameters\n",
    "comm = MPI.COMM_WORLD\n",
    "size = comm.Get_size()\n",
    "rank = comm.Get_rank()\n",
    "root_node = 0\n",
    "\n",
    "# rename outfile\n",
    "if not outfile.endswith(\".pickle\"):\n",
    "    outfile += \".pickle\"\n",
    "\n",
    "simulation_parameters = dict(\n",
    "    infile=path,\n",
    "    outfile=outfile,\n",
    "    scf_xcf_orientation=scf_xcf_orientation,\n",
    "    ref_xcf_orientations=ref_xcf_orientations,\n",
    "    kset=kset,\n",
    "    kdirs=kdirs,\n",
    "    ebot=ebot,\n",
    "    eset=eset,\n",
    "    esetp=esetp,\n",
    "    parallel_size=size,\n",
    ")\n",
    "\n",
    "# if ebot is not given put it 0.1 eV under the smallest energy\n",
    "if simulation_parameters[\"ebot\"] is None:\n",
    "    try:\n",
    "        eigfile = simulation_parameters[\"infile\"][:-3] + \"EIG\"\n",
    "        simulation_parameters[\"ebot\"] = read_siesta_emin(eigfile) - 0.1\n",
    "    except:\n",
    "        print(\"Could not determine ebot.\")\n",
    "        print(\"Parameter was not given and .EIG file was not found.\")\n",
    "# digestion of the input\n",
    "# read sile\n",
    "fdf = sisl.get_sile(simulation_parameters[\"infile\"])\n",
    "# read in hamiltonian\n",
    "dh = fdf.read_hamiltonian()\n",
    "simulation_parameters[\"cell\"] = fdf.read_geometry().cell\n",
    "\n",
    "# unit cell index\n",
    "uc_in_sc_idx = dh.lattice.sc_index([0, 0, 0])\n",
    "\n",
    "if rank == root_node:\n",
    "    print_parameters(simulation_parameters)\n",
    "    times[\"setup_time\"] = timer()\n",
    "    print(f\"Setup done. Elapsed time: {times['setup_time']} s\")\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hamiltonian and exchange field rotated. Elapsed time: 3.976168875 s\n",
      "================================================================================================================================================================\n"
     ]
    }
   ],
   "source": [
    "hh, ss = build_hh_ss(dh)\n",
    "NO = dh.no\n",
    "\n",
    "# symmetrizing Hamiltonian and overlap matrix to make them hermitian\n",
    "for i in range(dh.lattice.sc_off.shape[0]):\n",
    "    j = dh.lattice.sc_index(-dh.lattice.sc_off[i])\n",
    "    h1, h1d = hh[i], hh[j]\n",
    "    hh[i], hh[j] = (h1 + h1d.T.conj()) / 2, (h1d + h1.T.conj()) / 2\n",
    "    s1, s1d = ss[i], ss[j]\n",
    "    ss[i], ss[j] = (s1 + s1d.T.conj()) / 2, (s1d + s1.T.conj()) / 2\n",
    "\n",
    "\n",
    "# identifying TRS and TRB parts of the Hamiltonian\n",
    "TAUY = np.kron(np.eye(NO), TAU_Y)\n",
    "hTR = np.array([TAUY @ hh[i].conj() @ TAUY for i in range(dh.lattice.nsc.prod())])\n",
    "hTRS = (hh + hTR) / 2\n",
    "hTRB = (hh - hTR) / 2\n",
    "\n",
    "# extracting the exchange field\n",
    "traced = [spin_tracer(hTRB[i]) for i in range(dh.lattice.nsc.prod())]  # equation 77\n",
    "XCF = np.array(\n",
    "    [\n",
    "        np.array([f[\"x\"] / 2 for f in traced]),\n",
    "        np.array([f[\"y\"] / 2 for f in traced]),\n",
    "        np.array([f[\"z\"] / 2 for f in traced]),\n",
    "    ]\n",
    ")  # equation 77\n",
    "\n",
    "\n",
    "# Check if exchange field has scalar part\n",
    "max_xcfs = abs(np.array(np.array([f[\"c\"] / 2 for f in traced]))).max()\n",
    "if max_xcfs > 1e-12:\n",
    "    warnings.warn(\n",
    "        f\"Exchange field has non negligible scalar part. Largest value is {max_xcfs}\"\n",
    "    )\n",
    "\n",
    "if rank == root_node:\n",
    "    times[\"H_and_XCF_time\"] = timer()\n",
    "    print(\n",
    "        f\"Hamiltonian and exchange field rotated. Elapsed time: {times['H_and_XCF_time']} s\"\n",
    "    )\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Site and pair dictionaries created. Elapsed time: 4.091688166 s\n",
      "================================================================================================================================================================\n"
     ]
    }
   ],
   "source": [
    "pairs, magnetic_entities = setup_pairs_and_magnetic_entities(\n",
    "    magnetic_entities, pairs, dh, simulation_parameters\n",
    ")\n",
    "\n",
    "if rank == root_node:\n",
    "    times[\"site_and_pair_dictionaries_time\"] = timer()\n",
    "    print(\n",
    "        f\"Site and pair dictionaries created. Elapsed time: {times['site_and_pair_dictionaries_time']} s\"\n",
    "    )\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k set created. Elapsed time: 22.801980041 s\n",
      "================================================================================================================================================================\n"
     ]
    }
   ],
   "source": [
    "kset = make_kset(\n",
    "    dirs=simulation_parameters[\"kdirs\"], NUMK=simulation_parameters[\"kset\"]\n",
    ")  # generate k space sampling\n",
    "wkset = np.ones(len(kset)) / len(kset)  # generate weights for k points\n",
    "kpcs = np.array_split(kset, size)  # split the k points based on MPI size\n",
    "\n",
    "\n",
    "if rank == root_node:\n",
    "    times[\"k_set_time\"] = timer()\n",
    "    print(f\"k set created. Elapsed time: {times['k_set_time']} s\")\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# make energy contour\n",
    "# we are working in eV now  !\n",
    "# and sisl shifts E_F to 0 !\n",
    "cont = make_contour(\n",
    "    emin=simulation_parameters[\"ebot\"],\n",
    "    enum=simulation_parameters[\"eset\"],\n",
    "    p=simulation_parameters[\"esetp\"],\n",
    ")\n",
    "eran = cont.ze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will contain the three hamiltonians in the reference directions needed to calculate the energy variations upon rotation\n",
    "hamiltonians = []\n",
    "\n",
    "# iterate over the reference directions (quantization axes)\n",
    "for i, orient in enumerate(simulation_parameters[\"ref_xcf_orientations\"]):\n",
    "    # obtain rotated exchange field\n",
    "    R = RotMa2b(simulation_parameters[\"scf_xcf_orientation\"], orient[\"o\"])\n",
    "    rot_XCF = np.einsum(\"ij,jklm->iklm\", R, XCF)\n",
    "    rot_H_XCF = sum(\n",
    "        [np.kron(rot_XCF[i], tau) for i, tau in enumerate([tau_x, tau_y, tau_z])]\n",
    "    )\n",
    "    rot_H_XCF_uc = rot_H_XCF[uc_in_sc_idx]\n",
    "\n",
    "    # obtain total Hamiltonian with the rotated exchange field\n",
    "    rot_H = hTRS + rot_H_XCF  # equation 76\n",
    "\n",
    "    hamiltonians.append(\n",
    "        dict(\n",
    "            orient=orient[\"o\"],\n",
    "            H=rot_H,\n",
    "            GS=np.zeros(\n",
    "                (simulation_parameters[\"eset\"], rot_H.shape[1], rot_H.shape[2]),\n",
    "                dtype=\"complex128\",\n",
    "            ),\n",
    "            GS_tmp=np.zeros(\n",
    "                (simulation_parameters[\"eset\"], rot_H.shape[1], rot_H.shape[2]),\n",
    "                dtype=\"complex128\",\n",
    "            ),\n",
    "        )\n",
    "    )  # store orientation and rotated Hamiltonian\n",
    "\n",
    "    # these are the rotations (for now) perpendicular to the quantization axis\n",
    "    for u in orient[\"vw\"]:\n",
    "        Tu = np.kron(np.eye(NO, dtype=int), tau_u(u))  # section 2.H\n",
    "\n",
    "        Vu1, Vu2 = calc_Vu(rot_H_XCF_uc, Tu)\n",
    "\n",
    "        for mag_ent in magnetic_entities:\n",
    "            idx = mag_ent[\"spin_box_indeces\"]\n",
    "            # fill up the perturbed potentials (for now) based on the on-site projections\n",
    "            mag_ent[\"Vu1\"][i].append(Vu1[:, idx][idx, :])\n",
    "            mag_ent[\"Vu2\"][i].append(Vu2[:, idx][idx, :])\n",
    "\n",
    "if rank == root_node:\n",
    "    times[\"reference_rotations_time\"] = timer()\n",
    "    print(\n",
    "        f\"Rotations done perpendicular to quantization axis. Elapsed time: {times['reference_rotations_time']} s\"\n",
    "    )\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if rank == root_node:\n",
    "    print(\"Starting matrix inversions.\")\n",
    "    print(f\"Total number of k points: {kset.shape[0]}\")\n",
    "    print(f\"Number of energy samples per k point: {simulation_parameters['eset']}\")\n",
    "    print(f\"Total number of directions: {len(hamiltonians)}\")\n",
    "    print(\n",
    "        f\"Total number of matrix inversions: {kset.shape[0] * len(hamiltonians) * simulation_parameters['eset']}\"\n",
    "    )\n",
    "    print(f\"The shape of the Hamiltonian and the Greens function is {NO}x{NO}={NO*NO}\")\n",
    "    # https://stackoverflow.com/questions/70746660/how-to-predict-memory-requirement-for-np-linalg-inv\n",
    "    # memory is O(64 n**2) for complex matrices\n",
    "    memory_size = getsizeof(hamiltonians[0][\"H\"].base) / 1024\n",
    "    print(\n",
    "        f\"Memory taken by a single Hamiltonian is: {getsizeof(hamiltonians[0]['H'].base) / 1024} KB\"\n",
    "    )\n",
    "    print(f\"Expected memory usage per matrix inversion: {memory_size * 32} KB\")\n",
    "    print(\n",
    "        f\"Expected memory usage per k point for parallel inversion: {memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32} KB\"\n",
    "    )\n",
    "    print(\n",
    "        f\"Expected memory usage on root node: {len(np.array_split(kset, size)[0]) * memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32 / 1024} MB\"\n",
    "    )\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )\n",
    "\n",
    "comm.Barrier()\n",
    "# ----------------------------------------------------------------------\n",
    "\n",
    "# make energy contour\n",
    "# we are working in eV now  !\n",
    "# and sisl shifts E_F to 0 !\n",
    "cont = make_contour(\n",
    "    emin=simulation_parameters[\"ebot\"],\n",
    "    enum=simulation_parameters[\"eset\"],\n",
    "    p=simulation_parameters[\"esetp\"],\n",
    ")\n",
    "eran = cont.ze\n",
    "\n",
    "# ----------------------------------------------------------------------\n",
    "# sampling the integrand on the contour and the BZ\n",
    "for k in kpcs[rank]:\n",
    "    wk = wkset[rank]  # weight of k point in BZ integral\n",
    "    # iterate over reference directions\n",
    "    for i, hamiltonian_orientation in enumerate(hamiltonians):\n",
    "        # calculate Greens function\n",
    "        H = hamiltonian_orientation[\"H\"]\n",
    "        HK, SK = hsk(H, ss, dh.sc_off, k)\n",
    "\n",
    "        # solve Greens function sequentially for the energies, because of memory bound\n",
    "        Gk = sequential_GK(HK, SK, eran, simulation_parameters[\"eset\"])\n",
    "\n",
    "        # saving this for total charge\n",
    "        hamiltonian_orientation[\"GS_tmp\"] += Gk @ SK * wk\n",
    "\n",
    "        # store the Greens function slice of the magnetic entities (for now) based on the on-site projections\n",
    "        for mag_ent in magnetic_entities:\n",
    "            mag_ent[\"Gii_tmp\"][i] += (\n",
    "                Gk[:, mag_ent[\"spin_box_indeces\"], :][:, :, mag_ent[\"spin_box_indeces\"]]\n",
    "                * wk\n",
    "            )\n",
    "\n",
    "        for pair in pairs:\n",
    "            # add phase shift based on the cell difference\n",
    "            phase = np.exp(1j * 2 * np.pi * k @ pair[\"Ruc\"].T)\n",
    "\n",
    "            # get the pair orbital sizes from the magnetic entities\n",
    "            ai = magnetic_entities[pair[\"ai\"]][\"spin_box_indeces\"]\n",
    "            aj = magnetic_entities[pair[\"aj\"]][\"spin_box_indeces\"]\n",
    "\n",
    "            # store the Greens function slice of the magnetic entities (for now) based on the on-site projections\n",
    "            pair[\"Gij_tmp\"][i] += Gk[:, ai][..., aj] * phase * wk\n",
    "            pair[\"Gji_tmp\"][i] += Gk[:, aj][..., ai] / phase * wk\n",
    "\n",
    "# summ reduce partial results of mpi nodes\n",
    "for i in range(len(hamiltonians)):\n",
    "    # for total charge\n",
    "    comm.Reduce(hamiltonians[i][\"GS_tmp\"], hamiltonians[i][\"GS\"], root=root_node)\n",
    "\n",
    "    for mag_ent in magnetic_entities:\n",
    "        comm.Reduce(mag_ent[\"Gii_tmp\"][i], mag_ent[\"Gii\"][i], root=root_node)\n",
    "\n",
    "    for pair in pairs:\n",
    "        comm.Reduce(pair[\"Gij_tmp\"][i], pair[\"Gij\"][i], root=root_node)\n",
    "        comm.Reduce(pair[\"Gji_tmp\"][i], pair[\"Gji\"][i], root=root_node)\n",
    "\n",
    "if rank == root_node:\n",
    "    times[\"green_function_inversion_time\"] = timer()\n",
    "    print(\n",
    "        f\"Calculated Greens functions. Elapsed time: {times['green_function_inversion_time']} s\"\n",
    "    )\n",
    "    print(\n",
    "        \"================================================================================================================================================================\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if rank == root_node:\n",
    "    # Calculate total charge\n",
    "    for hamiltonian in hamiltonians:\n",
    "        GS = hamiltonian[\"GS\"]\n",
    "        traced = np.trace((GS), axis1=1, axis2=2)\n",
    "        print(\"Total charge: \", int_de_ke(traced, cont.we))\n",
    "\n",
    "    # iterate over the magnetic entities\n",
    "    for tracker, mag_ent in enumerate(magnetic_entities):\n",
    "        # iterate over the quantization axes\n",
    "        for i, Gii in enumerate(mag_ent[\"Gii\"]):\n",
    "            storage = []\n",
    "            # iterate over the first and second order local perturbations\n",
    "            for Vu1, Vu2 in zip(mag_ent[\"Vu1\"][i], mag_ent[\"Vu2\"][i]):\n",
    "                # The Szunyogh-Lichtenstein formula\n",
    "                traced = np.trace((Vu2 @ Gii + 0.5 * Gii @ Vu1 @ Gii), axis1=1, axis2=2)\n",
    "                # evaluation of the contour integral\n",
    "                storage.append(int_de_ke(traced, cont.we))\n",
    "\n",
    "            # fill up the magnetic entities dictionary with the energies\n",
    "            magnetic_entities[tracker][\"energies\"].append(storage)\n",
    "        # convert to np array\n",
    "        magnetic_entities[tracker][\"energies\"] = np.array(\n",
    "            magnetic_entities[tracker][\"energies\"]\n",
    "        )\n",
    "    print(\"Magnetic entities integrated.\")\n",
    "\n",
    "    # iterate over the pairs\n",
    "    for tracker, pair in enumerate(pairs):\n",
    "        # iterate over the quantization axes\n",
    "        for i, (Gij, Gji) in enumerate(zip(pair[\"Gij\"], pair[\"Gji\"])):\n",
    "            site_i = magnetic_entities[pair[\"ai\"]]\n",
    "            site_j = magnetic_entities[pair[\"aj\"]]\n",
    "\n",
    "            storage = []\n",
    "            # iterate over the first order local perturbations in all possible orientations for the two sites\n",
    "            for Vui in site_i[\"Vu1\"][i]:\n",
    "                for Vuj in site_j[\"Vu1\"][i]:\n",
    "                    # The Szunyogh-Lichtenstein formula\n",
    "                    traced = np.trace((Vui @ Gij @ Vuj @ Gji), axis1=1, axis2=2)\n",
    "                    # evaluation of the contour integral\n",
    "                    storage.append(int_de_ke(traced, cont.we))\n",
    "            # fill up the pairs dictionary with the energies\n",
    "            pairs[tracker][\"energies\"].append(storage)\n",
    "        # convert to np array\n",
    "        pairs[tracker][\"energies\"] = np.array(pairs[tracker][\"energies\"])\n",
    "\n",
    "    print(\"Pairs integrated.\")\n",
    "\n",
    "    # calculate magnetic parameters\n",
    "    for mag_ent in magnetic_entities:\n",
    "        Kxx, Kyy, Kzz, consistency = calculate_anisotropy_tensor(mag_ent)\n",
    "        mag_ent[\"K\"] = np.array([Kxx, Kyy, Kzz]) * sisl.unit_convert(\"eV\", \"meV\")\n",
    "        mag_ent[\"K_consistency\"] = consistency\n",
    "\n",
    "    for pair in pairs:\n",
    "        J_iso, J_S, D, J = calculate_exchange_tensor(pair)\n",
    "        pair[\"J_iso\"] = J_iso * sisl.unit_convert(\"eV\", \"meV\")\n",
    "        pair[\"J_S\"] = J_S * sisl.unit_convert(\"eV\", \"meV\")\n",
    "        pair[\"D\"] = D * sisl.unit_convert(\"eV\", \"meV\")\n",
    "        pair[\"J\"] = J * sisl.unit_convert(\"eV\", \"meV\")\n",
    "\n",
    "    print(\"Magnetic parameters calculated.\")\n",
    "\n",
    "    times[\"end_time\"] = timer()\n",
    "    print(\n",
    "        \"##################################################################### GROGU OUTPUT #############################################################################\"\n",
    "    )\n",
    "\n",
    "    print_parameters(simulation_parameters)\n",
    "    print_atoms_and_pairs(magnetic_entities, pairs)\n",
    "    print_runtime_information(times)\n",
    "\n",
    "    pairs, magnetic_entities = remove_clutter_for_save(pairs, magnetic_entities)\n",
    "    # create output dictionary with all the relevant data\n",
    "    results = dict(\n",
    "        parameters=simulation_parameters,\n",
    "        magnetic_entities=magnetic_entities,\n",
    "        pairs=pairs,\n",
    "        runtime=times,\n",
    "    )\n",
    "\n",
    "    save_pickle(simulation_parameters[\"outfile\"], results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "========================================\n",
    " \n",
    "Atom Angstrom\n",
    "# Label,        x           y           z          Sx           Sy           Sz         #Q           Lx           Ly           Lz           Jx           Jy           Jz\n",
    "--------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
    "Te1         1.8955      1.0943      13.1698     -0.0000     0.0000      -0.1543    # 5.9345       -0.0000      0.0000       -0.0537      -0.0000      0.0000       -0.2080   \n",
    "Te2         1.8955      1.0943      7.4002      0.0000      -0.0000     -0.1543    # 5.9345       0.0000       -0.0000      -0.0537      0.0000       -0.0000      -0.2080   \n",
    "Ge3         -0.0000     2.1887      10.2850     0.0000      0.0000      -0.1605    # 3.1927       -0.0000      0.0000       0.0012       0.0000       0.0000       -0.1593   \n",
    "Fe4         -0.0000     0.0000      11.6576     0.0001      -0.0001     2.0466     # 8.3044       0.0000       -0.0000      0.1606       0.0001       -0.0001      2.2072    \n",
    "Fe5         -0.0000     0.0000      8.9124      -0.0001     0.0001      2.0466     # 8.3044       -0.0000      0.0000       0.1606       -0.0001      0.0001       2.2072    \n",
    "Fe6         1.8955      1.0944      10.2850     0.0000      0.0000      1.5824     # 8.3296       -0.0000      -0.0000      0.0520       -0.0000      0.0000       1.6344    \n",
    "==================================================================================================================================\n",
    " \n",
    "Exchange meV\n",
    "--------------------------------------------------------------------------------\n",
    "# at1     at2   i  j  k    #    d (Ang)\n",
    "--------------------------------------------------------------------------------\n",
    "Fe4     Fe5     0  0  0    #    2.7452\n",
    "Isotropic -82.0854\n",
    "DMI 0.12557 -0.00082199  6.9668e-08\n",
    "Symmetric-anisotropy -0.60237    -0.83842 -0.00032278 -1.2166e-05 -3.3923e-05\n",
    "--------------------------------------------------------------------------------\n",
    "Fe4     Fe6     0  0  0    #    2.5835\n",
    "Isotropic -41.9627\n",
    "DMI 1.1205     -1.9532   0.0018386\n",
    "Symmetric-anisotropy 0.26007 -0.00013243     0.12977   -0.069979   -0.042066\n",
    "--------------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "On-site meV\n",
    "----------------------------------------\n",
    "Fe4\n",
    "0.16339\t0.16068\t0\t0\t0\t0\n",
    "========================================\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}