<jupyter_start><jupyter_text>Create Causal Graph node<jupyter_code># create generic discrete probability function
class cg_node():
    def __init__(self,n_inputs,name):
        
        self.n_inputs = n_inputs
        self.name = name
        
        if n_inputs == 0:
            self.label = 'exogenous'
        else:
            self.label = 'endogenous'
            
        return
    
    def p_init(self,input_data,var_data):
        
        self.n_data = len(input_data)
        
        self.input_data = input_data
        self.var_data = var_data
        
        if self.n_inputs == 0:
            p_ave = np.zeros(3)
            n_count = self.n_data
            for i in range(0,3):
                p_ave[i] = np.sum(var_data == i-1)/n_count
        
        elif self.n_inputs == 1:
            n_count = np.zeros(3)
            p_ave = np.zeros((3,3))
            
            for i in range(0,3):
                n_count[i] = np.sum(input_data == i-1)
                for j in range(0,3):
                    p_ave[j,i] = np.sum((input_data[:,0] == i-1)*(var_data == j-1))/n_count[i]
            
            
        elif self.n_inputs == 2:
            n_count = np.zeros((3,3))
            p_ave = np.zeros((3,3,3))
            
            for i in range(0,3):
                for j in range(0,3):
                    n_count[i,j] = np.sum((input_data[:,0] == i-1)*(input_data[:,1] == j-1))
                    for k in range(0,3):
                        p_ave[k,i,j] = np.sum(
                            (input_data[:,0] == i-1)*(input_data[:,1] == j-1)*(var_data == k-1))/n_count[i,j]
                        
        elif self.n_inputs == 3:
            n_count = np.zeros((3,3,3))
            p_ave = np.zeros((3,3,3,3))
            
            for i in range(0,3):
                for j in range(0,3):
                    for k in range(0,3):
                        n_count[i,j,k] = np.sum(
                            (input_data[:,0] == i-1)*(input_data[:,1] == j-1)*(input_data[:,2] == k-1))
                        for m in range(0,3):
                            p_ave[m,i,j,k] = np.sum((input_data[:,0] == i-1)*(input_data[:,1] == j-1)
                                *(input_data[:,2] == k-1)*(var_data == m-1))/n_count[i,j,k]
                            
        elif self.n_inputs == 4:
            n_count = np.zeros((3,3,3,3))
            p_ave = np.zeros((3,3,3,3,3))
            
            for i in range(0,3):
                for j in range(0,3):
                    for k in range(0,3):
                        for m in range(0,3):
                            n_count[i,j,k,m] = np.sum((input_data[:,0] == i-1)*(input_data[:,1] == j-1)
                                *(input_data[:,2] == k-1)*(input_data[:,3] == m-1))
                            for q in range(0,3):
                                p_ave[q,i,j,k,m] = np.sum((input_data[:,0] == i-1)*(input_data[:,1] == j-1)
                                    *(input_data[:,2] == k-1)*(input_data[:,3] == m-1)
                                    *(var_data == q-1))/n_count[i,j,k,m]
                        
            
        else:
            print('error -- too many inputs')
            return
            
        self.n_count = torch.tensor(n_count/self.n_data)
        self.prob_dist = torch.tensor(p_ave)
        
        return
    
    def sample(self,data_in=[]):
        
        if self.n_inputs == 0:
            samp_out = pyro.sample('',pyro.distributions.Multinomial(probs = self.prob_dist)).bool()
        elif self.n_inputs == 1:
            p_temp = torch.squeeze(self.prob_dist[:,data_in[0]])
            samp_out = pyro.sample('',pyro.distributions.Multinomial(probs = p_temp)).bool()
        elif self.n_inputs == 2:
            p_temp = torch.squeeze(self.prob_dist[:,data_in[0],data_in[1]])
            samp_out = pyro.sample('',pyro.distributions.Multinomial(probs = p_temp)).bool()
        elif self.n_inputs == 3:
            p_temp = torch.squeeze(self.prob_dist[:,data_in[0],data_in[1],data_in[2]])
            samp_out = pyro.sample('',pyro.distributions.Multinomial(probs = p_temp)).bool()
        else:
            print('error -- too many inputs')
            samp_out = []
        
        return samp_out<jupyter_output><empty_output><jupyter_text>Causal graph class<jupyter_code>class cg_graph():
    
    def __init__(self,str_list=[],bel_graph=[]):
        
        edge_list = []

        entity_list = []
        
        if str_list:

            for item in str_list:

                sub_ind = item.find('=')

                sub_temp = item[:sub_ind-1]
                obj_temp = item[sub_ind+3:]
                
                rel_temp = item[sub_ind:sub_ind+2]

                if sub_temp not in entity_list:
                    entity_list.append(sub_temp)
                if obj_temp not in entity_list:
                    entity_list.append(obj_temp)
                
                # ignore hasVariant, partOf relations
                
                if rel_temp.find('crease') > 0:
                    edge_list.append([sub_temp,obj_temp,rel_temp])
                
                # check for duplicate edges
                #nodes_temp = [sub_temp,obj_temp]
                #list_temp = [[item[0],item[1]] for item in edge_list]
                #if nodes_temp in list_temp:
                    #ind_temp = list_temp.index(nodes_temp)
                    #edge_list[ind_temp][2] += ',' + rel_temp
                #else:
                    #edge_list.append([sub_temp,obj_temp,rel_temp])
                
        elif bel_graph:
            
            for item in bel_graph.edges:
                edge_temp = bel_graph.get_edge_data(item[0],item[1],item[2])
                sub_temp = str(item[0]).replace('"','')
                obj_temp = str(item[1]).replace('"','')
                rel_temp = edge_temp['relation']
                
                if sub_temp not in entity_list:
                    entity_list.append(sub_temp)
                if obj_temp not in entity_list:
                    entity_list.append(obj_temp)
                
                # ignore hasVariant, partOf relations
                
                if rel_temp.find('crease') > 0:
                    edge_list.append([sub_temp,obj_temp,rel_temp])
                
                # check for duplicate edges
                #nodes_temp = [sub_temp,obj_temp]
                #list_temp = [[item[0],item[1]] for item in edge_list]
                #if nodes_temp in list_temp:
                    #ind_temp = list_temp.index(nodes_temp)
                    #edge_list[ind_temp][2] += ',' + rel_temp
                #else:
                    #edge_list.append([sub_temp,obj_temp,rel_temp])
                

        n_nodes = len(entity_list)
        self.n_nodes = n_nodes

        adj_mat = np.zeros((n_nodes,n_nodes),dtype=int)

        for item in edge_list:
            out_ind = entity_list.index(item[0])
            in_ind = entity_list.index(item[1])
            adj_mat[out_ind,in_ind] = 1
            
        self.edge_list = edge_list
        self.entity_list = entity_list
        self.adj_mat = adj_mat
        
        self.graph = nx.DiGraph(adj_mat)
        
        # check to make sure that it's a DAG
        if nx.algorithms.dag.is_directed_acyclic_graph(self.graph):
            print('The causal graph is a acyclic')
        else:
            print('The causal graph has cycles -- this is a problem')
        
        node_dict = {}
        
        for i in range(0,n_nodes):
            node_dict[entity_list[i]] = cg_node(np.sum(adj_mat[:,i]),entity_list[i])
        
        self.node_dict = node_dict
        
        self.cond_list = []
        
        self.sample_dict = {}
        
        self.parent_ind_list = []
        self.child_ind_list = []
        self.parent_name_dict = {}
        self.child_name_dict = {}
        
        self.parent_ind_list = [np.where(self.adj_mat[:,i] > 0)[0] for i in range(0,n_nodes)]
        self.child_ind_list = [np.where(self.adj_mat[i,:] > 0)[0] for i in range(0,n_nodes)]
        
        for i in range(0,n_nodes):
            self.parent_name_dict[entity_list[i]] = [entity_list[item] for item in self.parent_ind_list[i]]
            self.child_name_dict[entity_list[i]] = [entity_list[item] for item in self.child_ind_list[i]]

        # create rank-3 delta tensor
        tensor_temp = torch.zeros((3,3,3)).double()
        for i in range(0,3):
            tensor_temp[i,i,i] = 1
            
        self.tensor_temp = tensor_temp
        return
    
    
    def prob_init(self,data_in):
        # initialize all of the nodes
        
        exog_list = []
        prob_dict = {}
        
        for name in self.node_dict:
            i = self.entity_list.index(name)
            data_in_temp = data_in[:,self.parent_ind_list[i]]
            data_out_temp = data_in[:,i]
            
            self.node_dict[name].p_init(data_in_temp,data_out_temp)
            
            if self.node_dict[name].n_inputs == 0:
                exog_list.append(name)
            prob_dict[name] = self.node_dict[name].prob_dist
        
        self.exog_list = exog_list
        self.prob_dict = prob_dict

        return
        
        
    def sample_vars(self,names,flag=0):
        # do a multi-variable sample
        
        # sample only those variables w/o sample data
        if np.any([item in self.sample_dict for item in names]):
            sample_list = [item for item in names if item not in self.sample_dict]
            if sample_list:
                self.sample_vars(sample_list,flag+1)
        
        # sample exogenous variables
        elif np.any([item in self.exog_list for item in names]):
            in_exog = [item for item in names if item in self.exog_list]
            not_in_exog = [item for item in names if item not in self.exog_list]
            
            for item in in_exog:
                self.sample_dict[item] = self.node_dict[item].sample()
            if not_in_exog:              
                self.sample_vars(not_in_exog,flag+1)
            
        # if you have samples from all of the parents, sample names
        # otherwise, sample the parents
        elif names:
            
            sample_list = []
            sample_list2 = []
            
            for item in names:
                parent_list = self.parent_name_dict[item]
                
                if np.all([item2 in self.sample_dict for item2 in parent_list]):
                    self.sample_dict[item] = self.node_dict[item].sample(
                        [self.sample_dict[item2] for item2 in parent_list])
                else:
                    sample_list = sample_list + [item2 for item2 in parent_list if item2 not in sample_list]
                    sample_list2 = sample_list2 + [item]
            if sample_list:
                self.sample_vars(sample_list,flag+1)
            if sample_list2:
                self.sample_vars(sample_list2,flag+1)
            
        # if you're back at the root node, return the samples
        # otherwise, don't return anything - the values are stored in self.sample_dict
        if flag == 0:
            
            tensor_sample = torch.Tensor([-1,0,1]).int()
            
            output =[tensor_sample[self.sample_dict[item]] for item in names]
            self.sample_dict = {}
            return output
        else:
            return
            

    def gen_path_nodes(self,sources,destinations):
        
        source_inds = [self.entity_list.index(item) for item in sources]
        dest_inds = [self.entity_list.index(item) for item in destinations]
        
        nodes = []
        
        for i in source_inds:
            for j in dest_inds:      
                for path in nx.all_simple_paths(self.graph, source=i, target=j):
                    for ind in path:
                        if self.entity_list[ind] not in nodes:
                            nodes.append(self.entity_list[ind])
                            
        return nodes
        
    
    def joint_dist_add(self,add_node,nodes_temp,prob_temp):
        # create a new joint distribution with add_node now included
        
        # deliberately let out 'a' - for the variable being added
        str_temp = 'bcdefghijklmnopqrstuvwxyz'
        
        # find parent indices
        par_inds = [nodes_temp.index(item2) for item2 in self.parent_name_dict[add_node]]

        n_inds = len(par_inds)

        if n_inds == 1:
            str1 = 'ay,' + str_temp[par_inds[0]] + 'yz,'
            str2 = str_temp[:len(nodes_temp)].replace(str_temp[par_inds[0]],'z')

            str_sum = str1 + str2
            prob_out = torch.einsum(str_sum,self.prob_dict[add_node],self.tensor_temp,prob_temp)

        elif n_inds == 2:                        
            str1 = 'awy,' + str_temp[par_inds[0]] + 'wx,' + str_temp[par_inds[1]] + 'yz,'

            str2 = str_temp[:len(nodes_temp)].replace(
                str_temp[par_inds[0]],'x').replace(str_temp[par_inds[1]],'z')

            str_sum = str1 + str2
            prob_out = torch.einsum(str_sum,self.prob_dict[add_node],self.tensor_temp,self.tensor_temp,prob_temp)

        elif n_inds == 3:
            str1 = ('auwy,' + str_temp[par_inds[0]] + 'uv,' + str_temp[par_inds[1]] 
                    + 'wx,' + str_temp[par_inds[2]] + 'yz,')

            str2 = str_temp[:len(nodes_temp)].replace(
                str_temp[par_inds[0]],'u').replace(str_temp[par_inds[1]],'x').replace(
                str_temp[par_inds[2]],'z')

            str_sum = str1 + str2
            prob_out = torch.einsum(str_sum,
                self.prob_dict[add_node],self.tensor_temp,self.tensor_temp,self.tensor_temp,prob_temp)
            
        else:
            print('too many parents')
            prob_out = prob_temp
            
        return prob_out
        
    
    def calc_prob(self,names):
        # calculate the joint probability over a list of named nodes
        
        # find all paths from exogenous nodes        
        path_nodes = self.gen_path_nodes(self.exog_list,names)
        
        for item in names:
            if item in self.exog_list and item not in path_nodes:
                path_nodes.append(item)
        
        print(path_nodes)

        # get joint exogenous probability distribution
        nodes_temp = []
        for item in self.exog_list:
            if item in path_nodes:
                nodes_temp.append(item)
                
        #print(nodes_temp)
        #print(self.exog_list)
        prob_temp = self.prob_dict[nodes_temp[0]]
        for item in nodes_temp[1:]:
            prob_temp = torch.einsum('...i,j',prob_temp,self.prob_dict[item])
        
        # identify all of the children of nodes_temp in path_nodes
        child_nodes = []
        for item in nodes_temp:
            for item2 in self.child_name_dict[item]:
                if item2 in path_nodes and item2 not in child_nodes:
                    child_nodes.append(item2)
            
        flag = 0
        
        # iterate through node children until target nodes are reached
        while flag == 0:
            #print(nodes_temp)
            
            # determine which nodes to add
            # all children of nodes_temp in path_nodes, not in nodes_temp, and that have all their parents in
            # nodes_temp
            add_nodes = []
            for item in nodes_temp:
                for item2 in self.child_name_dict[item]:
                    if (item2 in path_nodes
                        and item2 not in nodes_temp
                        and np.all([item3 in nodes_temp for item3 in self.parent_name_dict[item2]])
                        and item2 not in add_nodes):
                        add_nodes.append(item2)
            
            #print(add_nodes)

            # add nodes to the joint distribution
            for item in add_nodes:
                prob_temp = self.joint_dist_add(item,nodes_temp,prob_temp)

                # add the new node to nodes_temp
                nodes_temp = [item] + nodes_temp
                
            # determine which nodes to subtract
            # all nodes in nodes_temp not in names and that have all their children in nodes_temp
            sub_nodes = []
            for item in nodes_temp:
                
                child_path_list = []
                
                for item2 in self.child_name_dict[item]:
                    if item2 in path_nodes:
                        child_path_list.append(item2)

                if item not in names and np.all([item2 in nodes_temp for item2 in child_path_list]):
                    sub_nodes.append(item)
                    
            #print(sub_nodes)
                    
            # sum over the sub_nodes probabilities
            
            if sub_nodes:
                remove_indices = [nodes_temp.index(item) for item in sub_nodes]
                prob_temp = torch.sum(prob_temp,dim=remove_indices)
            
                # remove summed nodes from nodes_temp
                for item in sub_nodes:
                    nodes_temp.remove(item)
                    
            if sorted(nodes_temp) == sorted(names):
                flag = 1
                permute_inds = [nodes_temp.index(item) for item in names] 
                prob_temp = prob_temp.permute(permute_inds)

            #print()
            
            
        return prob_temp
    
    
    def calc_cond_prob(self,cond_prob,uncond_prob):
        
        # check to make sure the lists don't overlap
        if np.any([item in uncond_prob for item in cond_prob]):
            print('error -- overlapping lists')
            return
        
        n_cond = len(cond_prob)
        n_uncond = len(cond_prob)
        
        p_joint = self.calc_prob(cond_prob + uncond_prob)
        p_uncond = self.calc_prob(uncond_prob)
                
        if n_uncond == 1:
            p_cond = torch.einsum('...i,ijk,j',p_joint,self.tensor_temp,1/p_uncond)
        elif n_uncond == 2:
            p_cond = torch.einsum('...il,ijk,lmn,jm',p_joint,self.tensor_temp,self.tensor_temp,1/p_uncond)
        elif n_uncond == 3:
            p_cond = torch.einsum('...ilp,ijk,lmn,pqr,jmq',
                p_joint,self.tensor_temp,self.tensor_temp,self.tensor_temp,1/p_uncond)
        else:
            print('too many conditioned variables')
            p_cond = p_joint
            
        return p_cond
            
    
    def calc_do(self,names,do_vars,do_vals):
        # calculate the final probability distribution of the variable in question given do variables
        
        # add do_vars to list of exogenous variables
        self.exog_list += do_vars
        
        # sever links from do_vars to parents
        child_temp = self.child_name_dict
        
        for item in self.child_name_dict:
            for item2 in do_vars:
                if item2 in self.child_name_dict[item]:
                    self.child_name_dict[item].remove(item2)
        
        names_prob_dict = {}
        for item in do_vars:
            names_prob_dict[item] = self.prob_dict[item]
        
        # specify distributions for those do_vars
        for item in do_vars:
            p_temp = np.zeros(3)
            p_temp[do_vals[item]+1] = 1
            self.prob_dict[item] = torch.Tensor(p_temp)
            
        
        prob_out = self.calc_prob(names)
        
        # restore original list of exogenous variables
        for item in do_vars:
            self.exog_list.remove(item)
        
        # restore original probability distributions
        for item in do_vars:
            self.prob_dict[item] = names_prob_dict[item]
        
        # restore original child dictionary
        self.child_name_dict = child_temp        
        
        return prob_out
    
        
    def calc_do_cond(self,do_vars,do_vals,cond_prob,uncond_prob):
        
        # check to make sure the lists don't overlap
        if (np.any([item in uncond_prob for item in cond_prob]) 
            or np.any([item in do_vars for item in cond_prob])
            or np.any([item in do_vars for item in uncond_prob])):
            print('error -- overlapping lists')
            return
        
        n_cond = len(cond_prob)
        n_uncond = len(cond_prob)
            
        
        p_joint = self.calc_do(cond_prob + uncond_prob,do_vars,do_vals)
        p_uncond = self.calc_do(uncond_prob,do_vars,do_vals)
                
        if n_uncond == 1:
            p_cond = torch.einsum('...i,ijk,j',p_joint,self.tensor_temp,1/p_uncond)
        elif n_uncond == 2:
            p_cond = torch.einsum('...il,ijk,lmn,jm',p_joint,self.tensor_temp,self.tensor_temp,1/p_uncond)
        elif n_uncond == 3:
            p_cond = torch.einsum('...ilp,ijk,lmn,pqr,jmq',
                p_joint,self.tensor_temp,self.tensor_temp,self.tensor_temp,1/p_uncond)
        else:
            print('too many conditioned variables')
            p_cond = p_joint
        
        return p_cond
            
            
            
    
    def calc_counterfact():
        return
    
    def calc_cde(self,names,do_vars,do_vals,ctrl_vars,ctrl_vals):
        tot_vars = do_vars + ctrl_vars
        tot_vals = {}
        tot_ctrl_vals = {}
        for item in do_vars:
            tot_vals[item] = do_vals[item]
            tot_ctrl_vals = ctrl_vals[item]
            
        for item in ctrl_vars:
            tot_vals[item] = 0
            tot_ctrl_vals[item] = ctrl_vals[item]
            
        return self.calc_do(names,tot_vars,tot_vals) - self.calc_do(names,tot_vars,tot_ctrl_vals)
        
        return
    
    def calc_te(self,names,do_vars,do_vals):
        
        do_vals_0 = {}
        for item in do_vals:
            do_vals_0[item] = 0
            
        return self.calc_do(names,do_vars,do_vals) - self.calc_do(do_vars,do_vals_0)
    
    def calc_nde(self,names,do_vars,do_vals):
        
        # identify parents of names
        parent_list = []
        for item in names:
            for item2 in self.parent_name_dict[item]:
                if item2 not in parent_list:
                    parent_list.append(item2)
                    
        par_do_vals = []
        for item in do_vars:
            if item in parent_list:
                par_do_vals.append(item)
        
        if not par_do_vals:
            print('no direct effect')
            prob_out = 0
        
        else:
            # calculate probability of non-do_var parents given do_vars = 0
            
            non_do_parents = [item for item in parent_list if item not in do_vars]
            do_vals_0 = {}
            for item in do_vars:
                do_vals_0[item] = 0
            
            non_par_do_vars = []
            for item in do_vars:
                if item not in parent_list:
                    non_par_do_vars.append(item)
            
            nodes_temp = non_do_parents
            
            prob_temp = self.calc_do(non_do_parents,do_vars,do_vals_0)
            # do outer product to get overall distribution
            for item in non_par_do_vars:
                p_temp = np.zeros(3)
                p_temp[do_vals[item]] = 1
                p_add = torch.Tensor(p_temp)
                prob_temp = torch.einsum('...i,j',prob_temp,p_add)
                nodes_temp.append(item)
            
            n_sum = len(nodes_temp)
            n_names = len(names)
            for item in names:
                prob_temp = self.joint_dist_add(item,nodes_temp,prob_temp)
                
            # calculate overall joint probability distribution
            sum_axes = range(n_names,n_names+n_sum)
            prob_do_vals = np.sum(prob_temp,axis=tuple(sum_axes))
            
            prob_out = prob_do_vals - self.calc_do(names,do_vars,do_vals_0)

        return prob_out
    
    def calc_nie():
        # identify parents of names
        parent_list = []
        for item in names:
            for item2 in self.parent_name_dict[item]:
                if item2 not in parent_list:
                    parent_list.append(item2)
                    
        par_do_vals = []
        for item in parent_list:
            if item in do_vars:
                par_do_vals.append(item)
        
        if np.all([item in do_vars for item in parent_list]):
            print('no indirect effect')
            prob_out = 0
        
        else:
            # calculate probability of non-do_var parents given do_vars = do_vals
            
            non_do_parents = [item for item in parent_list if item not in do_vars]
            do_vals_0 = {}
            for item in do_vars:
                do_vals_0[item] = 0
            
            non_par_do_vars = []
            for item in do_vars:
                if item not in parent_list:
                    non_par_do_vars.append(item)
            
            nodes_temp = non_do_parents
            
            prob_temp = self.calc_do(non_do_parents,do_vars,do_vals)
            # do outer product to get overall distribution
            for item in non_par_do_vars:
                p_temp = np.zeros(3)
                p_temp[1] = 1
                p_add = torch.Tensor(p_temp)
                prob_temp = torch.einsum('...i,j',prob_temp,p_add)
                nodes_temp.append(item)
            
            n_sum = len(nodes_temp)
            n_names = len(names)
            for item in names:
                prob_temp = self.joint_dist_add(item,nodes_temp,prob_temp)
                
            # calculate overall joint probability distribution
            sum_axes = range(n_names,n_names+n_sum)
            prob_do_vals = np.sum(prob_temp,axis=tuple(sum_axes))
            
            prob_out = prob_do_vals - self.calc_do(names,do_vars,do_vals_0)

        return prob_out

    def cond_mut_info(self,target,test,cond,data_in):
        
        cond_temp = cond
        
        if not cond:
            # find parents of target
            for item in target:
                for item2 in self.parent_name_dict[item]:
                    if item2 not in cond_temp:
                        cond_temp.append(item2)
        
        
        target_inds = [self.entity_list.index(item) for item in target]
        test_inds = [self.entity_list.index(item) for item in test]
        cond_inds = [self.entity_list.index(item) for item in cond_temp]
        
        n_total = len(data_in)
        
        null_joint = data_in[:,target_inds + cond_inds]
        null_cond = data_in[:,cond_inds]
        
        hypth_joint = data_in[:,target_inds + test_inds + cond_inds]
        hypth_cond = data_in[:,test_inds + cond_inds]
        
        null_entropy = 0
        null_list = []
        
        hypth_entropy = 0
        hypth_list = []
        for i in range(0,n_total):

            if np.all([np.any(null_joint[i,:] != item) for item in null_list]):
                num_sum = np.sum([np.all(null_joint[i,:] == item) for item in null_joint])
                denom_sum = np.sum([np.all(null_cond[i,:] == item) for item in null_cond])
                null_entropy = null_entropy - num_sum*np.log(num_sum/denom_sum)
                null_list.append(null_joint[i,:])
            
            if np.all([np.any(hypth_joint[i,:] != item) for item in hypth_list]):
                num_sum = np.sum([np.all(hypth_joint[i,:] == item) for item in hypth_joint])
                denom_sum = np.sum([np.all(hypth_cond[i,:] == item) for item in hypth_cond])
                hypth_entropy = hypth_entropy - num_sum*np.log(num_sum/denom_sum)
                hypth_list.append(hypth_joint[i,:])
                
        return (null_entropy - hypth_entropy)/n_total
        
    def g_test(self,name,data_in):
        # do the G-test on a single variable of interest
        
        p_name = self.calc_prob(name)*len(data_in)
        name_ind = self.entity_list.index(name[0])
        name_data = data_in[:,name_ind]
        
        p_data = torch.Tensor([np.sum(name_data == -1),np.sum(name_data == 0),np.sum(name_data == 1)])
        
        print(p_name)
        print(p_data)
        
        g_val = 2*torch.sum(p_data*torch.log(p_data/p_name))
        
        dof = len(data_in) - 1
        
        p_val = 1-sp.stats.chi2.cdf(g_val.item(), dof)
        
        return g_val,p_val
    
    def g_test_emp(self,name,data_in):
        # do the G-test on a single variable of interest
        
        #p_name = self.calc_prob(name)*len(data_in)
        # generate an empirical distribution for variable name
        model_data = np.zeros(len(data_in))
        for i in range(0,len(data_in)):
            model_data[i] = self.sample_vars(name)[0].item()
            
        p_model = torch.Tensor([np.sum(model_data == -1),np.sum(model_data == 0),np.sum(model_data == 1)])
        print(p_model)
        
        name_ind = self.entity_list.index(name[0])
        name_data = data_in[:,name_ind]
        
        p_data = torch.Tensor([np.sum(name_data == -1),np.sum(name_data == 0),np.sum(name_data == 1)])
        print(p_data)
        
        g_val = 2*torch.sum(p_data*torch.log(p_data/p_model))
        
        dof = len(data_in) - 1
        
        p_val = 1-sp.stats.chi2.cdf(g_val.item(), dof)
        
        return g_val,p_val
    
    def write_to_cf(self,filename,spacing):
        # write the causal graph to a text file to import into causal fusion
        
        pos_dict = nx.drawing.layout.planar_layout(self.graph)
        
        write_dict = {}
        write_dict['name'] = 'causal_graph'
        
        # write nodes
        write_dict['nodes'] = []
        for i in range(0,len(self.entity_list)):
            name = self.entity_list[i]
            
            write_dict['nodes'].append({})
            
            write_dict['nodes'][-1]['id'] = 'node' + str(i)
            write_dict['nodes'][-1]['name'] = name
            write_dict['nodes'][-1]['label'] = name
            write_dict['nodes'][-1]['type'] = 'basic'
            write_dict['nodes'][-1]['metadata'] = {}
            write_dict['nodes'][-1]['metadata']['x'] = spacing*pos_dict[i][0]
            write_dict['nodes'][-1]['metadata']['y'] = spacing*pos_dict[i][1]
            write_dict['nodes'][-1]['metadata']['label'] = ''
            write_dict['nodes'][-1]['metadata']['shape'] = 'ellipse'
            write_dict['nodes'][-1]['metadata']['fontSize'] = 14
            write_dict['nodes'][-1]['metadata']['sizeLabelMode'] = 5
            write_dict['nodes'][-1]['metadata']['font'] = {}
            write_dict['nodes'][-1]['metadata']['font']['size'] = 14
            write_dict['nodes'][-1]['metadata']['size'] = 14
            write_dict['nodes'][-1]['metadata']['labelNodeId'] = 'node' + str(i) + 'ID'
            write_dict['nodes'][-1]['metadata']['labelNodeOffset'] = {}
            write_dict['nodes'][-1]['metadata']['labelNodeOffset']['x'] = 0
            write_dict['nodes'][-1]['metadata']['labelNodeOffset']['y'] = 0
            write_dict['nodes'][-1]['metadata']['labelOffset'] = {}
            write_dict['nodes'][-1]['metadata']['labelOffset']['x'] = 0
            write_dict['nodes'][-1]['metadata']['labelOffset']['y'] = 0
            write_dict['nodes'][-1]['metadata']['shadow'] = {}
            write_dict['nodes'][-1]['metadata']['shadow']['color'] = '#00000080'
            write_dict['nodes'][-1]['metadata']['shadow']['size'] = 0
            write_dict['nodes'][-1]['metadata']['shadow']['x'] = 0
            write_dict['nodes'][-1]['metadata']['shadow']['y'] = 0
            
        # write edges
        write_dict['edges'] = []
        
        for i in range(0,len(self.edge_list)):
            
            item = self.edge_list[i]
            from_node = self.entity_list.index(item[0])
            to_node = self.entity_list.index(item[1])
            
            write_dict['edges'].append({})
            
            write_dict['edges'][-1]['id'] = 'node' + str(from_node) + '->node' + str(to_node)
            write_dict['edges'][-1]['from'] = item[0]
            write_dict['edges'][-1]['to'] = item[1]
            write_dict['edges'][-1]['type'] = 'directed'
            write_dict['edges'][-1]['metadata'] = {}
            write_dict['edges'][-1]['metadata']['isLabelDraggable'] = True
            write_dict['edges'][-1]['metadata']['label'] = ''
            
        
        write_dict['task'] = {}
        
        write_dict['metadata'] = {}
        
        write_dict['project_id'] = '123456789'
        write_dict['_fileType'] = 'graph'
                
        with open(filename + '.json', 'w') as json_file:
            json.dump(write_dict, json_file)<jupyter_output><empty_output><jupyter_text>Test out cg_graph<jupyter_code>bel_temp = pb.from_bel_script('sag_bel_graph.txt')
