def rmh_master_variance_hyperparams(comm, shape_prev, rate_prev, MPIROOT=0,
                                    prior_mean_log=2.65,
                                    prior_prec_log=1. / 0.652 ** 2,
                                    prior_shape=1., prior_rate=0.,
                                    brent_scale=6., fallback_upper=1e4,
                                    propDf=5.):
    '''
    Master side of Metropolis-Hastings step for variance hyperparameters given
    all other parameters.

    Have normal likelihood, so variance likelihood has the same form as gamma
    distribution. Using a log-normal prior for the shape hyperparameter and
    gamma prior for rate hyperparameter.

    Proposing from normal approximation to the conditional posterior
    (conditional independence chain). Parameters are log-transformed.

    Builds normal approximation based upon local data, then combines this with
    others on the master process. These approximations are used to generate a
    proposal, which is then broadcast back to the workers. The workers then
    evaluate the log-target ratio and combine these on the master to execute
    the MH step. The resulting draw is __not__ brought back to the workers until
    the next synchronization.

    Returns a 2-tuple consisting of the new (shape, rate) and a boolean
    indicating acceptance.
    '''
    # Build normal approximation to posterior of transformed hyperparameters.
    # Aggregating local results from workers.
    # This assumes that rmh_worker_nbinom_hyperparams() has been called on all
    # workers.
    T, buf = np.zeros((2, 3))
    comm.Reduce(buf, T, root=MPIROOT)
    n = T[2]

    shape_hat, rate_hat = map_estimator_gamma(x=None, T=T, log=True,
                                              prior_shape=prior_shape,
                                              prior_rate=prior_rate,
                                              prior_mean_log=prior_mean_log,
                                              prior_prec_log=prior_prec_log,
                                              brent_scale=brent_scale,
                                              fallback_upper=fallback_upper)
    theta_hat = np.log(np.array([shape_hat, rate_hat]))

    # Propose using a bivariate normal approximate to the joint conditional
    # posterior of (shape, rate)

    # Compute posterior information matrix for parameters
    info = info_posterior_gamma(shape=shape_hat, rate=rate_hat,
                                x=None, T=T, log=True,
                                prior_shape=prior_shape,
                                prior_rate=prior_rate,
                                prior_mean_log=prior_mean_log,
                                prior_prec_log=prior_prec_log)

    # Cholesky decompose information matrix for bivariate draw and
    # density calculations
    U = linalg.cholesky(info, lower=False)

    # Propose shape and rate parameter jointly
    z_prop = (np.random.randn(2) /
              np.sqrt(np.random.gamma(shape=propDf / 2., scale=2.,
                                      size=2) / propDf))
    theta_prop = theta_hat + linalg.solve_triangular(U, z_prop)
    shape_prop, rate_prop = np.exp(theta_prop)

    # Demean and decorrelate previous draws
    theta_prev = np.log(np.array([shape_prev, rate_prev]))
    z_prev = np.dot(U, theta_prev - theta_hat)

    # Compute log-ratio of target densities.
    log_target_ratio = \
            - n * (special.gammaln(shape_prop) - special.gammaln(shape_prev)) \
            + n * (shape_prop * np.log(rate_prop) - shape_prev *
                   np.log(rate_prev)) \
            + (shape_prop - shape_prev) * T[1] \
            - (rate_prop - rate_prev) * T[0]

    # Add log-prior ratio
    if prior_prec_log > 0:
        # Add the log-normal prior on the shape parameter
        log_target_ratio += (dlnorm(shape_prop, mu=prior_mean_log,
                                    sigmasq=1. / prior_prec_log, log=True) -
                             dlnorm(shape_prev, mu=prior_mean_log,
                                    sigmasq=1. / prior_prec_log, log=True))
    # Add the gamma prior on the rate parameter
    if prior_rate > 0:
        log_target_ratio += (dgamma(rate_prop, shape=prior_shape,
                                    rate=prior_rate, log=True) -
                             dgamma(rate_prev, shape=prior_shape,
                                    rate=prior_rate, log=True))
    else:
        log_target_ratio += np.log(rate_prop / rate_prev) * (shape_prop - 1.)

    # Compute log-ratio of proposal densities

    # These are transformed bivariate t's with equivalent covariance
    # matrices, so the resulting Jacobian terms cancel. We are left to
    # contend with the z's and the Jacobian terms resulting from
    # exponentiation.
    log_prop_ratio = -np.sum(np.log(1. + z_prop ** 2 / propDf) -
                             np.log(1. + z_prev ** 2 / propDf))
    log_prop_ratio *= (propDf + 1.) / 2.
    log_prop_ratio += -np.sum(theta_prop - theta_prev)

    # Execute MH update
    return mh_update(
        prop=(shape_prop, rate_prop), prev=(shape_prev, rate_prev),
        log_target_ratio=log_target_ratio,
        log_prop_ratio=log_prop_ratio)
