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


with Model() as model:
    number_events = 0
    background_rate = 0.01 #
    observation_time = 46.1/365. # 
    duty_cycle = 0.2
    horizon = 73.2    
    grb_rate = 1.0
    
    #trunc = pm.Bound(pm.DensityDist, lower=0)
    signal_rate = pm.DensityDist('signal_rate', 
                        logp=lambda value: log_signal_rate(value, number_events, background_rate, observation_time),
                       testval=10)
    
    n_galaxy = number_mweg(horizon)
    
    cbc_rate = pm.Deterministic('cbc_rate', signal_rate * n_galaxy * horizon)
    
    efficiency = pm.Uniform('efficiency', 0,1)
   
    costheta = pm.Deterministic('cos_angle', 1.0 - ((grb_rate/cbc_rate*efficiency)))
    
    angle = pm.Deterministic("angle", theano.tensor.arccos(costheta))
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 = 100000
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))
 
  0%|          | 0/100000 [00:00<?, ?it/s]
  0%|          | 1/100000 [00:00<24:33:01,  1.13it/s]
  1%|          | 551/100000 [00:00<17:05:32,  1.62it/s]
  1%|          | 1130/100000 [00:01<11:53:46,  2.31it/s]
  2%|▏         | 1625/100000 [00:01<8:17:14,  3.30it/s] 
  2%|▏         | 2207/100000 [00:01<5:46:05,  4.71it/s]
  3%|▎         | 2767/100000 [00:01<4:00:58,  6.73it/s]
  3%|▎         | 3335/100000 [00:01<2:47:46,  9.60it/s]
  4%|▍         | 3861/100000 [00:01<1:56:53, 13.71it/s]
  4%|▍         | 4464/100000 [00:01<1:21:23, 19.56it/s]
  5%|▌         | 5016/100000 [00:01<56:43, 27.90it/s]  
  6%|▌         | 5600/100000 [00:01<39:32, 39.78it/s]
  6%|▌         | 6181/100000 [00:01<27:35, 56.66it/s]
  7%|▋         | 6768/100000 [00:02<19:16, 80.61it/s]
  7%|▋         | 7336/100000 [00:02<13:29, 114.46it/s]

  8%|▊         | 7902/100000 [00:02<09:28, 162.04it/s]
  9%|▊         | 8501/100000 [00:02<06:39, 228.83it/s]
  9%|▉         | 9072/100000 [00:02<04:43, 321.28it/s]
 10%|▉         | 9640/100000 [00:02<03:22, 447.31it/s]
 10%|█         | 10250/100000 [00:02<02:24, 619.54it/s]
 11%|█         | 10822/100000 [00:02<01:45, 842.44it/s]
 11%|█▏        | 11380/100000 [00:02<01:18, 1129.53it/s]
 12%|█▏        | 11936/100000 [00:03<00:59, 1475.81it/s]
 12%|█▏        | 12481/100000 [00:03<00:46, 1884.82it/s]
 13%|█▎        | 13079/100000 [00:03<00:36, 2372.01it/s]
 14%|█▎        | 13638/100000 [00:03<00:30, 2853.66it/s]
 14%|█▍        | 14192/100000 [00:03<00:25, 3321.23it/s]
 15%|█▍        | 14783/100000 [00:03<00:22, 3823.69it/s]
 15%|█▌        | 15345/100000 [00:03<00:20, 4122.68it/s]
 16%|█▌        | 15945/100000 [00:03<00:18, 4548.58it/s]
 17%|█▋        | 16614/100000 [00:03<00:16, 5030.46it/s]
 17%|█▋        | 17299/100000 [00:03<00:15, 5465.66it/s]
 18%|█▊        | 17931/100000 [00:04<00:14, 5693.98it/s]
 19%|█▊        | 18562/100000 [00:04<00:13, 5863.66it/s]
 19%|█▉        | 19202/100000 [00:04<00:13, 6014.47it/s]
 20%|█▉        | 19832/100000 [00:04<00:13, 6001.61it/s]
 20%|██        | 20500/100000 [00:04<00:12, 6188.10it/s]
 21%|██        | 21135/100000 [00:04<00:12, 6153.07it/s]
 22%|██▏       | 21830/100000 [00:04<00:12, 6369.24it/s]
 23%|██▎       | 22536/100000 [00:04<00:11, 6561.33it/s]
 23%|██▎       | 23201/100000 [00:04<00:11, 6434.36it/s]
 24%|██▍       | 23861/100000 [00:04<00:11, 6481.55it/s]
 25%|██▍       | 24514/100000 [00:05<00:11, 6392.66it/s]
 25%|██▌       | 25157/100000 [00:05<00:11, 6343.02it/s]
 26%|██▌       | 25860/100000 [00:05<00:11, 6533.82it/s]
 27%|██▋       | 26559/100000 [00:05<00:11, 6663.97it/s]
 27%|██▋       | 27243/100000 [00:05<00:10, 6713.85it/s]
 28%|██▊       | 27917/100000 [00:05<00:11, 6334.56it/s]
 29%|██▊       | 28606/100000 [00:05<00:11, 6490.15it/s]
 29%|██▉       | 29261/100000 [00:05<00:10, 6472.39it/s]
 30%|██▉       | 29912/100000 [00:05<00:10, 6446.53it/s]
 31%|███       | 30570/100000 [00:05<00:10, 6484.94it/s]
 31%|███▏      | 31270/100000 [00:06<00:10, 6630.06it/s]
 32%|███▏      | 31936/100000 [00:06<00:10, 6399.41it/s]
 33%|███▎      | 32580/100000 [00:06<00:10, 6152.75it/s]
 33%|███▎      | 33210/100000 [00:06<00:10, 6194.68it/s]
 34%|███▍      | 33856/100000 [00:06<00:10, 6270.32it/s]
 34%|███▍      | 34486/100000 [00:06<00:10, 6174.22it/s]
 35%|███▌      | 35106/100000 [00:06<00:11, 5659.41it/s]
 36%|███▌      | 35682/100000 [00:06<00:11, 5518.93it/s]
 36%|███▌      | 36242/100000 [00:06<00:11, 5320.60it/s]
 37%|███▋      | 36806/100000 [00:07<00:11, 5411.50it/s]
 37%|███▋      | 37380/100000 [00:07<00:11, 5503.00it/s]
 38%|███▊      | 37935/100000 [00:07<00:11, 5485.30it/s]
 38%|███▊      | 38487/100000 [00:07<00:11, 5484.85it/s]
 39%|███▉      | 39084/100000 [00:07<00:10, 5619.16it/s]
 40%|███▉      | 39649/100000 [00:07<00:11, 5106.99it/s]
 40%|████      | 40171/100000 [00:07<00:14, 4238.75it/s]
 41%|████      | 40667/100000 [00:07<00:13, 4430.15it/s]
 41%|████▏     | 41251/100000 [00:07<00:12, 4774.62it/s]
 42%|████▏     | 41783/100000 [00:08<00:11, 4923.68it/s]
 42%|████▏     | 42322/100000 [00:08<00:11, 5053.49it/s]
 43%|████▎     | 42913/100000 [00:08<00:10, 5283.09it/s]
 43%|████▎     | 43454/100000 [00:08<00:10, 5206.99it/s]
 44%|████▍     | 44002/100000 [00:08<00:10, 5283.83it/s]
 45%|████▍     | 44537/100000 [00:08<00:11, 4934.84it/s]
 45%|████▌     | 45051/100000 [00:08<00:11, 4992.52it/s]
 46%|████▌     | 45632/100000 [00:08<00:10, 5210.72it/s]
 46%|████▌     | 46171/100000 [00:08<00:10, 5262.50it/s]
 47%|████▋     | 46721/100000 [00:09<00:09, 5330.91it/s]
 47%|████▋     | 47258/100000 [00:09<00:10, 5142.33it/s]
 48%|████▊     | 47809/100000 [00:09<00:09, 5246.32it/s]
 48%|████▊     | 48402/100000 [00:09<00:09, 5432.82it/s]
 49%|████▉     | 48950/100000 [00:09<00:09, 5434.46it/s]
 49%|████▉     | 49497/100000 [00:09<00:09, 5410.67it/s]
 50%|█████     | 50060/100000 [00:09<00:09, 5472.45it/s]
 51%|█████     | 50663/100000 [00:09<00:08, 5627.24it/s]
 51%|█████     | 51228/100000 [00:09<00:08, 5606.03it/s]
 52%|█████▏    | 51799/100000 [00:09<00:08, 5636.27it/s]
 52%|█████▏    | 52395/100000 [00:10<00:08, 5729.32it/s]
 53%|█████▎    | 52970/100000 [00:10<00:08, 5577.86it/s]
 54%|█████▎    | 53530/100000 [00:10<00:08, 5430.62it/s]
 54%|█████▍    | 54081/100000 [00:10<00:08, 5453.40it/s]
 55%|█████▍    | 54681/100000 [00:10<00:08, 5605.33it/s]
 55%|█████▌    | 55244/100000 [00:10<00:08, 5430.41it/s]
 56%|█████▌    | 55790/100000 [00:10<00:08, 5099.99it/s]
 56%|█████▋    | 56306/100000 [00:10<00:08, 4860.61it/s]
 57%|█████▋    | 56799/100000 [00:10<00:08, 4827.79it/s]
 57%|█████▋    | 57289/100000 [00:11<00:08, 4847.98it/s]
 58%|█████▊    | 57856/100000 [00:11<00:08, 5067.01it/s]
 58%|█████▊    | 58368/100000 [00:11<00:08, 4858.85it/s]
 59%|█████▉    | 58922/100000 [00:11<00:08, 5043.11it/s]
 60%|█████▉    | 59546/100000 [00:11<00:07, 5349.38it/s]
 60%|██████    | 60133/100000 [00:11<00:07, 5494.23it/s]
 61%|██████    | 60690/100000 [00:11<00:07, 5242.63it/s]
 61%|██████    | 61222/100000 [00:11<00:07, 5113.35it/s]
 62%|██████▏   | 61740/100000 [00:11<00:08, 4664.56it/s]
 62%|██████▏   | 62219/100000 [00:12<00:08, 4233.61it/s]
 63%|██████▎   | 62703/100000 [00:12<00:08, 4397.62it/s]
 63%|██████▎   | 63320/100000 [00:12<00:07, 4811.68it/s]
 64%|██████▍   | 63884/100000 [00:12<00:07, 5031.71it/s]
 64%|██████▍   | 64405/100000 [00:12<00:07, 4946.06it/s]
 65%|██████▍   | 64912/100000 [00:12<00:07, 4732.16it/s]
 65%|██████▌   | 65396/100000 [00:12<00:07, 4731.57it/s]
 66%|██████▌   | 65877/100000 [00:12<00:07, 4684.85it/s]
 66%|██████▋   | 66351/100000 [00:12<00:08, 4182.71it/s]
 67%|██████▋   | 66784/100000 [00:13<00:08, 4146.90it/s]
 67%|██████▋   | 67249/100000 [00:13<00:07, 4282.88it/s]
 68%|██████▊   | 67739/100000 [00:13<00:07, 4447.88it/s]
 68%|██████▊   | 68265/100000 [00:13<00:06, 4662.54it/s]
 69%|██████▉   | 68894/100000 [00:13<00:06, 5052.53it/s]
 69%|██████▉   | 69415/100000 [00:13<00:06, 4942.76it/s]
 70%|██████▉   | 69921/100000 [00:13<00:06, 4739.12it/s]
 70%|███████   | 70405/100000 [00:13<00:06, 4686.80it/s]
 71%|███████   | 70881/100000 [00:13<00:06, 4584.82it/s]
 71%|███████▏  | 71400/100000 [00:13<00:06, 4750.90it/s]
 72%|███████▏  | 71913/100000 [00:14<00:05, 4858.23it/s]
 72%|███████▏  | 72404/100000 [00:14<00:05, 4655.52it/s]
 73%|███████▎  | 72875/100000 [00:14<00:06, 4407.81it/s]
 73%|███████▎  | 73382/100000 [00:14<00:05, 4587.34it/s]
 74%|███████▍  | 73939/100000 [00:14<00:05, 4842.16it/s]
 75%|███████▍  | 74518/100000 [00:14<00:05, 5090.16it/s]
 75%|███████▌  | 75143/100000 [00:14<00:04, 5388.62it/s]
 76%|███████▌  | 75726/100000 [00:14<00:04, 5509.08it/s]
 76%|███████▋  | 76286/100000 [00:14<00:04, 5525.33it/s]
 77%|███████▋  | 76891/100000 [00:14<00:04, 5672.41it/s]
 77%|███████▋  | 77464/100000 [00:15<00:04, 5557.05it/s]
 78%|███████▊  | 78082/100000 [00:15<00:03, 5718.25it/s]
 79%|███████▊  | 78668/100000 [00:15<00:03, 5757.95it/s]
 79%|███████▉  | 79247/100000 [00:15<00:03, 5539.16it/s]
 80%|███████▉  | 79805/100000 [00:15<00:03, 5481.20it/s]
 80%|████████  | 80368/100000 [00:15<00:03, 5523.37it/s]
 81%|████████  | 80973/100000 [00:15<00:03, 5669.92it/s]
 82%|████████▏ | 81543/100000 [00:15<00:03, 5618.44it/s]
 82%|████████▏ | 82126/100000 [00:15<00:03, 5678.19it/s]
 83%|████████▎ | 82696/100000 [00:16<00:03, 4991.47it/s]
 83%|████████▎ | 83212/100000 [00:16<00:03, 4812.93it/s]
 84%|████████▎ | 83706/100000 [00:16<00:03, 4592.87it/s]
 84%|████████▍ | 84281/100000 [00:16<00:03, 4887.79it/s]
 85%|████████▍ | 84913/100000 [00:16<00:02, 5243.20it/s]
 85%|████████▌ | 85454/100000 [00:16<00:02, 5204.75it/s]
 86%|████████▌ | 86028/100000 [00:16<00:02, 5353.38it/s]
 87%|████████▋ | 86626/100000 [00:16<00:02, 5525.49it/s]
 87%|████████▋ | 87187/100000 [00:16<00:02, 5262.55it/s]
 88%|████████▊ | 87722/100000 [00:17<00:02, 5062.35it/s]
 88%|████████▊ | 88236/100000 [00:17<00:02, 4930.05it/s]
 89%|████████▊ | 88739/100000 [00:17<00:02, 4958.05it/s]
 89%|████████▉ | 89239/100000 [00:17<00:02, 4940.79it/s]
 90%|████████▉ | 89746/100000 [00:17<00:02, 4977.20it/s]
 90%|█████████ | 90246/100000 [00:17<00:01, 4941.81it/s]
 91%|█████████ | 90751/100000 [00:17<00:01, 4972.43it/s]
 91%|█████████▏| 91270/100000 [00:17<00:01, 5035.27it/s]
 92%|█████████▏| 91775/100000 [00:17<00:01, 5010.58it/s]
 92%|█████████▏| 92315/100000 [00:17<00:01, 5119.13it/s]
 93%|█████████▎| 92828/100000 [00:18<00:01, 5118.83it/s]
 93%|█████████▎| 93341/100000 [00:18<00:01, 5088.61it/s]
 94%|█████████▍| 93867/100000 [00:18<00:01, 5137.25it/s]
 94%|█████████▍| 94382/100000 [00:18<00:01, 5070.19it/s]
 95%|█████████▍| 94890/100000 [00:18<00:01, 4941.49it/s]
 95%|█████████▌| 95401/100000 [00:18<00:00, 4988.32it/s]
 96%|█████████▌| 95901/100000 [00:18<00:00, 4849.63it/s]
 96%|█████████▋| 96459/100000 [00:18<00:00, 5045.45it/s]
 97%|█████████▋| 96969/100000 [00:18<00:00, 5060.67it/s]
 98%|█████████▊| 97557/100000 [00:18<00:00, 5278.76it/s]
 98%|█████████▊| 98164/100000 [00:19<00:00, 5493.18it/s]
 99%|█████████▉| 98751/100000 [00:19<00:00, 5600.81it/s]
 99%|█████████▉| 99316/100000 [00:19<00:00, 5607.50it/s]
