import pymc3 as pm
import numpy as np
import theano

from theano.compile.ops import as_op
def background_rate_f(b, T, n):
    """
    
    """
    out = 0
    #n = int(n)
    for i in range(n+1):
        out += ((b*T)**i * np.exp(- b*T)) / np.math.factorial(i)
    return out

def log_background_rate(b, T, n):
    return np.log(background_rate_f(b, T, n))

def signal_rate_part(s, n, b, T):
    top_a = T * ((s + b) * T)**n 
    top_b = np.exp(-(s + b)*T)
    p = (top_a * top_b) / np.math.factorial(n)
    return theano.tensor.switch(theano.tensor.le(s, 0), 0, p)

#@as_op(itypes=[T.dscalar, T.dscalar, T.dscalar, T.dscalar], otypes=[T.dscalar])
def log_signal_rate(s,n,b,T):
    #if theano.tensor.lt(0, s): return np.array([[0.0]])
    p = -log_background_rate(b,T,n) + np.log(signal_rate_part(s,n,b,T))
    
    return p
def number_mweg(horizon):
    return 4./3 * np.pi * horizon**3 *(2.26)**-3* (0.0116) #* horizon**3
import theano.tensor as T
from pymc3 import DensityDist, Uniform, Normal
from pymc3 import Model
from pymc3 import distributions

def grb_model(number_events, background_rate, 
              observation_time, horizon, grb_rate,
             efficiency_prior = "uniform"):
    with Model() as model:
        signal_rate = pm.DensityDist('signal_rate', 
                            logp=lambda value: log_signal_rate(value, number_events, background_rate, observation_time),
                           testval=50)

        n_galaxy = number_mweg(horizon)

        cbc_rate = pm.Deterministic('cbc_rate', signal_rate * n_galaxy)

        # Allow the efficiency prior to be switched-out
        if efficiency_prior == "uniform":
            efficiency = pm.Uniform('efficiency', 0,1)
        elif efficiency_prior == "jeffreys":
            efficiency = pm.Beta('efficiency', 0.5, 0.5)
        elif isinstance(efficiency_prior, float):
            efficiency = efficiency_prior

        costheta = pm.Deterministic('cos_angle', 1.0 - ((grb_rate/(cbc_rate*efficiency))))

        angle = pm.Deterministic("angle", theano.tensor.arccos(costheta))
        
        return model
priors = ["uniform", "jeffreys", 1.0, 0.5]
# O1 Scenarios
number_events = 0 # There were no BNS detections in O1
background_rate = 0.01 # We take the FAR to be 1/100 yr
observation_time = 46.1/365. # The number of days of analysis conducted by gstLAL
horizon = 73.2  # The O1 BNS horizon distance in O1 BNS paper
grb_rate = 10.0
o1_models = []
for prior in priors:
    o1_models.append( grb_model(number_events, background_rate, observation_time, horizon, grb_rate, prior))
# O2 Scenarios
number_events = 1 # Assume O2 will see one BNS observation
background_rate = 0.01 # We take the FAR to be 1/100 yr
observation_time = 0.5 # The number of days of analysis conducted by gstLAL
duty_cycle = 0.4
observation_time *= duty_cycle
horizon = 80.0  # The O1 BNS horizon distance in O1 BNS paper
grb_rate = 10.0
o2_models = []
for prior in priors:
    o2_models.append( grb_model(number_events, background_rate, observation_time, horizon, grb_rate, prior))
samples = 1000000
o1_traces = []
for model in o1_models:
    with model:
        step = pm.Metropolis()
        o1_traces.append(pm.sample(samples, step))
    
o2_traces = []    
for model in o2_models:
    with model:
        step = pm.Metropolis()
        o2_traces.append(pm.sample(samples, step))
 100%|██████████| 1000000/1000000 [02:23<00:00, 6975.32it/s]
100%|██████████| 1000000/1000000 [02:22<00:00, 7003.99it/s]
100%|██████████| 1000000/1000000 [01:17<00:00, 12827.91it/s]
100%|██████████| 1000000/1000000 [01:13<00:00, 13516.73it/s]
100%|██████████| 1000000/1000000 [02:31<00:00, 6622.06it/s]
100%|██████████| 1000000/1000000 [02:30<00:00, 6666.50it/s]
100%|██████████| 1000000/1000000 [01:13<00:00, 13671.24it/s]
100%|██████████| 1000000/1000000 [01:12<00:00, 13714.12it/s]

