Coverage for causalspyne/data_gen.py: 100%

1""" 

2generate Linear Gaussian 

3""" 

4 

5import numpy as np 

6from numpy.random import default_rng 

7 

8from causalspyne.noise_idiosyncratic import Idiosyncratic 

9from causalspyne.edge_models import EdgeModelLinear 

10 

11 

12class DataGen: 

13 def __init__(self, dag, edge_model=None, idiosynchratic=None, rng=default_rng(0)): 

14 self.dag = dag 

15 self.idiosyncratic = idiosynchratic 

16 if idiosynchratic is None: 

17 self.idiosyncratic = Idiosyncratic(rng) 

18 self.edge_model = edge_model 

19 if edge_model is None: 

20 self.edge_model = EdgeModelLinear(self.dag) 

21 

22 def gen(self, num_samples): 

23 """ 

24 Generate linear Gaussian data from a given DAG. 

25 

26 Parameters: 

27 - num_samples: int, number of samples to generate. 

28 - noise_std: float, standard deviation of the Gaussian noise. 

29 

30 Returns: 

31 - data: np.ndarray, generated data of shape (num_samples, num_nodes). 

32 """ 

33 list_ind_nodes_topo_order = self.dag.topological_sort() 

34 

35 # Number of nodes 

36 num_nodes = len(list_ind_nodes_topo_order) 

37 

38 # Initialize the data matrix 

39 data = np.zeros((num_samples, num_nodes)) 

40 

41 # Generate data for each node in topological order 

42 for node in list_ind_nodes_topo_order: 

43 list_parents_inds = self.dag.get_list_parents_inds(node) 

44 noise = self.idiosyncratic.gen(num_samples) 

45 data[:, node] = noise 

46 if list_parents_inds: 

47 bias = self.edge_model.run(node, data[:, list_parents_inds]) 

48 data[:, node] += bias 

49 return data