QScore.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "52a244d5-1552-403e-831e-c1b3c8d4c7e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from iotbx.data_manager import DataManager\n",
    "from cctbx.array_family import flex\n",
    "import numpy as np\n",
    "from qscore_utils import (\n",
    "          sphere_points,\n",
    "          #stack_fill,\n",
    "          trilinear_interpolation,\n",
    "          rowwise_corrcoef,\n",
    "          )\n",
    "from tqdm.notebook import tqdm\n",
    "from multiprocessing import Pool\n",
    "import math\n",
    "import cctbx\n",
    "from cctbx.array_family import flex\n",
    "from scitbx_array_family_flex_ext import mean as flex_mean\n",
    "from scitbx_array_family_flex_ext import max as flex_max\n",
    "from scitbx_array_family_flex_ext import min as flex_min\n",
    "\n",
    "def flex_std(array):\n",
    "    return array.standard_deviation_of_the_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8f117d63-fa2e-4ec7-9dd0-154d80b5e32f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from qscore import radial_shell_mp, radial_shell_worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a2c1a82c-2816-4d88-972c-ba223b0541f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_depth(lst, level=1):\n",
    "    \"\"\"Recursively determine the depth of the nested list\"\"\"\n",
    "    if not isinstance(lst, list) or not lst:\n",
    "        return level\n",
    "    return max(get_depth(item, level + 1) for item in lst)\n",
    "\n",
    "def flatten_nested_list(nested_list, target_dim=None):\n",
    "    current_depth = get_depth(nested_list)\n",
    "    \n",
    "    if target_dim is None:\n",
    "        target_dim = 1\n",
    "        \n",
    "    if target_dim <= 0 or target_dim > current_depth:\n",
    "        raise ValueError(f\"Target dimension must be between 1 and {current_depth} (inclusive).\")\n",
    "\n",
    "    flatten_count = current_depth - target_dim\n",
    "    flattened_list = nested_list\n",
    "    for _ in range(flatten_count):\n",
    "        flattened_list = [item for sublist in flattened_list for item in (sublist if isinstance(sublist, list) else [sublist])]\n",
    "    return flattened_list\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d82ab364-2e10-4dc6-873c-9aa021956d56",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sphere_points(ctr, rad, N):\n",
    "    if ctr.ndim==1:\n",
    "        ctr = ctr[None,:]\n",
    "    h = -1.0 + (2.0 * np.arange(N) / float(N-1))[:, np.newaxis]\n",
    "    phis = np.arccos(h)\n",
    "    thetas = np.zeros_like(phis)\n",
    "    thetas[1:-1, :] = (3.6 / np.sqrt(N * (1.0 - h[1:-1]**2))) % (2 * np.pi)\n",
    "    thetas = np.cumsum(thetas, axis=0)\n",
    "\n",
    "    x = np.sin(phis) * np.cos(thetas)\n",
    "    y = np.sin(phis) * np.sin(thetas)\n",
    "    z = np.cos(phis)\n",
    "\n",
    "    # Stack x, y, z to form points and multiply by rad\n",
    "    points = rad * np.stack([x, y, z], axis=-1)\n",
    "\n",
    "    # Reshape points to (1, N, 3)\n",
    "    points = points.reshape(1, N, 3)\n",
    "\n",
    "    # Add center coordinates to all points\n",
    "    # ctr shape: (M, 3), points shape: (1, N, 3)\n",
    "    # Resultant shape: (M, N, 3)\n",
    "    pts = ctr[:, np.newaxis, :] + points\n",
    "\n",
    "    return pts\n",
    "\n",
    "def sphere_points_flex(ctr, rad, N):\n",
    "    assert len(ctr.focus())==2 and ctr.focus()[1]==3, \"Provide an array of cartesian coordinates, shape (N,3)\"\n",
    "    h = -1.0 + (2.0 * np.arange(N) / float(N-1))[:, np.newaxis]\n",
    "    phis = np.arccos(h)\n",
    "    thetas = np.zeros_like(phis)\n",
    "    thetas[1:-1, :] = (3.6 / np.sqrt(N * (1.0 - h[1:-1]**2))) #% (2 * np.pi)\n",
    "    thetas = np.cumsum(thetas, axis=0)\n",
    "\n",
    "    x = np.sin(phis) * np.cos(thetas)\n",
    "    y = np.sin(phis) * np.sin(thetas)\n",
    "    z = np.cos(phis)\n",
    "\n",
    "    # Stack x, y, z to form points and multiply by rad\n",
    "    points = rad * np.stack([x, y, z], axis=-1)\n",
    "\n",
    "    # Reshape points to (1, N, 3)\n",
    "    points = points.reshape(1, N, 3)\n",
    "\n",
    "    # Add center coordinates to all points\n",
    "    # ctr shape: (M, 3), points shape: (1, N, 3)\n",
    "    # Resultant shape: (M, N, 3)\n",
    "    pts = ctr[:, np.newaxis, :] + points\n",
    "\n",
    "    return pts\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7139a8b0-387d-4861-8ba3-830af1fdc3a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def floorf_flex(value, divisor):\n",
    "    # Assuming value is a flex array and divisor is a scalar\n",
    "    if divisor == 0.0:\n",
    "        raise ValueError(\"Cannot divide by zero!\")\n",
    "\n",
    "    quotient = value / divisor\n",
    "\n",
    "    # Getting boolean selectors\n",
    "    pos_sel = quotient >= 0\n",
    "    neg_sel = ~pos_sel \n",
    "\n",
    "    # Applying floor for positive values and ceil for negative values\n",
    "    quotient.set_selected(pos_sel, flex.floor(quotient.select(pos_sel)))\n",
    "    quotient.set_selected(neg_sel, flex.ceil(quotient.select(neg_sel)) - 1)\n",
    "\n",
    "    return quotient\n",
    "\n",
    "\n",
    "    \n",
    "def mod_flex(value, divisor):\n",
    "    # Getting boolean selectors\n",
    "    pos_sel = value >= 0\n",
    "    neg_sel = ~pos_sel \n",
    "    pos_val = value.select(pos_sel)\n",
    "    neg_val = value.select(neg_sel)\n",
    "    pos_quotient = floorf_flex(pos_val,divisor)\n",
    "    pos_out = pos_val-pos_quotient * divisor\n",
    "    neg_quotient = floorf_flex(neg_val,divisor)-1\n",
    "    neg_out = neg_val-neg_quotient * divisor\n",
    "    \n",
    "    out = flex.double(len(value),0.0)\n",
    "    out = out.set_selected(pos_sel,pos_out)\n",
    "    out = out.set_selected(neg_sel,neg_out)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f40a03f3-772f-46ce-a12e-3f6a65dc9471",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sphere_points_combined(ctr_np,ctr_flex,rad,N):\n",
    "\n",
    "    h_np = -1.0 + (2.0 * np.arange(N) / float(N-1))[:, np.newaxis]\n",
    "    h_flex = -1.0 + (2.0 * flex.double_range(N) / (N-1))\n",
    "    #h_flex.reshape(flex.grid(N,1))\n",
    "    assert np.all(np.isclose(h_np.flatten(),h_flex.as_numpy_array()))\n",
    "    phis_flex = flex.acos(h_flex)\n",
    "    phis_np = np.arccos(h_np)\n",
    "\n",
    "    thetas_np = np.zeros_like(phis_np)\n",
    "    thetas_flex = flex.double(len(phis_flex),0.0)\n",
    "    a_flex = (3.6 / flex.sqrt(N * (1.0 - h_flex[1:-1]**2)))\n",
    "    a_np = (3.6 / np.sqrt(N * (1.0 - h_np[1:-1]**2)))\n",
    "    assert np.all(np.isclose(a_np.flatten(),a_flex.as_numpy_array()))\n",
    "\n",
    "    thetas_np[1:-1, :] = a_np\n",
    "    thetas_flex = thetas_flex.set_selected(flex.uint32_range(1,N-1),a_flex)\n",
    "    assert np.all(np.isclose(thetas_np.flatten(),thetas_flex.as_numpy_array()))\n",
    "\n",
    "\n",
    "    thetas_np = np.cumsum(thetas_np,axis=0)\n",
    "\n",
    "    def cumsum_flex(arr):\n",
    "        result = []\n",
    "        running_sum = 0.0\n",
    "        for i,x in enumerate(arr):\n",
    "            running_sum += x\n",
    "            result.append(running_sum)\n",
    "        return flex.double(result)\n",
    "\n",
    "    thetas_flex = cumsum_flex(thetas_flex)\n",
    "\n",
    "    assert np.all(np.isclose(thetas_np.flatten(),thetas_flex.as_numpy_array()))\n",
    "\n",
    "\n",
    "    x_np = np.sin(phis_np) * np.cos(thetas_np)\n",
    "    y_np = np.sin(phis_np) * np.sin(thetas_np)\n",
    "    z_np = np.cos(phis_np)\n",
    "\n",
    "    x_flex = flex.sin(phis_flex) * flex.cos(thetas_flex)\n",
    "    y_flex = flex.sin(phis_flex) * flex.sin(thetas_flex)\n",
    "    z_flex = flex.cos(phis_flex)\n",
    "\n",
    "\n",
    "    # Stack x, y, z to form points and multiply by rad\n",
    "    points_np = rad * np.stack([x_np, y_np, z_np], axis=-1)\n",
    "\n",
    "    points_flex = rad * flex.vec3_double(x_flex,y_flex,z_flex)\n",
    "    \n",
    "    assert np.all(np.isclose(points_np.flatten(),points_flex.as_numpy_array().flatten()))\n",
    "\n",
    "\n",
    "    # put all together\n",
    "\n",
    "    # numpy\n",
    "    points_np = points_np.reshape(1, N, 3)\n",
    "    # add to ctr\n",
    "    points_np = ctr_np[:, np.newaxis, :] + points_np\n",
    "\n",
    "\n",
    "    # flex\n",
    "    def broadcast_add(ctr_flex, points_flex):\n",
    "        N = points_flex.size()\n",
    "        M = ctr_flex.size()\n",
    "        # Preallocate an array of shape (M, N, 3)\n",
    "        result = flex.vec3_double(flex.grid(M, N))\n",
    "\n",
    "        for i in range(M):\n",
    "            center = ctr_flex[i]\n",
    "            for j in range(N):\n",
    "                point = points_flex[j]\n",
    "                result[i, j] = (center[0] + point[0], center[1] + point[1], center[2] + point[2])\n",
    "\n",
    "        result = result.as_1d().as_double()\n",
    "        result.reshape(flex.grid(len(ctr_flex),len(points_flex),3))\n",
    "        return result\n",
    "    \n",
    "    \n",
    "    points_flex = broadcast_add(ctr_flex,points_flex)\n",
    "\n",
    "    assert np.all(np.isclose(points_np,points_flex.as_numpy_array()))\n",
    "    return points_np, points_flex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8fcfc7ea-311e-4a80-94ec-f579412b7355",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ctr_np = np.random.random((100,3))\n",
    "a = sphere_points(ctr_np,2,1000)\n",
    "b, _ = sphere_points_combined(ctr_np,flex.vec3_double(ctr_np),2,1000)\n",
    "assert np.all(np.isclose(a,b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "c1641ce5-7bcb-4303-a3fc-a894ff8fe307",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2f8dc6b0-89ab-414b-a111-ae949e2200bd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "floorf_flex(a_flex,2.0*3.6).as_numpy_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "83a5d730-37f7-4009-a2ae-30ccca1bb05f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos_quotient.as_numpy_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "99921b66-063b-4141-a093-cbb19e3cad73",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_flex.as_numpy_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bb75f0d2-9e4b-4327-bd02-3c91eef1aeb2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "a_np = (3.6 / np.sqrt(N * (1.0 - h_np[1:-1]**2))) % (2*math.pi)\n",
    "a_flex = mod_flex(h_flex,2*math.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d885b23d-5160-4ddf-9435-553fd10ba632",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<scitbx_array_family_flex_ext.double at 0x7fde78341cc0>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "floorf_flex(pos_val,divisor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e81643-6fbe-4ac9-8588-9db65c26e685",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "96bb6fe5-2d4d-4aae-9b83-44eeec86f2b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def __rep(flex_a):\n",
    "    a = flex_a.as_numpy_array()\n",
    "    print(a.__class__.__repr__(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e56a076e-81b0-482b-ad53-cbb524c077fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "1c27a0e6-e323-4b3c-b44a-8f67da1cb427",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = 3.6/flex.sqrt(N * 1.0-h[1:-1]**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "92b8b48b-b176-472f-ae68-3dd5f2613328",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "unsupported operand type(s) for //: 'double' and 'float'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[38], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m manual_modulo(a,\u001b[38;5;241m3.6\u001b[39m)\n",
      "Cell \u001b[0;32mIn[36], line 7\u001b[0m, in \u001b[0;36mmanual_modulo\u001b[0;34m(a, b)\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe divisor (b) cannot be zero.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# Calculate the quotient\u001b[39;00m\n\u001b[0;32m----> 7\u001b[0m q \u001b[38;5;241m=\u001b[39m a \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m b\n\u001b[1;32m      9\u001b[0m \u001b[38;5;66;03m# Calculate the remainder\u001b[39;00m\n\u001b[1;32m     10\u001b[0m remainder \u001b[38;5;241m=\u001b[39m a \u001b[38;5;241m-\u001b[39m (b \u001b[38;5;241m*\u001b[39m q)\n",
      "\u001b[0;31mTypeError\u001b[0m: unsupported operand type(s) for //: 'double' and 'float'"
     ]
    }
   ],
   "source": [
    "manual_modulo(a,3.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "ba2831d0-94c4-4dd5-9227-86f56051bdcc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'double' object has no attribute '_'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[39], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m a\u001b[38;5;241m.\u001b[39m_\u001b[38;5;241m%\u001b[39m_\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'double' object has no attribute '_'"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9686949b-3e40-4a0d-9ca4-03f812d84770",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def manual_modulo(a, b):\n",
    "    # Ensure we don't divide by zero\n",
    "    if b == 0:\n",
    "        raise ValueError(\"The divisor (b) cannot be zero.\")\n",
    "    \n",
    "    # Calculate the quotient\n",
    "    q = a // b\n",
    "\n",
    "    # Calculate the remainder\n",
    "    remainder = a - (b * q)\n",
    "    \n",
    "    return remainder\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "004dcec9-077b-4657-8b57-d5433c72df02",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'slice' object is not iterable",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[26], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmath\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m thetas[\u001b[38;5;241m1\u001b[39m:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m3.6\u001b[39m \u001b[38;5;241m/\u001b[39m flex\u001b[38;5;241m.\u001b[39msqrt(N \u001b[38;5;241m*\u001b[39m (\u001b[38;5;241m1.0\u001b[39m \u001b[38;5;241m-\u001b[39m h[\u001b[38;5;241m1\u001b[39m:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m2\u001b[39m)))\n",
      "\u001b[0;31mTypeError\u001b[0m: 'slice' object is not iterable"
     ]
    }
   ],
   "source": [
    "import math\n",
    "thetas[1:-1] = (3.6 / flex.sqrt(N * (1.0 - h[1:-1]**2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5100d39e-2e89-4b51-af71-9b6a079f0584",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "861f33c1-1522-4e57-8211-a19931a233e1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100,)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " % flex.double(2 * math.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6a828fc8-8a00-4d57-8e77-1d53586ce129",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 1)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4502a411-9bfc-4d5f-be11-a1185bbc7bc4",
   "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": 5,
   "id": "cce29459-cdde-44ac-a880-2c13df50c1b7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'flex_utils'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mflex_utils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m cdist_flex, query_atom_neighbors\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'flex_utils'"
     ]
    }
   ],
   "source": [
    "from flex_utils import cdist_flex, query_atom_neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "200b77f4-ed02-433a-8ce8-2ab3b8ac6e50",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "inds,dists = query_atom_neighbors(model,radius=3.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "34d634f9-6383-4814-aa0e-bb2bc2fe52a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "selection = model.selection(\"chain A and resseq 1:10\").as_numpy_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0ad5e66c-e629-42fb-a5b4-8da6c57ac3fb",
   "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_bool=selection\n",
    "num_processes=1\n",
    "\n",
    "atoms_xyz = model.get_sites_cart().as_numpy_array()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67535fd-1d8c-43ab-8441-9e62ff53d3e3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad81e701-5c5b-4e90-ad2c-5cd3ebbaa93b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1f6426d2-25f7-49a3-bcb5-c5b0d6428c92",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 6, got 5)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m       selection_bool \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mfull(atoms_xyz\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m],\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m      3\u001b[0m atoms_xyz \u001b[38;5;241m=\u001b[39m atoms_xyz[selection_bool]\n\u001b[0;32m----> 5\u001b[0m probe_xyz,keep_mask \u001b[38;5;241m=\u001b[39m radial_shell_mp(atoms_xyz,\n\u001b[1;32m      6\u001b[0m                                           n_probes\u001b[38;5;241m=\u001b[39mn_probes,\n\u001b[1;32m      7\u001b[0m                                           radii\u001b[38;5;241m=\u001b[39mradii,\n\u001b[1;32m      8\u001b[0m                                           rtol \u001b[38;5;241m=\u001b[39m rtol,\n\u001b[1;32m      9\u001b[0m                                           num_processes\u001b[38;5;241m=\u001b[39mnum_processes)\n",
      "File \u001b[0;32m~/Projects/Qscore/qscore_standalone/qscore.py:78\u001b[0m, in \u001b[0;36mradial_shell_mp\u001b[0;34m(atoms_xyz, n_probes, radii, rtol, num_processes)\u001b[0m\n\u001b[1;32m     76\u001b[0m   results \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m     77\u001b[0m   \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args:\n\u001b[0;32m---> 78\u001b[0m     result \u001b[38;5;241m=\u001b[39m radial_shell_worker(arg)\n\u001b[1;32m     79\u001b[0m     results\u001b[38;5;241m.\u001b[39mappend(result)\n\u001b[1;32m     81\u001b[0m probe_xyz_all \u001b[38;5;241m=\u001b[39m [result[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;28;01mfor\u001b[39;00m result \u001b[38;5;129;01min\u001b[39;00m results]\n",
      "File \u001b[0;32m~/Projects/Qscore/qscore_standalone/qscore.py:14\u001b[0m, in \u001b[0;36mradial_shell_worker\u001b[0;34m(args)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mradial_shell_worker\u001b[39m(args):\n\u001b[0;32m---> 14\u001b[0m     i,atoms_xyz,n_probes,n_probes_target,radius_shell,rtol \u001b[38;5;241m=\u001b[39m args\n\u001b[1;32m     15\u001b[0m     numPts \u001b[38;5;241m=\u001b[39m n_probes_target\n\u001b[1;32m     16\u001b[0m     RAD \u001b[38;5;241m=\u001b[39m radius_shell\n",
      "\u001b[0;31mValueError\u001b[0m: not enough values to unpack (expected 6, got 5)"
     ]
    }
   ],
   "source": [
    "\n",
    "if selection_bool is None:\n",
    "      selection_bool = np.full(atoms_xyz.shape[0],True)\n",
    "atoms_xyz = atoms_xyz[selection_bool]\n",
    "\n",
    "probe_xyz,keep_mask = radial_shell_mp(atoms_xyz,\n",
    "                                          n_probes=n_probes,\n",
    "                                          radii=radii,\n",
    "                                          rtol = rtol,\n",
    "                                          num_processes=num_processes)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7493c8ad-bed4-45ac-8524-bd62860dbfa6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd6b8260-098a-431d-8981-7c5b16fde943",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert to list to have common cctbx/np branch point\n",
    "n_shells, n_atoms,n_probes,_ = probe_xyz.shape\n",
    "\n",
    "probe_xyz_l = probe_xyz.tolist()\n",
    "keep_mask_l = keep_mask.tolist()\n",
    "\n",
    "\n",
    "\n",
    "# Flatten the nested list to match the expected shape\n",
    "\n",
    "probe_xyz_l_flat = flatten_nested_list(probe_xyz_l, 3)\n",
    "keep_mask_l_flat = flatten_nested_list(keep_mask_l,1)\n",
    "\n",
    "#######################\n",
    "# cctbx/np branch point\n",
    "#######################\n",
    "\n",
    "# TEST LIST CONVERSION\n",
    "\n",
    "# Convert the flattened nested list back to a numpy array\n",
    "probe_xyz_l_flat_np = np.array(probe_xyz_l_flat)\n",
    "keep_mask_l_flat_np = np.array(keep_mask_l_flat)\n",
    "\n",
    "# Reshape the original numpy array to match the shape of probe_xyz_l_flat_np\n",
    "reshaped_probe_xyz = probe_xyz.reshape((n_shells*n_atoms*n_probes, 3))\n",
    "reshaped_keep_mask = keep_mask.reshape((n_shells*n_atoms*n_probes))\n",
    "\n",
    "# Compare the two arrays\n",
    "assert np.array_equal(probe_xyz_l_flat_np, reshaped_probe_xyz)\n",
    "assert np.array_equal(keep_mask_l_flat_np, reshaped_keep_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e413422e-fca1-4bc8-809f-3182eb61c20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "do_np = False\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",
    "\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",
    "    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",
    "# 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 = 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)+flex_std(M)*10,flex_max(M))\n",
    "    minD_cctbx = max(flex_mean(M)-flex_std(M)*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)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 300,
   "id": "8babb9e2-b0fa-40ce-a706-273644ca08c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.09960622198608986,\n",
       " 0.29913070189881485,\n",
       " 0.41570096586160304,\n",
       " 0.3816933601922859]"
      ]
     },
     "execution_count": 300,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_cctbx[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0baaa533-4fb9-4c8d-a77d-0227f88e2b30",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "id": "3d6f7d07-206a-4792-950e-a75f54f2c683",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.23491911 0.21801951 0.25036227 0.26736135]\n",
      "[0.2349191050696235, 0.21801951046549, 0.2503622662386241, 0.26736135025163027]\n",
      "[1.56185174 1.56185174 1.56185174 1.56185174]\n",
      "[1.561851742665422, 1.561851742665422, 1.561851742665422, 1.561851742665422]\n"
     ]
    }
   ],
   "source": [
    "print(drow[:4])\n",
    "print(list(drow_cctbx[:4]))\n",
    "print(grow[:4])\n",
    "print(list(grow_cctbx[:4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "ab55ff4e-8e54-402f-87d3-6010701ed258",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.2349191050696235, 0.21801951046549, 0.2503622662386241, 0.26736135025163027]"
      ]
     },
     "execution_count": 279,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 271,
   "id": "8fda9c7d-e424-4e6e-b707-d078cdf02d38",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 261,
   "id": "7e585841-f6fe-44ca-8aaa-901b5e896874",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "240"
      ]
     },
     "execution_count": 261,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(inds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "id": "53c8c711-49e5-4917-9524-449781c77135",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "72fcee26-c980-4758-a3c9-b21694deec1f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64, 240)"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_vals_2d_cctbx.focus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "a9f2eeeb-cb82-4a13-9fa0-038f4c8e90da",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.09960622198608968"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "    # do a row\n",
    "    #shape = (n_atoms,n_shells*n_probes)\n",
    "    #start,stop =  (i * shape[1], (i+1) * shape[1])\n",
    "    #g_val = flex.double(flat_list) # the reference\n",
    "    #d_val = d_vals_2d_cctbx[start:stop]\n",
    "    #print(list(g_val)[:4])\n",
    "    #print(list(d_val)[:4])\n",
    "    #qval = flex.linear_correlation(a,b).coefficient()\n",
    "    #q_cctbx.append(qval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "69cebccc-6d7c-4baa-9784-71de979dea4f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64, 240)"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_vals_2d_cctbx.focus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "8f80e51c-56e4-4076-bf08-5889bcac41d8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20 64 12\n"
     ]
    }
   ],
   "source": [
    "n_shells,n_atoms,n_probes = 20,64,12\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8d0cea-3b8a-4c37-963b-fb25fc9622e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare an output array with zeros\n",
    "D_3d = np.zeros((N,M,L))\n",
    "D_3d[C_3d] = D\n",
    "D_2d = D_3d.transpose(1,0,2).reshape(D_3d.shape[1], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "166ead24-80ef-41e7-923f-34097003077e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# prepare an output array with zeros\n",
    "d_vals = np.zeros((n_shells, n_atoms, n_probes))\n",
    "\n",
    "d_vals[keep_mask] = masked_density\n",
    "d_vals_2d = d_vals.transpose(1,0,2).reshape(d_vals.shape[1], -1)\n",
    "\n",
    "# cctbx\n",
    "def reshape_density_to_full(masked_density_cctbx,keep_mask_flat_cctbx,n_shells,n_atoms,n_probes):\n",
    "    N,M,L = n_shells, n_atoms, n_probes\n",
    "    C = keep_mask_flat_cctbx\n",
    "    D = masked_density_cctbx\n",
    "    \n",
    "    result = []\n",
    "    index_d = 0\n",
    "    for i in range(N):\n",
    "        result_m = []\n",
    "        for j in range(M):\n",
    "            result_l = []\n",
    "            for k in range(L):\n",
    "                # Calculate the equivalent flat index for C\n",
    "                flat_index = i * M * L + j * L + k\n",
    "                if C[flat_index]:\n",
    "                    result_l.append(D[index_d])\n",
    "                    index_d += 1\n",
    "                else:\n",
    "                    result_l.append(fill_value)\n",
    "            result_m.append(result_l)\n",
    "        result.append(result_m)\n",
    "    return result\n",
    "\n",
    "def reshape_density_to_rows(d_vals_cctbx,n_shells,n_atoms,n_probes):\n",
    "    N,M,L = n_shells, n_atoms, n_probes\n",
    "    \n",
    "    \n",
    "    D_3d = d_vals_cctbx\n",
    "    result = []\n",
    "    for j in range(M):\n",
    "        D_2d_row = []\n",
    "        for i in range(N):\n",
    "            for k in range(L):\n",
    "                D_2d_row.append(D_3d[i][j][k])\n",
    "        result.append(D_2d_row)\n",
    "    return result\n",
    "\n",
    "# test full reshape to (n_shells,n_atoms,n_probes)\n",
    "d_vals_cctbx = reshape_density_to_full(masked_density_cctbx,\n",
    "                                      keep_mask_flat_cctbx,\n",
    "                                      n_shells,n_atoms,n_probes)\n",
    "\n",
    "assert np.all(np.isclose(np.array(d_vals_cctbx),d_vals))\n",
    "\n",
    "# test reshape to rows (n_atoms,n_shells*n_probes)\n",
    "d_vals_2d_cctbx = reshape_density_to_rows(d_vals_cctbx,n_shells,n_atoms,n_probes)\n",
    "\n",
    "assert np.all(np.isclose(np.array(d_vals_2d_cctbx),d_vals_2d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "57294432-dfe9-45fa-8f83-948ea9d9c31c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fill_value = 0.0  # or any value you choose as the fill value\n",
    "result = [[] for _ in range(M)]\n",
    "index_d = 0\n",
    "\n",
    "for n in range(N):\n",
    "    for m in range(M):\n",
    "        l_start,l_end = 0,L\n",
    "        for l in range(L):\n",
    "            # Calculate the equivalent flat index for C\n",
    "            flat_index = n * M * L + m * L + l\n",
    "            \n",
    "            # Determine the value to append\n",
    "            if C[flat_index]:\n",
    "                value_to_append = D[index_d]\n",
    "                index_d += 1\n",
    "            else:\n",
    "                value_to_append = fill_value\n",
    "\n",
    "            # Append directly to the appropriate row in the result\n",
    "            result[j].extend([value_to_append])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "b6331eee-6e94-4cb3-a3cc-2e8196c3ffdf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "d_vals_cctbx_flat = flex.double(n_shells*n_atoms*n_probes,0.0)\n",
    "d_vals_cctbx_flat = d_vals_cctbx_flat.set_selected(keep_mask_flat_cctbx,masked_density_cctbx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "id": "ec1f79fb-9060-4284-98f4-de40e9833186",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from itertools import chain\n",
    "d_vals_2d_cctbx = []\n",
    "indices_atoms = []\n",
    "for m in range(M):\n",
    "    row_indices = []\n",
    "    for n in range(N):\n",
    "        l_start = 0\n",
    "        l_end = L-1\n",
    "        \n",
    "        flat_start = n * M * L + m * L + l_start\n",
    "        flat_end = n * M * L + m * L + l_start\n",
    "        l_chunk_indices = list(range(flat_start,flat_end))\n",
    "        row_indices.append(l_chunk_indices)\n",
    "    indices_atoms.append(list(chain.from_iterable(row_indices)))\n",
    "        \n",
    "# d_vals_2d_cctbx = []\n",
    "# for row in indices_atoms:\n",
    "#     chunk = d_vals_cctbx_flat.select(row)\n",
    "#     d_vals_2d_cctbx.append(chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "eb2e36f9-f2db-4f78-b16b-6eff397f3ced",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(20, 64, 12)\n"
     ]
    }
   ],
   "source": [
    "a = flex.double(d_vals)\n",
    "print(a.focus())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d3e2a6-dfad-4d1c-96f9-6c3a86a13724",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "id": "46dd84f2-3984-46b1-9460-30d038f0eb3e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.float64",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[182], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m a\u001b[38;5;241m.\u001b[39mselect(masked_density)\n",
      "\u001b[0;31mTypeError\u001b[0m: No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.float64"
     ]
    }
   ],
   "source": [
    "a.select(masked_density)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "10b85cd7-bd40-4911-abdb-7126fefe0860",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "operands could not be broadcast together with shapes (64,220) (64,240) ",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[170], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m np\u001b[38;5;241m.\u001b[39misclose(np\u001b[38;5;241m.\u001b[39marray(d_vals_2d_cctbx),d_vals_2d)\n",
      "File \u001b[0;32m<__array_function__ internals>:180\u001b[0m, in \u001b[0;36misclose\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/software/miniconda3/lib/python3.11/site-packages/numpy/core/numeric.py:2375\u001b[0m, in \u001b[0;36misclose\u001b[0;34m(a, b, rtol, atol, equal_nan)\u001b[0m\n\u001b[1;32m   2373\u001b[0m yfin \u001b[38;5;241m=\u001b[39m isfinite(y)\n\u001b[1;32m   2374\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mall\u001b[39m(xfin) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mall\u001b[39m(yfin):\n\u001b[0;32m-> 2375\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m within_tol(x, y, atol, rtol)\n\u001b[1;32m   2376\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2377\u001b[0m     finite \u001b[38;5;241m=\u001b[39m xfin \u001b[38;5;241m&\u001b[39m yfin\n",
      "File \u001b[0;32m~/software/miniconda3/lib/python3.11/site-packages/numpy/core/numeric.py:2356\u001b[0m, in \u001b[0;36misclose.<locals>.within_tol\u001b[0;34m(x, y, atol, rtol)\u001b[0m\n\u001b[1;32m   2354\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwithin_tol\u001b[39m(x, y, atol, rtol):\n\u001b[1;32m   2355\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m errstate(invalid\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mignore\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[0;32m-> 2356\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m less_equal(\u001b[38;5;28mabs\u001b[39m(x\u001b[38;5;241m-\u001b[39my), atol \u001b[38;5;241m+\u001b[39m rtol \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mabs\u001b[39m(y))\n",
      "\u001b[0;31mValueError\u001b[0m: operands could not be broadcast together with shapes (64,220) (64,240) "
     ]
    }
   ],
   "source": [
    "np.isclose(np.array(d_vals_2d_cctbx),d_vals_2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "5d8d7642-fd5a-4e00-b25c-7a571f824d2b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64, 220)"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(d_vals_2d_cctbx).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "a451422e-8a08-44c7-ad47-8e38ada0b096",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "select( (double)self, (bool)flags) -> double :\n",
       "\n",
       "    C++ signature :\n",
       "        scitbx::af::shared<double> select(scitbx::af::versa<double, scitbx::af::flex_grid<scitbx::af::small<long, 10ul> > >,scitbx::af::const_ref<bool, scitbx::af::trivial_accessor>)\n",
       "\n",
       "select( (double)self, (object)indices [, (object)reverse=False]) -> double :\n",
       "\n",
       "    C++ signature :\n",
       "        scitbx::af::shared<double> select(scitbx::af::versa<double, scitbx::af::flex_grid<scitbx::af::small<long, 10ul> > >,scitbx::af::const_ref<unsigned int, scitbx::af::trivial_accessor> [,bool=False])\n",
       "\n",
       "select( (double)self, (size_t)indices [, (object)reverse=False]) -> double :\n",
       "\n",
       "    C++ signature :\n",
       "        scitbx::af::shared<double> select(scitbx::af::versa<double, scitbx::af::flex_grid<scitbx::af::small<long, 10ul> > >,scitbx::af::const_ref<unsigned long, scitbx::af::trivial_accessor> [,bool=False])\n",
       "\n",
       "select( (double)arg1, (unsigned)selection) -> double :\n",
       "\n",
       "    C++ signature :\n",
       "        scitbx::af::shared<double> select(scitbx::af::versa<double, scitbx::af::flex_grid<scitbx::af::small<long, 10ul> > >,std::vector<unsigned int, std::allocator<unsigned int> >)\n",
       "\n",
       "select( (double)arg1, (unsigned)selection) -> double :\n",
       "\n",
       "    C++ signature :\n",
       "        scitbx::af::shared<double> select(scitbx::af::versa<double, scitbx::af::flex_grid<scitbx::af::small<long, 10ul> > >,std::set<unsigned int, std::less<unsigned int>, std::allocator<unsigned int> >)\n",
       "\u001b[0;31mType:\u001b[0m      method"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d_vals_cctbx_flat.select??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "c9dc18ff-0594-4a28-ae2a-d688339f8349",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1280, 11)"
      ]
     },
     "execution_count": 142,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(d_vals_2d_cctbx).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d7f22d1c-eb86-4424-921d-8028af055088",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# array A is shape (N,M,L)\n",
    "B = A.reshape(-1)\n",
    "\n",
    "# array C.shape==B.shape (N*M*L). C is a boolean mask\n",
    "\n",
    "# select values\n",
    "D = B[C]\n",
    "\n",
    "\n",
    "\n",
    "# I now want to transform D to list of len N*M where each value is a list of len L, where the values are taken from D where there were selected (C==True), or are a fill value if they correspond to a False value in C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8b2c370a-4502-48b2-8a8a-dddac84ae150",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.23491911,  0.21801951,  0.25036227, ..., -0.17700099,\n",
       "       -0.21437779, -0.20423411])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0222d0e9-5e6b-4e5d-9f50-ccb0d4730ef1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "M = mm.map_data()\n",
    "maxD = min(flex_mean(M)+flex_std(M)*10,flex_min(M))\n",
    "minD = max(flex_mean(M)-flex_std(M)*1,flex_min(M))\n",
    "A = maxD-minD\n",
    "B = minD\n",
    "u = 0\n",
    "sigma = 0.6\n",
    "x = radii\n",
    "y = A * np.exp(-0.5*((x-u)/sigma)**2) + B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cab4a494-9070-4d6e-a93d-b3ed60f861bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "flex.v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c09afa-7ee9-41fa-990b-c7f9876af46d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2242890d-ba40-447f-8af5-b837e3fcf73a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Qscore(volume,\n",
    "                atoms_xyz,\n",
    "                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_bool=None,\n",
    "                num_processes=cpu_count()):\n",
    "\n",
    "    # handle selection at the very beginning\n",
    "    if selection_bool is None:\n",
    "      selection_bool = np.full(atoms_xyz.shape[0],True)\n",
    "    atoms_xyz = atoms_xyz[selection_bool]\n",
    "\n",
    "\n",
    "              \n",
    "    probe_xyz,keep_mask = radial_shell_mp(atoms_xyz,\n",
    "                                          n_probes=n_probes,\n",
    "                                          radii=radii,\n",
    "                                          rtol = rtol,\n",
    "                                          num_processes=num_processes)\n",
    "    \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",
    "    # apply mask to the flattened probe_xyz\n",
    "    masked_probe_xyz_flat = probe_xyz_flat[keep_mask_flat]\n",
    "    #masked_probe_xyz_flat_flex = flex.vec3_double(masked_probe_xyz_flat)\n",
    "    \n",
    "    # apply trilinear interpolation only to the relevant probes\n",
    "    masked_density = trilinear_interpolation(volume, masked_probe_xyz_flat, voxel_size=voxel_size) # (n_valid_probes,)\n",
    "    #masked_density = mm.density_at_sites_cart(masked_probe_xyz_flat_flex).as_numpy_array()\n",
    "    \n",
    "    # prepare an output array with zeros\n",
    "    d_vals = np.zeros((n_shells, n_atoms, n_probes))\n",
    "    \n",
    "    # reshape interpolated values to (n_shells, n_atoms, n_probes) using the mask\n",
    "    d_vals[keep_mask] = masked_density\n",
    "    \n",
    "    \n",
    "    \n",
    "    n_atoms = probe_xyz.shape[1]\n",
    "    n_probes = probe_xyz.shape[2]\n",
    "    M = volume\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 = radii\n",
    "    y = A * np.exp(-0.5*((x-u)/sigma)**2) + B \n",
    "    \n",
    "    # stack the reference to shape (n_shells,n_atoms,n_probes)\n",
    "    g_vals = np.repeat(y[:,None],n_probes,axis=1)\n",
    "    x_repeat = np.repeat(x,n_probes)\n",
    "    g_vals = np.expand_dims(g_vals,1)\n",
    "    \n",
    "    g_vals = np.tile(g_vals,(n_atoms,1))\n",
    "    \n",
    "    # Reshape to 2d for masked rowwise correlation calculation\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",
    "    q = rowwise_corrcoef(g_vals_2d,d_vals_2d,mask=mask_2d)\n",
    "    return q,probe_xyz,keep_mask,d_vals, g_vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4e0008-c712-4d29-9730-25f07e6aa363",
   "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
}