import matplotlib.pyplot as plt
%matplotlib inline
plt.style.use("/home/daniel/papers/thesis/thesis-style.mpl")
priors = ["U(0,1)", "jeffreys", "$\delta(1)$", "$\delta(0.5)$"]
f, ax = plt.subplots(4,len(priors), figsize = (15, 15))
for i in range(len(priors)):

    o1_trace = o1_traces[i]
    o2_trace = o2_traces[i]

    ax[0,i].hist(o1_trace[2000:]['signal_rate'], histtype="step", label="O1 Signal Rate",
                 lw=2, normed=True, bins=100);
    ax[0,i].hist(o2_trace[2000:]['signal_rate'], histtype="step", label="O2 Signal Rate",
                 lw=2, normed=True, bins=100);
    #ax[0,0].vlines(np.percentile(trace['signal_rate'][2000:], 90), 0, 0.085,
    #              label = "90% lower limit");
    ax[0,i].set_title("$\epsilon \sim$ {}".format(priors[i]));
    ax[0,i].set_xlabel("$s$ /yr")
    if i==3:ax[0,i].legend()

    ax[1,i].hist(o1_trace[2000:]['cbc_rate'], histtype="step", lw=2, bins=100,
                 normed=True, label="O1 CBC Rate");
    ax[1,i].hist(o2_trace[2000:]['cbc_rate'], histtype="step", lw=2, bins=100,
                 normed=True, label="O2 CBC Rate");
    ax[1,i].set_xlabel(u"$s \mathrm{/Gpc}^3\mathrm{/yr}$");
    #ax[1,0].vlines(np.percentile(o1_trace['cbc_rate'][2000:], 90), 0, 5e-5,
    #              label = "90% lower limit, O1");
    #ax[1,0].vlines(np.percentile(o2_trace['cbc_rate'][2000:], 90), 0, 5e-5,
    #              label = "90% lower limit, O2", color="orange");
    if i==3:ax[1,i].legend()
    try:
        ax[2,i].hist(o1_trace[2000:]['efficiency'], histtype="step", lw=2, label="O1 Efficiency");
        ax[2,i].hist(o2_trace[2000:]['efficiency'], histtype="step", lw=2, label="O2 Efficiency");
        #ax[1,0].set_xlabel(u"$s$ /Gpc³/yr")
        if i==0: ax[2,i].legend()
        ax[2,i].set_ylim(0,2.5e5)
    except KeyError:
            ax[2,i].spines['top'].set_visible(False)
            ax[2,i].spines['bottom'].set_visible(False)
            ax[2,i].spines['left'].set_visible(False)
            ax[2,i].spines['right'].set_visible(False)
            ax[2,i].grid(False)
            ax[2,i].set_yticklabels([])
            ax[2,i].set_xticklabels([])
            ax[2,i].yaxis.set_ticks_position('none')
            ax[2,i].xaxis.set_ticks_position('none')

    ax[3,i].hist(np.rad2deg(o1_trace[2000:]['angle'][np.isfinite(o1_trace[2000:]['angle'])]),
               histtype="step", lw=2, label="O1 Angle", log=True);
    ax[3,i].hist(np.rad2deg(o2_trace[2000:]['angle'][np.isfinite(o2_trace[2000:]['angle'])]),
               histtype="step", lw=2, label="O2 Angle", log=True);
    #ax[3,i].vlines(np.percentile(np.rad2deg(o1_trace[2000:]['angle'][np.isfinite(o1_trace[2000:]['angle'])]), 90), 0, 3e5,
    #              label = "90\% lower limit, O1");
    #ax[3,i].vlines(np.percentile(np.rad2deg(o2_trace[2000:]['angle'][np.isfinite(o2_trace[2000:]['angle'])]), 90), 0, 3e5,
    #              label = "90\% lower limit, O2");
    #ax[1,0].set_xlabel(u"$s$ /Gpc³/yr")
    if i==3: ax[3,i].legend()
    if i > 0: 
        ax[3,i].set_yticklabels([]), ax[2,i].set_yticklabels([])
        ax[0,i].set_yticklabels([]), ax[1,i].set_yticklabels([])
    ax[3,i].set_xlabel("Beaming angle [deg]")    
f.subplots_adjust(0.07, 0.05, .99, .95, wspace=0.05)
#f.savefig("/home/daniel/papers/thesis/figures/grbbeamingpymc.pdf")
png

O1 Trace Summaries

from pymc3 import summary

for i, trace in enumerate(o1_traces):
    print "# {}".format(priors[i])
    print summary(trace, varnames=['angle'])
 # U(0,1)

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.029, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.026          0.048          0.076          0.135          0.618

None
# jeffreys

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.025, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.025          0.048          0.082          0.174          2.233

None
# $\delta(1)$

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.023, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.020          0.033          0.047          0.073          0.246

None
# $\delta(0.5)$

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.032, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.029          0.047          0.066          0.103          0.351

None

O2 Trace Summaries

from pymc3 import summary

for i, trace in enumerate(o2_traces):
    print "# {}".format(priors[i])
    summary(trace, varnames=['angle'])
 # U(0,1)

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.025, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.022          0.037          0.052          0.080          0.271

# jeffreys

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.023, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.022          0.036          0.055          0.103          1.025

# $\delta(1)$

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  nan              nan              nan              [0.020, nan]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.018          0.026          0.033          0.044          0.088

# $\delta(0.5)$

angle:

  Mean             SD               MC Error         95% HPD interval
  -------------------------------------------------------------------
  
  0.054            0.029            0.000            [0.021, 0.105]

  Posterior quantiles:
  2.5            25             50             75             97.5
  |--------------|==============|==============|--------------|
  
  0.026          0.037          0.047          0.062          0.125