Coverage for src/causalspyne/main.py: 54%

1""" 

2generate DAG and its marginal DAG 

3""" 

4 

5from datetime import datetime 

6try: 

7 from contextlib import chdir 

8except Exception: 

9 from causalspyne.py3_9_10_compatibility import chdir 

10 

11from pathlib import Path 

12 

13import matplotlib.pyplot as plt 

14 

15from causalspyne.gen_dag_2level import GenDAG2Level 

16from causalspyne.dag_gen import GenDAG 

17from causalspyne.dag_gen_topo_order import RootConfounderDAG 

18from causalspyne.dag_viewer import DAGView 

19from causalspyne.dag2ancestral import DAG2Ancestral 

20 

21from causalspyne.draw_dags import draw_dags_nx 

22from causalspyne.utils_random import coerce_rng 

23 

24 

25def gen_partially_observed( 

26 degree=2, 

27 list_confounder2hide=None, 

28 size_micro_node_dag=4, 

29 max_num_local_nodes=4, 

30 min_num_local_nodes=3, 

31 num_macro_nodes=4, 

32 num_sample=200, 

33 output_dir="output/", 

34 rng=None, 

35 dft_noise="Gaussian", 

36 graphviz=False, 

37 plot=True, 

38 strategy_cls=None, 

39): 

40 """ 

41 sole function as user interface 

42 strategy_cls: optional skeleton generator class (default: Erdos_Renyi_PLP) 

43 """ 

44 if list_confounder2hide is None: 

45 list_confounder2hide = [0.5, 0.9] 

46 rng = coerce_rng(rng) 

47 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_") 

48 output_dir = Path(output_dir) 

49 output_dir.mkdir(parents=True, exist_ok=True) 

50 

51 simple_dag_gen = GenDAG(num_nodes=size_micro_node_dag, 

52 degree=degree, rng=rng, strategy_cls=strategy_cls) 

53 

54 dag_gen = GenDAG2Level( 

55 dag_generator=simple_dag_gen, 

56 num_macro_nodes=num_macro_nodes, 

57 num_micro_nodes=size_micro_node_dag, 

58 max_num_local_nodes=max_num_local_nodes, 

59 min_num_local_nodes=min_num_local_nodes, 

60 rng=rng, 

61 ) 

62 dag = dag_gen.run() 

63 dag.to_binary_csv(benchpress=False, 

64 name=output_dir / f"ground_truth_dag_{timestamp}d.csv") 

65 

66 subview = DAGView(dag=dag, rng=rng, dft_noise=dft_noise) 

67 return re_hide(subview, dag, num_sample, list_confounder2hide, output_dir, 

68 graphviz, timestamp, plot=plot) 

69 

70 

71def gen_root_confounder_hidden( 

72 degree=2, 

73 size_micro_node_dag=4, 

74 max_num_local_nodes=4, 

75 num_macro_nodes=4, 

76 num_sample=200, 

77 output_dir="output/", 

78 rng=None, 

79 dft_noise="Gaussian", 

80 graphviz=False, 

81 plot=True, 

82): 

83 """ 

84 FCI benchmark where the root macro node (guaranteed confounder) is 

85 entirely hidden. Uses RootConfounderDAG for the macro backbone so the 

86 root always confounds at least two other macro nodes, then hides all 

87 micro nodes belonging to that root macro cluster. 

88 """ 

89 rng = coerce_rng(rng) 

90 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_") 

91 output_dir = Path(output_dir) 

92 output_dir.mkdir(parents=True, exist_ok=True) 

93 

94 simple_dag_gen = GenDAG( 

95 num_nodes=size_micro_node_dag, 

96 degree=degree, 

97 rng=rng, 

98 strategy_cls=RootConfounderDAG, 

99 ) 

100 

101 dag_gen = GenDAG2Level( 

102 dag_generator=simple_dag_gen, 

103 num_macro_nodes=num_macro_nodes, 

104 num_micro_nodes=size_micro_node_dag, 

105 max_num_local_nodes=max_num_local_nodes, 

106 rng=rng, 

107 ) 

108 dag = dag_gen.run() 

109 dag.to_binary_csv(benchpress=False, 

110 name=output_dir / f"ground_truth_dag_{timestamp}d.csv") 

111 

112 # hide all micro nodes of every root macro node 

113 root_names = dag_gen.get_root_macro_names() 

114 list_global_inds_to_hide = [] 

115 for name in root_names: 

116 list_global_inds_to_hide.extend(dag_gen.get_macro_node_global_inds(name)) 

117 

118 subview = DAGView(dag=dag, rng=rng, dft_noise=dft_noise) 

