cstree

class cslearn.cstree.CStree(cards=[], labels=None)

Bases: object

A class representing a CStree, see [1]. It is initialized by a list of cardinalities of the variables at each level and a list of labels. The labels are optional and defaults (None) to [0,1,…,p-1].

Parameters:

• cards (list) – A list of integers representing the cardinality of each level (indexed from 0).

• labels (list, optional) – A list of strings or integers representing the labels of each level. Defaults to [0,1,…,p-1].

References

[1] E. Duarte and L. Solus. Representation of context-specific causal models with observational and interventional data, 2021, https://arxiv.org/abs/2101.09271.

Example

>>> # Figure 1. from (Duarte & Solus 2022)
>>> import cslearn.cstree as ct
>>> import cslearn.stage as st
>>> tree = ct.CStree([2, 2, 2, 2], labels=["X"+str(i) for i in range(1, 5)])
>>> tree.update_stages({
...     0: [{"context": {0: 0}},
...         {"context": {0: 1}}],
...     1: [{"context": {1: 0}, "color": "green"},
...         {"context": {0: 0, 1: 1}},
...         {"context": {0: 1, 1: 1}}],
...     2: [{"context": {0: 0, 2: 0}, "color": "blue"},
...         {"context": {0: 0, 2: 1}, "color": "orange"},
...         {"context": {0: 1, 2: 0}, "color": "red"},
...         {"context": {0: 1, 1: 1, 2: 1}},
...         {"context": {0: 1, 1: 0, 2: 1}}]})
>>> tree.sample_stage_parameters()

csi_relations(level='all')

Returns the context specific indepencende (CSI) relations for the CStree.

Examples

>>> rels = tree.csi_relations()
>>> for cont, rels in rels.items():
...     for rel in rels:
...         print(rel)
X1 ⊥ X3 | X2=0
X2 ⊥ X4 | X1=0, X3=0
X2 ⊥ X4 | X1=0, X3=1
X2 ⊥ X4 | X1=1, X3=0

csi_relations_per_level()

Get the context specific independence relations per level.

Returns:

The CSI relations per level. The keys are the levels, and the values are lists of CSI relations.

Return type:

dict

Example

>>> rels = tree.csi_relations_per_level()
>>> print("CSI relations per level")
>>> for key, rel in rels.items():
...     print("level {}: ".format(key))
...     for r in rel:
...         if (len(r.ci.a) > 0) and (len(r.ci.b) > 0):
...             # avoiding the singletons
...             print("the CSI")
...             print(r)
CSI relations per level
level 0:
level 1:
the CSI
0 ⊥ 2 | 1=0
level 2:
the CSI
1 ⊥ 3 | 0=0, 2=0
the CSI
1 ⊥ 3 | 0=0, 2=1
the CSI
1 ⊥ 3 | 0=1, 2=0
level 3:

estimate_stage_parameters(data, method='BDeu', alpha_tot=1)

Estimate the parameters of the stages of the CStree under a Dirichlet model.

Parameters:

• data (pd.DataFrame) – A pandas dataframe with the data.

• method (str) – The estimation method. Currently only “BDeu” [2] is implemented.

• alpha_tot (float) – Total Dirichlet pseudo-count, split proportionally across stages by their size.

Reference:

[2] C. Hughes, P. Strong, and A. Shenvi. Score equivalence for staged trees, 2023, https://arxiv.org/abs/2206.15322

Example

>>> t = ct.sample_cstree([2,2,2,2], max_cvars=1, prob_cvar=0.5, prop_nonsingleton=1)
>>> t.sample_stage_parameters()
>>> df = t.sample(500)
>>> t.estimate_stage_parameters(df, alpha_tot=1.0, method="BDeu")
>>> for lev, stagings in t.stages.items():
...    print("Level {}".format(lev))
...    for stage in stagings:
...        print(stage)
...    print()

fit(data: DataFrame, gibbs_samples=10000, poss_cvars=None, grasp_depth=3, grasp_score='local_score_BDeu', max_cvars=2, alpha_tot=1.0, method='BDeu', save_as='')

Learn a CStree from data using GRaSP + Gibbs order sampling (CSlearn).

Runs the three-phase CSlearn pipeline: 1. DAG pre-screening via GRaSP (or uses caller-supplied poss_cvars). 2. Gibbs order MCMC to find the MAP variable ordering. 3. Exact staging search to find the optimal CStree for that ordering.

Updates self in place and returns it.

Parameters:

• data (pd.DataFrame) – Training data; first column may be cardinalities (matching the format returned by sample()).