def rmh_master_glm_coef(comm, b_prev, MPIROOT=0, propDf=5., method='newton',
                        cov=emulate.cov_sqexp, n_iter_refine=2,
                        final_info_refine=1, prior_log_density=None,
                        prior_args=tuple(), prior_kwargs={}):
    '''
    Master component of single Metropolis-Hastings step for GLM coefficients
    using a normal approximation to their posterior distribution. Proposes
    linearly-transformed vector of independent t_propDf random variables.

    Builds normal approximation based upon local data, then combines this with
    others on the master process. These approximations are used to generate a
    proposal, which is then broadcast back to the workers. The workers then
    evaluate the log-target ratio and combine these on the master to execute
    the MH step. The resulting draw is __not__ brought back to the workers until
    the next synchronization.

    Returns a 2-tuple consisting of the resulting coefficients and a boolean
    indicating acceptance.
    '''
    # Compute dimensions
    p = np.size(b_prev)

    if method == 'emulate':
        # Gather emulators from workers
        emulator = emulate.aggregate_emulators_mpi(
            comm=comm, emulator=None, MPIROOT=MPIROOT,
            info=lambda e: linalg.cho_solve((e['slope_mean'], True),
                                            np.eye(e['slope_mean'].shape[0])))

        # Find root of combined approximate score function
        b_hat = emulator['center']
        b_hat = optimize.fsolve(
            func=emulate.evaluate_emulator, x0=b_hat, args=(emulator, cov))

        # Compute Cholesky decomposition of approximate combined information
        # for proposal
        U = linalg.cholesky(emulator['info'], lower=False)
    else:
        # Build normal approximation to posterior of transformed hyperparameters.
        # Aggregating local results from workers.
        # This assumes that rmh_worker_nbinom_glm_coef() has been called on all
        # workers.
        b_hat, prec = posterior_approx_distributed(
            comm=comm, dim_param=p, MPIROOT=MPIROOT)

        # Refine approximation with single Newton-Raphson step
        b_hat, prec = refine_distributed_approx(
            comm=comm, est=b_hat, prec=prec, dim_param=p, n_iter=n_iter_refine,
            final_info=final_info_refine, MPIROOT=MPIROOT)

        # Cholesky decompose precision matrix for draws and density calculations
        U = linalg.cholesky(prec, lower=False)

    # Propose from linearly-transformed t with appropriate mean and covariance
    z_prop = (np.random.randn(p) /
              np.sqrt(np.random.gamma(shape=propDf / 2., scale=2., size=p) /
                      propDf))
    b_prop = b_hat + linalg.solve_triangular(U, z_prop, lower=False)

    # Demean and decorrelate previous draw of b
    z_prev = np.dot(U, b_prev - b_hat)

    # Broadcast b_prop to workers
    comm.Bcast([b_prop, MPI.DOUBLE], root=MPIROOT)

    # Compute log-ratio of target densities.
    # Start by obtaining likelihood component from workers.
    log_target_ratio = np.array(0.)
    buf = np.array(0.)
    comm.Reduce([buf, MPI.DOUBLE], [log_target_ratio, MPI.DOUBLE],
                op=MPI.SUM, root=MPIROOT)
    if prior_log_density is not None:
        log_target_ratio += (
            prior_log_density(b_prop, *prior_args, **prior_kwargs) -
            prior_log_density(b_prev, *prior_args, **prior_kwargs))

    # Compute log-ratio of proposal densities. This is very easy with the
    # demeaned and decorrelated values z.
    log_prop_ratio = -(propDf + 1.) / 2. * \
            np.sum(np.log(1. + z_prop ** 2 / propDf) - \
                   np.log(1. + z_prev ** 2 / propDf))

    return mh_update(prop=b_prop, prev=b_prev,
                     log_target_ratio=log_target_ratio,
                     log_prop_ratio=log_prop_ratio)
