class Array:
    def __init__(self, M, N):
        self.M = M
        self.N = N

        self.rnames = np.array(['m_%d' % m for m in range(self.M)])
        self.cnames = np.array(['n_%d' % n for n in range(self.N)])
        
        self.array = sparse.lil_matrix((self.M, self.N), dtype = np.bool)

        # Offset initially set to None so simpler offset-free model
        # code can be used by default
        self.offset = None

        self.row_covariates = {}
        self.col_covariates = {}
        self.edge_covariates = {}

    def __setitem__(self, index, x):
        self.array.__setitem__(index, x)

    def new_row_covariate(self, name, dtype = np.float64):
        self.row_covariates[name] = NodeCovariate(self.rnames, dtype)
        return self.row_covariates[name]

    def new_col_covariate(self, name, dtype = np.float64):
        self.col_covariates[name] = NodeCovariate(self.cnames, dtype)
        return self.col_covariates[name]

    def new_edge_covariate(self, name):
        self.edge_covariates[name] = EdgeCovariate(self.rnames, self.cnames)
        return self.edge_covariates[name]

    def initialize_offset(self):
        self.offset = EdgeCovariate(self.rnames, self.cnames)
        return self.offset

    def subarray(self, rinds = None, cinds = None):
        if rinds is None:
            rinds = np.arange(self.M)
        if cinds is None:
            cinds = np.arange(self.N)

        sub_M = len(rinds)
        sub_N = len(cinds)
        sub = Array(sub_M, sub_N)
        sub.rnames = self.rnames[rinds]
        sub.cnames = self.cnames[cinds]

        sub.array = self.array[rinds][:,cinds]

        for row_covariate in self.row_covariates:
            src = self.row_covariates[row_covariate]
            sub.row_covariates[row_covariate] = src.subset(rinds)
        for col_covariate in self.col_covariates:
            src = self.col_covariates[col_covariate]
            sub.col_covariates[col_covariate] = src.subset(cinds)
        for edge_covariate in self.edge_covariates:
            src = self.edge_covariates[edge_covariate]
            sub.edge_covariates[edge_covariate] = src.subset(rinds, cinds)

        if self.offset:
            sub.offset = self.offset.subset(rinds, cinds)

        return sub

    # Syntactic sugar to make array generation look like object mutation
    def generate(self, model, **opts):
        self.array = model.generate(self, **opts)

    def is_sparse(self):
        return sparse.issparse(self.array)
        
    def as_dense(self):
        if self.is_sparse():
            return np.asarray(self.array.todense())
        else:
            return self.array

    def offset_extremes(self):
        if self.offset is None:
            self.initialize_offset()

        # (Separately) sort rows and columns of adjacency matrix by
        # increasing sum
        A = self.as_dense()
        r_ord = np.argsort(A.sum(1))
        c_ord = np.argsort(A.sum(0))
        A = A[r_ord][:,c_ord]

        # Convenience function to set blocks of the offset
        def set_offset_block(r, c, val):
            # Replacing this with a vectorized operation will require
            # a conversion to CSR. It's not clear if this is worth the
            # overhead.
            for i in r_ord[r]:
                for j in c_ord[c]:
                    self.offset[i,j] = val

        # Recursively examine for submatrices that will send
        # corresponding EMLE parameter estimates to infinity
        to_screen = [(np.arange(self.M), np.arange(self.N))]
        while len(to_screen) > 0:
            r_act, c_act = to_screen.pop()

            if len(r_act) == 0 or len(c_act) == 0:
                continue

            A_act = A[r_act][:,c_act]
            n_act = A_act.shape

            # Cumulative sums allow for efficient bad substructure search
            C_inc = np.cumsum(np.cumsum(A_act,1),0)
            def reverse(x):
                return x[range(n_act[0]-1,-1,-1)][:,range(n_act[1]-1,-1,-1)]
            C_dec = reverse(np.cumsum(np.cumsum(reverse(1-A_act),1),0))

            if C_inc[-1,-1] == 0:
                set_offset_block(r_act, c_act, -np.inf)
                continue
            if C_dec[0,0] == 0:
                set_offset_block(r_act, c_act, np.inf)
                continue

            align = np.zeros((n_act[0]+1,n_act[1]+1))
            align[1:(n_act[0]+1)][:,1:(n_act[1]+1)] += (C_inc == 0)
            align[0:n_act[0]][:,0:n_act[1]] *= (C_dec == 0)
            splits = zip(*np.where(align))
            if len(splits) == 0:
                continue
            i_split, j_split = splits[0]

            set_offset_block(r_act[:i_split], c_act[:j_split], -np.inf)
            set_offset_block(r_act[i_split:], c_act[j_split:], np.inf)

            to_screen.append((r_act[i_split:], c_act[:j_split]))
            to_screen.append((r_act[:i_split], c_act[j_split:])) def array_from_data(z, xs = []):
    """Populate an Array with data z and covariates xs."""
    M, N = z.shape
    arr = Array(M, N)
    arr.array[:] = z

    for i, x in enumerate(xs):
        arr.new_edge_covariate('x_%d' % i)[:] = x

    return arr def approximate_conditional_nll(A, w, sort_by_wopt_var = True):
    """Return approximate row/column-conditional NLL of binary matrix.
    
Return the approximate nll of an observed binary matrix given
specified Bernoulli weights, conditioned on having the observed
margins.

Inputs:
  A: observed data, (m x n) binary matrix
  w: weight matrix, (m x n) matrix with values in (0, +infty)
Output:
  ncll: negative conditional log-likelihood
"""
    assert(A.shape == w.shape)

    r = A.sum(1, dtype=np.int)
    c = A.sum(0, dtype=np.int)

    r, c, arrays, _ = _prune(r, c, A, w)
    A, w = arrays

    # Sizing
    m, n = len(r), len(c)
    if (m == 0) or (n == 0):
        return 0.0

    # Sort the row margins (descending)
    rndx = np.argsort(-r)
    rsort = r[rndx]

    # Balance the weights
    a_scale, b_scale = canonical_scalings(w, r, c)
    wopt = apply_scale(w, a_scale, b_scale)
    if np.isnan(wopt).any():
        wopt = w

    # Reorder the columns
    if sort_by_wopt_var:
        cndx = np.lexsort((-wopt.var(0), c))
    else:
        cndx = np.argsort(c)
    csort = c[cndx]
    wopt = wopt[:,cndx]

    # Compute G
    G = _compute_G(r, m, n, wopt)

    return _compute_cnll(A, r, rsort, rndx, csort, cndx, m, n, G) def approximate_from_margins_weights(r, c, w, T = None,
                                     sort_by_wopt_var = True):
    """Return approximate samples from row/column-conditional binary matrices.
                                     
Return a binary matrix (or a list of binary matrices) sampled
approximately according to the specified Bernoulli weights,
conditioned on having the specified margins.
Inputs:
  r: row margins, length m
  c: column margins, length n
  w: weight matrix, (m x n) matrix with values in (0, +infty)
  T: number of matrices to sample
  sort_by_wopt_var: when enabled, column ordering depends on w
Output:
  B_sample_sparse: (T default) sparse representation of (m x n) binary matrix
                   (T >= 1) list of (sparse binary matrices, logQ, logP)

More explicitly, consider independent Bernoulli random variables
B(i,j) arranged as an m x n matrix B given the m-vector of row sums
r and the n-vector of column sums c of the sample, i.e., given that
sum(B_sample, 1) = r and sum(B_sample, 0) = c.

An error is generated if no binary matrix agrees with r and c.

B(i,j) is Bernoulli(p(i,j)) where p(i,j) = w(i,j)/(1+w(i,j)), i.e.,
w(i,j) = p(i,j)/(1-p(i,j)).  [The case p(i,j) = 1 must be handled by
the user in a preprocessing step, by converting to p(i,j) = 0 and
decrementing the row and column sums appropriately.]

The sparse representation used for output is a matrix giving the
locations of the ones in the sample. If d = sum(r) = sum(c), then
B_sample_sparse has dimensions (d x 2). If something goes wrong (due
to undetected improper input), some of the rows of B_sample_sparse
may [-1,-1], indicating no entry of B_sample.

B_sample can be recovered from B_sample_sparse via:

    B_sample = np.zeros((m,n), dtype=np.bool)
    for i, j in B_sample_sparse:
        if i == -1: break 
        B_sample[i,j] = 1
"""
    r_prune, c_prune, arrays_prune, unprune = _prune(r, c, w)
    w_prune = arrays_prune[0]

    _check_margins(r_prune, c_prune)

    ### Preprocessing

    # Sizing (making copies of m and n, as they are mutated during sampling)
    r_init = r_prune.copy()
    m, n = len(r_prune), len(c_prune)
    if (m == 0) or (n == 0):
        if T:
            return [unprune([np.empty((0,2)), 0, 0]) for t in xrange(T)]
        else:
            return np.empty((0,0))
    m_init, n_init = m, n
    assert((m,n) == w_prune.shape)

    # Sort the row margins (descending)
    rndx_init = np.argsort(-r_prune)
    rsort = r_prune[rndx_init]

    # Balance the weights
    a_scale, b_scale = canonical_scalings(w_prune, r_prune, c_prune)
    wopt = apply_scale(w_prune, a_scale, b_scale)

    # Reorder the columns
    if sort_by_wopt_var:
        cndx = np.lexsort((-wopt.var(0), c_prune))
    else:
        cndx = np.argsort(c_prune)
    csort = c_prune[cndx]
    wopt = wopt[:,cndx]

    # Precompute log weights
    logw = np.log(w_prune)

    # Compute G
    G = _compute_G(r_prune, m, n, wopt)

    # Generate the inverse index for the row orders to facilitate fast
    # sorting during the updating
    irndx_init = np.argsort(rndx_init)

    # Compute the conjugate of c
    cconj_init = conjugate(csort, m)

    # Get the running total of number of ones to assign
    count_init = np.sum(rsort)

    def do_sample():
        sample_prune = _compute_sample(logw,
                                       count_init, m_init, n_init,
                                       r_init, rndx_init, irndx_init,
                                       csort, cndx, cconj_init,
                                       G)
        return unprune(sample_prune)
    
    if T:
        return [do_sample() for t in xrange(T)]
    else:
        return do_sample()[0] def clear_cache():
    _dict_canonical_scalings = {}
    _dict_conjugate = {} def invert_test(theta_grid, test_val, crit):
    theta_l_min, theta_l_max = theta_grid.min(), theta_grid.max()

    C_alpha = theta_grid[test_val > crit]
    if len(C_alpha) == 0:
        return 0, 0

    C_alpha_l, C_alpha_u = C_alpha.min(), C_alpha.max()
    if C_alpha_l == theta_l_min:
        C_alpha_l = -np.inf
    if C_alpha_u == theta_l_max:
        C_alpha_u = np.inf

    return C_alpha_l, C_alpha_u def ci_conservative_generic(X, K, theta_grid, alpha_level,
                            suff, log_likelihood, sample, t,
                            corrected = True, two_sided = True,
                            verbose = False):
    L = len(theta_grid)
    
    # Generate samples from the mixture proposal distribution
    Y = [sample(theta_grid[np.random.randint(L)]) for k in range(K)]

    # Test statistic at observation, for each grid point
    t_X = t(X).reshape((L, 1))
    if verbose:
        print 'X: t_min = %.2f, t_max = %.2f' % (t_X.min(), t_X.max())
        
    # Statistics for the samples from the proposal distribution only
    # need to be calculated once...
    t_Y = np.empty((L, K+1))
    t_Y[:,K] = 0.0
    for k in range(K):
        t_Y[:,k] = t(Y[k])
        if verbose:
            print 'Y_%d: t_min = %.2f, t_max = %.2f' % \
                (k, t_Y[:,k].min(), t_Y[:,k].max())
    if two_sided:
        I_t_Y_plus = t_Y >= t_X
        I_t_Y_plus[:,K] = True
    I_t_Y_minus = -t_Y >= -t_X
    I_t_Y_minus[:,K] = True

    # Probabilities under each component of the proposal distribution
    # only need to be calculated once...
    log_Q_X = np.empty(L)
    log_Q_Y = np.empty((L,K))
    for l in range(L):
        theta_l = theta_grid[l]

        log_Q_X[l] = log_likelihood(X, theta_l)
        for k in range(K):
            log_Q_Y[l,k] = log_likelihood(Y[k], theta_l)
        if verbose:
            print '%.2f: %.2g, %.2g' % \
              (theta_l, np.exp(log_Q_X[l]), np.exp(log_Q_Y[l].max()))
    log_Q_sum_X = logsumexp(log_Q_X)
    log_Q_sum_Y = np.empty(K)
    for k in range(K):
        log_Q_sum_Y[k] = logsumexp(log_Q_Y[:,k])

    # Step over the grid, calculating approximate p-values
    if two_sided:
        log_p_plus = np.empty(L)
    log_p_minus = np.empty(L)
    for l in range(L):
        theta_l = theta_grid[l]
        log_w_l = np.empty(K + 1)

        # X contribution
        if corrected:
            log_w_l[K] = theta_l * (suff * X).sum() - log_Q_sum_X
        else:
            log_w_l[K] = -np.inf

        # Y contribution
        for k in range(K):
            log_w_l[k] = theta_l * (suff * Y[k]).sum() - log_Q_sum_Y[k]

        if two_sided:
            log_p_num_plus = logsumexp(log_w_l[I_t_Y_plus[l]])
        log_p_num_minus = logsumexp(log_w_l[I_t_Y_minus[l]])
        log_p_denom = logsumexp(log_w_l)

        if verbose:
            print '%.2f: %.2g (%.2g, %.2g)' % \
              (theta_l, log_w_l[K], log_w_l[0:K].min(), log_w_l[0:K].max())

        if two_sided:
            log_p_plus[l] = log_p_num_plus - log_p_denom
        log_p_minus[l] = log_p_num_minus - log_p_denom

    if two_sided:
        # p_pm = min(1, 2 * min(p_plus, p_minus))
        log_p_vals = np.fmin(0, np.log(2) + np.fmin(log_p_plus, log_p_minus))
    else:
        log_p_vals = np.fmin(0, log_p_minus)

    return invert_test(theta_grid, log_p_vals, np.log(alpha_level)) class Seed:
    def __init__(self, seed = 0):
        self.seed = seed
        np.random.seed(self.seed)

    def next(self):
        self.seed += 1
        np.random.seed(self.seed)

    def final(self):
        self.next()
        print 'URN from Seed:', np.random.random() class FixedMargins(IndependentBernoulli):
    def __init__(self, base_model, r_name = 'r', c_name = 'c', coverage = 0):
        self.base_model = base_model
        self.r_name = r_name
        self.c_name = c_name
        self.coverage = coverage
        self.conf = tree()

    def check_separated(self, network, samples = 100, beta_scale = 8.0):
        A = network.as_dense()

        beta = self.base_model.beta
        B = len(beta)

        overlap = np.zeros(B, dtype=np.bool)
        for b, b_n in enumerate(beta):
            b_mat = network.edge_covariates[b_n].matrix()

            T_0 = b_mat[A].sum()
            T_samp = np.empty(2 * samples)

            beta[b_n] = -beta_scale
            for s in range(samples):
                A_samp = self.generate(network)
                T_samp[s] = b_mat[A_samp].sum()

            beta[b_n] = beta_scale
            for s in range(samples):
                A_samp = self.generate(network)
                T_samp[samples + s] = b_mat[A_samp].sum()

            beta[b_n] = 0.0

            overlap[b] = T_samp.min() < T_0 < T_samp.max()

        self.base_model.fit_info['separated'] = np.any(~overlap)

    def generate(self, network, **opts):
        if not 'coverage' in opts:
            opts['coverage'] = self.coverage

        if not self.r_name in network.row_covariates:
            print 'Row covariate "%s" not found.' % self.r_name
            r = np.asarray(network.as_dense()).sum(1, dtype = np.int)
        else:
            r = network.row_covariates[self.r_name][:]

        if not self.c_name in network.col_covariates:
            print 'Column covariate "%s" not found.' % self.c_name
            c = np.asarray(network.as_dense()).sum(0, dtype = np.int)
        else:
            c = network.col_covariates[self.c_name][:]

        return self.base_model.generate_margins(network, r, c, **opts)

    def nll(self, network, **opts):
        A = network.as_dense()
        w = np.exp(self.edge_probabilities(network, logit = True, **opts))

        return acnll(A, w)

    def edge_probabilities(self, network, **opts):
        return self.base_model.edge_probabilities(network, **opts)

    def confidence_wald(self, network, alpha_level = 0.05, strict = True,
                        **fit_options):
        self.fit(network, **fit_options)
        method = strict and 'wald' or 'wald_inverse'

        z_score = norm().ppf(1.0 - alpha_level / 2.0)
        for b in self.base_model.beta:
            b_name = 'theta_{%s}' % b
            if b_name in self.base_model.hess:
                if strict:
                    I_inv_b = 1.0 / self.base_model.hess[b_name]
                else:
                    I_inv_b = self.base_model.hess_inv[b_name]
                self.conf[b][method] = \
                    (self.base_model.beta[b] - z_score * np.sqrt(I_inv_b),
                     self.base_model.beta[b] + z_score * np.sqrt(I_inv_b))

    def confidence_boot(self, network, **opts):
        self.base_model.fit = self.fit
        self.base_model.generate = self.generate
        self.base_model.confidence_boot(network, **opts)
        lift_tree(self.base_model.conf, self.conf)

    def confidence_cons(self, network, b, **opts):
        self.base_model.fit = self.fit
        self.base_model.generate = self.generate
        self.base_model.confidence_cons(network, b, **opts)
        lift_tree(self.base_model.conf, self.conf)

    def reset_confidence(self):
        self.conf = tree() class StationaryLogistic(Stationary):
    def __init__(self):
        Stationary.__init__(self)
        self.beta = {}
        self.fit = self.fit_convex_opt

    def edge_probabilities(self, network, submatrix = None,
                           ignore_offset = False, logit = False):
        M = network.M
        N = network.N
        if submatrix:
            i_sub, j_sub = submatix
            M, N = len(i_sub), len(j_sub)
        
        if (not ignore_offset) and network.offset:
            logit_P = network.offset.matrix().copy()
            if submatrix:
                logit_P = logit_P[i_sub][:,j_sub]
        else:
            logit_P = np.zeros((M,N))
        for b in self.beta:
            ec_b = network.edge_covariates[b].matrix()
            if submatrix:
                ec_b = ec_b[i_sub][:,j_sub]
            logit_P += self.beta[b] * ec_b
        logit_P += self.kappa

        if logit:
            return logit_P
        else:
            return inv_logit(logit_P)

    # The network is needed for its covariates, not for the observed
    # pattern of edges, etc.
    def fisher_information(self, network, inverse = False):
        M = network.M
        N = network.N
        B = len(self.beta)

        P = self.edge_probabilities(network)

        x = np.empty((B,M,N))
        for i, b in enumerate(self.beta):
            x[i] = network.edge_covariates[b].matrix()

        P_bar = P * (1.0 - P)

        I = np.zeros((1 + B, 1 + B))
        I[0,0] = P_bar.sum()
        for b in range(B):
            v = (x[b] * P_bar).sum()
            I[1 + b,0] = v
            I[0,1 + b] = v
        for b_1 in range(B):
            for b_2 in range(B):
                v = (x[b_1] * x[b_2] * P_bar).sum()
                I[1 + b_1,1 + b_2] = v
                I[1 + b_2,1 + b_1] = v

        names_theta = ['theta_{%s}' % b for b in self.beta]
        I_names = np.array(['kappa'] + names_theta)
        self.I = {}
        for i in range(1 + B):
            for j in range(1 + B):
                if i == j:
                    self.I[I_names[i]] = I[i,i]
                else:
                    self.I[(I_names[i],I_names[j])] = I[i,j]

        if inverse:
            self._invert_fisher_information(I, I_names)

    def _invert_fisher_information(self, I, I_names):
        self.I_inv = {}

        I_r_keep = (I.sum(1) != 0)
        L = I_r_keep.sum()
        I_keep = I[I_r_keep][:,I_r_keep]

        try:
            I_inv = inv(I_keep)
        except:
            print 'Warning: unable to invert Fisher information matrix.'
            return

        I_inv_names = I_names[I_r_keep]
        for i in range(L):
            for j in range(L):
                if i == j:
                    self.I_inv[I_inv_names[i]] = I_inv[i,i]
                else:
                    self.I_inv[(I_inv_names[i],I_inv_names[j])] = I_inv[i,j]

    def check_separated(self, network):
        A = network.as_dense()
        B = len(self.beta)

        overlap = np.zeros(B, dtype=np.bool)
        for b, b_n in enumerate(self.beta):
            b_mat = network.edge_covariates[b_n].matrix()

            b_0 = b_mat[-A]
            b_1 = b_mat[A]

            l_0 = b_0.min()
            u_0 = b_0.max()
            l_1 = b_1.min()
            u_1 = b_1.max()

            if (((l_0 < u_1) and (l_1 < u_0)) or
                (l_0 == u_0 == l_1 <  u_1) or
                (l_1 <  u_1 == l_0 == u_0) or
                (l_1 == u_1 == l_0 <  u_0) or
                (l_0 <  u_0 == l_1 == u_1)):
                overlap[b] = True

        self.fit_info['separated'] = np.any(~overlap)

    def baseline(self, network):
        return np.mean(self.edge_probabilities(network))

    def baseline_logit(self, network):
        return np.mean(self.edge_probabilities(network, logit = True))

    def fit_convex_opt(self, network, verbose = False, fix_beta = False):
        B = len(self.beta)

        self.fit_info['nll_evals'] = 0
        self.fit_info['grad_nll_evals'] = 0
        self.fit_info['grad_nll_final'] = np.empty(B + 1)
        
        start_time = time()

        # Calculate observed sufficient statistics
        T = np.empty(B + 1)
        A = network.as_dense()
        for b, b_n in enumerate(self.beta):
            T[b] = (A * network.edge_covariates[b_n].matrix()).sum()
        T[B] = A.sum(dtype=np.int)

        # Initialize theta
        theta = np.zeros(B + 1)
        if fix_beta:
            for b, b_n in enumerate(self.beta):
                theta[b] = self.beta[b_n]
        theta[B] = logit(A.sum(dtype=np.int) / (1.0 * network.M * network.N))
        if network.offset:
            theta[B] -= logit_mean(network.offset.matrix())

        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for nan-containing vector.'
                return np.Inf
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            self.kappa = theta[B]
            nll = self.nll(network)
            self.fit_info['nll_evals'] += 1
            return nll
        def grad(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing gradient for nan-containing vector.'
                return np.zeros(B + 1)
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            self.kappa = theta[B]
            ET = np.empty(B + 1)
            P = self.edge_probabilities(network)
            if fix_beta:
                ET[0:B] = 0.0
            else:
                for b, b_n in enumerate(self.beta):
                    ET[b] = (P * network.edge_covariates[b_n].matrix()).sum()
            ET[B] = P.sum()
            g = ET - T
            if fix_beta:
                g[0:B] = 0.0
            self.fit_info['grad_nll_evals'] += 1
            self.fit_info['grad_nll_final'][:] = g
            if verbose:
                abs_grad = np.abs(g)
                print '|ET - T|: %.2f, %.2f, %.2f (min, mean, max)' % \
                    (np.min(abs_grad), np.mean(abs_grad), np.max(abs_grad))
            return g

        bounds = [(-8,8)] * B + [(-15,15)]
        theta_opt = opt.fmin_l_bfgs_b(obj, theta, grad, bounds = bounds)[0]
        if (np.any(theta_opt == [b[0] for b in bounds]) or
            np.any(theta_opt == [b[1] for b in bounds])):
            print 'Warning: some constraints active in model fitting.'
            if verbose:
                for b, b_n in enumerate(self.beta):
                    if theta_opt[b] in (bounds[b][0], bounds[b][1]):
                        print '%s: %.2f (T = %.2f)' % (b_n, theta_opt[b], T[b])
                if theta_opt[B] in (bounds[B][0], bounds[B][1]):
                    print 'kappa: %.2f (T = %.2f)' % (theta_opt[B], T[B])
            
        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = theta_opt[b]
        self.kappa = theta_opt[B]

        self.fit_info['wall_time'] = time() - start_time

    def fit_conditional(self, network,
                        fit_grid = False, T = 0, one_sided = False,
                        evaluate = False, verbose = False):
        B = len(self.beta)
        if fit_grid and not B in (1,2):
            print 'Can only grid search B = 1, 2. Defaulting to minimizer.'
            fit_grid = False

        self.fit_info['cnll_evals'] = 0

        start_time = time()

        A = network.as_dense()
        r, c = A.sum(1, dtype=np.int), A.sum(0, dtype=np.int)

        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for nan-containing vector.'
                return np.Inf
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            logit_P = StationaryLogistic.edge_probabilities(self, network,
                                                            logit = True)
            w = np.exp(logit_P)

            if T == 0:
                if (np.all(w == 0.0) or np.all(w == np.Inf) or
                    np.isnan(w).any()):
                    cnll = np.Inf
                else:
                    cnll = acnll(A, w, sort_by_wopt_var = False)
            else:
                z = acsample(r, c, w, T, sort_by_wopt_var = True)
                if not verbose:
                    logkappa = log_partition_is(z)
                else:
                    logkappa, logcvsq = log_partition_is(z, cvsq = True)
                    print 'est. cv^2 = %.2f (T = %d)' % (np.exp(logcvsq), T)
                cnll = logkappa - np.log(w[A]).sum()

            self.fit_info['cnll_evals'] += 1
            if verbose:
                print cnll, theta
            return cnll

        if fit_grid:
            if B == 1:
                G = 100
                theta_0 = np.linspace(-6.0, 6.0, G)
                cnll = np.empty(G)
                for g in range(G):
                    cnll[g] = obj(np.array([theta_0[g]]))
                theta_0_opt_ind = np.argmin(cnll)
                self.beta[self.beta.keys()[0]] = theta_0[theta_0_opt_ind]
            if B == 2:
                G = 15
                theta_0 = np.linspace(-6.0, 6.0, G)
                theta_1 = np.linspace(-6.0, 6.0, G)
                cnll = np.empty((G,G))
                for g_0 in range(G):
                    for g_1 in range(G):
                        t_0, t_1 = theta_0[g_0], theta_1[g_1]
                        cnll[g_0,g_1] = obj(np.array([t_0,t_1]))
                cnll_min = np.min(cnll)
                theta_opt_ind = np.where(cnll == cnll_min)
                theta_0_opt = theta_0[theta_opt_ind[0][0]]
                theta_1_opt = theta_1[theta_opt_ind[1][0]]
                self.beta[self.beta.keys()[0]] = theta_0_opt
                self.beta[self.beta.keys()[1]] = theta_1_opt
        else:
            if evaluate:
                theta = np.zeros(B)
                for b, b_n in enumerate(self.beta):
                    theta[b] = self.beta[b_n]

                cnll = obj(theta)
                self.fit_info['wall_time'] = time() - start_time

                return cnll
            else:
                # Initialize theta
                theta = np.zeros(B)

                if T > 0:
                    # Use Kiefer-Wolfowitz stochastic approximation
                    scale = 1.0 / obj(np.repeat(0, B))
                    for n in range(1, 40):
                        a_n = 2.0 * scale * n ** (-1.0)
                        c_n = 0.5 * n ** (-1.0 / 3)
                        grad = np.empty(B)
                        if one_sided:
                            y_0 = obj(theta)
                            for b in range(B):
                                e = np.zeros(B)
                                e[b] = 1.0
                                y_p = obj(theta + 2.0 * c_n * e)
                                grad[b] = (y_p - y_0) / c_n
                        else:
                            for b in range(B):
                                e = np.zeros(B)
                                e[b] = 1.0
                                y_p = obj(theta + c_n * e)
                                y_m = obj(theta - c_n * e)
                                grad[b] = (y_p - y_m) / c_n
                        theta -= a_n * grad
                    theta_opt = theta
                    for b, b_n in enumerate(self.beta):
                        self.beta[b_n] = theta_opt[b]
                else:
                    #theta_opt = opt.fmin(obj, theta)
                    res = opt.minimize(obj, theta, method = 'L-BFGS-B',
                                       bounds = [(-10.0, 10.0)] * B)

                    for b, b_n in enumerate(self.beta):
                        self.beta[b_n] = res.x[b]

                    hess_inv = res.hess_inv.todense()
                    hess = inv(hess_inv)
                    self.hess_inv = {}
                    self.hess = {}
                    names = ['theta_{%s}' % b for b in self.beta]
                    for i in range(B):
                        for j in range(B):
                            if i == j:
                                self.hess_inv[names[i]] = hess_inv[i,i]
                                self.hess[names[i]] = hess[i,i]
                            else:
                                ij_name = (names[i],names[j])
                                self.hess_inv[ij_name] = hess_inv[i,j]
                                self.hess[ij_name] = hess[i,j]

        self.fit_info['wall_time'] = time() - start_time
        
    def fit_saddlepoint(self, network, verbose = False):
        B = len(self.beta)
        M = network.M
        N = network.N
        A = network.as_dense()

        self.fit_info['cnll_evals'] = 0
        
        start_time = time()

        # Identify non-extreme sub-matrix on which saddlepoint
        # approximation is well-defined.
        i_nonextreme = range(M)
        j_nonextreme = range(N)
        while True:
            found = False
            A_j_nonextreme = A[:,j_nonextreme]
            for i in i_nonextreme:
                r = A_j_nonextreme[i,:].sum()
                if r == 0 or r == len(j_nonextreme):
                    found = True
                    i_nonextreme.remove(i)
            if found:
                continue

            found = False
            A_i_nonextreme = A[i_nonextreme,:]
            for j in j_nonextreme:
                c = A_i_nonextreme[:,j].sum()
                if c == 0 or c == len(i_nonextreme):
                    found = True
                    j_nonextreme.remove(j)
            if found:
                continue

            break

        # Initialize theta
        theta = np.zeros(B + 1)
        theta[B] = logit(np.mean(A[i_nonextreme][:,j_nonextreme]))
        if network.offset:
            theta[B] -= logit_mean(network.offset.matrix())

        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for nan-containing vector.'
                return np.Inf
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            self.kappa = theta[B]
            nll = self.nll(network)
            A_non = A[i_nonextreme][:,j_nonextreme]
            r_non = A_non.sum(1)
            c_non = A_non.sum(0)
            P = self.edge_probabilities(network)
            P_non = P[i_nonextreme][:,j_nonextreme]
            log_p_denom = log_p_margins_saddlepoint(r_non, c_non, P_non)
            cnll = nll + log_p_denom
            self.fit_info['cnll_evals'] += 1
            if verbose:
                print theta, nll, log_p_denom, cnll
            return cnll

        bounds = [(-8,8)] * B + [(-15,15)]
        theta_opt = opt.fmin_l_bfgs_b(obj, theta, approx_grad = True,
                                      bounds = bounds)[0]
        if (np.any(theta_opt == [b[0] for b in bounds]) or
            np.any(theta_opt == [b[1] for b in bounds])):
            print 'Warning: some constraints active in model fitting.'
            if verbose:
                for b, b_n in enumerate(self.beta):
                    if theta_opt[b] in (bounds[b][0], bounds[b][1]):
                        print '%s: %.2f' % (b_n, theta_opt[b])
                if theta_opt[B] in (bounds[B][0], bounds[B][1]):
                    print 'kappa: %.2f' % theta_opt[B]

        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = theta_opt[b]
        self.kappa = theta_opt[B]

        self.fit_info['wall_time'] = time() - start_time

    def fit_brazzale(self, network, target_n,
                     alpha_level = 0.05, verbose = False):
        M = network.M
        N = network.N
        A = network.as_dense()

        # Construct (flattened) R data frame with data
        y_vec = robjects.FloatVector(A.flatten())
        row_vec = robjects.IntVector(np.repeat(range(M), N))
        col_vec = robjects.IntVector(range(N) * M)
        dat_cols = { 'y': y_vec, 'row': row_vec, 'col': col_vec}
        for b_n in self.beta:
            x_b = network.edge_covariates[b_n].matrix()
            x_b_vec = robjects.FloatVector(x_b.flatten())
            if b_n == target_n:
                b_n = '..target..'
            dat_cols[b_n] = x_b_vec
        if network.offset:
            o = network.offset.matrix()
            o_vec = robjects.FloatVector(o.flatten())
            dat_cols['o'] = o_vec
        dat = robjects.DataFrame(dat_cols)
        robjects.globalenv['dat'] = dat
        robjects.r('dat$row <- factor(dat$row)')
        robjects.r('dat$col <- factor(dat$col)')

        # Use extreme offset values to prune down to non-extreme substructure
        # FIXME: this pruning is overly aggressive for conditional inference!
        if network.offset:
            robjects.r('dat <- dat[is.finite(dat$o),]')
        if verbose:
            robjects.r('write.csv(dat, file = "debug.csv")')
        robjects.r('nr <- length(unique(dat$row))')
        robjects.r('nc <- length(unique(dat$col))')

        # Do inference, defaulting to theta_hat = 0.0 if anything goes wrong
        if robjects.r('dim(dat)')[0] == 0:
            self.beta[target_n] = 0.0
        else:
            robjects.r('dat.glm <- glm(y ~ ..target.. + ' + \
              'C(row, how.many=(nr-1)) + C(col, how.many=(nc-1)), ' + \
              'data=dat, family=binomial("logit"))')
            try:
                robjects.r('dat.cond <- cond(dat.glm, ..target.., ' + \
                           'from=-8.0, to=8.0, pts=30)')
                robjects.globalenv['alpha'] = alpha_level
                robjects.r('dat.cond.summ <- summary(dat.cond, alpha=alpha)')
                if verbose:
                    print robjects.r('dat.cond.summ')
                theta_opt = robjects.r('dat.cond.summ$coefficients[2,1]')[0]
                # "Modified likelihood root (cont. corr.)"
                theta_ci_l = robjects.r('dat.cond.summ$conf.int[1,5]')[0]
                theta_ci_u = robjects.r('dat.cond.summ$conf.int[3,5]')[0]
                self.beta[target_n] = theta_opt
                self.conf[target_n]['brazzale'] = (theta_ci_l, theta_ci_u)
            except:
                self.beta[target_n] = 0.0
                self.conf[target_n]['brazzale'] = (0.0, 0.0)

        self.fit_convex_opt(network, fix_beta = True)

    def fit_c_conditional(self, network, verbose = False, evaluate = False):
        M = network.M
        N = network.N
        B = len(self.beta)

        self.fit_info['c_cnll_evals'] = 0

        start_time = time()

        # Pre-compute column statistics
        T_c = {}
        A = network.as_dense()
        c = A.sum(0, dtype=np.int)
        for b, b_n in enumerate(self.beta):
            T_b = A * network.edge_covariates[b_n].matrix()
            for j in range(N):
                T_c[(j,b)] = T_b[:,j]
        z = {}
        for j in range(N):
            z[j] = A[:,j].flatten()

        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for vector with nan.'
                return np.Inf
            c_cnll = 0
            alpha = np.zeros(M)
            alpha[0:(M-1)] = theta[B:(B + (M-1))]
            for j in range(N):
                c_j = c[j]

                if ((c_j == 0) or (c_j == M)):
                    continue

                logit_P = np.zeros(M)

                logit_P += alpha
                for b in range(B):
                    logit_P += theta[b] * T_c[(j,b)]
                P = inv_logit(logit_P)
                if P.sum() == 0:
                    continue

                log_P = np.log(P)
                log1m_P = np.log1p(-P)

                # Probability of column sum equal to c_j
                log_p_c = np.zeros(c_j + 1)
                for k in range(0, M):
                    log_bk, log1m_bk = log_P[k], log1m_P[k]

                    if k <= c_j:
                        i_init = k
                        if k < c_j:
                            log_p_c[k+1] = log_p_c[k] + log_bk
                    else:
                        i_init = c_j

                    for i in range(i_init, 0, -1):
                        x = log_p_c[i-1] + log_bk
                        y = log_p_c[i] + log1m_bk
                        x, y = min(x, y), max(x, y)
                        log_p_c[i] = y + np.log1p(np.exp(x-y))

                    log_p_c[0] += log1m_bk

                c_cnll += log_p_c[c_j]

                # Marginal probability of columns with given column sum
                z_j = z[j]
                c_cnll -= log_P[z_j].sum() + log1m_P[-z_j].sum()

            self.fit_info['c_cnll_evals'] += 1
            if verbose:
                print c_cnll, theta
            return c_cnll

        if evaluate:
            theta = np.zeros(B + (M-1))
            for b, b_n in enumerate(self.beta):
                theta[b] = self.beta[b_n]
            theta[B:(B + (M-1))] = network.row_covariates['alpha_out'][0:(M-1)]
            c_cnll = obj(theta)
            self.fit_info['wall_time'] = time() - start_time

            return c_cnll
        else:
            # Initialize theta
            theta = np.zeros(B + (M-1))

            bounds = [(-8,8)] * B + [(-6,6)] * (M-1)
            theta_opt = opt.fmin_l_bfgs_b(obj, theta, bounds = bounds,
                                          approx_grad = True)[0]
            if (np.any(theta_opt == [b[0] for b in bounds]) or
                np.any(theta_opt == [b[1] for b in bounds])):
                print 'Warning: some constraints active in model fitting.'
                if verbose:
                    for b, b_n in enumerate(self.beta):
                        if theta_opt[b] in (bounds[b][0], bounds[b][1]):
                            print '%s: %.2f (T = %.2f)' % \
                                (b_n, theta_opt[b], T[b])
                    for i in range(M-1):
                        r_i = B + i
                        if theta_opt[r_i] in (bounds[r_i][0], bounds[r_i][1]):
                            print 'alpha_%d: %.2f (T = %.2f)' % \
                                (i, theta_opt[r_i], T[r_i])

            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta_opt[b]

            self.fit_convex_opt(network, fix_beta = True)

            self.fit_info['wall_time'] = time() - start_time
        
    def fit_composite(self, network, T = 100, one_sided = False,
                      verbose = False):
        M = network.M
        B = len(self.beta)

        start_time = time()

        A = network.as_dense()
        r, c = A.sum(1), A.sum(0)

        # Initialize theta
        theta = np.zeros(B)

        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for nan-containing vector.'
                return np.Inf
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            logit_P = StationaryLogistic.edge_probabilities(self, network,
                                                            logit = True)
            w = np.exp(logit_P)

            cnll = 0.0
            for t in range(T):
                # Pick a random pair of rows 
                a, b = 0, 0
                while a == b:
                    a, b = np.random.randint(M), np.random.randint(M)

                # Extract subnetwork
                A_sub = A[[a,b],:]
                w_sub = w[[a,b],:]

                # Extract nontrivial columns
                nontrivial = A_sub.sum(0) == 1
                A_sub = A_sub[:,nontrivial]
                w_sub = w_sub[:,nontrivial]

                # Check if all permutations are trivial, so no change to cnll
                active = A_sub.shape[1]
                if active > 8:
                    print 'Skipping row pair with excessive active columns.'
                    continue
                if not (0 < A_sub[0,:].sum() < active):
                    continue

                logkappa = 0.0
                seen = set()
                for perm in permutations(range(active)):
                    A_sub_perm = A_sub[:,perm]

                    # Skip redundant permutations
                    rep = A_sub_perm[0,:].dumps()
                    if rep in seen: continue
                    seen.add(rep)

                    logp_perm = np.log(w_sub[A_sub_perm]).sum()
                    logkappa = np.logaddexp(logkappa, logp_perm)
                
                cnll += logkappa - np.log(w_sub[A_sub]).sum()

            if verbose:
                print cnll, theta
            return cnll

        # Use Kiefer-Wolfowitz stochastic approximation
        scale = 1.0 / obj(np.repeat(0, B))
        for n in range(1, 40):
            a_n = 2.0 * scale * n ** (-1.0)
            c_n = 0.5 * n ** (-1.0 / 3)
            grad = np.empty(B)
            if one_sided:
                y_0 = obj(theta)
                for b in range(B):
                    e = np.zeros(B)
                    e[b] = 1.0
                    y_p = obj(theta + 2.0 * c_n * e)
                    grad[b] = (y_p - y_0) / c_n
            else:
                for b in range(B):
                    e = np.zeros(B)
                    e[b] = 1.0
                    y_p = obj(theta + c_n * e)
                    y_m = obj(theta - c_n * e)
                    grad[b] = (y_p - y_m) / c_n
            theta -= a_n * grad
        theta_opt = theta

        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = theta_opt[b]

        self.fit_convex_opt(network, fix_beta = True)

        self.fit_info['wall_time'] = time() - start_time

    def fit_logistic(self, network):
        import statsmodels.api as sm

        M = network.M
        N = network.N
        B = len(self.beta)

        y = network.as_dense().reshape((M*N,))
        Phi = np.zeros((M*N,B + 1))
        Phi[:,B] = 1.0
        for b, b_n in enumerate(self.beta):
            Phi[:,b] =  network.edge_covariates[b_n].matrix().reshape((M*N,))

        try:
            if network.offset:
                offset = network.offset.matrix().reshape((M*N,))
                fit = sm.GLM(y, Phi, sm.families.Binomial(), offset).fit()
            else:
                fit = sm.Logit(y, Phi).fit()
            coefs = fit.params
        except:
            print 'Warning: logistic fit failed.'
            coefs = np.zeros(B + 1)

        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = coefs[b]
        self.kappa = coefs[B]

    # Doing this fit endows the model with a posterior
    # variance/covariance matrix estimated via the Laplace
    # approximation (see Bishop for details)
    def fit_logistic_l2(self, network, prior_precision = 1.0,
                        variance_covariance = False):
        if network.offset:
            print 'Regularized logistic regression with offset not supported.'
            raise
        
        from sklearn.linear_model import LogisticRegression

        M = network.M
        N = network.N
        B = len(self.beta)

        y = network.as_dense().reshape((M*N,))
        Phi = np.zeros((M*N,B))
        for b, b_n in enumerate(self.beta):
            Phi[:,b] =  network.edge_covariates[b_n].matrix().reshape((M*N,))
        lr = LogisticRegression(fit_intercept = True,
                                C = 1.0 / prior_precision,
                                penalty = 'l2')
        try:
            lr.fit(Phi, y)
            coefs, intercept = lr.coef_[0], lr.intercept_[0]
        except:
            print 'Warning: regularized logistic fit failed.'
            coefs, intercept = np.zeros(B), 0

        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = coefs[b]
        self.kappa = intercept

        if variance_covariance:
            S_0_inv = prior_precision * np.eye(B + 1)
            Phi_kappa = np.hstack([Phi, np.ones((M*N,1))])
            w = np.empty(B + 1)
            w[0:B] = coefs
            w[B] = intercept
            C = 0.0
            for i in range(M*N):
                y = inv_logit(np.dot(w, Phi_kappa[i,:]))
                C += y * (1.0 - y) * (np.outer(Phi_kappa[i,:], Phi_kappa[i,:]))
            S_N = np.linalg.inv(S_0_inv + C)

            # Convert variance/covariance matrix into a dictionary
            # indexed by pairs of parameters
            self.variance_covariance = {}
            parameters = self.beta.keys() + ['kappa']
            for r, p_1 in enumerate(parameters):
                for s, p_2 in enumerate(parameters):
                    self.variance_covariance[(p_1,p_2)] = S_N[r,s]

    def confidence_wald(self, network, alpha_level = 0.05, strict = True,
                        **fit_options):
        self.fit(network, **fit_options)
        self.fisher_information(network, inverse = (not strict))
        method = strict and 'wald' or 'wald_inverse'

        z_score = norm().ppf(1.0 - alpha_level / 2.0)
        if 'kappa' in self.I:
            if strict:
                I_inv_kappa = 1.0 / self.I['kappa']
            else:
                I_inv_kappa = self.I_inv['kappa']
            self.conf['kappa'][method] = \
                (self.kappa - z_score * np.sqrt(I_inv_kappa),
                 self.kappa + z_score * np.sqrt(I_inv_kappa))
        for b in self.beta:
            b_name = 'theta_{%s}' % b
            if b_name in self.I:
                if strict:
                    I_inv_b = 1.0 / self.I['kappa']
                else:
                    I_inv_b = self.I_inv[b_name]
                self.conf[b][method] = \
                    (self.beta[b] - z_score * np.sqrt(I_inv_b),
                     self.beta[b] + z_score * np.sqrt(I_inv_b))

    def confidence_boot(self, network, n_bootstrap = 100, alpha_level = 0.05,
                        **fit_options):
        # Point estimate
        self.fit(network, **fit_options)
        theta_hats = { b: self.beta[b] for b in self.beta }
        theta_hats['_kappa'] = self.kappa

        # Parametric bootstrap to characterize uncertainty in point estimate
        network_samples = [self.generate(network) for k in range(n_bootstrap)]
        network_original = network.array.copy()
        theta_hat_bootstraps = { b: np.empty(n_bootstrap) for b in self.beta }
        theta_hat_bootstraps['_kappa'] = np.empty(n_bootstrap)
        for k in range(n_bootstrap):
            network.array = network_samples[k]
            self.fit(network, **fit_options)
            for b in theta_hat_bootstraps:
                if b == '_kappa': continue
                theta_hat_bootstraps[b][k] = self.beta[b]
            theta_hat_bootstraps['_kappa'][k] = self.kappa
        network.array = network_original

        # Construct (asymptotically valid) confidence interval
        p_l, p_u = alpha_level / 2.0, 1.0 - alpha_level / 2.0
        for b in theta_hats:
            if b == '_kappa':
                display_name = 'kappa'
            else:
                display_name = b
            theta_hat = theta_hats[b]
            theta_hat_bootstrap = theta_hat_bootstraps[b]
            self.conf[display_name]['percentile'] = \
                (np.percentile(theta_hat_bootstrap, 100.0 * p_l),
                 np.percentile(theta_hat_bootstrap, 100.0 * p_u))
            self.conf[display_name]['pivotal'] = \
                (2*theta_hat - np.percentile(theta_hat_bootstrap, 100.0 * p_u),
                 2*theta_hat - np.percentile(theta_hat_bootstrap, 100.0 * p_l))
            z_score = norm().ppf(p_u)
            theta_hat_se = np.sqrt(np.mean((theta_hat_bootstrap-theta_hat)**2))
            self.conf[display_name]['normal'] = \
                (theta_hat - z_score * theta_hat_se,
                 theta_hat + z_score * theta_hat_se)

    # Implementation of ideas from "Conservative Hypothesis Tests and
    # Confidence Intervals using Importance Sampling" (Harrison, 2012).
    def confidence_cons(self, network, b, alpha_level = 0.05, n_MC = 100,
                        L = 601, beta_l_min = -6.0, beta_l_max = 6.0,
                        test = 'score', corrected = True, verbose = False):
        M = network.M
        N = network.N
        A = network.as_dense()
        x = network.edge_covariates[b].matrix()

        # Row and column margins; the part of the data we can use to design Q
        r, c = A.sum(1), A.sum(0)

        # Evaluate log-likelihood at specified parameter value
        def log_likelihood(z, theta):
            return -acnll(z, np.exp(theta * x), sort_by_wopt_var = False)

        # Generate sample from k-th component of mixture proposal distribution
        def sample(theta):
            Y_sparse = acsample(r, c, np.exp(theta * x), T = 0,
                                sort_by_wopt_var = False)
            Y_dense = np.zeros((M,N), dtype = np.bool)
            for i, j in Y_sparse:
                if i == -1: break
                Y_dense[i,j] = 1
            return Y_dense

        theta_grid = np.linspace(beta_l_min, beta_l_max, L)

        # Test statistic for CI
        if test == 'score':
            name = 'conservative-score'
            two_sided = True
            def t(z):
                return np.array([(z * x).sum() for theta_l in theta_grid])
        elif test == 'lr':
            name = 'conservative-lr'
            two_sided = False
            def t(z):
                ll = np.array([log_likelihood(z, theta_l)
                               for theta_l in theta_grid])
                return ll - ll.max()

        (l, u) = ci_conservative_generic(A, n_MC, theta_grid, alpha_level,
                                         x, log_likelihood, sample, t,
                                         corrected = corrected,
                                         two_sided = two_sided,
                                         verbose = verbose)

        self.conf[b][name] = (l, u) class NonstationaryLogistic(StationaryLogistic):
    def __init__(self):
        StationaryLogistic.__init__(self)
        self.fit = self.fit_convex_opt
        
    def edge_probabilities(self, network, submatrix = None,
                           ignore_offset = False, logit = False):
        M = network.M
        N = network.N
        if submatrix:
            i_sub, j_sub = submatrix

        alpha_out = network.row_covariates['alpha_out']
        alpha_in = network.col_covariates['alpha_in']
        if submatrix:
            alpha_out = alpha_out[i_sub]
            alpha_in = alpha_in[j_sub]
        
        if (not ignore_offset) and network.offset:
            logit_P = network.offset.matrix().copy()
            if submatrix:
                logit_P = logit_P[i_sub][:,j_sub]
        else:
            logit_P = np.zeros((M,N))
        np.add(logit_P, alpha_out[:].reshape((-1,1)), logit_P)
        np.add(logit_P, alpha_in[:].reshape((1,-1)), logit_P)
        for b in self.beta:
            ec_b = network.edge_covariates[b].matrix()
            if submatrix:
                ec_b = ec_b[i_sub][:,j_sub]
            logit_P += self.beta[b] * ec_b
        logit_P += self.kappa

        if logit:
            return logit_P
        else:
            return inv_logit(logit_P)

    def baseline(self, network):
        M = network.M
        N = network.N
        P = self.edge_probabilities(network)
        def params_to_Q(params):
            a = params[0:M]
            b = params[M:(M + N)]
            c = params[(M+N)]
            logit_Q = np.zeros((M,N))
            for i in range(M):
                logit_Q[i,:] += a[i]
            for j in range(N):
                logit_Q[:,j] += b[j]
            logit_Q += c
            return inv_logit(logit_Q)
        params = np.zeros(M + N + 1)
        def obj(params):
            Q = params_to_Q(params)
            return np.reshape((P - Q), (M*N,))
        best_params = opt.leastsq(obj, params)[0]
        Q = params_to_Q(best_params)
        return Q

    def baseline_logit(self, network):
        M = network.M
        N = network.N
        logit_P = self.edge_probabilities(network, logit = True)
        a, b = logit_P.mean(1), logit_P.mean(0)
        a_mean, b_mean = a.mean(), b.mean()
        a -= a_mean
        b -= b_mean
        c = a_mean + b_mean
        logit_Q = np.zeros((M,N))
        for i in range(M):
            logit_Q[i,:] += a[i]
        for j in range(N):
            logit_Q[:,j] += b[j]
        logit_Q += c
        return logit_Q

    def fit_convex_opt(self, network, verbose = False, fix_beta = False):
        M = network.M
        N = network.N
        B = len(self.beta)

        self.fit_info['nll_evals'] = 0
        self.fit_info['grad_nll_evals'] = 0
        self.fit_info['grad_nll_final'] = np.empty(B + 1 + (M-1) + (N-1))

        start_time = time()

        if network.offset:
            O = network.offset.matrix()
        alpha_zero(network)

        # Calculate observed sufficient statistics
        T = np.empty(B + 1 + (M-1) + (N-1))
        A = network.as_dense()
        r = A.sum(1, dtype=np.int)[0:(M-1)]
        c = A.sum(0, dtype=np.int)[0:(N-1)]
        A_sum = A.sum(dtype=np.int)
        for b, b_n in enumerate(self.beta):
            T[b] = (A * network.edge_covariates[b_n].matrix()).sum()
        T[B] = A_sum
        T[(B + 1):(B + 1 + (M-1))] = r
        T[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))] = c

        # Initialize theta
        theta = np.zeros(B + 1 + (M-1) + (N-1))
        if fix_beta:
            for b, b_n in enumerate(self.beta):
                theta[b] = self.beta[b_n]
        theta[B] = logit((A_sum + 1.0) / (1.0 * M * N + 2.0))
        if network.offset:
            theta[B] -= logit_mean(O)
        theta[(B + 1):(B + 1 + (M-1) + (N-1))] = -theta[B]
        for i in range(M-1):
            theta[B + 1 + i] += \
              logit((r[i] + 1.0) / (N + 2.0))
            if network.offset:
                o_row = logit_mean(O[i,:])
                if np.isfinite(o_row):
                    theta[B + 1 + i] -= o_row
        for j in range(N-1):
            theta[B + 1 + (M-1) + j] += \
              logit((c[j] + 1.0) / (M + 2.0))
            if network.offset:
                o_col = logit_mean(O[:,j])
                if np.isfinite(o_col):
                    theta[B + 1 + (M-1) + j] -= o_col

        alpha_out = network.row_covariates['alpha_out']
        alpha_in = network.col_covariates['alpha_in']
        def obj(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing objective for nan-containing vector.'
                return np.Inf
            alpha_out[0:M-1] = theta[(B + 1):(B + 1 + (M-1))]
            alpha_in[0:N-1] = theta[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))]
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            self.kappa = theta[B]
            nll = self.nll(network)
            self.fit_info['nll_evals'] += 1
            return nll
        def grad(theta):
            if np.any(np.isnan(theta)):
                print 'Warning: computing gradient for nan-containing vector.'
                return np.zeros(B + 1 + (M-1) + (N-1))
            alpha_out[0:M-1] = theta[(B + 1):(B + 1 + (M-1))]
            alpha_in[0:N-1] = theta[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))]
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta[b]
            self.kappa = theta[B]
            ET = np.empty(B + 1 + (M-1) + (N-1))
            P = self.edge_probabilities(network)
            Er = P.sum(1)[0:(M-1)]
            Ec = P.sum(0)[0:(N-1)]
            ET[(B + 1):(B + 1 + (M-1))] = Er
            ET[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))] = Ec
            if fix_beta:
                ET[0:B] = 0.0
            else:
                for b, b_n in enumerate(self.beta):
                    ET[b] = (P * network.edge_covariates[b_n].matrix()).sum()
            ET[B] = P.sum()
            g = ET - T
            if fix_beta:
                g[0:B] = 0.0
            self.fit_info['grad_nll_evals'] += 1
            self.fit_info['grad_nll_final'][:] = g
            if verbose:
                abs_grad = np.abs(g)
                print '|ET - T|: %.2f, %.2f, %.2f (min, mean, max)' % \
                    (np.min(abs_grad), np.mean(abs_grad), np.max(abs_grad))
            return g

        bounds = [(-8,8)] * B + [(-15,15)] + [(-8,8)] * ((M-1) + (N-1))
        theta_opt = opt.fmin_l_bfgs_b(obj, theta, grad, bounds = bounds)[0]
        if (np.any(theta_opt == [b[0] for b in bounds]) or
            np.any(theta_opt == [b[1] for b in bounds])):
            print 'Warning: some constraints active in model fitting.'
            if verbose:
                for b, b_n in enumerate(self.beta):
                    if theta_opt[b] in (bounds[b][0], bounds[b][1]):
                        print '%s: %.2f (T = %.2f)' % (b_n, theta_opt[b], T[b])
                if theta_opt[B] in (bounds[B][0], bounds[B][1]):
                    print 'kappa: %.2f (T = %.2f)' % (theta_opt[B], T[B])
                for i in range(M-1):
                    r_i = B + 1 + i
                    if theta_opt[r_i] in (bounds[r_i][0], bounds[r_i][1]):
                        print 'alpha_%d: %.2f (T = %.2f)' % \
                            (i, theta_opt[r_i], T[r_i])
                for j in range(N-1):
                    c_j = B + 1 + (M-1) + j
                    if theta_opt[c_j] in (bounds[c_j][0], bounds[c_j][1]):
                        print 'beta_%d: %.2f (T = %.2f)' % \
                            (j, theta_opt[c_j], T[c_j])

        alpha_out[0:M-1] = theta_opt[(B + 1):(B + 1 + (M-1))]
        alpha_in[0:N-1] = theta_opt[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))]
        alpha_out_mean = np.mean(alpha_out[:])
        alpha_in_mean = np.mean(alpha_in[:])
        alpha_out[:] -= alpha_out_mean
        alpha_in[:] -= alpha_in_mean
        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = theta_opt[b]
        self.kappa = theta_opt[B] + alpha_out_mean + alpha_in_mean

        self.fit_info['wall_time'] = time() - start_time

    def fit_conditional(self, network, **kwargs):
        StationaryLogistic.fit_conditional(self, network, **kwargs)

        start_time = time()

        self.fit_convex_opt(network, fix_beta = True)

        self.fit_info['wall_time'] += time() - start_time

    def fit_irls(self, network, verbose = False, perturb = 1e-4):
        M = network.M
        N = network.N
        B = len(self.beta)
        P = B + 1 + (M-1) + (N-1)

        start_time = time()

        alpha_zero(network)
        alpha_out = network.row_covariates['alpha_out']
        alpha_in = network.col_covariates['alpha_in']
        
        # Construct response and design matrices
        y = np.asarray(network.as_dense(), dtype='float64')
        y = y.reshape((M*N,1))
        X = np.zeros((M*N,P))
        for b, b_n in enumerate(self.beta):
            X[:,b] =  network.edge_covariates[b_n].matrix().reshape((M*N,))
        X[:,B] = 1.0
        for r in range(M-1):
            X_row = np.zeros((M,N))
            X_row[r,:] = 1.0
            X[:,B + 1 + r] = X_row.reshape((M*N,))
        for c in range(N-1):
            X_col = np.zeros((M,N))
            X_col[:,c] = 1.0
            X[:,B + 1 + (M-1) + c] = X_col.reshape((M*N,))

        theta = np.zeros((P,1))

        def fitted_p(theta):
            theta_vec = np.reshape(theta, (P,))
            alpha_out[0:M-1] = theta_vec[(B + 1):(B + 1 + (M-1))]
            alpha_in[0:N-1] = theta_vec[(B + 1 + (M-1)):P]
            for b, b_n in enumerate(self.beta):
                self.beta[b_n] = theta_vec[b]
            self.kappa = theta_vec[B]
            return self.edge_probabilities(network).reshape((M*N,1))

        for iter in range(10):
            p = fitted_p(theta)
            X_tilde = X * p
            del p
            for j in range(P):
                X_tilde[:,j] += np.random.uniform(-perturb, perturb, M*N)
            X_t = np.transpose(X)
            X_t_X_tilde = np.dot(X_t, X_tilde)
            del X_tilde
            hat = solve(X_t_X_tilde, X_t, overwrite_a = True)
            p = fitted_p(theta)
            theta += np.dot(hat, (y - p))

        theta_vec = np.reshape(theta, (P,))
        alpha_out[0:M-1] = theta_vec[(B + 1):(B + 1 + (M-1))]
        alpha_in[0:N-1] = theta_vec[(B + 1 + (M-1)):P]
        alpha_out_mean = np.mean(alpha_out[:])
        alpha_in_mean = np.mean(alpha_in[:])
        alpha_out[:] -= alpha_out_mean
        alpha_in[:] -= alpha_in_mean
        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = theta_vec[b]
        self.kappa = theta_vec[B] + alpha_out_mean + alpha_in_mean

        self.fit_info['wall_time'] = time() - start_time
        
    def fit_logistic(self, network):
        import statsmodels.api as sm

        M = network.M
        N = network.N
        B = len(self.beta)

        # Set up outcome and design matrix for fit
        y = network.as_dense().reshape((M*N,))
        Phi = np.zeros((M*N,B + 1 + (M-1) + (N-1)))
        for b, b_n in enumerate(self.beta):
            Phi[:,b] =  network.edge_covariates[b_n].matrix().reshape((M*N,))
        Phi[:,B] = 1.0
        for r in range(M-1):
            phi_row = np.zeros((M,N))
            phi_row[r,:] = 1.0
            Phi[:,B + 1 + r] = phi_row.reshape((M*N,))
        for c in range(N-1):
            phi_col = np.zeros((M,N))
            phi_col[:,c] = 1.0
            Phi[:,B + 1 + (M-1) + c] = phi_col.reshape((M*N,))

        # Do fit, defaulting to beta = 0 in case of problems
        try:
            if network.offset:
                offset = network.offset.matrix().reshape((M*N,))
                fit = sm.GLM(y, Phi, sm.families.Binomial(), offset).fit()
            else:
                fit = sm.Logit(y, Phi).fit()
            coefs = fit.params
        except:
            print 'Warning: logistic fit failed.'
            coefs = np.zeros(B + 1 + (M-1) + (N-1))

        alpha_zero(network)
        alpha_out = network.row_covariates['alpha_out']
        alpha_in = network.col_covariates['alpha_in']
        alpha_out[0:M-1] = coefs[(B + 1):(B + 1 + (M-1))]
        alpha_in[0:N-1] = coefs[(B + 1 + (M-1)):(B + 1 + (M-1) + (N-1))]
        alpha_out_mean = np.mean(alpha_out[:])
        alpha_in_mean = np.mean(alpha_in[:])
        alpha_out[:] -= alpha_out_mean
        alpha_in[:] -= alpha_in_mean
        for b, b_n in enumerate(self.beta):
            self.beta[b_n] = coefs[b]
        self.kappa = coefs[B] + alpha_out_mean + alpha_in_mean

    # The network is needed for its covariates and degree
    # heterogeneity terms, not for the observed pattern of edges, etc.
    #
    # Typically, the inverse Fisher information matrix will be more
    # useful (it gives a lower bound on the variances/covariances of
    # an unbised estimator), so that is calculated by default.
    def fisher_information(self, network, inverse = True):
        M = network.M
        N = network.N
        B = len(self.beta)

        P = self.edge_probabilities(network)

        x = np.empty((B,M,N))
        for i, b in enumerate(self.beta):
            x[i] = network.edge_covariates[b].matrix()

        P_bar = P * (1.0 - P)

        I = np.zeros(((M-1) + (N-1) + 1 + B, (M-1) + (N-1) + 1 + B))
        for i in range(M-1):
            for j in range(N-1):
                v = P_bar[i,j]
                I[i,(M-1)+j] = v
                I[(M-1)+j,i] = v
        for i in range(M-1):
            v = (P_bar[i,:]).sum()
            I[i,i] = v
            I[(M-1) + (N-1),i] = v
            I[i,(M-1) + (N-1)] = v
            for b in range(B):
                v = (x[b,i,:] * P_bar[i,:]).sum()
                I[(M-1) + (N-1) + 1 + b,i] = v
                I[i,(M-1) + (N-1) + 1 + b] = v
        for j in range(N-1):
            v = (P_bar[:,j]).sum()
            I[(M-1) + j,(M-1) + j] = v
            I[(M-1) + (N-1),(M-1) + j] = v
            I[(M-1) + j,(M-1) + (N-1)] = v
            for b in range(B):
                v = (x[b,:,j] * P_bar[:,j]).sum()
                I[(M-1) + (N-1) + 1 + b,(M-1) + j] = v
                I[(M-1) + j,(M-1) + (N-1) + 1 + b] = v
        I[(M-1) + (N-1),(M-1) + (N-1)] = P_bar.sum()
        for b in range(B):
            v = (x[b] * P_bar).sum()
            I[(M-1) + (N-1) + 1 + b,(M-1) + (N-1)] = v
            I[(M-1) + (N-1),(M-1) + (N-1) + 1 + b] = v
        for b_1 in range(B):
            for b_2 in range(B):
                v = (x[b_1] * x[b_2] * P_bar).sum()
                I[(M-1) + (N-1) + 1 + b_1,(M-1) + (N-1) + 1 + b_2] = v
                I[(M-1) + (N-1) + 1 + b_2,(M-1) + (N-1) + 1 + b_1] = v

        names_alpha = ['alpha_{%s}' % n for n in network.rnames[0:(M-1)]]
        names_beta = ['beta_{%s}' % n for n in network.cnames[0:(N-1)]]
        names_theta = ['theta_{%s}' % b for b in self.beta]
        I_names = np.array(names_alpha + names_beta + ['kappa'] + names_theta)
        self.I = {}
        for i in range((M-1) + (N-1) + 1 + B):
            for j in range((M-1) + (N-1) + 1 + B):
                if i == j:
                    self.I[I_names[i]] = I[i,i]
                else:
                    self.I[(I_names[i],I_names[j])] = I[i,j]

        if inverse:
            self._invert_fisher_information(I, I_names) def alpha_zero(network):
    alpha_f(network, lambda l: np.tile(0.0, l)) def l2(x):
