Coverage for src/causalspyne/benchmark_fci.py: 100%

1""" 

2Core logic for the paired FCI benchmark (root vs intermediate hidden). 

3 

4Factored out of examples/ so it can be unit-tested and reused. 

5The key invariant: both scenarios are generated from the same integer seed, 

6so they share the same DAG adjacency matrix and the same full data array; 

7only the hidden column indices differ. 

8""" 

9 

10from __future__ import annotations 

11 

12import numpy as np 

13 

14from causalspyne.main import gen_partially_observed 

15from causalspyne.dag_gen_topo_order import RootConfounderDAG 

16 

17 

18SCENARIO_HIDDEN = { 

19 "root": [0], # topologically first confounder 

20 "intermediate": [1.0], # topologically last confounder 

21} 

22 

23 

24def run_paired_scenarios( 

25 seed: int, 

26 num_macro_nodes: int = 4, 

27 size_micro_node_dag=None, 

28 max_num_local_nodes: int = 7, 

29 min_num_local_nodes: int = 3, 

30 degree: float = 2.0, 

31 num_sample: int = 200, 

32 output_dir: str = "/tmp/benchmark_fci", 

33 strategy_cls=None, 

34) -> dict: 

35 """ 

36 Run both scenarios (root hidden, intermediate hidden) on the same DAG. 

37 

38 Returns a dict with keys 'root' and 'intermediate', each containing: 

39 - 'subview': the DAGView object (observed data + metadata) 

40 - 'full_data': np.ndarray of shape (num_sample, num_nodes_total) 

41 — the data BEFORE any columns are hidden 

42 - 'adj': binary adjacency matrix of the ground-truth DAG 

43 - 'hidden': list of global node indices that were hidden 

44 

45 Invariant (tested in tests/test_benchmark_fci.py): 

46 results['root']['adj'] == results['intermediate']['adj'] 

47 results['root']['full_data'] == results['intermediate']['full_data'] 

48 results['root']['hidden'] != results['intermediate']['hidden'] 

49 """ 

50 if strategy_cls is None: 

51 strategy_cls = RootConfounderDAG 

52 

53 results = {} 

54 for scenario_name, hidden_spec in SCENARIO_HIDDEN.items(): 

55 import io, contextlib 

56 buf = io.StringIO() 

57 with contextlib.redirect_stdout(buf): 

58 subview = gen_partially_observed( 

59 size_micro_node_dag=size_micro_node_dag, 

60 max_num_local_nodes=max_num_local_nodes, 

61 min_num_local_nodes=min_num_local_nodes, 

62 num_macro_nodes=num_macro_nodes, 

63 degree=degree, 

64 list_confounder2hide=hidden_spec, 

65 num_sample=num_sample, 

66 output_dir=f"{output_dir}/{scenario_name}/seed_{seed}", 

67 rng=seed, # integer → fresh RNG → reproducible 

68 plot=False, 

69 strategy_cls=strategy_cls, 

70 ) 

71 

72 results[scenario_name] = { 

73 "subview": subview, 

74 "full_data": subview._data_arr, # set by DAGView.run(), pre-hide 

75 "adj": (subview.dag.mat_adjacency != 0).astype(int), 

76 "hidden": list(subview.list_global_inds_nodes2hide), 

77 } 

78 

79 return results