<jupyter_start><jupyter_text>Extract SNe<jupyter_code>#!pip install psycopg2-binary
import psycopg2
import json
credentials_file = "../../usecases/alercereaduser_v2.json"
with open(credentials_file) as jsonfile:
    params = json.load(jsonfile)["params"]
conn = psycopg2.connect(dbname=params['dbname'], user=params['user'], host=params['host'], password=params['password'])
query = "select * from class"
classes = pd.read_sql_query(query, conn)
classmapper = dict(zip(classes.name.tolist(), classes.id.tolist()))
classmapper_invert = dict(zip(classes.id.tolist(), classes.name.tolist()))
classmapper
query='''
select oid
from objects
where classrf in ('%i', '%i', '%i', '%i')
''' % (classmapper["SNIa"], classmapper["SNIbc"], classmapper["SNII"], classmapper["SLSN"])
sn = pd.read_sql_query(query, conn)
display(sn.head())
sn.shape
#data = pd.read_csv("../data_examples/csv/raw_detections.csv", index_col=False)
#'ZTF19aazzpje' #'ZTF20aaelulu' #'ZTF19aazzpje' # SN
#'ZTF18aaiscil' # weird
#'ZTF18abdgukn' # AGN
#'ZTF17aaaedmd' # LPV
#"ZTF19abaejrh"
#"ZTF20aaelulu"<jupyter_output><empty_output><jupyter_text>Work with test parquets<jupyter_code>det = pd.read_parquet("./det_sample.parquet")
det["mjd"] = det.jd - 2400000.5
non_det = pd.read_parquet("./non_det_sample.parquet")
non_det["mjd"] = non_det.jd - 2400000.5
# check for oid in non_det missing from det
len(set(non_det.objectId.unique()).difference(set(det.objectId.unique())))
# count number of detections
ndet = det.groupby("objectId").apply(len)
# select a given number of oids
seloid = ndet[ndet > 5].index.values
#seloid = sn.oid.values[:1000]
# filter detections
det = det.loc[det.objectId.isin(seloid)].copy()
non_det = non_det.loc[non_det.objectId.isin(seloid)][["objectId", "fid", "mjd", "diffmaglim"]].copy()
det.head()
det.shape
det.objectId.unique()
non_det.head()
non_det.shape
non_det.objectId.unique()
# test c2
#data.at[180, "isdiffpos"] = 'f'
# test c3
#data.at[150, "distnr"] = 0.1
pd.options.display.max_rows = 999
display(det[["objectId", "fid", "magpsf", "distnr", "chinr", "sharpnr"]])
def apply_correction_df(data):
    
    # create copy of dataframe
    df = data.copy()
    df.set_index("candid", inplace=True)
    
    # map isdiffpos
    df['isdiffpos'] = df['isdiffpos'].map({'t': 1., '1': 1, '0': -1, 'f': -1.})

    # apply correction where distnr < 1.4
    df["corrected"] = df["distnr"] < 1.4

    # correct
    correction_results = df.apply(lambda x: correction(x.magnr, x.magpsf, x.sigmagnr, x.sigmapsf, x.isdiffpos, oid=df.objectId.iloc[0], candid=x.name) if x["corrected"] else (np.nan, np.nan, np.nan), axis=1, result_type="expand")
    try:
        df["magpsf_corr"], df["sigmapsf_corr"], df["sigmapsf_corr_ext"] = correction_results[0], correction_results[1], correction_results[2]
    except:
        display(correction_results)
        print(correction_results)
    
    # check if suspicious behavior in the light curve (c1 | c2 | c3)
    corr_magstats = df.loc[df.index.min()]["corrected"]
    mask = ((df["corrected"] == False) & (df.isdiffpos == -1)) | (corr_magstats & (df["corrected"] == False)) | ((corr_magstats == False) & df["corrected"])
    df["dubious"] = mask
    
    return df