graph_test = cg_graph(bel_graph=bel_temp)
print(dir(graph_test))
print()
for item in graph_test.edge_list:
    print(item)
    print()
graph_test.write_to_cf('sag_graph',300)
str_list = ['temp =| cloudy','cloudy => rainy','temp => icream','rainy =| icream']
graph_test = cg_graph(str_list=str_list)
for item in graph_test.edge_list:
    print(item)
    
print()
for item in graph_test.entity_list:
    print(item)
for item in data:
    print(item)
    print(data[item])
    print()
print(data['metadata'])
print()
for item in data['metadata']:
    print(item)
    print(data['metadata'][item])
    print(type(item))
    print()
for item in graph_test.node_dict:
    print([graph_test.node_dict[item].name,graph_test.node_dict[item].n_inputs])
graph_test.prob_init(tot_data)
print(graph_test.exog_list)
for item in graph_test.node_dict:
    print(item)
    print(graph_test.node_dict[item].n_count)
    print(graph_test.node_dict[item].prob_dist)
    print()
graph_test.sample_vars(['icream','rainy'])
x = graph_test.calc_prob(['rainy','icream'])
y = graph_test.calc_prob(['icream'])
y2 = graph_test.calc_prob(['rainy'])
print(x)
print()
print(y)
print(torch.sum(x,dim=0))