• gibbs_samples (int) – Number of Gibbs iterations for order sampling. Defaults to 10000.

• poss_cvars (dict, optional) – Pre-computed dict mapping each variable label to its list of possible context variables. If None (default), GRaSP is run on data to produce this dict.

• grasp_depth (int) – GRaSP search depth. Defaults to 3.

• grasp_score (str) – Scoring function for GRaSP. Defaults to "local_score_BDeu".

• max_cvars (int) – Maximum context-set size β. Defaults to 2.

• alpha_tot (float) – Total BDeu pseudo-count. Defaults to 1.0.

• method (str) – Scoring method, currently only "BDeu" is supported.

• save_as (str) – If non-empty, save intermediate results (poss_cvars, score tables, learned tree, …) to {save_as}-<name>.pkl files.

Returns:

self, updated with the learned structure and parameters.

Return type:

CStree

Example

>>> import pandas as pd
>>> import cslearn.cstree as ct
>>> tree = ct.sample_cstree([2, 2, 2, 2], max_cvars=1, prob_cvar=0.5)
>>> tree.sample_stage_parameters()
>>> df = tree.sample(500)
>>> learned = ct.CStree(tree.cards).fit(df)

get_stage(node: tuple)

Get the stage of a node in the CStree.

Parameters:

node (tuple or list) – A node in the CStree. It could be e.g. (0, 1, 0, 1).

Example

>>> # tree is the fig. 1 CStree
>>> stage = tree.get_stage([0, 0])
>>> print(stage)
[{0, 1}, 0]; probs: [0.38 0.62]; color: green

plot(full=False)

Plot the CStree.

Parameters:

full (bool) – If True, the complete tree is drawn with all nodes. Defaults to False, in which case only nodes visited during sampling are shown (pass full=True for a fresh tree with no sampling history).

Returns:

A pygraphviz graph.

Return type:

pygraphviz.agraph.AGraph

Examples

>>> tree.sample_stage_parameters()
>>> agraph = tree.plot(full=True)
>>> agraph.draw("cstree.png")

pmf(x, label_order=None)

Calculate the probability mass function of a given outcome.

Parameters:

• x (list) – Outcome values, one per variable. Order follows label_order if provided, otherwise self.labels.

• label_order (list, optional) – Permutation of self.labels describing the order of values in x. Defaults to self.labels.

Returns:

The probability mass of the outcome.

Return type:

float

pmf_log(x, label_order=None)

Calculate the log probability mass function of a given outcome.

Parameters:

x (list) – A list of values representing the outcome.

Returns:

The log probability mass of the outcome.

Return type:

float

predict(partial_observations: DataFrame, return_probs: bool = False) → DataFrame

Predict the most likely completion of partially-observed rows.

Given a CStree over random variables \(X_{[p]}\) and a DataFrame of partial observations, returns the MAP completion of the unobserved variables for each row.

Parameters:

• partial_observations – DataFrame whose columns are the labels of the observed variables and whose rows are individual observations. Pass a single-row DataFrame with no columns (pd.DataFrame(index=[0])) to predict all variables unconditionally.

• return_probs – If True, append a PROB column containing the normalized conditional probability of the MAP completion.

Returns:

DataFrame with one column per predicted variable (in label order). If return_probs=True, a PROB column is appended.

Example

>>> import random
>>> import numpy as np
>>> import pandas as pd
>>> import cslearn.cstree as ct
>>> np.random.seed(22)
>>> random.seed(22)
>>> t = ct.sample_cstree([3, 2, 2, 3], max_cvars=2, prob_cvar=0.5, prop_nonsingleton=1)
>>> t.sample_stage_parameters(alpha=2)
>>> t.sample(100)
>>> t.predict(pd.DataFrame({0: [1]}))
   1  2  3
0  0  1  1

sample(n: int)

Draws random samples from the CStree. It also dynamically generates nodes in the underlying tree and associated parameters on the fly in order to avoid creating the whole tree, which is O(2^p) (if singeltone stages are allowed), just to sample data. This is also useful when plotting the tree, as it will only contain the nodes that are actually used in the sampling.

Parameters:

n (int) – number of random samples.

Returns:

A pandas dataframe containing the data. The header contains labels of the variables and row 0 contains the cardinalities. The other rows contain the samples.

Return type:

pandas.DataFrame

Examples

>>> df = tree.sample(5)
>>> print(df)
   0  1  2  3
0  2  2  2  2
1  0  0  1  1
2  0  0  0  0
3  0  0  0  1
4  0  0  1  1
5  0  0  0  1