def apply_mag_stats(df):
    
    response = {}
    
    # minimum and maximum candid
    idxmin = df.index.min()
    idxmax = df.index.max()

    # corrected at the first detection?
    response['corrected'] = df.loc[idxmin]["corrected"]
    
    # stellar_object at first detection?
    nearZTF = (df.loc[idxmin].distnr >= 0) & (df.loc[idxmin].distnr < 1.4) # near a ZTF object
    nearPS1 = (df.loc[idxmin].distpsnr1 >= 0) & (df.loc[idxmin].distpsnr1 < 1.4) # near a ZTF object
    stellarPS1 = (df.loc[idxmin].sgscore1 > 0.4)  # nearest object in PS1 is stelar
    stellarZTF = (df.loc[idxmin].chinr < 2) & (df.loc[idxmin].sharpnr > -0.13) & (df.loc[idxmin].sharpnr < 0.1) # nearest object in ZT is stellar
    response["nearZTF"] = nearZTF
    response["nearPS1"] = nearPS1
    response["stellarZTF"] = nearZTF
    response["stellarPS1"] = nearPS1
    response["stellar"] = (nearZTF & nearPS1 & stellarPS1) | (nearZTF & ~nearPS1 & stellarZTF)
    
    # number of detections and dubious detections
    response["ndet"] = df.shape[0]
    response["ndubious"] = df.dubious.sum()

    # reference id
    response["nrfid"] = len(df.rfid.dropna().unique())

    # psf magnitude statatistics
    response["magpsf_mean"] = df.magpsf.mean()
    response["magpsf_median"] = df.magpsf.median()
    response["magpsf_max"] = df.magpsf.max()
    response["magpsf_min"] = df.magpsf.min()
    response["magpsf_first"] = df.loc[idxmin].magpsf
    response["sigmapsf_first"] = df.loc[idxmin].sigmapsf
    response["magpsf_last"] = df.loc[idxmax].magpsf
    
    # psf corrected magnitude statatistics
    response["magpsf_corr_mean"] = df.magpsf_corr.mean()
    response["magpsf_corr_median"] = df.magpsf_corr.median()
    response["magpsf_corr_max"] = df.magpsf_corr.max()
    response["magpsf_corr_min"] = df.magpsf_corr.min()
    response["magpsf_corr_first"] = df.loc[idxmin].magpsf_corr
    response["magpsf_corr_last"] = df.loc[idxmax].magpsf_corr
    
    # corrected psf magnitude statistics
    response["magap_mean"] = df.magap.mean()
    response["magap_median"] = df.magap.median()
    response["magap_max"] = df.magap.max()
    response["magap_min"] = df.magap.min()
    response["magap_first"] = df.loc[idxmin].magap
    response["magap_last"] = df.loc[idxmax].magap
           
    # time statistics
    response["first_mjd"] = df.loc[idxmin].mjd
    response["last_mjd"] = df.loc[idxmax].mjd
       
    return pd.Series(response)
def apply_object_stats(df):
    
    response = {}

    response["nearZTF"] = df.nearZTF.all()
    response["nearPS1"] = df.nearPS1.all()
    response["stellar"] = df.stellar.all()
    response["corrected"] = df.corrected.all()
    response["ndet"] = df.ndet.sum() # sum of detections in all bands
    response["ndubious"] = df.ndubious.sum() # sum of dubious corrections in all bands
    
    return pd.Series(response)
def do_dmdt(nd):
    
    response = {}
    objectId = nd.objectId.iloc[0]
    fid = nd.fid.iloc[0]
    if fid in magstats.loc[objectId].index:
        mjd_first = magstats.loc[objectId].loc[fid].first_mjd
        if nd.mjd.min() < mjd_first:
            mask = nd.mjd < mjd_first
            magpsf_first = magstats.loc[objectId].loc[fid].magpsf_first
            sigmapsf_first = magstats.loc[objectId].loc[fid].sigmapsf_first
            dm = magpsf_first - nd.loc[mask].diffmaglim
            dm_sigma = magpsf_first + sigmapsf_first - nd.loc[mask].diffmaglim
            dt = mjd_first - nd.loc[mask].mjd
            dmsigdt = (dm_sigma / dt)
            idxmin = dmsigdt.idxmin()
            response["dmdt_first"] = dmsigdt.loc[idxmin]
            response["dm_first"] = magpsf_first - nd.diffmaglim.loc[idxmin]
            response["sigmadm_first"] = sigmapsf_first - nd.diffmaglim.loc[idxmin]
            response["dt_first"] = dt.loc[idxmin]
        else:
            response["dmdt_first"] = np.nan
            response["dm_first"] = np.nan
            response["sigmadm_first"] = np.nan
            response["dt_first"] = np.nan
    else:
        response["dmdt_first"] = np.nan
        response["dm_first"] = np.nan
        response["sigmadm_first"] = np.nan
        response["dt_first"] = np.nan

    return pd.Series(response)<jupyter_output><empty_output><jupyter_text>Do the magic<jupyter_code>if len(non_det.index.names) == 2:
    non_det.reset_index(level=[0, 1], inplace=True)
non_det.head()
dflarge = applyParallel(det.groupby(["objectId", "fid"]), apply_correction_df)
dflarge.head(2)
magstats = dflarge.groupby(["objectId", "fid"]).apply(apply_mag_stats)
magstats.head(2)
objstats = magstats.groupby(["objectId"]).apply(apply_object_stats)
objstats.head(2)
dmdt = pd.DataFrame(non_det.groupby(["objectId", "fid"]).apply(do_dmdt))
magstats = magstats.join(dmdt)
magstats
non_det.set_index(["objectId", "fid"], inplace=True)<jupyter_output><empty_output><jupyter_text>Plot some examples<jupyter_code>#for oid in dflarge.objectId.unique():

