1_split_query_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pandas    0.23.4\n",
      "numpy     1.15.4\n",
      "watermark 1.8.1\n",
      "ELEANOR LUTZ 2019-06-03 \n",
      "\n",
      "CPython 3.7.1\n",
      "IPython 7.2.0\n",
      "\n",
      "compiler   : MSC v.1900 64 bit (AMD64)\n",
      "system     : Windows\n",
      "release    : 10\n",
      "machine    : AMD64\n",
      "processor  : Intel64 Family 6 Model 63 Stepping 2, GenuineIntel\n",
      "CPU cores  : 12\n",
      "interpreter: 64bit\n"
     ]
    }
   ],
   "source": [
    "# Watermark is not required for this code, but is included for information. \n",
    "import watermark\n",
    "%load_ext watermark\n",
    "%watermark -a \"ELEANOR LUTZ\" -d -v -iv -m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from astropy.time import Time, TimeDelta\n",
    "import os.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_angle(x0, y0, x1, y1):\n",
    "    ''' Calculate the angle from horizontal, counterclockwise '''\n",
    "    angle = np.rad2deg(np.arctan2(y1-y0, x1-x0))\n",
    "    return angle\n",
    "\n",
    "def split_asteroids(readname, savename, min_distance, max_distance, \n",
    "                    min_diameter, max_diameter, start_date, stop_date, \n",
    "                    randomize=False):\n",
    "    print(savename.split('/')[-1])\n",
    "    df = pd.read_csv(readname, low_memory=False)\n",
    "    print('Original dataset contains', len(df), 'items')\n",
    "\n",
    "    # SPLIT BY DISTANCE FROM SUN, AU\n",
    "    df['q'] = pd.to_numeric(df['q'])\n",
    "    df = df[df['q'] < max_distance]\n",
    "    df = df[df['q'] >= min_distance]\n",
    "    print('Dataset now contains', len(df), 'items', min_distance, '~', \n",
    "          max_distance, 'AU from the sun')\n",
    "\n",
    "    # SPLIT BY DIAMETER, KM\n",
    "    if min_diameter != 'null':\n",
    "        df = df[df['diameter'].astype(float) >= min_diameter]\n",
    "        df = df[df['diameter'].astype(float) < max_diameter]\n",
    "    else:\n",
    "        df = df[pd.isnull(df['diameter'])]\n",
    "        exclude = ['GRK', 'TJN', 'MBA']\n",
    "        df = df[~df['class'].isin(exclude)]\n",
    "        \n",
    "    print('Dataset now contains', len(df), 'items', min_diameter, '~', \n",
    "          max_diameter, 'km in diameter') \n",
    "    \n",
    "    assert df['spkid'].isna().sum() == 0\n",
    "    df['horizons'] = \"DES=+\"+df['spkid'].astype(str)\n",
    "    \n",
    "    # REMOVE DUPLICATES\n",
    "    count = len(df)\n",
    "    df = df.drop_duplicates(keep='first', subset='spkid')\n",
    "    print('Dropped', count-len(df), 'duplicated rows by spkid')\n",
    "    print(len(df[df.duplicated('horizons') == True]), \n",
    "          'duplicated rows remaining by horizons')\n",
    "    \n",
    "    # ADD DATETIME LIMITS\n",
    "    print('Dropped', len(df[df['per'].isna()]), \"NaN values in period data\")\n",
    "    df = df[np.isfinite(df['per'])]\n",
    "    \n",
    "    if randomize != False: \n",
    "        df_named = df[~pd.isnull(df['name'])].copy()\n",
    "        df_notnamed = df[pd.isnull(df['name'])].copy()\n",
    "        if len(df_named) < randomize:\n",
    "            df_sample = df_notnamed.sample(n=randomize-len(df_named))\n",
    "            df = pd.concat([df_sample, df_named])\n",
    "        else:\n",
    "            df = df_named.sample(n=randomize)\n",
    "        assert len(df) == randomize\n",
    "        print(len(df_named), 'named asteroids included in randomized set')\n",
    "    \n",
    "    df1 = df[df['per'] < 40*365].copy()\n",
    "    df1['timedelta'] = TimeDelta(df1['per']*0.25*24*60*60, format='sec')\n",
    "    df1['begin_time'] = Time(Time(start_date, format=\"iso\") - df1['timedelta']).value\n",
    "    df1.drop('timedelta', axis=1, inplace=True)\n",
    "    print(len(df1), 'values truncated because orbital period is shorter than 40 years')\n",
    "    \n",
    "    df2 = df[df['per'] >= 40*365].copy()\n",
    "    df2['begin_time'] = Time(stop_date, format=\"iso\").value\n",
    "    df = pd.concat([df1, df2])\n",
    "    df['end_time'] = Time(start_date, format=\"iso\").value\n",
    "    \n",
    "    df['end_time'] = pd.to_datetime(df['end_time'], dayfirst=False).dt.strftime('%Y-%m-%d-%H-%M-%S')\n",
    "    df['begin_time'] = pd.to_datetime(df['begin_time'], dayfirst=False).dt.strftime('%Y-%m-%d-%H-%M-%S')\n",
    "    \n",
    "    print('Dataset has', len(df), 'items total,', \n",
    "          len(df[df['name'] != np.nan]), 'with proper names.')\n",
    "    \n",
    "    if not os.path.isfile(savename):\n",
    "        df.to_csv(savename, index=False)\n",
    "    else:\n",
    "        print('---NOT SAVED BECAUSE FILE ALREADY EXISTS---\\n')\n",
    "    \n",
    "def split_planets(readname, savename, start_date, stop_date):\n",
    "    \n",
    "    print(savename.split('/')[-1])\n",
    "    df = pd.read_csv(readname, low_memory=False)\n",
    "    print('Original dataset contains', len(df), 'items')\n",
    "\n",
    "    # ADD DATETIME LIMITS\n",
    "    print('Dropped', len(df[df['per'].isna()]), \"NaN values in period data\")\n",
    "    df = df[np.isfinite(df['per'])]\n",
    "    \n",
    "    df['timedelta'] = TimeDelta(df['per']*1*24*60*60, format='sec')\n",
    "    df['begin_time'] = Time(Time(start_date, format=\"iso\") - df['timedelta']).value\n",
    "    df['end_time'] = Time(start_date, format=\"iso\").value\n",
    "    df.drop('timedelta', axis=1, inplace=True)\n",
    "    \n",
    "    df['end_time'] = pd.to_datetime(df['end_time'], dayfirst=False).dt.strftime('%Y-%m-%d-%H-%M-%S')\n",
    "    df['begin_time'] = pd.to_datetime(df['begin_time'], dayfirst=False).dt.strftime('%Y-%m-%d-%H-%M-%S')\n",
    "            \n",
    "    # DUPLICATES\n",
    "    print(len(df[df.duplicated('horizons') == True]), 'duplicated rows remaining by horizons')\n",
    "    print('Dataset has', len(df), 'items total,', \n",
    "          len(df[df['name'] != np.nan]), 'with proper names.')\n",
    "    \n",
    "    if not os.path.isfile(savename):\n",
    "        df.to_csv(savename, index=False)\n",
    "    else:\n",
    "        print('---NOT SAVED BECAUSE FILE ALREADY EXISTS---\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "planets.csv\n",
      "Original dataset contains 8 items\n",
      "Dropped 0 NaN values in period data\n",
      "0 duplicated rows remaining by horizons\n",
      "Dataset has 8 items total, 8 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n",
      "large_asteroids.csv\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: ErfaWarning: ERFA function \"taiutc\" yielded 1 of \"dubious year (Note 4)\" [astropy._erfa.core]\n",
      "WARNING: ErfaWarning: ERFA function \"d2dtf\" yielded 2 of \"dubious year (Note 5)\" [astropy._erfa.core]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original dataset contains 794562 items\n",
      "Dataset now contains 794562 items 0 ~ 240 AU from the sun\n",
      "Dataset now contains 2714 items 20 ~ inf km in diameter\n",
      "Dropped 0 duplicated rows by spkid\n",
      "0 duplicated rows remaining by horizons\n",
      "Dropped 0 NaN values in period data\n",
      "2648 values truncated because orbital period is shorter than 40 years\n",
      "Dataset has 2714 items total, 2714 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n",
      "small_asteroids.csv\n",
      "Original dataset contains 794562 items\n",
      "Dataset now contains 794562 items 0 ~ 240 AU from the sun\n",
      "Dataset now contains 7468 items 10 ~ 20 km in diameter\n",
      "Dropped 0 duplicated rows by spkid\n",
      "0 duplicated rows remaining by horizons\n",
      "Dropped 0 NaN values in period data\n",
      "7464 values truncated because orbital period is shorter than 40 years\n",
      "Dataset has 7468 items total, 7468 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n",
      "large_comets.csv\n",
      "Original dataset contains 3568 items\n",
      "Dataset now contains 3568 items 0 ~ 240 AU from the sun\n",
      "Dataset now contains 12 items 10 ~ inf km in diameter\n",
      "Dropped 0 duplicated rows by spkid\n",
      "0 duplicated rows remaining by horizons\n",
      "Dropped 0 NaN values in period data\n",
      "5 values truncated because orbital period is shorter than 40 years\n",
      "Dataset has 12 items total, 12 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n",
      "any_outer_asteroids.csv\n",
      "Original dataset contains 794562 items\n",
      "Dataset now contains 27258 items 3 ~ 240 AU from the sun\n",
      "Dataset now contains 9199 items null ~ null km in diameter\n",
      "Dropped 0 duplicated rows by spkid\n",
      "0 duplicated rows remaining by horizons\n",
      "Dropped 0 NaN values in period data\n",
      "78 named asteroids included in randomized set\n",
      "3173 values truncated because orbital period is shorter than 40 years\n",
      "Dataset has 5000 items total, 5000 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n",
      "any_inner_asteroids.csv\n",
      "Original dataset contains 794562 items\n",
      "Dataset now contains 554983 items 0 ~ 2.5 AU from the sun\n",
      "Dataset now contains 53819 items null ~ null km in diameter\n",
      "Dropped 0 duplicated rows by spkid\n",
      "0 duplicated rows remaining by horizons\n",
      "Dropped 1 NaN values in period data\n",
      "370 named asteroids included in randomized set\n",
      "2995 values truncated because orbital period is shorter than 40 years\n",
      "Dataset has 3000 items total, 3000 with proper names.\n",
      "---NOT SAVED BECAUSE FILE ALREADY EXISTS---\n",
      "\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "Select asteroids by size and distance from the sun \n",
    "Note: objects are selected by perihelion distance,\n",
    "so they may not be in visible range after getting the exact \n",
    "orbital locations from HORIZONS.\n",
    "'''\n",
    "start_date = '2000-01-01 00:00:00'\n",
    "stop_date = '1990-01-01 00:00:00'\n",
    "\n",
    "# PLANETS\n",
    "readname = './data/planets.csv'\n",
    "savename = './data/planets.csv'\n",
    "split_planets(readname, savename, start_date, stop_date)\n",
    "\n",
    "# ASTEROIDS \n",
    "readname = './data/all_asteroids_wrangled.csv'\n",
    "readname_comets = './data/all_comets_wrangled.csv'\n",
    "\n",
    "# ASTEROIDS >20KM in DIAMETER\n",
    "savename = './data/large_asteroids.csv'\n",
    "min_diameter, max_diameter = 20, np.inf\n",
    "min_distance, max_distance = 0, 240\n",
    "split_asteroids(readname, savename, min_distance, max_distance, \n",
    "                min_diameter, max_diameter, start_date, stop_date)\n",
    "\n",
    "# ASTEROIDS 10~20KM in DIAMETER\n",
    "savename = './data/small_asteroids.csv'\n",
    "min_diameter, max_diameter = 10, 20\n",
    "min_distance, max_distance = 0, 240\n",
    "split_asteroids(readname, savename, min_distance, max_distance, \n",
    "                min_diameter, max_diameter, start_date, stop_date)\n",
    "\n",
    "# COMETS >10KM in DIAMETER\n",
    "savename = './data/large_comets.csv'\n",
    "min_diameter, max_diameter = 10, np.inf\n",
    "min_distance, max_distance = 0, 240\n",
    "split_asteroids(readname_comets, savename, min_distance, max_distance, \n",
    "                min_diameter, max_diameter, start_date, stop_date)\n",
    "\n",
    "# ASTEROIDS\n",
    "savename = './data/any_outer_asteroids.csv'\n",
    "min_diameter, max_diameter = 'null', 'null'\n",
    "min_distance, max_distance = 3, 240\n",
    "split_asteroids(readname, savename, min_distance, max_distance, \n",
    "               min_diameter, max_diameter, start_date, stop_date, randomize=5000)\n",
    "\n",
    "# ASTEROIDS\n",
    "savename = './data/any_inner_asteroids.csv'\n",
    "min_diameter, max_diameter = 'null', 'null'\n",
    "min_distance, max_distance = 0, 2.5\n",
    "split_asteroids(readname, savename, min_distance, max_distance, \n",
    "                 min_diameter, max_diameter, start_date, stop_date, randomize=3000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MBA    1734\n",
       "OMB     346\n",
       "GRK     296\n",
       "TJN     259\n",
       "TNO      41\n",
       "CEN      32\n",
       "MCA       4\n",
       "AMO       1\n",
       "AST       1\n",
       "Name: class, dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "MBA    5191\n",
       "OMB    1035\n",
       "GRK     722\n",
       "TJN     495\n",
       "CEN       9\n",
       "MCA       5\n",
       "AST       4\n",
       "AMO       3\n",
       "IMB       2\n",
       "TNO       2\n",
       "Name: class, dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "OMB    3034\n",
       "TNO    1736\n",
       "CEN     204\n",
       "AST      26\n",
       "Name: class, dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "MCA    992\n",
       "IMB    920\n",
       "APO    543\n",
       "AMO    378\n",
       "OMB     89\n",
       "ATE     73\n",
       "CEN      3\n",
       "IEO      2\n",
       "Name: class, dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "JFc    6\n",
       "HTC    4\n",
       "COM    1\n",
       "CTc    1\n",
       "Name: class, dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Output the different classes of asteroids and comets for reference\n",
    "\n",
    "ast = pd.read_csv('./data/large_asteroids.csv', low_memory=False)\n",
    "display(ast['class'].value_counts())\n",
    "\n",
    "ast = pd.read_csv('./data/small_asteroids.csv', low_memory=False)\n",
    "display(ast['class'].value_counts())\n",
    "\n",
    "ast = pd.read_csv('./data/any_outer_asteroids.csv', low_memory=False)\n",
    "display(ast['class'].value_counts())\n",
    "\n",
    "ast = pd.read_csv('./data/any_inner_asteroids.csv', low_memory=False)\n",
    "display(ast['class'].value_counts())\n",
    "\n",
    "com = pd.read_csv('./data/large_comets.csv', low_memory=False)\n",
    "display(com['class'].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "296 GRKs\n"
     ]
    }
   ],
   "source": [
    "df_asts = pd.read_csv('./data/large_asteroids.csv')\n",
    "df_tjn = df_asts[df_asts['class'] == 'TJN'].copy()\n",
    "df_non_tjn = df_asts[df_asts['class'] != 'TJN'].copy()\n",
    "count, t = 0, 0\n",
    "indices = []\n",
    "\n",
    "for index, row in df_tjn.iterrows():\n",
    "    filename = \"./data/large_asteroids/\"+row['horizons']+\".csv\"\n",
    "    try:\n",
    "        df = pd.read_csv(filename)\n",
    "        xs, ys = df[\"X\"].tolist(), df[\"Y\"].tolist()\n",
    "        theta = [get_angle(0, 0, x, y) for x, y in zip(xs, ys)]\n",
    "        theta = [np.radians(x) for x in theta]\n",
    "        if theta[-1] > 0.6294830920687847:\n",
    "            df_tjn.loc[index, 'class'] = 'GRK'\n",
    "            t += 1\n",
    "    except:\n",
    "        count += 1\n",
    "\n",
    "df = pd.concat([df_tjn, df_non_tjn])\n",
    "df.to_csv('./data/large_asteroids.csv', index=False)\n",
    "print(len(df[df['class'] == 'GRK']), 'GRKs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "722 GRKs\n"
     ]
    }
   ],
   "source": [
    "df_asts = pd.read_csv('./data/small_asteroids.csv')\n",
    "df_tjn = df_asts[df_asts['class'] == 'TJN'].copy()\n",
    "df_non_tjn = df_asts[df_asts['class'] != 'TJN'].copy()\n",
    "count, t = 0, 0\n",
    "indices = []\n",
    "\n",
    "for index, row in df_tjn.iterrows():\n",
    "    filename = \"./data/small_asteroids/\"+row['horizons']+\".csv\"\n",
    "    try:\n",
    "        df = pd.read_csv(filename)\n",
    "        xs, ys = df[\"X\"].tolist(), df[\"Y\"].tolist()\n",
    "        theta = [get_angle(0, 0, x, y) for x, y in zip(xs, ys)]\n",
    "        theta = [np.radians(x) for x in theta]\n",
    "        \n",
    "        # Angle of Jupiter position (found after HORIZONS search)\n",
    "        if theta[-1] > 0.6294830920687847:\n",
    "            df_tjn.loc[index, 'class'] = 'GRK'\n",
    "            t += 1\n",
    "    except:\n",
    "        count += 1\n",
    "\n",
    "df = pd.concat([df_tjn, df_non_tjn])\n",
    "df.to_csv('./data/small_asteroids.csv', index=False)\n",
    "print(len(df[df['class'] == 'GRK']), 'GRKs')"
   ]
  },
  {
   "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
}