2023.04.12_kabsch_algo_rotate_align.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "78b09b93",
   "metadata": {},
   "source": [
    "Sources: https://towardsdatascience.com/the-definitive-procedure-for-aligning-two-sets-of-3d-points-with-the-kabsch-algorithm-a7ec2126c87e "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04786b2",
   "metadata": {},
   "source": [
    "https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.align_vectors.html "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c59366",
   "metadata": {},
   "source": [
    "# Given two Nx3 matrices, Rotate to minimize RMSD. Apply for computing RMSD for two wlcsim structures. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "18ad6e5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# import scipy\n",
    "# from scipy import spatial\n",
    "from scipy.spatial.transform import Rotation as R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "93e761d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def avg_col_subtract(a):\n",
    "    \"\"\"\n",
    "    Subtract from each element the average of the whole column to center \n",
    "    the matrix at the origin\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    \n",
    "    a: Nx3 np.array or np.matrix\n",
    "        Original matrix\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    b: Nx3 np.array or np.matrix\n",
    "        Mean subtracted matrix \n",
    "        \n",
    "    Cites\n",
    "    -----\n",
    "    See bullet points 1 and 2: https://towardsdatascience.com/the-definitive-procedure-for-aligning-two-sets-of-3d-points-with-the-kabsch-algorithm-a7ec2126c87e \n",
    "    \"\"\"\n",
    "    b = a - a.mean(axis=0)\n",
    "    return b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1bd5e595",
   "metadata": {},
   "outputs": [],
   "source": [
    "array_a = np.matrix([[1,2,3],[4,5,6],[6,23,6],[2,55,6],[41,35,6],[45,5,16]])\n",
    "array_b = array_a + 13\n",
    "rotate90 = np.matrix([[0, -1, 0],\n",
    "                   [1, 0, 0],\n",
    "                   [0, 0, 1]])\n",
    "array_c = np.matmul(array_a, rotate90)\n",
    "array_d = np.matmul(array_c, rotate90)\n",
    "\n",
    "a_centered = avg_col_subtract(array_a)\n",
    "b_centered = avg_col_subtract(array_b)\n",
    "c_centered = avg_col_subtract(array_c)\n",
    "d_centered = avg_col_subtract(array_d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e736834f",
   "metadata": {},
   "source": [
    "##### rssd is the square root of the weighted sum of the squared distances between the given sets of vectors after alignment.  It is equal to sqrt(2 * minimum_loss), where minimum_loss is the loss function evaluated for the found optimal rotation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c6b581b",
   "metadata": {},
   "source": [
    "The distance between a vector and itself should be 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "ed135853",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimated_rotation, rssd = R.align_vectors(a=array_a, b=array_a)\n",
    "rssd==0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63e465ee",
   "metadata": {},
   "source": [
    "The distance between a vector and a translated version of itself should be 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8c8b5c28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "53.590917643656574"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimated_rotation, rssd = R.align_vectors(a=array_a, b=array_b)\n",
    "rssd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "838e6d23",
   "metadata": {},
   "source": [
    "The above failed because the matrices were not centered at the origin. Corrected below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9a323365",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.3486991523486091e-06"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimated_rotation, rssd = R.align_vectors(a=a_centered, b=b_centered)\n",
    "rssd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29e91183",
   "metadata": {},
   "source": [
    "The distance between a vector and a rotated version of itself should be 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "69197bc9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimated_rotation, rssd = R.align_vectors(a=a_centered, b=c_centered)\n",
    "rssd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "3c6d1654",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimated_rotation, rssd = R.align_vectors(a=array_a, b=array_d)\n",
    "rssd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28bf6280",
   "metadata": {},
   "source": [
    "#### Check that differences can be detected:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e069e247",
   "metadata": {},
   "source": [
    "Move one point along z axis. Verify that the rssd monotonically increases as the translation along the z increases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "497b4abb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.3486991523486091e-06\n",
      "0.7443273964019765\n",
      "1.4879472943013732\n",
      "2.230892732984932\n",
      "2.9731975852677732\n",
      "3.7148965197470205\n",
      "4.456024960472314\n",
      "5.196619044134032\n",
      "5.9367155748766764\n",
      "6.6763519768948765\n"
     ]
    }
   ],
   "source": [
    "for delta in np.arange(0,10):\n",
    "    array_e = array_a + 13 + np.matrix([[0, 0, delta],[0, 0, 0],[0, 0, 0],[0, 0, 0],[0, 0, 0],[0, 0, 0]])\n",
    "    e_centered = avg_col_subtract(array_e)\n",
    "    estimated_rotation, rssd = R.align_vectors(a=a_centered, b=e_centered)\n",
    "    print(rssd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b982e05",
   "metadata": {},
   "source": [
    "Combine as single formula in which 2 np.arrays can be given as input ans the rssd is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "7c7da7bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://www.physicsforums.com/threads/rms-vs-rss-for-uncertainty.982466/ \n",
    "def rss_to_rmsd(rss,num_points):\n",
    "    \"\"\"\n",
    "    Get RMSD from Root Sum of Squares by dividing by 1/n \n",
    "    where n is the number of points\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    rss: float\n",
    "        root sum of squares computed from kabsch algorithm\n",
    "        \n",
    "    num_points: int\n",
    "        Number of point in the original dataset \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    rmsd: float\n",
    "        root mean square deviation (aka rmse root mean square error)\n",
    "    \n",
    "    Cite (maybe verify with better cite)\n",
    "    ----\n",
    "    https://www.physicsforums.com/threads/rms-vs-rss-for-uncertainty.982466/ \n",
    "    \"\"\"\n",
    "    rmsd = rss / np.sqrt(num_points)\n",
    "    return rmsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "e85e8dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kabsh_algo_rmsd(positions_a, positions_b):\n",
    "    \"\"\"\n",
    "    Given two sets of 3D positioning data, compute the RMSD after \n",
    "    optimally aligning using the kabsch algorithm \n",
    "    \n",
    "    Parameters\n",
    "    ---------\n",
    "    positions_a: Nx3 np.array or np.matrix\n",
    "        First set of positions \n",
    "    positions_b: Nx3 np.array or np.matrix\n",
    "        Second set of positions \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    rmsd : float \n",
    "        root-mean-square deviation \n",
    "    \"\"\"\n",
    "    number_points = np.shape(positions_a)[0]\n",
    "    points_a_centered = avg_col_subtract(positions_a)\n",
    "    points_b_centered = avg_col_subtract(positions_b)\n",
    "    estimated_rotation, rssd = R.align_vectors(a=points_a_centered, b=points_b_centered)\n",
    "    rmsd = rss_to_rmsd(rssd, number_points)\n",
    "    return rmsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "e3041f3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0\n",
      "0.30387038712652503\n",
      "0.6074519391988661\n",
      "0.9107581444493539\n",
      "1.213802831396854\n",
      "1.5166001534367755\n",
      "1.8191645723772087\n",
      "2.1215108409596013\n",
      "2.423653984413593\n",
      "2.7256092811022543\n"
     ]
    }
   ],
   "source": [
    "for delta in np.arange(0,10):\n",
    "    array_e = array_a + 13 + np.matrix([[0, 0, delta],[0, 0, 0],[0, 0, 0],[0, 0, 0],[0, 0, 0],[0, 0, 0]])\n",
    "    e_centered = avg_col_subtract(array_e)\n",
    "    rmsd = kabsh_algo_rmsd(a_centered, e_centered)\n",
    "    print(rmsd)"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}