sample_stage_parameters(alpha=1)

Set the parameters of the stages of the CStree to be random samples from a Dirichlet distribution with hyper parameter alpha.

Parameters:

alpha (float) – The hyper parameter for the Dirichlet distribution.

Example

>>> t = ct.sample_cstree([2,2,2,2], max_cvars=1, prob_cvar=0.5, prop_nonsingleton=1)
>>> t.sample_stage_parameters()

stage_proportion(stage: Stage)

The proportion of the outcome space this stage represents.

Parameters:

stage (st.Stage) – A stage of this CStree.

Returns:

A number between 0 and 1.

Return type:

float

Example

>>> # Assuming all variables are binary
>>> s = st.Stage([0, {0, 1}, 1])
>>> stree.stage_proportion(s)
0.25

to_LDAG()

Return the LDAG representation of the CStree.

Returns:

A DAG whose edges carry CSI-relation labels. Node labels match self.labels.

Return type:

nx.DiGraph

to_df(write_probs=False)

Converts the CStree to a Pandas dataframe. The labels of the dataframe are the labels of the levels/variables in the CStree. The first row contains the cardinalities of the variables. The other rows contain the stages of the CStree. E.g. row number 3 contains the level 1 stage with context X2=0, etc.

Parameters:

write_probs (bool) – If True, the probabilities of the stages are written to the dataframe. Defaults to False.

Returns:

A Pandas dataframe with the stages of the CStree.

Return type:

df (pd.DataFrame)

Example

>>> tree.to_df()
    X1      X2      X3      X4
0   2       2       2       2
1   0       -       -       -
2   1       -       -       -
3   *       0       -       -
4   0       1       -       -
5   1       1       -       -
6   0       *       0       -
7   0       *       1       -
8   1       *       0       -
9   1       1       1       -
10  1       0       1       -
11  -       -       -       -

to_joint_distribution(label_order=None, with_outcomes=True)

Return the joint distribution of the CStree.

Parameters:

label_order (list, optional) – A list of labels from self.labels in the desired order. Defaults to None which means self.labels.

Returns:

The joint distribution of the CStree.

Return type:

Pandas Dataframe

Example

>>> df = tree.to_joint_distribution()
>>> print(df)
    X1      X2      X3      X4      prob    log_prob
0   0       0       0       0       0.283088        -1.261999
1   0       0       0       1       0.219686        -1.515555
2   0       0       1       0       0.061755        -2.784584
3   0       0       1       1       0.307910        -1.177948
4   0       1       0       0       0.004319        -5.444837
5   0       1       0       1       0.003351        -5.698393
6   0       1       1       0       0.005404        -5.220655
7   0       1       1       1       0.026943        -3.614019
8   1       0       0       0       0.003439        -5.672556
9   1       0       0       1       0.025955        -3.651409
10  1       0       1       0       0.020695        -3.877848
11  1       0       1       1       0.000916        -6.995125
12  1       1       0       0       0.002284        -6.081627
13  1       1       0       1       0.017241        -4.060480
14  1       1       1       0       0.012563        -4.377006
15  1       1       1       1       0.004451        -5.414586

to_minimal_context_agraphs(layout='dot')

This returns a dict of minimal context DAGs as pgraphviz graphs.

Parameters:

layout (str, optional) – Graphviz engine. Defaults to “dot”.

Returns:

pygraphviz graphs

Return type:

pygraphviz.agraph.AGraph

to_minimal_context_csis()

This returns a dict of minimal context CSIs.

Example

>>> minlcsis = tree.to_minimal_context_csis()
>>> for key, csis in minlcsis.items():
...     for csi in csis:
...         print("{}: CSI {}".format(key, csi))
X2=0: CSI X1 ⊥ X3 | X2=0
X3=0: CSI X2 ⊥ X4 | X1, X3=0
X1=0: CSI X2 ⊥ X4 | X3, X1=0

to_minimal_context_graphs()

This returns a dict of minimal context NetworkX DAGs.

Returns:

The keys are the contexts, and the values are the NetworkX DAGs.

Return type:

dict

Example

>>> # tree is the Figure 1 CStree
>>> gs = tree.to_minimal_context_graphs()
>>> for key, graph in gs.items():
...     print("{}: Edges {}".format(key, graph.edges()))
X2=0: Edges [('X1', 'X4'), ('X3', 'X4')]
X3=0: Edges [('X1', 'X2'), ('X1', 'X4')]
X1=0: Edges [('X2', 'X3'), ('X3', 'X4')]

