crossmatch_lyrae_ztf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import axs\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "from pyspark.sql.functions import size as spark_size\n",
    "from pyspark.sql.types import ArrayType, DoubleType\n",
    "from pyspark.sql.functions import pandas_udf\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import cesium\n",
    "from cesium.time_series import TimeSeries\n",
    "from cesium.featurize import featurize_single_ts, featurize_time_series\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spark_start(project_path, metastore=None):\n",
    "    from pyspark.sql import SparkSession\n",
    "\n",
    "    warehouse_location = os.path.join(project_path, 'spark-warehouse')\n",
    "\n",
    "    local_dir = os.path.join(project_path, 'spark-tmp')\n",
    "\n",
    "    spark = ( \n",
    "            SparkSession.builder\n",
    "            .appName(\"LSD2\")\n",
    "            .config(\"spark.sql.warehouse.dir\", warehouse_location)\n",
    "            .config('spark.master', \"local[4]\")\n",
    "            .config('spark.driver.memory', '6G') # 128\n",
    "            .config('spark.local.dir', local_dir)\n",
    "            .config('spark.memory.offHeap.enabled', 'true')\n",
    "            .config('spark.memory.offHeap.size', '4G') # 256\n",
    "            .config(\"spark.sql.execution.arrow.enabled\", \"true\")\n",
    "            .config(\"spark.driver.maxResultSize\", \"6G\")\n",
    "            .config(\"spark.driver.extraJavaOptions\", f\"-Dderby.system.home={metastore}\")\n",
    "            .enableHiveSupport()\n",
    "            .getOrCreate()\n",
    "                    )   \n",
    "\n",
    "    return spark\n",
    "\n",
    "spark_session = spark_start(\"/epyc/users/ctslater\")\n",
    "\n",
    "catalog = axs.AxsCatalog(spark_session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ztf = catalog.load(\"ztf_1am_lc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sesar_axs = catalog.load(\"sesar_rrlyrae\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "matched = sesar_axs.crossmatch(ztf).drop(\"axsdist\")\n",
    "matched_filtered = (matched.select(\"ra\", \"dec\", \"matchid\", \"Per\", \"weightedmeanmag\", \"filterid\", \"mjd\", \"psfflux\")\n",
    "                    .where((spark_size(matched['mjd']) > 5) &\n",
    "                           ((matched['S3ab'] > 0.8) | (matched['S3c'] > 0.8))\n",
    "                            ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cesium in the python layer\n",
    "========"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.36 s, sys: 145 ms, total: 1.5 s\n",
      "Wall time: 51.6 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "results = matched_filtered.head(10)\n",
    "result = results[0]\n",
    "\n",
    "features_to_use = [\"amplitude\",\n",
    "                   \"percent_beyond_1_std\",\n",
    "                   \"maximum\",\n",
    "                   \"max_slope\",\n",
    "                   \"median\",\n",
    "                   \"median_absolute_deviation\",\n",
    "                   \"percent_close_to_median\",\n",
    "                   \"minimum\",\n",
    "                   \"skew\",\n",
    "                   \"std\",\n",
    "                   \"weighted_average\"]\n",
    "\n",
    "ls_features = [\"freq1_amplitude1\",\n",
    "                \"freq1_amplitude2\",\n",
    "                \"freq1_amplitude3\",\n",
    "                \"freq1_amplitude4\",\n",
    "                \"freq1_freq\",\n",
    "                \"freq1_lambda\",\n",
    "                \"freq1_rel_phase2\",\n",
    "                \"freq1_rel_phase3\",\n",
    "                \"freq1_rel_phase4\",\n",
    "                \"freq1_signif\",\n",
    "                \"freq2_amplitude1\",\n",
    "                \"freq2_amplitude2\",\n",
    "                \"freq2_amplitude3\",\n",
    "                \"freq2_amplitude4\",\n",
    "                \"freq2_freq\",\n",
    "                \"freq2_rel_phase2\",\n",
    "                \"freq2_rel_phase3\",\n",
    "                \"freq2_rel_phase4\"]\n",
    "\n",
    "ts = TimeSeries(t=np.array(result['mjd']), m=np.array(result['psfflux']))\n",
    "feat_out = featurize_single_ts(ts, features_to_use + ls_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "channel\n",
       "0    0.574923\n",
       "dtype: float64"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/feat_out['freq1_freq']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.57479179"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result['Per']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "feature                    channel\n",
       "amplitude                  0            9396.832794\n",
       "percent_beyond_1_std       0               0.307190\n",
       "maximum                    0           19601.496094\n",
       "max_slope                  0          825072.546836\n",
       "median                     0           12735.896484\n",
       "median_absolute_deviation  0            2143.480469\n",
       "percent_close_to_median    0               0.431373\n",
       "minimum                    0             807.830505\n",
       "skew                       0              -0.375430\n",
       "std                        0            3359.907709\n",
       "weighted_average           0           13263.989591\n",
       "freq1_amplitude1           0            3022.199473\n",
       "freq1_amplitude2           0            1063.568925\n",
       "freq1_amplitude3           0             230.753765\n",
       "freq1_amplitude4           0              16.869063\n",
       "freq1_freq                 0               1.739364\n",
       "freq1_lambda               0               7.425334\n",
       "freq1_rel_phase2           0              -0.952610\n",
       "freq1_rel_phase3           0              -1.643048\n",
       "freq1_rel_phase4           0               2.473324\n",
       "freq1_signif               0               7.547384\n",
       "freq2_amplitude1           0             892.550622\n",
       "freq2_amplitude2           0              59.022117\n",
       "freq2_amplitude3           0               3.508887\n",
       "freq2_amplitude4           0               2.409814\n",
       "freq2_freq                 0               1.353818\n",
       "freq2_rel_phase2           0               1.866372\n",
       "freq2_rel_phase3           0               1.857011\n",
       "freq2_rel_phase4           0               3.126963\n",
       "dtype: float64"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feat_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 92.3 ms, sys: 36.8 ms, total: 129 ms\n",
      "Wall time: 13min 13s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "109850351"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "ztf.where(spark_size(matched['mjd']) > 10).count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 284 ms, sys: 1.58 s, total: 1.87 s\n",
      "Wall time: 74.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ts = TimeSeries(t=np.array(result['mjd']), m=np.array(result['psfflux']))\n",
    "feat_out = featurize_single_ts(ts, features_to_use + ls_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/epyc/opt/anaconda/lib/python3.6/site-packages/cesium/features/common_functions.py:7: RuntimeWarning: invalid value encountered in true_divide\n",
      "  slopes = np.diff(x) / np.diff(t)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 2s, sys: 3min 25s, total: 4min 27s\n",
      "Wall time: 3.37 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "feat_out2 = featurize_time_series([np.array(res['mjd']) for res in results[:20]],\n",
    "                                  [np.array(res['psfflux']) for res in results[:20]],\n",
    "                                  features_to_use=features_to_use + ls_features\n",
    "                                  )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cesium UDF\n",
    "======="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "features_to_use = [\"amplitude\",\n",
    "                   \"percent_beyond_1_std\",\n",
    "                   \"maximum\",\n",
    "                   \"max_slope\",\n",
    "                   \"median\",\n",
    "                   \"median_absolute_deviation\",\n",
    "                   \"percent_close_to_median\",\n",
    "                   \"minimum\",\n",
    "                   \"skew\",\n",
    "                   \"std\",\n",
    "                   \"weighted_average\"]\n",
    "\n",
    "ls_features = [\"freq1_amplitude1\",\n",
    "                \"freq1_amplitude2\",\n",
    "                \"freq1_amplitude3\",\n",
    "                \"freq1_amplitude4\",\n",
    "                \"freq1_freq\",\n",
    "                \"freq1_lambda\",\n",
    "                \"freq1_rel_phase2\",\n",
    "                \"freq1_rel_phase3\",\n",
    "                \"freq1_rel_phase4\",\n",
    "                \"freq1_signif\",\n",
    "                \"freq2_amplitude1\",\n",
    "                \"freq2_amplitude2\",\n",
    "                \"freq2_amplitude3\",\n",
    "                \"freq2_amplitude4\",\n",
    "                \"freq2_freq\",\n",
    "                \"freq2_rel_phase2\",\n",
    "                \"freq2_rel_phase3\",\n",
    "                \"freq2_rel_phase4\"]\n",
    "    \n",
    "def featurize_udf(mjd, psfflux):\n",
    "    feat_outs = []\n",
    "    \n",
    "    for row_mjd, row_psfflux in zip(mjd, psfflux):\n",
    "        mjd_arr = np.array(row_mjd)\n",
    "        \n",
    "        # If the time between the first and last observation is <0.06, \n",
    "        # cesium throws an error.\n",
    "        if(np.max(mjd_arr) - np.min(mjd_arr) < 0.10):\n",
    "            feat_outs.append(np.zeros(len(features_to_use) + len(ls_features)))\n",
    "        else:\n",
    "            feat_out = featurize_time_series(mjd_arr, np.array(row_psfflux),\n",
    "                                             features_to_use=features_to_use + ls_features)\n",
    "            feat_outs.append(feat_out.values.flatten())\n",
    "\n",
    "    return pd.Series(feat_outs)\n",
    "\n",
    "featurize_lc = pandas_udf(featurize_udf, returnType=ArrayType(DoubleType()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 50.9 ms, sys: 14 ms, total: 65 ms\n",
      "Wall time: 5min 46s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "feature_results = matched_filtered.select(\n",
    "    featurize_lc(matched_filtered['mjd'], matched_filtered['psfflux'])).head(1000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Row(featurize_udf(mjd, psfflux)=[9396.832794189453, 0.30718954248366015, 19601.49609375, 825072.5468358505, 12735.896484375, 2143.48046875, 0.43137254901960786, 807.8305053710938, -0.37543006309128435, 3359.907708598801, 13263.989590912863, 3022.1994726892235, 1063.5689251811777, 230.7537648132592, 16.869063248696087, 1.7393636957614456, 7.425334084418403, -0.9526102393556031, -1.6430480738036535, 2.4733237980405356, 7.547383578188932, 892.550621721922, 59.02211695616729, 3.508887093596339, 2.409814231998434, 1.3538175947376925, 1.8663721976359344, 1.857010582095423, 3.1269629685411493])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_results[0]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3-AXS Spark",
   "language": "python",
   "name": "spark-smj"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}