In [53]:
%pylab inline
import pandas as pd
from scipy import stats
from utilHelpers import lineHist
from functools import partial
from scipy import integrate


Populating the interactive namespace from numpy and matplotlib

In [50]:
def rejectionSampling(envPDF, envRV, targetPDF, n):
    """Generates n random variates from the target distribution
    using acceptance/rejection sampling with the envelope distribution, envPDF.
    
    Parameters
    ----------
    envPDF : func
        Takes an array of x and returns their probabilities
        times a constant, such that envPDF(x) > targetPDF(x), always.
    envRV : func
        Takes single arg n and returns n random variates from envPDF.
    targetPDF : func
        Takes array of x and returns their probabilities
    n : int
        Number of total samples to draw from targetPDF
    
    Returns
    -------
    arr : array
        Array of samples distributed proportionately to target PDF
    acceptanceProbability : float
        Estimate for diagnosing the envelope distribution."""
    
    i = 0
    arr = zeros(n)
    acceptanceProbability = 1
    allDraws = None
    while i < n:
        """Draw as many samples as we expect to need,
        based on the acceptance probability (initially 1)"""
        neededN = n-i
        drawingN = int(ceil(neededN * (1/acceptanceProbability)))
        
        samples = envRV(drawingN)
        envPr = envPDF(samples)
        targetPr = targetPDF(samples)
        unifRVs = rand(drawingN)
        
        """Keep the samples that meet the acceptance criteria"""
        keepInds = unifRVs < (targetPr/envPr)
        actualN = keepInds.sum()
        if actualN < neededN:
            """Fill arr with all the samples that were accepted"""
            arr[i:(i+actualN)] = samples[keepInds]
        else:
            """Fill the rest of the arr with as many samples as needed"""
            arr[i:] = samples[keepInds][:neededN]
        
        """Keep track of all draws to compute acceptanceProbability"""
        if allDraws is None:
            allDraws = targetPr/envPr
        else:
            allDraws = concatenate((allDraws,targetPr/envPr))
        
        """Estimate the acceptance probability"""
        acceptanceProbability = allDraws.mean()
        
        i+=actualN
    
    return arr,acceptanceProbability

In [85]:
"""Define parameters"""
lam0 = 0.008 #baseline lambda
RRmax = 4 #max RR prior to decay
RRdecay = 1 #exp decay with tau = 1 years
n = 100000

In [86]:
"""Functions for the exponential distribution
(used as "envelope" distribution)"""
expPDF = lambda lam,t: lam*exp(-lam*t)
expCDF = lambda lam,t: 1 - exp(-lam*t)
expQ = lambda lam,pr: -log(1-pr)/lam

"""Function for a time-varying (tv) (exp decay) lambda parameter"""
RRCORFunc = lambda RR,decay,t: (RR-1)*exp(-t/decay) + 1

"""Non-normalized time-varying density function
For AR sampling to work with this function,
    (1) RRmax >= 1
    (2) RRdecay << 1/lam0"""
tvDF = lambda lam0,RR,decay,t: expPDF(lam0*RRCORFunc(RR,decay,t),t)

ic,err = integrate.quad(partial(tvDF,lam0,RRmax,RRdecay),0,inf)
print "Normalizing constant (upper bound on error): %1.3f (%1.3g)" % (ic,err)

"""Normalized PDF with a time-varying lambda parameter"""
tvPDF = lambda t: partial(tvDF,lam0,RRmax,RRdecay)(t)/ic

"""Exponential RV function with frozen parameter lam"""
expRV = lambda n: -log(rand(n)) / lam0

"""Scaled envelope function such that envPDF >= targetPDF"""
envPDF = lambda t: (tvPDF(0)/expPDF(lam0,0)) * expPDF(lam0,t)

expT = expRV(n)
targetT,acc = rejectionSampling(envPDF = envPDF,
                                envRV = expRV,
                                targetPDF = tvPDF,
                                n = n)

print "Acceptance probability: %1.3f" % acc

t = linspace(0,40/lam0,50000)
xlbase = array([0,4/lam0])

for xl in [xlbase, xlbase/10]:
    figure(figsize=(17,17))
    subplot(2,1,1)
    plot(t, lam0 * ones(t.shape),lw=2,color='k',label='Constant lambda')
    plot(t, lam0 * RRCORFunc(RRmax,RRdecay,t),lw=2,color='r',label='Time-varying lambda')
    ylabel('lambda')
    xlim(xl)
    ylim((0,lam0 * RRCORFunc(RRmax,RRdecay,0)))
    legend(loc=0)

    subplot(2,1,2)
    lineHist(expT, color='k', alpha = 0.3, edges=t, density=True, filled=True)
    lineHist(targetT, color='r', alpha = 0.3, edges=t, density=True, filled=True)
    plot(t, expPDF(lam0,t), color='k',lw=2,label='Constant lambda')
    plot(t, envPDF(t), color='g',lw=2,label='Envelope function')
    plot(t, tvPDF(t),color='r',lw=2,label='Time-varying lambda')
    xlim(xl)
    ylabel('PDF')
    legend(loc=0)
    xlabel('Time (years)')


Normalizing constant (upper bound on error): 1.023 (3.44e-11)
Acceptance probability: 0.256