VESPER_prob_correct.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "source": [
    "import concurrent.futures\r\n",
    "import copy\r\n",
    "import math\r\n",
    "import multiprocessing\r\n",
    "import os\r\n",
    "\r\n",
    "import mrcfile\r\n",
    "import numba\r\n",
    "import numpy as np\r\n",
    "import pyfftw\r\n",
    "import scipy.fft\r\n",
    "from numba.typed import List\r\n",
    "from scipy.ndimage import convolve\r\n",
    "from scipy.spatial.transform import Rotation as R\r\n",
    "from tqdm import tqdm"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "pyfftw.config.PLANNER_EFFORT = \"FFTW_MEASURE\"\r\n",
    "pyfftw.config.NUM_THREADS = multiprocessing.cpu_count()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "class mrc_obj:\r\n",
    "    def __init__(self, path):\r\n",
    "        mrc = mrcfile.open(path)\r\n",
    "        data = mrc.data\r\n",
    "        header = mrc.header\r\n",
    "        self.xdim = int(header.nx)\r\n",
    "        self.ydim = int(header.ny)\r\n",
    "        self.zdim = int(header.nz)\r\n",
    "        self.xwidth = mrc.voxel_size.x\r\n",
    "        self.ywidth = mrc.voxel_size.y\r\n",
    "        self.zwidth = mrc.voxel_size.z\r\n",
    "        self.cent = [\r\n",
    "            self.xdim * 0.5,\r\n",
    "            self.ydim * 0.5,\r\n",
    "            self.zdim * 0.5,\r\n",
    "        ]\r\n",
    "        self.orig = {\"x\": header.origin.x, \"y\": header.origin.y, \"z\": header.origin.z}\r\n",
    "        self.data = np.swapaxes(copy.deepcopy(data), 0, 2)\r\n",
    "        self.dens = data.flatten()\r\n",
    "        self.vec = np.zeros((self.xdim, self.ydim, self.zdim, 3), dtype=\"float32\")\r\n",
    "        self.dsum = None\r\n",
    "        self.Nact = None\r\n",
    "        self.ave = None\r\n",
    "        self.std_norm_ave = None\r\n",
    "        self.std = None"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "def mrc_set_vox_size(mrc, th=0.01, voxel_size=7.0):\r\n",
    "\r\n",
    "    # set shape and size\r\n",
    "    size = mrc.xdim * mrc.ydim * mrc.zdim\r\n",
    "    shape = (mrc.xdim, mrc.ydim, mrc.zdim)\r\n",
    "\r\n",
    "    # if th < 0 add th to all value\r\n",
    "    if th < 0:\r\n",
    "        mrc.dens = mrc.dens - th\r\n",
    "        th = 0.0\r\n",
    "\r\n",
    "    # Trim all the values less than threshold\r\n",
    "    mrc.dens[mrc.dens < th] = 0.0\r\n",
    "    mrc.data[mrc.data < th] = 0.0\r\n",
    "\r\n",
    "    # calculate dmax distance for non-zero entries\r\n",
    "    non_zero_index_list = np.array(np.nonzero(mrc.data)).T\r\n",
    "    # non_zero_index_list[:, [2, 0]] = non_zero_index_list[:, [0, 2]]\r\n",
    "    cent_arr = np.array(mrc.cent)\r\n",
    "    d2_list = np.linalg.norm(non_zero_index_list - cent_arr, axis=1)\r\n",
    "    dmax = max(d2_list)\r\n",
    "\r\n",
    "    # dmax = math.sqrt(mrc.cent[0] ** 2 + mrc.cent[1] ** 2 + mrc.cent[2] ** 2)\r\n",
    "\r\n",
    "    print(\"#dmax=\" + str(dmax / mrc.xwidth))\r\n",
    "    dmax = dmax * mrc.xwidth\r\n",
    "\r\n",
    "    # set new center\r\n",
    "    new_cent = [\r\n",
    "        mrc.cent[0] * mrc.xwidth + mrc.orig[\"x\"],\r\n",
    "        mrc.cent[1] * mrc.xwidth + mrc.orig[\"y\"],\r\n",
    "        mrc.cent[2] * mrc.xwidth + mrc.orig[\"z\"],\r\n",
    "    ]\r\n",
    "\r\n",
    "    tmp_size = 2 * dmax / voxel_size\r\n",
    "\r\n",
    "    # find the minimum size of the map\r\n",
    "    #     b = y = 2**math.ceil(math.log2(tmp_size))\r\n",
    "    #     while 1:\r\n",
    "    #         while y < tmp_size:\r\n",
    "    #             y = y * 3\r\n",
    "    #             continue\r\n",
    "    #         if y < b:\r\n",
    "    #             b = y\r\n",
    "    #         if y % 2 != 0:\r\n",
    "    #             break\r\n",
    "    #         y = y / 2\r\n",
    "\r\n",
    "    #     new_xdim = int(b)\r\n",
    "    new_xdim = pyfftw.next_fast_len(int(tmp_size))\r\n",
    "\r\n",
    "    # set new origins\r\n",
    "    new_orig = {\r\n",
    "        \"x\": new_cent[0] - 0.5 * new_xdim * voxel_size,\r\n",
    "        \"y\": new_cent[1] - 0.5 * new_xdim * voxel_size,\r\n",
    "        \"z\": new_cent[2] - 0.5 * new_xdim * voxel_size,\r\n",
    "    }\r\n",
    "\r\n",
    "    # create new mrc object\r\n",
    "    mrc_set = copy.deepcopy(mrc)\r\n",
    "    mrc_set.orig = new_orig\r\n",
    "    mrc_set.xdim = new_xdim\r\n",
    "    mrc_set.ydim = new_xdim\r\n",
    "    mrc_set.zdim = new_xdim\r\n",
    "    mrc_set.cent = new_cent\r\n",
    "    mrc_set.xwidth = mrc_set.ywidth = mrc_set.zwidth = voxel_size\r\n",
    "\r\n",
    "    print(\"Nvox= \" + str(mrc_set.xdim) + \", \" + str(mrc_set.ydim) + \", \" + str(mrc_set.zdim))\r\n",
    "    print(\"cent= \" + str(new_cent[0]) + \", \" + str(new_cent[1]) + \", \" + str(new_cent[2]))\r\n",
    "    print(\"ori= \" + str(new_orig[\"x\"]) + \", \" + str(new_orig[\"y\"]) + \", \" + str(new_orig[\"z\"]))\r\n",
    "\r\n",
    "    return mrc, mrc_set"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def calc(stp, endp, pos, mrc1_data, fsiv):\r\n",
    "    dtotal = 0.0\r\n",
    "    pos2 = np.zeros((3,))\r\n",
    "\r\n",
    "    for xp in range(stp[0], endp[0]):\r\n",
    "        rx = float(xp) - pos[0]\r\n",
    "        rx = rx ** 2\r\n",
    "        for yp in range(stp[1], endp[1]):\r\n",
    "            ry = float(yp) - pos[1]\r\n",
    "            ry = ry ** 2\r\n",
    "            for zp in range(stp[2], endp[2]):\r\n",
    "                rz = float(zp) - pos[2]\r\n",
    "                rz = rz ** 2\r\n",
    "                d2 = rx + ry + rz\r\n",
    "                v = mrc1_data[xp][yp][zp] * math.exp(-1.5 * d2 * fsiv)\r\n",
    "                dtotal += v\r\n",
    "                pos2[0] += v * float(xp)\r\n",
    "                pos2[1] += v * float(yp)\r\n",
    "                pos2[2] += v * float(zp)\r\n",
    "    return dtotal, pos2"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "def fastVEC(mrc1, mrc2, dreso=16.0):\r\n",
    "\r\n",
    "    xydim = mrc1.xdim * mrc1.ydim\r\n",
    "    Ndata = mrc2.xdim * mrc2.ydim * mrc2.zdim\r\n",
    "\r\n",
    "    print(len(mrc2.dens))\r\n",
    "\r\n",
    "    print(\"#Start VEC\")\r\n",
    "    gstep = mrc1.xwidth\r\n",
    "    fs = (dreso / gstep) * 0.5\r\n",
    "    fs = fs ** 2\r\n",
    "    fsiv = 1.0 / fs\r\n",
    "    fmaxd = (dreso / gstep) * 2.0\r\n",
    "    print(\"#maxd= {fmaxd}\".format(fmaxd=fmaxd))\r\n",
    "    print(\"#fsiv= \" + str(fsiv))\r\n",
    "\r\n",
    "    dsum = 0.0\r\n",
    "    Nact = 0\r\n",
    "\r\n",
    "    list_d = []\r\n",
    "\r\n",
    "    for x in tqdm(range(mrc2.xdim)):\r\n",
    "        for y in range(mrc2.ydim):\r\n",
    "            for z in range(mrc2.zdim):\r\n",
    "                stp = [0] * 3\r\n",
    "                endp = [0] * 3\r\n",
    "                ind2 = 0\r\n",
    "                ind = 0\r\n",
    "\r\n",
    "                pos = [0.0] * 3\r\n",
    "                pos2 = [0.0] * 3\r\n",
    "                ori = [0.0] * 3\r\n",
    "\r\n",
    "                tmpcd = [0.0] * 3\r\n",
    "\r\n",
    "                v, dtotal, rd = 0.0, 0.0, 0.0\r\n",
    "\r\n",
    "                pos[0] = (x * mrc2.xwidth + mrc2.orig[\"x\"] - mrc1.orig[\"x\"]) / mrc1.xwidth\r\n",
    "                pos[1] = (y * mrc2.xwidth + mrc2.orig[\"y\"] - mrc1.orig[\"y\"]) / mrc1.xwidth\r\n",
    "                pos[2] = (z * mrc2.xwidth + mrc2.orig[\"z\"] - mrc1.orig[\"z\"]) / mrc1.xwidth\r\n",
    "\r\n",
    "                ind = mrc2.xdim * mrc2.ydim * z + mrc2.xdim * y + x\r\n",
    "\r\n",
    "                # check density\r\n",
    "\r\n",
    "                if (\r\n",
    "                    pos[0] < 0\r\n",
    "                    or pos[1] < 0\r\n",
    "                    or pos[2] < 0\r\n",
    "                    or pos[0] >= mrc1.xdim\r\n",
    "                    or pos[1] >= mrc1.ydim\r\n",
    "                    or pos[2] >= mrc1.zdim\r\n",
    "                ):\r\n",
    "                    mrc2.dens[ind] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue\r\n",
    "\r\n",
    "                if mrc1.data[int(pos[0])][int(pos[1])][int(pos[2])] == 0:\r\n",
    "                    mrc2.dens[ind] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue\r\n",
    "\r\n",
    "                ori[0] = pos[0]\r\n",
    "                ori[1] = pos[1]\r\n",
    "                ori[2] = pos[2]\r\n",
    "\r\n",
    "                # Start Point\r\n",
    "                stp[0] = int(pos[0] - fmaxd)\r\n",
    "                stp[1] = int(pos[1] - fmaxd)\r\n",
    "                stp[2] = int(pos[2] - fmaxd)\r\n",
    "\r\n",
    "                # set start and end point\r\n",
    "                if stp[0] < 0:\r\n",
    "                    stp[0] = 0\r\n",
    "                if stp[1] < 0:\r\n",
    "                    stp[1] = 0\r\n",
    "                if stp[2] < 0:\r\n",
    "                    stp[2] = 0\r\n",
    "\r\n",
    "                endp[0] = int(pos[0] + fmaxd + 1)\r\n",
    "                endp[1] = int(pos[1] + fmaxd + 1)\r\n",
    "                endp[2] = int(pos[2] + fmaxd + 1)\r\n",
    "\r\n",
    "                if endp[0] >= mrc1.xdim:\r\n",
    "                    endp[0] = mrc1.xdim\r\n",
    "                if endp[1] >= mrc1.ydim:\r\n",
    "                    endp[1] = mrc1.ydim\r\n",
    "                if endp[2] >= mrc1.zdim:\r\n",
    "                    endp[2] = mrc1.zdim\r\n",
    "\r\n",
    "                # setup for numba acc\r\n",
    "                stp_t = List()\r\n",
    "                endp_t = List()\r\n",
    "                pos_t = List()\r\n",
    "                [stp_t.append(x) for x in stp]\r\n",
    "                [endp_t.append(x) for x in endp]\r\n",
    "                [pos_t.append(x) for x in pos]\r\n",
    "\r\n",
    "                # compute the total density\r\n",
    "                dtotal, pos2 = calc(stp_t, endp_t, pos_t, mrc1.data, fsiv)\r\n",
    "\r\n",
    "                mrc2.dens[ind] = dtotal\r\n",
    "                mrc2.data[x][y][z] = dtotal\r\n",
    "\r\n",
    "                if dtotal == 0:\r\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue\r\n",
    "\r\n",
    "                rd = 1.0 / dtotal\r\n",
    "\r\n",
    "                pos2[0] *= rd\r\n",
    "                pos2[1] *= rd\r\n",
    "                pos2[2] *= rd\r\n",
    "\r\n",
    "                tmpcd[0] = pos2[0] - pos[0]\r\n",
    "                tmpcd[1] = pos2[1] - pos[1]\r\n",
    "                tmpcd[2] = pos2[2] - pos[2]\r\n",
    "\r\n",
    "                dvec = math.sqrt(tmpcd[0] ** 2 + tmpcd[1] ** 2 + tmpcd[2] ** 2)\r\n",
    "\r\n",
    "                if dvec == 0:\r\n",
    "                    dvec = 1.0\r\n",
    "\r\n",
    "                rdvec = 1.0 / dvec\r\n",
    "\r\n",
    "                mrc2.vec[x][y][z][0] = tmpcd[0] * rdvec\r\n",
    "                mrc2.vec[x][y][z][1] = tmpcd[1] * rdvec\r\n",
    "                mrc2.vec[x][y][z][2] = tmpcd[2] * rdvec\r\n",
    "\r\n",
    "                dsum += dtotal\r\n",
    "                Nact += 1\r\n",
    "\r\n",
    "    print(\"#End LDP\")\r\n",
    "    print(dsum)\r\n",
    "    print(Nact)\r\n",
    "\r\n",
    "    mrc2.dsum = dsum\r\n",
    "    mrc2.Nact = Nact\r\n",
    "    mrc2.ave = dsum / float(Nact)\r\n",
    "    mrc2.std = np.linalg.norm(mrc2.dens[mrc2.dens > 0])\r\n",
    "    mrc2.std_norm_ave = np.linalg.norm(mrc2.dens[mrc2.dens > 0] - mrc2.ave)\r\n",
    "\r\n",
    "    print(\"#MAP AVE={ave} STD={std} STD_norm={std_norm}\".format(ave=mrc2.ave, std=mrc2.std, std_norm=mrc2.std_norm_ave))\r\n",
    "    # return False\r\n",
    "    return mrc2\r\n"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "def fastVEC_prob(mrc1, mrc2, dreso=16.0):\r\n",
    "\r\n",
    "    xydim = mrc1.xdim * mrc1.ydim\r\n",
    "    Ndata = mrc2.xdim * mrc2.ydim * mrc2.zdim\r\n",
    "\r\n",
    "    print(len(mrc2.dens))\r\n",
    "\r\n",
    "    print(\"#Start VEC\")\r\n",
    "    gstep = mrc1.xwidth\r\n",
    "    fs = (dreso / gstep) * 0.5\r\n",
    "    fs = fs ** 2\r\n",
    "    fsiv = 1.0 / fs\r\n",
    "    fmaxd = (dreso / gstep) * 2.0\r\n",
    "    print(\"#maxd= {fmaxd}\".format(fmaxd=fmaxd))\r\n",
    "    print(\"#fsiv= \" + str(fsiv))\r\n",
    "\r\n",
    "    dsum = 0.0\r\n",
    "    Nact = 0\r\n",
    "\r\n",
    "    list_d = []\r\n",
    "\r\n",
    "    for x in tqdm(range(mrc2.xdim)):\r\n",
    "        for y in range(mrc2.ydim):\r\n",
    "            for z in range(mrc2.zdim):\r\n",
    "                stp = [0] * 3\r\n",
    "                endp = [0] * 3\r\n",
    "                ind2 = 0\r\n",
    "                ind = 0\r\n",
    "\r\n",
    "                pos = [0.0] * 3\r\n",
    "                pos2 = [0.0] * 3\r\n",
    "                ori = [0.0] * 3\r\n",
    "\r\n",
    "                tmpcd = [0.0] * 3\r\n",
    "\r\n",
    "                v, dtotal, rd = 0.0, 0.0, 0.0\r\n",
    "\r\n",
    "                pos[0] = (x * mrc2.xwidth + mrc2.orig[\"x\"] - mrc1.orig[\"x\"]) / mrc1.xwidth\r\n",
    "                pos[1] = (y * mrc2.xwidth + mrc2.orig[\"y\"] - mrc1.orig[\"y\"]) / mrc1.xwidth\r\n",
    "                pos[2] = (z * mrc2.xwidth + mrc2.orig[\"z\"] - mrc1.orig[\"z\"]) / mrc1.xwidth\r\n",
    "\r\n",
    "                ind = mrc2.xdim * mrc2.ydim * z + mrc2.xdim * y + x\r\n",
    "\r\n",
    "                # check density\r\n",
    "\r\n",
    "                if (\r\n",
    "                    pos[0] < 0\r\n",
    "                    or pos[1] < 0\r\n",
    "                    or pos[2] < 0\r\n",
    "                    or pos[0] >= mrc1.xdim\r\n",
    "                    or pos[1] >= mrc1.ydim\r\n",
    "                    or pos[2] >= mrc1.zdim\r\n",
    "                ):\r\n",
    "                    mrc2.dens[ind] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue\r\n",
    "\r\n",
    "                if mrc1.data[int(pos[0])][int(pos[1])][int(pos[2])] == 0:\r\n",
    "                    mrc2.dens[ind] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    #mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue\r\n",
    "\r\n",
    "                ori[0] = pos[0]\r\n",
    "                ori[1] = pos[1]\r\n",
    "                ori[2] = pos[2]\r\n",
    "\r\n",
    "                # Start Point\r\n",
    "                stp[0] = int(pos[0] - fmaxd)\r\n",
    "                stp[1] = int(pos[1] - fmaxd)\r\n",
    "                stp[2] = int(pos[2] - fmaxd)\r\n",
    "\r\n",
    "                # set start and end point\r\n",
    "                if stp[0] < 0:\r\n",
    "                    stp[0] = 0\r\n",
    "                if stp[1] < 0:\r\n",
    "                    stp[1] = 0\r\n",
    "                if stp[2] < 0:\r\n",
    "                    stp[2] = 0\r\n",
    "\r\n",
    "                endp[0] = int(pos[0] + fmaxd + 1)\r\n",
    "                endp[1] = int(pos[1] + fmaxd + 1)\r\n",
    "                endp[2] = int(pos[2] + fmaxd + 1)\r\n",
    "\r\n",
    "                if endp[0] >= mrc1.xdim:\r\n",
    "                    endp[0] = mrc1.xdim\r\n",
    "                if endp[1] >= mrc1.ydim:\r\n",
    "                    endp[1] = mrc1.ydim\r\n",
    "                if endp[2] >= mrc1.zdim:\r\n",
    "                    endp[2] = mrc1.zdim\r\n",
    "\r\n",
    "                # setup for numba acc\r\n",
    "                stp_t = List()\r\n",
    "                endp_t = List()\r\n",
    "                pos_t = List()\r\n",
    "                [stp_t.append(x) for x in stp]\r\n",
    "                [endp_t.append(x) for x in endp]\r\n",
    "                [pos_t.append(x) for x in pos]\r\n",
    "\r\n",
    "                # compute the total density\r\n",
    "                dtotal, pos2 = calc(stp_t, endp_t, pos_t, mrc1.data, fsiv)\r\n",
    "\r\n",
    "                mrc2.dens[ind] = dtotal\r\n",
    "                mrc2.data[x][y][z] = dtotal\r\n",
    "\r\n",
    "                ''''if dtotal == 0:\r\n",
    "                    mrc2.vec[x][y][z][0] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][1] = 0.0\r\n",
    "                    mrc2.vec[x][y][z][2] = 0.0\r\n",
    "                    continue'''\r\n",
    "\r\n",
    "                rd = 1.0 / dtotal\r\n",
    "\r\n",
    "                pos2[0] *= rd\r\n",
    "                pos2[1] *= rd\r\n",
    "                pos2[2] *= rd\r\n",
    "\r\n",
    "                tmpcd[0] = pos2[0] - pos[0]\r\n",
    "                tmpcd[1] = pos2[1] - pos[1]\r\n",
    "                tmpcd[2] = pos2[2] - pos[2]\r\n",
    "\r\n",
    "                dvec = math.sqrt(tmpcd[0] ** 2 + tmpcd[1] ** 2 + tmpcd[2] ** 2)\r\n",
    "\r\n",
    "                if dvec == 0:\r\n",
    "                    dvec = 1.0\r\n",
    "\r\n",
    "                rdvec = 1.0 / dvec\r\n",
    "\r\n",
    "                #mrc2.vec[x][y][z][0] = tmpcd[0] * rdvec\r\n",
    "                #mrc2.vec[x][y][z][1] = tmpcd[1] * rdvec\r\n",
    "                #mrc2.vec[x][y][z][2] = tmpcd[2] * rdvec\r\n",
    "\r\n",
    "                dsum += dtotal\r\n",
    "                Nact += 1\r\n",
    "\r\n",
    "    print(\"#End LDP\")\r\n",
    "    print(dsum)\r\n",
    "    print(Nact)\r\n",
    "\r\n",
    "    mrc2.dsum = dsum\r\n",
    "    mrc2.Nact = Nact\r\n",
    "    mrc2.ave = dsum / float(Nact)\r\n",
    "    mrc2.std = np.linalg.norm(mrc2.dens[mrc2.dens > 0])\r\n",
    "    mrc2.std_norm_ave = np.linalg.norm(mrc2.dens[mrc2.dens > 0] - mrc2.ave)\r\n",
    "\r\n",
    "    print(\"#MAP AVE={ave} STD={std} STD_norm={std_norm}\".format(ave=mrc2.ave, std=mrc2.std, std_norm=mrc2.std_norm_ave))\r\n",
    "    # return False\r\n",
    "    return mrc2\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "def rot_pos(vec, angle, inv=False):\r\n",
    "    r = R.from_euler(\"zyx\", angle, degrees=True)\r\n",
    "    if inv:\r\n",
    "        r = r.inv()\r\n",
    "    rotated_vec = r.apply(vec)\r\n",
    "    return rotated_vec"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def rot_pos_mtx(mtx, vec):\r\n",
    "    mtx = mtx.astype(np.float32)\r\n",
    "    vec = vec.astype(np.float32)\r\n",
    "    return vec @ mtx"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def rot_mrc_combine(old_pos, new_pos, new_vec_shape):\r\n",
    "    # combined_arr = np.hstack((old_pos,new_pos))\r\n",
    "    combined_arr = np.concatenate((old_pos, new_pos), axis=1)\r\n",
    "\r\n",
    "    combined_arr = combined_arr[\r\n",
    "        (combined_arr[:, 0] >= 0)\r\n",
    "        & (combined_arr[:, 1] >= 0)\r\n",
    "        & (combined_arr[:, 2] >= 0)\r\n",
    "        & (combined_arr[:, 0] < new_vec_shape[0])\r\n",
    "        & (combined_arr[:, 1] < new_vec_shape[0])\r\n",
    "        & (combined_arr[:, 2] < new_vec_shape[0])\r\n",
    "    ]\r\n",
    "\r\n",
    "    return combined_arr"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def test_rot_mrc(orig_mrc_data, orig_mrc_vec, mtx):\r\n",
    "\r\n",
    "    mtx = mtx.astype(np.float32)\r\n",
    "\r\n",
    "    # init new vec and dens array\r\n",
    "    new_vec_array = np.zeros_like(orig_mrc_vec)\r\n",
    "    new_data_array = np.zeros_like(orig_mrc_data)\r\n",
    "\r\n",
    "    dim = orig_mrc_data.shape[0]\r\n",
    "    cent = 0.5 * dim\r\n",
    "\r\n",
    "    for x in range(dim):\r\n",
    "        for y in range(dim):\r\n",
    "            for z in range(dim):\r\n",
    "                pos = (np.array([x, y, z]) - cent).astype(np.float32)\r\n",
    "                pos2 = pos @ mtx\r\n",
    "                pos2 += cent\r\n",
    "\r\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:\r\n",
    "                    continue\r\n",
    "\r\n",
    "                pos2 = pos2.astype(np.int32)\r\n",
    "\r\n",
    "                if orig_mrc_data[pos2[0]][pos2[1]][pos2[2]] == 0:\r\n",
    "                    continue\r\n",
    "\r\n",
    "                vec = orig_mrc_vec[pos2[0]][pos2[1]][pos2[2]].astype(np.float32)\r\n",
    "                vec2 = vec @ mtx.T\r\n",
    "\r\n",
    "                pos = pos.astype(np.int32)\r\n",
    "                new_vec_array[pos[0]][pos[1]][pos[2]] = vec2\r\n",
    "                new_data_array[pos[0]][pos[1]][pos[2]] = orig_mrc_data[pos2[0]][pos2[1]][pos2[2]]\r\n",
    "\r\n",
    "    return new_vec_array, new_data_array"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "source": [
    "def rot_mrc(orig_mrc_data, orig_mrc_vec, angle):\r\n",
    "    dim = orig_mrc_vec.shape[0]\r\n",
    "\r\n",
    "    new_pos = np.array(np.meshgrid(np.arange(dim), np.arange(dim), np.arange(dim),)).T.reshape(-1, 3)\r\n",
    "\r\n",
    "    # set the center\r\n",
    "    cent = 0.5 * float(dim)\r\n",
    "\r\n",
    "    # get relative positions from center\r\n",
    "    new_pos = new_pos - cent\r\n",
    "    # print(new_pos)\r\n",
    "\r\n",
    "    # init the rotation by euler angle\r\n",
    "    r = R.from_euler(\"ZYX\", angle, degrees=True)\r\n",
    "    mtx = r.as_matrix()\r\n",
    "    mtx[np.isclose(mtx, 0, atol=1e-15)] = 0\r\n",
    "\r\n",
    "    # print(mtx)\r\n",
    "\r\n",
    "    old_pos = rot_pos_mtx(np.flip(mtx).T, new_pos) + cent\r\n",
    "\r\n",
    "    combined_arr = np.hstack((old_pos, new_pos))\r\n",
    "\r\n",
    "    # filter values outside the boundaries\r\n",
    "    in_bound_mask = (\r\n",
    "        (old_pos[:, 0] >= 0)\r\n",
    "        & (old_pos[:, 1] >= 0)\r\n",
    "        & (old_pos[:, 2] >= 0)\r\n",
    "        & (old_pos[:, 0] < dim)\r\n",
    "        & (old_pos[:, 1] < dim)\r\n",
    "        & (old_pos[:, 2] < dim)\r\n",
    "    )\r\n",
    "\r\n",
    "    combined_arr = combined_arr[in_bound_mask]\r\n",
    "\r\n",
    "    #     combined_arr = combined_arr[(combined_arr[:, 0] >= 0) &\r\n",
    "    #                                 (combined_arr[:, 1] >= 0) &\r\n",
    "    #                                 (combined_arr[:, 2] >= 0) &\r\n",
    "    #                                 (combined_arr[:, 0] < dim) &\r\n",
    "    #                                 (combined_arr[:, 1] < dim) &\r\n",
    "    #                                 (combined_arr[:, 2] < dim)]\r\n",
    "\r\n",
    "    # print(combined_arr)\r\n",
    "\r\n",
    "    # conversion indices to int\r\n",
    "    combined_arr = combined_arr.astype(np.int32)\r\n",
    "\r\n",
    "    # print(combined_arr.shape)\r\n",
    "\r\n",
    "    # combined_arr = rot_mrc_combine(old_pos, new_pos, new_vec_shape)\r\n",
    "\r\n",
    "    index_arr = combined_arr[:, 0:3]\r\n",
    "\r\n",
    "    # print(index_arr)\r\n",
    "    # print(np.count_nonzero(orig_mrc_data))\r\n",
    "\r\n",
    "    dens_mask = orig_mrc_data[index_arr[:, 0], index_arr[:, 1], index_arr[:, 2]] != 0.0\r\n",
    "\r\n",
    "    # print(dens_mask.shape)\r\n",
    "    # print(dens_mask)\r\n",
    "\r\n",
    "    non_zero_rot_list = combined_arr[dens_mask]\r\n",
    "\r\n",
    "    # print(non_zero_rot_list.shape)\r\n",
    "    #     with np.printoptions(threshold=np.inf):\r\n",
    "    #         print(non_zero_rot_list[:, 0:3])\r\n",
    "\r\n",
    "    non_zero_vec = orig_mrc_vec[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\r\n",
    "\r\n",
    "    non_zero_dens = orig_mrc_data[non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]]\r\n",
    "\r\n",
    "    # print(non_zero_dens)\r\n",
    "\r\n",
    "    # non_zero_dens[:, [2, 0]] = non_zero_dens[:, [0, 2]]\r\n",
    "    new_vec = rot_pos_mtx(np.flip(mtx), non_zero_vec)\r\n",
    "\r\n",
    "    # print(new_vec)\r\n",
    "\r\n",
    "    # init new vec and dens array\r\n",
    "    new_vec_array = np.zeros_like(orig_mrc_vec)\r\n",
    "    new_data_array = np.zeros_like(orig_mrc_data)\r\n",
    "\r\n",
    "    # print(new)\r\n",
    "\r\n",
    "    # fill in the new data\r\n",
    "    for vec, ind, dens in zip(new_vec, (non_zero_rot_list[:, 3:6] + cent).astype(int), non_zero_dens):\r\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][0] = vec[0]\r\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][1] = vec[1]\r\n",
    "        new_vec_array[ind[0]][ind[1]][ind[2]][2] = vec[2]\r\n",
    "        new_data_array[ind[0]][ind[1]][ind[2]] = dens\r\n",
    "\r\n",
    "    return new_vec_array, new_data_array"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "def find_best_trans_list(input_list):\r\n",
    "    \r\n",
    "    sum_arr = np.zeros_like(input_list[0])\r\n",
    "    for arr in input_list:\r\n",
    "        sum_arr = sum_arr + arr\r\n",
    "    best = np.amax(sum_arr)\r\n",
    "    trans = np.unravel_index(sum_arr.argmax(), sum_arr.shape)\r\n",
    "    \r\n",
    "    return best, trans"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "source": [
    "def get_score(\r\n",
    "    target_map_data, search_map_data, target_map_vec, search_map_vec, trans, ave1, ave2, std1, std2, pstd1, pstd2\r\n",
    "):\r\n",
    "\r\n",
    "    px, py, pz = 0, 0, 0\r\n",
    "    dim = target_map_data.shape[0]\r\n",
    "    total = 0\r\n",
    "\r\n",
    "    t = np.array(trans)\r\n",
    "    if trans[0] > 0.5 * dim:\r\n",
    "        t[0] -= dim\r\n",
    "    if trans[1] > 0.5 * dim:\r\n",
    "        t[1] -= dim\r\n",
    "    if trans[2] > 0.5 * dim:\r\n",
    "        t[2] -= dim\r\n",
    "        \r\n",
    "    #duplicate this part for all probability maps\r\n",
    "\r\n",
    "    target_pos = np.array(np.meshgrid(np.arange(dim), np.arange(dim), np.arange(dim),)).T.reshape(-1, 3)\r\n",
    "\r\n",
    "    search_pos = target_pos + t\r\n",
    "\r\n",
    "    total += np.count_nonzero(target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]])\r\n",
    "\r\n",
    "    combined_arr = np.hstack((target_pos, search_pos))\r\n",
    "\r\n",
    "    combined_arr = combined_arr[\r\n",
    "        (combined_arr[:, 3] >= 0)\r\n",
    "        & (combined_arr[:, 4] >= 0)\r\n",
    "        & (combined_arr[:, 5] >= 0)\r\n",
    "        & (combined_arr[:, 3] < dim)\r\n",
    "        & (combined_arr[:, 4] < dim)\r\n",
    "        & (combined_arr[:, 5] < dim)\r\n",
    "    ]\r\n",
    "\r\n",
    "    target_pos = combined_arr[:, 0:3]\r\n",
    "    search_pos = combined_arr[:, 3:6]\r\n",
    "\r\n",
    "    d1 = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]]\r\n",
    "    d2 = search_map_data[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]]\r\n",
    "    \r\n",
    "    d1 = np.where(d1 <= 0, 0.0, d1)\r\n",
    "    d2 = np.where(d2 <= 0, 0.0, d1)\r\n",
    "\r\n",
    "    print(np.sum(d1))\r\n",
    "    print(np.sum(d2))\r\n",
    "\r\n",
    "    pd1 = np.where(d1 <= 0, 0.0, d1 - ave1)\r\n",
    "    pd2 = np.where(d2 <= 0, 0.0, d2 - ave2)\r\n",
    "\r\n",
    "    cc = np.sum(np.multiply(d1, d2))\r\n",
    "    pcc = np.sum(np.multiply(pd1, pd2))\r\n",
    "\r\n",
    "    target_zero_mask = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]] == 0\r\n",
    "    target_non_zero_mask = target_map_data[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]] > 0\r\n",
    "    search_non_zero_mask = search_map_data[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]] > 0\r\n",
    "    search_non_zero_count = np.count_nonzero(np.multiply(target_zero_mask, search_non_zero_mask))\r\n",
    "\r\n",
    "    trimmed_target_vec = target_map_vec[target_pos[:, 0], target_pos[:, 1], target_pos[:, 2]]\r\n",
    "    trimmed_search_vec = search_map_vec[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]]\r\n",
    "\r\n",
    "    total += search_non_zero_count\r\n",
    "\r\n",
    "    sco_arr = np.zeros_like(search_map_data)\r\n",
    "    sco = np.einsum(\"ij,ij->i\", trimmed_target_vec, trimmed_search_vec)\r\n",
    "    sco_arr[search_pos[:, 0], search_pos[:, 1], search_pos[:, 2]] = sco\r\n",
    "    sco_sum = np.sum(sco_arr)\r\n",
    "    Nm = np.count_nonzero(np.multiply(target_non_zero_mask, search_non_zero_mask))\r\n",
    "\r\n",
    "    print(\r\n",
    "        \"Overlap= \"\r\n",
    "        + str(float(Nm) / float(total))\r\n",
    "        + \" \"\r\n",
    "        + str(Nm)\r\n",
    "        + \"/\"\r\n",
    "        + str(total)\r\n",
    "        + \" CC= \"\r\n",
    "        + str(cc / (std1 * std2))\r\n",
    "        + \" PCC= \"\r\n",
    "        + str(pcc / (pstd1 * pstd2))\r\n",
    "    )\r\n",
    "    print(\"Score= \", sco_sum)\r\n",
    "    return sco_arr"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "source": [
    "def fft_search_score_trans(target_X, target_Y, target_Z, search_vec, a, b, c, fft_object, ifft_object):\r\n",
    "    x2 = copy.deepcopy(search_vec[..., 0])\r\n",
    "    y2 = copy.deepcopy(search_vec[..., 1])\r\n",
    "    z2 = copy.deepcopy(search_vec[..., 2])\r\n",
    "\r\n",
    "    X2 = np.zeros_like(target_X)\r\n",
    "    np.copyto(a, x2)\r\n",
    "    np.copyto(X2, fft_object(a))\r\n",
    "    dot_X = target_X * X2\r\n",
    "    np.copyto(b, dot_X)\r\n",
    "    dot_x = np.zeros_like(x2)\r\n",
    "    np.copyto(dot_x, ifft_object(b))\r\n",
    "\r\n",
    "    Y2 = np.zeros_like(target_Y)\r\n",
    "    np.copyto(a, y2)\r\n",
    "    np.copyto(Y2, fft_object(a))\r\n",
    "    dot_Y = target_Y * Y2\r\n",
    "    np.copyto(b, dot_Y)\r\n",
    "    dot_y = np.zeros_like(y2)\r\n",
    "    np.copyto(dot_y, ifft_object(b))\r\n",
    "\r\n",
    "    Z2 = np.zeros_like(target_Z)\r\n",
    "    np.copyto(a, z2)\r\n",
    "    np.copyto(Z2, fft_object(a))\r\n",
    "    dot_Z = target_Z * Z2\r\n",
    "    np.copyto(b, dot_Z)\r\n",
    "    dot_z = np.zeros_like(z2)\r\n",
    "    np.copyto(dot_z, ifft_object(b))\r\n",
    "\r\n",
    "    #         if tuple(angle) == (0, 0, 30):\r\n",
    "    #             XX = [x12, y12, z12]\r\n",
    "\r\n",
    "    #         X2 = np.fft.rfftn(x2)\r\n",
    "    #         X12 = X1 * X2\r\n",
    "    #         x12 = np.fft.irfftn(X12, norm=\"forward\")\r\n",
    "\r\n",
    "    # #         if (tuple(angle) == (0,0,30)):\r\n",
    "    # #             XX = [X12,x12]\r\n",
    "\r\n",
    "    #         Y2 = np.fft.rfftn(y2)\r\n",
    "    #         Y12 = Y1 * Y2\r\n",
    "    #         y12 = np.fft.irfftn(Y12, norm=\"forward\")\r\n",
    "\r\n",
    "    #         Z2 = np.fft.rfftn(z2)\r\n",
    "    #         Z12 = Z1 * Z2\r\n",
    "    #         z12 = np.fft.irfftn(Z12, norm=\"forward\")\r\n",
    "\r\n",
    "    return find_best_trans_list([dot_x, dot_y, dot_z])"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "source": [
    "def fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object):\r\n",
    "    dot_product_list = []\r\n",
    "    for target_complex, query_real in zip(target_list, query_list):\r\n",
    "        \r\n",
    "        query_complex = np.zeros_like(target_complex)\r\n",
    "        np.copyto(a, query_real)\r\n",
    "        np.copyto(query_complex, fft_object(a))\r\n",
    "        dot_complex = target_complex * query_complex\r\n",
    "        np.copyto(b, dot_complex)\r\n",
    "        dot_real = np.zeros_like(query_real)\r\n",
    "        np.copyto(dot_real, ifft_object(b))\r\n",
    "        \r\n",
    "        dot_product_list.append(dot_real)\r\n",
    "        \r\n",
    "    return dot_product_list"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "source": [
    "def fft_search_score_trans_1d(target_X, search_data, a, b, fft_object, ifft_object, mode, ave=None):\r\n",
    "\r\n",
    "    x2 = copy.deepcopy(search_data)\r\n",
    "\r\n",
    "    if mode == \"Overlap\":\r\n",
    "        x2 = np.where(x2 > 0, 1.0, 0.0)\r\n",
    "    elif mode == \"CC\":\r\n",
    "        x2 = np.where(x2 > 0, x2, 0.0)\r\n",
    "    elif mode == \"PCC\":\r\n",
    "        x2 = np.where(x2 > 0, x2 - ave, 0.0)\r\n",
    "    elif mode == \"Laplacian\":\r\n",
    "        weights = np.array(\r\n",
    "            [\r\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\r\n",
    "                [[0.0, 1.0, 0.0], [1.0, -6.0, 1.0], [0.0, 1.0, 0.0]],\r\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\r\n",
    "            ]\r\n",
    "        )\r\n",
    "        x2 = convolve(x2, weights, mode=\"constant\")\r\n",
    "\r\n",
    "    X2 = np.zeros_like(target_X)\r\n",
    "    np.copyto(a, x2)\r\n",
    "    np.copyto(X2, fft_object(a))\r\n",
    "    dot_X = target_X * X2\r\n",
    "    np.copyto(b, dot_X)\r\n",
    "    dot_x = np.zeros_like(x2)\r\n",
    "    np.copyto(dot_x, ifft_object(b))\r\n",
    "\r\n",
    "    return find_best_trans_list([dot_x])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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=30, mode=\"VecProduct\", is_eval_mode=False):\r\n",
    "\r\n",
    "    if is_eval_mode:\r\n",
    "        print(\"#For Evaluation Mode\")\r\n",
    "        print(\"#Please use the same coordinate system and map size for map1 and map2.\")\r\n",
    "        print(\"#Example:\")\r\n",
    "        print(\"#In Chimera command line: open map1 and map2 as #0 and #1, then type\")\r\n",
    "        print(\"#> open map1.mrc\")\r\n",
    "        print(\"#> open map2.mrc\")\r\n",
    "        print(\"#> vop #1 resample onGrid #0\")\r\n",
    "        print(\"#> volume #2 save new.mrc\")\r\n",
    "        print(\"#Chimera will generate the resampled map2.mrc as new.mrc\")\r\n",
    "        return\r\n",
    "\r\n",
    "    #     x1 = copy.deepcopy(mrc_target.vec[:, :, :, 0])\r\n",
    "    #     y1 = copy.deepcopy(mrc_target.vec[:, :, :, 1])\r\n",
    "    #     z1 = copy.deepcopy(mrc_target.vec[:, :, :, 2])\r\n",
    "\r\n",
    "    # init the target map vectors\r\n",
    "    #does this part need to be changed?\r\n",
    "    x1 = copy.deepcopy(mrc_target.vec[:, :, :, 0])\r\n",
    "    p1 = copy.deepcopy(mrc_P1.data)\r\n",
    "    p2 = copy.deepcopy(mrc_P2.data)\r\n",
    "    p3 = copy.deepcopy(mrc_P3.data)\r\n",
    "    p4 = copy.deepcopy(mrc_P4.data)\r\n",
    "\r\n",
    "\r\n",
    "    if mode == \"VecProduct\":\r\n",
    "        y1 = copy.deepcopy(mrc_target.vec[:, :, :, 1])\r\n",
    "        z1 = copy.deepcopy(mrc_target.vec[:, :, :, 2])\r\n",
    "\r\n",
    "    # Postprocessing for other modes\r\n",
    "    if mode == \"Overlap\":\r\n",
    "        x1 = np.where(x1 > 0, 1.0, 0.0)\r\n",
    "    elif mode == \"CC\":\r\n",
    "        x1 = np.where(x1 > 0, x1, 0.0)\r\n",
    "    elif mode == \"PCC\":\r\n",
    "        x1 = np.where(x1 > 0, x1 - mrc_target.ave, 0.0)\r\n",
    "    elif mode == \"Laplacian\":\r\n",
    "        weights = np.array(\r\n",
    "            [\r\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\r\n",
    "                [[0.0, 1.0, 0.0], [1.0, -6.0, 1.0], [0.0, 1.0, 0.0]],\r\n",
    "                [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],\r\n",
    "            ]\r\n",
    "        )\r\n",
    "        x1 = convolve(x1, weights, mode=\"constant\")\r\n",
    "\r\n",
    "    d3 = mrc_target.xdim ** 3\r\n",
    "\r\n",
    "    rd3 = 1.0 / d3\r\n",
    "    \r\n",
    "    \r\n",
    "    X1 = np.fft.rfftn(x1)\r\n",
    "    X1 = np.conj(X1)\r\n",
    "    P1 = np.fft.rfftn(p1)\r\n",
    "    P1 = np.conj(P1)\r\n",
    "    P2 = np.fft.rfftn(p2)\r\n",
    "    P2 = np.conj(P2)\r\n",
    "    P3 = np.fft.rfftn(p3)\r\n",
    "    P3 = np.conj(P3)\r\n",
    "    P4 = np.fft.rfftn(p4)\r\n",
    "    P4 = np.conj(P4)\r\n",
    "    \r\n",
    "    #P4 = 1 - P2 - P3 - P1\r\n",
    "\r\n",
    "    if mode == \"VecProduct\":\r\n",
    "        Y1 = np.fft.rfftn(y1)\r\n",
    "        Y1 = np.conj(Y1)\r\n",
    "        Z1 = np.fft.rfftn(z1)\r\n",
    "        Z1 = np.conj(Z1)\r\n",
    "\r\n",
    "    x_angle = []\r\n",
    "    y_angle = []\r\n",
    "    z_angle = []\r\n",
    "\r\n",
    "    i = 0\r\n",
    "    while i < 360:\r\n",
    "        x_angle.append(i)\r\n",
    "        y_angle.append(i)\r\n",
    "        i += ang\r\n",
    "\r\n",
    "    i = 0\r\n",
    "    while i <= 180:\r\n",
    "        z_angle.append(i)\r\n",
    "        i += ang\r\n",
    "\r\n",
    "    angle_comb = np.array(np.meshgrid(x_angle, y_angle, z_angle)).T.reshape(-1, 3)\r\n",
    "\r\n",
    "    rot_vec_dict = {}\r\n",
    "    rot_data_dict = {}\r\n",
    "    #rot_vec_dict_p1 = {}\r\n",
    "    rot_data_dict_p1 = {}\r\n",
    "    #rot_vec_dict_p2 = {}\r\n",
    "    rot_data_dict_p2 = {}\r\n",
    "    #rot_vec_dict_p3 = {}\r\n",
    "    rot_data_dict_p3 = {}\r\n",
    "    #rot_vec_dict_p4 = {}\r\n",
    "    rot_data_dict_p4 = {}\r\n",
    "    \r\n",
    "\r\n",
    "    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count() + 4) as executor:\r\n",
    "        trans_vec = {executor.submit(rot_mrc, mrc_search.data, mrc_search.vec, angle,): angle for angle in angle_comb}\r\n",
    "        #trans_vec_p1 = {executor.submit(rot_mrc, mrc_search_p1.data, mrc_search_p1.vec, angle,): angle for angle in angle_comb}\r\n",
    "        #trans_vec_p2 = {executor.submit(rot_mrc, mrc_search_p2.data, mrc_search_p2.vec, angle,): angle for angle in angle_comb}\r\n",
    "        #trans_vec_p3 = {executor.submit(rot_mrc, mrc_search_p3.data, mrc_search_p3.vec, angle,): angle for angle in angle_comb}\r\n",
    "        #trans_vec_p4 = {executor.submit(rot_mrc, mrc_search_p4.data, mrc_search_p4.vec, angle,): angle for angle in angle_comb}\r\n",
    "        for future in concurrent.futures.as_completed(trans_vec):\r\n",
    "            angle = trans_vec[future]\r\n",
    "            rot_vec_dict[tuple(angle)] = future.result()[0]\r\n",
    "            rot_data_dict[tuple(angle)] = future.result()[1]\r\n",
    "            rot_data_dict_p1[tuple(angle)] = future.result()[1]\r\n",
    "            rot_data_dict_p2[tuple(angle)] = future.result()[1]\r\n",
    "            rot_data_dict_p3[tuple(angle)] = future.result()[1]\r\n",
    "            rot_data_dict_p4[tuple(angle)] = future.result()[1]\r\n",
    "\r\n",
    "    #     for angle in tqdm(angle_comb, desc=\"Rotation\"):\r\n",
    "    #         rot_result = rot_mrc(\r\n",
    "    #             mrc_search.data,\r\n",
    "    #             mrc_search.vec,\r\n",
    "    #             angle,\r\n",
    "    #         )\r\n",
    "    #         mrc_angle_dict[tuple(angle)] = rot_result\r\n",
    "\r\n",
    "    # fftw plans\r\n",
    "    a = pyfftw.empty_aligned((x1.shape), dtype=\"float32\")\r\n",
    "    b = pyfftw.empty_aligned((a.shape[0], a.shape[1], a.shape[2] // 2 + 1), dtype=\"complex64\")\r\n",
    "    c = pyfftw.empty_aligned((x1.shape), dtype=\"float32\")\r\n",
    "\r\n",
    "    fft_object = pyfftw.FFTW(a, b, axes=(0, 1, 2))\r\n",
    "    ifft_object = pyfftw.FFTW(b, c, direction=\"FFTW_BACKWARD\", axes=(0, 1, 2), normalise_idft=False)\r\n",
    "\r\n",
    "    angle_score = []\r\n",
    "\r\n",
    "    for angle in tqdm(angle_comb, desc=\"FFT\"):\r\n",
    "        rot_mrc_vec = rot_vec_dict[tuple(angle)]\r\n",
    "        rot_mrc_data = rot_data_dict[tuple(angle)]\r\n",
    "        #rot_mrc_vec_p1 = rot_vec_dict_p1[tuple(angle)]\r\n",
    "        rot_mrc_data_p1 = rot_data_dict_p1[tuple(angle)]\r\n",
    "        #rot_mrc_vec_p2 = rot_vec_dict_p2[tuple(angle)]\r\n",
    "        rot_mrc_data_p2 = rot_data_dict_p2[tuple(angle)]\r\n",
    "        #rot_mrc_vec_p3 = rot_vec_dict_p3[tuple(angle)]\r\n",
    "        rot_mrc_data_p3 = rot_data_dict_p3[tuple(angle)]\r\n",
    "        #rot_mrc_vec_p4 = rot_vec_dict_p4[tuple(angle)]\r\n",
    "        rot_mrc_data_p4 = rot_data_dict_p4[tuple(angle)]\r\n",
    "        \r\n",
    "      \r\n",
    "        if mode == \"VecProduct\":\r\n",
    "            \r\n",
    "            x2 = copy.deepcopy(rot_mrc_vec[..., 0])\r\n",
    "            y2 = copy.deepcopy(rot_mrc_vec[..., 1])\r\n",
    "            z2 = copy.deepcopy(rot_mrc_vec[..., 2])\r\n",
    "            p21 = copy.deepcopy(rot_mrc_data_p1)\r\n",
    "            p22 = copy.deepcopy(rot_mrc_data_p2)\r\n",
    "            p23 = copy.deepcopy(rot_mrc_data_p3)\r\n",
    "            p24 = copy.deepcopy(rot_mrc_data_p4)\r\n",
    "            \r\n",
    "            #p24 = 1 - (p21 + p22 + p23)\r\n",
    "                    \r\n",
    "            target_list = [X1, Y1, Z1, P1, P2, P3, P4]\r\n",
    "            query_list = [x2, y2, z2, p21, p22, p23, p24]\r\n",
    "            \r\n",
    "            fft_result_list = fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object)\r\n",
    "                    \r\n",
    "            best, trans = find_best_trans_list(fft_result_list)\r\n",
    "            \r\n",
    "        else:\r\n",
    "            best, trans = fft_search_score_trans_1d(\r\n",
    "                X1, rot_mrc_data, a, b, fft_object, ifft_object, mode, mrc_target.ave\r\n",
    "            )\r\n",
    "            if mode == \"CC\":\r\n",
    "                rstd2 = 1.0 / mrc_target.std ** 2\r\n",
    "                best = best * rstd2\r\n",
    "            if mode == \"PCC\":\r\n",
    "                rstd3 = 1.0 / mrc_target.std_norm_ave ** 2\r\n",
    "                best = best * rstd3\r\n",
    "\r\n",
    "        angle_score.append([tuple(angle), best * rd3, trans])\r\n",
    "\r\n",
    "    # calculate the ave and std\r\n",
    "    score_arr = np.array([row[1] for row in angle_score])\r\n",
    "    ave = np.mean(score_arr)\r\n",
    "    std = np.std(score_arr)\r\n",
    "    print(\"Std= \" + str(std) + \" Ave= \" + str(ave))\r\n",
    "\r\n",
    "    # sort the list and get topN\r\n",
    "    sorted_topN = sorted(angle_score, key=lambda x: x[1], reverse=True)[:TopN]\r\n",
    "\r\n",
    "    for x in sorted_topN:\r\n",
    "        print(x)\r\n",
    "\r\n",
    "    refined_score = []\r\n",
    "    if ang > 5.0:\r\n",
    "        for t_mrc in sorted_topN:\r\n",
    "            ang_list = np.array(\r\n",
    "                np.meshgrid(\r\n",
    "                    [t_mrc[0][0] - 5, t_mrc[0][0], t_mrc[0][0] + 5],\r\n",
    "                    [t_mrc[0][1] - 5, t_mrc[0][1], t_mrc[0][1] + 5],\r\n",
    "                    [t_mrc[0][2] - 5, t_mrc[0][2], t_mrc[0][2] + 5],\r\n",
    "                )\r\n",
    "            ).T.reshape(-1, 3)\r\n",
    "            # print(ang_list)\r\n",
    "            for ang in ang_list:\r\n",
    "                rotated_vec, rotated_data = rot_mrc(mrc_search.data, mrc_search.vec, ang)\r\n",
    "\r\n",
    "                if mode == \"VecProduct\":\r\n",
    "                    \r\n",
    "                    x2 = copy.deepcopy(rotated_vec[..., 0])\r\n",
    "                    y2 = copy.deepcopy(rotated_vec[..., 1])\r\n",
    "                    z2 = copy.deepcopy(rotated_vec[..., 2])\r\n",
    "                    p21 = copy.deepcopy(rot_mrc_data_p1)\r\n",
    "                    p22 = copy.deepcopy(rot_mrc_data_p2)\r\n",
    "                    p23 = copy.deepcopy(rot_mrc_data_p3)\r\n",
    "                    p24 = copy.deepcopy(rot_mrc_data_p4)\r\n",
    "                    \r\n",
    "                    target_list = [X1, Y1, Z1, P1, P2, P3, P4]\r\n",
    "                    query_list = [x2, y2, z2, p21, p22, p23, p24]\r\n",
    "                    \r\n",
    "                    fft_result_list = fft_search_best_dot(target_list, query_list, a, b, c, fft_object, ifft_object)\r\n",
    "                    \r\n",
    "                    best, trans = find_best_trans_list(fft_result_list)\r\n",
    "                    \r\n",
    "                    # best, trans = fft_search_score_trans(X1, Y1, Z1, rotated_vec, a, b, c, fft_object, ifft_object)\r\n",
    "                else:\r\n",
    "                    best, trans = fft_search_score_trans_1d(\r\n",
    "                        X1, rotated_data, a, b, fft_object, ifft_object, mode, ave=mrc_target.ave\r\n",
    "                    )\r\n",
    "                    if mode == \"CC\":\r\n",
    "                        rstd2 = 1.0 / mrc_target.std ** 2\r\n",
    "                        best = best * rstd2\r\n",
    "                    if mode == \"PCC\":\r\n",
    "                        rstd3 = 1.0 / mrc_target.std_norm_ave ** 2\r\n",
    "                        best = best * rstd3\r\n",
    "\r\n",
    "                refined_score.append([tuple(ang), best * rd3, trans, rotated_vec, rotated_data])\r\n",
    "\r\n",
    "        refined_list = sorted(refined_score, key=lambda x: x[1], reverse=True)[:TopN]\r\n",
    "    else:\r\n",
    "        refined_list = sorted_topN\r\n",
    "\r\n",
    "    for i, t_mrc in enumerate(refined_list):\r\n",
    "        print(\"R=\" + str(t_mrc[0]) + \" T=\" + str(t_mrc[2]))\r\n",
    "        sco = get_score(\r\n",
    "            mrc_target.data,\r\n",
    "            t_mrc[4],\r\n",
    "            mrc_target.vec,\r\n",
    "            t_mrc[3],\r\n",
    "            t_mrc[2],\r\n",
    "            mrc_target.ave,\r\n",
    "            mrc_search.ave,\r\n",
    "            mrc_target.std,\r\n",
    "            mrc_search.std,\r\n",
    "            mrc_target.std_norm_ave,\r\n",
    "            mrc_search.std_norm_ave,\r\n",
    "        )\r\n",
    "        \r\n",
    "        # Write result to PDB files\r\n",
    "        show_vec(mrc_target.orig, t_mrc[3], t_mrc[4], sco, mrc_search.xwidth, t_mrc[2], \"model_\" + str(i + 1) + \".pdb\")\r\n",
    "\r\n",
    "    #     num_jobs = math.ceil(360 / ang) * math.ceil(360 / ang) * (180 // ang + 1)\r\n",
    "\r\n",
    "    return refined_list"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "source": [
    "def show_vec(origin, sampled_mrc_vec, sampled_mrc_data, sampled_mrc_score, sample_width, trans, name):\r\n",
    "\r\n",
    "    dim = sampled_mrc_data.shape[0]\r\n",
    "\r\n",
    "    origin = np.array([origin[\"x\"], origin[\"y\"], origin[\"z\"]])\r\n",
    "    trans = np.array(trans)\r\n",
    "\r\n",
    "    if trans[0] > 0.5 * dim:\r\n",
    "        trans[0] -= dim\r\n",
    "    if trans[1] > 0.5 * dim:\r\n",
    "        trans[1] -= dim\r\n",
    "    if trans[2] > 0.5 * dim:\r\n",
    "        trans[2] -= dim\r\n",
    "\r\n",
    "    add = origin - trans * sample_width\r\n",
    "\r\n",
    "    natm = 1\r\n",
    "    nres = 1\r\n",
    "\r\n",
    "    pdb_file = open(name, \"w\")\r\n",
    "    #     pdb_file.write(\"MODEL\\n\")\r\n",
    "    for x in range(dim):\r\n",
    "        for y in range(dim):\r\n",
    "            for z in range(dim):\r\n",
    "\r\n",
    "                if sampled_mrc_data[x][y][z] != 0.0:\r\n",
    "                    tmp = np.array([x, y, z])\r\n",
    "                    tmp = tmp * sample_width + add\r\n",
    "                    atom_header = \"ATOM{:>7d}  CA  ALA{:>6d}    \".format(natm, nres)\r\n",
    "                    atom_content = \"{:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}\".format(\r\n",
    "                        tmp[0], tmp[1], tmp[2], 1.0, sampled_mrc_score[x][y][z]\r\n",
    "                    )\r\n",
    "                    pdb_file.write(atom_header + atom_content + \"\\n\")\r\n",
    "                    natm += 1\r\n",
    "\r\n",
    "                    tmp = np.array([x, y, z])\r\n",
    "                    tmp = (tmp + sampled_mrc_vec[x][y][z]) * sample_width + add\r\n",
    "                    atom_header = \"ATOM{:>7d}  CB  ALA{:>6d}    \".format(natm, nres)\r\n",
    "                    atom_content = \"{:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}\".format(\r\n",
    "                        tmp[0], tmp[1], tmp[2], 1.0, sampled_mrc_score[x][y][z]\r\n",
    "                    )\r\n",
    "                    pdb_file.write(atom_header + atom_content + \"\\n\")\r\n",
    "                    natm += 1\r\n",
    "\r\n",
    "                    nres += 1\r\n",
    "\r\n",
    "\r\n",
    "#     pdb_file.write(\"ENDMDL\\n\")"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "source": [
    "mrc1 = mrc_obj(\"./Data/emd_8097.mrc\")\r\n",
    "mrc2 = mrc_obj(\"./Data/grid1A.mrc\")\r\n",
    "\r\n",
    "#probability MRCs\r\n",
    "\r\n",
    "mrc1_p1 = mrc_obj(\"./Data/emd_8097_helix.mrc\")\r\n",
    "mrc1_p2 = mrc_obj(\"./Data/emd_8097_strand.mrc\")\r\n",
    "mrc1_p3 = mrc_obj(\"./Data/emd_8097_coil.mrc\")\r\n",
    "mrc1_p4 = mrc_obj(\"./Data/emd_8097_DNA.mrc\")\r\n",
    "\r\n",
    "mrc2_p1 = mrc_obj(\"./Data/emd_8097_chainA_helix_r.mrc\")\r\n",
    "mrc2_p2 = mrc_obj(\"./Data/emd_8097_chainA_strand_r.mrc\")\r\n",
    "mrc2_p3 = mrc_obj(\"./Data/emd_8097_chainA_coil_r.mrc\")\r\n",
    "mrc2_p4 = mrc_obj(\"./Data/emd_8097_chainA_DNA_r.mrc\")\r\n",
    "\r\n",
    "#\r\n",
    "\r\n",
    "mrc1, mrc_N1 = mrc_set_vox_size(mrc1)\r\n",
    "mrc2, mrc_N2 = mrc_set_vox_size(mrc2)\r\n",
    "mrc1_p1, mrc_N1_p1 = mrc_set_vox_size(mrc1_p1)\r\n",
    "mrc1_p2, mrc_N1_p2 = mrc_set_vox_size(mrc1_p2)\r\n",
    "mrc1_p3, mrc_N1_p3 = mrc_set_vox_size(mrc1_p3)\r\n",
    "mrc1_p4, mrc_N1_p4 = mrc_set_vox_size(mrc1_p4)\r\n",
    "mrc2_p1, mrc_N2_p1 = mrc_set_vox_size(mrc2_p1)\r\n",
    "mrc2_p2, mrc_N2_p2 = mrc_set_vox_size(mrc2_p2)\r\n",
    "mrc2_p3, mrc_N2_p3 = mrc_set_vox_size(mrc2_p3)\r\n",
    "mrc2_p4, mrc_N2_p4 = mrc_set_vox_size(mrc2_p4)\r\n",
    "\r\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]\r\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))\r\n",
    "\r\n",
    "print(max_dim)\r\n",
    "\r\n",
    "for mrc in mrc_list:\r\n",
    "    if mrc.xdim != max_dim:\r\n",
    "        mrc.xdim = mrc.ydim = mrc.zdim = max_dim\r\n",
    "        mrc.orig[\"x\"] = mrc.cent[0] - 0.5 * 7 * mrc.xdim\r\n",
    "        mrc.orig[\"y\"] = mrc.cent[1] - 0.5 * 7 * mrc.xdim\r\n",
    "        mrc.orig[\"z\"] = mrc.cent[2] - 0.5 * 7 * mrc.xdim\r\n",
    "\r\n",
    "for mrc in mrc_list:\r\n",
    "    mrc.dens = np.zeros((max_dim ** 3, 1), dtype=\"float32\")\r\n",
    "    mrc.data = np.zeros((max_dim, max_dim, max_dim), dtype=\"float32\")\r\n",
    "    mrc.vec = np.zeros((max_dim, max_dim, max_dim, 3), dtype=\"float32\")"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#dmax=141.34333822979175\n",
      "Nvox= 48, 48, 48\n",
      "cent= 118.23501896858215, 107.53501725196838, 102.1850163936615\n",
      "ori= -49.76498103141785, -60.464982748031616, -65.8149836063385\n",
      "#dmax=144.43683740652867\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -29.0, -29.0, -29.0\n",
      "#dmax=203.23385544736388\n",
      "Nvox= 60, 60, 60\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -92.0, -92.0, -92.0\n",
      "#dmax=203.23385544736388\n",
      "Nvox= 60, 60, 60\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -92.0, -92.0, -92.0\n",
      "#dmax=203.23385544736388\n",
      "Nvox= 60, 60, 60\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -92.0, -92.0, -92.0\n",
      "#dmax=203.23385544736388\n",
      "Nvox= 60, 60, 60\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -92.0, -92.0, -92.0\n",
      "#dmax=148.0168909280289\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -29.0, -29.0, -29.0\n",
      "#dmax=148.0168909280289\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -29.0, -29.0, -29.0\n",
      "#dmax=148.0168909280289\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -29.0, -29.0, -29.0\n",
      "#dmax=148.0168909280289\n",
      "Nvox= 42, 42, 42\n",
      "cent= 118.0, 118.0, 118.0\n",
      "ori= -29.0, -29.0, -29.0\n",
      "60\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "source": [
    "print(mrc_N1_p1.data.shape, mrc_N1_p2.data.shape, mrc_N1_p3.data.shape, mrc_N1_p4.data.shape)\r\n",
    "print(mrc_N2_p1.data.shape, mrc_N2_p2.data.shape, mrc_N2_p3.data.shape, mrc_N2_p4.data.shape)\r\n",
    "print(mrc_N1.data.shape, mrc_N2.data.shape)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(60, 60, 60) (60, 60, 60) (60, 60, 60) (60, 60, 60)\n",
      "(60, 60, 60) (60, 60, 60) (60, 60, 60) (60, 60, 60)\n",
      "(60, 60, 60) (60, 60, 60)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "source": [
    "%%time\r\n",
    "mrc_N1 = fastVEC(mrc1, mrc_N1)\r\n",
    "mrc_N1_p1 = fastVEC_prob(mrc1_p1, mrc_N1_p1)\r\n",
    "mrc_N1_p2 = fastVEC_prob(mrc1_p2, mrc_N1_p2)\r\n",
    "mrc_N1_p3 = fastVEC_prob(mrc1_p3, mrc_N1_p3)\r\n",
    "mrc_N1_p4 = fastVEC_prob(mrc1_p4, mrc_N1_p4)\r\n",
    "\r\n",
    "\r\n",
    "\r\n",
    "mrc_N2 = fastVEC(mrc2, mrc_N2)\r\n",
    "mrc_N2_p1 = fastVEC_prob(mrc2_p1, mrc_N2_p1)\r\n",
    "mrc_N2_p2 = fastVEC_prob(mrc2_p2, mrc_N2_p2)\r\n",
    "mrc_N2_p3 = fastVEC_prob(mrc2_p3, mrc_N2_p3)\r\n",
    "mrc_N2_p4 = fastVEC_prob(mrc2_p4, mrc_N2_p4)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "216000\n",
      "#Start VEC\n",
      "#maxd= 29.9065372581333\n",
      "#fsiv= 0.017889068239927756\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:06<00:00,  9.26it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "32939.518123120644\n",
      "2257\n",
      "#MAP AVE=14.594381091325053 STD=773.427490234375 STD_norm=342.72186279296875\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:35<00:00,  1.71it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "3788320.042065102\n",
      "14079\n",
      "#MAP AVE=269.07593167590755 STD=33703.18359375 STD_norm=10796.236328125\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:34<00:00,  1.72it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "3379571.7982607177\n",
      "14073\n",
      "#MAP AVE=240.14579679249042 STD=34818.83984375 STD_norm=20019.025390625\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:35<00:00,  1.71it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "4126323.61167042\n",
      "14024\n",
      "#MAP AVE=294.2330014026255 STD=35957.0078125 STD_norm=8877.2626953125\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:29<00:00,  2.00it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "337430.0382834219\n",
      "12070\n",
      "#MAP AVE=27.956092649827827 STD=3369.867919921875 STD_norm=1386.6455078125\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:03<00:00, 19.38it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "94415.25910992023\n",
      "565\n",
      "#MAP AVE=167.1066532918942 STD=4309.1591796875 STD_norm=1670.7587890625\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:05<00:00, 10.93it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "588638.846717525\n",
      "1533\n",
      "#MAP AVE=383.97837359264514 STD=15830.130859375 STD_norm=4956.65869140625\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:05<00:00, 11.31it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "826476.6565615083\n",
      "1539\n",
      "#MAP AVE=537.0218691107916 STD=23738.8359375 STD_norm=10940.58203125\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:04<00:00, 13.32it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "222678.4840001846\n",
      "1247\n",
      "#MAP AVE=178.57135846045276 STD=6804.87646484375 STD_norm=2557.80126953125\n",
      "216000\n",
      "#Start VEC\n",
      "#maxd= 32.0\n",
      "#fsiv= 0.015625\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 60/60 [00:04<00:00, 12.84it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "318994.2742775019\n",
      "1192\n",
      "#MAP AVE=267.61264620595796 STD=12156.642578125 STD_norm=7900.44580078125\n",
      "Wall time: 2min 44s\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "source": [
    "mrc_N2_p1.data.shape"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(60, 60, 60)"
      ]
     },
     "metadata": {},
     "execution_count": 53
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "source": [
    "%%time\r\n",
    "score_list = search_map_fft(mrc_N1_p1, mrc_N1_p2, mrc_N1_p3, mrc_N1_p4, mrc_N1,\r\n",
    "                            mrc_N2, mrc_N2_p1, mrc_N2_p2, mrc_N2_p3, mrc_N2_p4, ang=30, mode=\"VecProduct\")\r\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "FFT: 100%|██████████| 1008/1008 [00:21<00:00, 46.25it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Std= 7356838.34875397 Ave= 100958095.19505231\n",
      "[(180, 330, 90), 117315915.92770371, (59, 17, 57)]\n",
      "[(180, 330, 120), 115906891.92770371, (53, 14, 57)]\n",
      "[(0, 30, 90), 115272619.95614815, (3, 17, 50)]\n",
      "[(180, 330, 150), 115240104.3911111, (47, 8, 57)]\n",
      "[(0, 30, 120), 115205559.63733333, (55, 16, 50)]\n",
      "[(180, 300, 120), 114716292.13392593, (54, 12, 3)]\n",
      "[(0, 210, 0), 114493935.7677037, (46, 0, 57)]\n",
      "[(180, 330, 180), 114493935.7677037, (46, 0, 57)]\n",
      "[(180, 150, 60), 114290434.35140741, (52, 50, 50)]\n",
      "[(0, 30, 0), 114278074.74725926, (17, 0, 50)]\n",
      "R=(-5, 205, 0) T=(49, 4, 5)\n",
      "32939.52\n",
      "2576.552\n",
      "Overlap= 0.04907063197026022 132/2690 CC= 0.01662168 PCC= -0.15864502\n",
      "Score=  29.745495\n",
      "R=(0, 205, 0) T=(49, 4, 5)\n",
      "32939.52\n",
      "2288.0342\n",
      "Overlap= 0.04587495375508694 124/2703 CC= 0.014298498 PCC= -0.11468694\n",
      "Score=  32.148216\n",
      "R=(0, 210, 0) T=(49, 4, 5)\n",
      "32939.52\n",
      "2662.6562\n",
      "Overlap= 0.049107142857142856 132/2688 CC= 0.017535577 PCC= -0.18064873\n",
      "Score=  35.5312\n",
      "R=(5, 205, 0) T=(49, 4, 5)\n",
      "32939.52\n",
      "2220.9531\n",
      "Overlap= 0.04384672070744289 119/2714 CC= 0.013933921 PCC= -0.11681828\n",
      "Score=  29.17415\n",
      "R=(5, 210, 0) T=(49, 4, 5)\n",
      "32939.52\n",
      "2519.1516\n",
      "Overlap= 0.04707190511489993 127/2698 CC= 0.016382286 PCC= -0.1631198\n",
      "Score=  31.607115\n",
      "R=(-5, 205, 5) T=(49, 4, 5)\n",
      "32939.52\n",
      "2470.6006\n",
      "Overlap= 0.048662704309063894 131/2692 CC= 0.015688704 PCC= -0.1347137\n",
      "Score=  33.537533\n",
      "R=(-5, 210, 5) T=(49, 4, 5)\n",
      "32939.52\n",
      "2822.1877\n",
      "Overlap= 0.051962616822429905 139/2675 CC= 0.018817224 PCC= -0.19399786\n",
      "Score=  32.137627\n",
      "R=(0, 210, 5) T=(49, 4, 5)\n",
      "32939.52\n",
      "2709.602\n",
      "Overlap= 0.05041075429424944 135/2678 CC= 0.01776202 PCC= -0.18145022\n",
      "Score=  34.336487\n",
      "R=(5, 210, 5) T=(49, 4, 5)\n",
      "32939.52\n",
      "2510.1902\n",
      "Overlap= 0.04714179658500371 127/2694 CC= 0.016540295 PCC= -0.15935645\n",
      "Score=  29.617031\n",
      "R=(175, 335, 180) T=(49, 4, 5)\n",
      "32939.52\n",
      "2576.552\n",
      "Overlap= 0.04907063197026022 132/2690 CC= 0.01662168 PCC= -0.15864502\n",
      "Score=  29.745495\n",
      "Wall time: 54.5 s\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%load_ext line_profiler"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%lprun -f rot_mrc rot_mrc(mrc_N2.data, mrc_N2.vec, [0,0,30])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.9.7 64-bit ('bio': conda)"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "interpreter": {
   "hash": "11dcb162716791435bf48ce72b845ec24ce37a51f37dce84df4d98a21bd70cb4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}