{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fadc2af2-f6a4-4da6-9a61-373756cdf1fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a074862a-1efc-4539-a274-06d866fbc030",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "palettable colormaps are not installed\n",
      "don't have phil's colormaps\n"
     ]
    }
   ],
   "source": [
    "from firefly.data_reader.GAIAreader import GaiaReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3e4c6849-652e-49a8-bbff-0571b64ff274",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "2cfcad47-cf1c-445f-8863-b7679e4c9c0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2.090136210123698e-06 min elapsed\n",
      "Make sure each field_array (3) has a field_filter_flag (0), assuming True.\n",
      "Make sure each field_array (3) has a field_colormap_flag (0), assuming True.\n",
      "\n",
      "1.2596448262532552e-06 min elapsed\n",
      "Make sure each field_array (3) has a field_filter_flag (0), assuming True.\n",
      "Make sure each field_array (3) has a field_colormap_flag (0), assuming True.\n",
      "datadir: /home/abg6257/projects/gaia -- is not a sub-directory of firefly/static/data. \n",
      "This may produce confusing or inoperable results. As such, we will create a symlink for you when you writeToDisk.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "gaiadir = os.path.join(os.environ['HOME'],'projects','gaia','GaiaSource')\n",
    "my_reader = GaiaReader(\n",
    "    nthreads=52,\n",
    "    write_startup=True,\n",
    "    clean_datadir=True,\n",
    "    use_mps=True,\n",
    "    min_to_refine=1e5/10,\n",
    "    nrecurse=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "cd6f8056-7599-4ba8-9661-96a18bb159c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Progress: |--------------------------------------------------| 1.5% complete\n",
      "173.17736530303955 s elapsed\n"
     ]
    }
   ],
   "source": [
    "data_dicts = my_reader.retrieve_gaia_data(return_arrays=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "78150b3f-9600-415c-8d27-cc07205dccb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.concatenate([data_dict['x'] for data_dict in data_dicts])\n",
    "ys = np.concatenate([data_dict['y'] for data_dict in data_dicts])\n",
    "zs = np.concatenate([data_dict['z'] for data_dict in data_dicts])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "e0701bdc-ff82-4f0c-a0b2-0d4dcb67b812",
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = np.sqrt(xs**2+ys**2+zs**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "d0dffa3b-e2c1-489e-934d-07ce7686099d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_xs = xs[rs < 300]\n",
    "sub_ys = ys[rs < 300]\n",
    "sub_zs = zs[rs < 300]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "c10633ae-8ea0-42f6-b2b4-cd6c8581afe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "79d0d15c-20f0-4f03-9acb-03fb4980761a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 211648 characters]\n",
      "text/plain": [
       "<Figure size 720x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from matplotlib.colors import LogNorm\n",
    "bins = np.linspace(-75,75,10**3)\n",
    "plt.hist2d(sub_zs,sub_xs,bins)\n",
    "plt.gca().set_aspect(1)\n",
    "plt.gcf().set_dpi(120)\n",
    "plt.gcf().set_facecolor('white')\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "cd773966-46a3-4262-aaad-629f250968fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-3770.70975328,  -552.31160977, -1523.9746656 , ...,\n",
       "         280.40817132,  2123.44044177,  2031.51468223])"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_dicts[0]['x']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "707f60a4-f77a-44c3-b0c0-22085e214b5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('/home/abg6257/projects/gaia/GaiaDataSettings.json',\n",
       " '/home/abg6257/projects/gaia/GaiaDataSettings.json')"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#my_reader.settings['GUIExcludeList'] += ['colorbarcontainer']\n",
    "my_reader.settings['collapseGUIAtStart'] = False\n",
    "for pg in my_reader.particleGroups:\n",
    "    my_reader.settings['sizeMult'][pg.UIname] = 0.1\n",
    "    my_reader.settings['depthTest'][pg.UIname] = True\n",
    "    my_reader.settings['blendingMode'][pg.UIname] = 'normal'\n",
    "    \n",
    "#my_reader.settings['showParts']['DR3-RV'] = False\n",
    "my_reader.settings['camera'] = [0,0,0]\n",
    "my_reader.settings['cameraRotation'] = [-1.5,0,-1.5]\n",
    "my_reader.settings['startFly'] = True\n",
    "my_reader.dumpSettingsToJSON()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ef4db67-60da-450f-a361-1971e3360ad2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "cda5a6f7-ab0a-4410-8e75-96144211fcee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DR3-noRV - 966,113,753 parts (31,451 nodes) - 3 fields\n",
      "DR3-RV - 31,431,382 parts (8,506 nodes) - 3 fields\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "''"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_reader.writeToDisk(nthreads=52)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "be652e71-1029-4605-a242-2b3ad8e64d47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/projects/b1094/ageller/Firefly_ABG/']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_reader.copyFireflySourceToTarget('/projects/b1094/ageller/Firefly_ABG/',dump_data=False,overwrite=False,nthreads=52)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "75e531d8-39ce-4e19-a832-fa261b220f2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "88dd98a8-fc29-49b9-a71f-dd4b3cbd3a69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<code object writeToDisk at 0x2af7bf0710e0, file \"/projects/b1026/agurvich/Firefly/src/firefly/data_reader/reader.py\", line 241>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_reader.writeToDisk.__code__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "85dc6ce7-1d26-462b-83e8-25f5b22ec12c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "983b56d1-7d6f-4ef9-859d-a6caf4e0ab29",
   "metadata": {},
   "outputs": [],
   "source": [
    "import firefly.data_reader.octree as oct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "d6417124-ce9b-442a-989f-5620a283c767",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['BinaryWriter',\n",
       " 'OctNode',\n",
       " 'Octree',\n",
       " 'RawBinaryWriter',\n",
       " '__builtins__',\n",
       " '__cached__',\n",
       " '__file__',\n",
       " '__loader__',\n",
       " '__name__',\n",
       " '__package__',\n",
       " '__spec__',\n",
       " 'convertNodeFFRawFFTree',\n",
       " 'copy',\n",
       " 'group_files',\n",
       " 'init_octree_root_node',\n",
       " 'itertools',\n",
       " 'load_from_json',\n",
       " 'multiprocessing',\n",
       " 'np',\n",
       " 'octant_offsets',\n",
       " 'os',\n",
       " 'printProgressBar',\n",
       " 'refineNode',\n",
       " 'time',\n",
       " 'validate_files',\n",
       " 'write_to_json']"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(oct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "97560a1d-2a73-40f8-85e9-726e14cecc3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def convertNodeFFRawFFTree(\n",
      "    node_dict,\n",
      "    fname,\n",
      "    field_names):\n",
      "\n",
      "    if 'files' not in node_dict.keys(): return node_dict\n",
      "\n",
      "    ## create a new octnode object to translate the data with\n",
      "    node = OctNode(\n",
      "        node_dict['center'],\n",
      "        node_dict['width'],\n",
      "        field_names,\n",
      "        node_dict['name'],\n",
      "        has_velocities=node_dict['com_velocity'] is not None,\n",
      "        has_colors=node_dict['rgba_color'] is not None) \n",
      "    ## load the node's particles from the .ffraw files\n",
      "    node.set_buffers_from_disk(node_dict['files'],node_dict['buffer_size'])\n",
      "    return node.write_fftree(fname)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(oct.convertNodeFFRawFFTree))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "7b4524e8-00ef-4af0-a770-687c87586d15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__call__',\n",
       " '__class__',\n",
       " '__delattr__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__eq__',\n",
       " '__format__',\n",
       " '__func__',\n",
       " '__ge__',\n",
       " '__get__',\n",
       " '__getattribute__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__le__',\n",
       " '__lt__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__repr__',\n",
       " '__self__',\n",
       " '__setattr__',\n",
       " '__sizeof__',\n",
       " '__str__',\n",
       " '__subclasshook__']"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(my_reader.writeToDisk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1e8b81a-355a-4667-afbe-3ac37d5b215a",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
