EigSymComputation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch import linalg as LA\n",
    "from torch_sym3eig import Sym3Eig\n",
    "from data.io import readRaw, ReadScalars, ReadTensors\n",
    "\n",
    "# torch.set_default_dtype(torch.float64)\n",
    "\n",
    "use_cuda = 0\n",
    "use_brain = 1\n",
    "torchdeviceId = torch.device('cuda:0') if use_cuda else 'cpu'\n",
    "torchdtype = torch.double\n",
    "\n",
    "dim = 3\n",
    "if use_brain:\n",
    "  cases = []\n",
    "  cases.append('103818')\n",
    "  cases.append('105923')\n",
    "  out_tensors = []\n",
    "  masks = []\n",
    "  for run_case in cases:\n",
    "    subj = run_case[:6]\n",
    "    outroot='/usr/sci/projects/HCP/Kris/NSFCRCNS/TestResults/working_3d_python/simulation_results/'\n",
    "    outdir = f'{outroot}{run_case}/mineval_0.005_n_3000_s_1.5/'\n",
    "    indir = '/usr/sci/projects/HCP/Kris/NSFCRCNS/prepped_data/' + subj + '/'\n",
    "\n",
    "    tens_file = 'scaled_tensors.nhdr'\n",
    "    out_tens = ReadTensors(outdir + tens_file)\n",
    "    out_tensors.append(out_tens)\n",
    "    mask_file = 'filt_mask.nhdr'\n",
    "    mask = ReadScalars(outdir + mask_file)\n",
    "    masks.append(mask)\n",
    "    \n",
    "  tens_0 = np.zeros((out_tensors[0].shape[0],out_tensors[0].shape[1],out_tensors[0].shape[2],3,3))\n",
    "  tens_1 = np.zeros((out_tensors[1].shape[0],out_tensors[1].shape[1],out_tensors[1].shape[2],3,3))\n",
    "  for xx in range(tens_0.shape[0]):\n",
    "    for yy in range(tens_0.shape[1]):\n",
    "      for zz in range(tens_0.shape[2]):\n",
    "        if masks[0][xx,yy,zz]:\n",
    "          tens_0[xx,yy,zz,0,0] = out_tensors[0][xx,yy,zz,0]\n",
    "          tens_0[xx,yy,zz,0,1] = out_tensors[0][xx,yy,zz,1]\n",
    "          tens_0[xx,yy,zz,1,0] = out_tensors[0][xx,yy,zz,1]\n",
    "          tens_0[xx,yy,zz,0,2] = out_tensors[0][xx,yy,zz,2]\n",
    "          tens_0[xx,yy,zz,2,0] = out_tensors[0][xx,yy,zz,2]\n",
    "          tens_0[xx,yy,zz,1,1] = out_tensors[0][xx,yy,zz,3]\n",
    "          tens_0[xx,yy,zz,1,2] = out_tensors[0][xx,yy,zz,4]\n",
    "          tens_0[xx,yy,zz,2,1] = out_tensors[0][xx,yy,zz,4]\n",
    "          tens_0[xx,yy,zz,2,2] = out_tensors[0][xx,yy,zz,5]\n",
    "        else:\n",
    "          tens_0[xx,yy,zz,0,0] = 1\n",
    "          tens_0[xx,yy,zz,1,1] = 1\n",
    "          tens_0[xx,yy,zz,2,2] = 1\n",
    "        if masks[1][xx,yy,zz]:\n",
    "          tens_1[xx,yy,zz,0,0] = out_tensors[1][xx,yy,zz,0]\n",
    "          tens_1[xx,yy,zz,0,1] = out_tensors[1][xx,yy,zz,1]\n",
    "          tens_1[xx,yy,zz,1,0] = out_tensors[1][xx,yy,zz,1]\n",
    "          tens_1[xx,yy,zz,0,2] = out_tensors[1][xx,yy,zz,2]\n",
    "          tens_1[xx,yy,zz,2,0] = out_tensors[1][xx,yy,zz,2]\n",
    "          tens_1[xx,yy,zz,1,1] = out_tensors[1][xx,yy,zz,3]\n",
    "          tens_1[xx,yy,zz,1,2] = out_tensors[1][xx,yy,zz,4]\n",
    "          tens_1[xx,yy,zz,2,1] = out_tensors[1][xx,yy,zz,4]\n",
    "          tens_1[xx,yy,zz,2,2] = out_tensors[1][xx,yy,zz,5]\n",
    "        else:\n",
    "          tens_1[xx,yy,zz,0,0] = 1\n",
    "          tens_1[xx,yy,zz,1,1] = 1\n",
    "          tens_1[xx,yy,zz,2,2] = 1\n",
    "        \n",
    "\n",
    "  if use_cuda:\n",
    "    g0 = torch.inverse(torch.from_numpy(tens_0).cuda().float()).contiguous()\n",
    "    g1 = torch.inverse(torch.from_numpy(tens_1).cuda().float()).contiguous()\n",
    "\n",
    "  else:\n",
    "    g0 = torch.inverse(torch.from_numpy(tens_0).float()).contiguous()\n",
    "    g1 = torch.inverse(torch.from_numpy(tens_1).float()).contiguous()\n",
    "else:\n",
    "  N1, N2, N3 = 100, 100, 100\n",
    "\n",
    "  Gr = torch.randn(2,N1,N2,N3,dim,dim,dtype=torchdtype, device=torchdeviceId)\n",
    "  G = torch.einsum(\"...ks,...ts->...kt\",[Gr, Gr])\n",
    "  g0 = G[0]\n",
    "  g1 = G[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([145, 174, 145, 3, 3]), torch.Size([145, 174, 145, 3, 3]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g0.size(), g1.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.3359e-05) tensor(0.0022)\n",
      "tensor(1.1630e-05) tensor(0.0014)\n"
     ]
    }
   ],
   "source": [
    "# check if g0 and g1 are positive definite symmetric\n",
    "\n",
    "print(torch.norm(g0 - torch.einsum(\"...ij->...ji\",[g0])), torch.min(torch.det(g0)))\n",
    "print(torch.norm(g1 - torch.einsum(\"...ij->...ji\",[g1])), torch.min(torch.det(g1)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print out the nonPDS matrix and the determinant\n",
    "if torch.min(torch.det(g0))<0:\n",
    "    g = g0\n",
    "    detg = torch.det(g)\n",
    "    I1,I2,I3 = (detg0==torch.min(detg)).nonzero().reshape(-1)\n",
    "    print('g0 is not PD')\n",
    "    print(g0[I1,I2,I3])\n",
    "    print('The det is', g0[I1,I2,I3].det().item())\n",
    "elif torch.min(torch.det(g1))<0:\n",
    "    g = g1\n",
    "    detg = torch.det(g)\n",
    "    I1,I2,I3 = (detg0==torch.min(detg)).nonzero().reshape(-1)\n",
    "    print('g1 is not PD')\n",
    "    print(g1[I1,I2,I3])\n",
    "    print('The det is',g1[I1,I2,I3].det())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "A, B = g0, g1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.26 s, sys: 1.91 s, total: 8.17 s\n",
      "Wall time: 1.34 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "G = LA.cholesky(A)\n",
    "G_inv = torch.inverse(G)\n",
    "W = torch.einsum(\"...ik,...ks,...js->...ij\",[G_inv, B, G_inv])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(W.device)\n",
    "print(W.is_contiguous())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#W = torch.from_numpy(ReadScalars('/home/sci/hdai/Shared/W.nhdr')).permute((2,0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([410000, 3, 3])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#W = WFull[83613]\n",
    "W.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 1: Use torch.symeig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.12 s, sys: 78.5 ms, total: 4.2 s\n",
      "Wall time: 3.55 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sci/kris/Software/py3_venv_nsfcrcns/lib64/python3.6/site-packages/ipykernel_launcher.py:1: UserWarning: torch.symeig is deprecated in favor of torch.linalg.eigh and will be removed in a future PyTorch release.\n",
      "The default behavior has changed from using the upper triangular portion of the matrix by default to using the lower triangular portion.\n",
      "L, _ = torch.symeig(A, upper=upper)\n",
      "should be replaced with\n",
      "L = torch.linalg.eigvalsh(A, UPLO='U' if upper else 'L')\n",
      "and\n",
      "L, V = torch.symeig(A, eigenvectors=True)\n",
      "should be replaced with\n",
      "L, V = torch.linalg.eigh(A, UPLO='U' if upper else 'L') (Triggered internally at  ../aten/src/ATen/native/BatchLinearAlgebra.cpp:2500.)\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "Lamb1,Lamb1v = torch.symeig(W, eigenvectors=True)\n",
    "trKsquare1 = torch.sum(torch.log(Lamb1)**2,(-1))\n",
    "# Cpu time 1.67s\n",
    "# GPU time 22 min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([145, 174, 145, 3, 3])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.85 s, sys: 21 ms, total: 3.87 s\n",
      "Wall time: 3.43 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "W_max = torch.max(W.reshape((*W.shape[0:3],9)),3)[0]\n",
    "scaled_W = torch.einsum('...ij,...->...ij',W, 1.0 / W_max)\n",
    "Lamb1Alt = torch.einsum('...,...i->...i', W_max, torch.symeig(scaled_W, eigenvectors=True)[0])\n",
    "trKsquare1Alt = torch.sum(torch.log(Lamb1Alt)**2,(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 2: Use sym3eig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 849 ms, sys: 127 ms, total: 976 ms\n",
      "Wall time: 197 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "Lamb2,Lamb2v = Sym3Eig.apply(W.reshape((-1,3,3)))\n",
    "#Lamb2r = Sym3Eig.apply(W.reshape((-1,3,3)))\n",
    "#if use_brain:\n",
    "#  #Lamb2 = torch.abs(Lamb2r[0].reshape((145,174,145,3)))\n",
    "#  Lamb2 = Lamb2r[0].reshape((145,174,145,3))\n",
    "#else:\n",
    "#  Lamb2 = torch.abs(Lamb2r[0].reshape((100,100,100,3)))\n",
    "#s3v = s3vr.reshape((145,174,145,3,3))\n",
    "trKsquare2 = torch.sum(torch.log(Lamb2)**2,(-1))\n",
    "# CPU time: 203 ms\n",
    "# GPU time: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1., 1., 1.]) tensor([[1., -0., 0.],\n",
      "        [0., 1., 0.],\n",
      "        [0., 0., 1.]])\n",
      "tensor([1., 1., 1.]) tensor([[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "print(Lamb1[0,0,0,:],Lamb1v[0,0,0,:,:])\n",
    "print(Lamb2.reshape((145,174,145,3))[0,0,0,:],Lamb2v.reshape((145,174,145,3,3))[0,0,0,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method3: use torchvectorized vlinalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvectorized import vlinalg as TV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([100, 100, 100, 3, 3])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "def checkMag(lamb,vec):\n",
    "  print(lamb.shape, vec.shape)\n",
    "  #lambmax = torch.max(torch.stack((torch.max(lamb*lamb,-1)[0],torch.max(lamb,-1)[0])),0)[0]\n",
    "  lambmax = torch.max(lamb*lamb,-1)[0]\n",
    "  #nrm = (LA.norm(vec,dim=len(vec.shape)-1) **2)\n",
    "  #thresh = 2**8 * 1e-9 * (lambmax**2)\n",
    "  #nrm = vec[:,:,:,:,0] * vec[:,:,:,:,0] + vec[:,:,:,:,1] * vec[:,:,:,:,1] + vec[:,:,:,:,2] * vec[:,:,:,:,2]\n",
    "  #nrm = torch.sqrt(vec[:,:,:,:,0] * vec[:,:,:,:,0] + vec[:,:,:,:,1] * vec[:,:,:,:,1] + vec[:,:,:,:,2] * vec[:,:,:,:,2]).to(lamb.device)\n",
    "  nrm = torch.sqrt(vec[:,:,:,0,:] * vec[:,:,:,0,:] + vec[:,:,:,1,:] * vec[:,:,:,1,:] + vec[:,:,:,2,:] * vec[:,:,:,2,:]).to(lamb.device)\n",
    "  thresh = 2**4 * 1e-9 * lambmax\n",
    "  print(lambmax.device,nrm.device,thresh.device,vec.device)\n",
    "  #return((torch.where(nrm[:,:,:,0] <= thresh), torch.where(nrm[:,:,:,1] <= thresh), torch.where(nrm[:,:,:,2] <= thresh)))\n",
    "  #return((torch.where(nrm[:,:,:,0] <= thresh), torch.where(nrm[:,:,:,1] <= thresh), torch.where(nrm[:,:,:,2] <= thresh)))\n",
    "  idx1 = torch.where(nrm[:,:,:,0] <= thresh)\n",
    "  idx2 = torch.where(nrm[:,:,:,1] <= thresh)\n",
    "  idx3 = torch.where(nrm[:,:,:,2] <= thresh)\n",
    "  idx = (torch.cat((idx1[0], idx2[0], idx3[0])), \n",
    "         torch.cat((idx1[1], idx2[1], idx3[1])), \n",
    "         torch.cat((idx1[2], idx2[2], idx3[2])))\n",
    "  return(idx)\n",
    "                    \n",
    "def checkMag2(lamb):\n",
    "  #lambmax = torch.max(torch.stack((torch.max(lamb*lamb,-1)[0],torch.max(lamb,-1)[0])),0)[0]\n",
    "  lambdiff = lamb[:,:,:,2] / lamb[:,:,:,0]\n",
    "  return(torch.where(lamb[:,:,:,0] < 1)) # 1e-3: error 0.008, 1e-1: error 0.004, 1: error 0.0019\n",
    "def checkMag3(lamb,A):\n",
    "    # From hybrid method at http://www.mpi-hd.mpg.de/personalhomes/globes/3x3/\n",
    "    # as described here: https://arxiv.org/pdf/physics/0610206.pdf\n",
    "    abslamb = torch.abs(lamb)\n",
    "    t = abslamb[:,:,:,2]\n",
    "    idx = torch.where(abslamb[:,:,:,1] > t)\n",
    "    print(len(idx[0]))\n",
    "    t[idx[0],idx[1],idx[2]] = abslamb[idx[0],idx[1],idx[2],1]\n",
    "    idx = torch.where(abslamb[:,:,:,0] > t)\n",
    "    print(len(idx[0]))\n",
    "    t[idx[0],idx[1],idx[2]] = abslamb[idx[0],idx[1],idx[2],0]\n",
    "    \n",
    "    u = t\n",
    "    idx = torch.where(t>=1.0)\n",
    "    u[idx[0],idx[1],idx[2]] = t[idx[0],idx[1],idx[2]] * t[idx[0],idx[1],idx[2]]\n",
    "    sqru = u*u\n",
    "    error = 256*1e-9*sqru\n",
    "    \n",
    "    q01 = A[:,:,:,0,1] * A[:,:,:,1,2] - A[:,:,:,0,2] * A[:,:,:,1,1]\n",
    "    q11 = A[:,:,:,0,2] * A[:,:,:,0,1] - A[:,:,:,1,2] * A[:,:,:,0,0]\n",
    "    #q21 = A[:,:,:,0,1] * A[:,:,:,0,1]\n",
    "    q00 = q01 + A[:,:,:,0,2] * lamb[:,:,:,2]\n",
    "    q10 = q11 + A[:,:,:,1,2] * lamb[:,:,:,2]\n",
    "    q20 = (A[:,:,:,0,0] - lamb[:,:,:,2]) * (A[:,:,:,1,1] - lamb[:,:,:,2]) - A[:,:,:,0,1] * A[:,:,:,0,1]\n",
    "    nrm = q00*q00 + q10*q10 + q20*q20\n",
    "\n",
    "    idx0 = torch.where(nrm <= error)\n",
    "    \n",
    "    # Version 2 of this check includes the following to produce idx1:\n",
    "    \n",
    "    # Following 3 lines only needed for computing 3rd eigenvector\n",
    "    #q00_2 = q00 * nrm\n",
    "    #q10_2 = q10 * nrm\n",
    "    #q20_2 = q20 * nrm\n",
    "    \n",
    "    q01_2 = q01 + A[:,:,:,0,2] * lamb[:,:,:,1]\n",
    "    q11_2 = q11 + A[:,:,:,1,2] * lamb[:,:,:,1]\n",
    "    q21 = (A[:,:,:,0,0] - lamb[:,:,:,1]) * (A[:,:,:,1,1] - lamb[:,:,:,1]) - A[:,:,:,0,1] * A[:,:,:,0,1]\n",
    "    nrm_2 = q01_2*q01_2 + q11_2*q11_2 + q21*q21\n",
    "    \n",
    "    idx1 = torch.where(nrm_2 <= error)\n",
    "    \n",
    "    idx = (torch.cat((idx0[0], idx1[0])), \n",
    "           torch.cat((idx0[1], idx1[1])), \n",
    "           torch.cat((idx0[2], idx1[2])))\n",
    " \n",
    "    return(idx)   \n",
    "    ##nrm = torch.sqrt(q00*q00 + q10*q10 + q20*q20)\n",
    "    #return(torch.where(nrm <= 256*1e-9*sqru))\n",
    "    ##return(torch.where(nrm <= 16*1e-4*u))\n",
    "def checkMag4(lamb,vec):\n",
    "  print(lamb.shape, vec.shape)\n",
    "  #lambmax = torch.max(torch.stack((torch.max(lamb*lamb,-1)[0],torch.max(lamb,-1)[0])),0)[0]\n",
    "  lambmax = torch.max(lamb,dim=-1)[0]\n",
    "  u = lambmax\n",
    "  idx = torch.where(lambmax>=1.0)\n",
    "  u[idx[0],idx[1],idx[2]] = lambmax[idx[0],idx[1],idx[2]]*lambmax[idx[0],idx[1],idx[2]]\n",
    "  sqru = u*u\n",
    "  error = 256*1e-9*sqru \n",
    "\n",
    "  #nrm = (vec[:,:,:,:,0] * vec[:,:,:,:,0] + vec[:,:,:,:,1] * vec[:,:,:,:,1] + vec[:,:,:,:,2] * vec[:,:,:,:,2]).to(lamb.device)\n",
    "  nrm = (vec[:,:,:,0,:] * vec[:,:,:,0,:] + vec[:,:,:,1,:] * vec[:,:,:,1,:] + vec[:,:,:,2,:] * vec[:,:,:,2,:]).to(lamb.device)\n",
    "  #return((torch.where(nrm[:,:,:,0] <= thresh), torch.where(nrm[:,:,:,1] <= thresh), torch.where(nrm[:,:,:,2] <= thresh)))\n",
    "  #return((torch.where(nrm[:,:,:,0] <= thresh), torch.where(nrm[:,:,:,1] <= thresh), torch.where(nrm[:,:,:,2] <= thresh)))\n",
    "  # Do we need to check 3rd eigenvector?\n",
    "  idx1 = torch.where(nrm[:,:,:,0] <= error)\n",
    "  idx2 = torch.where(nrm[:,:,:,1] <= error)\n",
    "  idx3 = torch.where(nrm[:,:,:,2] <= error)\n",
    "  print(len(idx1[0]),len(idx2[0]),len(idx3[0]))\n",
    "  idx = (torch.cat((idx1[0], idx2[0], idx3[0])), \n",
    "         torch.cat((idx1[1], idx2[1], idx3[1])), \n",
    "         torch.cat((idx1[2], idx2[2], idx3[2])))\n",
    "  return(idx)\n",
    "                    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 100, 100, 3, 3])\n",
      "1000000\n",
      "torch.Size([3, 3, 100, 100, 100])\n",
      "tensor([[[[[ 8.5207e-01,  2.2181e-01,  3.0813e+00,  ...,  6.0611e+00,\n",
      "             7.1486e-01,  1.3881e+00],\n",
      "           [ 2.0328e+00,  1.1288e-01,  2.4111e+00,  ...,  2.4533e+00,\n",
      "             2.1193e+00,  3.5957e+00],\n",
      "           [ 7.6886e-01,  3.8845e+00,  3.1359e+00,  ...,  2.8681e-01,\n",
      "             1.7287e-01,  7.0561e-01],\n",
      "           ...,\n",
      "           [ 2.3204e-01,  4.0186e-01,  4.2897e+00,  ...,  1.6221e+01,\n",
      "             7.0160e-01,  3.1719e+00],\n",
      "           [ 2.1442e+00,  1.7903e+01,  2.0547e-01,  ...,  1.2468e+00,\n",
      "             3.5853e+00,  1.7907e+01],\n",
      "           [ 5.6605e-01,  1.8817e-01,  9.3034e-02,  ...,  2.0682e+00,\n",
      "             1.8257e+00,  2.8736e-01]],\n",
      "\n",
      "          [[ 4.4475e-01,  1.6184e-01,  2.6065e-01,  ...,  4.6684e+00,\n",
      "             1.3017e-01,  1.4833e+00],\n",
      "           [ 4.6167e-01,  1.2411e+00,  1.2574e+00,  ...,  3.4412e-01,\n",
      "             6.1798e+00,  1.2440e+00],\n",
      "           [ 3.6992e-01,  2.7505e-01,  1.2931e+00,  ...,  4.0363e-01,\n",
      "             8.9933e-01,  7.2631e+00],\n",
      "           ...,\n",
      "           [ 1.2628e-01,  4.2130e-01,  2.2439e+00,  ...,  1.2927e-01,\n",
      "             2.4000e+00,  3.0692e+00],\n",
      "           [ 3.5010e-01,  7.7817e+00,  5.5934e-01,  ...,  5.7913e-01,\n",
      "             2.4126e+00,  3.2553e+00],\n",
      "           [ 3.5412e+00,  4.6179e-01,  4.1850e+00,  ...,  9.2724e-01,\n",
      "             6.5389e-02,  2.0445e-01]],\n",
      "\n",
      "          [[ 2.5127e-01,  1.1403e+01,  2.0668e+00,  ...,  4.9997e-01,\n",
      "             1.1671e+00,  5.8686e-01],\n",
      "           [ 6.9564e+00,  7.0310e-01,  1.9865e+00,  ...,  2.9227e-01,\n",
      "             1.6382e-01,  2.4397e+00],\n",
      "           [ 2.8736e+00,  6.5839e-01,  5.0876e-01,  ...,  1.9854e+00,\n",
      "             1.3376e+00,  4.0979e-01],\n",
      "           ...,\n",
      "           [ 2.5764e+00,  7.5256e-01,  2.5369e+00,  ...,  2.1087e+01,\n",
      "             2.1380e+00,  5.9209e+00],\n",
      "           [ 5.0388e-01,  6.8572e+00,  9.8878e-01,  ...,  7.8460e-02,\n",
      "             1.9176e+00,  4.0071e-01],\n",
      "           [ 4.2107e+00,  1.8452e+00,  6.0228e+00,  ...,  1.2381e+01,\n",
      "             1.3066e+01,  1.0803e+00]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[ 1.5652e+00,  4.4360e+00,  3.1785e-01,  ...,  2.2050e-01,\n",
      "             3.6383e-01,  8.7226e-01],\n",
      "           [ 1.0833e+00,  4.0117e-01,  3.6729e+00,  ...,  6.1066e-02,\n",
      "             3.0340e+00,  1.4683e+00],\n",
      "           [ 9.9418e-01,  1.4804e+01,  2.2584e-01,  ...,  5.8651e-01,\n",
      "             6.3005e+00,  4.3736e-01],\n",
      "           ...,\n",
      "           [ 2.1338e-01,  1.4625e+00,  6.0166e-02,  ...,  2.1658e+00,\n",
      "             5.0454e-01,  4.9063e+00],\n",
      "           [ 1.3748e+00,  1.8939e+00,  3.8734e+00,  ...,  1.9729e+00,\n",
      "             5.3915e-01,  4.1740e+00],\n",
      "           [ 7.8642e-01,  9.6442e-02,  4.2122e-01,  ...,  4.5693e-01,\n",
      "             3.1830e-01,  9.0985e-01]],\n",
      "\n",
      "          [[ 4.8240e+00,  4.3829e+00,  2.8541e+00,  ...,  1.5046e+00,\n",
      "             8.3181e-01,  4.2806e-01],\n",
      "           [ 9.2992e-02,  1.5480e+01,  4.2041e-01,  ...,  2.8304e+01,\n",
      "             7.1264e+00,  2.1676e+01],\n",
      "           [ 2.4839e-01,  5.8047e-01,  8.2996e-01,  ...,  7.1360e-01,\n",
      "             7.4682e-01,  1.6377e-01],\n",
      "           ...,\n",
      "           [ 1.4970e+00,  8.6042e-01,  3.3989e-01,  ...,  7.8077e-01,\n",
      "             3.5700e+00,  1.4039e+00],\n",
      "           [ 4.0594e-01,  8.0816e-01,  2.6643e-01,  ...,  6.7495e-01,\n",
      "             4.2617e+00,  2.2861e+00],\n",
      "           [ 3.4411e-02,  3.5659e+00,  5.6488e-01,  ...,  1.2222e+00,\n",
      "             5.1714e-01,  4.5635e-01]],\n",
      "\n",
      "          [[ 8.1731e-01,  6.9522e-01,  2.9086e-01,  ...,  1.9690e+00,\n",
      "             4.0568e+00,  1.6526e-01],\n",
      "           [ 3.2512e-01,  2.5430e+00,  1.7964e+00,  ...,  3.2969e-01,\n",
      "             1.2198e+00,  1.1148e+00],\n",
      "           [ 5.3113e+01,  1.2359e+00,  9.1635e+00,  ...,  2.1809e+00,\n",
      "             4.3907e-01,  1.5562e-01],\n",
      "           ...,\n",
      "           [ 2.5819e-01,  9.8971e+00,  3.6653e-01,  ...,  1.1356e+00,\n",
      "             3.1281e-01,  7.8433e-01],\n",
      "           [ 1.3025e+00,  3.6021e+00,  3.5363e-01,  ...,  6.0872e+00,\n",
      "             2.1638e+00,  1.7360e+00],\n",
      "           [ 6.9851e-01,  5.7613e-01,  4.1829e-01,  ...,  4.0172e+00,\n",
      "             1.6532e-01,  4.2449e+00]]],\n",
      "\n",
      "\n",
      "         [[[-1.3386e-01,  2.1021e+00, -6.6335e+00,  ...,  3.1040e+00,\n",
      "            -1.1768e+00,  2.1971e+00],\n",
      "           [-2.1816e+00,  1.0031e-01,  1.4652e-01,  ...,  4.2235e+00,\n",
      "             2.9079e+00, -2.0941e+00],\n",
      "           [ 3.3355e-01, -8.9818e+00, -6.7758e-02,  ...,  1.0514e+00,\n",
      "             1.5560e-01,  2.9431e-01],\n",
      "           ...,\n",
      "           [-3.9257e-01, -8.5828e-01,  1.0944e+00,  ..., -1.6816e+00,\n",
      "             4.3414e-01,  2.6450e+00],\n",
      "           [ 8.3434e-01,  1.1334e+01, -1.7780e-02,  ..., -3.1703e-01,\n",
      "            -1.4412e+01, -2.0652e+00],\n",
      "           [-2.9206e-01, -1.1990e-01, -2.0663e-01,  ..., -7.4389e-01,\n",
      "            -2.6526e-01,  1.8398e+00]],\n",
      "\n",
      "          [[-2.8613e-01, -6.0279e-01, -8.6574e-02,  ..., -8.0184e+00,\n",
      "            -1.1041e+00, -2.0709e-01],\n",
      "           [-5.1084e-01,  1.5478e+00,  3.6895e-02,  ..., -9.8102e-01,\n",
      "            -3.7158e+00,  1.6124e-02],\n",
      "           [-5.3097e-01,  1.8759e-02, -4.9009e-01,  ..., -3.9091e-02,\n",
      "            -1.1891e+00, -1.2971e+01],\n",
      "           ...,\n",
      "           [ 1.7327e-01, -9.9427e-02, -2.3833e+00,  ..., -2.1477e-01,\n",
      "            -8.5823e-01, -1.2478e+00],\n",
      "           [-3.5884e-01,  6.9346e+00,  4.5583e-01,  ...,  3.2070e+00,\n",
      "            -6.1300e-01, -1.4962e+00],\n",
      "           [-2.1328e+00,  2.4708e-01,  1.2786e-01,  ..., -1.1880e+00,\n",
      "            -3.9927e-02, -3.1922e-01]],\n",
      "\n",
      "          [[-7.7270e-02, -1.2558e+01, -2.3543e-01,  ...,  3.2951e-02,\n",
      "             9.0375e-01, -2.2034e+00],\n",
      "           [ 1.3066e+01, -7.6354e-02, -7.0725e-01,  ..., -3.3620e-01,\n",
      "             4.1743e-02, -1.5523e+00],\n",
      "           [-3.9842e+00, -6.3313e-01, -1.5132e+00,  ...,  2.8874e-01,\n",
      "             1.1220e+00, -8.8789e+00],\n",
      "           ...,\n",
      "           [-1.4748e+01,  6.4863e-01, -5.2041e-01,  ...,  1.1830e+01,\n",
      "             1.5393e+00,  6.2634e+00],\n",
      "           [-3.5872e-01, -1.9324e+00,  2.2527e+00,  ...,  5.2912e-02,\n",
      "             1.4204e-01,  2.1752e-01],\n",
      "           [-1.4727e+00, -8.6774e-01,  7.0941e-01,  ..., -1.9571e+01,\n",
      "            -2.5651e+00, -5.3954e-01]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[-7.1223e-01,  2.0686e+01,  2.4464e-01,  ..., -3.8276e+00,\n",
      "             7.3672e-01,  1.1516e-01],\n",
      "           [-2.5677e+00, -1.4889e+00, -2.5832e+00,  ..., -5.8236e-02,\n",
      "             1.0698e+00, -1.7320e+00],\n",
      "           [ 1.2306e+00, -1.5591e+00, -6.3391e-02,  ...,  5.9483e-01,\n",
      "            -4.7441e+00, -6.3414e-02],\n",
      "           ...,\n",
      "           [ 1.2482e-01, -3.1911e+00,  2.0592e-01,  ...,  1.6788e+00,\n",
      "            -4.3832e-01,  2.4416e+01],\n",
      "           [ 6.0156e-01,  7.7004e-01,  4.4645e+00,  ...,  1.2013e+00,\n",
      "            -2.7986e+00, -3.0011e+00],\n",
      "           [ 1.0907e+00, -4.9171e-01,  4.4185e-01,  ...,  4.9650e-01,\n",
      "            -7.9858e-01, -5.1824e+00]],\n",
      "\n",
      "          [[ 1.8244e+00, -7.0667e+00,  4.1603e+00,  ...,  1.4044e+01,\n",
      "             3.4508e-01,  1.9935e-01],\n",
      "           [-4.8578e-01, -1.7316e+00, -6.5738e-01,  ..., -1.1440e+00,\n",
      "             5.3761e+00,  5.2564e+01],\n",
      "           [-7.3698e-01,  2.8655e-01, -2.5244e-01,  ..., -2.9168e-02,\n",
      "            -1.1557e+00, -9.4837e-02],\n",
      "           ...,\n",
      "           [-1.3771e+00,  4.6737e-02,  5.5482e-01,  ..., -1.2235e+00,\n",
      "            -3.3182e+00, -3.1221e-01],\n",
      "           [ 1.9648e-01, -4.1248e-01, -4.0282e-01,  ..., -1.1823e-01,\n",
      "             3.3350e+00,  1.5488e+00],\n",
      "           [ 8.5044e-02,  3.3600e+00,  6.5193e-01,  ..., -8.2252e-01,\n",
      "             2.6231e-01,  5.7302e-01]],\n",
      "\n",
      "          [[ 1.1324e+00,  3.8699e-01, -7.7977e-01,  ...,  2.8289e+00,\n",
      "            -1.5565e+00, -2.0587e-01],\n",
      "           [ 5.6135e-01, -7.8459e-01, -1.3527e+00,  ...,  2.5389e-01,\n",
      "             1.8650e-01, -2.3588e-01],\n",
      "           [-2.5179e+01, -1.2098e+00, -1.4712e+00,  ..., -2.4559e+00,\n",
      "             5.7427e-01,  6.6339e-02],\n",
      "           ...,\n",
      "           [ 6.9739e-02, -1.7953e+01,  5.8271e-01,  ..., -5.6486e-01,\n",
      "            -5.7335e-01, -2.5417e-01],\n",
      "           [ 1.6184e+00, -2.8124e+00, -2.3661e-01,  ...,  2.7265e+00,\n",
      "             6.8192e-01, -1.2957e+00],\n",
      "           [-8.1044e-01,  2.0750e-01, -1.2670e+00,  ..., -2.7811e+00,\n",
      "             6.3418e-02,  1.4007e+00]]],\n",
      "\n",
      "\n",
      "         [[[-1.3627e-01,  8.3903e-01, -5.9579e+00,  ..., -2.9331e+01,\n",
      "            -1.4369e+01, -9.1847e-01],\n",
      "           [-7.4420e+00, -4.2915e-01,  7.4875e-01,  ...,  2.0689e+02,\n",
      "             7.1138e+00,  7.4863e+00],\n",
      "           [ 3.1091e+00, -6.0690e+01,  3.0557e+00,  ...,  1.7822e+00,\n",
      "             9.9020e+00,  2.1804e+00],\n",
      "           ...,\n",
      "           [ 4.8884e-01, -9.2811e-01, -5.4130e+01,  ..., -1.2137e+01,\n",
      "             8.9329e-01,  8.2791e-01],\n",
      "           [-2.8343e+00,  5.6632e+00,  3.8624e+00,  ...,  1.6438e+00,\n",
      "            -1.5304e+03, -2.8707e+01],\n",
      "           [ 1.1878e+01, -7.4997e-01, -7.0171e-01,  ...,  3.6420e+00,\n",
      "            -1.4620e+00, -4.4087e+00]],\n",
      "\n",
      "          [[-1.4768e+01,  1.1499e+00, -3.0711e-01,  ...,  1.0912e+03,\n",
      "            -2.4664e-01,  1.6573e+00],\n",
      "           [-1.1825e+00, -8.0707e-01, -1.4928e+00,  ..., -3.1958e+00,\n",
      "            -4.7250e-01,  1.1392e+00],\n",
      "           [ 1.1152e+00,  4.1784e-01,  1.2152e+00,  ..., -7.5515e-01,\n",
      "            -1.0110e+01, -1.2910e+02],\n",
      "           ...,\n",
      "           [-2.5685e-02,  9.1247e-01, -3.1458e+00,  ..., -2.2440e-01,\n",
      "             5.4293e+00, -3.3518e+00],\n",
      "           [-2.5518e-01, -4.2597e+00,  1.1847e+00,  ..., -1.3561e+01,\n",
      "            -1.3348e+00, -1.0176e+01],\n",
      "           [-6.2466e+00, -4.6867e+00, -8.1644e+00,  ..., -2.8341e-01,\n",
      "             2.5770e-01, -1.1468e+00]],\n",
      "\n",
      "          [[-3.1530e-01,  1.1590e+00, -1.0482e+00,  ...,  7.3544e-01,\n",
      "            -3.5560e+00, -4.1439e+00],\n",
      "           [ 4.5551e+00,  1.9714e-01,  1.6407e+00,  ..., -9.4197e-01,\n",
      "            -3.3836e-01,  2.9692e+01],\n",
      "           [-4.2126e-01,  6.8536e-01, -1.5168e+00,  ..., -1.1914e+00,\n",
      "             4.4880e+00,  5.3863e+00],\n",
      "           ...,\n",
      "           [ 3.5457e+00,  1.8030e-01,  3.8512e+01,  ..., -3.1100e+02,\n",
      "            -6.2491e+00,  4.7857e+00],\n",
      "           [-4.8358e+00, -5.2445e+00,  5.7894e-01,  ...,  3.2422e-02,\n",
      "            -5.1265e+00,  8.7068e+00],\n",
      "           [ 1.9081e+00,  2.1043e+00,  3.2230e-01,  ..., -7.4495e+01,\n",
      "            -4.8994e+00, -5.2297e-01]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[ 7.9986e+00, -1.8420e+01,  4.0260e-01,  ...,  3.7061e+00,\n",
      "            -3.3451e-01, -1.1079e+00],\n",
      "           [ 4.7886e+00, -2.9699e+00, -2.9712e+00,  ...,  3.7731e-02,\n",
      "             7.8726e+00, -6.8369e-01],\n",
      "           [-1.5111e+02,  4.1286e+00,  1.6248e+00,  ...,  2.6938e-01,\n",
      "             2.0062e+01,  3.2530e-01],\n",
      "           ...,\n",
      "           [ 1.0490e+00, -1.3966e+01, -4.5754e-02,  ..., -6.4903e+00,\n",
      "             1.2567e-01, -3.2279e+01],\n",
      "           [ 2.2636e+00, -1.5343e+00, -3.0278e+00,  ...,  1.4936e+01,\n",
      "            -3.5562e+01, -4.8893e-01],\n",
      "           [-9.8340e-01,  1.7678e+00, -6.4978e-01,  ..., -4.4638e-01,\n",
      "             2.5764e+00,  1.9458e+00]],\n",
      "\n",
      "          [[ 1.9290e+00, -2.1840e+00,  3.7668e+00,  ...,  6.0160e+00,\n",
      "            -1.8304e+00,  3.2396e-01],\n",
      "           [-2.0521e-01, -9.8394e-02,  7.7056e-02,  ...,  1.3331e+01,\n",
      "             1.4004e+01, -5.2487e+00],\n",
      "           [ 1.0530e-01,  6.6949e+00,  2.0346e+00,  ..., -8.3975e-01,\n",
      "            -2.4720e+00, -1.2127e-01],\n",
      "           ...,\n",
      "           [ 1.1123e+00,  5.9546e-01,  4.4977e-02,  ...,  4.1402e+00,\n",
      "             7.8337e-01, -1.9633e+00],\n",
      "           [-2.7858e-04, -3.2368e-01, -1.2717e-01,  ...,  1.7088e+00,\n",
      "             5.5709e+00,  8.3777e+01],\n",
      "           [-1.5352e-01,  5.2943e+00, -1.0004e+00,  ..., -4.2177e+00,\n",
      "            -3.3574e-02,  7.3321e+00]],\n",
      "\n",
      "          [[-1.5959e-01, -5.6042e+00, -2.4883e+00,  ...,  1.0726e+01,\n",
      "            -4.3627e+00, -3.6603e+00],\n",
      "           [-5.6674e+00, -4.3817e+00,  6.1289e+00,  ...,  4.9954e-01,\n",
      "            -1.5109e+00, -2.3830e+00],\n",
      "           [ 1.0492e+02, -4.7806e+00, -2.3080e+01,  ...,  4.9953e+00,\n",
      "             1.4317e+00, -2.2564e-01],\n",
      "           ...,\n",
      "           [-1.8180e+00, -2.1669e+01, -3.0700e-02,  ..., -1.1156e+00,\n",
      "             1.4181e+00,  2.5137e-01],\n",
      "           [-1.0734e+00,  8.0135e-02,  7.0852e-01,  ...,  2.3438e+00,\n",
      "            -4.2102e+00, -1.2215e+01],\n",
      "           [ 3.0094e+00,  7.3741e-01,  2.2728e+00,  ..., -2.0928e+01,\n",
      "            -4.1444e-01,  1.4133e+00]]],\n",
      "\n",
      "\n",
      "         ...,\n",
      "\n",
      "\n",
      "         [[[-1.3627e-01,  8.3903e-01, -5.9579e+00,  ..., -2.9331e+01,\n",
      "            -1.4369e+01, -9.1847e-01],\n",
      "           [-7.4420e+00, -4.2915e-01,  7.4875e-01,  ...,  2.0689e+02,\n",
      "             7.1138e+00,  7.4863e+00],\n",
      "           [ 3.1091e+00, -6.0690e+01,  3.0557e+00,  ...,  1.7822e+00,\n",
      "             9.9020e+00,  2.1804e+00],\n",
      "           ...,\n",
      "           [ 4.8884e-01, -9.2811e-01, -5.4130e+01,  ..., -1.2137e+01,\n",
      "             8.9329e-01,  8.2791e-01],\n",
      "           [-2.8343e+00,  5.6632e+00,  3.8624e+00,  ...,  1.6438e+00,\n",
      "            -1.5304e+03, -2.8707e+01],\n",
      "           [ 1.1878e+01, -7.4997e-01, -7.0171e-01,  ...,  3.6420e+00,\n",
      "            -1.4620e+00, -4.4087e+00]],\n",
      "\n",
      "          [[-1.4768e+01,  1.1499e+00, -3.0711e-01,  ...,  1.0912e+03,\n",
      "            -2.4664e-01,  1.6573e+00],\n",
      "           [-1.1825e+00, -8.0707e-01, -1.4928e+00,  ..., -3.1958e+00,\n",
      "            -4.7250e-01,  1.1392e+00],\n",
      "           [ 1.1152e+00,  4.1784e-01,  1.2152e+00,  ..., -7.5515e-01,\n",
      "            -1.0110e+01, -1.2910e+02],\n",
      "           ...,\n",
      "           [-2.5685e-02,  9.1247e-01, -3.1458e+00,  ..., -2.2440e-01,\n",
      "             5.4293e+00, -3.3518e+00],\n",
      "           [-2.5518e-01, -4.2597e+00,  1.1847e+00,  ..., -1.3561e+01,\n",
      "            -1.3348e+00, -1.0176e+01],\n",
      "           [-6.2466e+00, -4.6867e+00, -8.1644e+00,  ..., -2.8341e-01,\n",
      "             2.5770e-01, -1.1468e+00]],\n",
      "\n",
      "          [[-3.1530e-01,  1.1590e+00, -1.0482e+00,  ...,  7.3544e-01,\n",
      "            -3.5560e+00, -4.1439e+00],\n",
      "           [ 4.5551e+00,  1.9714e-01,  1.6407e+00,  ..., -9.4197e-01,\n",
      "            -3.3836e-01,  2.9692e+01],\n",
      "           [-4.2126e-01,  6.8536e-01, -1.5168e+00,  ..., -1.1914e+00,\n",
      "             4.4880e+00,  5.3863e+00],\n",
      "           ...,\n",
      "           [ 3.5457e+00,  1.8030e-01,  3.8512e+01,  ..., -3.1100e+02,\n",
      "            -6.2491e+00,  4.7857e+00],\n",
      "           [-4.8358e+00, -5.2445e+00,  5.7894e-01,  ...,  3.2422e-02,\n",
      "            -5.1265e+00,  8.7068e+00],\n",
      "           [ 1.9081e+00,  2.1043e+00,  3.2230e-01,  ..., -7.4495e+01,\n",
      "            -4.8994e+00, -5.2297e-01]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[ 7.9986e+00, -1.8420e+01,  4.0260e-01,  ...,  3.7061e+00,\n",
      "            -3.3451e-01, -1.1079e+00],\n",
      "           [ 4.7886e+00, -2.9699e+00, -2.9712e+00,  ...,  3.7731e-02,\n",
      "             7.8726e+00, -6.8369e-01],\n",
      "           [-1.5111e+02,  4.1286e+00,  1.6248e+00,  ...,  2.6938e-01,\n",
      "             2.0062e+01,  3.2530e-01],\n",
      "           ...,\n",
      "           [ 1.0490e+00, -1.3966e+01, -4.5754e-02,  ..., -6.4903e+00,\n",
      "             1.2567e-01, -3.2279e+01],\n",
      "           [ 2.2636e+00, -1.5343e+00, -3.0278e+00,  ...,  1.4936e+01,\n",
      "            -3.5562e+01, -4.8893e-01],\n",
      "           [-9.8340e-01,  1.7678e+00, -6.4978e-01,  ..., -4.4638e-01,\n",
      "             2.5764e+00,  1.9458e+00]],\n",
      "\n",
      "          [[ 1.9290e+00, -2.1840e+00,  3.7668e+00,  ...,  6.0160e+00,\n",
      "            -1.8304e+00,  3.2396e-01],\n",
      "           [-2.0521e-01, -9.8394e-02,  7.7056e-02,  ...,  1.3331e+01,\n",
      "             1.4004e+01, -5.2487e+00],\n",
      "           [ 1.0530e-01,  6.6949e+00,  2.0346e+00,  ..., -8.3975e-01,\n",
      "            -2.4720e+00, -1.2127e-01],\n",
      "           ...,\n",
      "           [ 1.1123e+00,  5.9546e-01,  4.4977e-02,  ...,  4.1402e+00,\n",
      "             7.8337e-01, -1.9633e+00],\n",
      "           [-2.7858e-04, -3.2368e-01, -1.2717e-01,  ...,  1.7088e+00,\n",
      "             5.5709e+00,  8.3777e+01],\n",
      "           [-1.5352e-01,  5.2943e+00, -1.0004e+00,  ..., -4.2177e+00,\n",
      "            -3.3574e-02,  7.3321e+00]],\n",
      "\n",
      "          [[-1.5959e-01, -5.6042e+00, -2.4883e+00,  ...,  1.0726e+01,\n",
      "            -4.3627e+00, -3.6603e+00],\n",
      "           [-5.6674e+00, -4.3817e+00,  6.1289e+00,  ...,  4.9954e-01,\n",
      "            -1.5109e+00, -2.3830e+00],\n",
      "           [ 1.0492e+02, -4.7806e+00, -2.3080e+01,  ...,  4.9953e+00,\n",
      "             1.4317e+00, -2.2564e-01],\n",
      "           ...,\n",
      "           [-1.8180e+00, -2.1669e+01, -3.0700e-02,  ..., -1.1156e+00,\n",
      "             1.4181e+00,  2.5137e-01],\n",
      "           [-1.0734e+00,  8.0135e-02,  7.0852e-01,  ...,  2.3438e+00,\n",
      "            -4.2102e+00, -1.2215e+01],\n",
      "           [ 3.0094e+00,  7.3741e-01,  2.2728e+00,  ..., -2.0928e+01,\n",
      "            -4.1444e-01,  1.4133e+00]]],\n",
      "\n",
      "\n",
      "         [[[ 2.9931e+00,  1.1374e+01,  1.4585e+01,  ..., -4.9674e+01,\n",
      "             5.1972e+01, -7.4041e-01],\n",
      "           [ 1.7461e+01, -2.2880e+00,  4.8047e+00,  ...,  1.6327e+03,\n",
      "             9.2905e+00, -4.9747e+00],\n",
      "           [ 4.5932e+00,  1.4812e+02,  6.5679e-01,  ...,  6.8291e+00,\n",
      "            -1.0663e+01, -9.3625e-01],\n",
      "           ...,\n",
      "           [-8.0414e+00,  3.2666e+00, -1.2927e+01,  ...,  1.9816e+00,\n",
      "             1.0199e+00,  3.2622e+00],\n",
      "           [-1.8672e+00,  4.5641e+00,  3.6917e+00,  ...,  2.8257e+00,\n",
      "             6.2615e+03, -2.6713e+01],\n",
      "           [-9.1428e+00,  1.6712e+00, -8.8826e+00,  ..., -8.7866e-01,\n",
      "            -2.8324e-01, -3.0714e+01]],\n",
      "\n",
      "          [[ 1.8836e+02, -5.5925e+00, -1.0277e-01,  ..., -4.3669e+03,\n",
      "             3.7313e+00, -1.5448e+00],\n",
      "           [ 3.4052e+00, -4.8210e-01,  2.7260e-01,  ...,  9.1248e+00,\n",
      "             1.2767e+00,  7.9493e-01],\n",
      "           [-1.7690e+00, -4.3525e-01, -1.0646e-01,  ...,  1.6896e+00,\n",
      "             1.4781e+01,  2.1572e+02],\n",
      "           ...,\n",
      "           [-5.1199e-02, -5.6584e-01,  3.1754e+00,  ...,  4.6160e-02,\n",
      "            -3.0346e+00,  3.1878e+00],\n",
      "           [-3.0410e-01, -3.7256e+00,  9.3609e-01,  ..., -8.9908e+01,\n",
      "             5.3569e+00,  1.2920e+02],\n",
      "           [ 3.4036e+00,  8.2764e-01,  3.4524e+00,  ..., -1.0707e+00,\n",
      "             9.8183e+00,  2.2471e+00]],\n",
      "\n",
      "          [[-1.5573e+00, -2.8219e+00,  2.0156e+00,  ...,  9.1040e-01,\n",
      "             1.1749e+00,  2.1309e+01],\n",
      "           [ 8.6322e+00,  1.6742e+00, -1.9549e+00,  ...,  1.6741e+00,\n",
      "             1.8708e+00, -4.1695e+01],\n",
      "           [ 6.0482e-01, -1.2056e+00,  4.8713e+00,  ...,  4.5827e+00,\n",
      "             4.2700e+00, -1.3353e+02],\n",
      "           ...,\n",
      "           [-1.8070e+01, -1.0799e-01,  4.0400e+01,  ..., -1.8661e+02,\n",
      "            -4.3692e+00,  2.1362e+00],\n",
      "           [-1.3277e+01,  2.3168e+00,  2.9771e+00,  ...,  5.9616e-01,\n",
      "            -1.4201e+00,  2.3726e+01],\n",
      "           [-1.0394e+00, -9.2863e-01, -3.4490e-02,  ...,  1.1449e+02,\n",
      "             4.5763e+00,  2.1181e+00]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[-2.8584e+00, -8.9599e+01,  3.7457e+00,  ..., -1.8016e+02,\n",
      "            -1.2524e+00, -1.3969e+00],\n",
      "           [-1.4843e+01,  1.8774e+01,  1.9785e+00,  ...,  1.4639e+00,\n",
      "             8.3420e+00,  5.0419e-01],\n",
      "           [-2.5013e+02, -2.7999e+00, -3.7604e-01,  ..., -6.4389e-01,\n",
      "            -1.4725e+01, -1.2585e+00],\n",
      "           ...,\n",
      "           [ 8.8769e-01,  3.0248e+01, -1.5936e-01,  ..., -5.3499e+00,\n",
      "            -3.0865e+00, -1.6868e+02],\n",
      "           [ 8.9316e-01,  1.4085e-01, -4.4451e+00,  ...,  3.3810e+01,\n",
      "             2.4208e+02,  5.0346e-01],\n",
      "           [-1.9822e+01, -1.0072e+01,  8.8749e-01,  ...,  8.9970e-01,\n",
      "            -6.6698e+00, -1.4960e+01]],\n",
      "\n",
      "          [[ 1.4258e+00, -1.2948e+00,  6.5798e+00,  ...,  5.8747e+01,\n",
      "            -4.7136e+00,  1.8099e-01],\n",
      "           [ 1.4072e+01, -3.2041e+00, -8.1823e-01,  ..., -1.2409e+00,\n",
      "             1.3555e+01, -1.2855e+01],\n",
      "           [ 1.7818e+00,  7.5326e+00, -2.0538e+00,  ...,  2.4392e+00,\n",
      "             1.7112e+01,  1.4991e-02],\n",
      "           ...,\n",
      "           [ 2.4177e-01, -1.4643e+01,  3.9722e+00,  ...,  1.0623e+00,\n",
      "             2.7937e-01,  1.7656e+01],\n",
      "           [ 7.2915e-01, -2.8233e+00, -1.8469e-01,  ...,  4.9263e+00,\n",
      "             6.5845e+00,  2.6325e+01],\n",
      "           [-3.5173e+00,  5.8572e+00, -9.2554e+00,  ...,  2.3191e+00,\n",
      "            -6.9075e+00,  6.9780e+00]],\n",
      "\n",
      "          [[-1.0129e+00, -4.9928e+00,  3.7077e+00,  ...,  1.8990e+01,\n",
      "             4.0222e+00, -3.0334e+01],\n",
      "           [-3.1848e+00,  3.9916e-01, -4.8936e+00,  ...,  2.0411e+00,\n",
      "            -3.2464e-01,  4.0080e-01],\n",
      "           [-4.8792e+01,  6.4411e+00,  8.6011e+00,  ..., -7.9544e+00,\n",
      "             2.5187e+02,  2.2459e+00],\n",
      "           ...,\n",
      "           [-5.0631e-01,  6.4202e+01, -2.3268e+00,  ...,  5.8535e-01,\n",
      "            -1.2262e+01,  3.3598e-02],\n",
      "           [-3.5679e+00, -8.2167e-01,  2.2679e-01,  ...,  1.0809e+00,\n",
      "             1.1955e+00,  9.1167e+00],\n",
      "           [-2.3627e+00, -3.8421e-01, -2.9878e+00,  ...,  1.2432e+01,\n",
      "             6.7321e-01, -3.0136e+00]]],\n",
      "\n",
      "\n",
      "         [[[ 5.2735e+00,  4.5604e+00,  2.0712e+01,  ...,  6.2546e+02,\n",
      "             5.4807e+02,  1.6230e+00],\n",
      "           [ 4.9884e+01,  7.0296e+00,  4.4687e+00,  ...,  6.7914e+04,\n",
      "             4.0308e+01,  1.7185e+01],\n",
      "           [ 1.0674e+02,  1.0000e+03,  3.8688e+00,  ...,  1.4987e+01,\n",
      "             1.7611e+03,  1.0103e+01],\n",
      "           ...,\n",
      "           [ 6.1364e+00,  6.1346e+00,  6.9611e+02,  ...,  9.2658e+00,\n",
      "             2.0856e+00,  1.0082e+01],\n",
      "           [ 4.9713e+00,  3.5854e+00,  2.4040e+02,  ...,  9.2797e+01,\n",
      "             6.7536e+05,  9.5427e+01],\n",
      "           [ 2.7635e+02,  5.4197e+00,  1.3338e+02,  ...,  9.2885e+00,\n",
      "             2.8908e+00,  7.0555e+01]],\n",
      "\n",
      "          [[ 4.4733e+03,  1.1126e+01,  4.7202e-01,  ...,  5.0305e+05,\n",
      "             8.0623e+00,  4.0409e+01],\n",
      "           [ 4.3202e+00,  9.9980e-01,  2.8434e+00,  ...,  3.2055e+01,\n",
      "             1.8412e+00,  2.0077e+00],\n",
      "           [ 3.4074e+00,  2.0374e+00,  1.3268e+00,  ...,  4.2073e+00,\n",
      "             1.1573e+02,  2.5425e+03],\n",
      "           ...,\n",
      "           [ 4.8591e-01,  3.1113e+00,  4.5394e+00,  ...,  7.6376e-01,\n",
      "             1.4230e+01,  9.4966e+00],\n",
      "           [ 6.9133e-01,  2.9125e+00,  2.6392e+00,  ...,  4.0036e+02,\n",
      "             8.1648e+01,  2.8553e+03],\n",
      "           [ 1.8301e+01,  6.3269e+01,  3.2943e+01,  ...,  4.4324e+00,\n",
      "             3.3864e+01,  2.0655e+01]],\n",
      "\n",
      "          [[ 9.6714e-01,  5.1120e-01,  3.0627e+00,  ...,  4.3347e+00,\n",
      "             1.4546e+01,  5.6341e+01],\n",
      "           [ 4.3254e+00,  1.4409e+00,  5.0149e+00,  ...,  4.1347e+00,\n",
      "             6.9994e+00,  8.2737e+02],\n",
      "           [ 3.5989e-01,  9.9909e-01,  5.3186e+00,  ...,  1.5756e+01,\n",
      "             1.9705e+01,  7.8495e+01],\n",
      "           ...,\n",
      "           [ 7.6644e+00,  2.1351e-01,  2.0153e+03,  ...,  4.7117e+03,\n",
      "             1.8890e+01,  4.9274e+00],\n",
      "           [ 1.5432e+02,  1.2466e+01,  1.4084e+00,  ...,  1.2321e+00,\n",
      "             1.6914e+01,  1.9515e+03],\n",
      "           [ 3.6474e+00,  2.5446e+00,  3.1245e-01,  ...,  4.8545e+02,\n",
      "             1.8838e+01,  2.9723e+00]],\n",
      "\n",
      "          ...,\n",
      "\n",
      "          [[ 5.4163e+01,  8.3061e+01,  4.5560e+00,  ...,  1.6113e+02,\n",
      "             4.6643e+00,  2.1104e+00],\n",
      "           [ 3.2761e+01,  4.6621e+01,  2.5638e+00,  ...,  7.4214e+00,\n",
      "             3.3766e+01,  4.8345e-01],\n",
      "           [ 2.8226e+04,  5.3943e+00,  1.7518e+01,  ...,  6.0155e+00,\n",
      "             1.1187e+02,  2.5832e+00],\n",
      "           ...,\n",
      "           [ 6.2625e+00,  1.3744e+02,  7.0696e-01,  ...,  3.2739e+01,\n",
      "             5.5581e+00,  2.2159e+02],\n",
      "           [ 5.0427e+00,  2.0338e+00,  5.7271e+00,  ...,  7.0222e+02,\n",
      "             4.1036e+03,  2.7018e-01],\n",
      "           [ 1.2317e+01,  3.9740e+01,  2.2803e+01,  ...,  2.1734e+00,\n",
      "             2.7855e+01,  6.2561e+00]],\n",
      "\n",
      "          [[ 2.8150e+00,  1.6234e+01,  1.5823e+01,  ...,  2.9130e+01,\n",
      "             1.6379e+01,  4.6197e-01],\n",
      "           [ 8.9518e+00,  1.6357e+00,  9.3175e-01,  ...,  1.9608e+01,\n",
      "             7.2059e+01,  1.3308e+00],\n",
      "           [ 6.3898e-01,  1.4615e+02,  2.3661e+01,  ...,  3.0530e+00,\n",
      "             5.8304e+01,  3.7269e-01],\n",
      "           ...,\n",
      "           [ 3.3247e+00,  6.8356e+01,  6.9021e+00,  ...,  3.4952e+01,\n",
      "             4.2604e+00,  3.5710e+01],\n",
      "           [ 8.0260e+00,  1.6473e+01,  1.5192e+00,  ...,  2.8582e+01,\n",
      "             2.5629e+01,  8.8280e+03],\n",
      "           [ 9.9601e+00,  8.7496e+00,  1.1147e+01,  ...,  1.6994e+01,\n",
      "             3.1180e+00,  1.3451e+02]],\n",
      "\n",
      "          [[ 6.4504e-01,  9.0517e+01,  5.5488e+01,  ...,  8.2565e+01,\n",
      "             1.2838e+01,  1.8982e+03],\n",
      "           [ 2.7113e+02,  1.7296e+01,  7.0959e+01,  ...,  3.0001e+01,\n",
      "             2.1924e+00,  7.7132e+00],\n",
      "           [ 2.1892e+02,  6.3420e+01,  2.6736e+02,  ...,  1.3631e+01,\n",
      "             4.5317e+03,  1.4004e+00],\n",
      "           ...,\n",
      "           [ 1.2920e+01,  1.8235e+02,  6.6371e+00,  ...,  1.4175e+00,\n",
      "             3.0071e+01,  2.1491e-01],\n",
      "           [ 2.6300e+01,  5.1001e+00,  8.0837e+00,  ...,  9.1433e-01,\n",
      "             1.7751e+01,  9.9583e+01],\n",
      "           [ 5.6400e+01,  3.8415e+00,  1.8986e+01,  ...,  1.1375e+02,\n",
      "             1.6310e+00,  3.2462e+00]]]]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "#W = torch.einsum(\"...ik,...ks,...js->...ij\",[G_inv, B, G_inv])\n",
    "print(W.shape)\n",
    "print(W.shape[0]*W.shape[1]*W.shape[2])\n",
    "print(W.permute((3,4,0,1,2)).shape)\n",
    "print(W.permute((3,4,0,1,2)).reshape((1,9,100,100,100)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 100, 100, 3]) torch.Size([100, 100, 100, 3, 3])\n",
      "363751 363751 363751\n",
      "1091253\n",
      "CPU times: user 1.88 s, sys: 5.94 ms, total: 1.88 s\n",
      "Wall time: 1.71 s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "if use_cuda:\n",
    "  W.cuda()\n",
    "\n",
    "#Lamb3 = TV.vSymEig(W.permute((3,4,0,1,2)).reshape((1,9,*W.shape[0:3])), eigenvectors=True)[0].double()\n",
    "#Wreshape = W.permute((3,4,0,1,2)).reshape((1,9,W.shape[0],W.shape[1],W.shape[2]))\n",
    "#Lamb3, vec = TV.vSymEig(Wreshape, eigenvectors=True)#[0].double()\n",
    "vals = TV.vSymEig(W.permute((3,4,0,1,2)).reshape((1,9,*W.shape[0:3])), eigenvectors=True)\n",
    "Lamb3=torch.abs(vals[0]).double()\n",
    "#Lamb3 = vals[0].double()\n",
    "#Lamb3[Lamb3<0]=1e-5\n",
    "vec=vals[1].double()\n",
    "#IdxMag = checkMag(Lamb3.reshape((3,*W.shape[0:3])).permute((1,2,3,0)).double(),vec.reshape((3,3,*W.shape[0:3])).permute((2,3,4,0,1)).double())\n",
    "#IdxMag = checkMag2(Lamb3.reshape((3,*W.shape[0:3])).permute((1,2,3,0)))\n",
    "#IdxMag = checkMag3(Lamb3.reshape((3,*W.shape[0:3])).permute((1,2,3,0)),W)\n",
    "IdxMag = checkMag4(Lamb3.reshape((3,*W.shape[0:3])).permute((1,2,3,0)).double(),vec.reshape((3,3,*W.shape[0:3])).permute((2,3,4,0,1)).double())\n",
    "##smallW = \n",
    "if use_cuda:\n",
    "  Lamb3[0,:,IdxMag[0],IdxMag[1],IdxMag[2]] = torch.symeig(W[IdxMag[0],IdxMag[1],IdxMag[2],:,:].cpu(), eigenvectors=False)[0].permute((1,0)).cuda()\n",
    "else:\n",
    "  Lamb3[0,:,IdxMag[0],IdxMag[1],IdxMag[2]] = torch.symeig(W[IdxMag[0],IdxMag[1],IdxMag[2],:,:], eigenvectors=False)[0].permute((1,0))\n",
    "    \n",
    "print(len(IdxMag[0]))\n",
    "##Lamb3[Lamb3<1e-7] = 1e-7\n",
    "##logLamb3 = torch.log(Lamb3).reshape((3,*W.shape[0:3])).permute((1,2,3,0)).double()\n",
    "##logLamb3[torch.where(torch.isnan(logLamb3))] = torch.log(Lamb1)[torch.where(torch.isnan(logLamb3))]\n",
    "trKsquare3 = torch.sum(torch.log(Lamb3)**2,(0,1))\n",
    "##trKsquare3 = torch.sum(logLamb3**2,(-1))\n",
    "# CPU performance:\n",
    "# Just abs: 0 recalc, time = 166ms, trkerror: 0.0204\n",
    "# checkMag: 85941 recalc, time = 1.08s, trkerror: 0.0155\n",
    "# checkMag2: 988635 recalc, time = 1.62s, trkerror: 0.0047\n",
    "# checkMag3: 240 recalc, time = 181ms, trkerror: 0.0239\n",
    "# checkMag3_v2: 210008 recalc, time = 528ms, trkerror: 0.0002\n",
    "# GPU performance:\n",
    "# checkMag: 85617 recalc, time=238ms, trkerror: 0.0069\n",
    "# checkMag2: 988724 recalc, time=1.58s, trkerror: 0.0083\n",
    "# checkMag3: 201 recalc, time=112 ms, trkerror: 0.0268\n",
    "# checkMag3_v2: 210142 recalc, time=458 ms, trkerror: 0.0061"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.2624], dtype=torch.float64)"
      ]
     },
     "execution_count": 152,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.log(torch.from_numpy(np.array([1.3])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28814\n"
     ]
    }
   ],
   "source": [
    "#print(W.shape,Lamb3.shape,len(IdxMag),torch.symeig(W[IdxMag[0],IdxMag[1],IdxMag[2],:,:])[0].shape,smallW.shape)\n",
    "#print(Lamb3[0,:,IdxMag[0],IdxMag[1],IdxMag[2]].shape)\n",
    "#print(torch.max(W[IdxMag[0],IdxMag[1],IdxMag[2],:,:]))\n",
    "print(len(IdxMag[0][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 960 ms, sys: 13.3 ms, total: 973 ms\n",
      "Wall time: 82.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "Lamb3Scaled = TV.vSymEig(scaled_W.permute((3,4,0,1,2)).reshape((1,9,*W.shape[0:3])), eigenvectors=False)[0].reshape((3,*W.shape[0:3])).permute((1,2,3,0))\n",
    "Lamb3Alt = torch.einsum('...,...i->...i', W_max, Lamb3Scaled)\n",
    "trKsquare3Alt = torch.sum(torch.log(Lamb3Alt)**2,(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 100, 100])\n",
      "torch.Size([100, 100, 100, 3, 3])\n",
      "tensor(-4757.7923, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "print(W_max.shape)\n",
    "print(scaled_W.shape)\n",
    "print(torch.min(Lamb3Alt))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Method 4: NISymmetricEigensolver3x3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<timed exec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-25-88696fa4db7f>\u001b[0m in \u001b[0;36mNISymmetricEigensolver3x3\u001b[0;34m(A)\u001b[0m\n\u001b[1;32m      8\u001b[0m   \u001b[0;31m# value of the components.  This guards against floating-point\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m   \u001b[0;31m# overflow when computing the eigenvalues.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m   \u001b[0mA_max\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m9\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m   \u001b[0mscaled_A\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meinsum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'...ij,...->...ij'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1.0\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mA_max\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m   \u001b[0;31m#          Real max0 = std::max(std::fabs(a00), std::fabs(a01));\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1."
     ]
    }
   ],
   "source": [
    "%%time\n",
    "Lamb4 = NISymmetricEigensolver3x3(W)\n",
    "trKsquare4 = torch.sum(torch.log(Lamb4)**2,(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the diff "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def checkNaN(A):\n",
    "    if (A != A).any():\n",
    "        print('NaN')\n",
    "        return(np.argwhere(np.isnan(A.numpy())))\n",
    "    else:\n",
    "        print('Good')\n",
    "\n",
    "def check_L2_diff_error(a,b):\n",
    "    L2Diff = LA.norm(a - b)\n",
    "    L2Error =  2*LA.norm(a - b)/( LA.norm(a)+LA.norm(b))\n",
    "    return L2Diff, L2Error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(13.3727, dtype=torch.float64), tensor(0.0003, dtype=torch.float64))"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check diff11alt\n",
    "check_L2_diff_error(trKsquare1,trKsquare1Alt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(6.4856e-13, dtype=torch.float64),\n",
       " tensor(1.6858e-16, dtype=torch.float64))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check diff12\n",
    "check_L2_diff_error(trKsquare1,trKsquare2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(113.0688, device='cuda:0', dtype=torch.float64),\n",
       " tensor(0.0022, device='cuda:0', dtype=torch.float64))"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check diff13\n",
    "check_L2_diff_error(trKsquare1,trKsquare3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'trKsquare2' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-61-5c5512b56041>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# check diff23\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mcheck_L2_diff_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrKsquare2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrKsquare3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'trKsquare2' is not defined"
     ]
    }
   ],
   "source": [
    "# check diff23\n",
    "check_L2_diff_error(trKsquare2,trKsquare3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'trKsquare2' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-62-7f8d564af293>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# check diff23\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mcheck_L2_diff_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrKsquare2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrKsquare3Alt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'trKsquare2' is not defined"
     ]
    }
   ],
   "source": [
    "# check diff23\n",
    "check_L2_diff_error(trKsquare2,trKsquare3Alt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(nan, dtype=torch.float64), tensor(nan, dtype=torch.float64))"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check diff14\n",
    "check_L2_diff_error(trKsquare1,trKsquare4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 100, 100, 100])\n",
      "torch.Size([3, 100, 100])\n"
     ]
    }
   ],
   "source": [
    "print(Lamb4.shape)\n",
    "print(trKsquare4.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check if there are NaN values in trKsquare3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Good\n",
      "Good\n",
      "NaN\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(None,\n",
       " None,\n",
       " array([[ 0, 58, 40],\n",
       "        [ 1, 16, 52],\n",
       "        [ 1, 26, 98],\n",
       "        [ 1, 72, 12],\n",
       "        [ 1, 86, 20],\n",
       "        [ 2, 23, 26],\n",
       "        [ 2, 32,  7],\n",
       "        [ 2, 45, 14],\n",
       "        [ 2, 96, 96],\n",
       "        [ 3, 11, 25],\n",
       "        [ 3, 22, 34],\n",
       "        [ 3, 31, 87],\n",
       "        [ 3, 99, 70],\n",
       "        [ 4, 45, 63],\n",
       "        [ 5, 58, 52],\n",
       "        [ 6, 47, 97],\n",
       "        [ 6, 77, 47],\n",
       "        [ 6, 91, 12],\n",
       "        [ 7, 77, 57],\n",
       "        [ 7, 77, 83],\n",
       "        [ 7, 82,  2],\n",
       "        [ 8, 15, 68],\n",
       "        [ 8, 23, 73],\n",
       "        [ 8, 55, 40],\n",
       "        [ 8, 99, 65],\n",
       "        [ 9, 58,  5],\n",
       "        [ 9, 94, 49],\n",
       "        [ 9, 99, 16],\n",
       "        [10, 39, 69],\n",
       "        [10, 96,  3],\n",
       "        [11, 21, 28],\n",
       "        [11, 28, 47],\n",
       "        [12, 58, 74],\n",
       "        [13, 51, 37],\n",
       "        [13, 99, 92],\n",
       "        [14, 27, 34],\n",
       "        [14, 28, 10],\n",
       "        [14, 52,  7],\n",
       "        [14, 72, 78],\n",
       "        [15, 35, 39],\n",
       "        [15, 80, 91],\n",
       "        [16, 15, 37],\n",
       "        [17,  7, 66],\n",
       "        [17, 69, 14],\n",
       "        [17, 98, 19],\n",
       "        [18, 16, 19],\n",
       "        [18, 28, 15],\n",
       "        [19,  1, 22],\n",
       "        [20, 32, 87],\n",
       "        [20, 98, 86],\n",
       "        [21, 36, 80],\n",
       "        [21, 99, 14],\n",
       "        [22, 29, 45],\n",
       "        [22, 75, 47],\n",
       "        [23, 64, 89],\n",
       "        [28, 46, 77],\n",
       "        [28, 64, 91],\n",
       "        [28, 96, 91],\n",
       "        [29, 22, 39],\n",
       "        [29, 51, 98],\n",
       "        [29, 82, 86],\n",
       "        [30, 41, 51],\n",
       "        [31, 59, 95],\n",
       "        [31, 97, 91],\n",
       "        [32, 13, 40],\n",
       "        [32, 29, 85],\n",
       "        [32, 92, 66],\n",
       "        [34, 24, 42],\n",
       "        [34, 59, 12],\n",
       "        [35, 91, 94],\n",
       "        [36, 41, 19],\n",
       "        [37, 31, 40],\n",
       "        [37, 75, 89],\n",
       "        [38, 85, 99],\n",
       "        [39, 57, 38],\n",
       "        [40, 10, 56],\n",
       "        [40, 69,  9],\n",
       "        [40, 87, 68],\n",
       "        [41, 10, 39],\n",
       "        [42, 32, 22],\n",
       "        [42, 40, 93],\n",
       "        [43, 38, 97],\n",
       "        [45,  2, 57],\n",
       "        [45, 94, 34],\n",
       "        [46, 76, 81],\n",
       "        [47, 50, 17],\n",
       "        [47, 83, 84],\n",
       "        [48, 91, 38],\n",
       "        [49, 22, 70],\n",
       "        [50, 50, 58],\n",
       "        [50, 59, 16],\n",
       "        [51, 33, 38],\n",
       "        [51, 57, 98],\n",
       "        [51, 75, 20],\n",
       "        [52,  4, 76],\n",
       "        [53, 38, 75],\n",
       "        [54, 49,  1],\n",
       "        [55, 57, 50],\n",
       "        [56, 45, 50],\n",
       "        [56, 59, 46],\n",
       "        [57,  3,  3],\n",
       "        [57, 54,  8],\n",
       "        [57, 82, 82],\n",
       "        [59, 31, 37],\n",
       "        [59, 64, 64],\n",
       "        [61, 64, 45],\n",
       "        [61, 87,  4],\n",
       "        [62, 89, 29],\n",
       "        [63,  4, 51],\n",
       "        [63, 33, 37],\n",
       "        [63, 62,  0],\n",
       "        [63, 62, 33],\n",
       "        [63, 68, 11],\n",
       "        [63, 79, 93],\n",
       "        [64, 60, 18],\n",
       "        [65,  1, 67],\n",
       "        [65, 43, 31],\n",
       "        [65, 89, 77],\n",
       "        [65, 98, 94],\n",
       "        [66, 10, 60],\n",
       "        [66, 25, 93],\n",
       "        [67, 91, 75],\n",
       "        [68,  0,  3],\n",
       "        [68, 73, 34],\n",
       "        [68, 81, 17],\n",
       "        [68, 91, 82],\n",
       "        [69, 23, 63],\n",
       "        [69, 25, 58],\n",
       "        [69, 42, 68],\n",
       "        [69, 48,  3],\n",
       "        [70,  8, 16],\n",
       "        [70, 96, 57],\n",
       "        [71,  8, 80],\n",
       "        [71, 62, 97],\n",
       "        [71, 70, 97],\n",
       "        [73, 24,  8],\n",
       "        [73, 29, 19],\n",
       "        [74, 26, 76],\n",
       "        [74, 47, 41],\n",
       "        [74, 75, 62],\n",
       "        [75, 25, 20],\n",
       "        [75, 32, 77],\n",
       "        [76, 70, 12],\n",
       "        [78,  8, 83],\n",
       "        [78, 37, 63],\n",
       "        [78, 59, 44],\n",
       "        [80, 24, 21],\n",
       "        [80, 32, 54],\n",
       "        [80, 33, 89],\n",
       "        [80, 56, 72],\n",
       "        [81,  1, 12],\n",
       "        [82, 38, 33],\n",
       "        [82, 55, 65],\n",
       "        [83,  4,  6],\n",
       "        [83, 34, 30],\n",
       "        [83, 70,  8],\n",
       "        [83, 75, 20],\n",
       "        [83, 89, 84],\n",
       "        [84, 65, 15],\n",
       "        [85, 60, 60],\n",
       "        [86, 33, 98],\n",
       "        [86, 60, 43],\n",
       "        [86, 76, 23],\n",
       "        [86, 86, 71],\n",
       "        [87, 44, 36],\n",
       "        [87, 53, 75],\n",
       "        [87, 65, 94],\n",
       "        [87, 73, 32],\n",
       "        [88,  6, 25],\n",
       "        [88, 25, 28],\n",
       "        [88, 28, 77],\n",
       "        [89,  0, 19],\n",
       "        [89, 17, 77],\n",
       "        [89, 18, 13],\n",
       "        [89, 27, 99],\n",
       "        [90, 39, 30],\n",
       "        [90, 52, 50],\n",
       "        [91, 19, 21],\n",
       "        [91, 71, 88],\n",
       "        [91, 94, 55],\n",
       "        [92, 13, 66],\n",
       "        [92, 25, 46],\n",
       "        [92, 43, 48],\n",
       "        [92, 52, 74],\n",
       "        [93, 79, 47],\n",
       "        [93, 93, 48],\n",
       "        [94,  0, 71],\n",
       "        [94, 63, 83],\n",
       "        [95, 21, 54],\n",
       "        [95, 46, 20],\n",
       "        [95, 61,  5],\n",
       "        [95, 71, 90],\n",
       "        [96, 40, 35],\n",
       "        [96, 97, 84],\n",
       "        [97,  8,  8],\n",
       "        [97, 40, 55],\n",
       "        [97, 80, 89],\n",
       "        [98, 51, 93],\n",
       "        [98, 56, 99],\n",
       "        [98, 75, 45],\n",
       "        [98, 97, 42],\n",
       "        [99, 47, 89],\n",
       "        [99, 94,  0]]))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkNaN(trKsquare1),checkNaN(trKsquare2),checkNaN(trKsquare3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Good\n",
      "Good\n",
      "Good\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(None, None, None)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check if there are NaN values in Lamb\n",
    "checkNaN(Lamb1), checkNaN(Lamb2), checkNaN(Lamb3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check NaN Eigenvectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Good\n",
      "NaN\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(None,\n",
       " array([[ 83613,      0,      0],\n",
       "        [ 83613,      0,      1],\n",
       "        [ 83613,      0,      2],\n",
       "        [ 83625,      0,      0],\n",
       "        [ 83625,      0,      1],\n",
       "        [ 83625,      0,      2],\n",
       "        [ 88047,      0,      0],\n",
       "        [ 88047,      0,      1],\n",
       "        [ 88047,      0,      2],\n",
       "        [ 88129,      0,      0],\n",
       "        [ 88129,      0,      1],\n",
       "        [ 88129,      0,      2],\n",
       "        [ 91978,      0,      0],\n",
       "        [ 91978,      0,      1],\n",
       "        [ 91978,      0,      2],\n",
       "        [ 91988,      0,      0],\n",
       "        [ 91988,      0,      1],\n",
       "        [ 91988,      0,      2],\n",
       "        [104444,      0,      0],\n",
       "        [104444,      0,      1],\n",
       "        [104444,      0,      2],\n",
       "        [104450,      0,      0],\n",
       "        [104450,      0,      1],\n",
       "        [104450,      0,      2],\n",
       "        [104693,      0,      0],\n",
       "        [104693,      0,      1],\n",
       "        [104693,      0,      2],\n",
       "        [141591,      0,      0],\n",
       "        [141591,      0,      1],\n",
       "        [141591,      0,      2],\n",
       "        [141595,      0,      0],\n",
       "        [141595,      0,      1],\n",
       "        [141595,      0,      2],\n",
       "        [158277,      0,      0],\n",
       "        [158277,      0,      1],\n",
       "        [158277,      0,      2],\n",
       "        [158283,      0,      0],\n",
       "        [158283,      0,      1],\n",
       "        [158283,      0,      2],\n",
       "        [186980,      0,      0],\n",
       "        [186980,      0,      1],\n",
       "        [186980,      0,      2],\n",
       "        [293781,      0,      0],\n",
       "        [293781,      0,      1],\n",
       "        [293781,      0,      2],\n",
       "        [293789,      0,      0],\n",
       "        [293789,      0,      1],\n",
       "        [293789,      0,      2],\n",
       "        [302024,      0,      0],\n",
       "        [302024,      0,      1],\n",
       "        [302024,      0,      2],\n",
       "        [302028,      0,      0],\n",
       "        [302028,      0,      1],\n",
       "        [302028,      0,      2],\n",
       "        [309895,      0,      0],\n",
       "        [309895,      0,      1],\n",
       "        [309895,      0,      2],\n",
       "        [309901,      0,      0],\n",
       "        [309901,      0,      1],\n",
       "        [309901,      0,      2],\n",
       "        [314197,      0,      0],\n",
       "        [314197,      0,      1],\n",
       "        [314197,      0,      2],\n",
       "        [314209,      0,      0],\n",
       "        [314209,      0,      1],\n",
       "        [314209,      0,      2],\n",
       "        [314242,      0,      0],\n",
       "        [314242,      0,      1],\n",
       "        [314242,      0,      2],\n",
       "        [314246,      0,      0],\n",
       "        [314246,      0,      1],\n",
       "        [314246,      0,      2],\n",
       "        [318260,      0,      0],\n",
       "        [318260,      0,      1],\n",
       "        [318260,      0,      2],\n",
       "        [318264,      0,      0],\n",
       "        [318264,      0,      1],\n",
       "        [318264,      0,      2],\n",
       "        [322402,      0,      0],\n",
       "        [322402,      0,      1],\n",
       "        [322402,      0,      2],\n",
       "        [322404,      0,      0],\n",
       "        [322404,      0,      1],\n",
       "        [322404,      0,      2],\n",
       "        [322561,      0,      0],\n",
       "        [322561,      0,      1],\n",
       "        [322561,      0,      2],\n",
       "        [322573,      0,      0],\n",
       "        [322573,      0,      1],\n",
       "        [322573,      0,      2]]))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkNaN(Lamb1v), checkNaN(Lamb2v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0167, 0.1000, 0.1000], dtype=torch.float64)\n",
      "tensor([[-0.5150, -0.8572,  0.0000],\n",
      "        [-0.8572,  0.5150,  0.0000],\n",
      "        [-0.0000,  0.0000,  1.0000]], dtype=torch.float64)\n",
      "tensor([0.1000, 0.1000, 0.0167], dtype=torch.float64)\n",
      "tensor([[    nan,     nan,     nan],\n",
      "        [ 0.0000,  0.0000, -1.0000],\n",
      "        [-0.5150, -0.8572,  0.0000]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "idx=83613\n",
    "print(Lamb1[idx])\n",
    "print(Lamb1v[idx])\n",
    "print(Lamb2[idx])\n",
    "print(Lamb2v[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[    nan,     nan,     nan],\n",
       "        [ 0.0000,  0.0000, -1.0000],\n",
       "        [-0.5150, -0.8572,  0.0000]], dtype=torch.float64)"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Lamb2v[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.1000, dtype=torch.float64) tensor([-0., 0., 1.], dtype=torch.float64)\n",
      "tensor(0.1000, dtype=torch.float64) tensor([ 0.,  0., -1.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "print(Lamb1[idx,2],Lamb1v[idx,2,:])\n",
    "print(Lamb2[idx,1],Lamb2v[idx,1,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-2.2102e-02, -3.6788e-02,  0.0000e+00],\n",
      "        [-3.6788e-02, -6.1231e-02,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00, -1.4901e-09]], dtype=torch.float64)\n",
      "tensor([-0.0221, -0.0368,  0.0000], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "AlI = W[idx] - Lamb2[idx,1] * torch.eye(3)\n",
    "print(AlI)\n",
    "print(AlI[:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0779, -0.0368,  0.1000],\n",
      "        [-0.0368,  0.0388,  0.1000],\n",
      "        [ 0.0000,  0.0000,  0.2000]], dtype=torch.float64)\n",
      "tensor([ 0.0779, -0.0368,  0.0000], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "Ale = W[idx] - Lamb2[idx,1] * Lamb2v[idx,1,:]\n",
    "print(Ale)\n",
    "print(Ale[:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.8572, -0.5150,  0.0000], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "newvI=torch.cross(AlI[:,0],Lamb2v[idx,1,:])\n",
    "newvI = newvI / torch.linalg.norm(newvI)\n",
    "print(newvI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4270, 0.9042, 0.0000], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "newve=torch.cross(Ale[:,0],Lamb2v[idx,1,:])\n",
    "newve = newve / torch.linalg.norm(newve)\n",
    "print(newve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.7476, -0.6642,  0.0000], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "v = torch.from_numpy(np.array([0.0466, -0.0414,0]))\n",
    "normv = v / torch.linalg.norm(v)\n",
    "print(normv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the index of a NaN value "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NaN\n",
      "203\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0, 58, 40)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IndNaN = checkNaN(trKsquare3)\n",
    "print(len(IndNaN))\n",
    "a,b,c = IndNaN[0]\n",
    "a,b,c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the PDS matrix at the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NaN\n",
      "tensor([[ 3.2045e+00,  8.4586e-02, -4.9083e+04],\n",
      "        [ 8.4586e-02,  1.9889e+00, -4.1348e+04],\n",
      "        [-4.9083e+04, -4.1348e+04,  1.5827e+09]], dtype=torch.float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.2045e+00,  8.4586e-02, -4.9083e+04,  8.4586e-02,  1.9889e+00,\n",
       "         -4.1348e+04, -4.9083e+04, -4.1348e+04,  1.5827e+09]],\n",
       "       dtype=torch.float64)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IndNaN = checkNaN(trKsquare3)\n",
    "test = W[a,b,c]\n",
    "\n",
    "# Compare the entries in W and reshaped W\n",
    "print(W[a,b,c])\n",
    "W.permute((3,4,0,1,2)).reshape((1,9,*W.shape[0:3]))[:,:,a,b,c]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000e+00, -7.3552e-16, -7.2760e-12],\n",
       "        [ 7.3552e-16,  0.0000e+00, -7.2760e-12],\n",
       "        [ 7.2760e-12,  7.2760e-12,  0.0000e+00]], dtype=torch.float64)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if the matrix is symetric\n",
    "test - test.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.4917e+08, dtype=torch.float64)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.det(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([3.6900e-02, 2.5541e+00, 1.5827e+09], dtype=torch.float64)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if the matrix is pd\n",
    "torch.symeig(test)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[[-5.1234e+00]]],\n",
       "\n",
       "\n",
       "         [[[ 7.7144e+00]]],\n",
       "\n",
       "\n",
       "         [[[ 1.5827e+09]]]]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Compare the eigenvalues \n",
    "\n",
    "# compute the eigenvalues using vSymEig\n",
    "eigTest = TV.vSymEig(test.reshape((1,9,1,1,1)), eigenvectors=False)[0]\n",
    "eigTest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0., 0.],\n",
       "        [0., 0., 0.],\n",
       "        [0., 0., 0.]], dtype=torch.float64)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "symtest = (test  + test.t())/2\n",
    "symtest-symtest.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([3.6900e-02, 2.5541e+00, 1.5827e+09], dtype=torch.float64)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.symeig(symtest)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[[-5.1234e+00]]],\n",
       "\n",
       "\n",
       "         [[[ 7.7144e+00]]],\n",
       "\n",
       "\n",
       "         [[[ 1.5827e+09]]]]])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eigTest = TV.vSymEig(symtest.reshape((1,9,1,1,1)), eigenvectors=False)[0]\n",
    "eigTest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the eigenvalues computed previously"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3.6900e-02, 2.5541e+00, 1.5827e+09], dtype=torch.float64)\n",
      "tensor([3.6900e-02, 2.5541e+00, 1.5827e+09], dtype=torch.float64)\n",
      "tensor([[-5.1234e+00,  7.7144e+00,  1.5827e+09]])\n"
     ]
    }
   ],
   "source": [
    "print(Lamb1[a,b,c])\n",
    "print(Lamb1[a,b,c])\n",
    "print(Lamb3[:,:,a,b,c])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.5827e+09, dtype=torch.float64)\n",
      "tensor(99630454.1432, dtype=torch.float64)\n",
      "tensor(23692036.4034, dtype=torch.float64)\n",
      "tensor(1301049.0521, dtype=torch.float64)\n",
      "tensor(584963.1336, dtype=torch.float64)\n",
      "tensor(34745428.6724, dtype=torch.float64)\n",
      "tensor(1547995.9492, dtype=torch.float64)\n",
      "tensor(1451792.2938, dtype=torch.float64)\n",
      "tensor(19738185.0698, dtype=torch.float64)\n",
      "tensor(2.3989e+08, dtype=torch.float64)\n",
      "tensor(3.5612e+08, dtype=torch.float64)\n",
      "tensor(7272763.7192, dtype=torch.float64)\n",
      "tensor(1.4080e+08, dtype=torch.float64)\n",
      "tensor(5.9377e+08, dtype=torch.float64)\n",
      "tensor(4600587.4629, dtype=torch.float64)\n",
      "tensor(2.2020e+08, dtype=torch.float64)\n",
      "tensor(5.4657e+08, dtype=torch.float64)\n",
      "tensor(1.0701e+08, dtype=torch.float64)\n",
      "tensor(18485300.0272, dtype=torch.float64)\n",
      "tensor(2.8805e+09, dtype=torch.float64)\n",
      "tensor(56554163.8674, dtype=torch.float64)\n",
      "tensor(2.2583e+08, dtype=torch.float64)\n",
      "tensor(29901546.3968, dtype=torch.float64)\n",
      "tensor(16539309.6501, dtype=torch.float64)\n",
      "tensor(8.5493e+08, dtype=torch.float64)\n",
      "tensor(73219.2085, dtype=torch.float64)\n",
      "tensor(59998719.5835, dtype=torch.float64)\n",
      "tensor(9.1606e+08, dtype=torch.float64)\n",
      "tensor(12221236.4433, dtype=torch.float64)\n",
      "tensor(34990628.4369, dtype=torch.float64)\n",
      "tensor(4546841.9975, dtype=torch.float64)\n",
      "tensor(1.0871e+09, dtype=torch.float64)\n",
      "tensor(1.4478e+09, dtype=torch.float64)\n",
      "tensor(4.0561e+08, dtype=torch.float64)\n",
      "tensor(16589404.6830, dtype=torch.float64)\n",
      "tensor(91032304.4662, dtype=torch.float64)\n",
      "tensor(23074249.5524, dtype=torch.float64)\n",
      "tensor(1.1989e+08, dtype=torch.float64)\n",
      "tensor(4.0253e+09, dtype=torch.float64)\n",
      "tensor(5.5618e+09, dtype=torch.float64)\n",
      "tensor(43922941.3406, dtype=torch.float64)\n",
      "tensor(2432463.3782, dtype=torch.float64)\n",
      "tensor(45532226.5222, dtype=torch.float64)\n",
      "tensor(7241005.6522, dtype=torch.float64)\n",
      "tensor(5.3779e+08, dtype=torch.float64)\n",
      "tensor(6.0400e+08, dtype=torch.float64)\n",
      "tensor(1.0875e+08, dtype=torch.float64)\n",
      "tensor(41058561.4726, dtype=torch.float64)\n",
      "tensor(1.7461e+08, dtype=torch.float64)\n",
      "tensor(4.1755e+08, dtype=torch.float64)\n",
      "tensor(7.6171e+08, dtype=torch.float64)\n",
      "tensor(1.9407e+08, dtype=torch.float64)\n",
      "tensor(8421326.5567, dtype=torch.float64)\n",
      "tensor(1.5064e+08, dtype=torch.float64)\n",
      "tensor(2.0585e+08, dtype=torch.float64)\n",
      "tensor(4.6277e+08, dtype=torch.float64)\n",
      "tensor(63959336.0619, dtype=torch.float64)\n",
      "tensor(20280446.9177, dtype=torch.float64)\n",
      "tensor(2.9545e+09, dtype=torch.float64)\n",
      "tensor(10895461.3731, dtype=torch.float64)\n",
      "tensor(3.5182e+08, dtype=torch.float64)\n",
      "tensor(6.5484e+08, dtype=torch.float64)\n",
      "tensor(53155996.9870, dtype=torch.float64)\n",
      "tensor(48487972.4501, dtype=torch.float64)\n",
      "tensor(1.7918e+08, dtype=torch.float64)\n",
      "tensor(3.1345e+08, dtype=torch.float64)\n",
      "tensor(1.6684e+10, dtype=torch.float64)\n",
      "tensor(2515281.6342, dtype=torch.float64)\n",
      "tensor(93580814.0830, dtype=torch.float64)\n",
      "tensor(3.8897e+08, dtype=torch.float64)\n",
      "tensor(34379118.4519, dtype=torch.float64)\n",
      "tensor(47094.4326, dtype=torch.float64)\n",
      "tensor(4.0195e+08, dtype=torch.float64)\n",
      "tensor(3.3133e+10, dtype=torch.float64)\n",
      "tensor(2.3373e+08, dtype=torch.float64)\n",
      "tensor(94293377.0188, dtype=torch.float64)\n",
      "tensor(7936225.2225, dtype=torch.float64)\n",
      "tensor(698371.9361, dtype=torch.float64)\n",
      "tensor(31752864.0726, dtype=torch.float64)\n",
      "tensor(4728291.0475, dtype=torch.float64)\n",
      "tensor(4485902.6973, dtype=torch.float64)\n",
      "tensor(2.2971e+08, dtype=torch.float64)\n",
      "tensor(41457289.1650, dtype=torch.float64)\n",
      "tensor(93914148.4625, dtype=torch.float64)\n",
      "tensor(85165510.3942, dtype=torch.float64)\n",
      "tensor(95304264.9636, dtype=torch.float64)\n",
      "tensor(1.7839e+08, dtype=torch.float64)\n",
      "tensor(48922.6124, dtype=torch.float64)\n",
      "tensor(2514462.9587, dtype=torch.float64)\n",
      "tensor(3.5126e+09, dtype=torch.float64)\n",
      "tensor(118359.8068, dtype=torch.float64)\n",
      "tensor(3.6773e+08, dtype=torch.float64)\n",
      "tensor(1.5227e+09, dtype=torch.float64)\n",
      "tensor(335885.7749, dtype=torch.float64)\n",
      "tensor(25202984.1829, dtype=torch.float64)\n",
      "tensor(2.5261e+08, dtype=torch.float64)\n",
      "tensor(14777719.7890, dtype=torch.float64)\n",
      "tensor(4178786.4857, dtype=torch.float64)\n",
      "tensor(3.4247e+08, dtype=torch.float64)\n",
      "tensor(1.8859e+08, dtype=torch.float64)\n",
      "tensor(2.5963e+09, dtype=torch.float64)\n",
      "tensor(1.5364e+10, dtype=torch.float64)\n",
      "tensor(59100147.5122, dtype=torch.float64)\n",
      "tensor(1235772.5725, dtype=torch.float64)\n",
      "tensor(2.4214e+08, dtype=torch.float64)\n",
      "tensor(24104491.7295, dtype=torch.float64)\n",
      "tensor(5.9153e+08, dtype=torch.float64)\n",
      "tensor(1.7116e+08, dtype=torch.float64)\n",
      "tensor(13986501.6832, dtype=torch.float64)\n",
      "tensor(802801.2973, dtype=torch.float64)\n",
      "tensor(12006.0601, dtype=torch.float64)\n",
      "tensor(59721874.2073, dtype=torch.float64)\n",
      "tensor(1.1529e+08, dtype=torch.float64)\n",
      "tensor(33037.3459, dtype=torch.float64)\n",
      "tensor(2.1541e+09, dtype=torch.float64)\n",
      "tensor(9957153.9734, dtype=torch.float64)\n",
      "tensor(2.7965e+09, dtype=torch.float64)\n",
      "tensor(17984867.3930, dtype=torch.float64)\n",
      "tensor(1.1825e+08, dtype=torch.float64)\n",
      "tensor(67617.6379, dtype=torch.float64)\n",
      "tensor(22790553.0917, dtype=torch.float64)\n",
      "tensor(2.3534e+08, dtype=torch.float64)\n",
      "tensor(235638.1073, dtype=torch.float64)\n",
      "tensor(17634104.5074, dtype=torch.float64)\n",
      "tensor(2.9267e+08, dtype=torch.float64)\n",
      "tensor(44821.3263, dtype=torch.float64)\n",
      "tensor(19519954.4446, dtype=torch.float64)\n",
      "tensor(6991741.3816, dtype=torch.float64)\n",
      "tensor(7071969.8522, dtype=torch.float64)\n",
      "tensor(2.3730e+08, dtype=torch.float64)\n",
      "tensor(99394891.8151, dtype=torch.float64)\n",
      "tensor(1.0261e+10, dtype=torch.float64)\n",
      "tensor(5535607.3641, dtype=torch.float64)\n",
      "tensor(8084034.8452, dtype=torch.float64)\n",
      "tensor(6.5833e+08, dtype=torch.float64)\n",
      "tensor(1.8996e+08, dtype=torch.float64)\n",
      "tensor(1354663.9189, dtype=torch.float64)\n",
      "tensor(1.8574e+08, dtype=torch.float64)\n",
      "tensor(3185240.5482, dtype=torch.float64)\n",
      "tensor(376369.6638, dtype=torch.float64)\n",
      "tensor(1.1496e+08, dtype=torch.float64)\n",
      "tensor(2.0038e+09, dtype=torch.float64)\n",
      "tensor(1291020.0907, dtype=torch.float64)\n",
      "tensor(42378520.0533, dtype=torch.float64)\n",
      "tensor(5.0439e+08, dtype=torch.float64)\n",
      "tensor(14370292.1338, dtype=torch.float64)\n",
      "tensor(1.1585e+08, dtype=torch.float64)\n",
      "tensor(1.7030e+08, dtype=torch.float64)\n",
      "tensor(622995.1548, dtype=torch.float64)\n",
      "tensor(9310681.9457, dtype=torch.float64)\n",
      "tensor(2.4258e+08, dtype=torch.float64)\n",
      "tensor(6.3453e+10, dtype=torch.float64)\n",
      "tensor(4601286.3463, dtype=torch.float64)\n",
      "tensor(7960556.7849, dtype=torch.float64)\n",
      "tensor(3544536.1808, dtype=torch.float64)\n",
      "tensor(1.9575e+08, dtype=torch.float64)\n",
      "tensor(10918825.9938, dtype=torch.float64)\n",
      "tensor(1.1486e+09, dtype=torch.float64)\n",
      "tensor(86725211.1447, dtype=torch.float64)\n",
      "tensor(8.2821e+09, dtype=torch.float64)\n",
      "tensor(3.3679e+09, dtype=torch.float64)\n",
      "tensor(982381.6314, dtype=torch.float64)\n",
      "tensor(67887469.6316, dtype=torch.float64)\n",
      "tensor(866457.3599, dtype=torch.float64)\n",
      "tensor(1.3943e+08, dtype=torch.float64)\n",
      "tensor(8.7910e+08, dtype=torch.float64)\n",
      "tensor(93157166.3544, dtype=torch.float64)\n",
      "tensor(1.5340e+08, dtype=torch.float64)\n",
      "tensor(6614242.1916, dtype=torch.float64)\n",
      "tensor(3.3717e+09, dtype=torch.float64)\n",
      "tensor(1022182.3117, dtype=torch.float64)\n",
      "tensor(32231927.8185, dtype=torch.float64)\n",
      "tensor(4.8949e+10, dtype=torch.float64)\n",
      "tensor(13175169.6219, dtype=torch.float64)\n",
      "tensor(13242762.5486, dtype=torch.float64)\n",
      "tensor(410235.7550, dtype=torch.float64)\n",
      "tensor(3221954.5210, dtype=torch.float64)\n",
      "tensor(40796813.9599, dtype=torch.float64)\n",
      "tensor(1992933.9256, dtype=torch.float64)\n",
      "tensor(6.5948e+08, dtype=torch.float64)\n",
      "tensor(2.1345e+11, dtype=torch.float64)\n",
      "tensor(12975991.0680, dtype=torch.float64)\n",
      "tensor(4.8463e+08, dtype=torch.float64)\n",
      "tensor(83363681.2184, dtype=torch.float64)\n",
      "tensor(29037958.7446, dtype=torch.float64)\n",
      "tensor(7.7401e+08, dtype=torch.float64)\n",
      "tensor(71631281.2737, dtype=torch.float64)\n",
      "tensor(1572878.9552, dtype=torch.float64)\n",
      "tensor(23433997.4457, dtype=torch.float64)\n",
      "tensor(150170.8658, dtype=torch.float64)\n",
      "tensor(21883377.8007, dtype=torch.float64)\n",
      "tensor(2.7314e+08, dtype=torch.float64)\n",
      "tensor(10329474.5134, dtype=torch.float64)\n",
      "tensor(1.0120e+08, dtype=torch.float64)\n",
      "tensor(22297841.1067, dtype=torch.float64)\n",
      "tensor(1.1461e+09, dtype=torch.float64)\n",
      "tensor(4.8984e+08, dtype=torch.float64)\n",
      "tensor(3166137.6737, dtype=torch.float64)\n",
      "tensor(7.3817e+08, dtype=torch.float64)\n",
      "tensor(2.4541e+09, dtype=torch.float64)\n",
      "tensor(10647256.4396, dtype=torch.float64)\n",
      "tensor(1.3331e+09, dtype=torch.float64)\n",
      "tensor(21687693.4415, dtype=torch.float64)\n",
      "tensor(12006.0601, dtype=torch.float64) tensor(2.1345e+11, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "min_max = torch.max(Lamb1[a,b,c])\n",
    "min_max_idx = IndNaN[0]\n",
    "max_max = torch.max(Lamb1[a,b,c])\n",
    "max_max_idx = IndNaN[0]\n",
    "for idx in IndNaN:\n",
    "  max_eval = torch.max(Lamb1[idx[0],idx[1],idx[2]])\n",
    "  print(max_eval)\n",
    "  if max_eval < min_max:\n",
    "    min_max = max_eval\n",
    "    min_max_idx = idx\n",
    "  if max_eval > max_max:\n",
    "    max_max = max_eval\n",
    "    max_max_idx = idx\n",
    "print(min_max, max_max)\n",
    "#  print(torch.max(Lamb3[:,:,idx[0],idx[1],idx[2]]))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([63, 62,  0])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min_max_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[   16.2573,    52.6822,   422.1447],\n",
      "        [   52.6822,   175.2501,  1430.2641],\n",
      "        [  422.1447,  1430.2641, 11817.8356]], dtype=torch.float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[   16.2573,    52.6822,   422.1447,    52.6822,   175.2501,  1430.2641,\n",
       "           422.1447,  1430.2641, 11817.8356]], dtype=torch.float64)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d,e,f = min_max_idx\n",
    "test2 = W[d,e,f]\n",
    "\n",
    "# Compare the entries in W and reshaped W\n",
    "print(W[d,e,f])\n",
    "W.permute((3,4,0,1,2)).reshape((1,9,*W.shape[0:3]))[:,:,d,e,f]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000e+00,  7.1054e-15,  1.1369e-13],\n",
       "        [-7.1054e-15,  0.0000e+00, -2.2737e-13],\n",
       "        [-1.1369e-13,  2.2737e-13,  0.0000e+00]], dtype=torch.float64)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if the matrix is symetric\n",
    "test2 - test2.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(8.9804e-06, dtype=torch.float64)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.det(test2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2.2786e-10, 3.2829e+00, 1.2006e+04], dtype=torch.float64)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if the matrix is pd\n",
    "torch.symeig(test2)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[[-1.4993e-09]]],\n",
       "\n",
       "\n",
       "         [[[ 3.2829e+00]]],\n",
       "\n",
       "\n",
       "         [[[ 1.2006e+04]]]]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eigTest2 = TV.vSymEig(test2.reshape((1,9,1,1,1)), eigenvectors=False)[0]\n",
    "eigTest2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigTest2, evtest2 = TV.vSymEig(test2.reshape((1,9,1,1,1)), eigenvectors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigTest2[eigTest2<0] = 1.0e-7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-9.6699e-10, -5.3043e-10, -1.8852e-12],\n",
       "        [ 1.1594e+00,  2.0761e+00,  4.7352e-02],\n",
       "        [ 2.2444e+01,  1.6583e+02,  1.1818e+04]])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recontest2 = torch.einsum('...ji,...j,...ik->...ik',evtest2,eigTest2,evtest2)\n",
    "recontest2.reshape((3,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[   16.2573,    52.6822,   422.1447],\n",
       "        [   52.6822,   175.2501,  1430.2641],\n",
       "        [  422.1447,  1430.2641, 11817.8356]], dtype=torch.float64)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 3, 1, 1, 1]) torch.Size([3, 3])\n"
     ]
    }
   ],
   "source": [
    "print(recontest2.shape,test2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.0921e-09,  2.1677e+00,  5.3571e+02])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.einsum('...ij,...j->...i',recontest2.reshape((3,3)),evtest2.reshape((3,3))[0,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([8.0309e-08, 5.9480e-08, 3.5459e-09])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1.e-7 * evtest2.reshape((3,3))[0,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From https://www.geometrictools.com/GTE/Mathematics/SymmetricEigensolver3x3.h\n",
    "def NISymmetricEigensolver3x3(A):\n",
    " # The input matrix must be symmetric, so only the unique elements\n",
    " # must be specified: a00, a01, a02, a11, a12, and a22.  The\n",
    " # eigenvalues are sorted in ascending order: eval0 <= eval1 <= eval2.\n",
    "\n",
    "  #Precondition the matrix by factoring out the maximum absolute\n",
    "  # value of the components.  This guards against floating-point\n",
    "  # overflow when computing the eigenvalues.\n",
    "  A_max = torch.max(A.reshape((*A.shape[0:3],9)),3)[0]\n",
    "  scaled_A = torch.einsum('...ij,...->...ij',A, 1.0 / A_max)\n",
    "  #          Real max0 = std::max(std::fabs(a00), std::fabs(a01));\n",
    "  #          Real max1 = std::max(std::fabs(a02), std::fabs(a11));\n",
    "  #          Real max2 = std::max(std::fabs(a12), std::fabs(a22));\n",
    "  #          Real maxAbsElement = std::max(std::max(max0, max1), max2);\n",
    "  if torch.any(A_max < 1e-12):\n",
    "    print ('small A_max detected')\n",
    "    #        if (maxAbsElement == (Real)0)\n",
    "    #        {\n",
    "    #            // A is the zero matrix.\n",
    "    #            eval[0] = (Real)0;\n",
    "    #            eval[1] = (Real)0;\n",
    "    #            eval[2] = (Real)0;\n",
    "    #            evec[0] = { (Real)1, (Real)0, (Real)0 };\n",
    "    #            evec[1] = { (Real)0, (Real)1, (Real)0 };\n",
    "    #            evec[2] = { (Real)0, (Real)0, (Real)1 };\n",
    "    #            return;\n",
    "    #        }\n",
    "\n",
    "    #        Real invMaxAbsElement = (Real)1 / maxAbsElement;\n",
    "    #        a00 *= invMaxAbsElement;\n",
    "    #        a01 *= invMaxAbsElement;\n",
    "    #        a02 *= invMaxAbsElement;\n",
    "    #        a11 *= invMaxAbsElement;\n",
    "    #        a12 *= invMaxAbsElement;\n",
    "    #        a22 *= invMaxAbsElement;\n",
    "\n",
    "  a01 = scaled_A[:,:,:,0,1]  \n",
    "  a02 = scaled_A[:,:,:,0,2]\n",
    "  a12 = scaled_A[:,:,:,1,2]\n",
    "  nrm = a01 * a01 + a02 * a02 + a12 * a12\n",
    "  #        Real norm = a01 * a01 + a02 * a02 + a12 * a12;\n",
    "  if torch.any(nrm < 1e-12):\n",
    "    print('Need diagonal case')\n",
    "                #// The matrix is diagonal.\n",
    "                #eval[0] = a00;\n",
    "                #eval[1] = a11;\n",
    "                #eval[2] = a22;\n",
    "                #evec[0] = { (Real)1, (Real)0, (Real)0 };\n",
    "                #evec[1] = { (Real)0, (Real)1, (Real)0 };\n",
    "                #evec[2] = { (Real)0, (Real)0, (Real)1 };\n",
    "  else:\n",
    "      #          // Compute the eigenvalues of A.\n",
    "\n",
    "      #          // In the PDF mentioned previously, B = (A - q*I)/p, where\n",
    "      #          // q = tr(A)/3 with tr(A) the trace of A (sum of the diagonal\n",
    "      #          // entries of A) and where p = sqrt(tr((A - q*I)^2)/6).\n",
    "      #          Real q = (a00 + a11 + a22) / (Real)3;\n",
    "    q = (scaled_A[:,:,:,0,0] + scaled_A[:,:,:,1,1] + scaled_A[:,:,:,2,2]) / 3.0\n",
    "\n",
    "      #          // The matrix A - q*I is represented by the following, where\n",
    "      #          // b00, b11 and b22 are computed after these comments,\n",
    "      #          //   +-           -+\n",
    "      #          //   | b00 a01 a02 |\n",
    "      #          //   | a01 b11 a12 |\n",
    "      #          //   | a02 a12 b22 |\n",
    "      #          //   +-           -+\n",
    "      #          Real b00 = a00 - q;\n",
    "      #          Real b11 = a11 - q;\n",
    "      #          Real b22 = a22 - q;\n",
    "    b00 = scaled_A[:,:,:,0,0] - q\n",
    "    b11 = scaled_A[:,:,:,1,1] - q\n",
    "    b22 = scaled_A[:,:,:,2,2] - q\n",
    "\n",
    "      #          // The is the variable p mentioned in the PDF.\n",
    "      #          Real p = std::sqrt((b00 * b00 + b11 * b11 + b22 * b22 + norm * (Real)2) / (Real)6);\n",
    "    p = torch.sqrt((b00 * b00 + b11 * b11 + b22 * b22 + nrm * 2.0) / 6.0)\n",
    "      #          // We need det(B) = det((A - q*I)/p) = det(A - q*I)/p^3.  The\n",
    "      #          // value det(A - q*I) is computed using a cofactor expansion\n",
    "      #          // by the first row of A - q*I.  The cofactors are c00, c01\n",
    "      #          // and c02 and the determinant is b00*c00 - a01*c01 + a02*c02.\n",
    "      #          // The det(B) is then computed finally by the division\n",
    "      #          // with p^3.\n",
    "      #          Real c00 = b11 * b22 - a12 * a12;\n",
    "      #          Real c01 = a01 * b22 - a12 * a02;\n",
    "      #          Real c02 = a01 * a12 - b11 * a02;\n",
    "      #          Real det = (b00 * c00 - a01 * c01 + a02 * c02) / (p * p * p);\n",
    "    c00 = b11 * b22 - a12 * a12\n",
    "    c01 = a01 * b22 - a12 * a02\n",
    "    c02 = a01 * a12 - b11 * a02\n",
    "    det = (b00 * c00 - a01 * c01 + a02 * c02) / (p * p * p)\n",
    "\n",
    "      #          // The halfDet value is cos(3*theta) mentioned in the PDF. The\n",
    "      #          // acos(z) function requires |z| <= 1, but will fail silently\n",
    "      #          // and return NaN if the input is larger than 1 in magnitude.\n",
    "      #          // To avoid this problem due to rounding errors, the halfDet\n",
    "      #          // value is clamped to [-1,1].\n",
    "      #          Real halfDet = det * (Real)0.5;\n",
    "      #          halfDet = std::min(std::max(halfDet, (Real)-1), (Real)1);\n",
    "    halfDet = 0.5 * det\n",
    "    halfDet[halfDet < -1] = -1\n",
    "    halfDet[halfDet > 1] = 1\n",
    "\n",
    "      #          // The eigenvalues of B are ordered as\n",
    "      #          // beta0 <= beta1 <= beta2.  The number of digits in\n",
    "      #          // twoThirdsPi is chosen so that, whether float or double,\n",
    "      #          // the floating-point number is the closest to theoretical\n",
    "      #          // 2*pi/3.\n",
    "      #          Real angle = std::acos(halfDet) / (Real)3;\n",
    "      #          Real const twoThirdsPi = (Real)2.09439510239319549;\n",
    "      #          Real beta2 = std::cos(angle) * (Real)2;\n",
    "      #          Real beta0 = std::cos(angle + twoThirdsPi) * (Real)2;\n",
    "      #          Real beta1 = -(beta0 + beta2);\n",
    "    angle = torch.acos(halfDet) / 3.0\n",
    "    twoThirdsPi = 2.09439510239319549\n",
    "    beta2 = torch.cos(angle) * 2.0\n",
    "    beta0 = torch.cos(angle + twoThirdsPi) * 2.0\n",
    "    beta1 = -(beta0 + beta2)\n",
    "\n",
    "      #          // The eigenvalues of A are ordered as\n",
    "      #          // alpha0 <= alpha1 <= alpha2.\n",
    "      #          eval[0] = q + p * beta0;\n",
    "      #          eval[1] = q + p * beta1;\n",
    "      #          eval[2] = q + p * beta2;\n",
    "    eval0 = q + p * beta0\n",
    "    eval1 = q + p * beta1\n",
    "    eval2 = q + p * beta2\n",
    "    \n",
    "\n",
    "      #          // Compute the eigenvectors so that the set\n",
    "      #          // {evec[0], evec[1], evec[2]} is right handed and\n",
    "      #          // orthonormal.\n",
    "      #          if (halfDet >= (Real)0)\n",
    "      #          {\n",
    "      #              ComputeEigenvector0(a00, a01, a02, a11, a12, a22, eval[2], evec[2]);\n",
    "      #              ComputeEigenvector1(a00, a01, a02, a11, a12, a22, evec[2], eval[1], evec[1]);\n",
    "      #              evec[0] = Cross(evec[1], evec[2]);\n",
    "      #          }\n",
    "      #          else\n",
    "      #          {\n",
    "      #              ComputeEigenvector0(a00, a01, a02, a11, a12, a22, eval[0], evec[0]);\n",
    "      #              ComputeEigenvector1(a00, a01, a02, a11, a12, a22, evec[0], eval[1], evec[1]);\n",
    "      #              evec[2] = Cross(evec[0], evec[1]);\n",
    "      #          }\n",
    "      #      }\n",
    "  print('Skipping eigenvector computation for now')\n",
    "\n",
    "      #      // The preconditioning scaled the matrix A, which scales the\n",
    "      #      // eigenvalues.  Revert the scaling.\n",
    "      #      eval[0] *= maxAbsElement;\n",
    "      #      eval[1] *= maxAbsElement;\n",
    "      #      eval[2] *= maxAbsElement;\n",
    "  evals = torch.stack((eval0 * A_max, eval1 * A_max, eval2 * A_max),3) \n",
    "  print(eval0.shape, evals.shape)\n",
    "\n",
    "  #          SortEigenstuff<Real>()(sortType, true, eval, evec);\n",
    "  return(torch.sort(evals,-1)[0])"
   ]
  },
  {
   "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.6.12"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}