{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "from glob import glob\n",
    "from notebooks.lib import read_xyz, read_load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "outputs": [],
   "source": [
    "def linear(x, a, b):\n",
    "    return x * a + b\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "c_direct_neighbours = read_load(\n",
    "    \"/Users/joshuacoles/Developer/checkouts/jc3091/CompB DLA/data-analysis/data/c-3d-direct-neighbours\", read_xyz)\n",
    "c_off_axis_neighbours = read_load(\n",
    "    \"/Users/joshuacoles/Developer/checkouts/jc3091/CompB DLA/data-analysis/data/c-3d-off-axis-neighbours\", read_xyz)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 640x480 with 1 Axes>",
      "image/png": "[encoded data, 53156 characters]"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<Figure size 640x480 with 1 Axes>",
      "image/png": "[encoded data, 143928 characters]"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for _, c in c_direct_neighbours.groupby('run'):\n",
    "    plt.plot(c.fd)\n",
    "\n",
    "for _, c in c_off_axis_neighbours.groupby('run'):\n",
    "    plt.plot(c.fd)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "for _, run in c_direct_neighbours.groupby('run'):\n",
    "    plt.plot(run.N, run.cr)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for _, run in c_direct_neighbours.groupby('run'):\n",
    "    plt.plot(run.N, run.cr)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/joshuacoles/Library/Caches/pypoetry/virtualenvs/data-analysis-B4Au_hWl-py3.10/lib/python3.10/site-packages/pandas/core/arraylike.py:402: RuntimeWarning: divide by zero encountered in log\n",
      "  result = getattr(ufunc, method)(*inputs, **kwargs)\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 640x480 with 1 Axes>",
      "image/png": "[encoded data, 44572 characters]"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "(array([1.68571639, 1.41007603]),\n array([[ 6.66524704e-07, -2.46788968e-06],\n        [-2.46788968e-06,  9.25357395e-06]]))"
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meaned_by_N = c_direct_neighbours.groupby('N').agg({'fd': ['mean', 'std']}) \\\n",
    "    .reset_index() \\\n",
    "    .replace([np.inf, -np.inf], np.nan)\n",
    "\n",
    "without_prefix = c_direct_neighbours[c_direct_neighbours.N > 200]\n",
    "\n",
    "p, pcov = scipy.optimize.curve_fit(linear, np.log(without_prefix.cr), np.log(without_prefix.N))\n",
    "linear_extent = np.linspace(0, np.max(np.log(c_direct_neighbours.cr)))\n",
    "\n",
    "plt.scatter(np.log(c_direct_neighbours.cr), np.log(c_direct_neighbours.N), s=1, marker='.', color=\"tab:blue\")\n",
    "plt.plot(linear_extent, linear(linear_extent, *p), color=\"tab:red\")\n",
    "\n",
    "plt.xlabel(\"$\\\\log r_{max}$\")\n",
    "plt.ylabel(\"$\\\\log N$\")\n",
    "plt.show()\n",
    "p, pcov"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.legend.Legend at 0x2c687ab30>"
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 1000x1000 with 1 Axes>",
      "image/png": "[encoded data, 62244 characters]"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "direct_meaned_by_N = c_direct_neighbours.groupby('N').agg({'fd': ['mean', 'std']}) \\\n",
    "    .reset_index() \\\n",
    "    .replace([np.inf, -np.inf], np.nan)\n",
    "\n",
    "offaxis_meaned_by_N = c_off_axis_neighbours.groupby('N').agg({'fd': ['mean', 'std']}) \\\n",
    "    .reset_index() \\\n",
    "    .replace([np.inf, -np.inf], np.nan)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "\n",
    "def nc_fd(df, label, color=None):\n",
    "    plt.fill_between(\n",
    "        df.N,\n",
    "        # TODO Check error math here\n",
    "        (df['fd']['mean'] - df['fd']['std'] / np.sqrt(20)),\n",
    "        (df['fd']['mean'] + df['fd']['std'] / np.sqrt(20)),\n",
    "        color=color,\n",
    "        alpha=0.2,\n",
    "        label=f\"{label}, standard error band\"\n",
    "    )\n",
    "\n",
    "    plt.plot(\n",
    "        df.N,\n",
    "        df['fd']['mean'],\n",
    "        color=color,\n",
    "        label=f\"{label}, fd mean\"\n",
    "    )\n",
    "\n",
    "\n",
    "nc_fd(direct_meaned_by_N[:], label=\"Direct\", color=\"tab:blue\")\n",
    "nc_fd(offaxis_meaned_by_N[:], label=\"Off-axis\", color=\"tab:green\")\n",
    "\n",
    "plt.plot([np.min(without_prefix.N), np.max(without_prefix.N)], [2.5, 2.5], color='tab:red', label='Theory')\n",
    "plt.fill_between(without_prefix.N, 2.5 - 0.01, 2.5 + 0.01, alpha=0.2, color='tab:red', label='Theory error band')\n",
    "plt.legend()"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