print()
print(y2)
print(torch.sum(x,dim=1))

# this is somehow reversed!!
z = graph_test.calc_cond_prob(['rainy'],['icream'])
print()
print(z)
print()
print(torch.matmul(z,y))

print(torch.sum(z,dim=1))

print(torch.sum(z,dim=0))
do_dict = {}
do_dict['rainy'] = torch.Tensor([1]).int()
for item in do_dict:
    print(do_dict[item])
    print(graph_test.prob_dict[item])


print(graph_test.exog_list)
a1 = graph_test.calc_do(['icream'],['rainy'],do_dict)
print(a1)


graph_test.calc_do_cond(['rainy'],do_dict,['temp'],['icream'])
for item in graph_test.prob_dict:
    print(item)
    print(graph_test.prob_dict[item])
    print()
print(graph_test.cond_mut_info(['rainy'],['temp'],['cloudy'],tot_data))
print(graph_test.gen_path_nodes(graph_test.exog_list,['temp']))
print(graph_test.prob_dict['icream'])

a = graph_test.g_test(['icream'],tot_data)
print(a)
print(graph_test.prob_dict['cloudy'][:,graph_test.node_dict['temp'].sample()])
print(graph_test.prob_dict['cloudy'])
print(graph_test.calc_prob(['cloudy']))
print(torch.matmul(graph_test.prob_dict['cloudy'],graph_test.prob_dict['temp']))
print(torch.matmul(graph_test.prob_dict['temp'],graph_test.prob_dict['cloudy']))
print()