119 # DAGView.hide_top_order expects topological-order positions, not global indices 

120 topo_positions = [dag.list_ind_nodes_sorted.index(g) 

121 for g in list_global_inds_to_hide] 

122 subview._data_arr = subview.data_gen.gen(num_sample) 

123 subview.hide_top_order(topo_positions) 

124 

125 return re_hide_by_inds(subview, dag, list_global_inds_to_hide, 

126 output_dir, graphviz, timestamp, plot=plot) 

127 

128 

129def re_hide_by_inds(subview, dag, list_global_inds_to_hide, 

130 output_dir, graphviz, timestamp, plot=True): 

131 """Like re_hide but uses pre-computed global indices instead of percentages.""" 

132 dag2ancestral = DAG2Ancestral(dag.mat_adjacency) 

133 pred_ancestral_graph_mat = dag2ancestral.run(list_global_inds_to_hide) 

134 

135 if plot: 

136 fig, (ax1, ax2, ax3) = plt.subplots(1, 3) 

137 str_hidden = "_".join(str(i) for i in list_global_inds_to_hide) 

138 mtitle = "hide_root_" + str_hidden 

139 fig.suptitle(mtitle) 

140 

141 dag.visualize(title="DAG", ax=ax1, graphviz=graphviz) 

142 ax1.set_title("DAG") 

143 

144 draw_dags_nx( 

145 pred_ancestral_graph_mat, 

146 dict_ind2name={ 

147 i: name for i, name in enumerate(sorted(subview.node_names)) 

148 }, 

149 title="ancestral", 

150 ax=ax2, 

151 graphviz=graphviz, 

152 ) 

153 ax2.set_title("ancestral") 

154 

155 subview.visualize(title="subDAG", ax=ax3, graphviz=graphviz) 

156 ax3.set_title("subDAG") 

157 

158 with chdir(output_dir): 

159 subview.to_csv() 

160 if plot: 

161 fig.savefig(f"graph_compare_{timestamp}dags.pdf", format="pdf") 

162 fig.savefig(f"graph_compare_{timestamp}dags.svg", format="svg") 

163 plt.close(fig) 

164 with open("hidden_nodes.csv", "w") as outfile: 

165 outfile.write(",".join(str(i) for i in list_global_inds_to_hide)) 

166 return subview 

167 

168 

169def ordered_ind_col2global_ind(inds_cols, subview_global_inds): 

170 """ 

171 given a predicted causal order in the form of column indices, transform it 

172 into global index of ground truth DAG 

173 """ 

174 list_global_inds = [subview_global_inds[ind_col] for ind_col in inds_cols] 

175 return list_global_inds 

176 

177 

178def re_hide(subview, dag, num_sample, list_confounder2hide, output_dir, 

179 graphviz, timestamp, plot=True): 

180 subview.run( 

181 num_samples=num_sample, confound=True, 

182 list_nodes2hide=list_confounder2hide 

183 ) 

184 str_node2hide = subview.str_node2hide 

185 

186 dag2ancestral = DAG2Ancestral(dag.mat_adjacency) 

187 list_confounder2hide_global_ind = subview.list_global_inds_nodes2hide 

188 pred_ancestral_graph_mat = dag2ancestral.run( 

189 list_confounder2hide_global_ind) 

190 

191 if plot: 

192 fig, (ax1, ax2, ax3) = plt.subplots(1, 3) 

193 mtitle = "hide_" + str_node2hide 

194 fig.suptitle(mtitle) # super-title 

195 

196 # ax1 

197 dag.visualize(title="DAG", ax=ax1, graphviz=graphviz) 

198 ax1.set_title("DAG") 

199 

200 # ax2 

201 draw_dags_nx( 

202 pred_ancestral_graph_mat, 

203 dict_ind2name={ 

204 i: name for i, name in enumerate(sorted(subview.node_names)) 

205 }, 

206 title="ancestral", 

207 ax=ax2, 

208 graphviz=graphviz, 

209 ) 

210 ax2.set_title("ancestral") 

211 # ax3 

212 subview.visualize( 

213 title="subDAG", ax=ax3, graphviz=graphviz 

214 ) 

215 ax3.set_title("subDAG") 

216 

217 with chdir(output_dir): 

218 subview.to_csv() 

219 if plot: 

220 fig.savefig(f"graph_compare_{timestamp}dags.pdf", format="pdf") 

221 fig.savefig(f"graph_compare_{timestamp}dags.svg", format="svg") 

222 plt.close(fig) 

223 with open("hidden_nodes.csv", "w") as outfile: 

224 outfile.write( 

225 ",".join(str(node) for node in 

226 subview._list_global_inds_unobserved) 

227 ) 

228 return subview