jdmaturen / shifted_beta_geometric_py Goto Github PK

def survivor(probabilities, t):
    """Survivor function S"""
    s = 1 - probabilities[0]
    for x in xrange(1, t + 1):
        s = s - probabilities[x]
    return s

def log_likelihood(alpha, beta, data, survivors=None):
    if alpha <= 0 or beta <= 0:
        return -1000
    if survivors is None:
        survivors = cancel_rates
    probabilities = generate_probabilities(alpha, beta, len(data))
    final_survivor_likelihood = survivor(probabilities, len(data) - 1)
    
    return sum([s * np.log(probabilities[t]) for t, s in enumerate(data)]) + survivors[-1] * np.log(final_survivor_likelihood)

def generate_probabilities(alpha, beta, x):
    """Generate probabilities in one pass for all t in x"""
    p = [alpha / (alpha + beta)]
    for t in xrange(1, x):
        pt = (beta + t - 1) / (alpha + beta + t) * p[t-1]
        p.append(pt)
    return p