def plot_object(oid):
    
    colors = {1: 'g', 2: 'r'}
    period = {}
    period['ZTF18aazxcwf'] = 0.614515 # RRL

        
    fig, ax = plt.subplots(ncols = 3, sharey = False, figsize=(30, 6))
    maskoid = dflarge.objectId == oid
    display(pd.DataFrame(objstats.loc[oid][["corrected", "nearZTF", "stellar"]]))
    display(magstats.loc[oid][["nearZTF", "nearPS1", "stellarZTF", "stellarPS1", "corrected", "nearZTF", "nearPS1", "stellar", "dmdt_first"]])
    #display(dflarge.loc[maskoid][["rfid", "fid", "mjd", "magpsf", "magnr", "isdiffpos", "magpsf_corr", "corrected", "dubious"]])
    
    for fid in dflarge.loc[maskoid].fid.unique():
        mask = maskoid & (dflarge.fid == fid)

        color = colors[fid]
        
        # non-detections
        if oid in non_det.index.levels[0]:
            if fid in non_det.loc[oid].index:
                ax[0].scatter(non_det.loc[oid].loc[fid].mjd, non_det.loc[oid].loc[fid].diffmaglim, marker = 'v', color=color, alpha=0.2)

        for pos in dflarge.loc[mask].isdiffpos.unique():
            
            mask = mask & (dflarge.isdiffpos == pos)
            mask_diff = mask & (dflarge.sigmapsf < 2)
            mask_corr = mask & (dflarge.sigmapsf_corr < 2)
            mask_corr_ext = mask & (dflarge.sigmapsf_corr_ext < 2)
                
            marker = 'o' if pos == 1 else '*'
            neg = '' if pos == 1 else ' (negative)'

            if oid not in period.keys():
                ax[0].errorbar(dflarge.loc[mask_diff].mjd, dflarge.loc[mask_diff].magpsf, yerr=dflarge.loc[mask_diff].sigmapsf, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf%s" % neg)
                ax[1].errorbar(dflarge.loc[mask_corr].jd, dflarge.loc[mask_corr].magpsf_corr, yerr=dflarge.loc[mask_corr].sigmapsf_corr, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf_corr%s" % neg)
                ax[2].errorbar(dflarge.loc[mask_corr_ext].jd, dflarge.loc[mask_corr_ext].magpsf_corr, yerr=dflarge.loc[mask_corr_ext].sigmapsf_corr_ext, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf_corr_ext%s" % neg)
            else:
                ax[0].errorbar(np.mod(dflarge.loc[mask_diff].mjd, period[oid]), dflarge.loc[mask_diff].magpsf, yerr=dflarge.loc[mask_diff].sigmapsf, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf%s" % neg)
                ax[1].errorbar(np.mod(dflarge.loc[mask_corr].mjd, period[oid]), dflarge.loc[mask_corr].magpsf_corr, yerr=dflarge.loc[mask_corr].sigmapsf_corr, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf_corr%s" % neg)
                ax[2].errorbar(np.mod(dflarge.loc[mask_corr_ext].mjd, period[oid]), dflarge.loc[mask_corr_ext].magpsf_corr, yerr=dflarge.loc[mask_corr_ext].sigmapsf_corr_ext, c=color, marker=marker, lw=0, elinewidth=1, alpha=1, label="magpsf_corr_ext%s" % neg)
    
    for i in range(3):
        ax[i].set_title("oid: %s, corrected: %s, dmdt: %s" % (oid, objstats.loc[oid].corrected, magstats.loc[oid].dmdt_first.values))
        ax[i].set_ylim(ax[i].get_ylim()[::-1])
        ax[i].legend()
        if oid not in period.keys():
            ax[i].set_xlabel("jd")
        else:
            ax[i].set_xlabel("phase")
for idx, oid in enumerate(dflarge.objectId.unique()):#sn.loc[sn.oid.isin(seloid)].oid):
    if idx < 50: # and oid == "ZTF19abdkgmf":
        print(oid)
        plot_object(oid)
dflarge.memory_usage(index=True).sum()
non_det.memory_usage(index=True).sum()
df.oid.isin(seloid).sum()
corr_f = []
corr_ext_f = []
masksigmapsf = dflarge.sigmapsf < 2
masksigmapsf_corr = dflarge.sigmapsf_corr < 2
masksigmapsf_corr_ext = dflarge.sigmapsf_corr_ext < 2
for oid in dflarge.objectId.unique():
    maskoid = dflarge.objectId == oid
    for fid in dflarge.loc[maskoid].fid.unique():
        mask = maskoid & (dflarge.fid == fid)
        #print(oid, fid, dflarge.loc[mask].isdiffpos.unique())
        corr_f.append(np.sum(mask & masksigmapsf_corr) / np.sum(mask & masksigmapsf))
        corr_ext_f.append(np.sum(mask & masksigmapsf_corr_ext) / np.sum(mask & masksigmapsf))
        #log["stellar%i" % fid] = dflarge.loc[mask].stellar_magstats.unique()
fig, ax = plt.subplots(ncols=2, figsize=(15, 6))
ax[0].hist(corr_f, bins=100);
ax[1].hist(corr_ext_f, bins=100);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1639648358738608, 'fraction_numerical': 0.014759254778611178, 'mean_word_length': 3.800619434765776, 'pattern_counts': {'":': 1, '<': 32, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '98558', 'license': 'MIT', 'n_tokens_mistral': 4943, 'n_tokens_neox': 4770, 'n_words': 765, 'path': 'notebooks/apply_correction_to_df_parquet_test.ipynb', 'repo_name': 'alercebroker/lc_correction', 'size': '12399'}