class Network(Array):
    def __init__(self, N = 0, names = None):
        Array.__init__(self, N, N)
        if names is None:
            self.names = np.array(['%d' % n for n in range(self.N)])
        else:
            self.names = names
        self.rnames = self.names
        self.cnames = self.names
        self.node_covariates = {}

    def new_node_covariate(self, name, as_int = False):
        if as_int:
            node_cov = NodeCovariate(self.names, dtype = np.int)
        else:
            node_cov = NodeCovariate(self.names)
        self.node_covariates[name] = node_cov
        self.row_covariates[name] = node_cov
        self.col_covariates[name] = node_cov
        return node_cov

    def new_node_covariate_int(self, name):
        return self.new_node_covariate(name, as_int = True)

    def subnetwork(self, inds):
        sub_array = self.subarray(inds, inds)

        sub = Network(len(inds), self.names[inds])
        sub.array = sub_array.array
        sub.row_covariates = sub_array.row_covariates
        sub.col_covariates = sub_array.col_covariates
        if sub_array.offset:
            sub.offset = sub_array.offset
        sub.edge_covariates = sub_array.edge_covariates
        for node_covariate in self.node_covariates:
            src = self.node_covariates[node_covariate]
            sub.node_covariates[node_covariate] = src.subset(inds)

        return sub

    def nodes(self):
        return self.names

    def edges(self):
        if self.is_sparse():
            nz_i, nz_j = self.array.nonzero()
            for n in range(self.array.nnz):
                yield (self.names[nz_i[n]], self.names[nz_j[n]])
        else:
            for i in range(self.N):
                for j in range(self.N):
                    if self.array[i,j]:
                        yield (self.names[i], self.names[j])

    def show(self):
        graph = nx.DiGraph()

        for n in self.nodes():
            graph.add_node(n)

        for n_i, n_j in self.edges():
            graph.add_edge(n_i, n_j)

        pos = nx.nx_pydot.graphviz_layout(graph, prog = 'neato')
        nx.draw(graph, pos)
        plt.show()

    def show_graphviz(self, file = 'out.pdf', splines = True, labels = True):
        outfile = open('temp_graphviz.dot', 'w')
        outfile.write('digraph G {\n')
        outfile.write('size="12,16";\n')
        outfile.write('orientation=landscape;\n')
        outfile.write('overlap=none;\n')
        outfile.write('repulsiveforce=12;\n')
        if splines:
            outfile.write('splines=true;\n')

        for name in self.nodes():
            outfile.write('%s [label=""];\n' % name)

        for edge in self.edges():
            outfile.write('%s -> %s;\n' % edge)
                        
        outfile.write('}\n')
        outfile.close()
        system('fdp -Tps2 temp_graphviz.dot -o temp_graphviz.ps')
        unlink('temp_graphviz.dot')
        system('ps2pdf temp_graphviz.ps %s' % file)
        unlink('temp_graphviz.ps')

    def show_heatmap(self, order_by = None,
                     order_by_row = None, order_by_col = None):
        if order_by:
            title = 'Network ordered by node covariate\n"%s"' % order_by
            o = np.argsort(self.node_covariates[order_by][:])
        elif order_by_row:
            title = 'Network ordered by row covariate\n"%s"' % order_by_row
            o = np.argsort(self.row_covariates[order_by_row][:])
        elif order_by_col:
            title = 'Network ordered by column covariate\n"%s"' % order_by_col
            o = np.argsort(self.col_covariates[order_by_col][:])
        else:
            title, o = 'Unordered adjacency matrix', np.arange(self.N)

        f, (ax_im, ax_ord) = plt.subplots(2, sharex = True)
        f.set_figwidth(3)
        f.set_figheight(6)
        A = self.as_dense()
        ax_im.imshow(A[o][:,o]).set_cmap('binary')
        ax_im.set_ylim(0, self.N - 1)
        ax_im.set_xticks([])
        ax_im.set_yticks([])
        ax_im.set_title(title)
        #plt.setp([ax_im.get_xticklabels(), ax_im.get_yticklabels()],
        #         visible = False)
        if order_by:
            ax_ord.scatter(np.arange(self.N), self.node_covariates[order_by][o])
            ax_ord.set_xlim(0, self.N - 1)
            ax_ord.set_ylim(self.node_covariates[order_by][o[0]],
                            self.node_covariates[order_by][o[-1]])
        plt.show()

    def show_offset(self, order_by = None):
        if order_by:
            title = 'Offsets ordered by node covariate\n"%s"' % order_by
            o = np.argsort(self.node_covariates[order_by][:])
        else:
            title, o = 'Unordered offsets', np.arange(self.N)

        f = plt.figure()
        ax = f.add_subplot(1, 1, 1)
        O = self.initialize_offset().matrix()
        ax.imshow(O[o][:,o])
        ax.set_xlim(0, self.N - 1)
        ax.set_ylim(0, self.N - 1)
        ax.set_title(title)
        plt.setp([ax.get_xticklabels(), ax.get_yticklabels()],
                 visible = False)
        plt.show()

    def show_degree_histograms(self):
        # Messy since otherwise row/column sums can overflow...
        r = np.array(self.array.asfptype().sum(1),dtype=np.int).flatten()
        c = np.array(self.array.asfptype().sum(0),dtype=np.int).flatten()
        
        plt.figure()
        plt.subplot(2,1,1)
        plt.title('out-degree')
        plt.hist(r, bins = max(r))
        plt.subplot(2,1,2)
        plt.title('in-degree')
        plt.hist(c, bins = max(c))
        plt.show() 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 alpha_norm(network, alpha_sd):
    alpha_f(network, lambda l: center(np.random.normal(0, alpha_sd, l))) def alpha_unif(network, alpha_sd):
    c = np.sqrt(12) / 2
    alpha_f(network,
            lambda l: center(np.random.uniform(-alpha_sd*c, alpha_sd * c, l))) def alpha_gamma(network, alpha_loc, alpha_scale):
    alpha_f(network,
            lambda l: center(np.random.gamma(alpha_loc, alpha_scale, l))) class RandomSubnetworks:
    def __init__(self, network, train_size, test_size = 0, method = 'node'):
        self.network = network
        self.train_size = train_size
        self.test_size = test_size
        self.method = method

        if self.method == 'node':
            if type(self.train_size) == tuple:
                self.rinds = np.arange(self.network.N)
                self.cinds = np.arange(self.network.N)
            else:
                self.inds = np.arange(self.network.N)
        else:
            self.rinds = np.arange(self.network.M)
            self.cinds = np.arange(self.network.N)
            if self.method == 'edge':
                edges = self.network.array.nonzero()
                edges_i, edges_j = edges[0], edges[1]
                self.edges_i = edges_i
                self.edges_j = edges_j
                self.num_edges = len(edges_i)
            elif self.method in ('link', 'link_f') :
                edges = self.network.array.nonzero()
                self.num_edges = len(edges[0])
                neighbors = { n: set() for n in range(self.network.N) }
                for i, j in zip(edges[0], edges[1]):
                    neighbors[i].add(j)
                    if self.method == 'link':
                        neighbors[j].add(i)
                self.neighbors = { n: list(neighbors[n]) for n in neighbors }
                self.num_nodes = len(self.neighbors)

    def sample(self, as_network = False):
        M = self.network.M
        N = self.network.N
        s = self.train_size

        def produce(r, c):
            if as_network:
                inds = np.unique(np.array(list(r) + list(c)))
                return self.network.subnetwork(inds)
            else:
                r_a = np.array(list(r))
                r_c = np.array(list(c))
                return self.network.subarray(r_a, r_c)

        def fallback():
            s_r, s_c = s
            np.random.shuffle(self.rinds)
            np.random.shuffle(self.cinds)
            if as_network:
                return produce(self.rinds[0:s_r], self.rinds[0:s_r])
            else:
                return produce(self.rinds[0:s_r], self.cinds[0:s_c])

        if self.method == 'node':
            return fallback()
        elif self.method == 'row':
            np.random.shuffle(self.rinds)
            return produce(self.rinds[0:s], self.cinds)
        elif self.method == 'col':
            np.random.shuffle(self.cinds)
            return produce(self.rinds, self.cinds[0:s])
        elif self.method in ('edge', 'link', 'link_f'):
            s_r, s_c = s

            if (self.num_edges < s_r) or (self.num_edges < s_c):
                return fallback()

            added_r = set()
            added_c = set()

            if self.method == 'edge':
                while (len(added_r) < s_r) and (len(added_c) < s_c):
                    e = np.random.randint(self.num_edges)
                    edge_i = self.edges_i[e]
                    edge_j = self.edges_j[e]
                    added_r.add(edge_i)
                    added_c.add(edge_j)
            elif self.method in ('link', 'link_f'):
                # The difference between 'link' and 'link_f' is in the
                # construction of neighbors, not the algorithm used to
                # follow links
                loc = np.random.randint(self.num_nodes)
                while (len(added_r) < s_r) and (len(added_c) < s_c):
                    added_r.add(loc)
                    if added_r.issuperset(self.neighbors[loc]):
                        loc = np.random.randint(self.num_nodes)
                        continue
                    pick_j = np.random.randint(len(self.neighbors[loc]))
                    loc_j = self.neighbors[loc][pick_j]
                    added_c.add(loc_j)
                    loc = loc_j

            remain_r = list(set(range(M)).difference(added_r))
            np.random.shuffle(remain_r)
            added_r.update(set(remain_r[0:(s_r - len(added_r))]))

            remain_c = list(set(range(N)).difference(added_c))
            np.random.shuffle(remain_c)
            added_c.update(set(remain_c[0:(s_c - len(added_c))]))

            return produce(added_r, added_c)

        if self.test_size == 0:
            if self.method in ('row', 'col'):
                return self.network.subnetwork(sub_train, self.method)
            else:
                return self.network.subnetwork(sub_train)
        else:
            sub_full = self.inds[0:(self.train_size + self.test_size)]
            return (self.network.subnetwork(sub_train),
                    self.network.subnetwork(sub_full)) 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 Results:
    def __init__(self, M_sizes, N_sizes, num_reps, title = None,
                 interactive = True):
        self.M_sizes = M_sizes
        self.N_sizes = N_sizes

        len_M = len(self.M_sizes)
        len_N = len(self.N_sizes)
        if not (len_M == len_N):
            print 'Warning: mismatch in M-dims and N-dims lengths.'
            num_both = min(len_M, len_N)
            self.M_sizes = self.M_sizes[0:num_both]
            self.N_sizes = self.N_sizes[0:num_both]

        # Remove redundant conditions
        sizes_seen = set()
        nonredundant = np.zeros_like(M_sizes, dtype=np.bool)
        for i, size in enumerate(zip(M_sizes, N_sizes)):
            if not size in sizes_seen:
                nonredundant[i] = True
            sizes_seen.add(size)
        self.M_sizes = self.M_sizes[nonredundant]
        self.N_sizes = self.N_sizes[nonredundant]

        self.num_conditions = len(self.M_sizes)
        self.num_reps = num_reps
        self.title = title
        self.interactive = interactive
        self.results = {}

        self.size_to_ind = {}
        for i, size in enumerate(zip(self.M_sizes, self.N_sizes)):
            self.size_to_ind[size] = i

    def merge(self, other):
        if not (np.all(self.M_sizes == other.M_sizes) and
                np.all(self.N_sizes == other.N_sizes)):
            print 'Warning: mismatched conditions in Results to merge.'

        self.num_reps = self.num_reps + other.num_reps
        for result_name in self.results:
            old_data = self.results[result_name]['data']
            other_data = other.results[result_name]['data']
            merged_data = np.hstack([old_data, other_data])
            self.results[result_name]['data'] = merged_data

    # Return a copy of the result structure, with new allocated storage
    def copy(self):
        dup = Results(self.M_sizes, self.N_sizes,
                      self.num_reps,
                      title = self.title, interactive = self.interactive)

        for result_name in self.results:
            result = self.results[result_name]
            f, f_type = result['f'], result['f_type']

            dup.new(result_name, f_type, f)

        return dup

    # Values expected for "f_type": 'a' (adjacency)
    #                               'n' (network)
    #                               'm' (models),
    #                               'nm' (network and models)
    def new(self, name, f_type, f):
        assert(f_type in ['a', 'n', 'm', 'nm', 'dof'])
        self.results[name] = { 'f': f, 'f_type': f_type,
                               'data': np.empty((self.num_conditions,
                                                 self.num_reps)) }

        return name

    def record(self, size, rep, network,
               data_model = None, fit_model = None):
        for result in self.results:
            f = self.results[result]['f']
            f_type = self.results[result]['f_type']
            if f_type == 'a':
                val = f(network.as_dense())
            elif f_type == 'n':
                val = f(network)
            elif f_type == 'm':
                val = f(data_model, fit_model)
            elif f_type == 'nm':
                val = f(network, data_model, fit_model)
            elif f_type == 'dof':
                val = f(network.M, network.N, len(fit_model.beta))

            data = self.results[result]['data']
            data[self.size_to_ind[size], rep] = val

    # To be called after all results have been recorded...
    def sweep(self, name, func, *args):
        self.results[name] = {'data': np.empty((self.num_conditions,1))}

        data = []
        for arg in args:
            data.append(self.results[arg]['data'])

        for n in range(self.num_conditions):
            self.results[name]['data'][n,0] = func(n, *data)

    def estimate_mse(self, name, true, estimate):
        self.sweep(name, lambda n, t, e: np.mean((t[n]-e[n])**2),
                   true, estimate)

    def estimate_mae(self, name, true, estimate):
        self.sweep(name, lambda n, t, e: np.median(np.abs(t[n]-e[n])),
                   true, estimate)

    def summary(self):
        for field in self.results:
            data = self.results[field]['data']
            average_data = np.mean(data, 1)
            print '%s: %s' % (field, repr(average_data))

    def plot(self, requests = None, general = {'xaxis': 'c'}):
        if general['xaxis'] == 'c':
            sizes = self.N_sizes
            sizes_label = 'N'
        elif general['xaxis'] == 'r':
            sizes = self.M_sizes
            sizes_label = 'M'
        else:
            print 'Warning: expecting \'xaxis\' to be \'r\' or \'c\'.'

        if requests is None:
            requests = self.results.keys()
        num_plots = len(requests)

        f, axarr = plt.subplots(num_plots, sharex = True)
        if num_plots == 1:
            axarr = [axarr]
            
        if self.title:
            axarr[0].set_title(self.title)

        for i, request in enumerate(requests):
            ax = axarr[i]

            if type(request) != str:
                names, options = request
            else:
                names, options = request, {}

            if type(names) == str:
                names = [names, names]
            plot_name, names = names[0], names[1:]
            for name in names:
                result = self.results[name]
                data = result['data']
                if 'plot_mean' in options and options['plot_mean']:
                    ax.plot(sizes, data.mean(1))
                else:
                    for rep in range(data.shape[1]):
                        ax.scatter(sizes, data[:,rep])

            if 'baseline' in options:
                ax.plot(sizes, np.repeat(options['baseline'], len(sizes)),
                        'k:')

            ax.set_ylabel(plot_name)
            if 'ymin' in options and 'ymax' in options:
                ax.set_ylim(options['ymin'], options['ymax'])
            elif 'ymin' in options:
                ax.set_ylim(ymin = options['ymin'])
            elif 'ymax' in options:
                ax.set_ylim(ymax = options['ymax'])

            if 'loglog' in options:
                ax.set_xscale('log')
                ax.set_yscale('log')
            elif 'semilogx' in options:
                ax.set_xscale('log')
            elif 'semilogy' in options:
                ax.set_yscale('log')

        axarr[-1].set_xlabel(sizes_label)
        f.subplots_adjust(hspace = 0)

        if self.interactive:
            plt.show() def add_array_stats(results, network = False):
    results.new('Density', 'a',
                lambda a: 1.0 * np.sum(a) / (a.shape[0] * a.shape[1]))
    results.new('Average row-sum', 'a',
                lambda a: 1.0 * np.sum(a) / a.shape[0])
    results.new('Average col-sum', 'a',
                lambda a: 1.0 * np.sum(a) / a.shape[1])
    results.new('Max row-sum', 'a', lambda a: np.max(a.sum(1)))
    results.new('Min row-sum', 'a', lambda a: np.min(a.sum(1)))
    results.new('Max col-sum', 'a', lambda a: np.max(a.sum(0)))
    results.new('Min col-sum', 'a', lambda a: np.min(a.sum(0)))
    if network:
        results.new('Self-loop density', 'a',
                    lambda a: np.mean(np.diagonal(a))) def rel_mse(est_1, est_2, truth):
    return robust_mse(est_1, truth) / robust_mse(est_2, truth) 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 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 sys
