Coverage for src/causalspyne/dag_gen_topo_order.py: 56%

1import numpy as np 

2 

3 

4def _can_reach(mat, src, dst): 

5 """Return True if there is a directed path from src to dst in mat.""" 

6 n = mat.shape[0] 

7 visited = set() 

8 stack = [src] 

9 while stack: 

10 node = stack.pop() 

11 if node == dst: 

12 return True 

13 if node in visited: 

14 continue 

15 visited.add(node) 

16 # children of node: columns where mat[:, node] != 0 

17 for child in np.where(mat[:, node] != 0)[0]: 

18 if child not in visited: 

19 stack.append(child) 

20 return False 

21 

22 

23class RandTopoOrderDAG: 

24 """ 

25 Generate a random DAG by sampling a random topological order, then 

26 independently including each forward edge with probability 

27 degree / (num_nodes - 1). 

28 

29 Guaranteed to produce a valid DAG: edges only go from earlier to 

30 later positions in the sampled order, so no cycle is possible. 

31 

32 Matrix convention matches Erdos_Renyi_PLP: entry (i, j) means j -> i. 

33 """ 

34 

35 def __init__(self, rng): 

36 self.rng = rng 

37 

38 def __call__(self, num_nodes, degree): 

39 prob = float(degree) / (num_nodes - 1) 

40 topo_order = self.rng.permutation(num_nodes) # topo_order[k] = node at position k 

41 mat = np.zeros((num_nodes, num_nodes)) 

42 for k in range(num_nodes): 

43 for l in range(k + 1, num_nodes): 

44 if self.rng.random() < prob: 

45 src, dst = topo_order[k], topo_order[l] 

46 mat[dst, src] = 1.0 # src -> dst 

47 return mat 

48 

49 

50class RootConfounderDAG(RandTopoOrderDAG): 

51 """ 

52 Extends RandTopoOrderDAG so that every root node (no incoming edges) 

53 is guaranteed to be a confounder (at least 2 outgoing edges). 

54 

55 After base generation, any root with fewer than 2 children gets extra 

56 edges added to nodes that cannot reach the root (cycle-safe by construction). 

57 """ 

58 

59 def __call__(self, num_nodes, degree): 

60 mat = super().__call__(num_nodes, degree) 

61 root_nodes = list(np.where(mat.sum(axis=1) == 0)[0]) 

62 for root in root_nodes: 

63 current_children = list(np.where(mat[:, root] > 0)[0]) 

64 needed = max(0, 2 - len(current_children)) 

65 if needed == 0: 

66 continue 

67 # exclude nodes that can reach root (adding root->n would create a cycle) 

68 candidates = [n for n in range(num_nodes) 

69 if n != root 

70 and n not in current_children 

71 and not _can_reach(mat, n, root)] 

72 if len(candidates) < needed: 

73 continue 

74 targets = self.rng.choice(candidates, size=needed, replace=False) 

75 for t in targets: 

76 mat[t, root] = 1.0 

77 return mat