VESPER_prob_rot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import copy\n",
    "import math\n",
    "import multiprocessing\n",
    "import os\n",
    "\n",
    "import mrcfile\n",
    "import numba\n",
    "import numpy as np\n",
    "import pyfftw\n",
    "import scipy.fft\n",
    "from numba.typed import List\n",
    "from scipy.ndimage import convolve\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(204, 215, 236, 4)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob_maps=np.load(\"interpolated_prediction_8097.npy\")\n",
    "prob_maps_chain=np.load(\"interpolated_prediction_8097_A.npy\")\n",
    "prob_maps.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyfftw.config.PLANNER_EFFORT = \"FFTW_MEASURE\"\n",
    "pyfftw.config.NUM_THREADS = multiprocessing.cpu_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class mrc_obj:\n",
    "    def __init__(self, path):\n",
    "        mrc = mrcfile.open(path)\n",
    "        data = mrc.data\n",
    "        header = mrc.header\n",
    "        self.xdim = int(header.nx)\n",
    "        self.ydim = int(header.ny)\n",
    "        self.zdim = int(header.nz)\n",
    "        self.xwidth = mrc.voxel_size.x\n",
    "        self.ywidth = mrc.voxel_size.y\n",
    "        self.zwidth = mrc.voxel_size.z\n",
    "        self.cent = [\n",
    "            self.xdim * 0.5,\n",
    "            self.ydim * 0.5,\n",
    "            self.zdim * 0.5,\n",
    "        ]\n",
    "        self.orig = {\"x\": header.origin.x, \"y\": header.origin.y, \"z\": header.origin.z}\n",
    "        self.data = np.swapaxes(copy.deepcopy(data), 0, 2)\n",
    "        self.dens = data.flatten()\n",
    "        self.vec = np.zeros((self.xdim, self.ydim, self.zdim, 3), dtype=\"float32\")\n",
    "        self.dsum = None\n",
    "        self.Nact = None\n",
    "        self.ave = None\n",
    "        self.std_norm_ave = None\n",
    "        self.std = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def mrc_set_vox_size(mrc, th=0.01, voxel_size=7.0):\n",
    "\n",
    "    # set shape and size\n",
    "    size = mrc.xdim * mrc.ydim * mrc.zdim\n",
    "    shape = (mrc.xdim, mrc.ydim, mrc.zdim)\n",
    "\n",
    "    # if th < 0 add th to all value\n",
    "    if th < 0:\n",
    "        mrc.dens = mrc.dens - th\n",
    "        th = 0.0\n",
    "\n",
    "    # Trim all the values less than threshold\n",
    "    mrc.dens[mrc.dens < th] = 0.0\n",
    "    mrc.data[mrc.data < th] = 0.0\n",
    "\n",
    "    # calculate dmax distance for non-zero entries\n",
    "    non_zero_index_list = np.array(np.nonzero(mrc.data)).T\n",
    "    # non_zero_index_list[:, [2, 0]] = non_zero_index_list[:, [0, 2]]\n",
    "    cent_arr = np.array(mrc.cent)\n",
    "    d2_list = np.linalg.norm(non_zero_index_list - cent_arr, axis=1)\n",
    "    dmax = max(d2_list)\n",
    "\n",
    "    # dmax = math.sqrt(mrc.cent[0] ** 2 + mrc.cent[1] ** 2 + mrc.cent[2] ** 2)\n",
    "\n",
    "    print(\"#dmax=\" + str(dmax / mrc.xwidth))\n",
    "    dmax = dmax * mrc.xwidth\n",
    "\n",
    "    # set new center\n",
    "    new_cent = [\n",
    "        mrc.cent[0] * mrc.xwidth + mrc.orig[\"x\"],\n",
    "        mrc.cent[1] * mrc.xwidth + mrc.orig[\"y\"],\n",
    "        mrc.cent[2] * mrc.xwidth + mrc.orig[\"z\"],\n",
    "    ]\n",
    "\n",
    "    tmp_size = 2 * dmax / voxel_size\n",
    "\n",
    "    # find the minimum size of the map\n",
    "    #     b = y = 2**math.ceil(math.log2(tmp_size))\n",
    "    #     while 1:\n",
    "    #         while y < tmp_size:\n",
    "    #             y = y * 3\n",
    "    #             continue\n",
    "    #         if y < b:\n",
    "    #             b = y\n",
    "    #         if y % 2 != 0:\n",
    "    #             break\n",
    "    #         y = y / 2\n",
    "\n",
    "    #     new_xdim = int(b)\n",
    "    new_xdim = pyfftw.next_fast_len(int(tmp_size))\n",
    "\n",
    "    # set new origins\n",
    "    new_orig = {\n",
    "        \"x\": new_cent[0] - 0.5 * new_xdim * voxel_size,\n",
    "        \"y\": new_cent[1] - 0.5 * new_xdim * voxel_size,\n",
    "        \"z\": new_cent[2] - 0.5 * new_xdim * voxel_size,\n",
    "    }\n",
    "\n",
    "    # create new mrc object\n",
    "    mrc_set = copy.deepcopy(mrc)\n",
    "    mrc_set.orig = new_orig\n",
    "    mrc_set.xdim = new_xdim\n",
    "    mrc_set.ydim = new_xdim\n",
    "    mrc_set.zdim = new_xdim\n",
    "    mrc_set.cent = new_cent\n",
    "    mrc_set.xwidth = mrc_set.ywidth = mrc_set.zwidth = voxel_size\n",
    "\n",
    "    print(\"Nvox= \" + str(mrc_set.xdim) + \", \" + str(mrc_set.ydim) + \", \" + str(mrc_set.zdim))\n",
    "    print(\"cent= \" + str(new_cent[0]) + \", \" + str(new_cent[1]) + \", \" + str(new_cent[2]))\n",
    "    print(\"ori= \" + str(new_orig[\"x\"]) + \", \" + str(new_orig[\"y\"]) + \", \" + str(new_orig[\"z\"]))\n",
    "\n",
    "    return mrc, mrc_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def calc(stp, endp, pos, mrc1_data, fsiv):\n",
    "    dtotal = 0.0\n",
    "    pos2 = np.zeros((3,))\n",
    "\n",
    "    for xp in range(stp[0], endp[0]):\n",
    "        rx = float(xp) - pos[0]\n",
    "        rx = rx ** 2\n",
    "        for yp in range(stp[1], endp[1]):\n",
    "            ry = float(yp) - pos[1]\n",
    "            ry = ry ** 2\n",
    "            for zp in range(stp[2], endp[2]):\n",
    "                rz = float(zp) - pos[2]\n",
    "                rz = rz ** 2\n",
    "                d2 = rx + ry + rz\n",
    "                v = mrc1_data[xp][yp][zp] * math.exp(-1.5 * d2 * fsiv)\n",
    "                dtotal += v\n",
    "                pos2[0] += v * float(xp)\n",
    "                pos2[1] += v * float(yp)\n",
    "                pos2[2] += v * float(zp)\n",
    "    return dtotal, pos2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fastVEC(mrc1, mrc2, dreso=16.0):\n",
    "\n",
    "    xydim = mrc1.xdim * mrc1.ydim\n",
    "    Ndata = mrc2.xdim * mrc2.ydim * mrc2.zdim\n",
    "\n",
    "    print(len(mrc2.dens))\n",
    "\n",
    "    print(\"#Start VEC\")\n",
    "    gstep = mrc1.xwidth\n",
    "    fs = (dreso / gstep) * 0.5\n",
    "    fs = fs ** 2\n",
    "    fsiv = 1.0 / fs\n",
    "    fmaxd = (dreso / gstep) * 2.0\n",
    "    print(\"#maxd= {fmaxd}\".format(fmaxd=fmaxd))\n",
    "    print(\"#fsiv= \" + str(fsiv))\n",
    "\n",
    "    dsum = 0.0\n",
    "    Nact = 0\n",
    "\n",
    "    list_d = []\n",
    "\n",
    "    for x in tqdm(range(mrc2.xdim)):\n",
    "        for y in range(mrc2.ydim):\n",
    "            for z in range(mrc2.zdim):\n",
    "                stp = [0] * 3\n",
    "                endp = [0] * 3\n",
    "                ind2 = 0\n",
    "                ind = 0\n",
    "\n",
    "                pos = [0.0] * 3\n",
    "                pos2 = [0.0] * 3\n",
    "                ori = [0.0] * 3\n",
    "\n",
    "                tmpcd = [0.0] * 3\n",
    "\n",
    "                v, dtotal, rd = 0.0, 0.0, 0.0\n",
    "\n",
    "                pos[0] = (x * mrc2.xwidth + mrc2.orig[\"x\"] - mrc1.orig[\"x\"]) / mrc1.xwidth\n",
    "                pos[1] = (y * mrc2.xwidth + mrc2.orig[\"y\"] - mrc1.orig[\"y\"]) / mrc1.xwidth\n",
    "                pos[2] = (z * mrc2.xwidth + mrc2.orig[\"z\"] - mrc1.orig[\"z\"]) / mrc1.xwidth\n",
    "\n",
    "                ind = mrc2.xdim * mrc2.ydim * z + mrc2.xdim * y + x\n",
    "\n",
    "                # check density\n",
    "\n",
    "                if (\n",
    "                    pos[0] < 0\n",
    "                    or pos[1] < 0\n",
    "                    or pos[2] < 0\n",
    "                    or pos[0] >= mrc1.xdim\n",
    "                    or pos[1] >= mrc1.ydim\n",
    "                    or pos[2] >= mrc1.zdim\n",
    "                ):\n",
    "                    mrc2.dens[ind] = 0.0\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue\n",
    "\n",
    "                if mrc1.data[int(pos[0])][int(pos[1])][int(pos[2])] == 0:\n",
    "                    mrc2.dens[ind] = 0.0\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue\n",
    "\n",
    "                ori[0] = pos[0]\n",
    "                ori[1] = pos[1]\n",
    "                ori[2] = pos[2]\n",
    "\n",
    "                # Start Point\n",
    "                stp[0] = int(pos[0] - fmaxd)\n",
    "                stp[1] = int(pos[1] - fmaxd)\n",
    "                stp[2] = int(pos[2] - fmaxd)\n",
    "\n",
    "                # set start and end point\n",
    "                if stp[0] < 0:\n",
    "                    stp[0] = 0\n",
    "                if stp[1] < 0:\n",
    "                    stp[1] = 0\n",
    "                if stp[2] < 0:\n",
    "                    stp[2] = 0\n",
    "\n",
    "                endp[0] = int(pos[0] + fmaxd + 1)\n",
    "                endp[1] = int(pos[1] + fmaxd + 1)\n",
    "                endp[2] = int(pos[2] + fmaxd + 1)\n",
    "\n",
    "                if endp[0] >= mrc1.xdim:\n",
    "                    endp[0] = mrc1.xdim\n",
    "                if endp[1] >= mrc1.ydim:\n",
    "                    endp[1] = mrc1.ydim\n",
    "                if endp[2] >= mrc1.zdim:\n",
    "                    endp[2] = mrc1.zdim\n",
    "\n",
    "                # setup for numba acc\n",
    "                stp_t = List()\n",
    "                endp_t = List()\n",
    "                pos_t = List()\n",
    "                [stp_t.append(x) for x in stp]\n",
    "                [endp_t.append(x) for x in endp]\n",
    "                [pos_t.append(x) for x in pos]\n",
    "\n",
    "                # compute the total density\n",
    "                dtotal, pos2 = calc(stp_t, endp_t, pos_t, mrc1.data, fsiv)\n",
    "\n",
    "                mrc2.dens[ind] = dtotal\n",
    "                mrc2.data[x][y][z] = dtotal\n",
    "\n",
    "                if dtotal == 0:\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue\n",
    "\n",
    "                rd = 1.0 / dtotal\n",
    "\n",
    "                pos2[0] *= rd\n",
    "                pos2[1] *= rd\n",
    "                pos2[2] *= rd\n",
    "\n",
    "                tmpcd[0] = pos2[0] - pos[0]\n",
    "                tmpcd[1] = pos2[1] - pos[1]\n",
    "                tmpcd[2] = pos2[2] - pos[2]\n",
    "\n",
    "                dvec = math.sqrt(tmpcd[0] ** 2 + tmpcd[1] ** 2 + tmpcd[2] ** 2)\n",
    "\n",
    "                if dvec == 0:\n",
    "                    dvec = 1.0\n",
    "\n",
    "                rdvec = 1.0 / dvec\n",
    "\n",
    "                mrc2.vec[x][y][z][0] = tmpcd[0] * rdvec\n",
    "                mrc2.vec[x][y][z][1] = tmpcd[1] * rdvec\n",
    "                mrc2.vec[x][y][z][2] = tmpcd[2] * rdvec\n",
    "\n",
    "                dsum += dtotal\n",
    "                Nact += 1\n",
    "\n",
    "    print(\"#End LDP\")\n",
    "    print(dsum)\n",
    "    print(Nact)\n",
    "\n",
    "    mrc2.dsum = dsum\n",
    "    mrc2.Nact = Nact\n",
    "    mrc2.ave = dsum / float(Nact)\n",
    "    mrc2.std = np.linalg.norm(mrc2.dens[mrc2.dens > 0])\n",
    "    mrc2.std_norm_ave = np.linalg.norm(mrc2.dens[mrc2.dens > 0] - mrc2.ave)\n",
    "\n",
    "    print(\"#MAP AVE={ave} STD={std} STD_norm={std_norm}\".format(ave=mrc2.ave, std=mrc2.std, std_norm=mrc2.std_norm_ave))\n",
    "    # return False\n",
    "    return mrc2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fastVEC_prob(mrc1, mrc2, dreso=16.0):\n",
    "\n",
    "    xydim = mrc1.xdim * mrc1.ydim\n",
    "    Ndata = mrc2.xdim * mrc2.ydim * mrc2.zdim\n",
    "\n",
    "    print(len(mrc2.dens))\n",
    "\n",
    "    print(\"#Start VEC\")\n",
    "    gstep = mrc1.xwidth\n",
    "    fs = (dreso / gstep) * 0.5\n",
    "    fs = fs ** 2\n",
    "    fsiv = 1.0 / fs\n",
    "    fmaxd = (dreso / gstep) * 2.0\n",
    "    print(\"#maxd= {fmaxd}\".format(fmaxd=fmaxd))\n",
    "    print(\"#fsiv= \" + str(fsiv))\n",
    "\n",
    "    dsum = 0.0\n",
    "    Nact = 0\n",
    "\n",
    "    list_d = []\n",
    "\n",
    "    for x in tqdm(range(mrc2.xdim)):\n",
    "        for y in range(mrc2.ydim):\n",
    "            for z in range(mrc2.zdim):\n",
    "                stp = [0] * 3\n",
    "                endp = [0] * 3\n",
    "                ind2 = 0\n",
    "                ind = 0\n",
    "\n",
    "                pos = [0.0] * 3\n",
    "                pos2 = [0.0] * 3\n",
    "                ori = [0.0] * 3\n",
    "\n",
    "                tmpcd = [0.0] * 3\n",
    "\n",
    "                v, dtotal, rd = 0.0, 0.0, 0.0\n",
    "\n",
    "                pos[0] = (x * mrc2.xwidth + mrc2.orig[\"x\"] - mrc1.orig[\"x\"]) / mrc1.xwidth\n",
    "                pos[1] = (y * mrc2.xwidth + mrc2.orig[\"y\"] - mrc1.orig[\"y\"]) / mrc1.xwidth\n",
    "                pos[2] = (z * mrc2.xwidth + mrc2.orig[\"z\"] - mrc1.orig[\"z\"]) / mrc1.xwidth\n",
    "\n",
    "                ind = mrc2.xdim * mrc2.ydim * z + mrc2.xdim * y + x\n",
    "\n",
    "                # check density\n",
    "\n",
    "                if (\n",
    "                    pos[0] < 0\n",
    "                    or pos[1] < 0\n",
    "                    or pos[2] < 0\n",
    "                    or pos[0] >= mrc1.xdim\n",
    "                    or pos[1] >= mrc1.ydim\n",
    "                    or pos[2] >= mrc1.zdim\n",
    "                ):\n",
    "                    mrc2.dens[ind] = 0.0\n",
    "                    #mrc2.vec[x][y][z][0] = 0.0\n",
    "                    #mrc2.vec[x][y][z][1] = 0.0\n",
    "                    #mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue\n",
    "\n",
    "                if mrc1.data[int(pos[0])][int(pos[1])][int(pos[2])] == 0:\n",
    "                    mrc2.dens[ind] = 0.0\n",
    "                    #mrc2.vec[x][y][z][0] = 0.0\n",
    "                    #mrc2.vec[x][y][z][1] = 0.0\n",
    "                    #mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue\n",
    "\n",
    "                ori[0] = pos[0]\n",
    "                ori[1] = pos[1]\n",
    "                ori[2] = pos[2]\n",
    "\n",
    "                # Start Point\n",
    "                stp[0] = int(pos[0] - fmaxd)\n",
    "                stp[1] = int(pos[1] - fmaxd)\n",
    "                stp[2] = int(pos[2] - fmaxd)\n",
    "\n",
    "                # set start and end point\n",
    "                if stp[0] < 0:\n",
    "                    stp[0] = 0\n",
    "                if stp[1] < 0:\n",
    "                    stp[1] = 0\n",
    "                if stp[2] < 0:\n",
    "                    stp[2] = 0\n",
    "\n",
    "                endp[0] = int(pos[0] + fmaxd + 1)\n",
    "                endp[1] = int(pos[1] + fmaxd + 1)\n",
    "                endp[2] = int(pos[2] + fmaxd + 1)\n",
    "\n",
    "                if endp[0] >= mrc1.xdim:\n",
    "                    endp[0] = mrc1.xdim\n",
    "                if endp[1] >= mrc1.ydim:\n",
    "                    endp[1] = mrc1.ydim\n",
    "                if endp[2] >= mrc1.zdim:\n",
    "                    endp[2] = mrc1.zdim\n",
    "\n",
    "                # setup for numba acc\n",
    "                stp_t = List()\n",
    "                endp_t = List()\n",
    "                pos_t = List()\n",
    "                [stp_t.append(x) for x in stp]\n",
    "                [endp_t.append(x) for x in endp]\n",
    "                [pos_t.append(x) for x in pos]\n",
    "\n",
    "                # compute the total density\n",
    "                dtotal, pos2 = calc(stp_t, endp_t, pos_t, mrc1.data, fsiv)\n",
    "\n",
    "                mrc2.dens[ind] = dtotal\n",
    "                mrc2.data[x][y][z] = dtotal\n",
    "\n",
    "                ''''if dtotal == 0:\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\n",
    "                    continue'''\n",
    "\n",
    "                ''''rd = 1.0 / dtotal\n",
    "\n",
    "                pos2[0] *= rd\n",
    "                pos2[1] *= rd\n",
    "                pos2[2] *= rd\n",
    "\n",
    "                tmpcd[0] = pos2[0] - pos[0]\n",
    "                tmpcd[1] = pos2[1] - pos[1]\n",
    "                tmpcd[2] = pos2[2] - pos[2]\n",
    "\n",
    "                dvec = math.sqrt(tmpcd[0] ** 2 + tmpcd[1] ** 2 + tmpcd[2] ** 2)\n",
    "\n",
    "                if dvec == 0:\n",
    "                    dvec = 1.0\n",
    "\n",
    "                rdvec = 1.0 / dvec'''\n",
    "\n",
    "                #mrc2.vec[x][y][z][0] = tmpcd[0] * rdvec\n",
    "                #mrc2.vec[x][y][z][1] = tmpcd[1] * rdvec\n",
    "                #mrc2.vec[x][y][z][2] = tmpcd[2] * rdvec\n",
    "\n",
    "                dsum += dtotal\n",
    "                Nact += 1\n",
    "\n",
    "    print(\"#End LDP\")\n",
    "    print(dsum)\n",
    "    print(Nact)\n",
    "\n",
    "    mrc2.dsum = dsum\n",
    "    mrc2.Nact = Nact\n",
    "    mrc2.ave = dsum / float(Nact)\n",
    "    mrc2.std = np.linalg.norm(mrc2.dens[mrc2.dens > 0])\n",
    "    mrc2.std_norm_ave = np.linalg.norm(mrc2.dens[mrc2.dens > 0] - mrc2.ave)\n",
    "\n",
    "    print(\"#MAP AVE={ave} STD={std} STD_norm={std_norm}\".format(ave=mrc2.ave, std=mrc2.std, std_norm=mrc2.std_norm_ave))\n",
    "    # return False\n",
    "    return mrc2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def rot_pos(vec, angle, inv=False):\n",
    "    r = R.from_euler(\"zyx\", angle, degrees=True)\n",
    "    if inv:\n",
    "        r = r.inv()\n",
    "    rotated_vec = r.apply(vec)\n",
    "    return rotated_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def rot_pos_mtx(mtx, vec):\n",
    "    mtx = mtx.astype(np.float32)\n",
    "    vec = vec.astype(np.float32)\n",
    "    return vec @ mtx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def rot_mrc_combine(old_pos, new_pos, new_vec_shape):\n",
    "    # combined_arr = np.hstack((old_pos,new_pos))\n",
    "    combined_arr = np.concatenate((old_pos, new_pos), axis=1)\n",
    "\n",
    "    combined_arr = combined_arr[\n",
    "        (combined_arr[:, 0] >= 0)\n",
    "        & (combined_arr[:, 1] >= 0)\n",
    "        & (combined_arr[:, 2] >= 0)\n",
    "        & (combined_arr[:, 0] < new_vec_shape[0])\n",
    "        & (combined_arr[:, 1] < new_vec_shape[0])\n",
    "        & (combined_arr[:, 2] < new_vec_shape[0])\n",
    "    ]\n",
    "\n",
    "    return combined_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True)\n",
    "def test_rot_mrc(orig_mrc_data, orig_mrc_vec, mtx):\n",
    "\n",
    "    mtx = mtx.astype(np.float32)\n",
    "\n",
    "    # init new vec and dens array\n",
    "    new_vec_array = np.zeros_like(orig_mrc_vec)\n",
    "    new_data_array = np.zeros_like(orig_mrc_data)\n",
    "\n",
    "    dim = orig_mrc_data.shape[0]\n",
    "    cent = 0.5 * dim\n",
    "\n",
    "    for x in range(dim):\n",
    "        for y in range(dim):\n",
    "            for z in range(dim):\n",
    "                pos = (np.array([x, y, z]) - cent).astype(np.float32)\n",
    "                pos2 = pos @ mtx\n",
    "                pos2 += cent\n",
    "\n",
    "                if pos2[0] < 0 or pos2[1] < 0 or pos2[2] < 0 or pos2[0] >= dim or pos2[1] >= dim or pos2[2] >= dim:\n",
    "                    continue\n",
    "\n",
    "                pos2 = pos2.astype(np.int32)\n",
    "\n",
    "                if orig_mrc_data[pos2[0]][pos2[1]][pos2[2]] == 0:\n",
    "                    continue\n",
    "\n",
    "                vec = orig_mrc_vec[pos2[0]][pos2[1]][pos2[2]].astype(np.float32)\n",
    "                vec2 = vec @ mtx.T\n",
    "\n",
    "                pos = pos.astype(np.int32)\n",
    "                new_vec_array[pos[0]][pos[1]][pos[2]] = vec2\n",
    "                new_data_array[pos[0]][pos[1]][pos[2]] = orig_mrc_data[pos2[0]][pos2[1]][pos2[2]]\n",
    "\n",
    "    return new_vec_array, new_data_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def rot_mrc_prob(orig_mrc_data, orig_mrc_vec,mrc_search_p1_data, mrc_search_p2_data, mrc_search_p3_data, mrc_search_p4_data,  angle):\n",
    "    dim = orig_mrc_vec.shape[0]\n",
    "\n",
    "    new_pos = np.array(np.meshgrid(np.arange(dim), np.arange(dim), np.arange(dim),)).T.reshape(-1, 3)\n",
    "\n",
    "    cent = 0.5 * float(dim)\n",
    "    new_pos = new_pos - cent\n",
    "\n",
    "    r = R.from_euler(\"ZYX\", angle, degrees=True)\n",
    "    mtx = r.as_matrix()\n",
    "    mtx[np.isclose(mtx, 0, atol=1e-15)] = 0\n",
    "\n",
    "    old_pos = rot_pos_mtx(np.flip(mtx).T, new_pos) + cent\n",
    "\n",
    "    combined_arr = np.hstack((old_pos, new_pos))\n",
    "    \n",
    "    in_bound_mask = (\n",
    "        (old_pos[:, 0] >= 0)\n",
    "        & (old_pos[:, 1] >= 0)\n",
    "        & (old_pos[:, 2] >= 0)\n",
    "        & (old_pos[:, 0] < dim)\n",
    "        & (old_pos[:, 1] < dim)\n",
    "        & (old_pos[:, 2] < dim)\n",
    "    )\n",
    "\n",
    "    combined_arr = combined_arr[in_bound_mask]\n",
    "\n",
    "    combined_arr = combined_arr.astype(np.int32)\n",
    "\n",
    "    index_arr = combined_arr[:, 0:3]\n",
    "\n",
    "  \n",
    "    dens_mask = orig_mrc_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "    dens_mask_p1 = mrc_search_p1_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "    dens_mask_p2 = mrc_search_p2_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "    dens_mask_p3 = mrc_search_p3_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "    dens_mask_p4 = mrc_search_p4_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "\n",
    "    non_zero_rot_list = combined_arr[dens_mask]\n",
    "    non_zero_rot_list_p1 = combined_arr[dens_mask_p1]\n",
    "    non_zero_rot_list_p2 = combined_arr[dens_mask_p2]\n",
    "    non_zero_rot_list_p3 = combined_arr[dens_mask_p3]\n",
    "    non_zero_rot_list_p4 = combined_arr[dens_mask_p4]\n",
    "\n",
    "\n",
    "\n",
    "    non_zero_vec = orig_mrc_vec[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\n",
    "\n",
    "    non_zero_dens = orig_mrc_data[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\n",
    "    \n",
    "    non_zero_dens_p1 = mrc_search_p1_data[non_zero_rot_list_p1[:, 0], non_zero_rot_list_p1[:, 1], non_zero_rot_list_p1[:, 2]]\n",
    "    \n",
    "    non_zero_dens_p2 = mrc_search_p2_data[non_zero_rot_list_p2[:, 0], non_zero_rot_list_p2[:, 1], non_zero_rot_list_p2[:, 2]]\n",
    "    \n",
    "    non_zero_dens_p3 = mrc_search_p3_data[non_zero_rot_list_p3[:, 0], non_zero_rot_list_p3[:, 1], non_zero_rot_list_p3[:, 2]]\n",
    "    \n",
    "    non_zero_dens_p4 = mrc_search_p4_data[non_zero_rot_list_p4[:, 0], non_zero_rot_list_p4[:, 1], non_zero_rot_list_p4[:, 2]]\n",
    "\n",
    "    new_vec = rot_pos_mtx(np.flip(mtx), non_zero_vec)\n",
    "\n",
    "\n",
    "    new_vec_array = np.zeros_like(orig_mrc_vec)\n",
    "    new_data_array = np.zeros_like(orig_mrc_data)\n",
    "    new_data_array_p1 = np.zeros_like(mrc_search_p1_data)\n",
    "    new_data_array_p2 = np.zeros_like(mrc_search_p2_data)\n",
    "    new_data_array_p3 = np.zeros_like(mrc_search_p3_data)\n",
    "    new_data_array_p4 = np.zeros_like(mrc_search_p4_data)\n",
    "\n",
    "    for vec, ind, dens in zip(new_vec, (non_zero_rot_list[:, 3:6] + cent).astype(int), non_zero_dens):\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][0] = vec[0]\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][1] = vec[1]\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][2] = vec[2]\n",
    "        new_data_array[ind[0]][ind[1]][ind[2]] = dens\n",
    "        \n",
    "    for ind, dens in zip((non_zero_rot_list_p1[:, 3:6] + cent).astype(int), non_zero_dens_p1):\n",
    "\n",
    "        new_data_array_p1[ind[0]][ind[1]][ind[2]] = dens\n",
    "        \n",
    "    for ind, dens in zip((non_zero_rot_list_p2[:, 3:6] + cent).astype(int), non_zero_dens_p2):\n",
    "        new_data_array_p2[ind[0]][ind[1]][ind[2]] = dens\n",
    "        \n",
    "    for ind, dens in zip((non_zero_rot_list_p3[:, 3:6] + cent).astype(int), non_zero_dens_p3):\n",
    "        new_data_array_p3[ind[0]][ind[1]][ind[2]] = dens\n",
    "        \n",
    "    for ind, dens in zip((non_zero_rot_list_p4[:, 3:6] + cent).astype(int), non_zero_dens_p4):\n",
    "        new_data_array_p4[ind[0]][ind[1]][ind[2]] = dens\n",
    "\n",
    "    return new_vec_array, new_data_array, new_data_array_p1, new_data_array_p2, new_data_array_p3, new_data_array_p4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot_mrc(orig_mrc_data, orig_mrc_vec, angle):\n",
    "    dim = orig_mrc_vec.shape[0]\n",
    "\n",
    "    new_pos = np.array(np.meshgrid(np.arange(dim), np.arange(dim), np.arange(dim),)).T.reshape(-1, 3)\n",
    "\n",
    "    # set the center\n",
    "    cent = 0.5 * float(dim)\n",
    "\n",
    "    # get relative positions from center\n",
    "    new_pos = new_pos - cent\n",
    "    # print(new_pos)\n",
    "\n",
    "    # init the rotation by euler angle\n",
    "    r = R.from_euler(\"ZYX\", angle, degrees=True)\n",
    "    mtx = r.as_matrix()\n",
    "    mtx[np.isclose(mtx, 0, atol=1e-15)] = 0\n",
    "\n",
    "    # print(mtx)\n",
    "\n",
    "    old_pos = rot_pos_mtx(np.flip(mtx).T, new_pos) + cent\n",
    "\n",
    "    combined_arr = np.hstack((old_pos, new_pos))\n",
    "\n",
    "    # filter values outside the boundaries\n",
    "    in_bound_mask = (\n",
    "        (old_pos[:, 0] >= 0)\n",
    "        & (old_pos[:, 1] >= 0)\n",
    "        & (old_pos[:, 2] >= 0)\n",
    "        & (old_pos[:, 0] < dim)\n",
    "        & (old_pos[:, 1] < dim)\n",
    "        & (old_pos[:, 2] < dim)\n",
    "    )\n",
    "\n",
    "    combined_arr = combined_arr[in_bound_mask]\n",
    "\n",
    "    #     combined_arr = combined_arr[(combined_arr[:, 0] >= 0) &\n",
    "    #                                 (combined_arr[:, 1] >= 0) &\n",
    "    #                                 (combined_arr[:, 2] >= 0) &\n",
    "    #                                 (combined_arr[:, 0] < dim) &\n",
    "    #                                 (combined_arr[:, 1] < dim) &\n",
    "    #                                 (combined_arr[:, 2] < dim)]\n",
    "\n",
    "    # print(combined_arr)\n",
    "\n",
    "    # conversion indices to int\n",
    "    combined_arr = combined_arr.astype(np.int32)\n",
    "\n",
    "    # print(combined_arr.shape)\n",
    "\n",
    "    # combined_arr = rot_mrc_combine(old_pos, new_pos, new_vec_shape)\n",
    "\n",
    "    index_arr = combined_arr[:, 0:3]\n",
    "\n",
    "    # print(index_arr)\n",
    "    # print(np.count_nonzero(orig_mrc_data))\n",
    "\n",
    "    dens_mask = orig_mrc_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\n",
    "\n",
    "    # print(dens_mask.shape)\n",
    "    # print(dens_mask)\n",
    "\n",
    "    non_zero_rot_list = combined_arr[dens_mask]\n",
    "\n",
    "    # print(non_zero_rot_list.shape)\n",
    "    #     with np.printoptions(threshold=np.inf):\n",
    "    #         print(non_zero_rot_list[:, 0:3])\n",
    "\n",
    "    non_zero_vec = orig_mrc_vec[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\n",
    "\n",
    "    non_zero_dens = orig_mrc_data[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\n",
    "\n",
    "    # print(non_zero_dens)\n",
    "\n",
    "    # non_zero_dens[:, [2, 0]] = non_zero_dens[:, [0, 2]]\n",
    "    new_vec = rot_pos_mtx(np.flip(mtx), non_zero_vec)\n",
    "\n",
    "    # print(new_vec)\n",
    "\n",
    "    # init new vec and dens array\n",
    "    new_vec_array = np.zeros_like(orig_mrc_vec)\n",
    "    new_data_array = np.zeros_like(orig_mrc_data)\n",
    "\n",
    "    # print(new)\n",
    "\n",
    "    # fill in the new data\n",
    "    for vec, ind, dens in zip(new_vec, (non_zero_rot_list[:, 3:6] + cent).astype(int), non_zero_dens):\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][0] = vec[0]\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][1] = vec[1]\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][2] = vec[2]\n",
    "        new_data_array[ind[0]][ind[1]][ind[2]] = dens\n",
    "\n",
    "    return new_vec_array, new_data_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_best_trans_list(input_list):\n",
    "    \n",
    "    sum_arr = np.zeros_like(input_list[0])\n",
    "    #for arr in input_list:\n",
    "    dot_array=input_list[0]+input_list[1]+input_list[2]\n",
    "    #avg_dot=np.sum(dot_array)/(dot_array.shape[0])\n",
    "    ave_dot = np.mean(dot_array)\n",
    "    std_dot = np.std(dot_array)\n",
    "    dot_array_z=(dot_array-ave_dot)/std_dot\n",
    "    \n",
    "    prob_array=input_list[3]+input_list[4]+input_list[5]+input_list[6]\n",
    "    #avg_dot=np.sum(dot_array)/(dot_array.shape[0])\n",
    "    ave_prob = np.mean(prob_array)\n",
    "    std_prob = np.std(prob_array)\n",
    "    prob_array_z=(prob_array-ave_prob)/std_prob\n",
    "    \n",
    "    \n",
    "    #sum_arr = sum_arr+input_list[0]+input_list[1]+input_list[2]+(input_list[3] + input_list[4] + input_list[5] + input_list[6])/5000000.00\n",
    "    sum_arr = sum_arr+dot_array_z+prob_array_z\n",
    "    prob_arr=input_list[3] + input_list[4] + input_list[5] + input_list[6]\n",
    "    best = np.amax(sum_arr)\n",
    "    best_prob=np.amax(prob_array_z)\n",
    "    trans = np.unravel_index(sum_arr.argmax(), sum_arr.shape)\n",
    "    \n",
    "    return best, trans, best_prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_score(\n",
    "    target_map_data, search_map_data, target_map_vec, search_map_vec, trans, ave1, ave2, std1, std2, pstd1, pstd2, sco_prob_added, sco_prob_only\n",
    "):\n",
    "    px, py, pz = 0, 0, 0\n",
    "    dim = target_map_data.shape[0]\n",
    "    total = 0\n",
    "\n",
    "    t = np.array(trans)\n",
    "    if trans[0] > 0.5 * dim:\n",
    "        t[0] -= dim\n",
    "    if trans[1] > 0.5 * dim:\n",
    "        t[1] -= dim\n",
    "    if trans[2] > 0.5 * dim:\n",
    "        t[2] -= dim\n",
    "        \n",
    "    #duplicate this part for all probability maps\n",
    "\n",
    "    target_pos = np.array(np.meshgrid(np.arange(dim), np.arange(dim), np.arange(dim),)).T.reshape(-1, 3)\n",
    "\n",
    "    search_pos = target_pos + t\n",
    "\n",
    "    total += np.count_nonzero(target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]])\n",
    "\n",
    "    combined_arr = np.hstack((target_pos, search_pos))\n",
    "\n",
    "    combined_arr = combined_arr[\n",
    "        (combined_arr[:, 3] >= 0)\n",
    "        & (combined_arr[:, 4] >= 0)\n",
    "        & (combined_arr[:, 5] >= 0)\n",
    "        & (combined_arr[:, 3] < dim)\n",
    "        & (combined_arr[:, 4] < dim)\n",
    "        & (combined_arr[:, 5] < dim)\n",
    "    ]\n",
    "\n",
    "    target_pos = combined_arr[:, 0:3]\n",
    "    search_pos = combined_arr[:, 3:6]\n",
    "\n",
    "    d1 = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]]\n",
    "    d2 = search_map_data[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]]\n",
    "    \n",
    "    d1 = np.where(d1 <= 0, 0.0, d1)\n",
    "    d2 = np.where(d2 <= 0, 0.0, d1)\n",
    "\n",
    "    print(np.sum(d1))\n",
    "    print(np.sum(d2))\n",
    "\n",
    "    pd1 = np.where(d1 <= 0, 0.0, d1 - ave1)\n",
    "    pd2 = np.where(d2 <= 0, 0.0, d2 - ave2)\n",
    "\n",
    "    cc = np.sum(np.multiply(d1, d2))\n",
    "    pcc = np.sum(np.multiply(pd1, pd2))\n",
    "\n",
    "    target_zero_mask = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]] == 0\n",
    "    target_non_zero_mask = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]] > 0\n",
    "    search_non_zero_mask = search_map_data[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]] > 0\n",
    "    search_non_zero_count = np.count_nonzero(np.multiply(target_zero_mask, search_non_zero_mask))\n",
    "\n",
    "    trimmed_target_vec = target_map_vec[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]]\n",
    "    trimmed_search_vec = search_map_vec[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]]\n",
    "\n",
    "    total += search_non_zero_count\n",
    "\n",
    "    sco_arr = np.zeros_like(search_map_data)\n",
    "    sco = np.einsum(\"ij,ij->i\", trimmed_target_vec, trimmed_search_vec)\n",
    "    sco_arr[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]] = sco\n",
    "    sco_sum = np.sum(sco_arr)\n",
    "    Nm = np.count_nonzero(np.multiply(target_non_zero_mask, search_non_zero_mask))\n",
    "\n",
    "    print(\n",
    "        \"Overlap= \"\n",
    "        + str(float(Nm) / float(total))\n",
    "        + \" \"\n",
    "        + str(Nm)\n",
    "        + \"/\"\n",
    "        + str(total)\n",
    "        + \" CC= \"\n",
    "        + str(cc / (std1 * std2))\n",
    "        + \" PCC= \"\n",
    "        + str(pcc / (pstd1 * pstd2))\n",
    "        + \" Scoreplusprob= \"\n",
    "        + str(sco_prob_added)\n",
    "        + \" Scoreprobonly= \"\n",
    "        + str(sco_prob_only)\n",
    "    )\n",
    "    print(\"Score= \", sco_sum)\n",
    "    return sco_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fft_search_score_trans(target_X, target_Y, target_Z, search_vec, a, b, c, fft_object, ifft_object):\n",
    "    x2 = copy.deepcopy(search_vec[..., 0])\n",
    "    y2 = copy.deepcopy(search_vec[..., 1])\n",
    "    z2 = copy.deepcopy(search_vec[..., 2])\n",
    "\n",
    "    X2 = np.zeros_like(target_X)\n",
    "    np.copyto(a, x2)\n",
    "    np.copyto(X2, fft_object(a))\n",
    "    dot_X = target_X * X2\n",
    "    np.copyto(b, dot_X)\n",
    "    dot_x = np.zeros_like(x2)\n",
    "    np.copyto(dot_x, ifft_object(b))\n",
    "\n",
    "    Y2 = np.zeros_like(target_Y)\n",
    "    np.copyto(a, y2)\n",
    "    np.copyto(Y2, fft_object(a))\n",
    "    dot_Y = target_Y * Y2\n",
    "    np.copyto(b, dot_Y)\n",
    "    dot_y = np.zeros_like(y2)\n",
    "    np.copyto(dot_y, ifft_object(b))\n",
    "\n",
    "    Z2 = np.zeros_like(target_Z)\n",
    "    np.copyto(a, z2)\n",
    "    np.copyto(Z2, fft_object(a))\n",
    "    dot_Z = target_Z * Z2\n",
    "    np.copyto(b, dot_Z)\n",
    "    dot_z = np.zeros_like(z2)\n",
    "    np.copyto(dot_z, ifft_object(b))\n",
    "\n",
    "    #         if tuple(angle) == (0, 0, 30):\n",
    "    #             XX = [x12, y12, z12]\n",
    "\n",
    "    #         X2 = np.fft.rfftn(x2)\n",
    "    #         X12 = X1 * X2\n",
    "    #         x12 = np.fft.irfftn(X12, norm=\"forward\")\n",
    "\n",
    "    # #         if (tuple(angle) == (0,0,30)):\n",
    "    # #             XX = [X12,x12]\n",
    "\n",
    "    #         Y2 = np.fft.rfftn(y2)\n",
    "    #         Y12 = Y1 * Y2\n",
    "    #         y12 = np.fft.irfftn(Y12, norm=\"forward\")\n",
    "\n",
    "    #         Z2 = np.fft.rfftn(z2)\n",
    "    #         Z12 = Z1 * Z2\n",
    "    #         z12 = np.fft.irfftn(Z12, norm=\"forward\")\n",
    "\n",
    "    return find_best_trans_list([dot_x, dot_y, dot_z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object):\n",
    "    dot_product_list = []\n",
    "    for target_complex, query_real in zip(target_list, query_list):\n",
    "        \n",
    "        query_complex = np.zeros_like(target_complex)\n",
    "        np.copyto(a, query_real)\n",
    "        np.copyto(query_complex, fft_object(a))\n",
    "        dot_complex = target_complex * query_complex\n",
    "        np.copyto(b, dot_complex)\n",
    "        dot_real = np.zeros_like(query_real)\n",
    "        np.copyto(dot_real, ifft_object(b))\n",
    "        \n",
    "        dot_product_list.append(dot_real)\n",
    "        \n",
    "    #My changes\n",
    "    #dot_product=dot_product_list[0] + dot_product_list[1] + dot_product_list[2]\n",
    "        \n",
    "    return dot_product_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_search_score_trans_1d(target_X, search_data, a, b, fft_object, ifft_object, mode, ave=None):\n",
    "\n",
    "    x2 = copy.deepcopy(search_data)\n",
    "\n",
    "    if mode == \"Overlap\":\n",
    "        x2 = np.where(x2 > 0, 1.0, 0.0)\n",
    "    elif mode == \"CC\":\n",
    "        x2 = np.where(x2 > 0, x2, 0.0)\n",
    "    elif mode == \"PCC\":\n",
    "        x2 = np.where(x2 > 0, x2 - ave, 0.0)\n",
    "    elif mode == \"Laplacian\":\n",
    "        weights = np.array(\n",
    "            [\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\n",
    "                [[0.0, 1.0, 0.0], [1.0, -6.0, 1.0], [0.0, 1.0, 0.0]],\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\n",
    "            ]\n",
    "        )\n",
    "        x2 = convolve(x2, weights, mode=\"constant\")\n",
    "\n",
    "    X2 = np.zeros_like(target_X)\n",
    "    np.copyto(a, x2)\n",
    "    np.copyto(X2, fft_object(a))\n",
    "    dot_X = target_X * X2\n",
    "    np.copyto(b, dot_X)\n",
    "    dot_x = np.zeros_like(x2)\n",
    "    np.copyto(dot_x, ifft_object(b))\n",
    "\n",
    "    return find_best_trans_list([dot_x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def search_map_fft(mrc_P1, mrc_P2, mrc_P3, mrc_P4, mrc_target, mrc_search,mrc_search_p1,mrc_search_p2,mrc_search_p3,mrc_search_p4, TopN=10, ang=10, mode=\"VecProduct\", is_eval_mode=False):\n",
    "\n",
    "    if is_eval_mode:\n",
    "        print(\"#For Evaluation Mode\")\n",
    "        print(\"#Please use the same coordinate system and map size for map1 and map2.\")\n",
    "        print(\"#Example:\")\n",
    "        print(\"#In Chimera command line: open map1 and map2 as #0 and #1, then type\")\n",
    "        print(\"#> open map1.mrc\")\n",
    "        print(\"#> open map2.mrc\")\n",
    "        print(\"#> vop #1 resample onGrid #0\")\n",
    "        print(\"#> volume #2 save new.mrc\")\n",
    "        print(\"#Chimera will generate the resampled map2.mrc as new.mrc\")\n",
    "        return\n",
    "\n",
    "    #     x1 = copy.deepcopy(mrc_target.vec[:, :, :, 0])\n",
    "    #     y1 = copy.deepcopy(mrc_target.vec[:, :, :, 1])\n",
    "    #     z1 = copy.deepcopy(mrc_target.vec[:, :, :, 2])\n",
    "\n",
    "    # init the target map vectors\n",
    "    #does this part need to be changed?\n",
    "    x1 = copy.deepcopy(mrc_target.vec[:, :, :, 0])\n",
    "    p1 = copy.deepcopy(mrc_P1.data)\n",
    "    p2 = copy.deepcopy(mrc_P2.data)\n",
    "    p3 = copy.deepcopy(mrc_P3.data)\n",
    "    p4 = copy.deepcopy(mrc_P4.data)\n",
    "\n",
    "\n",
    "    if mode == \"VecProduct\":\n",
    "        y1 = copy.deepcopy(mrc_target.vec[:, :, :, 1])\n",
    "        z1 = copy.deepcopy(mrc_target.vec[:, :, :, 2])\n",
    "\n",
    "    # Postprocessing for other modes\n",
    "    if mode == \"Overlap\":\n",
    "        x1 = np.where(x1 > 0, 1.0, 0.0)\n",
    "    elif mode == \"CC\":\n",
    "        x1 = np.where(x1 > 0, x1, 0.0)\n",
    "    elif mode == \"PCC\":\n",
    "        x1 = np.where(x1 > 0, x1 - mrc_target.ave, 0.0)\n",
    "    elif mode == \"Laplacian\":\n",
    "        weights = np.array(\n",
    "            [\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\n",
    "                [[0.0, 1.0, 0.0], [1.0, -6.0, 1.0], [0.0, 1.0, 0.0]],\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\n",
    "            ]\n",
    "        )\n",
    "        x1 = convolve(x1, weights, mode=\"constant\")\n",
    "\n",
    "    d3 = mrc_target.xdim ** 3\n",
    "\n",
    "    rd3 = 1.0 / d3\n",
    "    \n",
    "    \n",
    "    X1 = np.fft.rfftn(x1)\n",
    "    X1 = np.conj(X1)\n",
    "    P1 = np.fft.rfftn(p1)\n",
    "    P1 = np.conj(P1)\n",
    "    P2 = np.fft.rfftn(p2)\n",
    "    P2 = np.conj(P2)\n",
    "    P3 = np.fft.rfftn(p3)\n",
    "    P3 = np.conj(P3)\n",
    "    P4 = np.fft.rfftn(p4)\n",
    "    P4 = np.conj(P4)\n",
    "    \n",
    "    #P4 = 1 - P2 - P3 - P1\n",
    "\n",
    "    if mode == \"VecProduct\":\n",
    "        Y1 = np.fft.rfftn(y1)\n",
    "        Y1 = np.conj(Y1)\n",
    "        Z1 = np.fft.rfftn(z1)\n",
    "        Z1 = np.conj(Z1)\n",
    "\n",
    "    x_angle = []\n",
    "    y_angle = []\n",
    "    z_angle = []\n",
    "\n",
    "    i = 0\n",
    "    while i < 360:\n",
    "        x_angle.append(i)\n",
    "        y_angle.append(i)\n",
    "        i += ang\n",
    "\n",
    "    i = 0\n",
    "    while i <= 180:\n",
    "        z_angle.append(i)\n",
    "        i += ang\n",
    "\n",
    "    angle_comb = np.array(np.meshgrid(x_angle, y_angle, z_angle)).T.reshape(-1, 3)\n",
    "\n",
    "    rot_vec_dict = {}\n",
    "    rot_data_dict = {}\n",
    "    #rot_vec_dict_p1 = {}\n",
    "    rot_data_dict_p1 = {}\n",
    "    #rot_vec_dict_p2 = {}\n",
    "    rot_data_dict_p2 = {}\n",
    "    #rot_vec_dict_p3 = {}\n",
    "    rot_data_dict_p3 = {}\n",
    "    #rot_vec_dict_p4 = {}\n",
    "    rot_data_dict_p4 = {}\n",
    "    \n",
    "\n",
    "    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count() + 4) as executor:\n",
    "        trans_vec = {executor.submit(rot_mrc_prob, mrc_search.data, mrc_search.vec, mrc_search_p1.data, mrc_search_p2.data, mrc_search_p3.data, mrc_search_p4.data, angle,): angle for angle in angle_comb}\n",
    "        #trans_vec_p1 = {executor.submit(rot_mrc, mrc_search_p1.data, mrc_search_p1.vec, angle,): angle for angle in angle_comb}\n",
    "        #trans_vec_p2 = {executor.submit(rot_mrc, mrc_search_p2.data, mrc_search_p2.vec, angle,): angle for angle in angle_comb}\n",
    "        #trans_vec_p3 = {executor.submit(rot_mrc, mrc_search_p3.data, mrc_search_p3.vec, angle,): angle for angle in angle_comb}\n",
    "        #trans_vec_p4 = {executor.submit(rot_mrc, mrc_search_p4.data, mrc_search_p4.vec, angle,): angle for angle in angle_comb}\n",
    "        for future in concurrent.futures.as_completed(trans_vec):\n",
    "            angle = trans_vec[future]\n",
    "            rot_vec_dict[tuple(angle)] = future.result()[0]\n",
    "            rot_data_dict[tuple(angle)] = future.result()[1]\n",
    "            rot_data_dict_p1[tuple(angle)] = future.result()[2]\n",
    "            rot_data_dict_p2[tuple(angle)] = future.result()[3]\n",
    "            rot_data_dict_p3[tuple(angle)] = future.result()[4]\n",
    "            rot_data_dict_p4[tuple(angle)] = future.result()[5]\n",
    "\n",
    "    #     for angle in tqdm(angle_comb, desc=\"Rotation\"):\n",
    "    #         rot_result = rot_mrc(\n",
    "    #             mrc_search.data,\n",
    "    #             mrc_search.vec,\n",
    "    #             angle,\n",
    "    #         )\n",
    "    #         mrc_angle_dict[tuple(angle)] = rot_result\n",
    "\n",
    "    # fftw plans\n",
    "    a = pyfftw.empty_aligned((x1.shape), dtype=\"float32\")\n",
    "    b = pyfftw.empty_aligned((a.shape[0], a.shape[1], a.shape[2] // 2 + 1), dtype=\"complex64\")\n",
    "    c = pyfftw.empty_aligned((x1.shape), dtype=\"float32\")\n",
    "\n",
    "    fft_object = pyfftw.FFTW(a, b, axes=(0, 1, 2))\n",
    "    ifft_object = pyfftw.FFTW(b, c, direction=\"FFTW_BACKWARD\", axes=(0, 1, 2), normalise_idft=False)\n",
    "\n",
    "    angle_score = []\n",
    "\n",
    "    for angle in tqdm(angle_comb, desc=\"FFT\"):\n",
    "        rot_mrc_vec = rot_vec_dict[tuple(angle)]\n",
    "        rot_mrc_data = rot_data_dict[tuple(angle)]\n",
    "        #rot_mrc_vec_p1 = rot_vec_dict_p1[tuple(angle)]\n",
    "        rot_mrc_data_p1 = rot_data_dict_p1[tuple(angle)]\n",
    "        #rot_mrc_vec_p2 = rot_vec_dict_p2[tuple(angle)]\n",
    "        rot_mrc_data_p2 = rot_data_dict_p2[tuple(angle)]\n",
    "        #rot_mrc_vec_p3 = rot_vec_dict_p3[tuple(angle)]\n",
    "        rot_mrc_data_p3 = rot_data_dict_p3[tuple(angle)]\n",
    "        #rot_mrc_vec_p4 = rot_vec_dict_p4[tuple(angle)]\n",
    "        rot_mrc_data_p4 = rot_data_dict_p4[tuple(angle)]\n",
    "        \n",
    "      \n",
    "        if mode == \"VecProduct\":\n",
    "            \n",
    "            x2 = copy.deepcopy(rot_mrc_vec[..., 0])\n",
    "            y2 = copy.deepcopy(rot_mrc_vec[..., 1])\n",
    "            z2 = copy.deepcopy(rot_mrc_vec[..., 2])\n",
    "            p21 = copy.deepcopy(rot_mrc_data_p1)\n",
    "            p22 = copy.deepcopy(rot_mrc_data_p2)\n",
    "            p23 = copy.deepcopy(rot_mrc_data_p3)\n",
    "            p24 = copy.deepcopy(rot_mrc_data_p4)\n",
    "            \n",
    "            #p24 = 1 - (p21 + p22 + p23)\n",
    "                    \n",
    "            target_list = [X1, Y1, Z1, P1, P2, P3, P4]\n",
    "            query_list = [x2, y2, z2, p21, p22, p23, p24]\n",
    "            \n",
    "            fft_result_list = fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object)\n",
    "                    \n",
    "            best, trans, best_prob = find_best_trans_list(fft_result_list)\n",
    "            \n",
    "        else:\n",
    "            best, trans = fft_search_score_trans_1d(\n",
    "                X1, rot_mrc_data, a, b, fft_object, ifft_object, mode, mrc_target.ave\n",
    "            )\n",
    "            if mode == \"CC\":\n",
    "                rstd2 = 1.0 / mrc_target.std ** 2\n",
    "                best = best * rstd2\n",
    "            if mode == \"PCC\":\n",
    "                rstd3 = 1.0 / mrc_target.std_norm_ave ** 2\n",
    "                best = best * rstd3\n",
    "\n",
    "        angle_score.append([tuple(angle), best * rd3, trans])\n",
    "\n",
    "    # calculate the ave and std\n",
    "    score_arr = np.array([row[1] for row in angle_score])\n",
    "    ave = np.mean(score_arr)\n",
    "    std = np.std(score_arr)\n",
    "    print(\"Std= \" + str(std) + \" Ave= \" + str(ave))\n",
    "\n",
    "    # sort the list and get topN\n",
    "    sorted_topN = sorted(angle_score, key=lambda x: x[1], reverse=True)[:TopN]\n",
    "\n",
    "    for x in sorted_topN:\n",
    "        print(x)\n",
    "\n",
    "    refined_score = []\n",
    "    if ang > 5.0:\n",
    "        for t_mrc in sorted_topN:\n",
    "            ang_list = np.array(\n",
    "                np.meshgrid(\n",
    "                    [t_mrc[0][0] - 5, t_mrc[0][0], t_mrc[0][0] + 5],\n",
    "                    [t_mrc[0][1] - 5, t_mrc[0][1], t_mrc[0][1] + 5],\n",
    "                    [t_mrc[0][2] - 5, t_mrc[0][2], t_mrc[0][2] + 5],\n",
    "                )\n",
    "            ).T.reshape(-1, 3)\n",
    "            # print(ang_list)\n",
    "            for ang in ang_list:\n",
    "                rotated = rot_mrc_prob(mrc_search.data, mrc_search.vec, mrc_search_p1.data, mrc_search_p2.data, mrc_search_p3.data, mrc_search_p4.data, ang)\n",
    "                rotated_vec=rotated[0]\n",
    "                rotated_data=rotated[1]\n",
    "                rotated_data_p1=rotated[2]\n",
    "                rotated_data_p2=rotated[3]\n",
    "                rotated_data_p3=rotated[4]\n",
    "                rotated_data_p4=rotated[5]\n",
    "                if mode == \"VecProduct\":\n",
    "                    \n",
    "                    x2 = copy.deepcopy(rotated_vec[..., 0])\n",
    "                    y2 = copy.deepcopy(rotated_vec[..., 1])\n",
    "                    z2 = copy.deepcopy(rotated_vec[..., 2])\n",
    "                    p21 = copy.deepcopy(rotated[2])\n",
    "                    p22 = copy.deepcopy(rotated[3])\n",
    "                    p23 = copy.deepcopy(rotated[4])\n",
    "                    p24 = copy.deepcopy(rotated[5])\n",
    "                    \n",
    "                    target_list = [X1, Y1, Z1, P1, P2, P3, P4]\n",
    "                    query_list = [x2, y2, z2, p21, p22, p23, p24]\n",
    "                    \n",
    "                    fft_result_list = fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object)\n",
    "                    \n",
    "                    best, trans, best_prob = find_best_trans_list(fft_result_list)\n",
    "                    \n",
    "                    # best, trans = fft_search_score_trans(X1, Y1, Z1, rotated_vec, a, b, c, fft_object, ifft_object)\n",
    "                else:\n",
    "                    best, trans = fft_search_score_trans_1d(\n",
    "                        X1, rotated_data, a, b, fft_object, ifft_object, mode, ave=mrc_target.ave\n",
    "                    )\n",
    "                    if mode == \"CC\":\n",
    "                        rstd2 = 1.0 / mrc_target.std ** 2\n",
    "                        best = best * rstd2\n",
    "                    if mode == \"PCC\":\n",
    "                        rstd3 = 1.0 / mrc_target.std_norm_ave ** 2\n",
    "                        best = best * rstd3\n",
    "\n",
    "                refined_score.append([tuple(ang), best * rd3, trans, rotated_vec, rotated_data, best_prob * rd3])\n",
    "\n",
    "        refined_list = sorted(refined_score, key=lambda x: x[1], reverse=True)[:TopN]\n",
    "    else:\n",
    "        refined_list = sorted_topN\n",
    "\n",
    "    for i, t_mrc in enumerate(refined_list):\n",
    "        print(\"R=\" + str(t_mrc[0]) + \" T=\" + str(t_mrc[2]))\n",
    "        sco = get_score(\n",
    "            mrc_target.data,\n",
    "            t_mrc[4],\n",
    "            mrc_target.vec,\n",
    "            t_mrc[3],\n",
    "            t_mrc[2],\n",
    "            mrc_target.ave,\n",
    "            mrc_search.ave,\n",
    "            mrc_target.std,\n",
    "            mrc_search.std,\n",
    "            mrc_target.std_norm_ave,\n",
    "            mrc_search.std_norm_ave,\n",
    "            sco_prob_added=t_mrc[1],\n",
    "            sco_prob_only=t_mrc[5],\n",
    "        )\n",
    "        \n",
    "        # Write result to PDB files\n",
    "        show_vec(mrc_target.orig, t_mrc[3], t_mrc[4], sco, mrc_search.xwidth, t_mrc[2], \"model_\" + str(i + 1) + \".pdb\")\n",
    "\n",
    "    #     num_jobs = math.ceil(360 / ang) * math.ceil(360 / ang) * (180 // ang + 1)\n",
    "\n",
    "    return refined_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def show_vec(origin, sampled_mrc_vec, sampled_mrc_data, sampled_mrc_score, sample_width, trans, name):\n",
    "\n",
    "    dim = sampled_mrc_data.shape[0]\n",
    "\n",
    "    origin = np.array([origin[\"x\"], origin[\"y\"], origin[\"z\"]])\n",
    "    trans = np.array(trans)\n",
    "\n",
    "    if trans[0] > 0.5 * dim:\n",
    "        trans[0] -= dim\n",
    "    if trans[1] > 0.5 * dim:\n",
    "        trans[1] -= dim\n",
    "    if trans[2] > 0.5 * dim:\n",
    "        trans[2] -= dim\n",
    "\n",
    "    add = origin - trans * sample_width\n",
    "\n",
    "    natm = 1\n",
    "    nres = 1\n",
    "\n",
    "    pdb_file = open(name, \"w\")\n",
    "    #     pdb_file.write(\"MODEL\\n\")\n",
    "    for x in range(dim):\n",
    "        for y in range(dim):\n",
    "            for z in range(dim):\n",
    "\n",
    "                if sampled_mrc_data[x][y][z] != 0.0:\n",
    "                    tmp = np.array([x, y, z])\n",
    "                    tmp = tmp * sample_width + add\n",
    "                    atom_header = \"ATOM{:>7d}  CA  ALA{:>6d}    \".format(natm, nres)\n",
    "                    atom_content = \"{:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}\".format(\n",
    "                        tmp[0], tmp[1], tmp[2], 1.0, sampled_mrc_score[x][y][z]\n",
    "                    )\n",
    "                    pdb_file.write(atom_header + atom_content + \"\\n\")\n",
    "                    natm += 1\n",
    "\n",
    "                    tmp = np.array([x, y, z])\n",
    "                    tmp = (tmp + sampled_mrc_vec[x][y][z]) * sample_width + add\n",
    "                    atom_header = \"ATOM{:>7d}  CB  ALA{:>6d}    \".format(natm, nres)\n",
    "                    atom_content = \"{:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}\".format(\n",
    "                        tmp[0], tmp[1], tmp[2], 1.0, sampled_mrc_score[x][y][z]\n",
    "                    )\n",
    "                    pdb_file.write(atom_header + atom_content + \"\\n\")\n",
    "                    natm += 1\n",
    "\n",
    "                    nres += 1\n",
    "\n",
    "\n",
    "#     pdb_file.write(\"ENDMDL\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "236\n",
      "215\n",
      "204\n",
      "(236, 215, 204)\n",
      "(236, 215, 204)\n",
      "#dmax=158.78680675673277\n",
      "Nvox= 45, 45, 45\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -39.5, -50.0, -55.5\n",
      "#dmax=147.56439272399015\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -29.0, -39.5, -45.0\n",
      "#dmax=185.97378847568817\n",
      "Nvox= 54, 54, 54\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -71.0, -81.5, -87.0\n",
      "#dmax=185.97378847568817\n",
      "Nvox= 54, 54, 54\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -71.0, -81.5, -87.0\n",
      "#dmax=185.97378847568817\n",
      "Nvox= 54, 54, 54\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -71.0, -81.5, -87.0\n",
      "#dmax=184.87090090114236\n",
      "Nvox= 52, 52, 52\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -64.0, -74.5, -80.0\n",
      "#dmax=155.3391450987162\n",
      "Nvox= 44, 44, 44\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -36.0, -46.5, -52.0\n",
      "#dmax=155.3391450987162\n",
      "Nvox= 44, 44, 44\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -36.0, -46.5, -52.0\n",
      "#dmax=155.3391450987162\n",
      "Nvox= 44, 44, 44\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -36.0, -46.5, -52.0\n",
      "#dmax=155.3391450987162\n",
      "Nvox= 44, 44, 44\n",
      "cent= 118.0, 107.5, 102.0\n",
      "ori= -36.0, -46.5, -52.0\n",
      "54\n"
     ]
    }
   ],
   "source": [
    "mrc1 = mrc_obj(\"./emd_8097.mrc\")\n",
    "mrc2 = mrc_obj(\"./chainA_simulated_resample.mrc\")\n",
    "print(mrc2.xdim)\n",
    "print(mrc2.ydim)\n",
    "print(mrc2.zdim)\n",
    "print(mrc2.data.shape)\n",
    "\n",
    "#probability MRCs\n",
    "\n",
    "mrc1_p1 = mrc_obj(\"./emd_8097.mrc\")\n",
    "mrc1_p2 = mrc_obj(\"./emd_8097.mrc\")\n",
    "mrc1_p3 = mrc_obj(\"./emd_8097.mrc\")\n",
    "mrc1_p4 = mrc_obj(\"./emd_8097.mrc\")\n",
    "\n",
    "mrc2_p1 = mrc_obj(\"./chainA_simulated_resample.mrc\")\n",
    "mrc2_p2 = mrc_obj(\"./chainA_simulated_resample.mrc\")\n",
    "mrc2_p3 = mrc_obj(\"./chainA_simulated_resample.mrc\")\n",
    "mrc2_p4 = mrc_obj(\"./chainA_simulated_resample.mrc\")\n",
    "#\n",
    "#\n",
    "''''mrc2_p1.xdim=prob_maps_chain.shape[0]\n",
    "mrc2_p1.ydim=prob_maps_chain.shape[1]\n",
    "mrc2_p1.zdim=prob_maps_chain.shape[2]\n",
    "\n",
    "mrc2_p2.xdim=prob_maps_chain.shape[0]\n",
    "mrc2_p2.ydim=prob_maps_chain.shape[1]\n",
    "mrc2_p2.zdim=prob_maps_chain.shape[2]\n",
    "\n",
    "mrc2_p3.xdim=prob_maps_chain.shape[0]\n",
    "mrc2_p3.ydim=prob_maps_chain.shape[1]\n",
    "mrc2_p3.zdim=prob_maps_chain.shape[2]\n",
    "\n",
    "mrc2_p4.xdim=prob_maps_chain.shape[0]\n",
    "mrc2_p4.ydim=prob_maps_chain.shape[1]\n",
    "mrc2_p4.zdim=prob_maps_chain.shape[2]'''\n",
    "\n",
    "mrc2_p1.data=np.swapaxes(prob_maps_chain[:,:,:,0], 0, 2)\n",
    "print(mrc2_p1.data.shape)\n",
    "mrc2_p2.data=np.swapaxes(prob_maps_chain[:,:,:,1], 0, 2)\n",
    "mrc2_p3.data=np.swapaxes(prob_maps_chain[:,:,:,2], 0, 2)\n",
    "mrc2_p4.data=np.swapaxes(prob_maps_chain[:,:,:,3], 0, 2)\n",
    "\n",
    "''''mrc1_p1.xdim=prob_maps.shape[0]\n",
    "mrc1_p1.ydim=prob_maps.shape[1]\n",
    "mrc1_p1.zdim=prob_maps.shape[2]\n",
    "\n",
    "mrc1_p2.xdim=prob_maps.shape[0]\n",
    "mrc1_p2.ydim=prob_maps.shape[1]\n",
    "mrc1_p2.zdim=prob_maps.shape[2]\n",
    "\n",
    "mrc1_p3.xdim=prob_maps.shape[0]\n",
    "mrc1_p3.ydim=prob_maps.shape[1]\n",
    "mrc1_p3.zdim=prob_maps.shape[2]\n",
    "\n",
    "mrc1_p4.xdim=prob_maps.shape[0]\n",
    "mrc1_p4.ydim=prob_maps.shape[1]\n",
    "mrc1_p4.zdim=prob_maps.shape[2]'''\n",
    "\n",
    "mrc1_p1.data=np.swapaxes(prob_maps[:,:,:,0], 0, 2)\n",
    "mrc1_p2.data=np.swapaxes(prob_maps[:,:,:,1], 0, 2)\n",
    "mrc1_p3.data=np.swapaxes(prob_maps[:,:,:,2], 0, 2)\n",
    "mrc1_p4.data=np.swapaxes(prob_maps[:,:,:,3], 0, 2)\n",
    "\n",
    "\n",
    "mrc1, mrc_N1 = mrc_set_vox_size(mrc1)\n",
    "mrc2, mrc_N2 = mrc_set_vox_size(mrc2)\n",
    "mrc1_p1, mrc_N1_p1 = mrc_set_vox_size(mrc1_p1)\n",
    "mrc1_p2, mrc_N1_p2 = mrc_set_vox_size(mrc1_p2)\n",
    "mrc1_p3, mrc_N1_p3 = mrc_set_vox_size(mrc1_p3)\n",
    "mrc1_p4, mrc_N1_p4 = mrc_set_vox_size(mrc1_p4)\n",
    "mrc2_p1, mrc_N2_p1 = mrc_set_vox_size(mrc2_p1)\n",
    "mrc2_p2, mrc_N2_p2 = mrc_set_vox_size(mrc2_p2)\n",
    "mrc2_p3, mrc_N2_p3 = mrc_set_vox_size(mrc2_p3)\n",
    "mrc2_p4, mrc_N2_p4 = mrc_set_vox_size(mrc2_p4)\n",
    "\n",
    "mrc_list = [mrc_N1, mrc_N2, mrc_N1_p1, mrc_N1_p2, mrc_N1_p3, mrc_N1_p4, mrc_N2_p1, mrc_N2_p2, mrc_N2_p3, mrc_N2_p4]\n",
    "max_dim = np.max((mrc_N1.xdim, mrc_N2.xdim, mrc_N2_p1.xdim, mrc_N2_p2.xdim, mrc_N2_p3.xdim, mrc_N2_p4.xdim, mrc_N1_p1.xdim, mrc_N1_p2.xdim, mrc_N1_p3.xdim, mrc_N1_p4.xdim))\n",
    "\n",
    "print(max_dim)\n",
    "\n",
    "for mrc in mrc_list:\n",
    "    if mrc.xdim != max_dim:\n",
    "        mrc.xdim = mrc.ydim = mrc.zdim = max_dim\n",
    "        mrc.orig[\"x\"] = mrc.cent[0] - 0.5 * 7 * mrc.xdim\n",
    "        mrc.orig[\"y\"] = mrc.cent[1] - 0.5 * 7 * mrc.xdim\n",
    "        mrc.orig[\"z\"] = mrc.cent[2] - 0.5 * 7 * mrc.xdim\n",
    "\n",
    "for mrc in mrc_list:\n",
    "    mrc.dens = np.zeros((max_dim ** 3, 1), dtype=\"float32\")\n",
    "    mrc.data = np.zeros((max_dim, max_dim, max_dim), dtype=\"float32\")\n",
    "    mrc.vec = np.zeros((max_dim, max_dim, max_dim, 3), dtype=\"float32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(54, 54, 54) (54, 54, 54) (54, 54, 54) (54, 54, 54)\n",
      "(54, 54, 54) (54, 54, 54) (54, 54, 54) (54, 54, 54)\n",
      "(54, 54, 54) (54, 54, 54)\n"
     ]
    }
   ],
   "source": [
    "print(mrc_N1_p1.data.shape, mrc_N1_p2.data.shape, mrc_N1_p3.data.shape, mrc_N1_p4.data.shape)\n",
    "print(mrc_N2_p1.data.shape, mrc_N2_p2.data.shape, mrc_N2_p3.data.shape, mrc_N2_p4.data.shape)\n",
    "print(mrc_N1.data.shape, mrc_N2.data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|████████████▌                                                                                                                          | 5/54 [00:00<00:01, 48.25it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:20<00:00,  2.62it/s]\n",
      " 17%|██████████████████████▌                                                                                                                | 9/54 [00:00<00:00, 82.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "39456.039104347255\n",
      "2128\n",
      "#MAP AVE=18.54137175956168 STD=935.837646484375 STD_norm=379.76708984375\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [02:47<00:00,  3.10s/it]\n",
      " 17%|██████████████████████▌                                                                                                                | 9/54 [00:00<00:00, 82.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "4566983.062774386\n",
      "20738\n",
      "#MAP AVE=220.22292712770692 STD=34293.390625 STD_norm=13049.216796875\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [02:37<00:00,  2.92s/it]\n",
      " 15%|████████████████████                                                                                                                   | 8/54 [00:00<00:00, 71.52it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "6356082.24759428\n",
      "20720\n",
      "#MAP AVE=306.76072623524516 STD=45364.58203125 STD_norm=10399.474609375\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [02:55<00:00,  3.25s/it]\n",
      "  9%|████████████▌                                                                                                                          | 5/54 [00:00<00:01, 49.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "4829020.5048538055\n",
      "20734\n",
      "#MAP AVE=232.9034679682553 STD=39437.4453125 STD_norm=20751.294921875\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [01:17<00:00,  1.44s/it]\n",
      "  9%|████████████▌                                                                                                                          | 5/54 [00:00<00:01, 47.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "142667.1015578666\n",
      "9105\n",
      "#MAP AVE=15.669094075548228 STD=2178.646240234375 STD_norm=1584.62451171875\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:05<00:00, 10.61it/s]\n",
      " 15%|████████████████████                                                                                                                   | 8/54 [00:00<00:00, 71.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "95784.81157440397\n",
      "572\n",
      "#MAP AVE=167.45596429091603 STD=4328.62158203125 STD_norm=1642.323974609375\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:16<00:00,  3.31it/s]\n",
      " 13%|█████████████████▌                                                                                                                     | 7/54 [00:00<00:00, 62.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "697807.5983372995\n",
      "1658\n",
      "#MAP AVE=420.87309911779226 STD=17961.58984375 STD_norm=5378.68505859375\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:14<00:00,  3.78it/s]\n",
      " 13%|█████████████████▌                                                                                                                     | 7/54 [00:00<00:00, 61.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "259584.222122695\n",
      "1503\n",
      "#MAP AVE=172.7107266285396 STD=7194.404296875 STD_norm=2631.819580078125\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:14<00:00,  3.64it/s]\n",
      " 13%|█████████████████▌                                                                                                                     | 7/54 [00:00<00:00, 66.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "869663.5906798018\n",
      "1654\n",
      "#MAP AVE=525.79419025381 STD=24247.20703125 STD_norm=11430.7919921875\n",
      "157464\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [00:12<00:00,  4.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#End LDP\n",
      "279632.14645056607\n",
      "1355\n",
      "#MAP AVE=206.37058778639562 STD=10128.5078125 STD_norm=6699.16552734375\n",
      "Wall time: 11min 1s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "mrc_N1 = fastVEC(mrc1, mrc_N1)\n",
    "mrc_N1_p1 = fastVEC_prob(mrc1_p1, mrc_N1_p1)\n",
    "mrc_N1_p2 = fastVEC_prob(mrc1_p2, mrc_N1_p2)\n",
    "mrc_N1_p3 = fastVEC_prob(mrc1_p3, mrc_N1_p3)\n",
    "mrc_N1_p4 = fastVEC_prob(mrc1_p4, mrc_N1_p4)\n",
    "\n",
    "\n",
    "\n",
    "mrc_N2 = fastVEC(mrc2, mrc_N2)\n",
    "mrc_N2_p1 = fastVEC_prob(mrc2_p1, mrc_N2_p1)\n",
    "mrc_N2_p2 = fastVEC_prob(mrc2_p2, mrc_N2_p2)\n",
    "mrc_N2_p3 = fastVEC_prob(mrc2_p3, mrc_N2_p3)\n",
    "mrc_N2_p4 = fastVEC_prob(mrc2_p4, mrc_N2_p4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(54, 54, 54)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mrc_N2_p1.data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "FFT: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1008/1008 [01:57<00:00,  8.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Std= 1.4335097354898193e-05 Ave= 0.00014480944914044847\n",
      "[(90, 30, 150), 0.0002678336393649649, (43, 11, 50)]\n",
      "[(270, 150, 150), 0.00024159931042739083, (12, 43, 50)]\n",
      "[(270, 150, 0), 0.00021696437127717485, (39, 7, 50)]\n",
      "[(90, 30, 180), 0.00021696437127717485, (39, 7, 50)]\n",
      "[(90, 30, 0), 0.00021175155248158903, (15, 46, 50)]\n",
      "[(270, 150, 180), 0.00021175155248158903, (15, 46, 50)]\n",
      "[(300, 180, 0), 0.0002009928096257071, (47, 5, 53)]\n",
      "[(120, 0, 180), 0.0002009928096257071, (47, 5, 53)]\n",
      "[(300, 150, 0), 0.000199705497076167, (40, 6, 48)]\n",
      "[(120, 30, 180), 0.000199705497076167, (40, 6, 48)]\n",
      "R=(275, 155, 150) T=(12, 43, 50)\n",
      "39454.94\n",
      "7186.823\n",
      "Overlap= 0.12339011217282925 297/2407 CC= 0.0451234 PCC= -0.3716458 Scoreplusprob= 0.0002805142663561195 Scoreprobonly= 3.935801015663079e-05\n",
      "Score=  153.34596\n",
      "R=(90, 25, 145) T=(43, 12, 51)\n",
      "39455.6\n",
      "7103.6855\n",
      "Overlap= 0.11956971452213488 289/2417 CC= 0.045144115 PCC= -0.38654342 Scoreplusprob= 0.0002731933394129099 Scoreprobonly= 3.907201808926923e-05\n",
      "Score=  151.08632\n",
      "R=(275, 155, 155) T=(13, 44, 50)\n",
      "39454.13\n",
      "6997.6553\n",
      "Overlap= 0.11960132890365449 288/2408 CC= 0.04409015 PCC= -0.366747 Scoreplusprob= 0.000268336349717547 Scoreprobonly= 3.932505998949158e-05\n",
      "Score=  145.00034\n",
      "R=(90, 30, 150) T=(43, 11, 50)\n",
      "39456.04\n",
      "7369.4585\n",
      "Overlap= 0.12406639004149378 299/2410 CC= 0.04699857 PCC= -0.40397516 Scoreplusprob= 0.0002678336393649649 Scoreprobonly= 3.9785332731349566e-05\n",
      "Score=  147.0466\n",
      "R=(90, 25, 150) T=(43, 11, 51)\n",
      "39456.04\n",
      "7682.7334\n",
      "Overlap= 0.13127090301003344 314/2392 CC= 0.0487159 PCC= -0.41157275 Scoreplusprob= 0.00026698777002513767 Scoreprobonly= 3.9112987008598923e-05\n",
      "Score=  145.01299\n",
      "R=(85, 25, 145) T=(43, 12, 51)\n",
      "39455.6\n",
      "6904.495\n",
      "Overlap= 0.11616370400992146 281/2419 CC= 0.04426362 PCC= -0.37268525 Scoreplusprob= 0.0002654209155374527 Scoreprobonly= 3.905685574326556e-05\n",
      "Score=  146.3153\n",
      "R=(85, 25, 150) T=(42, 11, 51)\n",
      "39455.81\n",
      "6827.0796\n",
      "Overlap= 0.1134575569358178 274/2415 CC= 0.043926943 PCC= -0.38663238 Scoreplusprob= 0.0002644524392343746 Scoreprobonly= 3.904223242214478e-05\n",
      "Score=  141.18445\n",
      "R=(90, 30, 155) T=(42, 10, 50)\n",
      "39455.812\n",
      "7245.5576\n",
      "Overlap= 0.12109862671660425 291/2403 CC= 0.046495512 PCC= -0.41011828 Scoreplusprob= 0.0002637546139871033 Scoreprobonly= 3.939751946925511e-05\n",
      "Score=  142.45071\n",
      "R=(275, 155, 145) T=(11, 42, 50)\n",
      "39455.51\n",
      "6919.4463\n",
      "Overlap= 0.12020033388981637 288/2396 CC= 0.043487035 PCC= -0.34734103 Scoreplusprob= 0.00026355472660243967 Scoreprobonly= 3.924962678818309e-05\n",
      "Score=  139.59857\n",
      "R=(95, 25, 155) T=(42, 10, 50)\n",
      "39455.812\n",
      "7175.5703\n",
      "Overlap= 0.120997920997921 291/2405 CC= 0.045616183 PCC= -0.39495838 Scoreplusprob= 0.0002634105344180356 Scoreprobonly= 3.8992209096074326e-05\n",
      "Score=  141.0889\n",
      "Wall time: 4min 47s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "score_list = search_map_fft(mrc_N1_p1, mrc_N1_p2, mrc_N1_p3, mrc_N1_p4, mrc_N1,\n",
    "                            mrc_N2, mrc_N2_p1, mrc_N2_p2, mrc_N2_p3, mrc_N2_p4, ang=30, mode=\"VecProduct\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'line_profiler'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-27-df8a33df4eaf>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mget_ipython\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'load_ext'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'line_profiler'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\IPython\\core\\interactiveshell.py\u001b[0m in \u001b[0;36mrun_line_magic\u001b[1;34m(self, magic_name, line, _stack_depth)\u001b[0m\n\u001b[0;32m   2325\u001b[0m                 \u001b[0mkwargs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'local_ns'\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mget_local_scope\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mstack_depth\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2326\u001b[0m             \u001b[1;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbuiltin_trap\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2327\u001b[1;33m                 \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2328\u001b[0m             \u001b[1;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2329\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<decorator-gen-58>\u001b[0m in \u001b[0;36mload_ext\u001b[1;34m(self, module_str)\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\IPython\\core\\magic.py\u001b[0m in \u001b[0;36m<lambda>\u001b[1;34m(f, *a, **k)\u001b[0m\n\u001b[0;32m    185\u001b[0m     \u001b[1;31m# but it's overkill for just that one bit of state.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    186\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mmagic_deco\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0marg\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 187\u001b[1;33m         \u001b[0mcall\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0ma\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mk\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0ma\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mk\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    188\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    189\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0marg\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\IPython\\core\\magics\\extension.py\u001b[0m in \u001b[0;36mload_ext\u001b[1;34m(self, module_str)\u001b[0m\n\u001b[0;32m     31\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mmodule_str\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     32\u001b[0m             \u001b[1;32mraise\u001b[0m \u001b[0mUsageError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'Missing module name.'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 33\u001b[1;33m         \u001b[0mres\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshell\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mextension_manager\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload_extension\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodule_str\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     34\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     35\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mres\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;34m'already loaded'\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\IPython\\core\\extensions.py\u001b[0m in \u001b[0;36mload_extension\u001b[1;34m(self, module_str)\u001b[0m\n\u001b[0;32m     78\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mmodule_str\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[1;32min\u001b[0m \u001b[0msys\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodules\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     79\u001b[0m                 \u001b[1;32mwith\u001b[0m \u001b[0mprepended_to_syspath\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mipython_extension_dir\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 80\u001b[1;33m                     \u001b[0mmod\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mimport_module\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodule_str\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     81\u001b[0m                     \u001b[1;32mif\u001b[0m \u001b[0mmod\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__file__\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstartswith\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mipython_extension_dir\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     82\u001b[0m                         print((\"Loading extensions from {dir} is deprecated. \"\n",
      "\u001b[1;32m~\\anaconda3\\lib\\importlib\\__init__.py\u001b[0m in \u001b[0;36mimport_module\u001b[1;34m(name, package)\u001b[0m\n\u001b[0;32m    125\u001b[0m                 \u001b[1;32mbreak\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    126\u001b[0m             \u001b[0mlevel\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 127\u001b[1;33m     \u001b[1;32mreturn\u001b[0m \u001b[0m_bootstrap\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_gcd_import\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mname\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mlevel\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mpackage\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlevel\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    128\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    129\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\importlib\\_bootstrap.py\u001b[0m in \u001b[0;36m_gcd_import\u001b[1;34m(name, package, level)\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\importlib\\_bootstrap.py\u001b[0m in \u001b[0;36m_find_and_load\u001b[1;34m(name, import_)\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\importlib\\_bootstrap.py\u001b[0m in \u001b[0;36m_find_and_load_unlocked\u001b[1;34m(name, import_)\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'line_profiler'"
     ]
    }
   ],
   "source": [
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%lprun -f rot_mrc rot_mrc(mrc_N2.data, mrc_N2.vec, [0,0,30])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "11dcb162716791435bf48ce72b845ec24ce37a51f37dce84df4d98a21bd70cb4"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}