def rmh_master_nbinom_hyperparams(comm, r_prev, p_prev, MPIROOT=0,
                                  prior_mean_log=2.65,
                                  prior_prec_log=1. / 0.652 ** 2,
                                  prior_a=1., prior_b=1., propDf=5.,
                                  method='newton', cov=emulate.cov_sqexp,
                                  n_iter_refine=10, final_info_refine=1):
    '''
    Master side of Metropolis-Hastings step for negative-binomial
    hyperparameters given all other parameters.

    Using a log-normal prior for the r (convolution) hyperparameter and a
    conditionally-conjugate beta prior for p.

    Proposing from normal approximation to the conditional posterior
    (conditional independence chain). Parameters are log- and logit-transformed.

    Builds normal approximation based upon local data, then combines this with
    others on the master process. These approximations are used to generate a
    proposal, which is then broadcast back to the workers. The workers then
    evaluate the log-target ratio and combine these on the master to execute
    the MH step. The resulting draw is __not__ brought back to the workers until
    the next synchronization.

    Returns a 2-tuple consisting of the new (shape, rate) and a boolean
    indicating acceptance.
    '''
    if method=='emulate':
        # Gather emulators from workers
        emulator = emulate.aggregate_emulators_mpi(
            comm=comm, emulator=None, MPIROOT=MPIROOT,
            info=lambda e: linalg.cho_solve((e['slope_mean'], True),
                                            np.eye(e['slope_mean'].shape[0])))

        # Find root of combined approximate score function
        theta_hat = emulator['center']
        theta_hat = optimize.fsolve(
            func=emulate.evaluate_emulator, x0=theta_hat, args=(emulator, cov))

        # Compute Cholesky decomposition of approximate combined information
        # for proposal
        U = linalg.cholesky(emulator['info'], lower=False)
    elif method == 'newton':
        # Build normal approximation to posterior of transformed hyperparameters.
        # Aggregating local results from workers.
        # This assumes that rmh_worker_nbinom_hyperparameters() has been called
        # on all workers.
        theta_hat, prec = posterior_approx_distributed(
            comm=comm, dim_param=2, MPIROOT=MPIROOT)

        # Refine approximation with single Newton-Raphson step
        theta_hat, prec = refine_distributed_approx(
            comm=comm, est=theta_hat, prec=prec, dim_param=2,
            n_iter=n_iter_refine, final_info=final_info_refine, MPIROOT=MPIROOT)

        # Cholesky decompose precision matrix for draws and density calculations
        U = linalg.cholesky(prec, lower=False)
    else:
        print >> sys.stderr, "Error - method %s unknown" % method
        return

    # Propose r and p jointly
    z_prop = (np.random.randn(2) /
              np.sqrt(np.random.gamma(shape=propDf / 2., scale=2.,
                                      size=2) / propDf))
    theta_prop = theta_hat + linalg.solve_triangular(U, z_prop)
    r_prop, p_prop = np.exp(theta_prop)
    p_prop = p_prop / (1. + p_prop)

    # Demean and decorrelate previous draws
    theta_prev = np.log(np.array([r_prev, p_prev]))
    theta_prev[1] -= np.log(1. - p_prev)
    z_prev = np.dot(U, theta_prev - theta_hat)

    # Broadcast theta_prop to workers
    comm.Bcast([theta_prop, MPI.DOUBLE], root=MPIROOT)

    # Compute log-ratio of target densities.
    # Start by obtaining likelihood component from workers.
    log_target_ratio = np.array(0.)
    buf = np.array(0.)
    comm.Reduce([buf, MPI.DOUBLE], [log_target_ratio, MPI.DOUBLE],
                op=MPI.SUM, root=MPIROOT)

    if prior_prec_log > 0:
        # Add the log-normal prior on r
        log_target_ratio += (dlnorm(r_prop, mu=prior_mean_log,
                                    sigmasq=1. / prior_prec_log, log=True) -
                             dlnorm(r_prev, mu=prior_mean_log,
                                    sigmasq=1. / prior_prec_log, log=True))
    # Add the beta prior on p
    log_target_ratio += (dbeta(p_prop, a=prior_a, b=prior_b, log=True) -
                         dbeta(p_prev, a=prior_a, b=prior_b, log=True))

    # Compute log-ratio of proposal densities

    # These are transformed bivariate t's with equivalent covariance
    # matrices, so the resulting Jacobian terms cancel. We are left to
    # contend with the z's and the Jacobian terms resulting from the
    # exponential and logit transformations.
    log_prop_ratio = -np.sum(np.log(1. + z_prop ** 2 / propDf) -
                             np.log(1. + z_prev ** 2 / propDf))
    log_prop_ratio *= (propDf + 1.) / 2.
    log_prop_ratio += -(np.log(r_prop) - np.log(r_prev))
    log_prop_ratio += -(np.log(p_prop) + np.log(1. - p_prop)
                        - np.log(p_prev) - np.log(1. - p_prev))

    # Execute MH update
    return mh_update(prop=(r_prop, p_prop), prev=(r_prev, p_prev),
                     log_target_ratio=log_target_ratio,
                     log_prop_ratio=log_prop_ratio)