{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Build Damped Random Walk Kernel"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# GPR kernel \ntry: \n    import celerite\nexcept:\n    !pip install celerite\nimport autograd.numpy as np\nimport celerite\nfrom celerite import terms\nfrom celerite import GP\nfrom scipy.optimize import minimize\nimport matplotlib.pyplot as plt \nimport matplotlib\n\nclass DRW_kernel(terms.Term):\n    # build a Damped Random Walk kernel, a means the standard variance, c means the characteristic timescale\n    \n\tparameter_names = (\"log_a\", \"log_c\")\n\t\n\tdef get_real_coefficients(self, params):\n\t\tlog_a, log_c = params\n        \n\t\treturn np.exp(log_a)**2, 1.0/np.exp(log_c)\n\ndef neg_log_like(params, y, gp):\n    gp.set_parameter_vector(params)\n    return -gp.log_likelihood(y)\n\ndef grad_neg_log_like(params, y, gp):\n    gp.set_parameter_vector(params)\n    return -gp.grad_log_likelihood(y)[1]\n\ndef GP_lc(mjd, band, band_err, sigma_in, tau_in):\n    \n#     print(\"True SF infinte: \", sigma_in*np.sqrt(2))\n#     print(\"True Tau: \", tau_in)\n    kernel = DRW_kernel(log_a = np.log(sigma_in), log_c = np.log(tau_in))\n    gp = GP(kernel, mean=np.mean(band))\n    gp.compute(mjd, band_err)  \n\n#     print(\"Initial log likelihood: {0}\".format(gp.log_likelihood(band)))  \n#     print(\"parameter_dict:\\n{0}\\n\".format(gp.get_parameter_dict()))\n#     print(\"parameter_names:\\n{0}\\n\".format(gp.get_parameter_names()))\n#     print(\"parameter_vector:\\n{0}\\n\".format(gp.get_parameter_vector()))\n#     print(\"parameter_bounds:\\n{0}\\n\".format(gp.get_parameter_bounds()))\n    \n    # set parameter bounds\n    sigma_bounds = [0.01, 200]/np.sqrt(2)\n    tau_bounds = [1,60000]\n    loga_bounds = (np.log(min(sigma_bounds)), np.log(max(sigma_bounds)))\n    logc_bounds= (np.log(min(tau_bounds)), np.log(max(tau_bounds)) )\n    bounds = [loga_bounds, logc_bounds]\n#     print(bounds)\n    \n    #find the maximum likelihood parameters for this model\n    initial_params = gp.get_parameter_vector()\n#     bounds = gp.get_parameter_bounds(bounds)\n#     print(bounds)\n    soln = minimize(neg_log_like, initial_params, method=\"L-BFGS-B\", bounds=bounds, args=(band, gp))\n    gp.set_parameter_vector(soln.x)\n    \n    a_out = gp.get_parameter_dict()['kernel:log_a']\n    c_out = gp.get_parameter_dict()['kernel:log_c']\n    SF_inft = np.sqrt(2)*np.exp(a_out)\n    tau_out = np.exp(c_out)\n    \n#     print(\"Estimated SF infinte: \", SF_inft)\n#     print(\"Estimated Tau: \", tau_out)\n    return gp, SF_inft, tau_out\n\ndef predict_lc(gp, mjd, band_name, band, band_err):\n    continuous_mjd = np.linspace(min(mjd),max(mjd),2000)\n    pred_mean, pred_var = gp.predict(band, continuous_mjd, return_var=True)\n    pred_std = np.sqrt(pred_var)\n   \n    #plot the predicted light curve\n    color = \"#ff7f0e\"\n    plt.figure(figsize = (20,10))\n    plt.errorbar(mjd, band, yerr=band_err, fmt=\".k\", capsize=0)\n    plt.plot(continuous_mjd, pred_mean, color=color)\n    plt.fill_between(continuous_mjd, pred_mean+pred_std, pred_mean-pred_std, color=color, alpha=0.3, edgecolor=\"none\")\n    plt.xlabel(\"mjd (days)\",fontsize=20)\n    plt.ylabel(\"flux\",fontsize=20)\n    plt.title(band_name+' band maximum likelihood prediction',fontsize=20)\n    plt.show()\n    return pred_mean, pred_std","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Select AGN from the training dataset"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport csv\n\ndata = pd.read_csv('../input/PLAsTiCC-2018/training_set.csv')\nmeta = pd.read_csv('../input/PLAsTiCC-2018/training_set_metadata.csv')\nagn_meta = meta[(meta.ddf<1)& (meta.target == 88)]\nagn_data = data[(data.object_id.isin(agn_meta['object_id'].tolist()))]\ndel data, meta\nmodelpar = pd.read_csv('../input/plasticc-converted-datasets/plasticc_modelpar_088_AGN.csv')\n\nsf_name = {'u':'SFU','g':'SFG','r':'SFR','i':'SFI','z':'SFZ','y':'SFY'}\nband_name = {'u':0,'g':1,'r':2,'i':3,'z':4,'y':5}\ncombine = agn_data.groupby('object_id')\nid_list = []\nfor _id, info in combine:\n    id_list.append(_id)\n\n# create a file storing the EM results\ndrw_file = open('drw_training_agn_redshift.csv','w')\nwfile = csv.writer(drw_file)\nwfile.writerow(['object_id', 'passband','tau_in','SF_in','tau_out','SF_out'])\n\n# id_list = id_list[7:9]\nfor i in id_list:\n    agn_obj = agn_data[(agn_data.object_id == i)]\n    paras = modelpar[(modelpar.object_id == i)]\n    true_tau = paras.iloc[0].at['TAU']\n    z = paras.iloc[0].at['REDSHIFT']\n    for band in ['u','g','r','i','z','y']:     \n        try:\n            band_info = agn_obj[(agn_obj.passband==band_name[band])]\n            t_obs = band_info.mjd.values\n            t_obs = t_obs - band_info.mjd.min()\n            t_obs = t_obs/(z+1.0)\n            flux = band_info.flux.values - band_info.flux.mean()\n            flux_err = band_info.flux_err.values\n            true_SF = paras.iloc[0].at[sf_name[band]]\n            gp, SF_out, tau_out = GP_lc(t_obs, flux, flux_err, true_SF/np.sqrt(2), true_tau)\n            wfile.writerow([i, band, true_tau, true_SF, tau_out, SF_out])\n        #   predict_lc(gp, t_obs,band,flux,flux_err)\n        except:\n            wfile.writerow([i, band, true_tau, true_SF, None, None])\n            continue\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Select AGN from test sets"},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport csv\nimport os\nfrom multiprocessing import Process\nfrom multiprocessing import Manager\nfrom multiprocessing import Pool\nimport multiprocessing as mp\nimport time\n\nsf_name = {'u':'SFU','g':'SFG','r':'SFR','i':'SFI','z':'SFZ','y':'SFY'}\nband_name = {'u':0,'g':1,'r':2,'i':3,'z':4,'y':5}\n    \ndef process_work(id_list, agn_data, modelpar):\n    print('process id:', os.getpid())\n#     print('agn_data ',os.getpid(), agn_data)\n#     print('modelpar ',os.getpid(), modelpar)\n#     print('id list ',os.getpid(),id_list)\n    row = pd.DataFrame(columns = ['object_id', 'passband','tau_in','SF_in','tau_out','SF_out'])\n    for i in id_list:\n        agn_obj = agn_data[(agn_data.object_id == i)]\n        paras = modelpar[(modelpar.object_id == i)]\n        true_tau = paras.iloc[0].at['TAU']\n#         z = paras.iloc[0].at['REDSHIFT']\n        for band in ['u','g','r','i','z','y']:     \n            try:\n                band_info = agn_obj[(agn_obj.passband==band_name[band])]\n                t_obs = band_info.mjd.values\n                t_obs = t_obs - band_info.mjd.min()\n#                 t_obs = t_obs/(z+1.0)\n                flux = band_info.flux.values - band_info.flux.mean()\n                flux_err = band_info.flux_err.values\n                true_SF = paras.iloc[0].at[sf_name[band]]\n                gp, SF_out, tau_out = GP_lc(t_obs, flux, flux_err, true_SF/np.sqrt(2), true_tau)\n                row = row.append({'object_id':i, 'passband':band, 'tau_in':true_tau, 'SF_in':true_SF, 'tau_out':tau_out, 'SF_out':SF_out},ignore_index=True)\n#                 print(row)\n#                 print({'object_id':i, 'passband':band, 'tau_in':true_tau, 'SF_in':true_SF, 'tau_out':tau_out, 'SF_out':SF_out})\n            except:\n                row = row.append({'object_id':i, 'passband':band, 'tau_in':true_tau, 'SF_in':true_SF, 'tau_out':None, 'SF_out':None},ignore_index=True)\n#                 print(row)\n                continue\n                    \n    print(os.getpid(), ' finished!')\n#     print(os.getpid(), row)\n    return row\n\ndef calculate_drw(batch, agn_meta_obj, modelpar):\n\n    \n    data = pd.read_csv('../input/plasticcunblindeddatasets/plasticc_test_set_batch'+str(batch)+'.csv')\n    agn_data = data[(data.object_id.isin(agn_meta_obj))]\n    del data\n\n    if len(agn_data)!=0:\n        \n        combine = agn_data.groupby('object_id')\n        id_list = []\n        for _id, info in combine:\n            id_list.append(_id)\n        print('total id_list num: ', len(id_list)) \n        \n        # divide the dataframe into N sub dataframes to N processes\n        Process_num = 4\n        chunk = len(id_list)//Process_num\n        idx = 0\n        \n        params = []\n        \n        for p in np.arange(0,Process_num,1): \n            if p != Process_num-1:\n                sub_list = id_list[idx:idx+chunk]\n                idx += chunk\n            else:\n                sub_list = id_list[idx:]\n                \n            agn_chunk = agn_data[(agn_data.object_id.isin(sub_list))]\n            params.append([sub_list, agn_chunk, modelpar])\n        del agn_data\n        p = Pool(mp.cpu_count())\n        rows = p.starmap(process_work, params)\n        p.close()\n                   \n        df = pd.concat(rows, ignore_index = True)\n        df.to_csv('drw_test_agn_'+str(batch)+'.csv')\n    \nif __name__ == '__main__':\n    # create a file storing the EM results\n\n\n    meta = pd.read_csv('../input/plasticcunblindeddatasets/plasticc_test_metadata.csv')\n    agn_meta = meta[(meta.ddf_bool<1) & (meta.true_target == 88)]\n    agn_meta_obj = agn_meta['object_id'].tolist()\n    del meta, agn_meta\n    modelpar = pd.read_csv('../input/plasticc-converted-datasets/plasticc_modelpar_088_AGN.csv')\n\n    for batch in np.arange(2,12,1):\n        print('batch: ', batch)\n        time_start=time.time()\n        calculate_drw(batch, agn_meta_obj,modelpar)\n        time_end=time.time()\n        print('time cost',time_end-time_start,'s')\n\n\n    ","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}