<jupyter_start><jupyter_text>Experimental setup- Fix settings:    - intervention sparsity, sample size, 2,...,E- Simulate DAG- Simulate intervention targets for each env- Simulate data from each env- Compute CPDAG (MEC) from DAG- Learn full pool oracle PDAG    - Compute TP, FP edges- Learn pairwise pool oracle PDAG - Compute TP, FP edges- Learn minimal changes result    - Compute TP, FP edges<jupyter_code>SPARSITY = 1
SAMPLE_SIZE = 100
N_ENV = 5
N_NODES = 4
DENSITY = 0.5<jupyter_output><empty_output><jupyter_text>Sample DAG<jupyter_code># Sample parent assignments
G_adj = erdos_renyi_dag(N_NODES, DENSITY)
plot_dag(G_adj)<jupyter_output><empty_output><jupyter_text>Sample Data<jupyter_code># Sample targets
intv_targets = [np.random.choice(N_NODES, SPARSITY, replace=False) for _ in range(N_ENV)]
# Environment data
base_seed = 1234
Xs = [sample_cdnod_sim(
        G_adj,
        SAMPLE_SIZE,
        intervention_targets=targets,
        base_random_state=base_seed,
    ) for targets in intv_targets]
df = pd.DataFrame(
    np.hstack((
        np.vstack(Xs),
        np.hstack([[i]*X.shape[0] for i, X in enumerate(Xs)]).reshape(-1, 1)
    )),
    columns=[f'X{i+1}' for i in range(Xs[0].shape[1])] + ['Env']
)

sns.pairplot(df, hue='Env', palette='Set1')
plt.suptitle(intv_targets, y=1.04)
plt.show()<jupyter_output><empty_output><jupyter_text>Oracle tests<jupyter_code>G_cl = create_causal_learn_dag(G_adj)
G_mec = dag2cpdag(G_cl)
G_2dag = pdag2dag(G_mec)
G_2dag.graph
G_mec.graph
plot_dag(G_adj)
G_adj
V_adj = np.asarray([[0, 0, 0], [1, 0, 0], [1, 1, 0]])
plot_dag(V_adj)
V_cl = create_causal_learn_dag(V_adj)
V_mec = dag2cpdag(V_cl)
V_2dag = pdag2dag(V_mec)
V_mec.graph
V_2dag.graph<jupyter_output><empty_output><jupyter_text>Augmented PC<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
import pandas as pd

from sparse_shift.datasets import sample_cdnod_sim, \
    erdos_renyi_dag, connected_erdos_renyi_dag, barabasi_albert_dag
from sparse_shift.plotting import plot_dag
from sparse_shift.utils import create_causal_learn_dag
from sparse_shift.utils import dag2cpdag, cpdag2dags

# from causallearn.utils.PCUtils.BackgroundKnowledge import BackgroundKnowledge
# # from causallearn.search.ConstraintBased.PC import pc
# from causallearn.utils.cit import fisherz, kci
# from sparse_shift.causal_learn.PC import augmented_pc
# from sparse_shift.causal_learn.SkeletonDiscovery import augmented_skeleton_discovery
# from causallearn.utils.PCUtils import Meek
# from sparse_shift.causal_learn.GraphClass import CausalGraph
# from causallearn.graph.GraphNode import GraphNode
# from causallearn.utils.DAG2CPDAG import dag2cpdag
# from causallearn.utils.PDAG2DAG import pdag2dag
# def _run_recursive_test(data, cpdag, alpha, test, test_kwargs):
#     from itertools import combinations

#     n_vars = cpdag.shape[0]
#     edges = set(list(range(n_vars)))
#     depth = -1
#     while np.max(cp.dag.sum(1)) - 1 > depth:
#         depth += 1
        
#         for x in range(n_vars):
#             neigh_x = np.where(cpdag[:, i] != 0)[0]
            
#             for S in combinations(neigh_x, depth):
#                 # n_vars is E
#                 p = test(data, x, n_vars, S, **test_kwargs)
#                 if p > alpha:
#                     edges.remove(x)
#                     break

#     return list(edges)


def _construct_augmented_cpdag(cpdag):
    from sparse_shift.utils import create_causal_learn_cpdag
    from sparse_shift.causal_learn.GraphClass import CausalGraph
    from causallearn.graph.GraphNode import GraphNode
    from causallearn.graph.Edge import Edge
    from causallearn.graph.Endpoint import Endpoint

    n_x_vars = cpdag.shape[0]
    cl_cpdag = create_causal_learn_cpdag(cpdag)
    cl_cpdag.add_node(GraphNode(f'X{n_x_vars+1}'))

    nodes = cl_cpdag.get_nodes()

    for i in range(n_x_vars):
        cl_cpdag.add_edge(Edge(nodes[i], nodes[-1], Endpoint.ARROW, Endpoint.TAIL))
    
    cg = CausalGraph(G=cl_cpdag)
    
    return cg

            
