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

1""" 

2concrete class to generate simple DAGs 

3 

4""" 

5 

6import warnings 

7 

8from causalspyne.erdo_renyi_plp import Erdos_Renyi_PLP 

9from causalspyne.dag_interface import MatDAG 

10from causalspyne.weight import WeightGenWishart 

11from causalspyne.dag_manipulator import DAGManipulator 

12from causalspyne.utils_random import coerce_rng 

13 

14 

15class GenDAG: 

16 def __init__(self, num_nodes, degree, obj_gen_weight=None, 

17 rng=None, strategy_cls=None): 

18 """ 

19 degree: expected degree for each node 

20 strategy_cls: callable class that takes rng and returns a skeleton 

21 generator; defaults to Erdos_Renyi_PLP 

22 """ 

23 rng = coerce_rng(rng) 

24 self.num_nodes = num_nodes 

25 self.degree = degree 

26 if strategy_cls is None: 

27 strategy_cls = Erdos_Renyi_PLP 

28 self.strategy_gen_dag = strategy_cls(rng) 

29 self.obj_gen_weight = obj_gen_weight 

30 if obj_gen_weight is None: 

31 self.obj_gen_weight = WeightGenWishart(rng=rng) 

32 self.dag_manipulator = None 

33 self.rng = rng 

34 

35 def gen_dag(self, num_nodes=None, prefix="", *, target_num_confounder): 

36 """ 

37 generate DAG and wrap it around with interface 

38 """ 

39 if num_nodes is None: 

40 num_nodes = self.num_nodes 

41 mat_skeleton = self.strategy_gen_dag(num_nodes, self.degree) 

42 

43 mat_mask = (mat_skeleton != 0).astype(float) 

44 mat_weight = self.obj_gen_weight.gen(num_nodes) 

45 # Hardarmard product 

46 mat_weighted_adjacency = mat_mask * mat_weight 

47 

48 dag = MatDAG(mat_weighted_adjacency, name_prefix=prefix, rng=self.rng) 

49 self.dag_manipulator = DAGManipulator(dag, 

50 self.obj_gen_weight, self.rng) 

51 ind_arbitrary = dag.get_top_last() 

52 counter = 0 

53 for _ in range(dag.num_nodes): 

54 flag_success = self.dag_manipulator.mk_confound( 

55 ind_arbitrary_confound_input=ind_arbitrary) 

56 if not flag_success: 

57 counter += 1 

58 if dag.num_confounder >= target_num_confounder: 

59 break 

60 # FIXME: it can be the new ind_arbitrary has been tried out already 

61 ind_arbitrary = dag.climb(ind_arbitrary) 

62 if ind_arbitrary is None: 

63 break 

64 num_confounder = len(dag.list_confounder) 

65 if num_confounder < target_num_confounder and \ 

66 dag.num_nodes - target_num_confounder > 1: 

67 warnings.warn( 

68 f"\n failed to ensure {target_num_confounder} confounders for \ 

69 adjacency matrix \n{dag.mat_adjacency}, \ 

70 \n after {counter} failed trials, \ 

71 \n{num_confounder} confounders only") 

72 return dag