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 num_macro_nodes=4, 

31 num_sample=200, 

32 output_dir="output/", 

33 rng=None, 

34 dft_noise="Gaussian", 

35 graphviz=False, 

36 plot=True, 

37 strategy_cls=None, 

38): 

39 """ 

40 sole function as user interface 

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

42 """ 

43 if list_confounder2hide is None: 

44 list_confounder2hide = [0.5, 0.9] 

45 rng = coerce_rng(rng) 

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

47 output_dir = Path(output_dir) 

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

49 

50 simple_dag_gen = GenDAG(num_nodes=size_micro_node_dag, 

51 degree=degree, rng=rng, strategy_cls=strategy_cls) 

52 

53 dag_gen = GenDAG2Level( 

54 dag_generator=simple_dag_gen, 

55 num_macro_nodes=num_macro_nodes, 

56 num_micro_nodes=size_micro_node_dag, 

57 max_num_local_nodes=max_num_local_nodes, 

58 rng=rng, 

59 ) 

60 dag = dag_gen.run() 

61 dag.to_binary_csv(benchpress=False, 

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

63 

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

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

66 graphviz, timestamp, plot=plot) 

67 

68 

69def gen_root_confounder_hidden( 

70 degree=2, 

71 size_micro_node_dag=4, 

72 max_num_local_nodes=4, 

73 num_macro_nodes=4, 

74 num_sample=200, 

75 output_dir="output/", 

76 rng=None, 

77 dft_noise="Gaussian", 

78 graphviz=False, 

79 plot=True, 

80): 

81 """ 

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

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

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

85 micro nodes belonging to that root macro cluster. 

86 """ 

87 rng = coerce_rng(rng) 

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

89 output_dir = Path(output_dir) 

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

91 

92 simple_dag_gen = GenDAG( 

93 num_nodes=size_micro_node_dag, 

94 degree=degree, 

95 rng=rng, 

96 strategy_cls=RootConfounderDAG, 

97 ) 

98 

99 dag_gen = GenDAG2Level( 

100 dag_generator=simple_dag_gen, 

101 num_macro_nodes=num_macro_nodes, 

102 num_micro_nodes=size_micro_node_dag, 

103 max_num_local_nodes=max_num_local_nodes, 

104 rng=rng, 

105 ) 

106 dag = dag_gen.run() 

107 dag.to_binary_csv(benchpress=False, 

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

109 

110 # hide all micro nodes of every root macro node 

111 root_names = dag_gen.get_root_macro_names() 

112 list_global_inds_to_hide = [] 

113 for name in root_names: 

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

115 

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

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

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

119 for g in list_global_inds_to_hide] 

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

121 subview.hide_top_order(topo_positions) 

122 

123 return re_hide_by_inds(subview, dag, list_global_inds_to_hide, 

124 output_dir, graphviz, timestamp, plot=plot) 

125 

126 

127def re_hide_by_inds(subview, dag, list_global_inds_to_hide, 

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

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

130 dag2ancestral = DAG2Ancestral(dag.mat_adjacency) 

131 pred_ancestral_graph_mat = dag2ancestral.run(list_global_inds_to_hide) 

132 

133 if plot: 

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

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

136 mtitle = "hide_root_" + str_hidden 

137 fig.suptitle(mtitle) 

138 

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

140 ax1.set_title("DAG") 

141 

142 draw_dags_nx( 

143 pred_ancestral_graph_mat, 

144 dict_ind2name={ 

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

146 }, 

147 title="ancestral", 

148 ax=ax2, 

149 graphviz=graphviz, 

150 ) 

151 ax2.set_title("ancestral") 

152 

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

154 ax3.set_title("subDAG") 

155 

156 with chdir(output_dir): 

157 subview.to_csv() 

158 if plot: 

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

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

161 plt.close(fig) 

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

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

164 return subview 

165 

166 

167def ordered_ind_col2global_ind(inds_cols, subview_global_inds): 

168 """ 

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

170 into global index of ground truth DAG 

171 """ 

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

173 return list_global_inds 

174 

175 

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

177 graphviz, timestamp, plot=True): 

178 subview.run( 

179 num_samples=num_sample, confound=True, 

180 list_nodes2hide=list_confounder2hide 

181 ) 

182 str_node2hide = subview.str_node2hide 

183 

184 dag2ancestral = DAG2Ancestral(dag.mat_adjacency) 

185 list_confounder2hide_global_ind = subview.list_global_inds_nodes2hide 

186 pred_ancestral_graph_mat = dag2ancestral.run( 

187 list_confounder2hide_global_ind) 

188 

189 if plot: 

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

191 mtitle = "hide_" + str_node2hide 

192 fig.suptitle(mtitle) # super-title 

193 

194 # ax1 

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

196 ax1.set_title("DAG") 

197 

198 # ax2 

199 draw_dags_nx( 

200 pred_ancestral_graph_mat, 

201 dict_ind2name={ 

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

203 }, 

204 title="ancestral", 

205 ax=ax2, 

206 graphviz=graphviz, 

207 ) 

208 ax2.set_title("ancestral") 

209 # ax3 

210 subview.visualize( 

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

212 ) 

213 ax3.set_title("subDAG") 

214 

215 with chdir(output_dir): 

216 subview.to_csv() 

217 if plot: 

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

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

220 plt.close(fig) 

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

222 outfile.write( 

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

224 subview._list_global_inds_unobserved) 

225 ) 

226 return subview