class AugmentedPC:
    """
    Runs the PC algorithm on an augmented graph, starting from a known MEC (optional).
    """
    def __init__(self, cpdag, test='kci', alpha=0.05, test_kwargs={}):
        self.cpdag = cpdag
        self.test = test
        self.alpha = alpha
        self.test_kwargs = test_kwargs
        self.dags_ = np.asarray(cpdag2dags(cpdag))
        self.n_vars_ = cpdag.shape[0]
        self.alpha_ = alpha
        self.n_envs_ = 0
        self.n_dags_ = self.dags_.shape[0]
        self.aug_cpdag_ = _construct_augmented_cpdag(cpdag)

    def learn_cpdag(self, Xs):
        if len(Xs) == 1:
            self.learned_cpdag_ = self.cpdag
            return
        
        data = np.block([
            [Xs[e], np.reshape([e] * Xs[e].shape[0], (-1, 1))]
            for e in range(len(Xs))
        ])
        
        if self.test == 'fisherz':
            from causallearn.utils.cit import fisherz
            test_func = fisherz
        elif self.test == 'kci':
            from causallearn.utils.cit import kci
            test_func = kci
        else:
            raise ValueError(f'Test {self.test} not implemented.')

        
        # Run meek rules on found edges
        from sparse_shift.causal_learn.SkeletonDiscovery import augmented_skeleton_discovery
        cg_skel_disc = augmented_skeleton_discovery(data, self.alpha_, test_func,
            stable=True,
            background_knowledge=None, verbose=False,
            show_progress=True, cg=self.aug_cpdag_)

        from causallearn.utils.PCUtils import Meek
        cg_meek = Meek.meek(cg_skel_disc, background_knowledge=None)
        
        adj = np.zeros(cg_meek.G.graph.shape)
        adj[cg_meek.G.graph > 0] = 1
        adj[np.abs(cg_meek.G.graph + cg_meek.G.graph.T) > 0] = 1
        
        self.learned_cpdag_ = adj[:-1, :-1]

    def get_dags(self):
        return cpdag2dags(self.get_cpdag())
    
    def get_cpdag(self):
        return self.learned_cpdag_
SPARSITY = 1
SAMPLE_SIZE = 100
N_ENV = 2
N_NODES = 4
DENSITY = 0.5
# Sample parent assignments
true_dag = erdos_renyi_dag(N_NODES, DENSITY, seed=2)
plot_dag(true_dag)
true_cpdag = dag2cpdag(true_dag)
apc = AugmentedPC(true_cpdag)
# Sample targets
intv_targets = [np.random.choice(N_NODES, SPARSITY, replace=False) for _ in range(N_ENV)]

intv_targets = [np.asarray([]), np.asarray([0])]

# Environment data
base_seed = 1234
Xs = [sample_cdnod_sim(
        true_dag,
        SAMPLE_SIZE,
        intervention_targets=targets,
        base_random_state=base_seed,
    ) for targets in intv_targets]
apc = AugmentedPC(true_cpdag)

apc.learn_cpdag(Xs)
apc.get_cpdag()
from causallearn.graph.Edge import Edge
from causallearn.graph.Endpoint import Endpoint

cl_dag = create_causal_learn_dag(true_dag)

cl_cpdag = dag2cpdag(cl_dag)
cl_cpdag.add_node(GraphNode(f'X{N_NODES+1}'))

nodes = cl_cpdag.get_nodes()

for i in range(N_NODES):
    cl_cpdag.add_edge(Edge(nodes[i], nodes[-1], Endpoint.ARROW, Endpoint.TAIL))
    # break
cg = CausalGraph(G=cl_cpdag)
cg.to_nx_graph()
cg.draw_nx_graph(skel=False)
# Run PC w/ background knowledge
cg_skel_disc = augmented_skeleton_discovery(data, 0.05, kci, stable=True,
    background_knowledge=None, verbose=False,
    show_progress=True, cg=cg)

cg_meek = Meek.meek(cg_skel_disc, background_knowledge=None)

from causaldag import DAG, PDAG
adj = np.zeros(cg_meek.G.graph.shape)
adj[cg_meek.G.graph > 0] = 1
adj[np.abs(cg_meek.G.graph + cg_meek.G.graph.T) > 0] = 1
cg_meek.G.graph
adj
from causallearn.graph.GraphClass import CausalGraph
G = CausalGraph(4)
G.to_nx_graph()
G.draw_nx_graph(skel=False)
cg.to_nx_graph()
cg.__dir__()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10474462017111745, 'fraction_numerical': 0.012963443090484833, 'mean_word_length': 3.446685878962536, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '419481', 'license': 'MIT', 'n_tokens_mistral': 2996, 'n_tokens_neox': 2847, 'n_words': 574, 'path': 'notebooks/07-rp-experimental_pipeline_causal_learn.ipynb', 'repo_name': 'rflperry/sparse_shift', 'size': '7714'}