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: int = 3, 

28 max_num_local_nodes: int = 4, 

29 degree: float = 2.0, 

30 num_sample: int = 200, 

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

32 strategy_cls=None, 

33) -> dict: 

34 """ 

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

36 

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

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

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

40 — the data BEFORE any columns are hidden 

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

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

43 

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

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

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

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

48 """ 

49 if strategy_cls is None: 

50 strategy_cls = RootConfounderDAG 

51 

52 results = {} 

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

54 import io, contextlib 

55 buf = io.StringIO() 

56 with contextlib.redirect_stdout(buf): 

57 subview = gen_partially_observed( 

58 size_micro_node_dag=size_micro_node_dag, 

59 max_num_local_nodes=max_num_local_nodes, 

60 num_macro_nodes=num_macro_nodes, 

61 degree=degree, 

62 list_confounder2hide=hidden_spec, 

63 num_sample=num_sample, 

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

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

66 plot=False, 

67 strategy_cls=strategy_cls, 

68 ) 

69 

70 results[scenario_name] = { 

71 "subview": subview, 

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

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

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

75 } 

76 

77 return results