update_stages(stages: dict)

Update/set the stages of the CStree.

Parameters:

stages (dict) – A dictionary of stage dicts. The keys are the levels, and the values are lists of dicts representing stages. The dicts representing stages should have the following keys: “context” and “color”. The “context” should have a dict as value, where the keys are the levels and the values are the values of the variables at that level. The “color” key is optional but should have a string as value, representing the color of the stage.

Example

>>> tree.update_stages({
...     0: [{"context": {0: 0}},
...         {"context": {0: 1}}],
...     1: [{"context": {1: 0}, "color": "green"},
...         {"context": {0: 0, 1: 1}},
...         {"context": {0: 1, 1: 1}}],
...     2: [{"context": {0: 0, 2: 0}, "color": "blue"},
...         {"context": {0: 0, 2: 1}, "color": "orange"},
...         {"context": {0: 1, 2: 0}, "color": "red"},
...         {"context": {0: 1, 1: 1, 2: 1}},
...         {"context": {0: 1, 1: 0, 2: 1}}]
...     })

cslearn.cstree.df_to_cstree(df, read_probs=True)

Convert a dataframe to a CStree. The dataframe should have the following format: The labels should be the level labels. The first row should be the cardinalities, the second row should be the first stage, the third row the second stage etc. (see cslearn.cstree.CStree.to_df()).

Parameters:

df (Pandas DataFrame) – The dataframe to convert.

Returns:

A CStree.

Return type:

CStree

Example

>>> df = tree.to_df()
>>> print(df)
>>> t2 = ct.df_to_cstree(df)
>>> df2 = t2.to_df()
>>> print("The same tree:")
>>> print(df)
    X1 X2 X3 X4
0   2  2  2  2
1   0  -  -  -
2   1  -  -  -
3   *  0  -  -
4   0  1  -  -
5   1  1  -  -
6   0  *  0  -
7   0  *  1  -
8   1  *  0  -
9   1  1  1  -
10  1  0  1  -
11  -  -  -  -
The same tree:
    X1 X2 X3 X4
0   2  2  2  2
1   0  -  -  -
2   1  -  -  -
3   *  0  -  -
4   0  1  -  -
5   1  1  -  -
6   0  *  0  -
7   0  *  1  -
8   1  *  0  -
9   1  1  1  -
10  1  0  1  -
11  -  -  -  -

cslearn.cstree.plot(graph, layout='dot')

Plots a graph using graphviz. Creates a pygraphviz graph from a NetworkX graph.

Parameters:

• graph (nx.Graph) – The graph to plot.

• layout (str, optional) – Graphviz layout engine. Defaults to “dot”.

Returns:

A pygraphviz graph.

Return type:

pygraphviz.agraph.AGraph

cslearn.cstree.sample_cstree(cards: list[int], max_cvars: int, prob_cvar: float, prop_nonsingleton: float = 1.0, labels: list | None = None) → CStree

Sample a random CStree with given cardinalities.

Parameters:

• cards (list) – cardinalities of the levels.

• max_cvars (int) – maximum number of context variables.

• prob_cvar (float) – probability a potential context variable (in the algorithm) will a context variable.

• prop_nonsingleton (float) – proportion of non-singleton stages. Defaults to 1.0, meaning no singletons, which is also what the algorithms in this library are designed for.

• labels (list, optional) – A list of strings or integers representing the labels of each level. Defaults to None which means [0,1,…,p-1].

Returns:

A CStree.

Return type:

CStree

Examples

>>> np.random.seed(1)
>>> random.seed(1)
>>> tree = ct.sample_cstree([2,2,2,2], max_cvars=1, prob_cvar=0.5, prop_nonsingleton=1)
>>> tree.to_df()
    0   1       2       3
0       2       2       2       2
1       0       -       -       -
2       1       -       -       -
3       1       *       -       -
4       0       *       -       -
5       *       *       0       -
6       *       *       1       -

cslearn.cstree.write_minimal_context_graphs_to_files(context_dags, prefix='mygraphs')

Write minimal context graphs to files.

Each context graph is written to a separate file; the context is included in the filename and as a figure label.

Parameters:

• context_dags (dict) – A dictionary of context graphs. The keys are the contexts, and the values are the graphs.

• prefix (str, optional) – Filename prefix. Defaults to “mygraphs”.

Example

>>> # tree is the Figure 1 CStree
>>> gs = tree.to_minimal_context_graphs()
>>> ct.write_minimal_context_graphs_to_files(gs, prefix="mygraphs")