a = graph_test.g_test_emp(['icream'],tot_data)
print(a)
x = pyro.sample('',pyro.distributions.Multinomial(probs = torch.Tensor([0.3,0.2,0.5]))).bool()
#samp_out = pyro.sample('',pyro.distributions.Multinomial(probs = self.prob_dist))
print(x)
y = torch.Tensor([-1,0,1])[x]
print(y)

z = torch.Tensor([1]).int()
print(z)
print(z+1)
print(graph_test.calc_prob(['icream']))
print(dir(graph_test))
print()
print(graph_test.entity_list)
print(graph_test.adj_mat)
sp.stats?
def indep_vars(n_samples):
    
    T_list = []
    C_list = []
    P_list = []
    
    for i in range(0,n_samples):
        
        #x = pyro.sample("x_{}".format(i), pyro.distributions.Normal(20,5))
        
        #T_temp = pyro.distributions.Normal(20,5).sample()
        #C_temp = 0.5*pyro.distributions.Beta(1,1+T_temp/10).sample() + 0.5*pyro.distributions.Uniform(0,1).sample()
        #P_temp = (0.5*pyro.distributions.Exponential(1).sample() 
            #+ 0.5*pyro.distributions.Exponential(1/(C_temp+1)).sample())
        
        T_list.append(pyro.sample("T_{}".format(i), pyro.distributions.Normal(20,5)))
        
        C_list.append(0.5*pyro.sample("C1_{}".format(i),pyro.distributions.Beta(1,1+T_list[-1]/10)) 
            + 0.5*pyro.sample("C2_{}".format(i),pyro.distributions.Uniform(0,1)))
        P_list.append(0.5*pyro.sample("P1_{}".format(i), pyro.distributions.Exponential(1))
            + 0.5*pyro.sample("P2.{}".format(i),pyro.distributions.Exponential(1/(C_list[-1]+1))))
        
    return T_list,C_list,P_list

