Qscore1_combined.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ce892245-b876-495d-a101-e4d277c27a85",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import iotbx.pdb\n",
    "from scitbx.array_family import flex\n",
    "from cctbx import crystal\n",
    "import cctbx\n",
    "import math\n",
    "from iotbx.data_manager import DataManager\n",
    "from itertools import chain\n",
    "from collections import defaultdict\n",
    "from scipy.spatial import KDTree \n",
    "import sys\n",
    "import numpy as np\n",
    "sys.path.append(\"../\")\n",
    "from qscore_utils import *\n",
    "# dm = DataManager()\n",
    "# dm.process_model_file(\"../data/1yjp.pdb\")\n",
    "# model = dm.get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "62a0ff67-c1d0-4336-950a-4cec63a5c057",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from qscore_combined import Qscore1_combined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a74f8b92-dce2-451d-b450-7a8bb9cd3e89",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# from scipy.spatial.distance import cdist\n",
    "# def nd_to_1d_indices(indices, shape):\n",
    "#     # Normalize indices to always use slice objects\n",
    "#     normalized_indices = []\n",
    "#     for dim, idx in enumerate(indices):\n",
    "#         if idx is None:\n",
    "#             normalized_indices.append(slice(0,shape[dim]))\n",
    "#         else:\n",
    "#             normalized_indices.append(idx)\n",
    "    \n",
    "#     # If any index is a slice, recursively call function for each value in slice\n",
    "#     for dim, (i, s) in enumerate(zip(normalized_indices, shape)):\n",
    "#         if isinstance(i, slice):\n",
    "#             result_indices = []\n",
    "#             start, stop, step = i.indices(s)\n",
    "#             for j in range(start, stop, step):\n",
    "#                 new_indices = list(normalized_indices)\n",
    "#                 new_indices[dim] = j\n",
    "#                 result_indices.extend(nd_to_1d_indices(new_indices, shape))\n",
    "#             return result_indices\n",
    "    \n",
    "#     # If no slices, calculate single 1D index\n",
    "#     index = 0\n",
    "#     stride = 1\n",
    "#     for i, dim in reversed(list(zip(normalized_indices, shape))):\n",
    "#         index += i * stride\n",
    "#         stride *= dim\n",
    "#     return [index]\n",
    "\n",
    "# def optimized_nd_to_1d_indices(i, shape):\n",
    "#     # For fixed input of (None, i, None), we directly compute based on given structure\n",
    "#     result_indices = []\n",
    "    \n",
    "#     # Pre-compute for 1st dimension which is always a slice\n",
    "#     start1, stop1 = 0, shape[0]\n",
    "    \n",
    "#     # Pre-compute for 3rd dimension which is always a slice\n",
    "#     start3, stop3 = 0, shape[2]\n",
    "#     stride3 = 1\n",
    "    \n",
    "#     # Directly compute for 2nd dimension which is variable\n",
    "#     stride2 = shape[2]\n",
    "#     index2 = i * stride2 * shape[0]\n",
    "    \n",
    "#     for val1 in range(start1, stop1):\n",
    "#         for val3 in range(start3, stop3):\n",
    "#             result_indices.append(val1 * stride2 + index2 + val3 * stride3)\n",
    "            \n",
    "#     return result_indices\n",
    "\n",
    "\n",
    "\n",
    "# def find_neighbors(A, B, radius):\n",
    "#     \"\"\"\n",
    "#     Find neighbors within a given radius for each point in set B from set A.\n",
    "\n",
    "#     Parameters:\n",
    "#     - A (np.array): Set of 3D points of shape (M, 3).\n",
    "#     - B (np.array): Set of 3D points of shape (N, 3).\n",
    "#     - radius (float): The query radius.\n",
    "\n",
    "#     Returns:\n",
    "#     - neighbors (list of lists): For each point in B, a list of its neighbors' indices in A.\n",
    "#     - distances (list of lists): For each point in B, a list of distances to its neighbors in A.\n",
    "#     \"\"\"\n",
    "    \n",
    "#     # Calculate the distance matrix\n",
    "#     d_matrix = cdist(B, A)\n",
    "    \n",
    "#     # Find neighbors and their distances for each point in B\n",
    "#     neighbors = [list(np.where(row < radius)[0]) for row in d_matrix]\n",
    "#     distances = [list(row[row < radius]) for row in d_matrix]\n",
    "    \n",
    "#     return neighbors, distances\n",
    "\n",
    "\n",
    "# def query_ball_point(points, query_point, radius, return_length=True):\n",
    "#     \"\"\"\n",
    "#     For a given query point, find all points within a given radius.\n",
    "\n",
    "#     Parameters:\n",
    "#     - points (np.array): Set of 3D points of shape (N, 3).\n",
    "#     - query_point (np.array): The 3D query point of shape (1, 3).\n",
    "#     - radius (float): The query radius.\n",
    "#     - return_length (bool, optional): If True, return only the count of neighbors.\n",
    "\n",
    "#     Returns:\n",
    "#     - indices (list): If return_length=False, indices of points within the radius. Else, count of such points.\n",
    "#     \"\"\"\n",
    "    \n",
    "#     # Calculate the distance from all points to the query point\n",
    "#     distances = cdist(points, [query_point]).flatten()\n",
    "    \n",
    "#     # Find indices of points within the radius\n",
    "#     indices = list(np.where(distances < radius)[0])\n",
    "    \n",
    "#     if return_length:\n",
    "#         return len(indices)\n",
    "#     else:\n",
    "#         return indices\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74bc0c5b-1fc3-4568-b35b-3e6716ffb9ca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b983ab1a-d1b4-4976-929b-87ff36afec20",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef8633ad-e905-40d0-a9ce-5c3900a79819",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from multiprocessing import Pool, cpu_count\n",
    "# def radial_shell_combined_mp(model,n_probes=64,n_probes_target=8,radii=None,rtol=0.9,num_processes=cpu_count(),selection=None):\n",
    "\n",
    "#     # Create argument tuples for each chunk\n",
    "#     args = [(i,model,n_probes,n_probes_target,radius_shell,rtol,selection) for i,radius_shell in enumerate(radii)]\n",
    "    \n",
    "#     # Create a pool of worker processes\n",
    "#     if num_processes >1:\n",
    "#         with Pool(num_processes) as p:\n",
    "#             # Use the pool to run the trilinear_interpolation_worker function in parallel\n",
    "#             results = p.map(radial_shell_worker_combined, args)\n",
    "#     else:\n",
    "#         results = []\n",
    "#         for arg in args:\n",
    "#             result = radial_shell_worker_ccombined(arg)\n",
    "#             results.append(result)\n",
    "\n",
    "    \n",
    "#     # stack numpy\n",
    "#     probe_xyz_all = [result[0] for result in results]\n",
    "#     keep_mask_all = [result[2] for result in results]\n",
    "    \n",
    "\n",
    "#     n_shells = len(radii)\n",
    "#     n_atoms = probe_xyz_all[0].shape[0]\n",
    "#     out_shape = (n_shells,n_atoms,n_probes,3 )\n",
    "#     out_size = np.prod(out_shape)\n",
    "#     shell_size = np.prod(out_shape[1:])\n",
    "#     out_probes = np.full((n_shells,n_atoms,n_probes_target,3),-1.0)\n",
    "#     out_mask = np.full((n_shells,n_atoms,n_probes_target),False)\n",
    "    \n",
    "#     for i,p in enumerate(probe_xyz_all):\n",
    "#         out_probes[i,:,:n_probes_target,:] =p[:,:n_probes_target]\n",
    "\n",
    "    \n",
    "#     for i,k in enumerate(keep_mask_all):\n",
    "#         start = i*shell_size\n",
    "#         stop = start+shell_size\n",
    "#         out_mask[i] = k[:,:n_probes_target]\n",
    "#     probe_xyz_np = out_probes\n",
    "#     keep_mask_np = out_mask\n",
    "    \n",
    "#     # stack flex\n",
    "#     probe_xyz_all = [result[1] for result in results]\n",
    "#     keep_mask_all = [result[3] for result in results]\n",
    "    \n",
    "    \n",
    "#     n_shells = len(radii)\n",
    "#     n_atoms = probe_xyz_all[0].focus()[0]\n",
    "#     out_shape = (n_shells,n_atoms,n_probes_target,3 )\n",
    "#     out_size = math.prod(out_shape)\n",
    "#     shell_size = math.prod(out_shape[1:])\n",
    "#     out_probes = flex.double(out_size,-1.0)\n",
    "#     out_mask = flex.bool(n_atoms*n_shells*n_probes_target,False)\n",
    "#     for i,p in enumerate(probe_xyz_all):\n",
    "#         start = i*shell_size\n",
    "#         stop = start+shell_size\n",
    "#         out_probes = out_probes.set_selected(flex.uint32_range(start,stop),p.as_1d())\n",
    "#     out_probes.reshape(flex.grid(*out_shape))\n",
    "\n",
    "    \n",
    "#     for i,k in enumerate(keep_mask_all):\n",
    "#         start = i*(n_atoms*n_probes_target)\n",
    "#         stop = start+(n_atoms*n_probes_target)\n",
    "#         out_mask = out_mask.set_selected(flex.uint32_range(start,stop),k.as_1d())\n",
    "#     out_mask.reshape(flex.grid(n_shells,n_atoms,n_probes_target))\n",
    "    \n",
    "#     probe_xyz_flex = out_probes\n",
    "#     keep_mask_flex = out_mask\n",
    "    \n",
    "#     return probe_xyz_np, probe_xyz_flex, keep_mask_np, keep_mask_flex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c8d54193-e34b-4686-9f7c-37085d313870",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dm = DataManager()\n",
    "dm.process_model_file(\"../data/6cvm.cif\")\n",
    "dm.process_real_map_file(\"../data/emd_7770.map\")\n",
    "mmm = dm.get_map_model_manager()\n",
    "assert flex.sum(mmm.map_manager().density_at_sites_cart(mmm.model().get_sites_cart()))>0,(\n",
    "    \"Negative or zero density around atoms. Map/model likely not aligned\")\n",
    "model = mmm.model()\n",
    "mm = mmm.map_manager()\n",
    "M = mm.map_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4eedc8ea-a1f0-4886-bbbb-def018ae620b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from qscore_combined import radial_shell_combined_mp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "def67a54-6a08-4a1e-916d-e106d7535064",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18d6fda37def471c903d714c57aa6fdb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "model = mmm.model()\n",
    "mm = mmm.map_manager()\n",
    "M = mm.map_data()\n",
    "radii = np.linspace(0,2,16)\n",
    "n_probes = 8\n",
    "selection = model.selection(\"chain A and resseq 1:10\")\n",
    "num_processes = 1\n",
    "rtol = 0.9\n",
    "(probe_xyz, \n",
    " probe_xyz_cctbx, \n",
    " keep_mask, \n",
    " keep_mask_cctbx) = radial_shell_combined_mp(model,\n",
    "                                              n_probes=n_probes,\n",
    "                                              radii=radii,\n",
    "                                              rtol = rtol,\n",
    "                                              selection=selection,\n",
    "                                              num_processes=num_processes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0adca75d-0102-422b-b7e3-ad9734958d2d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef1bcd9e10f34c6f8a3eee65c63502f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "NameError",
     "evalue": "name 'probe_xyz_flex' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m Qscore1_combined(mmm,num_processes\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "File \u001b[0;32m~/Projects/Qscore/qscore_standalone/qscore_combined.py:303\u001b[0m, in \u001b[0;36mQscore1_combined\u001b[0;34m(mmm, voxel_size, n_probes, min_probes, radii, rtol, selection, num_processes)\u001b[0m\n\u001b[1;32m    300\u001b[0m \u001b[38;5;66;03m# init cctbx arrays\u001b[39;00m\n\u001b[1;32m    302\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m do_flex:\n\u001b[0;32m--> 303\u001b[0m     n_shells,n_atoms,n_probes,_ \u001b[38;5;241m=\u001b[39m probe_xyz_flex\u001b[38;5;241m.\u001b[39mfocus()\n\u001b[1;32m    304\u001b[0m     probe_xyz_cctbx \u001b[38;5;241m=\u001b[39m flex\u001b[38;5;241m.\u001b[39mdouble(probe_xyz)\n\u001b[1;32m    305\u001b[0m     keep_mask_cctbx \u001b[38;5;241m=\u001b[39m flex\u001b[38;5;241m.\u001b[39mbool(keep_mask)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'probe_xyz_flex' is not defined"
     ]
    }
   ],
   "source": [
    "Qscore1_combined(mmm,num_processes=1,selection = model.selection(\"chain A and resseq 1:10\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8ea0b54-f780-4184-b499-aacb646400aa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "6e3f5707-964c-4bcb-845d-231f96229ec1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mask_clash=True\n",
    "voxel_size=1.0\n",
    "n_probes=8\n",
    "min_probes=1\n",
    "radii=np.arange(0.1,2.1,0.1)\n",
    "rtol=0.9\n",
    "ignore_min_probes=False\n",
    "selection = model.selection(\"chain A and resseq 1:10\")\n",
    "num_processes = 6\n",
    "\n",
    "\n",
    "\n",
    "probe_xyz_np, probe_xyz_flex, keep_mask_np, keep_mask_flex = radial_shell_combined_mp(model,radii=radii,n_probes=16,n_probes_target=8,selection=selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "610318d4-ac0f-44ac-8ad1-87e408320c2a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert np.all(np.isclose(probe_xyz_np, probe_xyz_flex.as_numpy_array()))\n",
    "assert np.all(keep_mask_np==keep_mask_flex.as_numpy_array())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f5c0d2d7-f8e9-468a-9b75-a01c2c7f4dff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "probe_xyz = probe_xyz_np\n",
    "keep_mask = keep_mask_np\n",
    "probe_xyz_cctbx = probe_xyz_flex\n",
    "keep_mask_cctbx = keep_mask_flex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "75415f76-af8b-4144-8205-0b128c26fa8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "do_np = True\n",
    "do_flex = True\n",
    "do_test = False\n",
    "\n",
    "# to test must do both\n",
    "if do_test:\n",
    "    do_np = True\n",
    "    do_flex = True\n",
    "# PROCEED\n",
    "\n",
    "\n",
    "\n",
    "# flatten numpy arrays\n",
    "if do_np:\n",
    "    n_shells,n_atoms,n_probes,_ = probe_xyz.shape\n",
    "    probe_xyz_flat = probe_xyz.reshape((n_atoms*n_shells*n_probes,3))\n",
    "    keep_mask_flat = keep_mask.reshape(-1) # (n_shells*n_atoms*n_probes,)\n",
    "    \n",
    "\n",
    "# init cctbx arrays\n",
    "\n",
    "if do_flex:\n",
    "    n_shells,n_atoms,n_probes,_ = probe_xyz_flex.focus()\n",
    "    probe_xyz_cctbx = flex.double(probe_xyz)\n",
    "    keep_mask_cctbx = flex.bool(keep_mask)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# APPLY MASK BEFORE INTERPOLATION\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    masked_probe_xyz_flat = probe_xyz_flat[keep_mask_flat]\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    keep_mask_cctbx_fullflat = []\n",
    "\n",
    "    for val in keep_mask_cctbx:\n",
    "        for _ in range(3):  # since A has an additional dimension of size 3\n",
    "            keep_mask_cctbx_fullflat.append(val)\n",
    "\n",
    "    mask = flex.bool(keep_mask_cctbx_fullflat)\n",
    "    #indices = flex.int([i for i in range(1, keep_mask_cctbx.size() + 1) for _ in range(3)])\n",
    "    sel = probe_xyz_cctbx.select(mask)\n",
    "    #sel_indices = indices.select(mask)\n",
    "    masked_probe_xyz_flat_cctbx = flex.vec3_double(sel)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "2b2076fc-7da3-448d-917e-e3e14d859416",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# INTERPOLATE\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    M = mm.map_data()\n",
    "    volume = M.as_numpy_array()\n",
    "    voxel_size = np.array(mm.pixel_sizes())\n",
    "    masked_density = trilinear_interpolation(volume, masked_probe_xyz_flat, voxel_size=voxel_size)\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    masked_density_cctbx = mm.density_at_sites_cart(masked_probe_xyz_flat_cctbx)\n",
    "\n",
    "    # test equivalent\n",
    "    if do_test:\n",
    "        assert np.all(np.isclose(masked_density,masked_density_cctbx.as_numpy_array()))\n",
    "\n",
    "\n",
    "# reshape interpolated values to (n_shells,n_atoms, n_probes)\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    d_vals = np.zeros((n_shells, n_atoms, n_probes))\n",
    "    d_vals[keep_mask] = masked_density\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    keep_mask_cctbx.reshape(flex.grid(n_shells*n_atoms*n_probes))\n",
    "    d_vals_cctbx = flex.double(keep_mask_cctbx.size(),0.0)\n",
    "    d_vals_cctbx = d_vals_cctbx.set_selected(keep_mask_cctbx,masked_density_cctbx)\n",
    "    d_vals_cctbx.reshape(flex.grid(n_shells,n_atoms,n_probes))\n",
    "\n",
    "    if do_test:\n",
    "        # test\n",
    "        assert np.all(np.isclose(d_vals,d_vals_cctbx.as_numpy_array()))\n",
    "\n",
    "\n",
    "# reshape to (M,N*L) for rowwise correlation\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    d_vals_2d = d_vals.transpose(1,0,2).reshape(d_vals.shape[1], -1)\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    def custom_reshape_indices(flex_array):\n",
    "        N,M,L = flex_array.focus()\n",
    "        result = flex.double(flex.grid(M, N * L))\n",
    "\n",
    "        for i in range(N):\n",
    "            for j in range(M):\n",
    "                for k in range(L):\n",
    "                    # Calculate the original flat index\n",
    "                    old_index = i * M * L + j * L + k\n",
    "                    # Calculate the new flat index after transpose and reshape\n",
    "                    new_index = j * N * L + i * L + k\n",
    "                    result[new_index] = flex_array[old_index]\n",
    "\n",
    "        return result\n",
    "\n",
    "    d_vals_2d_cctbx = custom_reshape_indices(d_vals_cctbx)\n",
    "\n",
    "\n",
    "    # test\n",
    "    if do_test:\n",
    "        assert np.all(np.isclose(d_vals_2d,d_vals_2d_cctbx.as_numpy_array()))\n",
    "\n",
    "\n",
    "# create the reference data\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    M = mm.map_data().as_numpy_array()\n",
    "    maxD = min(M.mean()+M.std()*10,M.max())\n",
    "    minD = max(M.mean()-M.std()*1,M.min())\n",
    "    A = maxD-minD\n",
    "    B = minD\n",
    "    u = 0\n",
    "    sigma = 0.6\n",
    "    x = np.array(radii)\n",
    "    y = A * np.exp(-0.5*((x-u)/sigma)**2) + B \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#cctbx\n",
    "if do_flex:\n",
    "    M = mm.map_data()\n",
    "    maxD_cctbx = min(flex.mean(M)+M.standard_deviation_of_the_sample()*10,flex.max(M))\n",
    "    minD_cctbx = max(flex.mean(M)-M.standard_deviation_of_the_sample()*1,flex.min(M))\n",
    "    A_cctbx = maxD_cctbx-minD_cctbx\n",
    "    B_cctbx = minD_cctbx\n",
    "    u = 0\n",
    "    sigma = 0.6\n",
    "    x = flex.double(radii)\n",
    "    y_cctbx = A_cctbx * flex.exp(-0.5*((flex.double(x)-u)/sigma)**2) + B_cctbx\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    # test\n",
    "    if do_test:\n",
    "        assert np.all(np.isclose(np.array(y_cctbx),y))\n",
    "\n",
    "\n",
    "# Stack and reshape data for correlation calc\n",
    "\n",
    "# numpy\n",
    "if do_np:\n",
    "    # stack the reference to shape (n_shells,n_atoms,n_probes)\n",
    "    g_vals = np.repeat(y[:,None],n_probes,axis=1)\n",
    "    g_vals = np.expand_dims(g_vals,1)\n",
    "    g_vals = np.tile(g_vals,(n_atoms,1))\n",
    "\n",
    "\n",
    "    # reshape\n",
    "    g_vals_2d = g_vals.transpose(1,0,2).reshape(g_vals.shape[1], -1)\n",
    "    d_vals_2d = d_vals.transpose(1,0,2).reshape(d_vals.shape[1], -1)\n",
    "    mask_2d = keep_mask.transpose(1,0,2).reshape(keep_mask.shape[1], -1)\n",
    "\n",
    "\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    # 1. Repeat y for n_probes (equivalent to np.repeat)\n",
    "    g_vals_cctbx = [[val] * n_probes for val in y_cctbx]\n",
    "\n",
    "    # 2. Add a new dimension (equivalent to np.expand_dims)\n",
    "    g_vals_expanded = [[item] for item in g_vals_cctbx]\n",
    "\n",
    "    # 3. Tile for each atom (equivalent to np.tile)\n",
    "    g_vals_tiled = []\n",
    "    for item in g_vals_expanded:\n",
    "        g_vals_tiled.append(item * n_atoms)\n",
    "\n",
    "\n",
    "    g_vals_cctbx = flex.double(np.array(g_vals_tiled) )\n",
    "\n",
    "\n",
    "    # test \n",
    "    if do_test:\n",
    "        assert np.all(np.isclose(g_vals_cctbx.as_numpy_array(),g_vals))\n",
    "\n",
    "\n",
    "# # CALCULATE Q\n",
    "\n",
    "# # numpy\n",
    "if do_np:\n",
    "    q = rowwise_corrcoef(g_vals_2d,d_vals_2d,mask=mask_2d)\n",
    "\n",
    "\n",
    "\n",
    "# cctbx\n",
    "if do_flex:\n",
    "    d_vals_cctbx = d_vals_cctbx.as_1d()\n",
    "    g_vals_cctbx = g_vals_cctbx.as_1d()\n",
    "    keep_mask_cctbx_double = keep_mask_cctbx.as_1d().as_double()\n",
    "    q_cctbx = []\n",
    "    for atomi in range(n_atoms):\n",
    "\n",
    "        #inds = nd_to_1d_indices((None,atomi,None),(n_shells,n_atoms,n_probes))\n",
    "        inds = optimized_nd_to_1d_indices(atomi,(n_shells,n_atoms,n_probes))\n",
    "        inds = flex.uint32(inds)\n",
    "        d_row = d_vals_cctbx.select(inds)\n",
    "        \n",
    "        if do_test:\n",
    "            assert np.all(np.isclose(d_row.as_numpy_array(),d_vals_2d[atomi]))\n",
    "\n",
    "        g_row = g_vals_cctbx.select(inds)\n",
    "        if do_test:\n",
    "            assert np.all(np.isclose(g_row.as_numpy_array(),g_vals_2d[atomi]))\n",
    "\n",
    "        mask = keep_mask_cctbx.select(inds)\n",
    "        if do_test:\n",
    "            assert np.all(np.isclose(mask.as_numpy_array(),mask_2d[atomi]))\n",
    "\n",
    "        d = d_row.select(mask)\n",
    "        g = g_row.select(mask)\n",
    "        qval = flex.linear_correlation(d,g).coefficient()\n",
    "        q_cctbx.append(qval)\n",
    "    \n",
    "\n",
    "    q_cctbx = flex.double(q_cctbx)\n",
    "    if do_test:\n",
    "        assert np.all(np.isclose(q,np.array(q_cctbx)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e8f140b8-6d1b-485c-bc49-2494865d62e1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1.64840792e-17,  2.45751023e-17,  0.00000000e+00,  5.99714904e-18,\n",
       "        0.00000000e+00,  5.81031901e-17,  3.71147537e-17, -7.09197339e-17,\n",
       "        3.77665092e-17,  2.08951328e-17, -6.37055373e-17,  3.74357560e-17,\n",
       "        5.81556119e-18,  0.00000000e+00, -9.36591433e-18,  1.23085050e-17,\n",
       "        0.00000000e+00, -2.01572320e-17,  4.67672103e-17,  1.34623744e-17,\n",
       "       -2.99346027e-17, -1.77864806e-17,  2.93908615e-17, -3.90477276e-17,\n",
       "        0.00000000e+00, -2.41990440e-17,  0.00000000e+00,  2.71549486e-18,\n",
       "        7.40005440e-17,  0.00000000e+00,  0.00000000e+00, -2.11092161e-17,\n",
       "       -1.31652185e-17,  4.53198060e-17,  1.10948691e-17,  0.00000000e+00,\n",
       "        1.04422268e-17,  0.00000000e+00,  3.56932413e-17,  2.18366163e-17,\n",
       "        0.00000000e+00, -4.20066059e-17,  5.98342913e-17, -5.01339265e-17,\n",
       "       -1.32897009e-17, -3.91604854e-17,  2.44047688e-17,  0.00000000e+00,\n",
       "       -1.36006992e-17,  0.00000000e+00, -3.53883341e-17, -2.05801223e-17,\n",
       "        2.32039736e-17, -1.88671411e-17,  2.01010301e-17, -5.23683555e-17,\n",
       "        2.55174589e-17, -3.81943229e-17,  8.53420669e-18,  3.79204985e-17,\n",
       "        3.60934780e-17,  0.00000000e+00, -3.90623907e-18,  3.58268366e-17])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3b19c4-c740-4445-ac48-719cbdbd703a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}