alignment_based_on_molecule.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using PorousMaterials # Pkg.clone(\"https://github.com/SimonEnsemble/PorousMaterials.jl\", \"v0.1.1\")\n",
    "using DelimitedFiles\n",
    "using LinearAlgebra\n",
    "using JLD2\n",
    "using CoherentPointDrift\n",
    "import Bio3DView\n",
    "using ProgressMeter\n",
    "using Printf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Database of cages\n",
    "\n",
    "Cage database CDB41 (41 cages) from [here](https://github.com/marcinmiklitz/CDB41). Cleaned of solvent. Thanks to Kim Jelfs and M. Miklitz for kindly providing the cages.\n",
    "\n",
    "Reference:\n",
    "> M. Miklitz, S. Jiang, R. Clowes, M. E. Briggs, A. I. Cooper and K. E. Jelfs, Computational Screening of Porous Organic Molecules for Xenon/Krypton Separation, J. Phys. Chem. C, 2017. DOI: 10.1021/acs.jpcc.7b03848\n",
    "\n",
    "Other set of 33 cages from the .xyz files deposited as ESI of DOI 10.1038/s41467-018-05271-9\n",
    "\n",
    "> R. L. Greenaway, V. Santolini, M. J. Bennison, B. M. Alston, C. J. Pugh, M. A. Little, M. Miklitz, E. G. B. Eden-Rump, R. Clowes, A. Shakil, H. J. Cuthbertson, H. Armstrong, M. E. Briggs, K. E. Jelfs & A. I. Cooper. High-throughput discovery of organic cages and catenanes using computational screening fused with robotic synthesis. Nature Communications, 2018. DOI: 10.1038/s41467-018-05271-9\n",
    "                                                              \n",
    "I manually looked at Fig 4 in the main text and selected those that were synthesized. Some cages required visualization to compare to Fig 4 since e.g. there are multiple versions of B23, a few of which were apparently not synthesized but only computational predictions.\n",
    "\n",
    "The `.xyz` files describing the molecular structure of all of these cages are in the directory `all_cages`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "cages = readdlm(\"all_cages/all_cages.txt\", String)[:];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Center cages\n",
    "\n",
    "The center of mass will be set to be the origin.\n",
    "Each cage will be rotated such that it is aligned with its principle moments of inertia.\n",
    "i.e. the moment of inertia matrix of an \"aligned\" cage is diagonal.\n",
    "\n",
    "\n",
    "First, a function to compute the center of mass of a group of `atoms` at Cartesian coordinates `x`, stored in the columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "center_of_mass (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function center_of_mass(atoms::Array{Symbol, 1}, x::Array{Float64, 2})\n",
    "    ams = read_atomic_masses()\n",
    "    com = [0.0, 0.0, 0.0]\n",
    "    total_mass = 0.0\n",
    "    for (i, atom) in enumerate(atoms)\n",
    "        com += ams[atom] * x[:, i]\n",
    "        total_mass += ams[atom]\n",
    "    end\n",
    "    return com / total_mass\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "centered_cage_coords (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function centered_cage_coords(cage::AbstractString)\n",
    "    # read in raw .xyz from `all_cages`\n",
    "    atoms, x = read_xyz(\"all_cages/\" * cage * \".xyz\")\n",
    "    \n",
    "    # compute center of mass\n",
    "    x_com = center_of_mass(atoms, x)\n",
    "    \n",
    "    # shift coords so that cage is centered\n",
    "    x = x .- x_com\n",
    "    \n",
    "    # write centered cage coords\n",
    "    write_xyz(atoms, x, \"centered_cages/\" * cage * \".xyz\")\n",
    "    \n",
    "    return atoms, x\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "rotational_inertia_aligned_cage (generic function with 1 method)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function moment_of_inertia(atoms::Array{Symbol, 1}, x::Array{Float64, 2})\n",
    "    @assert(isapprox(center_of_mass(atoms, x), [0.0, 0.0, 0.0], atol=0.0001),\n",
    "        \"cage must be centered to compute moment of inertia matrix!\")\n",
    "    \n",
    "    # read in atomic masses\n",
    "    ams = read_atomic_masses()\n",
    "    \n",
    "    # moment of inertia matrix `mi`\n",
    "    mi = zeros(Float64, 3, 3)\n",
    "    for i = 1:3\n",
    "        for j = 1:3\n",
    "            for (a, atom) in enumerate(atoms)\n",
    "                if i == j\n",
    "                    ids = collect(setdiff(Set([1, 2, 3]), Set([i])))\n",
    "                    mi[i, j] += ams[atom] * sum(x[ids, a] .^ 2)\n",
    "                else\n",
    "                    mi[i, j] -= ams[atom] * x[i, a] * x[j, a]\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    # should be symmetric!\n",
    "    @assert isapprox(mi, mi')\n",
    "        \n",
    "    return mi\n",
    "end\n",
    "\n",
    "function diagonlize_moment_of_inertia(atoms::Array{Symbol, 1}, x::Array{Float64, 2})\n",
    "    @assert(isapprox(center_of_mass(atoms, x), [0.0, 0.0, 0.0], atol=0.0001),\n",
    "        \"cage not centered!\")\n",
    "  \n",
    "    # compute moment of inertia matrix\n",
    "    mi = moment_of_inertia(atoms, x)\n",
    "\n",
    "    # diagnomize moment of inertia matrix\n",
    "    λ, v = eigen(mi) # columns of v are eigenvectors\n",
    "    \n",
    "    # reorder eigenvalues and eigenvectors from large to small\n",
    "    ids = sortperm(λ, rev=true) # sort eigenvalues from large to small.\n",
    "    λ = λ[ids]\n",
    "    v = v[:, ids] # eigenvectors in columns\n",
    "    \n",
    "    # tests out of paranoia\n",
    "    @assert(isapprox(mi * v, v * diagm(0 => λ)), \"eigenvectors not right\")\n",
    "    @assert(isapprox([norm(v[:, i]) for i = 1:3], ones(3)), \"eigenvectors not unit vectors\")\n",
    "    @assert(isapprox(mi, v * diagm(0 => λ) * v'))\n",
    "    # `v` is a rotation matrix b/c `mi` is symmetric and the columns are orthonormal\n",
    "    # i.e. `v` is a unitary matrix. now use it to rotate x to align with principle axes of rotation\n",
    "    @assert(isapprox(v' * v, Diagonal{Float64}(I, 3)), \"v is not unitary...\")\n",
    "    @assert (λ[1] >= λ[2]) & (λ[2] >= λ[3])\n",
    "    \n",
    "    return λ, v, mi\n",
    "end\n",
    "\n",
    "function rotational_inertia_aligned_cage(cage::AbstractString)\n",
    "    atoms, x = centered_cage_coords(cage)\n",
    "    λ, v, mi = diagonlize_moment_of_inertia(atoms, x)\n",
    "    \n",
    "    # rotate cage so as to align it with the principle axes of inertia\n",
    "    x = v' * x\n",
    "    \n",
    "    # now ensure its moment of inertia matrix is diagonal\n",
    "    mi = moment_of_inertia(atoms, x)\n",
    "    @assert(isapprox(diagm(0 => diag(mi)), mi), \"not aligned with principle axes\")\n",
    "\n",
    "    # assert moments of inertia are ordered largest to smallest\n",
    "    @assert(mi[1, 1] >= mi[2, 2])\n",
    "    @assert(mi[2, 2] >= mi[3, 3])\n",
    "    return atoms, x\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "read_final_aligned_cage (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if ! isdir(\"molecule_aligned_cages\")\n",
    "    mkdir(\"molecule_aligned_cages\")\n",
    "end\n",
    "write_final_aligned_cage(atoms::Array{Symbol, 1}, x::Array{Float64, 2}, cage::AbstractString) = write_xyz(atoms, x, \"molecule_aligned_cages/\" * cage * \".xyz\")\n",
    "read_final_aligned_cage(cage::AbstractString) = read_xyz(\"molecule_aligned_cages/\" * cage * \".xyz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "find_pair_to_align (generic function with 1 method)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function read_jld_file_results(cage_y::AbstractString, cage_x::AbstractString; verbose::Bool=false)\n",
    "    jldfilename = jldopen(@sprintf(\"cpd_results/align_%s_to_%s_molecule.jld2\", cage_y, cage_x))\n",
    "    R = read(jldfilename, \"R\")\n",
    "    σ² = read(jldfilename, \"σ²\")\n",
    "    ℓ = read(jldfilename, \"ℓ\")\n",
    "    @assert read(jldfilename, \"cage_y\") == cage_y\n",
    "    @assert read(jldfilename, \"cage_x\") == cage_x\n",
    "    close(jldfilename)\n",
    "    return R, σ², ℓ\n",
    "end\n",
    "\n",
    "function find_pair_to_align(cage_ys::Array{String, 1}, cage_xs::Array{String, 1})\n",
    "    # look at all possible pairs\n",
    "    pairz = Tuple{String, String}[]\n",
    "    ℓs = Float64[]\n",
    "    for cage_y in cage_ys\n",
    "        for cage_x in cage_xs\n",
    "            if cage_y == cage_x\n",
    "                continue\n",
    "            end\n",
    "            \n",
    "            R, σ², ℓ = read_jld_file_results(cage_y, cage_x)\n",
    "            \n",
    "            push!(ℓs, ℓ)\n",
    "            push!(pairz, (cage_y, cage_x))\n",
    "        end\n",
    "    end\n",
    "    best_pair = pairz[argmin(ℓs)]\n",
    "    return best_pair[1], best_pair[2]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mAligning cage CC9 to CC10...\u001b[39m\n"
     ]
    },
    {
     "ename": "SystemError",
     "evalue": "SystemError: opening file molecule_aligned_cages/CC10.xyz: No such file or directory",
     "output_type": "error",
     "traceback": [
      "SystemError: opening file molecule_aligned_cages/CC10.xyz: No such file or directory",
      "",
      "Stacktrace:",
      " [1] #systemerror#39(::Nothing, ::Function, ::String, ::Bool) at ./error.jl:106",
      " [2] systemerror at ./error.jl:106 [inlined]",
      " [3] #open#293(::Nothing, ::Nothing, ::Nothing, ::Nothing, ::Nothing, ::Function, ::String) at ./iostream.jl:283",
      " [4] open at ./iostream.jl:275 [inlined]",
      " [5] read_xyz(::String) at /home/simoncor/.julia/dev/PorousMaterials/src/Misc.jl:17",
      " [6] read_final_aligned_cage(::String) at ./In[12]:5",
      " [7] top-level scope at In[14]:15"
     ]
    }
   ],
   "source": [
    "aligned = Dict(cage => false for cage in cages)\n",
    "\n",
    "cage_y, cage_x = find_pair_to_align(cages, cages)\n",
    "aligned[cage_x] = true\n",
    "\n",
    "while sum([! aligned[cage] for cage in cages]) > 0\n",
    "    \n",
    "    aligned_cages = cages[[aligned[cage] for cage in cages]]\n",
    "    unaligned_cages = cages[[! aligned[cage] for cage in cages]]\n",
    "    \n",
    "    cage_y, cage_x = find_pair_to_align(unaligned_cages, aligned_cages)\n",
    "    printstyled(@sprintf(\"Aligning cage %s to %s...\\n\", cage_y, cage_x), color=:yellow)\n",
    "    @assert aligned[cage_x] && ! aligned[cage_y]\n",
    "    \n",
    "    atoms_x, x = read_final_aligned_cage(cage_x)\n",
    "    \n",
    "    atoms_y, y = rotational_inertia_aligned_cage(cage_y)\n",
    "    \n",
    "    # transformation is applied not to reference cage but the unaligned cage!\n",
    "    R, t, σ², ℓ = CoherentPointDrift.rigid_point_set_registration(x, y, verbose=false,\n",
    "        w=0.0, σ²_tol=0.05, q_tol=1.0, max_nb_em_steps=30, print_ending=true)\n",
    "\n",
    "    write_final_aligned_cage(atoms_y, R * y, cage_y)\n",
    "    \n",
    "    aligned[cage_y] = true\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aligned = Dict(cage => false for cage in cages)\n",
    "\n",
    "# me = \"B1\"\n",
    "# aligned[me] = true\n",
    "# atoms, x = centered_cage_coords(me)\n",
    "# write_final_aligned_cage(atoms, x, me)\n",
    "\n",
    "# while sum([! aligned[cage] for cage in cages]) > 0\n",
    "#     # cage is unaligned. find best to align to among already-aligned cages\n",
    "#     printstyled(\"Finding best cage to align with me,\" * me * \"...\", color=:green)\n",
    "\n",
    "#     unaligned_cages = cages[[! aligned[cage] for cage in cages]]\n",
    "#     filter(x -> x != me, unaligned_cages)\n",
    "    \n",
    "#     align_to_me = find_best_cage_to_align_to(me, unaligned_cages)# , cage_names=aligned_cages)\n",
    "#     printstyled(align_to_me * \" aligns with me best.\\n\", color=:red)\n",
    "    \n",
    "#     atoms_x, x = rotational_inertia_aligned_cage(me)\n",
    "#     atoms_y, y = rotational_inertia_aligned_cage(align_to_me)\n",
    "    \n",
    "#     println(\"\\tCoherent point drift params: \")\n",
    "#     R, t, σ², ℓ = CoherentPointDrift.rigid_point_set_registration(x, y, verbose=false,\n",
    "#             w=0.0, σ²_tol=0.1, q_tol=1.0, max_nb_em_steps=25, print_ending=true)\n",
    "    \n",
    "#     write_final_aligned_cage(atoms_y, R * y, align_to_me)\n",
    "#     aligned[align_to_me]= true\n",
    "#     me = align_to_me\n",
    "# end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.0.1",
   "language": "julia",
   "name": "julia-1.0"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.0.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}