hexamer_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import glob\n",
    "import time\n",
    "from matplotlib.colors import LogNorm\n",
    "import h5py\n",
    "import matplotlib.pylab as pylab\n",
    "import scipy.spatial.distance as scp\n",
    "import pyemma \n",
    "import tqdm\n",
    "from mpl_toolkits import mplot3d\n",
    "import extq\n",
    "import warnings\n",
    "import matplotlib as mpl\n",
    "from scipy import stats\n",
    "\n",
    "import matplotlib.transforms as mtransforms\n",
    "%matplotlib notebook "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pmf(x,y,weights,kT,xedges,yedges,\n",
    "                title=None,xlabel='x',ylabel='y',label=None,ax=None,vmin=None,vmax=None,\n",
    "                colorbar=True,cmap='cividis',contours=True,pcolor=True,clines=np.linspace(0,50,51),**kwargs):\n",
    "    \n",
    "    \"\"\"Compute and plot a 2D PMF\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate a PMF\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    kt: float\n",
    "        Value for Boltzmann constant and temperature (choice of units)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "    plotting parameters\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    \n",
    "    # bin values\n",
    "    hist, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = -kT * np.log(hist)\n",
    "    hist -= np.nanmin(hist)\n",
    "    if vmin is not None or vmax is not None:\n",
    "        if vmin is not None:\n",
    "            hist[hist < vmin] = np.nan\n",
    "        if vmax is not None:\n",
    "            hist[hist > vmax] = np.nan\n",
    "    # plot the PMF in bins\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    fig = ax.get_figure()\n",
    "    if pcolor:\n",
    "        pcm = ax.pcolormesh(xedges, yedges, hist.T, shading=\"flat\", cmap=cmap, **kwargs)\n",
    "        if colorbar:\n",
    "            cbar = fig.colorbar(pcm, ax=ax)\n",
    "        if label is not None:\n",
    "            cbar.set_label(label, rotation=-90, va=\"baseline\")\n",
    "    ax.set_xlim(xedges[0], xedges[-1])\n",
    "    ax.set_ylim(yedges[0], yedges[-1])\n",
    "    ax.set_xlabel(xlabel)\n",
    "    ax.set_ylabel(ylabel)\n",
    "    ax.set_title(title)\n",
    "    #Draw contours if desired\n",
    "    if contours:\n",
    "        centerx =(xedges[1:]+xedges[:-1])/2\n",
    "        centery =(yedges[1:]+yedges[:-1])/2\n",
    "        if pcolor:\n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1)\n",
    "        else: \n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=0.5)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_pmf(x,y,weights,kT,xedges,yedges):\n",
    "    \n",
    "    \"\"\"Compute  2D PMF (no plot)\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate a PMF\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    kt: float\n",
    "        Value for Boltzmann constant and temperature (choice of units)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    \n",
    "    # bin values\n",
    "    hist, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = -kT * np.log(hist)\n",
    "    hist -= np.nanmin(hist)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_3dpmf(x,y,z,weights,kT,xedges,yedges,zedges):\n",
    "    \"\"\"Compute  3D PMF (no plot)\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate a 3D PMF\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y, z : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    kt: float\n",
    "        Value for Boltzmann constant and temperature (choice of units)\n",
    "    xedges, yedges, zedges : (nx+1,) or (ny+1,) or (nz+1,) ndarray\n",
    "        The edges of the PMF \n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny, nz) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    x = _flatten(x)\n",
    "    y = _flatten(y)\n",
    "    z = _flatten(z)\n",
    "    weights = _flatten(weights)\n",
    "    \n",
    "    hist, _ = np.histogramdd(\n",
    "        (list(x), list(y), list(z)), weights=weights, bins=(xedges, yedges, zedges)\n",
    "    )\n",
    "    \n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = -kT * np.log(hist)\n",
    "    hist -= np.nanmin(hist)\n",
    "\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_avg(x, y, v, weights, xedges, yedges, label=None, ax=None,\n",
    "                title=None, xlabel='x', ylabel='y',\n",
    "                cmap='plasma',q=False, contours=False,clines=np.linspace(0,50,51),**kwargs):\n",
    "    \"\"\"Compute and plot a 2D average\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate an average\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    v : ndarray or list/tuple of ndarray\n",
    "        Value to be averaged for each frame\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "    plotting parameters\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    # bin values\n",
    "    numer, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=v.ravel() * weights.ravel()\n",
    "    )\n",
    "    denom, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = numer / denom\n",
    "        denom = -kT * np.log(denom)\n",
    "    denom -= np.nanmin(denom)\n",
    "    # plot\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    fig = ax.get_figure()\n",
    "    pcm = ax.pcolormesh(xedges, yedges, hist.T, shading=\"flat\", cmap=cmap,**kwargs)\n",
    "    ax.set_xlim(xedges[0], xedges[-1])\n",
    "    ax.set_ylim(yedges[0], yedges[-1])\n",
    "    ax.set_xlabel(xlabel)\n",
    "    ax.set_ylabel(ylabel)\n",
    "    ax.set_title(title)\n",
    "    cbar = fig.colorbar(pcm, ax=ax)\n",
    "    if label is not None:\n",
    "        cbar.set_label(label, rotation=-90, va=\"baseline\")\n",
    "    #Draw contours of the PMF if desired\n",
    "    if contours:\n",
    "        centerx =(xedges[1:]+xedges[:-1])/2\n",
    "        centery =(yedges[1:]+yedges[:-1])/2\n",
    "        if q:\n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1.0)\n",
    "            ax.contour(centerx, centery, hist.T, levels=[0.5], colors='purple', linestyles='solid', linewidths=2.5)\n",
    "        else:\n",
    "            ax.contour(centerx, centery, denom.T, levels=clines, colors='black', linestyles='solid', linewidths=0.7)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_avg(x, y, v, weights, xedges, yedges):\n",
    "    \"\"\"Compute a 2D average (no plot)\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate an average\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    v : ndarray or list/tuple of ndarray\n",
    "        Value to be averaged for each frame\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    # bin values\n",
    "    numer, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=v.ravel() * weights.ravel()\n",
    "    )\n",
    "    denom, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = numer / denom\n",
    "        denom = -kT * np.log(denom)\n",
    "    denom -= np.nanmin(denom)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.ndimage\n",
    "\n",
    "\n",
    "def kdesum2d(\n",
    "    x,\n",
    "    y,\n",
    "    w1,\n",
    "    w2,\n",
    "    *,\n",
    "    xmin=None,\n",
    "    xmax=None,\n",
    "    ymin=None,\n",
    "    ymax=None,\n",
    "    xstd=None,\n",
    "    ystd=None,\n",
    "    nx=100,\n",
    "    ny=100,\n",
    "    cut=4.0,\n",
    "    plot=False,\n",
    "    axis=None,\n",
    "    xlabel=None,\n",
    "    ylabel=None,\n",
    "    title=None,\n",
    "    norm=False,\n",
    "    pcolor=False,\n",
    "    cbar=True,\n",
    "    **kwargs,\n",
    "):\n",
    "    \"\"\"Compute a 2D kernel density estimate.\n",
    "\n",
    "    This function histograms the data, then uses a Gaussian filter to\n",
    "    approximate a kernel density estimate with a Gaussian kernel.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    w : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame. The output is the sum of these\n",
    "        values in each bin, after smoothing.\n",
    "    xmin, xmax, ymin, ymax : float, optional\n",
    "        Limits of kernel density estimate. If None, takes the min/max\n",
    "        of the data along the coordinate.\n",
    "    xstd, ystd : float, optional\n",
    "        Standard deviation of the Gaussian filter. If None, these are\n",
    "        set to (xmax - xmin) / nx and (ymax - ymin) / ny, respectively.\n",
    "        Increase this to smooth the results more.\n",
    "    nx, ny : int, optional\n",
    "        Number of bins in each dimension. This should be set as high as\n",
    "        reasonable, since xstd/ystd takes care of the smoothing.\n",
    "    cut : float, optional\n",
    "        Number of standard deviations at which to truncate the Gaussian\n",
    "        filter. The default, 4, usually doesn't need to be changed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    kde : (nx, ny) ndarray\n",
    "        Kernel density estimate, given as bins.\n",
    "    xedges : (nx+1,) ndarray\n",
    "        Bin edges along the x dimension.\n",
    "    yedges : (ny+1,) ndarray\n",
    "        Bin edges along the y dimension.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # flatten input to 1D arrays\n",
    "    x = _flatten(x)\n",
    "    y = _flatten(y)\n",
    "    w1 = _flatten(w1)\n",
    "    w2 = _flatten(w2)\n",
    "\n",
    "    # limits\n",
    "    _xmin = np.min(x)\n",
    "    _xmax = np.max(x)\n",
    "    _ymin = np.min(y)\n",
    "    _ymax = np.max(y)\n",
    "    if xmin is None:\n",
    "        xmin = _xmin\n",
    "    if xmax is None:\n",
    "        xmax = _xmax\n",
    "    if ymin is None:\n",
    "        ymin = _ymin\n",
    "    if ymax is None:\n",
    "        ymax = _ymax\n",
    "\n",
    "    # separation between grid points\n",
    "    xsep = (xmax - xmin) / nx\n",
    "    ysep = (ymax - ymin) / ny\n",
    "\n",
    "    # number of grid points to pad the boundaries,\n",
    "    # since the Gaussian filter extends beyond the boundaries\n",
    "    # usually overestimates the padding, but whatever\n",
    "    ax = max(0, int(np.ceil((xmin - _xmin) / xsep + 1e-6)))\n",
    "    bx = max(0, int(np.ceil((_xmax - xmax) / xsep + 1e-6)))\n",
    "    ay = max(0, int(np.ceil((ymin - _ymin) / ysep + 1e-6)))\n",
    "    by = max(0, int(np.ceil((_ymax - ymax) / ysep + 1e-6)))\n",
    "\n",
    "    # output bin edges\n",
    "    xedges = np.linspace(xmin, xmax, nx + 1)\n",
    "    yedges = np.linspace(ymin, ymax, ny + 1)\n",
    "\n",
    "    # bin edges, with the added padding\n",
    "    xedges_padded = np.concatenate(\n",
    "        [\n",
    "            xmin + xsep * np.arange(-ax, 0),\n",
    "            xedges,\n",
    "            xmax + xsep * np.arange(1, bx + 1),\n",
    "        ]\n",
    "    )\n",
    "    yedges_padded = np.concatenate(\n",
    "        [\n",
    "            ymin + ysep * np.arange(-ay, 0),\n",
    "            yedges,\n",
    "            ymax + ysep * np.arange(1, by + 1),\n",
    "        ]\n",
    "    )\n",
    "    assert np.allclose(xedges_padded[1:] - xedges_padded[:-1], xsep)\n",
    "    assert np.allclose(yedges_padded[1:] - yedges_padded[:-1], ysep)\n",
    "    assert xedges_padded[0] <= _xmin and _xmax <= xedges_padded[-1]\n",
    "    assert yedges_padded[0] <= _ymin and _ymax <= yedges_padded[-1]\n",
    "\n",
    "    # construct 2D histogram on padded edges\n",
    "    hist_padded1, _, _ = np.histogram2d(\n",
    "        x, y, weights=w1, bins=(xedges_padded, yedges_padded)\n",
    "    )\n",
    "    \n",
    "     # construct 2D histogram on padded edges\n",
    "    hist_padded2, _, _ = np.histogram2d(\n",
    "        x, y, weights=w2, bins=(xedges_padded, yedges_padded)\n",
    "    )\n",
    "\n",
    "    # Gaussian kernel parameters\n",
    "    if xstd is None:\n",
    "        xstd = xsep\n",
    "    if ystd is None:\n",
    "        ystd = ysep\n",
    "\n",
    "    # apply Gaussian filter to histogram\n",
    "    kde_padded1 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded1,\n",
    "        sigma=(xstd / xsep, ystd / ysep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "    \n",
    "    kde_padded2 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded2,\n",
    "        sigma=(xstd / xsep, ystd / ysep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "\n",
    "\n",
    "    # remove the padding\n",
    "    assert ax + nx + bx == kde_padded1.shape[0]\n",
    "    assert ay + ny + by == kde_padded1.shape[1]\n",
    "    kde1 = kde_padded1[ax : ax + nx, ay : ay + ny]\n",
    "    assert ax + nx + bx == kde_padded2.shape[0]\n",
    "    assert ay + ny + by == kde_padded2.shape[1]\n",
    "    kde2 = kde_padded2[ax : ax + nx, ay : ay + ny]\n",
    "\n",
    "    if plot:\n",
    "        xc = 0.5 * (xedges[1:] + xedges[:-1])\n",
    "        yc = 0.5 * (yedges[1:] + yedges[:-1])\n",
    "        vx=kde1\n",
    "        vy=kde2\n",
    "            # split direction and magnitude\n",
    "        v = np.sqrt(vx ** 2 + vy ** 2)\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.filterwarnings(\"ignore\")\n",
    "            if norm:\n",
    "                ux = vx / v\n",
    "                uy = vy / v\n",
    "            else:\n",
    "                ux = vx\n",
    "                uy = vy\n",
    "            logv = np.log(v)\n",
    "        ux[v == 0.0] = 0.0\n",
    "        uy[v == 0.0] = 0.0\n",
    "\n",
    "        cmap='cividis'\n",
    "        # plot\n",
    "        if axis is None:\n",
    "            axis = plt.gca()\n",
    "        fig = axis.get_figure()\n",
    "        if pcolor:\n",
    "            cmap='coolwarm'\n",
    "            vmax = np.max(ux)\n",
    "            pcm = axis.pcolormesh(xedges, yedges, ux.T, shading=\"flat\", cmap=cmap,**kwargs,vmin=-1*vmax,vmax=vmax)\n",
    "            cbar = fig.colorbar(pcm, ax=axis, format='%.1e')\n",
    "        else:\n",
    "            pcm = axis.quiver(\n",
    "                xc, yc, ux.T, uy.T, v.T, angles=\"xy\", pivot=\"middle\", cmap=cmap,width = 0.0055, edgecolor='black', **kwargs\n",
    "            )\n",
    "            if cbar:\n",
    "                cbar = fig.colorbar(pcm, ax=axis, format='%.1e')\n",
    "                #cbar.set_label(None, rotation=-90, va=\"baseline\")\n",
    "        axis.set_xlim(xedges[0], xedges[-1])\n",
    "        axis.set_ylim(yedges[0], yedges[-1])\n",
    "        axis.set_xlabel(xlabel)\n",
    "        axis.set_ylabel(ylabel)\n",
    "        axis.set_title(title)\n",
    "        \n",
    "        \n",
    "    return kde1, kde2, xedges, yedges\n",
    "\n",
    "\n",
    "def _flatten(a):\n",
    "    if isinstance(a, np.ndarray):\n",
    "        # avoid creating a new array (and using twice the memory)\n",
    "        return np.ravel(a)\n",
    "    else:\n",
    "        return np.ravel(np.concatenate(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.ndimage\n",
    "\n",
    "\n",
    "def kdesum2d1(\n",
    "    x,\n",
    "    y,\n",
    "    w1,\n",
    "    w2,\n",
    "    *,\n",
    "    xmin=None,\n",
    "    xmax=None,\n",
    "    ymin=None,\n",
    "    ymax=None,\n",
    "    xstd=None,\n",
    "    ystd=None,\n",
    "    nx=100,\n",
    "    ny=100,\n",
    "    cut=4.0,\n",
    "    plot=False,\n",
    "    axis=None,\n",
    "    xlabel=None,\n",
    "    ylabel=None,\n",
    "    title=None,\n",
    "    norm=False,\n",
    "    pcolor=False,\n",
    "    cbar=True,\n",
    "    label=None,\n",
    "    **kwargs,\n",
    "):\n",
    "    \"\"\"Compute a 2D kernel density estimate - with adjusted plotting parameters.\n",
    "\n",
    "    This function histograms the data, then uses a Gaussian filter to\n",
    "    approximate a kernel density estimate with a Gaussian kernel.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    w : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame. The output is the sum of these\n",
    "        values in each bin, after smoothing.\n",
    "    xmin, xmax, ymin, ymax : float, optional\n",
    "        Limits of kernel density estimate. If None, takes the min/max\n",
    "        of the data along the coordinate.\n",
    "    xstd, ystd : float, optional\n",
    "        Standard deviation of the Gaussian filter. If None, these are\n",
    "        set to (xmax - xmin) / nx and (ymax - ymin) / ny, respectively.\n",
    "        Increase this to smooth the results more.\n",
    "    nx, ny : int, optional\n",
    "        Number of bins in each dimension. This should be set as high as\n",
    "        reasonable, since xstd/ystd takes care of the smoothing.\n",
    "    cut : float, optional\n",
    "        Number of standard deviations at which to truncate the Gaussian\n",
    "        filter. The default, 4, usually doesn't need to be changed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    kde : (nx, ny) ndarray\n",
    "        Kernel density estimate, given as bins.\n",
    "    xedges : (nx+1,) ndarray\n",
    "        Bin edges along the x dimension.\n",
    "    yedges : (ny+1,) ndarray\n",
    "        Bin edges along the y dimension.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # flatten input to 1D arrays\n",
    "    x = _flatten(x)\n",
    "    y = _flatten(y)\n",
    "    w1 = _flatten(w1)\n",
    "    w2 = _flatten(w2)\n",
    "\n",
    "    # limits\n",
    "    _xmin = np.min(x)\n",
    "    _xmax = np.max(x)\n",
    "    _ymin = np.min(y)\n",
    "    _ymax = np.max(y)\n",
    "    if xmin is None:\n",
    "        xmin = _xmin\n",
    "    if xmax is None:\n",
    "        xmax = _xmax\n",
    "    if ymin is None:\n",
    "        ymin = _ymin\n",
    "    if ymax is None:\n",
    "        ymax = _ymax\n",
    "\n",
    "    # separation between grid points\n",
    "    xsep = (xmax - xmin) / nx\n",
    "    ysep = (ymax - ymin) / ny\n",
    "\n",
    "    # number of grid points to pad the boundaries,\n",
    "    # since the Gaussian filter extends beyond the boundaries\n",
    "    # usually overestimates the padding, but whatever\n",
    "    ax = max(0, int(np.ceil((xmin - _xmin) / xsep + 1e-6)))\n",
    "    bx = max(0, int(np.ceil((_xmax - xmax) / xsep + 1e-6)))\n",
    "    ay = max(0, int(np.ceil((ymin - _ymin) / ysep + 1e-6)))\n",
    "    by = max(0, int(np.ceil((_ymax - ymax) / ysep + 1e-6)))\n",
    "\n",
    "    # output bin edges\n",
    "    xedges = np.linspace(xmin, xmax, nx + 1)\n",
    "    yedges = np.linspace(ymin, ymax, ny + 1)\n",
    "\n",
    "    # bin edges, with the added padding\n",
    "    xedges_padded = np.concatenate(\n",
    "        [\n",
    "            xmin + xsep * np.arange(-ax, 0),\n",
    "            xedges,\n",
    "            xmax + xsep * np.arange(1, bx + 1),\n",
    "        ]\n",
    "    )\n",
    "    yedges_padded = np.concatenate(\n",
    "        [\n",
    "            ymin + ysep * np.arange(-ay, 0),\n",
    "            yedges,\n",
    "            ymax + ysep * np.arange(1, by + 1),\n",
    "        ]\n",
    "    )\n",
    "    assert np.allclose(xedges_padded[1:] - xedges_padded[:-1], xsep)\n",
    "    assert np.allclose(yedges_padded[1:] - yedges_padded[:-1], ysep)\n",
    "    assert xedges_padded[0] <= _xmin and _xmax <= xedges_padded[-1]\n",
    "    assert yedges_padded[0] <= _ymin and _ymax <= yedges_padded[-1]\n",
    "\n",
    "    # construct 2D histogram on padded edges\n",
    "    hist_padded1, _, _ = np.histogram2d(\n",
    "        x, y, weights=w1, bins=(xedges_padded, yedges_padded)\n",
    "    )\n",
    "    \n",
    "     # construct 2D histogram on padded edges\n",
    "    hist_padded2, _, _ = np.histogram2d(\n",
    "        x, y, weights=w2, bins=(xedges_padded, yedges_padded)\n",
    "    )\n",
    "\n",
    "    # Gaussian kernel parameters\n",
    "    if xstd is None:\n",
    "        xstd = xsep\n",
    "    if ystd is None:\n",
    "        ystd = ysep\n",
    "\n",
    "    # apply Gaussian filter to histogram\n",
    "    kde_padded1 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded1,\n",
    "        sigma=(xstd / xsep, ystd / ysep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "    \n",
    "    kde_padded2 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded2,\n",
    "        sigma=(xstd / xsep, ystd / ysep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "\n",
    "\n",
    "    # remove the padding\n",
    "    assert ax + nx + bx == kde_padded1.shape[0]\n",
    "    assert ay + ny + by == kde_padded1.shape[1]\n",
    "    kde1 = kde_padded1[ax : ax + nx, ay : ay + ny]\n",
    "    assert ax + nx + bx == kde_padded2.shape[0]\n",
    "    assert ay + ny + by == kde_padded2.shape[1]\n",
    "    kde2 = kde_padded2[ax : ax + nx, ay : ay + ny]\n",
    "\n",
    "    if plot:\n",
    "        xc = 0.5 * (xedges[1:] + xedges[:-1])\n",
    "        yc = 0.5 * (yedges[1:] + yedges[:-1])\n",
    "        vx=kde1\n",
    "        vy=kde2\n",
    "            # split direction and magnitude\n",
    "        v = np.sqrt(vx ** 2 + vy ** 2)\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.filterwarnings(\"ignore\")\n",
    "            if norm:\n",
    "                ux = vx / v\n",
    "                uy = vy / v\n",
    "            else:\n",
    "                ux = vx\n",
    "                uy = vy\n",
    "            logv = np.log(v)\n",
    "        ux[v == 0.0] = 0.0\n",
    "        uy[v == 0.0] = 0.0\n",
    "\n",
    "        cmap='cividis'\n",
    "        # plot\n",
    "        if axis is None:\n",
    "            axis = plt.gca()\n",
    "        fig = axis.get_figure()\n",
    "        if pcolor:\n",
    "            cmap='coolwarm'\n",
    "            vmax = np.max(ux)\n",
    "            pcm = axis.pcolormesh(xedges, yedges, ux.T, shading=\"flat\", cmap=cmap,**kwargs,vmin=-1*vmax,vmax=vmax)\n",
    "            cbar = fig.colorbar(pcm, ax=axis, format='%.1e')\n",
    "        else:\n",
    "            pcm = axis.quiver(\n",
    "                xc, yc, ux.T, uy.T, v.T, angles=\"xy\", pivot=\"middle\", cmap=cmap,width = 0.0055, edgecolor='black', **kwargs\n",
    "            )\n",
    "            if cbar:\n",
    "                cbformat = mpl.ticker.ScalarFormatter()   # create the formatter\n",
    "                cbformat.set_powerlimits((-7,-5)) \n",
    "                cbar = fig.colorbar(pcm, ax=axis, format=cbformat)\n",
    "                cbar.set_label(label, rotation=-90, va=\"baseline\", fontsize=16)\n",
    "                cbar.ax.tick_params(width=1.5, labelsize=14)\n",
    "        axis.set_xlim(xedges[0], xedges[-1])\n",
    "        axis.set_ylim(yedges[0], yedges[-1])\n",
    "        axis.set_xlabel(xlabel)\n",
    "        axis.set_ylabel(ylabel)\n",
    "        axis.set_title(title)\n",
    "        \n",
    "        \n",
    "    return kde1, kde2, xedges, yedges\n",
    "\n",
    "\n",
    "def _flatten(a):\n",
    "    if isinstance(a, np.ndarray):\n",
    "        # avoid creating a new array (and using twice the memory)\n",
    "        return np.ravel(a)\n",
    "    else:\n",
    "        return np.ravel(np.concatenate(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kdesum3d(\n",
    "    x,\n",
    "    y,\n",
    "    z,\n",
    "    w1,\n",
    "    w2,\n",
    "    w3,\n",
    "    *,\n",
    "    xmin=0,\n",
    "    xmax=100,\n",
    "    ymin=0,\n",
    "    ymax=100,\n",
    "    zmin=0,\n",
    "    zmax=100,\n",
    "    xstd=None,\n",
    "    ystd=None,\n",
    "    zstd=None,\n",
    "    nx=50,\n",
    "    ny=50,\n",
    "    nz=50,\n",
    "    cut=4.0,\n",
    "    plot=False,\n",
    "    axis=None,\n",
    "    xlabel=None,\n",
    "    ylabel=None,\n",
    "    zlabel=None,\n",
    "    title=None,\n",
    "    norm=False,\n",
    "    **kwargs,\n",
    "):\n",
    "    \"\"\"Compute a 3D kernel density estimate.\n",
    "\n",
    "    This function histograms the data, then uses a Gaussian filter to\n",
    "    approximate a kernel density estimate with a Gaussian kernel.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y, z : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    w1, w2 : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame. The output is the sum of these\n",
    "        values in each bin, after smoothing.\n",
    "    xmin, xmax, ymin, ymax, zmin, zmax : float, optional\n",
    "        Limits of kernel density estimate. If None, takes the min/max\n",
    "        of the data along the coordinate.\n",
    "    xstd, ystd, zstd : float, optional\n",
    "        Standard deviation of the Gaussian filter. If None, these are\n",
    "        set to (xmax - xmin) / nx and (ymax - ymin) / ny, respectively.\n",
    "        Increase this to smooth the results more.\n",
    "    nx, ny, nz : int, optional\n",
    "        Number of bins in each dimension. This should be set as high as\n",
    "        reasonable, since xstd/ystd takes care of the smoothing.\n",
    "    cut : float, optional\n",
    "        Number of standard deviations at which to truncate the Gaussian\n",
    "        filter. The default, 4, usually doesn't need to be changed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    kde : (nx, ny) ndarray\n",
    "        Kernel density estimate, given as bins.\n",
    "    xedges : (nx+1,) ndarray\n",
    "        Bin edges along the x dimension.\n",
    "    yedges : (ny+1,) ndarray\n",
    "        Bin edges along the y dimension.\n",
    "    zedges: (nz+1,) ndarray\n",
    "        Bin edges along the z dimension.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # flatten input to 1D arrays\n",
    "    x = _flatten(x)\n",
    "    y = _flatten(y)\n",
    "    z = _flatten(z)\n",
    "    w1 = _flatten(w1)\n",
    "    w2 = _flatten(w2)\n",
    "    w3 = _flatten(w3)\n",
    "\n",
    "    # limits\n",
    "    _xmin = np.min(x)\n",
    "    _xmax = np.max(x)\n",
    "    _ymin = np.min(y)\n",
    "    _ymax = np.max(y)\n",
    "    _zmin = np.min(z)\n",
    "    _zmax = np.max(z)\n",
    "    if xmin is None:\n",
    "        xmin = _xmin\n",
    "    if xmax is None:\n",
    "        xmax = _xmax\n",
    "    if ymin is None:\n",
    "        ymin = _ymin\n",
    "    if ymax is None:\n",
    "        ymax = _ymax\n",
    "    if zmin is None:\n",
    "        zmin = _zmin\n",
    "    if zmax is None:\n",
    "        zmax = _zmax\n",
    "\n",
    "    # separation between grid points\n",
    "    xsep = (xmax - xmin) / nx\n",
    "    ysep = (ymax - ymin) / ny\n",
    "    zsep = (zmax - zmin) / nz\n",
    "\n",
    "    # number of grid points to pad the boundaries,\n",
    "    # since the Gaussian filter extends beyond the boundaries\n",
    "    # usually overestimates the padding, but whatever\n",
    "    ax = max(0, int(np.ceil((xmin - _xmin) / xsep + 1e-6)))\n",
    "    bx = max(0, int(np.ceil((_xmax - xmax) / xsep + 1e-6)))\n",
    "    ay = max(0, int(np.ceil((ymin - _ymin) / ysep + 1e-6)))\n",
    "    by = max(0, int(np.ceil((_ymax - ymax) / ysep + 1e-6)))\n",
    "    az = max(0, int(np.ceil((zmin - _zmin) / zsep + 1e-6)))\n",
    "    bz = max(0, int(np.ceil((_zmax - zmax) / zsep + 1e-6)))\n",
    "\n",
    "    # output bin edges\n",
    "    xedges = np.linspace(xmin, xmax, nx + 1)\n",
    "    yedges = np.linspace(ymin, ymax, ny + 1)\n",
    "    zedges = np.linspace(zmin, zmax, nz + 1)\n",
    "\n",
    "    # bin edges, with the added padding\n",
    "    xedges_padded = np.concatenate(\n",
    "        [\n",
    "            xmin + xsep * np.arange(-ax, 0),\n",
    "            xedges,\n",
    "            xmax + xsep * np.arange(1, bx + 1),\n",
    "        ]\n",
    "    )\n",
    "    yedges_padded = np.concatenate(\n",
    "        [\n",
    "            ymin + ysep * np.arange(-ay, 0),\n",
    "            yedges,\n",
    "            ymax + ysep * np.arange(1, by + 1),\n",
    "        ]\n",
    "    )\n",
    "    zedges_padded = np.concatenate(\n",
    "        [\n",
    "            zmin + zsep * np.arange(-az, 0),\n",
    "            zedges,\n",
    "            zmax + zsep * np.arange(1, bz + 1),\n",
    "        ]\n",
    "    )\n",
    "    assert np.allclose(xedges_padded[1:] - xedges_padded[:-1], xsep)\n",
    "    assert np.allclose(yedges_padded[1:] - yedges_padded[:-1], ysep)\n",
    "    assert np.allclose(zedges_padded[1:] - zedges_padded[:-1], zsep)\n",
    "    assert xedges_padded[0] <= _xmin and _xmax <= xedges_padded[-1]\n",
    "    assert yedges_padded[0] <= _ymin and _ymax <= yedges_padded[-1]\n",
    "    assert zedges_padded[0] <= _zmin and _zmax <= zedges_padded[-1]\n",
    "\n",
    "    # construct 3D histogram on padded edges\n",
    "    hist_padded1, _ = np.histogramdd(\n",
    "        (list(x), list(y), list(z)), weights=w1, bins=(xedges_padded, yedges_padded, zedges_padded)\n",
    "    )\n",
    "    \n",
    "     # construct 3D histogram on padded edges\n",
    "    hist_padded2, _ = np.histogramdd(\n",
    "        (list(x), list(y), list(z)), weights=w2, bins=(xedges_padded, yedges_padded, zedges_padded)\n",
    "    )\n",
    "    \n",
    "    # construct 3D histogram on padded edges\n",
    "    hist_padded3, _ = np.histogramdd(\n",
    "        (list(x), list(y), list(z)), weights=w3, bins=(xedges_padded, yedges_padded, zedges_padded)\n",
    "    )\n",
    "\n",
    "\n",
    "    # Gaussian kernel parameters\n",
    "    if xstd is None:\n",
    "        xstd = xsep\n",
    "    if ystd is None:\n",
    "        ystd = ysep\n",
    "    if zstd is None:\n",
    "        zstd = zsep\n",
    "\n",
    "    # apply Gaussian filter to histogram\n",
    "    kde_padded1 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded1,\n",
    "        sigma=(xstd / xsep, ystd / ysep, zstd / zsep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "    \n",
    "    kde_padded2 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded2,\n",
    "        sigma=(xstd / xsep, ystd / ysep, zstd / zsep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "    kde_padded3 = scipy.ndimage.gaussian_filter(\n",
    "        hist_padded3,\n",
    "        sigma=(xstd / xsep, ystd / ysep, zstd / zsep),  # in units of grid points\n",
    "        mode=\"constant\",\n",
    "        truncate=cut,\n",
    "    )\n",
    "\n",
    "\n",
    "\n",
    "    # remove the padding\n",
    "    assert ax + nx + bx == kde_padded1.shape[0]\n",
    "    assert ay + ny + by == kde_padded1.shape[1]\n",
    "    assert az + nz + bz == kde_padded1.shape[2]\n",
    "    kde1 = kde_padded1[ax : ax + nx, ay : ay + ny, az : az + nz]\n",
    "    assert ax + nx + bx == kde_padded2.shape[0]\n",
    "    assert ay + ny + by == kde_padded2.shape[1]\n",
    "    assert az + nz + bz == kde_padded2.shape[2]\n",
    "    kde2 = kde_padded2[ax : ax + nx, ay : ay + ny, az : az + nz]\n",
    "    assert ax + nx + bx == kde_padded3.shape[0]\n",
    "    assert ay + ny + by == kde_padded3.shape[1]\n",
    "    assert az + nz + bz == kde_padded3.shape[2]\n",
    "    kde3 = kde_padded3[ax : ax + nx, ay : ay + ny, az : az + nz]\n",
    "\n",
    "    if plot:\n",
    "        xc = 0.5 * (xedges[1:] + xedges[:-1])\n",
    "        yc = 0.5 * (yedges[1:] + yedges[:-1])\n",
    "        zc = 0.5 * (zedges[1:] + zedges[:-1])\n",
    "        vx=kde1\n",
    "        vy=kde2\n",
    "        vz=kde3\n",
    "        # split direction and magnitude\n",
    "        v = np.sqrt(vx ** 2 + vy ** 2 + vz ** 2)\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.filterwarnings(\"ignore\")\n",
    "            if norm:\n",
    "                ux = vx / v\n",
    "                uy = vy / v\n",
    "                uz = vz / v\n",
    "            else:\n",
    "                ux = vx\n",
    "                uy = vy\n",
    "                uz = vz\n",
    "            logv = np.log(v)\n",
    "        ux[v == 0.0] = 0.0\n",
    "        uy[v == 0.0] = 0.0\n",
    "        uz[v == 0.0] = 0.0\n",
    "\n",
    "        cmap='cividis'\n",
    "        # plot\n",
    "        if axis is None:\n",
    "            axis = plt.gca()\n",
    "        fig = axis.get_figure()\n",
    "        start=None\n",
    "        stop=None\n",
    "        Skip=1\n",
    "        \n",
    "        X,Y,Z=np.meshgrid(xc,yc,zc)\n",
    "        y = X[0::Skip,0::Skip,0::Skip].reshape(np.product(X[0::Skip,0::Skip,0::Skip].shape))\n",
    "        x = Y[0::Skip,0::Skip,0::Skip].reshape(np.product(Y[0::Skip,0::Skip,0::Skip].shape))\n",
    "        z = Z[0::Skip,0::Skip,0::Skip].reshape(np.product(Z[0::Skip,0::Skip,0::Skip].shape))\n",
    "        \n",
    "        ux = _flatten(ux)\n",
    "        uy = _flatten(uy)\n",
    "        uz = _flatten(uz)\n",
    "        v = _flatten(v)\n",
    "    \n",
    "        for x1,y1,z1,u1,v1,w1,l in zip(x[start:stop],y[start:stop],z[start:stop],ux[start:stop],uy[start:stop],uz[start:stop],v[start:stop]):\n",
    "            axis.quiver(x1, y1, z1, u1, v1, w1, pivot = 'middle', length=min(8.0,l*1000000), normalize=True, arrow_length_ratio=0.5)\n",
    "            \n",
    "        axis.set_xlim(xedges[0], xedges[-1])\n",
    "        axis.set_ylim(yedges[0], yedges[-1])\n",
    "        axis.set_zlim(zedges[0], zedges[-1])\n",
    "        axis.set_xlabel(xlabel)\n",
    "        axis.set_ylabel(ylabel)\n",
    "        axis.set_zlabel(ylabel)\n",
    "        axis.set_title(title)\n",
    "        \n",
    "        \n",
    "    return kde1, kde2, kde3, xedges, yedges, zedges\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pmf1(x,y,weights,kT,xedges,yedges,\n",
    "                title=None,xlabel='x',ylabel='y',label=None,ax=None,vmin=None,vmax=None,\n",
    "                colorbar=True,cmap='cividis',contours=True,pcolor=True,clines=np.linspace(0,50,51),**kwargs):\n",
    "    \"\"\"Compute and plot a 2D PMF - with adjusted plotting parameters\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate a PMF\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    kt: float\n",
    "        Value for Boltzmann constant and temperature (choice of units)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "    plotting parameters\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    # bin values\n",
    "    hist, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = -kT * np.log(hist)\n",
    "    hist -= np.nanmin(hist)\n",
    "    if vmin is not None or vmax is not None:\n",
    "        if vmin is not None:\n",
    "            hist[hist < vmin] = np.nan\n",
    "        if vmax is not None:\n",
    "            hist[hist > vmax] = np.nan\n",
    "    # plot the PMF in bins\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    fig = ax.get_figure()\n",
    "    if pcolor:\n",
    "        pcm = ax.pcolormesh(xedges, yedges, hist.T, shading=\"flat\", cmap=cmap, **kwargs)\n",
    "        if colorbar:\n",
    "            cbar = fig.colorbar(pcm, ax=ax)\n",
    "        if label is not None:\n",
    "            cbar.set_label(label, rotation=-90, va=\"baseline\", fontsize=16)\n",
    "            cbar.ax.tick_params(width=1.5, labelsize=14)\n",
    "    ax.set_xlim(xedges[0], xedges[-1])\n",
    "    ax.set_ylim(yedges[0], yedges[-1])\n",
    "    #ax.set_aspect(\"equal\")\n",
    "    ax.set_xlabel(xlabel, fontsize=16)\n",
    "    ax.set_ylabel(ylabel, fontsize=16)\n",
    "    ax.set_title(title, fontsize=16)\n",
    "    ax.tick_params(width=1.5, labelsize=14)\n",
    "    #Draw contours if desired\n",
    "    if contours:\n",
    "        centerx =(xedges[1:]+xedges[:-1])/2\n",
    "        centery =(yedges[1:]+yedges[:-1])/2\n",
    "        if pcolor:\n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1)\n",
    "        else: \n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_avg1(x, y, v, weights, xedges, yedges, label=None, ax=None,\n",
    "                title=None, xlabel='x', ylabel='y',\n",
    "                cmap='plasma',q=False, contours=False,clines=np.linspace(0,50,51),**kwargs):\n",
    "    \"\"\"Compute and plot a 2D average - with adjusted plotting parameters\n",
    "\n",
    "    This function histograms the data, weighted by the change of measure, \n",
    "    to generate an average\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : ndarray or list/tuple of ndarray\n",
    "        Coordinates of each frame.\n",
    "    v : ndarray or list/tuple of ndarray\n",
    "        Value to be averaged for each frame\n",
    "    weights : ndarray or list/tuple of ndarray\n",
    "        Weight or value of each frame (change of measure)\n",
    "    xedges, yedges : (nx+1,ny+1) ndarray\n",
    "        The edges of the PMF \n",
    "    plotting parameters\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    hist : (nx, ny) ndarray\n",
    "        Histogram of the 2D PMF\n",
    "    \"\"\"\n",
    "    # bin values\n",
    "    numer, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=v.ravel() * weights.ravel()\n",
    "    )\n",
    "    denom, _, _ = np.histogram2d(\n",
    "        x.ravel(), y.ravel(), bins=(xedges, yedges), weights=weights.ravel()\n",
    "    )\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        hist = numer / denom\n",
    "        denom = -kT * np.log(denom)\n",
    "    denom -= np.nanmin(denom)\n",
    "    # plot\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    fig = ax.get_figure()\n",
    "    pcm = ax.pcolormesh(xedges, yedges, hist.T, shading=\"flat\", cmap=cmap,**kwargs)\n",
    "    ax.set_xlim(xedges[0], xedges[-1])\n",
    "    ax.set_ylim(yedges[0], yedges[-1])\n",
    "    #ax.set_aspect(\"equal\")\n",
    "    ax.set_xlabel(xlabel, fontsize=16)\n",
    "    ax.set_ylabel(ylabel, fontsize=16)\n",
    "    ax.set_title(title, fontsize=16)\n",
    "    ax.tick_params(width=1.5, labelsize=14)\n",
    "    cbar = fig.colorbar(pcm, ax=ax)\n",
    "    if label is not None:\n",
    "            cbar.set_label(label, rotation=-90, va=\"baseline\", fontsize=16)\n",
    "            cbar.ax.tick_params(width=1.5, labelsize=14)\n",
    "    #Draw contours of the PMF if desired\n",
    "    if contours:\n",
    "        centerx =(xedges[1:]+xedges[:-1])/2\n",
    "        centery =(yedges[1:]+yedges[:-1])/2\n",
    "        if q:\n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1.0)\n",
    "            ax.contour(centerx, centery, hist.T, levels=[0.5], colors='purple', linestyles='solid', linewidths=2.5)\n",
    "        else:\n",
    "            ax.contour(centerx, centery, hist.T, levels=clines, colors='black', linestyles='solid', linewidths=1.0)\n",
    "    return hist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load features from coordinates\n",
    "\n",
    "Use PyEMMA to import your desired raw features - pairwise distances, dihedral angles, etc. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feat3 = pyemma.coordinates.featurizer('PATH_TO_PDB')\n",
    "\n",
    "# Add atom ids for input features\n",
    "cas3 = (4761, 4764, 163, 182, 239, 564, 583, 616, 1102, 1122, 1169, 1186, 2436, 2470, 2481, 2500, 2510, 2521, 2540, 2578, 2597, 2805, 2822, 2872)\n",
    "gks3 = (4761, 4764, 569, 1176, 2489, 2526)\n",
    "\n",
    "#Add the distances between binding pocket's alpha carbons and IPH1 (protein-protein and protein-ligand)\n",
    "feat3.add_distances(cas3)\n",
    "#Distances between gatekeeping gamma carbons and IPH1 (protein-protein and protein-ligand)\n",
    "feat3.add_distances(gks3)\n",
    "feat3.dimension()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load in that data from all of your trajectories\n",
    "dfiles = sorted(glob.glob('PATH_TO_ALL_DATA_FILES'))\n",
    "raw_feat = pyemma.coordinates.load(dfiles_f, features=feat3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load any supplementary data (CVs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ccvfiles = sorted(glob.glob('PATH_TO_CVs'))\n",
    "\n",
    "supp_ctrajs = []\n",
    "for ccvfile in ccvfiles:\n",
    "    cv = np.loadtxt(ccvfile)\n",
    "    supp_ctrajs.append(cv)\n",
    "supp_CCV_arr = np.vstack(supp_ctrajs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define sets A and B for DGA analysis \n",
    "In this case, we are defining our free and bound states. Definitions are found in the paper. \n",
    "InD means all those in the domain - namely, not in either A or B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvfiles = sorted(glob.glob('PATH_TO_CVFILES'))\n",
    "\n",
    "ahcvfiles = sorted(glob.glob('PATH_TO_HBOND_FILES'))\n",
    "\n",
    "cv_trajs = []\n",
    "InD = []\n",
    "InBound = []\n",
    "InFree = []\n",
    "\n",
    "#Define the parameters for each set\n",
    "lower_cutoff=5\n",
    "upper_cutoff=540\n",
    "n_stdevs = 2\n",
    "abr_c = np.asarray((0.087238, 52.503, 31.924))\n",
    "abr_r = np.asarray((0.022239, 5.023, 5.264)) * n_stdevs\n",
    "\n",
    "print(abr_c, abr_r)\n",
    "\n",
    "#Assemble the cv trajs, as well as the in/out of domain matrices\n",
    "\n",
    "#Define unbound as less than 10 contacts with the protein, and in the sphere centered at (abr_c)\n",
    "for ahcvfile, cvfile, rcvfile in zip(ahcvfiles, cvfiles, rcvfiles):\n",
    "    cv = np.loadtxt(cvfile)[::2,:]\n",
    "    ahcv = np.loadtxt(ahcvfile)[::2]\n",
    "    rcv = np.loadtxt(rcvfile, usecols=(1))[:]\n",
    "    cv_trajs.append(cv)\n",
    "    if len(cv) != 4001: print(cvfile, len(cv))\n",
    "    if len(ahcv) != 4001: print(ahcvfile, len(ahcv))\n",
    "    if len(rcv) != 4001: print(rcvfile, len(rcv))\n",
    "    in_bound = np.zeros(len(cv))\n",
    "    in_free = np.zeros(len(cv))\n",
    "    ind = np.ones(len(cv))\n",
    "    point_r = np.sum(((cv[:,3:6] - abr_c) / abr_r)**2, axis=1)\n",
    "    in_bound[np.logical_and(ahcv > 1.1, np.logical_and(point_r[:] <= 1, cv[:,6] > upper_cutoff))] = 1\n",
    "    in_bound = in_bound.astype('int')\n",
    "    in_free[np.logical_and(cv[:,6] < lower_cutoff/2, np.logical_and(cv[:,1] < lower_cutoff, cv[:,2] < lower_cutoff))] = 1\n",
    "    in_free = in_free.astype('int')\n",
    "    ind = (ind - in_free - in_bound).astype('int')\n",
    "    \n",
    "    InD.append(ind)\n",
    "    InBound.append(in_bound)\n",
    "    InFree.append(in_free)\n",
    "\n",
    "InD_arr = np.reshape(np.hstack(InD), (-1,1))\n",
    "InBound_arr = np.reshape(np.hstack(InBound), (-1,1))\n",
    "InFree_arr = np.reshape(np.hstack(InFree), (-1,1))\n",
    "\n",
    "CV_arr = np.vstack(cv_trajs)\n",
    "\n",
    "Bound_cv = CV_arr[(InBound_arr==1).flatten()]\n",
    "Free_cv = CV_arr[(InFree_arr==1).flatten()]\n",
    "D_cv = CV_arr[(InD_arr==1).flatten()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate the distances to the bound/free states\n",
    "\n",
    "These will be used to calculate the smoothing and guess functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_bound = []\n",
    "counter=0.0\n",
    "tot = len(RawFeat_arr)\n",
    "skip = 3\n",
    "\n",
    "def reduce_func(D_chunk, start):\n",
    "    min2 = np.min(D_chunk,axis=1)\n",
    "    return min2\n",
    "\n",
    "for chunk in skl.pairwise_distances_chunked(RawFeat_arr[:,:], Bound_arr[::skip,:], n_jobs=-1, reduce_func=reduce_func):\n",
    "    d_bound.append(chunk)\n",
    "    counter+=float(len(chunk))\n",
    "    perc = counter / tot * 100\n",
    "    print(str(perc) + '% Completed')\n",
    "\n",
    "d_bounds = np.reshape(np.concatenate(d_bound), (-1,1))\n",
    "d_bounds[(InBound_arr==1).flatten()]=0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_free = []\n",
    "counter=0.0\n",
    "tot = len(RawFeat_arr)\n",
    "def reduce_func(D_chunk, start):\n",
    "    min2 = np.min(D_chunk,axis=1)\n",
    "    return min2\n",
    "\n",
    "for chunk in skl.pairwise_distances_chunked(RawFeat_arr[:,:], Free_arr[:,:], n_jobs=-1, reduce_func=reduce_func):\n",
    "    d_free.append(chunk)\n",
    "    counter+=float(len(chunk))\n",
    "    perc = counter / tot * 100\n",
    "    print(str(perc) + '% Completed')\n",
    "\n",
    "d_frees = np.reshape(np.concatenate(d_free), (-1,1))\n",
    "d_frees[(InFree_arr==1).flatten()]=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate the smoothing and guess functions, build the basis\n",
    "Smooth the raw features earlier generated with PyEMMA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smooth = (d_frees * d_bounds) / (d_frees + d_bounds)**2\n",
    "guess_bound = (d_frees)**2 / (d_frees**2 + d_bounds**2)\n",
    "guess_free = (d_bounds)**2 / (d_frees**2 + d_bounds**2)\n",
    "\n",
    "smooth[(InD_arr==0).flatten()]=0\n",
    "guess_bound[(InBound_arr==1).flatten()]=1\n",
    "guess_bound[(InFree_arr==1).flatten()]=0\n",
    "guess_free[(InFree_arr==1).flatten()]=1\n",
    "guess_free[(InBound_arr==1).flatten()]=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#First, Smooth after adding a constant function\n",
    "const = np.ones((len(RawFeat_arr),1))\n",
    "BasisL_arr = np.hstack((const,RawFeat_arr))\n",
    "basis_smooth = smooth * BasisL_arr\n",
    "\n",
    "# #Now, whiten\n",
    "basis_white,L = svd_whiten(basis_smooth,1,frac_retain=1.0)\n",
    "basis_white[(InD_arr==0).flatten()]=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #Now, whiten the raw data for COM calculation\n",
    "basis_COM,L = svd_whiten(RawFeat_arr,1,frac_retain=1.0)\n",
    "basis_COM[(InD_arr==0).flatten()]=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert to the correct data format (list of arrays instead one long array)\n",
    "\n",
    "BasisF_arr = basis_white\n",
    "BasisC_arr = basis_COM\n",
    "GuessF_arr = guess_bound\n",
    "GuessUF_arr = guess_free\n",
    "#InD already exists as a list of arrays\n",
    "BasisF= []\n",
    "BasisC = []\n",
    "GuessUF= []\n",
    "GuessF= []\n",
    "\n",
    "traj_length = len(raw_feat[0])\n",
    "f_start = 0\n",
    "for i in range(len(raw_feat)):\n",
    "    BasisF.append(BasisF_arr[f_start:f_start+traj_length,:])\n",
    "    BasisC.append(BasisC_arr[f_start:f_start+traj_length,:])\n",
    "    GuessUF.append(GuessUF_arr[f_start:f_start+traj_length,:].flatten())\n",
    "    GuessF.append(GuessF_arr[f_start:f_start+traj_length,:].flatten())\n",
    "    f_start+=traj_length\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Start calculating DGA quantities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Calculate Change of Measure\n",
    "\n",
    "C_COMInds = []\n",
    "for i,LagTime in enumerate(lags[]):\n",
    "    com=extq.dga.reweight(BasisC,LagTime)\n",
    "    print(np.min(com), np.max(com), LagTime)\n",
    "    C_COMInds.append(com)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Calculate forward and backward committors for both the binding and unbinding process\n",
    "\n",
    "C_QIndsUnfold = []\n",
    "C_QIndsFold = []\n",
    "C_QbIndsUnfold = []\n",
    "C_QbIndsFold = []\n",
    "\n",
    "for i,LagTime in enumerate(lags):\n",
    "    C_QIndsUnfold.append(extq.dga.forward_committor(BasisF,C_COMInds[i],C_InD,GuessUF,LagTime))\n",
    "    C_QbIndsUnfold.append(extq.dga.backward_committor(BasisF,C_COMInds[i],C_InD,GuessF,LagTime))\n",
    "    C_QIndsFold.append(extq.dga.forward_committor(BasisF,C_COMInds[i],C_InD,GuessF,LagTime))\n",
    "    C_QbIndsFold.append(extq.dga.backward_committor(BasisF,C_COMInds[i],C_InD,GuessUF,LagTime))\n",
    "    print(LagTime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Explicitly make sure committors obey boundary conditions\n",
    "\n",
    "C_QIndsUnfold = np.asarray(C_QIndsUnfold)\n",
    "C_QbIndsUnfold = np.asarray(C_QbIndsUnfold)\n",
    "C_QIndsFold= np.asarray(C_QIndsFold)\n",
    "C_QbIndsFold = np.asarray(C_QbIndsFold)\n",
    "\n",
    "C_QIndsUnfold[C_QIndsUnfold > 1] = 1\n",
    "C_QIndsUnfold[C_QIndsUnfold < 0] = 0\n",
    "\n",
    "C_QbIndsUnfold[C_QbIndsUnfold > 1] = 1\n",
    "C_QbIndsUnfold[C_QbIndsUnfold < 0] = 0\n",
    "\n",
    "C_QbIndsFold[C_QbIndsFold > 1] = 1\n",
    "C_QbIndsFold[C_QbIndsFold < 0] = 0\n",
    "\n",
    "C_QIndsFold[C_QIndsFold > 1] = 1\n",
    "C_QIndsFold[C_QIndsFold < 0] = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate rates\n",
    "\n",
    "urates = []\n",
    "frates = []\n",
    "for i, Lag in enumerate(lags):\n",
    "    urates.append(extq.tpt.rate(C_QIndsUnfold[i],C_QbIndsUnfold[i],C_COMInds[i],C_InD,C_QIndsUnfold[i],Lag))\n",
    "    frates.append(extq.tpt.rate(C_QIndsFold[i],C_QbIndsFold[i],C_COMInds[i],C_InD,C_QIndsFold[i],Lag))\n",
    "    print(Lag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert to unimolecular rate constants\n",
    "\n",
    "kF_denoms = []\n",
    "kU_denoms = []\n",
    "\n",
    "for i, lag in enumerate(lags):\n",
    "    trjlen = np.asarray(C_QbIndsFold[i]).shape[1]\n",
    "    kF_denom = np.sum(np.asarray(C_QbIndsFold[i])[:,:(trjlen-lag)]* \n",
    "                      np.asarray(C_COMInds[i])[:,:(trjlen-lag)]) / np.sum(np.asarray(C_COMInds[i])[:,:(trjlen-lag)])\n",
    "    kU_denom = np.sum(np.asarray(C_QbIndsUnfold[i])[:,:(trjlen-lag)]* \n",
    "                      np.asarray(C_COMInds[i])[:,:(trjlen-lag)]) / np.sum(np.asarray(C_COMInds[i])[:,:(trjlen-lag)])\n",
    "    kF_denoms.append(kF_denom)\n",
    "    kU_denoms.append(kU_denom)\n",
    "    print(kF_denom, kU_denom, trjlen)\n",
    "kF_denoms = np.asarray(kF_denoms)\n",
    "kU_denoms = np.asarray(kU_denoms)\n",
    "\n",
    "kF = np.asarray(frates) / kF_denoms\n",
    "kU = np.asarray(urates) / kU_denoms\n",
    "K_D = kU / kF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert to bimolecular rate constants\n",
    "\n",
    "#Load the distance between phenol and ZN COM\n",
    "phen_r = np.load('c_extended40_final/phen_r.npy')\n",
    "\n",
    "kF_denoms = []\n",
    "kU_denoms = []\n",
    "bi_frates = []\n",
    "urate_mods = []\n",
    "\n",
    "#Diffusion constant in cm^2/s\n",
    "D=2.6529e-5 + 0.0396e-5\n",
    "\n",
    "#Radii in nm\n",
    "b=3.3\n",
    "c=5.3\n",
    "phen_r_new = np.reshape(phen_r, C_QbIndsFold[0].shape)\n",
    "\n",
    "for i, lag in enumerate(lags):\n",
    "    trjlen = np.asarray(C_QbIndsFold[i]).shape[1]\n",
    "    in_range = np.logical_and(phen_r_new > b - 0.05, phen_r_new < b + 0.05)\n",
    "    q_bind = np.mean(C_QbIndsFold[i][in_range])\n",
    "    f = q_bind / (1.0 - ((b/c)*(1.0-q_bind)))\n",
    "    bi_frate = 4.0 * np.pi * D * b * 1e-7 * f * 6.022e23 / 1000\n",
    "    urate_mod = 1.0 - ((b/c)*f)\n",
    "    print(np.shape(C_QbIndsFold[i][in_range]), np.mean(C_QbIndsFold[i][in_range]), bi_frate)\n",
    "    kF_denom = np.sum(np.asarray(C_QbIndsFold[i])[:,:(trjlen-lag)]* \n",
    "                      np.asarray(C_COMInds[i])[:,:(trjlen-lag)]) / np.sum(np.asarray(C_COMInds[i])[:,:(trjlen-lag)])\n",
    "    kU_denom = np.sum(np.asarray(C_QbIndsUnfold[i])[:,:(trjlen-lag)]* \n",
    "                      np.asarray(C_COMInds[i])[:,:(trjlen-lag)]) / np.sum(np.asarray(C_COMInds[i])[:,:(trjlen-lag)])\n",
    "    kF_denoms.append(kF_denom)\n",
    "    kU_denoms.append(kU_denom)\n",
    "    bi_frates.append(bi_frate)\n",
    "    urate_mods.append(urate_mod)\n",
    "    print(kF_denom, kU_denom, trjlen)\n",
    "kF_denoms = np.asarray(kF_denoms)\n",
    "kU_denoms = np.asarray(kU_denoms)\n",
    "\n",
    "kF = np.asarray(bi_frates)\n",
    "kU = np.asarray(urates) * np.asarray(urate_mods) / (kU_denoms * 10 * 1e-12)\n",
    "K_D = kU / kF\n",
    "print(K_D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Calculate reactive currents along any dimension (here we do 4)\n",
    "\n",
    "JABx = []\n",
    "JABy = []\n",
    "JABz = []\n",
    "qJABz = []\n",
    "JBAx = []\n",
    "JBAy = []\n",
    "JBAz = []\n",
    "qJBAz = []\n",
    "\n",
    "cv1 = [col[:,4] for col in cv_trajs]\n",
    "cv2 = [col[:,5] for col in cv_trajs]\n",
    "cv3 = [col[:,5] for col in supp_cv_trajs]\n",
    "\n",
    "for i, (LagTime,com, qf, qb) in enumerate(zip(lags, C_COMInds, C_QIndsUnfold, C_QbIndsUnfold)):\n",
    "    JABx.append((extq.tpt.current(qf, qb, com, C_InD, cv1, LagTime)))\n",
    "    JABy.append((extq.tpt.current(qf, qb, com, C_InD, cv2, LagTime)))\n",
    "    JABz.append((extq.tpt.current(qf, qb, com, C_InD, cv3, LagTime)))\n",
    "    qJABz.append((extq.tpt.current(qf, qb, com, C_InD, qf, LagTime)))\n",
    "    print(LagTime)\n",
    "for i, (LagTime,com, qf, qb) in enumerate(zip(lags, C_COMInds, C_QIndsFold, C_QbIndsFold)):\n",
    "    JBAx.append((extq.tpt.current(qf, qb, com, C_InD, cv1, LagTime)))\n",
    "    JBAy.append((extq.tpt.current(qf, qb, com, C_InD, cv2, LagTime)))\n",
    "    JBAz.append((extq.tpt.current(qf, qb, com, C_InD, cv3, LagTime)))\n",
    "    qJBAz.append((extq.tpt.current(qf, qb, com, C_InD, qf, LagTime)))\n",
    "    print(LagTime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now let's plot some representative figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kT=1\n",
    "xlim = np.linspace(0,100,51)\n",
    "ylim = np.linspace(0,100,51)\n",
    "xlim_flux = np.linspace(0,100,26)\n",
    "ylim_flux = np.linspace(0,100,26)\n",
    "zlim_flux = np.linspace(0,100,26)\n",
    "centerx_flux =(xlim_flux[1:]+xlim_flux[:-1])/2\n",
    "centery_flux =(ylim_flux[1:]+ylim_flux[:-1])/2\n",
    "centerz_flux =(zlim_flux[1:]+zlim_flux[:-1])/2\n",
    "\n",
    "kT=1\n",
    "f, ax = plt.subplots(1,2,figsize=(9.5,5), sharey=True)\n",
    "i=4\n",
    "plot_pmf(np.asarray(cv1),np.asarray(cv2),np.asarray(C_COMInds[i]),kT,xlim,ylim,ax=ax[0],vmax=12.5,\n",
    "                title=None, label=None,\n",
    "                xlabel='Pathway 1 Contacts', ylabel='Pathway 4 Contacts', contours=True, pcolor=False)\n",
    "plot_pmf(np.asarray(cv1),np.asarray(cv2),np.asarray(C_COMInds[i]),kT,xlim,ylim,ax=ax[1],vmax=12.5,\n",
    "                title=None, label=None,\n",
    "                xlabel='Pathway 1 Contacts', ylabel='Pathway 4 Contacts', contours=True, pcolor=False)\n",
    "kde, kde1, e1, e2 = kdesum2d(cv1,cv2,np.asarray(JBAx[i]),np.asarray(JBAy[i]),xmin=0,xmax=100,ymin=0,ymax=100,nx=22,ny=22,xstd=5,ystd=5,plot=True,axis=ax[0],\n",
    "                xlabel='Pathway 1 Contacts', ylabel='Pathway 4 Contacts', title='Binding', cbar = True)\n",
    "kde, kde1, e1, e2 = kdesum2d(cv1,cv2,np.asarray(JABx[i]),np.asarray(JABy[i]),xmin=0,xmax=100,ymin=0,ymax=100,nx=22,ny=22,xstd=5,ystd=5,plot=True,axis=ax[1],\n",
    "                xlabel='Pathway 1 Contacts', title='Unbinding', cbar=True)\n",
    "ax[0].set_ylabel('Pathway 4 Contacts', fontsize=16)\n",
    "ax[0].set_title('Binding', fontsize=16)\n",
    "ax[1].set_title('Unbinding', fontsize=16)\n",
    "for axis in ax.ravel():\n",
    "    axis.tick_params(width=1.5, labelsize=14)\n",
    "    axis.set_xlabel('Pathway 1 Contacts', fontsize=16)\n",
    "\n",
    "f.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "f, axes = plt.subplots(5,2,figsize=(7,15))\n",
    "\n",
    "for i, (lag, ax, q, com) in enumerate(zip(lags, axes.flatten(), C_QIndsUnfold, C_COMInds)):\n",
    "    if i < len(lags):\n",
    "        plot_avg(np.asarray(cv1),np.asarray(cv2),np.asarray(q),np.asarray(com),xlim,ylim,ax=ax,vmin=0,vmax=1,\n",
    "                title=\"Lag=\" + str(lag*10) + \" ps\",label=r'$q_{\\mathrm{unbind}}$',\n",
    "                xlabel=r'$N_{\\mathrm{PW1}}$', ylabel=r'$N_{\\mathrm{PW1}}$', cmap='coolwarm',\n",
    "                contours=True, clines=np.linspace(0,1.0,11),q=True)\n",
    "f.tight_layout()\n",
    "\n",
    "f.savefig('Paper_figs/q_varyt.png', dpi=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "q = np.asarray(C_QIndsFold[4])\n",
    "q_arr = np.asarray(q).flatten()\n",
    "fig = plt.figure(figsize=(6,6))\n",
    "ax = plt.axes(projection='3d')\n",
    "h= ax.scatter(CV_arr[::40,4], CV_arr[::40,5], q_arr[::40], c=q_arr[::40], alpha=0.15, cmap = 'coolwarm')\n",
    "ax.tick_params(width=1.5, labelsize=12)\n",
    "ax.set_xlabel(r'$N_{\\mathrm{PW1}}$', fontsize=16)\n",
    "ax.set_ylabel(r'$N_{\\mathrm{PW4}}$', fontsize=16)\n",
    "ax.set_zlabel(r'$q_{\\mathrm{bind}}$', fontsize=16)\n",
    "ax.view_init(elev=27., azim=-141)\n",
    "\n",
    "\n",
    "fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Bin and smooth the reactive current in 3 dimensions\n",
    "\n",
    "i=4\n",
    "nx=51\n",
    "ny=51\n",
    "nz=51\n",
    "xlim_flux = np.linspace(0,100,nx)\n",
    "ylim_flux = np.linspace(0,100,ny)\n",
    "zlim_flux = np.linspace(0,100,nz)\n",
    "centerx_flux =(xlim_flux[1:]+xlim_flux[:-1])/2\n",
    "centery_flux =(ylim_flux[1:]+ylim_flux[:-1])/2\n",
    "centerz_flux =(zlim_flux[1:]+zlim_flux[:-1])/2\n",
    "cv1 = np.asarray([col[:,4] for col in cv_trajs])\n",
    "cv4 = np.asarray([col[:,5] for col in cv_trajs])\n",
    "q = np.asarray(C_QIndsUnfold[i])*100.0\n",
    "q_arr = np.asarray(q).flatten()\n",
    "\n",
    "kde1, kde2, kde3, e1, e2, e3 = kdesum3d(cv1,cv4,q,np.asarray(JABx[i]),np.asarray(JABy[i]),np.asarray(qJABz[i])*100.0,\n",
    "                                        nx=nx-1,ny=ny-1,nz=nz-1,xstd=8,ystd=8,zstd=8,plot=False)\n",
    "\n",
    "j = kde3.swapaxes(2,0)[20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plot slices of this as a function of committor\n",
    "\n",
    "f, axes = plt.subplots(4,4,figsize=(9.5,9.5), sharex=True, sharey=True)\n",
    "for i, (x, ax) in enumerate(zip(kde3.swapaxes(2,0)[8:],axes.ravel())):\n",
    "    pcm = ax.pcolormesh(xlim_flux, ylim_flux, x, shading=\"flat\", vmax=(np.mean(x) + 3*np.std(x)))\n",
    "    plt.colorbar(pcm, ax=ax, format=\"%.0e\")\n",
    "    ax.set_title(\"q=\" + \"{:.2f}\".format((centerz_flux[i + 8])/100))\n",
    "    j,k = divmod(i,4)\n",
    "    if k==0: \n",
    "        ax.set_ylabel(r'$N_{\\mathrm{PW4}}$')\n",
    "    if j==3:\n",
    "        ax.set_xlabel(r'$N_{\\mathrm{PW1}}$')\n",
    "    print( (np.sum(x)))\n",
    "    ax.axvline(x=80, c='C1')\n",
    "    ax.axvline(x=44, c='C1')\n",
    "    ax.axhline(y=70, c='C3')\n",
    "    ax.axhline(y=32, c='C3')\n",
    "f.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Assign different regions to different pathways\n",
    "\n",
    "assignments = np.zeros(np.shape(kde3.swapaxes(2,0)[0]))\n",
    "X, Y = np.meshgrid(centerx_flux, centery_flux)\n",
    "print(X.shape)\n",
    "print(assignments.shape)\n",
    "x_1 = 64\n",
    "y_1 = 70\n",
    "\n",
    "x_2 = 40\n",
    "y_2 = 32\n",
    "\n",
    "x_3 = 20\n",
    "y_3 = 16\n",
    "pw_4a = (Y > y_1)\n",
    "pw_1a = np.logical_and(pw_4a==False, X > x_1)\n",
    "pw_4b = np.logical_and(pw_1a==False, np.logical_and(Y > y_2, pw_4a==False))\n",
    "pw_1b = np.logical_and(pw_4b==False, np.logical_and(pw_1a==False, np.logical_and(X > x_2, pw_4a==False)))\n",
    "pw_3 = np.logical_and(X < x_3, Y < y_3)\n",
    "\n",
    "pw_2 = np.logical_and(pw_3==False, np.logical_and(pw_1b==False, np.logical_and(pw_1a==False, \n",
    "                                                                               np.logical_and(pw_4a==False, pw_4b==False))))\n",
    "assignments[pw_4a] = 1\n",
    "assignments[pw_1a] = 2\n",
    "assignments[pw_4b] = 3\n",
    "assignments[pw_1b] = 4\n",
    "assignments[pw_3] = 5\n",
    "assignments[pw_2] = 6\n",
    "f, ax =plt.subplots(figsize=(6,5))\n",
    "pcm = ax.pcolormesh(xlim_flux, ylim_flux, assignments,shading=\"flat\")\n",
    "plt.colorbar(pcm, ax=ax, format=\"%.0e\")\n",
    "f.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Actually calculate the weights by calculating the appropriate dot products\n",
    "pw1a_ws = []\n",
    "pw4a_ws = []\n",
    "pw1b_ws = []\n",
    "pw4b_ws = []\n",
    "pw2_ws = []\n",
    "pw3_ws = []\n",
    "\n",
    "for j in kde3.swapaxes(2,0)[8:28]:\n",
    "    total_flux = np.sum(j)\n",
    "    pw1a_w = np.sum(pw_1a*j) / total_flux\n",
    "    pw4a_w = np.sum(pw_4a*j) / total_flux\n",
    "    pw1b_w = np.sum(pw_1b*j) / total_flux\n",
    "    pw4b_w = np.sum(pw_4b*j) / total_flux\n",
    "    pw2_w = np.sum(pw_2*j) / total_flux\n",
    "    pw3_w = np.sum(pw_3*j) / total_flux\n",
    "    \n",
    "    pw1a_ws.append(pw1a_w)\n",
    "    pw4a_ws.append(pw4a_w)\n",
    "    pw1b_ws.append(pw1b_w)\n",
    "    pw4b_ws.append(pw4b_w)\n",
    "    pw2_ws.append(pw2_w)\n",
    "    pw3_ws.append(pw3_w)\n",
    "    \n",
    "pw1a_ws = np.array(pw1a_ws)\n",
    "pw4a_ws = np.array(pw4a_ws)\n",
    "pw1b_ws = np.array(pw1b_ws)\n",
    "pw4b_ws = np.array(pw4b_ws)\n",
    "pw2_ws = np.array(pw2_ws)\n",
    "pw3_ws = np.array(pw3_ws)\n",
    "print(pw1a_ws + pw1b_ws + pw2_ws + pw3_ws + pw4a_ws + pw4b_ws)\n",
    "print(pw1a_ws[3], pw1b_ws[3], pw2_ws[3], pw3_ws[3], pw4a_ws[3], pw4b_ws[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plot relative weights as a function of the dividing surface q (for one value of lag time). \n",
    "#You could also do this for one dividing surface as a function of lag time. \n",
    "\n",
    "f, ax = plt.subplots(figsize=(6,6))\n",
    "ax.plot(centerz_flux[8:28] / 100, pw1a_ws, '--o', label=\"PW1\")\n",
    "ax.plot(centerz_flux[8:28] /100, pw1b_ws, '--o', label=\"PW1a\")\n",
    "ax.plot(centerz_flux[8:28]/100, pw2_ws, '--o', label=\"PW2\")\n",
    "ax.plot(centerz_flux[8:28]/100, pw3_ws, '--o', label=\"PW3\")\n",
    "ax.plot(centerz_flux[8:28]/100, pw4a_ws, '--o', label=\"PW4\")\n",
    "ax.plot(centerz_flux[8:28]/100, pw4b_ws, '--o', label=\"PW4a\")\n",
    "ax.legend(loc='upper right', fontsize=18)\n",
    "ax.set_xlabel('Q Slice', fontsize=23)\n",
    "ax.set_ylabel('Relative weight', fontsize=23)\n",
    "ax.set_title('WT', fontsize=23)\n",
    "ax.set_ylim((0,1))\n",
    "f.tight_layout()"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}