def dep_vars(T_list,C_list,P_list):
    
    n_pts = len(T_list)
    
    I_list = []
    
    for i in range(0,n_pts):
        
        T_temp = T_list[i]
        C_temp = C_list[i]
        P_temp = P_list[i]
        
        if P_temp > 2.5 or T_temp < 15:
            I_list.append(1e-6*pyro.sample("I_{}".format(i),pyro.distributions.Bernoulli(1)))
        else:
            I_list.append(pyro.sample("I_{}".format(i),
                pyro.distributions.Beta(2*(2.5-P_temp)*(T_temp-12)/(2.5*12),2)))
            #I_temp = torch.tensor(0.5)
        
    return I_list
n_data = 10000
temp,cloud,precip = indep_vars(n_data)
icream = dep_vars(temp,cloud,precip)
# trinarize data relative to baseline
T_base = 20
C_base = 0.38
P_base = 1.2
I_base = 0.23

T_sig = 1.0
C_sig = 0.02
P_sig = 0.06
I_sig = 0.01

n_count_tri = np.zeros((3,3,3))
p_ave_tri = np.zeros((3,3,3,3))

def cond_test(val,base,sig):
    
    conds = [val < base-sig,val > base-sig and val < base+sig,val > base+sig]
    
    return conds.index(True)-1
        
T_ind = []
C_ind = []
P_ind = []
I_ind = []
for ind in range(0,n_data):
    T_ind.append(cond_test(temp[ind],T_base,T_sig))
    C_ind.append(cond_test(cloud[ind],C_base,C_sig))
    P_ind.append(cond_test(precip[ind],P_base,P_sig))
    I_ind.append(cond_test(icream[ind],I_base,I_sig))
tot_data = np.asarray([T_ind,C_ind,P_ind,I_ind]).T
print(np.shape(tot_data))
print(tot_data[0:5,:])
print(type(tot_data))<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup