{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Dataset anisotropy\n## Some Useful Functions","metadata":{}},{"cell_type":"code","source":"import os, gc, sys, time, datetime, math, random,  psutil\nimport numpy as np,  matplotlib.pyplot as plt, pandas as pd\nfrom   pathlib   import Path        \nfrom   tqdm.auto import tqdm\nimport pyarrow, pyarrow.parquet as pq     # read by chanks\n\nPATH = Path(\"/kaggle/input/icecube-neutrinos-in-deep-ice\")  # path to dataset\nfiles_trn = [item for item in (PATH  / \"train\").glob('*')]  # all train files\nprint(f\"{len(files_trn):3d} train files\")\n\n#===============================================================================\ndef info(text, pref=\"\", end=\"\\n\"):\n    \"\"\" Information about the progress of calculations (time and memory) \"\"\"\n    gc.collect()\n    ram, t = psutil.virtual_memory().used / 1024**3,  time.time()    \n    print(f\"{pref}{(t-info.beg)/60:5.1f}m[{t-info.last:+5.1f}s] {ram:6.3f}Gb > {text}\",end=end)\n    info.last = time.time(); \ninfo.beg = info.last = time.time()\n\n#-------------------------------------------------------------------------------\n\ndef get_sensors():\n    \"\"\" Get sensor positions \"\"\"    \n    df = pd.read_csv(PATH / \"sensor_geometry.csv\")      \n    df['line_id'] = (df.sensor_id // 60 + 1).astype(np.int16) # string id\n    df['core']    = (df.line_id > 78).astype(np.float32)      # sensor from DeepCore\n    df.x = ( df.x * 1e-3 ).astype(np.float32)                 # distances in kilometers\n    df.y = ( df.y * 1e-3 ).astype(np.float32)\n    df.z = ( df.z * 1e-3 ).astype(np.float32)\n    df.sensor_id = df.sensor_id.astype(np.int16)\n    return df[['sensor_id', 'line_id', 'core', 'x', 'y', 'z']]\n\n#-------------------------------------------------------------------------------\n\ndef get_target_angles(batches=set(), verbose = True):\n    \"\"\" Get target angles for batches from  batches id from set batches \"\"\"    \n    res = None\n    file = pq.ParquetFile(PATH / \"train_meta.parquet\")    \n    for b in file.iter_batches(batch_size=200_000, columns=['event_id','batch_id','azimuth','zenith']):    \n        if len(batches) == 0:\n            break\n        batch_df = b.to_pandas()\n        batch_id = batch_df.batch_id[0]\n        if batch_id in batches:      \n            batch_df.event_id= batch_df.event_id.astype(np.int64)      \n            batch_df.azimuth = batch_df.azimuth.astype(np.float32)\n            batch_df.zenith  = batch_df.zenith.astype(np.float32)    \n            batch_df = batch_df[ ['event_id','batch_id','azimuth','zenith'] ]\n            res =  batch_df if res is None else res.append(batch_df)            \n            batches.remove(batch_id)            \n            if verbose: info(f\"load target for batch {batch_id};  {len(batches)} candies left\")\n    return res\n#-------------------------------------------------------------------------------\n\ndef prepare_batch(df, drop_aux = True, doms_agg = False, verbose=True):\n    \"\"\" \n    Preparing a loaded batch, shifting and normalizing times \n        * drop_aux - whether it is necessary to throw out auxiliary pulses (auxiliary==True)\n        * doms_agg - whether it is necessary to aggregate pulses on sensors\n        * verbose  - debug output\n    \"\"\"    \n    df['event_id'] = df.index.astype(np.int64)\n    df = df.reset_index(drop=True)  # sensor_id, t, charge, aux, event_id    \n    df.rename(columns={\"time\": \"t\", \"auxiliary\": \"aux\", 'charge': 'q'}, inplace=True)\n    df.q = df.q.astype(np.float32)\n\n    if drop_aux:\n        df = df[ ~df.aux ]\n    \n    if doms_agg:\n        df = df.groupby(['event_id', 'sensor_id']).agg(\n            aux = ( 'aux', \"mean\"),\n            q   = ( 'q',   \"sum\"),\n            t   = ( 't',   \"min\"),\n            cnt =( 't',   \"count\"),)\n        df = df.reset_index()    \n    else:\n        df['cnt'] = 1\n    df.cnt = df.cnt.astype(np.int32)\n        \n    if verbose: info(f\"prepare_batch: loaded  {df.shape}\")\n        \n    times = df.groupby('event_id').agg( t_min = ('t', 'min') )\n    df = df.merge(times, left_on='event_id', right_index=True, how='left')\n    df.t = (( df.t - df.t_min ) * 0.299792458e-3 ).astype(np.float32)             \n    \n    if verbose: info(\"prepare_batch: shift_times\")    \n\n    return df[['event_id', 'sensor_id', 'aux', 'q', 't','cnt' ]]\n\n#-------------------------------------------------------------------------------\n\ndef agg_features(df):\n    \"\"\" Calculate some features aggregated for a given event \"\"\"\n    df = df.groupby('event_id').agg(\n        lines = ( 'line_id',   'nunique'),\n        doms  = ( 'sensor_id', 'nunique'),\n        pulses= ( 't',         'count'),\n        core  = ( 'core',      'mean'),        \n        t     = ( 't',         'mean'),\n        stdT  = ( 't',         'std'),\n        q     = ( 'q',         'median'),\n        aux   = ( 'aux',       'mean'),\n        x     = ( 'x',         'mean'),\n        y     = ( 'y',         'mean'),\n        z     = ( 'z',         'mean'),\n        stdX  = ( 'x',         'std'),\n        stdY  = ( 'y',         'std'),\n        stdZ  = ( 'z',         'std'),        \n    )\n    return df\n#-------------------------------------------------------------------------------\n\ndef angles2vector(df):\n    \"\"\" Add unit vector components from (azimuth,zenith) to the DataFrame df \"\"\"\n    df['nx'] = np.sin(df.zenith) * np.cos(df.azimuth)\n    df['ny'] = np.sin(df.zenith) * np.sin(df.azimuth)\n    df['nz'] = np.cos(df.zenith) \n    return df\n\n#-------------------------------------------------------------------------------\n\ndef vector2angles(n, eps=1e-8):\n    \"\"\"  Get spherical angles of vector n: (B,3) \"\"\"                \n    n = n / np.linalg.norm(n, axis=1, keepdims=True).clip(eps)    \n                                \n    azimuth = np.arctan2( n[:,1],  n[:,0])    \n    azimuth[azimuth < 0] += 2*np.pi\n                                \n    zenith = np.arccos( n[:,2].clip(-1,1) )                                \n    \n    return azimuth, zenith\n\n#-------------------------------------------------------------------------------\n\ndef angle_errors(n1, n2, eps=1e-8):\n    \"\"\" Calculate angles between two vectors:: n1,n2: (B,3) return: (B,) \"\"\"\n    n1 = n1 / np.linalg.norm(n1, axis=1, keepdims=True).clip(eps)\n    n2 = n2 / np.linalg.norm(n2, axis=1, keepdims=True).clip(eps)\n    \n    cos = (n1*n2).sum(axis=1)                     # angles between vectors\n    angle_err = np.arccos( cos.clip(-1,1) )    \n        \n    r1   =  n1[:,0]*n1[:,0] + n1[:,1]*n1[:,1]    # angles between vectors in (x,y)    \n    r2   =  n2[:,0]*n2[:,0] + n2[:,1]*n2[:,1]\n    cosX = (n1[:,0]*n2[:,0] + n1[:,1]*n2[:,1]) / (np.sqrt(r1*r2) + eps)    \n    azimuth_err = np.arccos( cosX.clip(-1,1) )\n                                \n    zerros = r1 < eps                            # azimuth angle not defined\n    azimuth_err[zerros] = np.random.random((len(n1[zerros]),))*np.pi\n    \n    zenith1  = np.arccos( n1[:,2].clip(-1,1) )\n    zenith2  = np.arccos( n2[:,2].clip(-1,1) )\n    zenith_err = zenith2 - zenith1    \n        \n    return angle_err, azimuth_err, zenith_err\n\n#-------------------------------------------------------------------------------\n\ndef float2index(values, min = None, max = None, num=100):\n    \"\"\" Series of float transform to int index of num bins \"\"\"    \n    min = np.min(values) if min is None else min\n    max = np.max(values) if max is None else max\n    return np.floor(num*(values - min)/(max-min)).clip(0,num-1).astype(int)\n\ndef index2float(index, min, max, num=100):\n    \"\"\" Int index transform to series with min and max values  \"\"\"            \n    return ((index / num)*(max-min) + min).astype(np.float32)\n\n#-------------------------------------------------------------------------------\n\ndef hist(x, x_sub=None, pref=\"\", digits=1, w=12, h=3, bins=50, bins_sub=100):\n    \"\"\" Distribution of values of x and its subset x_sub (in a narrower range) \"\"\"\n    r = lambda x: '{x:.{digits}f}'.format(x=round(x,digits), digits=digits)\n    x_sub = x if x_sub is None else x_sub\n    plt.figure(figsize=(w,h), facecolor ='w')         \n    plt.suptitle(f\"{pref}median={r(x.median())}; mean={r(x.mean())} ± {r(x.std())}  [min,max]=[{r(x.min())}, {r(x.max())}]; cnt={len(x)} ({100*len(x_sub)/len(x):.0f}%)\", fontsize=14)\n    plt.subplot(1,2,1)\n    plt.hist(x, bins=bins, log=True, color=\"lightblue\", ec=\"black\");  plt.grid(ls=\":\",alpha=1); plt.ylabel(\"log10(N)\")\n    plt.subplot(1,2,2)\n    plt.hist(x_sub, bins=bins_sub, density=True, color=\"lightblue\", ec=\"black\");    plt.grid(ls=\":\",alpha=1); plt.ylabel(\"Density\")\n    plt.show()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-14T16:47:24.760289Z","iopub.execute_input":"2023-03-14T16:47:24.760727Z","iopub.status.idle":"2023-03-14T16:47:25.105430Z","shell.execute_reply.started":"2023-03-14T16:47:24.760688Z","shell.execute_reply":"2023-03-14T16:47:25.104421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load train batches and target angles","metadata":{}},{"cell_type":"code","source":"class CFG:\n    drop_aux    = False       # discard pulses with auxiliary == True; the method Line-fit needs\n    doms_agg    = False       # aggregate pulses on sensors    \n    drop_core   = False\n    batches     = range(1,6) # list of batch id-s   \n    \ninfo.beg = info.last = time.time()\ninfo(\"Begin\")\n\n# 1. load sensors\nsensors_df = get_sensors()    \ndisplay(sensors_df.head(2))\n\n# 2. load pulses from batches\npulses_df = None              \nfor batch_id in tqdm(CFG.batches):\n    df = pd.read_parquet(PATH  / \"train\" / f\"batch_{batch_id}.parquet\")     \n    df = prepare_batch(df, drop_aux = CFG.drop_aux, doms_agg = CFG.doms_agg)\n    df = df.merge(sensors_df, left_on='sensor_id', right_on='sensor_id', how=\"left\")\n    pulses_df = df if pulses_df is None else pulses_df.append(df)\ndel df; \ndisplay(pulses_df.head(2))\n\n# 3. load target angles and get unit direction vector\ntrue_df = get_target_angles(batches=set( CFG.batches ))\ntrue_df = angles2vector(true_df)\ndisplay(true_df.head(2))\ninfo(\"loaded target\")\n\nif CFG.drop_core:\n    pulses_df = pulses_df[pulses_df.core==0]    # !!!!\n    info(f\"the End: sensors_df={sensors_df.shape},  pulses_df={pulses_df.shape},  true_df: {true_df.shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-03-14T16:47:25.107131Z","iopub.execute_input":"2023-03-14T16:47:25.107686Z","iopub.status.idle":"2023-03-14T16:49:19.696427Z","shell.execute_reply.started":"2023-03-14T16:47:25.107649Z","shell.execute_reply":"2023-03-14T16:49:19.695288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Aggregate by cos zenith","metadata":{}},{"cell_type":"code","source":"agg_df =  agg_features(pulses_df)\nagg_df = agg_df.merge(true_df, left_index=True, right_on=\"event_id\",  how='left')\n\nagg_df['cos_ze']    = np.cos(agg_df.zenith)\nagg_df['cos_ze_id'] = float2index(agg_df.cos_ze, -1, 1, num=25)\n\ndisplay(agg_df.head(2))\ninfo(f\"agregated: agg_df={agg_df.shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-03-14T16:49:19.697842Z","iopub.execute_input":"2023-03-14T16:49:19.698163Z","iopub.status.idle":"2023-03-14T16:51:35.662044Z","shell.execute_reply.started":"2023-03-14T16:49:19.698130Z","shell.execute_reply":"2023-03-14T16:51:35.660926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plot data","metadata":{}},{"cell_type":"code","source":"from matplotlib.ticker import FormatStrFormatter\n\ndef plot_errorbar(df, col, r=1, c=3, pos=1, mean=None, lims=None):    \n    plt.subplot(r,c,pos); \n    plt.grid(ls=\":\")\n    if lims is not None:\n        plt.ylim(lims[0], lims[1])\n    plt.plot(df.cos_ze, df.m);  plt.xlabel('cos(zenith)'); \n    plt.title(f\"{col}:    mean={mean:.3f}\")\n    \n    plt.errorbar(df.cos_ze, df.m, yerr = df.s,  fmt='o',  ecolor = 'lightblue', color='blue', capsize=3, ms=3)    \n    ax1 = plt.gca();  ax1.xaxis.grid(True)\n    ax2 = ax1.twinx()\n    \n    prob = df.cnt / df.cnt.sum();  ax2.set_ylim(0.03, 0.06)\n    ax2.set_yticks(np.linspace(0.03,0.06, 4))\n    #print(prob.max()/prob.min(), prob.max(), prob.min())\n    ax2.plot(df.cos_ze, prob, ls=\"--\", color=\"gray\", alpha=1, linewidth=0.8); \n    if pos == c:\n        ax2.set_ylabel(\"density events\")\n        ax2.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n    else:\n        ax2.set_yticks([])\n    return ax1, ax2\n\ndef plot_cols(cols, lims=None):    \n    plt.figure(figsize=(13,3), facecolor ='w')    \n    if lims is None:\n        lims = [None] * len(cols)\n    if len(lims)==1:\n        lims = lims * len(cols)\n    for i, col in enumerate(cols):    \n        mean = agg_df[col].mean()        \n        df = agg_df.groupby('cos_ze_id').agg( m=(col,'mean'), s=(col, 'std'), cnt=('t','count'), cos_ze = ('cos_ze', 'mean') )   \n        df.s = df.s/np.sqrt(df.cnt)    \n        ax1, ax2 = plot_errorbar(df, col=col, pos=i+1, mean=mean, lims=lims[i])            \n        ax1.hlines(mean, df.cos_ze.min(), df.cos_ze.max(), ls='-', color='b', lw=0.5)\n    plt.show()\n\nplot_cols(['lines', 'doms', 'pulses'], [(30, 33), (55,70), None] )\nplot_cols(['q', 't', 'aux'],           [(0.91, 0.93), (1.55, 1.65), (0.6, 0.7)] )\nplot_cols(['x', 'y', 'z'],              [(-0.03, 0.02)])\nplot_cols(['stdX', 'stdY', 'stdZ'],     [(0.22, 0.29)] )\n\n","metadata":{"execution":{"iopub.status.busy":"2023-03-14T16:51:35.664870Z","iopub.execute_input":"2023-03-14T16:51:35.665363Z","iopub.status.idle":"2023-03-14T16:51:38.576831Z","shell.execute_reply.started":"2023-03-14T16:51:35.665320Z","shell.execute_reply":"2023-03-14T16:51:38.575513Z"},"trusted":true},"execution_count":null,"outputs":[]}]}