CorrFunction_NMRRelaxation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.optimize import curve_fit\n",
    "import mdtraj as md"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_NHVecs(traj_file, top_file, start_snap=0, end_snap=-1):\n",
    "    \"\"\"\n",
    "    Uses mdtraj to load the trajectory and get the atomic indices and coordinates to calculate the correlation functions.\n",
    "    For each, trajectory load the trajectory using mdtraj, get the atomic index for the the N-H atoms and calculate the vector between the two.\n",
    "    Append the vector to the NHVecs list for all the trajectories. \n",
    "    NHVecs should return a list of shape: (# Trajectories, # Snapshots, # Residues w/N-H Vectors, 3)\n",
    "    \"\"\"\n",
    "    traj = md.load(traj_file, top=top_file)\n",
    "    top = traj.topology\n",
    "    \n",
    "    ##AtomSelection Indices\n",
    "    Nit = top.select('name N and not resname PRO') ## PRO residue do not have N-H vectors\n",
    "    Hyd = top.select('name H and not resname PRO')\n",
    "    NH_Pair = [[i,j] for i,j in zip(Nit,Hyd)]\n",
    "    NH_Pair_Name = [[top.atom(i),top.atom(j)] for i,j in NH_Pair]\n",
    "    NH_Res = [\"{}-{}{}\".format(str(i).split('-')[0],str(i).split('-')[1], str(j).split('-')[1]) for i,j in NH_Pair_Name]\n",
    "    \n",
    "    ##Generate the N-H vectors in Laboratory Frame\n",
    "    NHVecs_tmp = np.take(traj.xyz, Hyd, axis=1) - np.take(traj.xyz, Nit, axis=1)\n",
    "    sh = list(NHVecs_tmp.shape)\n",
    "    sh[2] = 1\n",
    "    NHVecs_tmp = NHVecs_tmp / np.linalg.norm(NHVecs_tmp, axis=2).reshape(sh)\n",
    "    \n",
    "    return NHVecs_tmp[start_snap:end_snap]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_NHVecs(nhvecs, dt, tau):\n",
    "    \"\"\"\n",
    "    This function will split the trajectory in chunks defined by tau. \n",
    "    nhvecs = array of N-H bond vectors,\n",
    "    dt = timestep of the simulation\n",
    "    tau = length of chunks\n",
    "    \"\"\"\n",
    "    nFiles = len(nhvecs) ## number of trajectories\n",
    "    nFramesPerChunk = int(tau/dt) ###tau/timestep \n",
    "    used_frames = np.zeros(nFiles,dtype=int)\n",
    "    remainingFrames = np.zeros(nFiles,dtype=int)\n",
    "    for i in range(nFiles):\n",
    "        nFrames = nhvecs[i].shape[0]\n",
    "        used_frames[i] = int(nFrames/nFramesPerChunk)*nFramesPerChunk\n",
    "        remainingFrames[i] = nFrames % nFramesPerChunk\n",
    "    \n",
    "    nFramesTot=int(used_frames.sum())\n",
    "    out = np.zeros((nFramesTot,NHVecs[0].shape[1],NHVecs[0].shape[2]), dtype=NHVecs[0].dtype)\n",
    "    start = 0\n",
    "    for i in range(nFiles):\n",
    "        end = int(start+used_frames[i])\n",
    "        endv = int(used_frames[i])\n",
    "        out[start:end,...] = nhvecs[i][0:endv,...]\n",
    "        start = end\n",
    "        \n",
    "    sh = out.shape\n",
    "    vecs = out.reshape((int(nFramesTot/nFramesPerChunk), nFramesPerChunk, sh[-2], sh[-1]))\n",
    "    \n",
    "    return vecs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_Ct(nhvecs):\n",
    "    \"\"\"\n",
    "    Calculates the correlation function of the N-H bond vectors found in nhvecs. \n",
    "    Direct space calculation. This could be changed to Fourier space calculation for increased speed. \n",
    "    \n",
    "    LICENSE INFO:\n",
    "    \n",
    "    MIT License\n",
    "\n",
    "    Copyright (c) 2017 Po-chia Chen\n",
    "\n",
    "    Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "    of this software and associated documentation files (the \"Software\"), to deal\n",
    "    in the Software without restriction, including without limitation the rights\n",
    "    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "    copies of the Software, and to permit persons to whom the Software is\n",
    "    furnished to do so, subject to the following conditions:\n",
    "\n",
    "    The above copyright notice and this permission notice shall be included in all\n",
    "    copies or substantial portions of the Software.\n",
    "\n",
    "    THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "    SOFTWARE.\n",
    "\n",
    "    \"\"\"\n",
    "    sh = nhvecs.shape\n",
    "    nReplicates=sh[0] ; nDeltas=int(sh[1]/2) ; nResidues=sh[2]\n",
    "    Ct  = np.zeros( (nDeltas, nResidues), dtype=nhvecs.dtype )\n",
    "    dCt = np.zeros( (nDeltas, nResidues), dtype=nhvecs.dtype )\n",
    "    \n",
    "    for delta in range(1,1+nDeltas):\n",
    "        nVals=sh[1]-delta\n",
    "        # = = Create < vi.v'i > with dimensions (nRep, nFr, nRes, 3) -> (nRep, nFr, nRes) -> ( nRep, nRes ), then average across replicates with SEM.\n",
    "        tmp = -0.5 + 1.5 * np.square( np.einsum( 'ijkl,ijkl->ijk', nhvecs[:,:-delta,...] , nhvecs[:,delta:,...] ) )\n",
    "        tmp  = np.einsum( 'ijk->ik', tmp ) / nVals\n",
    "        Ct[delta-1]  = np.mean( tmp, axis=0 )\n",
    "        dCt[delta-1] = np.std( tmp, axis=0 ) / ( np.sqrt(nReplicates) - 1.0 )\n",
    "    \n",
    "    return Ct, dCt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _bound_check(func, params):\n",
    "    \"\"\"\n",
    "    \n",
    "    Checks if the fit returns a sum of the amplitudes greater than 1.\n",
    "    \n",
    "    MIT License\n",
    "\n",
    "    Copyright (c) 2017 Po-chia Chen\n",
    "\n",
    "    Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "    of this software and associated documentation files (the \"Software\"), to deal\n",
    "    in the Software without restriction, including without limitation the rights\n",
    "    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "    copies of the Software, and to permit persons to whom the Software is\n",
    "    furnished to do so, subject to the following conditions:\n",
    "\n",
    "    The above copyright notice and this permission notice shall be included in all\n",
    "    copies or substantial portions of the Software.\n",
    "\n",
    "    THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "    SOFTWARE.\n",
    "\n",
    "    \"\"\"\n",
    "    if len(params) == 1:\n",
    "        return False\n",
    "    elif len(params) %2 == 0 :\n",
    "        s = sum(params[0::2])\n",
    "        return (s>1)\n",
    "    else:\n",
    "        s = params[0]+sum(params[1::2])\n",
    "        return (s>1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_chi(y1, y2, dy=[]):\n",
    "    \"\"\"\n",
    "    Calculates the chi^2 difference between the predicted model and the actual data. \n",
    "    \n",
    "    LICENSE INFO:\n",
    "    MIT License\n",
    "\n",
    "    Copyright (c) 2017 Po-chia Chen\n",
    "\n",
    "    Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "    of this software and associated documentation files (the \"Software\"), to deal\n",
    "    in the Software without restriction, including without limitation the rights\n",
    "    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "    copies of the Software, and to permit persons to whom the Software is\n",
    "    furnished to do so, subject to the following conditions:\n",
    "\n",
    "    The above copyright notice and this permission notice shall be included in all\n",
    "    copies or substantial portions of the Software.\n",
    "\n",
    "    THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "    SOFTWARE.\n",
    "\n",
    "    \"\"\"\n",
    "    if dy != []:\n",
    "        return np.sum( (y1-y2)**2.0/dy )/len(y1)\n",
    "    else:\n",
    "        return np.sum( (y1-y2)**2.0 )/len(y1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Functions 1,3,5,7,9 are the functions that the sum of coefficients are equal to 1. They have one less parameter.\n",
    "## Functions 2,4,6,8,10 are the functions where the sum of coefficients are not restricted.\n",
    "\n",
    "def func_exp_decay1(t, tau_a):\n",
    "    return np.exp(-t/tau_a)\n",
    "def func_exp_decay2(t, A, tau_a):\n",
    "    return A*np.exp(-t/tau_a)\n",
    "def func_exp_decay3(t, A, tau_a, tau_b):\n",
    "    return A*np.exp(-t/tau_a) + (1-A)*np.exp(-t/tau_b)\n",
    "def func_exp_decay4(t, A, tau_a, B, tau_b ):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b)\n",
    "def func_exp_decay5(t, A, tau_a, B, tau_b, tau_g ):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + (1-A-B)*np.exp(-t/tau_g)\n",
    "def func_exp_decay6(t, A, tau_a, B, tau_b, G, tau_g ):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + G*np.exp(-t/tau_g)\n",
    "def func_exp_decay7(t, A, tau_a, B, tau_b, G, tau_g, tau_d):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + G*np.exp(-t/tau_g) + (1-A-B-G)*np.exp(-t/tau_d)\n",
    "def func_exp_decay8(t, A, tau_a, B, tau_b, G, tau_g, D, tau_d):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + G*np.exp(-t/tau_g) + D*np.exp(-t/tau_d)\n",
    "def func_exp_decay9(t, A, tau_a, B, tau_b, G, tau_g, D, tau_d, tau_e):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + G*np.exp(-t/tau_g) + D*np.exp(-t/tau_d) + (1-A-B-G-D)*np.exp(-t/tau_e)\n",
    "def func_exp_decay10(t, A, tau_a, B, tau_b, G, tau_g, D, tau_d, E, tau_e):\n",
    "    return A*np.exp(-t/tau_a) + B*np.exp(-t/tau_b) + G*np.exp(-t/tau_g) + D*np.exp(-t/tau_d) + E*np.exp(-t/tau_e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _return_parameter_names(num_pars):\n",
    "    \"\"\"\n",
    "    Function that returns the names of the parameters for writing to the dataframe after the fit.\n",
    "    num_pars is the number of parameters in the fit. 1,3,5,7,9 are the num_params that constrain the fit.\n",
    "    while the even numbers are the parameters for the functions that don't constrain the fits.\n",
    "    \n",
    "    LICENSE INFO:\n",
    "    MIT License\n",
    "\n",
    "    Copyright (c) 2017 Po-chia Chen\n",
    "\n",
    "    Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "    of this software and associated documentation files (the \"Software\"), to deal\n",
    "    in the Software without restriction, including without limitation the rights\n",
    "    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "    copies of the Software, and to permit persons to whom the Software is\n",
    "    furnished to do so, subject to the following conditions:\n",
    "\n",
    "    The above copyright notice and this permission notice shall be included in all\n",
    "    copies or substantial portions of the Software.\n",
    "\n",
    "    THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "    SOFTWARE.\n",
    "\n",
    "    \"\"\"\n",
    "    if num_pars==1:\n",
    "        return ['C_a', 'tau_a']\n",
    "    elif num_pars==2:\n",
    "         return ['C_a', 'tau_a']\n",
    "    elif num_pars==3:\n",
    "         return ['C_a', 'tau_a', 'tau_b']\n",
    "    elif num_pars==4:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b']\n",
    "    elif num_pars==5:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b', 'tau_g']\n",
    "    elif num_pars==6:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b', 'C_g', 'tau_g']\n",
    "    elif num_pars==7:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b', 'C_g', 'tau_g', 'tau_d']\n",
    "    elif num_pars==8:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b', 'C_g', 'tau_g', 'C_d', 'tau_d']\n",
    "    elif num_pars==9:\n",
    "         return ['C_a', 'tau_a', 'C_b', 'tau_b', 'C_g', 'tau_g', 'C_d', 'tau_d', 'tau_e']\n",
    "    elif num_pars==10:\n",
    "         return [ 'C_a', 'tau_a', 'C_b', 'tau_b', 'C_g', 'tau_g', 'C_d', 'tau_d', 'C_e', 'tau_e']\n",
    "\n",
    "    return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_Expstyle_fit2(num_pars, x, y, dy=np.empty([]), tau_mem=50.):\n",
    "    \"\"\"\n",
    "    Performs the exponential fit on the function defined by num_pars using scipy optimize curve fit.\n",
    "    Provides initial guesses for the amplitudes and the correlation times.\n",
    "    Takes the number of parameters, x values, y values, error in the y (dy), and tau_mem.\n",
    "    Tau_mem to help scale the initial guesses \n",
    "    Can also be set to np.inf if you want no bounds. \n",
    "    \n",
    "    Returns, the Chi-squared value of the fit to the model along with the parameter values (popt),\n",
    "    the parameter error (popv) and the model itself.\n",
    "    \"\"\"\n",
    "    \n",
    "    b1_guess = y[0]/num_pars/2 \n",
    "    t1_guess = [tau_mem/1280.0, tau_mem/640.0, tau_mem/64.0, tau_mem/8.0]\n",
    "    \n",
    "    if num_pars==1:\n",
    "        func=func_exp_decay1\n",
    "        guess=(t1_guess[2])\n",
    "        bound=(0.,np.inf)\n",
    "    elif num_pars==2:\n",
    "        func=func_exp_decay2\n",
    "        guess=(b1_guess, t1_guess[2])\n",
    "        bound=([0.0, x[0]],[1., np.inf])\n",
    "    elif num_pars==3:\n",
    "        func=func_exp_decay3\n",
    "        guess=(b1_guess, t1_guess[3], t1_guess[2])\n",
    "        bound=([0.0,x[0],x[0]],[1., np.inf, np.inf])\n",
    "    elif num_pars==4:\n",
    "        func=func_exp_decay4\n",
    "        guess=(b1_guess, t1_guess[3], b1_guess, t1_guess[2])\n",
    "        bound=([0.0, x[0], 0.0, x[0]],[1., np.inf, 1., np.inf])\n",
    "    elif num_pars==5:\n",
    "        func=func_exp_decay5\n",
    "        guess=(b1_guess, t1_guess[3], b1_guess, t1_guess[2], t1_guess[1])\n",
    "        bound=([0.0, x[0], 0.0, x[0],x[0]],[1., np.inf, 1., np.inf, np.inf])\n",
    "    elif num_pars==6:\n",
    "        func=func_exp_decay6\n",
    "        guess=(b1_guess, t1_guess[3], b1_guess, t1_guess[2], b1_guess, t1_guess[1])\n",
    "        bound=([0.0, x[0], 0.0, x[0], 0.0, x[0]],[1., np.inf, 1., np.inf, 1., np.inf])\n",
    "    elif num_pars==7:\n",
    "        func=func_exp_decay7\n",
    "        guess=(b1_guess, t1_guess[2], b1_guess, t1_guess[1], b1_guess, t1_guess[0],\n",
    "               t1_guess[3])\n",
    "        bound=([0.0, x[0], 0.0, x[0], 0.0, x[0], x[0]],[1., np.inf, 1., np.inf, 1., np.inf, np.inf])\n",
    "    elif num_pars==8:\n",
    "        func=func_exp_decay8\n",
    "        guess=(b1_guess, t1_guess[3], b1_guess, t1_guess[2], b1_guess, t1_guess[1],\n",
    "               b1_guess, t1_guess[0])\n",
    "        bound=([0.0, x[0], 0.0, x[0], 0.0, x[0], 0.0, x[0]],[1., np.inf, 1., np.inf, 1., np.inf, 1., np.inf])\n",
    "\n",
    "    if dy != []:\n",
    "        popt, popv = curve_fit(func, x, y, p0=guess, sigma=dy, bounds=bound, method='trf', loss='soft_l1')\n",
    "    else:\n",
    "        popt, popv = curve_fit(func, x, y, p0=guess, bounds=bound, loss='soft_l1')\n",
    "\n",
    "    ymodel=[ func(x[i], *popt) for i in range(len(x)) ]\n",
    "    #print ymodel\n",
    "\n",
    "    bExceed=_bound_check(func, popt)\n",
    "    if bExceed:\n",
    "        print >> sys.stderr, \"= = = WARNING, curve fitting in do_LSstyle_fit returns a sum>1.//\"\n",
    "        return 9999.99, popt, np.sqrt(np.diag(popv)), ymodel\n",
    "    else:\n",
    "        return calc_chi(y, ymodel, dy), popt, popv, ymodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findbest_Expstyle_fits2(x, y, taum=150.0, dy=[], bPrint=True, par_list=[2,3,5,7], threshold=1.0):\n",
    "    \"\"\"\n",
    "        Function tries to find the best set of parameters to describe the correlation fucntion for each residues\n",
    "        Takes the x,y values for the fit and the errors, dy. par_list is the number of parameters to check,\n",
    "        threshold is the cutoff for the chi2. This is the old way of checking, but can be re-implemented.\n",
    "        Runs the fit for a given parameter by calling do_Expstyle_fit3. The initial fit is chosen, but \n",
    "        subsequent fits are chosen with a strict criteria based on the ratio of the number of parameters from \n",
    "        the current best fit and the latest fit.\n",
    "        \n",
    "        Returns the chi^2, names of the parameters, parameters, errors, model, and covariance matrix of the best fit.\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    chi_min=np.inf\n",
    "    # Search forwards\n",
    "    print('Starting New Fit')\n",
    "    for npars in par_list:\n",
    "        print(npars)\n",
    "        names = _return_parameter_names(npars)\n",
    "        try:\n",
    "            chi, params, covarMat, ymodel = do_Expstyle_fit2(npars, x, y, dy, taum)\n",
    "        except:\n",
    "            print(\" ...fit returns an error! Continuing.\")\n",
    "            break\n",
    "        bBadFit=False\n",
    "        \n",
    "        errors = np.sqrt(np.diag(covarMat))\n",
    "        \n",
    "        step_check = 0\n",
    "        while step_check < npars:\n",
    "                \n",
    "            ## Check the error to make sure there is no overfitting\n",
    "            chkerr = errors[step_check]/params[step_check]\n",
    "            if (chkerr>0.10):\n",
    "                print( \" --- fit shows overfitting with %d parameters.\" % npars)\n",
    "                print(  \"  --- Occurred with parameter %s: %g +- %g \" % (names[step_check], params[step_check],\n",
    "                                                                         errors[step_check]))\n",
    "                bBadFit=True\n",
    "                break\n",
    "            \n",
    "            step_check += 1\n",
    "        \n",
    "        ## Chi^2 model fitting check. \n",
    "        ## SclChk can be increased to make it easier to fit higher order models, or lower for a stronger criteria\n",
    "        ## First model check is always set to 1.0 so its accepted\n",
    "        SclChk = 0.5\n",
    "        chi_check = chi/chi_min\n",
    "        if npars == par_list[0]:\n",
    "            threshold = 1.0\n",
    "        else:\n",
    "            threshold = (1-npar_min/npars)*SclChk\n",
    "                \n",
    "        print(\"--- The chi_check for {} parameters is {}\".format(npars, chi_check))\n",
    "        print(\"--- The threshold for this check is {}\".format(threshold))\n",
    "        if (not bBadFit) and (chi/chi_min < threshold):\n",
    "            chi_min=chi ; par_min=params ; err_min=errors ; npar_min=npars ; ymod_min=ymodel; covar_min = covarMat;\n",
    "        else:\n",
    "            break; \n",
    "            \n",
    "    tau_min = par_min[1::2]\n",
    "    sort_tau = np.argsort(tau_min)[::-1]\n",
    "    nsort_params = np.array([[2*tau_ind, 2*tau_ind+1] for tau_ind in sort_tau]).flatten()\n",
    "    \n",
    "    err_min = err_min[nsort_params] \n",
    "    par_min = par_min[nsort_params]\n",
    "    sort_covarMat = covar_min[:,nsort_params][nsort_params]\n",
    "    names = _return_parameter_names(npar_min)    \n",
    "    \n",
    "    if bPrint:       \n",
    "        print( \"= = Found %d parameters to be the minimum necessary to describe curve: chi(%d) = %g vs. chi(%d) = %g)\" % (npar_min, npar_min, chi_min,  npars, chi))\n",
    "        print( \"Parameter %d %s: %g +- %g \" % (npar_min, len(names), len(par_min), len(err_min)))\n",
    "        for i in range(npar_min):\n",
    "            print( \"Parameter %d %s: %g +- %g \" % (i, names[i], par_min[i], err_min[i]))\n",
    "        print('\\n')   \n",
    "    return chi_min, names, par_min, err_min, ymod_min, sort_covarMat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitstoDF(resnames, chi_list, pars_list, errs_list, names_list):\n",
    "    ## Set Up columns indices and names for the data frame\n",
    "    \"\"\"\n",
    "    Function that takes the residue names, chi^2, parameters, errors and names of the fits and returns a data frame\n",
    "    of the parameters.\n",
    "    \"\"\"\n",
    "    mparnames = _return_parameter_names(8) ## Always return the longest possible number of \n",
    "    mtau_names = np.array(mparnames)[1::2]\n",
    "    mc_names = np.array(mparnames)[::2]\n",
    "    colnames = np.array(['Resname','NumExp'])\n",
    "    tau_errnames = np.array([[c,\"{}_err\".format(c)] for c in mtau_names]).flatten()\n",
    "    mc_errnames = np.array([[c, \"{}_err\".format(c)] for c in mc_names]).flatten()\n",
    "    colnames = np.hstack([colnames,mc_errnames])\n",
    "    colnames = np.hstack([colnames,tau_errnames])\n",
    "    colnames = np.hstack([colnames,np.array(['Chi_Fit'])])\n",
    "    FitDF = pd.DataFrame(index=np.arange(len(pars_list)), columns=colnames).fillna(0.0)\n",
    "    FitDF['Resname'] = resnames\n",
    "    FitDF['Chi_Fit'] = chi_list\n",
    "    \n",
    "    for i in range(len(pars_list)):\n",
    "        npar = len(pars_list[i])\n",
    "        if (npar%2)==1:\n",
    "            ccut = npar-2\n",
    "            tau_f, terr = pars_list[i][1:ccut+1:2], errs_list[i][1:ccut+1:2]\n",
    "            tau_f = np.hstack([tau_f, pars_list[i][-1]])\n",
    "            terr = np.hstack([terr, errs_list[i][-1]])\n",
    "            sort_tau = np.argsort(tau_f)\n",
    "            coeff, cerr= pars_list[i][0:ccut:2], errs_list[i][0:ccut:2]\n",
    "            Clast = 1; Clasterr = 0.0;\n",
    "            for n,m in zip(coeff, cerr):\n",
    "                Clast -= n\n",
    "                Clasterr += m\n",
    "            \n",
    "            coeff = np.hstack([coeff, np.array(Clast)])\n",
    "            cerr = np.hstack([cerr, np.array(Clasterr)])\n",
    "    \n",
    "            tne = np.array([[c,\"{}_err\".format(c)] for c in mparnames[1:npar+1:2]]).flatten()\n",
    "            cne = np.array([[c, \"{}_err\".format(c)] for c in mparnames[0:npar:2]]).flatten()\n",
    "                \n",
    "        else:\n",
    "            tau_f, terr = pars_list[i][1::2], errs_list[i][1::2] \n",
    "            coeff, cerr= pars_list[i][0::2], errs_list[i][0::2]\n",
    "            sort_tau = np.argsort(tau_f)[::-1]\n",
    "            tne = np.array([[c,\"{}_err\".format(c)] for c in names_list[i][1::2]]).flatten()\n",
    "            cne = np.array([[c, \"{}_err\".format(c)] for c in names_list[i][0::2]]).flatten()\n",
    "    \n",
    "        NumExp=np.array(len(tau_f))\n",
    "        tau_err = np.array([[t,e] for t,e in zip(tau_f[sort_tau],terr[sort_tau])]).flatten()\n",
    "        c_err = np.array([[c,e] for c,e in zip(coeff[sort_tau], cerr[sort_tau])]).flatten()\n",
    "        namesarr = np.hstack([np.array('NumExp'),cne,tne])\n",
    "        valarr = np.hstack([NumExp,c_err,tau_err])\n",
    "    \n",
    "        FitDF.loc[i,namesarr] = valarr\n",
    "        \n",
    "    FitDF['AUC_a'] = FitDF.C_a*FitDF.tau_a; FitDF['AUC_b'] = FitDF.C_b*FitDF.tau_b; \n",
    "    FitDF['AUC_g'] = FitDF.C_g*FitDF.tau_g; FitDF['AUC_d'] = FitDF.C_d*FitDF.tau_d;\n",
    "    FitDF['AUC_Total'] = FitDF[['AUC_a','AUC_b','AUC_g','AUC_d']].sum(axis=1)\n",
    "    \n",
    "    return FitDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitCorrF(CorrDF, dCorrDF, tau_mem, pars_l, fixfit=False, threshold=1.0):\n",
    "    \"\"\"\n",
    "        Input Variables:\n",
    "            CorrDF: Dataframe containing the correlation functions. Columns are the NH-bond vectors, rows are timesteps. \n",
    "            dCorrDF: Error in the correlation function at time t\n",
    "            tau_mem: Cut-Off time to remove noise at the tail of the correlation function \n",
    "            pars_l : parameters list. \n",
    "            fixfit : Boolean to decide if you want to use a specific exponential function \n",
    "        Main function to fit the correlation function. \n",
    "        Loops over all residues with N-H vectors and calculates the fit, appends the best fit from findbest_Expstyle_fits2.\n",
    "        Passes the set of lists to fitstoDF to return a data frame of the best fits for each residue. \n",
    "        \n",
    "        Takes the correlation function CorrDF and errors in the correlation function, maximum tau mem to cut correlation\n",
    "        function off from, the list of parameters you want to fit too. If you don't want to test the fit and use \n",
    "        a fixed parameter set, set fixfit to True and pass a list of length 1 into pars_l.\n",
    "    \"\"\"\n",
    "    NH_Res = CorrDF.columns\n",
    "    chi_list=[] ; names_list=[] ; pars_list=[] ; errs_list=[] ; ymodel_list=[]; covarMat_list = [];\n",
    "    for i in CorrDF.columns:\n",
    "        \n",
    "        tstop = np.where(CorrDF.index.values==tau_mem)[0][0]\n",
    "            \n",
    "        x = CorrDF.index.values[:tstop]\n",
    "        y = CorrDF[i].values[:tstop]\n",
    "        dy = dCorrDF[i].values[:tstop]\n",
    "        ## If there is no error provided i.e. no std. dev. over correlation functions is provided then set dy to empty set\n",
    "        if np.all(np.isnan(dy)):\n",
    "            dy = []\n",
    "        \n",
    "        ## if not fixfit then find find the best expstyle fit. Otherwise force the fit to nparams \n",
    "        if (not fixfit)&(len(pars_l)>1):\n",
    "            print(\"Finding the best fit for residue {}\".format(i))\n",
    "            \n",
    "            chi, names, pars, errs, ymodel, covarMat = findbest_Expstyle_fits2(x, y, tau_mem, dy,  \n",
    "                                                            par_list=pars_l, threshold=threshold)\n",
    "        \n",
    "        elif (fixfit)&(len(pars_l)==1):\n",
    "            print(\"Performing a fixed fit for {} exponentials\".format(int(pars_l[0]/2)))\n",
    "            \n",
    "            chi, pars, covarMat, ymodel = do_Expstyle_fit2(pars_l[0], x, y, dy, tau_mem)\n",
    "            names = _return_parameter_names(len(pars))\n",
    "            errs = np.sqrt(np.diag(covarMat))\n",
    "            \n",
    "        else:\n",
    "            print(\"The list of parameters is empty. Breaking out.\")\n",
    "            break;\n",
    "            \n",
    "        chi_list.append(chi)\n",
    "        names_list.append(names)\n",
    "        pars_list.append(pars)\n",
    "        errs_list.append(errs)\n",
    "        ymodel_list.append(ymodel)\n",
    "        \n",
    "        \n",
    "    FitDF = fitstoDF(NH_Res, chi_list, pars_list, errs_list, names_list)\n",
    "    \n",
    "    return FitDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def J_direct_transform(om, consts, taus):\n",
    "    \n",
    "    \"\"\"\n",
    "        Calculation of the spectral density from the parameters of the fit by direct fourier transform\n",
    "    \"\"\"\n",
    "    ## Calculation for the direct spectral density \n",
    "    ndecay=len(consts) ; noms=1;###lnden(om)\n",
    "    Jmat = np.zeros( (ndecay, noms ) )\n",
    "    for i in range(ndecay):\n",
    "        Jmat[i] = consts[i]*(taus[i]*1e-9)/(\n",
    "            1 + np.power((taus[i]*1e-9)*(om),2.))\n",
    "    return Jmat.sum(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_NMR_Relax(J, fdd, fcsa, gammaH, gammaN):\n",
    "    \"\"\"\n",
    "        Function to calculate the R1, R2 and NOE from the spectral densities and the physical parameters for the \n",
    "        dipole-dipole and csa contributions, fdd and fcsa. \n",
    "    \"\"\"\n",
    "    R1 = fdd * (J['Diff'] + 3*J['15N'] + 6*J['Sum']) + fcsa * J['15N']\n",
    "    \n",
    "    R2 = (0.5 * fdd * (4*J['0'] + J['Diff'] + 3*J['15N'] + 6*J['1H'] + 6*J['Sum']) \n",
    "          + (1./6.) * fcsa*(4*J['0'] + 3*J['15N']) )\n",
    "    \n",
    "    NOE = 1 + ((fdd*gammaH)/(gammaN*R1))*(6*J['Sum'] - J['Diff'])\n",
    "    \n",
    "    return R1, R2, NOE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Begin Implementation of Code:\n",
    "## Definition of global file locations\n",
    "    1. Notebook can be run in the local directory, in which case, skip over the first cell\n",
    "    2. File locations of trajectories to be loaded using mdtraj for calculation of N-H bond vectors. These should be changed by the user.\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Global Variables for the calculation of the NH Vecs and the correlation functions\n",
    "FileLoc = \"\" ## Main Directory Location\n",
    "\n",
    "RUN = [\"Run{}\".format(i) for i in range(1,5)]\n",
    "JOBS = ['PROD1','PROD2','PROD3']\n",
    "## For use if replicate trajectories are stored as follows  \n",
    "TRAJLIST_LOC = [\"{}/Analysis/{}\".format(J,R) for J in JOBS for R in RUN]\n",
    "FTOPN = \"Q15.gro\" ## Name of topology for the trajectory\n",
    "FMDN = \"Q15.noH20.xtc\"  ## Name of the trajectory, should be centered and stripped of solute"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of physical constants and parameters\n",
    "    1. Several parameters should be changed if necessary\n",
    "        a. B0 --> Set to experimental magnetic field you want to compare against\n",
    "        b. dSigmaN --> -170e-6 is a well-established value, but can be changed\n",
    "    2. Units are in s in the parameters, but the timesteps should be in ns. Converted in J_direct_transform.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Parameters and Physical Constants for calculation of Relaxation Rates\n",
    "\n",
    "H_gyro = 2*np.pi*42.57748*1e6     ## Gyromagnetic Ratio: Hydrogen ([rad]/[s][T]) \n",
    "N_gyro = -2*np.pi*4.317267*1e6     ## Gyromagnetic Ratio: Nitrogen ([rad]/[s][T])\n",
    "B0 = 18.8                        ## Field Strength = 18.8 [Teslas] \n",
    "\n",
    "## Need 5 Frequencies: ## J[0], J[wH], J[wN], J[wH-wN], J[wH+wN]\n",
    "Larmor1H = H_gyro*B0              ## Larmor Frequency: Hydrogen ([rad]/[s])\n",
    "Larmor15N = N_gyro*B0             ## Larmor Frequency: Hydrogen ([rad]/[s])\n",
    "omDiff = Larmor1H - Larmor15N    ## Diff in Larmor Frequencies of Spin IS\n",
    "omSum  = Larmor1H + Larmor15N    ## Sum of Larmor Frequencies of Spin IS\n",
    "\n",
    "mu_0 = 4*np.pi*1e-7    ; ## Permeability of Free Space: ([H]/[m]) \n",
    "hbar = 1.0545718e-34  ; ## Reduced Plank's constant: [J] * [s] = [kg] * [m^2] * [s^-1] \n",
    "\n",
    "R_NH = 1.02e-10                     ## distance between N-H atoms in Angstroms\n",
    "dSigmaN = -170e-6               ##  CSA of the S-spin atom\n",
    "\n",
    "FDD = (1./10.)*np.power((mu_0*hbar*H_gyro*N_gyro)/(4*np.pi*np.power(R_NH,3)),2)\n",
    "#FCSA = 498637299.69233465\n",
    "FCSA = (2.0/15.0)*(Larmor15N**2)*(dSigmaN**2)        ## CSA factor "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load trajectories and calculate the NH-Vecs in the laboratory frame \n",
    "### Skip to calculation of correlation functions if already performed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Change directory to examples to test code\n",
    "%cd EXAMPLES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate the NHVecs; Can be adapted to loop over multiple trajectories using TRAJLIST_LOC\n",
    "NHVecs = []\n",
    "start=0; end=-1;  ## \n",
    "NHV = calc_NHVecs(FMDN, FTOPN, start, end)\n",
    "NHVecs.append(NHV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 10 ## timestep of simulations: (ps)\n",
    "tau_split = np.array(NHVecs).shape[1]*dt ## Number of snapshots to calculate the correlation function over.\n",
    "\n",
    "## Split the vecs based off the tau_split you want and the time step. \n",
    "vecs_split = split_NHVecs(NHVecs, dt, tau_split) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate the correlation functions and the standard deviation in the correlation function.\n",
    "## Save the correlation functions in a dataframe and then to a csv file for later use.\n",
    "Ct, dCt = calc_Ct(vecs_split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Convert to dataframe with index set as timesteps in ns\n",
    "CtOutFname = 'NH_Ct.csv'\n",
    "dCtOutFname = 'NH_dCt.csv'\n",
    "CtDF = pd.DataFrame(Ct, index = np.arange(1, Ct.shape[0]+1)*dt/1000) \n",
    "dCtDF = pd.DataFrame(dCt, index = np.arange(1, dCt.shape[0]+1)*dt/1000)\n",
    "CtDF.to_csv(CtOutFname)\n",
    "dCtDF.to_csv(dCtOutFname)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Begin fitting of the correlation functions\n",
    "    1. Load the correlation functions from before \n",
    "    2. Calculate the correlation functions\n",
    "        a. For a single exponential model, fixfit=True\n",
    "        b. Find the best exponential model, fixfit=False (default)\n",
    "    3. Pass the fitted parameters for each residue to calculate the spectral density\n",
    "    4. Calculate the NMR relaxation parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load the correlation functions from the saved csv files \n",
    "CtInName = 'NH_Ct.csv'\n",
    "dCtInName = 'NH_dCt.csv'\n",
    "CtDF = pd.read_csv(CtInName, index_col=0)\n",
    "dCtDF = pd.read_csv(dCtInName, index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau_mem=2.5 ## Cut off to remove noise from the tail of the correlation function in the fit (ns) \n",
    "fixfit = True ## find the best model\n",
    "parameters_list = [4] ## for fixfit = False \n",
    "thresh=1.0 ## \n",
    "FitDF = fitCorrF(CtDF, dCtDF, tau_mem, parameters_list, fixfit, thresh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate spectral density from the FitDF by calling the J_direct_transform function for each of the 5 frequencies.\n",
    "## Loop over the rows of the FitDF dataframe from fitCorrF function and calcuate the spectral densities.\n",
    "## Save the spectral densities to a dictionary and append to a list.\n",
    "Jarr = []\n",
    "\n",
    "for i,fit in FitDF.iterrows():\n",
    "    c = fit[['C_a','C_b','C_g','C_d']].values\n",
    "    t = fit[['tau_a','tau_b','tau_g','tau_d']].values\n",
    "    Jdict = {'0':0, '1H':0,'15N':0,'Sum':0,'Diff':0} \n",
    "    J0 = J_direct_transform(0, c, t)\n",
    "    JH = J_direct_transform(Larmor1H, c, t)\n",
    "    JN = J_direct_transform(Larmor15N, c, t)\n",
    "    JSum = J_direct_transform(omSum,  c, t)\n",
    "    JDiff = J_direct_transform(omDiff,  c, t)\n",
    "    Jdict['1H'] = JH ; Jdict['15N'] = JN; Jdict['0'] = J0; \n",
    "    Jdict['Sum'] = JSum; Jdict['Diff'] = JDiff;\n",
    "    Jarr.append(Jdict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate NMR relaxation parameters for each residue by calling calc_NMR_relax \n",
    "## Save the T1, T2 and NOE parameters to a dataframe\n",
    "NMRRelaxDF = pd.DataFrame(np.zeros((len(Jarr),3)),index=range(1,len(Jarr)+1), columns=['T1','T2','NOE'])\n",
    "for index in range(1,len(Jarr)+1):\n",
    "    r1, r2, noe = calc_NMR_Relax(Jarr[index-1], FDD, FCSA, H_gyro, N_gyro)\n",
    "    NMRRelaxDF.loc[index,'T1'] = 1/r1; \n",
    "    NMRRelaxDF.loc[index,'T2'] = 1/r2; \n",
    "    NMRRelaxDF.loc[index,'NOE'] = noe; \n",
    "\n",
    "NMRRelaxDF['Resname'] = FitDF['Resname'].values\n",
    "NMRRelaxDF['RESNUM'] = NMRRelaxDF['Resname'].str.extract('([0-9]+)',expand=False).astype('int')+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Merge the NMR relaxation dataframes with the FitDF dataframe\n",
    "FitRelaxDF = FitDF.merge(NMRRelaxDF, how='left', left_on='Resname',right_on='Resname').set_index(NMRRelaxDF.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Save FitRelaxDF to a csv file\n",
    "FitRelaxDF.to_csv('NMRRelaxtionDF.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}