import json
import argparse
import numpy as np
import matplotlib.backends.backend_pdf as pltpdf
                    import scipy.io
                    import scipy.io
from Network import Network
from Models import StationaryLogistic, NonstationaryLogistic
from Models import alpha_zero, alpha_norm, alpha_unif, alpha_gamma
from Experiment import RandomSubnetworks, Seed, \
     Results, add_array_stats, rel_mse
from BinaryMatrix import approximate_conditional_nll as acnll
from Utility import l2, logit, pick, unpick
        from scipy.stats import chi2
    if params['alpha_norm_sd'] > 0.0:
        alpha_norm(arr, params['alpha_norm_sd'])
    elif params['alpha_unif_sd'] > 0.0:
        alpha_unif(arr, params['alpha_unif_sd'])
    elif params['alpha_gamma_sd'] > 0.0:
        # Choosing location somewhat arbitrarily to give unit skewness
        alpha_gamma(arr, 4.0, params['alpha_gamma_sd'])
    else:
        alpha_zero(arr)

    # Generate covariates and associated coefficients
    data_model = NonstationaryLogistic()
    covariates = []
    for b in range(params['B']):
        name = 'x_%d' % b
        covariates.append(name)

        if name in params['theta_fixed']:
            data_model.beta[name] = params['theta_fixed'][name]
        else:
            data_model.beta[name] = np.random.normal(0, params['theta_sd'])

        if params['cov_unif_sd'] > 0.0:
            c = np.sqrt(12) / 2
            def f_x(i_1, i_2):
                return np.random.uniform(-c * params['cov_unif_sd'],
                                         c * params['cov_unif_sd'])
        elif params['cov_norm_sd'] > 0.0:
            def f_x(i_1, i_2):
                return np.random.normal(0, params['cov_norm_sd'])
        elif params['cov_disc_sd'] > 0.0:
            def f_x(i_1, i_2):
                return (params['cov_disc_sd'] *
                        (np.sign(np.random.random() - 0.5)))
        else:
            print 'Error: no covariate distribution specified.'
            sys.exit()

        arr.new_edge_covariate(name).from_binary_function_ind(f_x)

    # Generate large network, if necessary
    if not params['sampling'] == 'new':
        data_model.match_kappa(arr, params['kappa_target'])
        arr.generate(data_model)

    if params['fit_nonstationary']:
        fit_model = NonstationaryLogistic()
    else:
        fit_model = StationaryLogistic()
    for c in covariates:
        fit_model.beta[c] = None

    # Set up recording of results from experiment
    results = Results(params['sub_sizes_r'], params['sub_sizes_c'],
                      params['num_reps'], interactive = params['interactive'])
    add_array_stats(results)
    if params['plot_sig']:
        crit = lambda dof: -0.5 * chi2.ppf(0.95, dof)

        umle_f = lambda n, f: f.nll(n, ignore_offset = True)
        umle_d = lambda n, d: d.nll(n, ignore_offset = True)
        umle_n = lambda n: NonstationaryLogistic().nll(n, ignore_offset = True)
        results.new('UMLE F-N', 'nm',
                    lambda n, d, f: umle_f(n, f) - umle_n(n))
        results.new('UMLE F-D', 'nm',
                    lambda n, d, f: umle_f(n, f) - umle_d(n, d))

        cmle_a_f = lambda n, f: acnll(n.as_dense(), np.exp(f.edge_probabilities(n)))
        cmle_a_d = lambda n, d: acnll(n.as_dense(), np.exp(d.edge_probabilities(n)))
        cmle_a_n = lambda n: acnll(n.as_dense(), np.ones_like(n.as_dense()))
        results.new('CMLE-A F-N', 'nm',
                    lambda n, d, f: cmle_a_f(n, f) - cmle_a_n(n))
        results.new('CMLE-A F-D', 'nm',
                    lambda n, d, f: cmle_a_f(n, f) - cmle_a_d(n, d))

        cmle_is_f = lambda n, f: f.fit_conditional(n, evaluate = True, T = 50)
        cmle_is_d = lambda n, d: d.fit_conditional(n, evaluate = True, T = 50)
        cmle_is_n = lambda n: NonstationaryLogistic().fit_conditional(n, evaluate = True, T = 50)
        results.new('CMLE-IS F-N', 'nm',
                    lambda n, d, f: cmle_is_f(n, f) - cmle_is_n(n))
        results.new('CMLE-IS F-D', 'nm',
                    lambda n, d, f: cmle_is_f(n, f) - cmle_is_d(n, d))

        c_cmle_f = lambda n, f: f.fit_c_conditional(n, evaluate = True)
        c_cmle_d = lambda n, d: d.fit_c_conditional(n, evaluate = True)
        c_cmle_n = lambda n: NonstationaryLogistic().fit_c_conditional(n, evaluate = True)
        results.new('C-CMLE F-N', 'nm',
                    lambda n, d, f: c_cmle_f(n, f) - c_cmle_n(n))
        results.new('C-CMLE F-D', 'nm',
                    lambda n, d, f: c_cmle_f(n, f) - c_cmle_d(n, d))

        results.new('UMLE sig.', 'dof',
                    lambda M, N, B: crit((M - 1) + (N - 1) + 1 + B))
        results.new('CMLE sig.', 'dof', lambda M, N, B: crit(B))
        results.new('C-CMLE sig.', 'dof', lambda M, N, B: crit((M - 1) + B))

    if params['sampling'] == 'new':
        results.new('Subnetwork kappa', 'm', lambda d, f: d.kappa)
    def true_est_theta_c(c):
        return (lambda d, f: d.beta[c]), (lambda d, f: f.beta[c])
    for c in covariates:
        # Need to do this hackily to avoid for-loop/lambda-binding weirdness.
        f_true, f_est = true_est_theta_c(c)
        results.new('True theta_{%s}' % c, 'm', f_true)
        results.new('Est. theta_{%s}' % c, 'm', f_est)
    if params['pre_offset'] or params['post_fit']:
        results.new('# Active', 'n',
                    lambda n: np.isfinite(n.offset.matrix()).sum())
    else:
        results.new('# Active', 'n', lambda n: n.M * n.N)
    if params['fisher_information']:
        def info_theta_c(c):
            def f_info_theta_c(d, f):
                return d.I_inv['theta_{%s}' % c]
            return f_info_theta_c
        for c in covariates:
            results.new('Info theta_{%s}' % c, 'm', info_theta_c(c))
    if params['baseline']:
        def rel_mse_p_ij(n, d, f):
            P = d.edge_probabilities(n)

            return rel_mse(f.edge_probabilities(n), f.baseline(n), P)