100%|█████████▉| 99880/100000 [00:19<00:00, 5338.90it/s]
100%|██████████| 100000/100000 [00:17<00:00, 5830.96it/s]
100%|██████████| 100000/100000 [00:08<00:00, 11366.43it/s]
100%|██████████| 100000/100000 [00:08<00:00, 11976.58it/s]
100%|██████████| 100000/100000 [00:15<00:00, 6282.39it/s]
100%|██████████| 100000/100000 [00:18<00:00, 5517.46it/s]
100%|██████████| 100000/100000 [00:09<00:00, 10935.50it/s]
100%|██████████| 100000/100000 [00:08<00:00, 12000.64it/s]

import matplotlib.pyplot as plt
%matplotlib inline
o1_trace.varnames
['signal_rate', 'cbc_rate', 'cos_angle', 'angle']
f, ax = plt.subplots(4,len(priors), figsize = (20, 14))
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("Efficiency ~ {}".format(priors[i]));
    ax[0,i].set_xlabel("$s$ /yr")
    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$ /Gpc³/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");
    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")
        ax[2,i].legend()
    except KeyError:
        pass

    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")
    ax[3,i].legend()
png
for i,o1_trace in enumerate(o1_traces):
    print priors[i], np.mean(o1_trace[2000:]['signal_rate'])
 uniform 7.92106015126
jeffreys 7.88566943524
1.0 7.86542454918
0.5 7.95348073587

for i,o1_trace in enumerate(o1_traces):
    print priors[i]
    print "10%", np.percentile(np.rad2deg(o1_trace[2000:]['angle'][np.isfinite(o1_trace[2000:]['angle'])]), 10)
    print "90%", np.percentile(np.rad2deg(o1_trace[2000:]['angle'][np.isfinite(o1_trace[2000:]['angle'])]), 90)
 uniform
10% 0.0745379907886
90% 0.52152504982
jeffreys
10% 0.0433182183062
90% 0.512836344577
1.0
10% 0.1589282988
90% 0.751545566894
0.5
10% 0.112591172589
90% 0.524950142617