def apply_scale(w, a, b):
def logabsdiffexp(x, y):
    def f(x, y):
def logit_mean(x):
def digest(x):
def tree():
def lift_tree(src, dst):
def init_latex_rendering():
def draw_confidence(a, m, S, levels = [0.95]):
import json
import signal
import numpy as np
import matplotlib.pyplot as plt
from time import time
from scipy.stats import chi2
from Array import Array, array_from_data
from BinaryMatrix import approximate_conditional_nll as cond_a_nll_b
from BinaryMatrix import approximate_from_margins_weights as cond_a_sample_b
from BinaryMatrix import clear_cache
from Confidence import invert_test, ci_conservative_generic
from Experiment import Seed
from Models import FixedMargins, StationaryLogistic, NonstationaryLogistic
from Models import alpha_zero
from Utility import logsumexp, logabsdiffexp
#!/usr/bin/env python





# Parameters
params = { 'case': { 'fixed_example': 'data/c_elegans_full_soma_dist.json',
                     #'M': 20,
                     #'N': 10,
                     #'r': 1,
                     #'c': 1,
                     #'kappa': -3.0,
                     #'alpha_min': 0.0,
                     #'beta_min': 0.0,
                     #'v_scale': 0.25,
                     #'v_loc': -np.sqrt(3),
                     #'v_discrete': False,
                     #'v_uniform': False,
                     #'v_normal': True
                   },
           'theta': -1.0, #-2.24,
           'alpha_level': 0.05,
           'n_MC_levels': [10], #[10, 50, 100, 500],
           'wopt_sort': False,
           'is_T': 50,
           'n_rep': 100,
           'L': 201,
           'theta_l': -2.0,
           'theta_u': 0.0,
           'random_seed': 137,
           'verbose': True,
           'plot': False,
           'clear_cache': False }

