VESPER_prototype.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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.spatial.transform import Rotation as R\r\n",
    "from tqdm.notebook import tqdm\r\n",
    "\r\n",
    "# from Bio.PDB import *\r\n",
    "# import numba as nb"
   ],
   "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",
    "\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 = mrc_set.ydim = 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",
    "    mrc_set.dens = np.zeros((mrc_set.xdim ** 3, 1))\r\n",
    "    mrc_set.vec = np.zeros((new_xdim, new_xdim, new_xdim, 3), dtype=\"float32\")\r\n",
    "    mrc_set.data = np.zeros((new_xdim, new_xdim, new_xdim))\r\n",
    "\r\n",
    "    print(\r\n",
    "        \"Nvox= \"\r\n",
    "        + str(mrc_set.xdim)\r\n",
    "        + \", \"\r\n",
    "        + str(mrc_set.ydim)\r\n",
    "        + \", \"\r\n",
    "        + str(mrc_set.zdim)\r\n",
    "    )\r\n",
    "    print(\r\n",
    "        \"cent= \" + str(new_cent[0]) + \", \" + str(new_cent[1]) + \", \" + str(new_cent[2])\r\n",
    "    )\r\n",
    "    print(\r\n",
    "        \"ori= \"\r\n",
    "        + str(new_orig[\"x\"])\r\n",
    "        + \", \"\r\n",
    "        + str(new_orig[\"y\"])\r\n",
    "        + \", \"\r\n",
    "        + str(new_orig[\"z\"])\r\n",
    "    )\r\n",
    "\r\n",
    "    return mrc, mrc_set"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def calc(stp, endp, pos, mrc1_data, fsiv):\r\n",
    "    dtotal = 0\r\n",
    "    pos2 = [0.0] * 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": 7,
   "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] = (\r\n",
    "                    x * mrc2.xwidth + mrc2.orig[\"x\"] - mrc1.orig[\"x\"]\r\n",
    "                ) / mrc1.xwidth\r\n",
    "                pos[1] = (\r\n",
    "                    y * mrc2.xwidth + mrc2.orig[\"y\"] - mrc1.orig[\"y\"]\r\n",
    "                ) / mrc1.xwidth\r\n",
    "                pos[2] = (\r\n",
    "                    z * mrc2.xwidth + mrc2.orig[\"z\"] - mrc1.orig[\"z\"]\r\n",
    "                ) / 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(\r\n",
    "        \"#MAP AVE={ave} STD={std} STD_norm={std_norm}\".format(\r\n",
    "            ave=mrc2.ave, std=mrc2.std, std_norm=mrc2.std_norm_ave\r\n",
    "        )\r\n",
    "    )\r\n",
    "    # return False\r\n",
    "    return mrc2"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "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": 13,
   "source": [
    "@numba.jit(nopython=True)\r\n",
    "def rot_pos_mtx(mtx, vec):\r\n",
    "    mtx = mtx.astype(np.float64)\r\n",
    "    vec = vec.astype(np.float64)\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": 106,
   "source": [
    "def rot_mrc(orig_mrc_data, orig_mrc_vec, new_vec_shape, angle):\r\n",
    "    new_pos = np.array(\r\n",
    "        np.meshgrid(\r\n",
    "            np.arange(new_vec_shape[0]),\r\n",
    "            np.arange(new_vec_shape[1]),\r\n",
    "            np.arange(new_vec_shape[2]),\r\n",
    "        )\r\n",
    "    ).T.reshape(-1, 3)\r\n",
    "\r\n",
    "    # set the center\r\n",
    "    cent = int(new_vec_shape[0] / 2.0)\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",
    "    \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",
    "    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",
    "    combined_arr = combined_arr.astype(np.int32)\r\n",
    "\r\n",
    "    #print(combined_arr)\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",
    "\r\n",
    "    # print(np.count_nonzero(orig_mrc_data))\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_dens = orig_mrc_vec[\r\n",
    "        non_zero_rot_list[:, 0], non_zero_rot_list[:, 1], non_zero_rot_list[:, 2]\r\n",
    "    ]\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_dens[:, 0:3])\r\n",
    "\r\n",
    "    # print(new_vec)\r\n",
    "\r\n",
    "    # init new vec array\r\n",
    "    new_vec_array = np.zeros((new_vec_shape[0], new_vec_shape[1], new_vec_shape[2], 3))\r\n",
    "\r\n",
    "    # print(new)\r\n",
    "\r\n",
    "    # fill in the new data\r\n",
    "    for vec, ind in zip(new_vec, (non_zero_rot_list[:, 3:6] + cent).astype(int)):\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",
    "\r\n",
    "    return new_vec_array"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "source": [
    "def find_best_trans(x, y, z):\r\n",
    "\r\n",
    "    xyz_arr = x + y + z\r\n",
    "    best = np.amax(xyz_arr)\r\n",
    "    trans = np.unravel_index(xyz_arr.argmax(), xyz_arr.shape)\r\n",
    "\r\n",
    "    return best, trans"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "source": [
    "def search_map_fft(mrc_target, mrc_search, TopN=10, ang=30, 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",
    "    #     x1 = np.swapaxes(x1 , 0, 2)\r\n",
    "    #     y1 = np.swapaxes(y1 , 0, 2)\r\n",
    "    #     z1 = np.swapaxes(z1 , 0, 2)\r\n",
    "\r\n",
    "    #     x2 = copy.deepcopy(mrc_search.vec[:, :, :, 0])\r\n",
    "\r\n",
    "    d3 = mrc_target.xdim ** 3\r\n",
    "\r\n",
    "    rd3 = 1.0 / d3\r\n",
    "\r\n",
    "    # print(f'{rd3:.20f}')\r\n",
    "\r\n",
    "    X1 = np.fft.rfftn(x1)\r\n",
    "    X1 = np.conj(X1)\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",
    "    #     X2 = np.fft.rfftn(x2)\r\n",
    "\r\n",
    "    #     X12 = np.multiply(X1,X2)\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",
    "    mrc_angle_dict = {}\r\n",
    "\r\n",
    "    with concurrent.futures.ThreadPoolExecutor(\r\n",
    "        max_workers=os.cpu_count() + 4\r\n",
    "    ) as executor:\r\n",
    "        trans_vec = {\r\n",
    "            executor.submit(\r\n",
    "                rot_mrc,\r\n",
    "                mrc_search.data,\r\n",
    "                mrc_search.vec,\r\n",
    "                (mrc_search.xdim, mrc_search.ydim, mrc_search.zdim),\r\n",
    "                angle,\r\n",
    "            ): angle\r\n",
    "            for angle in angle_comb\r\n",
    "        }\r\n",
    "        for future in concurrent.futures.as_completed(trans_vec):\r\n",
    "            angle = trans_vec[future]\r\n",
    "            mrc_angle_dict[tuple(angle)] = future.result()\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",
    "    #             (mrc_search.xdim, mrc_search.ydim, mrc_search.zdim),\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(\r\n",
    "        (a.shape[0], a.shape[1], a.shape[2] // 2 + 1), dtype=\"complex64\"\r\n",
    "    )\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(\r\n",
    "        b, c, direction=\"FFTW_BACKWARD\", axes=(0, 1, 2), normalise_idft=False\r\n",
    "    )\r\n",
    "\r\n",
    "    angle_score = []\r\n",
    "\r\n",
    "    XX = []\r\n",
    "\r\n",
    "    for angle in tqdm(angle_comb, desc=\"FFT\"):\r\n",
    "        rot_mrc_vec = mrc_angle_dict[tuple(angle)]\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",
    "\r\n",
    "        #         x2 = np.swapaxes(x2 , 0, 2)\r\n",
    "        #         y2 = np.swapaxes(y2 , 0, 2)\r\n",
    "        #         z2 = np.swapaxes(z2 , 0, 2)\r\n",
    "\r\n",
    "        X2 = np.zeros_like(X1)\r\n",
    "        np.copyto(a, x2)\r\n",
    "        np.copyto(X2, fft_object(a))\r\n",
    "        X12 = X1 * X2\r\n",
    "        np.copyto(b, X12)\r\n",
    "        x12 = np.zeros_like(x1)\r\n",
    "        np.copyto(x12, ifft_object(b))\r\n",
    "\r\n",
    "        Y2 = np.zeros_like(Y1)\r\n",
    "        np.copyto(a, y2)\r\n",
    "        np.copyto(Y2, fft_object(a))\r\n",
    "        Y12 = Y1 * Y2\r\n",
    "        np.copyto(b, Y12)\r\n",
    "        y12 = np.zeros_like(y1)\r\n",
    "        np.copyto(y12, ifft_object(b))\r\n",
    "\r\n",
    "        Z2 = np.zeros_like(Z1)\r\n",
    "        np.copyto(a, z2)\r\n",
    "        np.copyto(Z2, fft_object(a))\r\n",
    "        Z12 = Z1 * Z2\r\n",
    "        np.copyto(b, Z12)\r\n",
    "        z12 = np.zeros_like(z1)\r\n",
    "        np.copyto(z12, 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",
    "        best, trans = find_best_trans(x12, y12, z12)\r\n",
    "\r\n",
    "        angle_score.append([tuple(angle), best * rd3, trans])\r\n",
    "\r\n",
    "    #     num_jobs = math.ceil(360 / ang) * math.ceil(360 / ang) * (180 // ang + 1)\r\n",
    "\r\n",
    "    return angle_score"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "mrc1 = mrc_obj(\"./data/emd_8097.mrc\")\r\n",
    "mrc2 = mrc_obj(\"./data/ChainA_simulated_resample.mrc\")\r\n",
    "\r\n",
    "mrc1, mrc_N1 = mrc_set_vox_size(mrc1)\r\n",
    "mrc2, mrc_N2 = mrc_set_vox_size(mrc2)\r\n",
    "\r\n",
    "if mrc_N1.xdim > mrc_N2.xdim:\r\n",
    "    mrc_N2.xdim = mrc_N2.ydim = mrc_N2.zdim = mrc_N1.xdim\r\n",
    "\r\n",
    "    mrc_N2.orig[\"x\"] = mrc_N2.cent[0] - 0.5 * 7 * mrc_N2.xdim\r\n",
    "    mrc_N2.orig[\"y\"] = mrc_N2.cent[1] - 0.5 * 7 * mrc_N2.xdim\r\n",
    "    mrc_N2.orig[\"z\"] = mrc_N2.cent[2] - 0.5 * 7 * mrc_N2.xdim\r\n",
    "\r\n",
    "else:\r\n",
    "    mrc_N1.xdim = mrc_N1.ydim = mrc_N1.zdim = mrc_N2.xdim\r\n",
    "\r\n",
    "    mrc_N1.orig[\"x\"] = mrc_N1.cent[0] - 0.5 * 7 * mrc_N1.xdim\r\n",
    "    mrc_N1.orig[\"y\"] = mrc_N1.cent[1] - 0.5 * 7 * mrc_N1.xdim\r\n",
    "    mrc_N1.orig[\"z\"] = mrc_N1.cent[2] - 0.5 * 7 * mrc_N1.xdim"
   ],
   "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=128.68973561744724\n",
      "Nvox= 48, 48, 48\n",
      "cent= 118.23501896858215, 107.53501725196838, 102.1850163936615\n",
      "ori= -49.76498103141785, -60.464982748031616, -65.8149836063385\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "%%time\r\n",
    "mrc_N1 = fastVEC(mrc1, mrc_N1)\r\n",
    "# mrc_N2 = fastVEC(mrc2, mrc_N2)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "110592\n",
      "#Start VEC\n",
      "#maxd= 29.9065372581333\n",
      "#fsiv= 0.017889068239927756\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "  0%|          | 0/48 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bed33de87b14765b31d1fa7ff2ecf85",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "32939.518123120644\n",
      "2257\n",
      "#MAP AVE=14.594381091325053 STD=773.4275314049188 STD_norm=342.72184914544823\n",
      "110592\n",
      "#Start VEC\n",
      "#maxd= 29.9065372581333\n",
      "#fsiv= 0.017889068239927756\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "  0%|          | 0/48 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b65b206a9984a6da73182be6696a596",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#End LDP\n",
      "77063.58651876262\n",
      "557\n",
      "#MAP AVE=138.35473342686288 STD=3523.8520239494 STD_norm=1324.9230607740571\n",
      "Wall time: 9.35 s\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "source": [
    "%%time\n",
    "mrc = rot_mrc(mrc_N2.data, mrc_N2.vec, (48, 48, 48), [0, 0, 30])"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Wall time: 16 ms\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "source": [
    "%%time\n",
    "score_list = search_map_fft(mrc_N1, mrc_N2, ang=30)"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "FFT:   0%|          | 0/1008 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f8bba9af9002429f925397d437c1fd45",
       "version_major": 2,
       "version_minor": 0
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Wall time: 13.9 s\n"
     ]
    }
   ],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "source": [
    "[score for score in score_list if score[0] == (0, 0, 30)]"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "[[(0, 0, 30), 67.83423303674768, (9, 9, 45)]]"
      ]
     },
     "metadata": {},
     "execution_count": 19
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "source": [
    "sorted(score_list, key=lambda x: x[1], reverse=True)[:10]"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "[[(90, 30, 150), 146.59449146412035, (37, 11, 44)],\n",
       " [(270, 150, 150), 130.95534939236111, (12, 37, 44)],\n",
       " [(270, 150, 0), 114.34474464699073, (33, 7, 45)],\n",
       " [(90, 30, 180), 114.34474464699073, (33, 7, 45)],\n",
       " [(90, 30, 0), 111.62222403067129, (15, 40, 45)],\n",
       " [(120, 0, 90), 110.02421513310185, (9, 8, 47)],\n",
       " [(300, 150, 0), 108.50299298321758, (34, 6, 42)],\n",
       " [(120, 30, 180), 108.50299298321758, (34, 6, 42)],\n",
       " [(270, 150, 180), 105.15455005787037, (15, 39, 44)],\n",
       " [(300, 180, 0), 104.73916739004629, (41, 5, 47)]]"
      ]
     },
     "metadata": {},
     "execution_count": 110
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "xx12 = np.fft.irfftn(XX12, norm=\"forward\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "xx12[0][0][1]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "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, (48,48,48), [0,0,30])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%lprun -f search_map_fft search_map_fft(mrc_N1, mrc_N2, ang=60)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# x1 = copy.deepcopy(mrc_N1.vec[:, :, :, 0])\n",
    "# y1 = copy.deepcopy(mrc_N1.vec[:, :, :, 1])\n",
    "\n",
    "# a = pyfftw.empty_aligned((x1.shape), dtype=\"float32\")\n",
    "# b = pyfftw.empty_aligned(\n",
    "#     (a.shape[0], a.shape[1], a.shape[2] // 2 + 1), dtype=\"complex64\"\n",
    "# )\n",
    "# a = x1\n",
    "# b = np.zeros(b.shape, dtype=\"complex64\")\n",
    "\n",
    "# fft_object = pyfftw.FFTW(a, b, axes=(0, 1, 2))\n",
    "\n",
    "# fft_object(a)\n",
    "\n",
    "# a = y1\n",
    "# fft_object(a)"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:bio]",
   "language": "python",
   "name": "conda-env-bio-py"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}