terminated = False
def sigint_handler(signum, frame):
    print 'Terminating after current trial completes.'
    global terminated
    terminated = True
signal.signal(signal.SIGINT, sigint_handler)

def cond_a_nll(X, w):
    return cond_a_nll_b(X, w, sort_by_wopt_var = params['wopt_sort'])

def cond_a_sample(r, c, w, T = 0):
    return cond_a_sample_b(r, c, w, T, sort_by_wopt_var = params['wopt_sort'])

# Generates (or loads) a particular realization of P_{ij} and then
# repeatedly samples independent Bernoulli random variables according
# to these cell probabilities.
def generate_data(case, theta, seed):
    # Advance random seed for parameter and covariate construction
    seed.next()

    case = params['case']
    alpha = beta = kappa = offset = 0
    conditional_sample = False
    if 'fixed_example' in case:
        # Load parameters and covariates
        with open(case['fixed_example'], 'r') as example_file:
            example = json.load(example_file)

            v = np.array(example['nu'])
            M, N = v.shape

            if 'alpha' in example:
                alpha = np.array(example['alpha']).reshape((M,1))
            if 'beta' in example:
                beta = np.array(example['beta']).reshape((1,N))
            if 'kappa' in example:
                kappa = example['kappa']
            if 'offset' in example:
                offset = example['offset']

            if ('r' in example) and ('c' in example):
                conditional_sample = True
                r = example['r']
                c = example['c']

    else:
        # Generate parameters and covariates
        M, N = case['M'], case['N']
        if 'alpha_min' in case:
            alpha = np.random.uniform(size = (M,1)) + case['alpha_min']
        if 'beta_min' in case:
            beta = np.random.uniform(size = (1,N)) + case['beta_min']
        if 'kappa' in case:
            kappa = case['kappa']
        if case['v_discrete']:
            v = np.sign(np.random.random(size = (M,N)) - 0.5) 
        elif case['v_uniform']:
            v = np.random.uniform(size = (M,N))
        elif case['v_normal']:
            v = np.random.normal(size = (M,N))
        if 'v_scale' in case:
            v *= case['v_scale']
        if 'v_loc' in case:
            v += case['v_loc']
        
        if ('r' in case) and ('c' in case):
            conditional_sample = True
            r = case['r']
            c = case['c']

    # Generate Bernoulli probabilities from logistic regression model
    logit_P = np.zeros((M,N)) + kappa
    logit_P += alpha
    logit_P += beta
    logit_P += theta * v
    logit_P += offset

    if conditional_sample:

        arr = Array(M, N)