|
3 | 3 | from __future__ import annotations |
4 | 4 |
|
5 | 5 | import logging |
| 6 | +from functools import partial |
6 | 7 | from itertools import combinations |
| 8 | +from multiprocessing import Pool |
7 | 9 | from typing import Generator, Set, Union |
8 | 10 |
|
9 | 11 | import networkx as nx |
| 12 | +import pandas as pd |
10 | 13 |
|
| 14 | +from causal_testing.estimation.abstract_estimator import Estimator |
| 15 | +from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator |
| 16 | +from causal_testing.estimation.logistic_regression_estimator import LogisticRegressionEstimator |
| 17 | +from causal_testing.estimation.multinomial_regression_estimator import MultinomialRegressionEstimator |
| 18 | +from causal_testing.specification.variable import Input, Output, Variable |
11 | 19 | from causal_testing.testing.base_test_case import BaseTestCase |
12 | | - |
13 | | -from .variable import Variable |
| 20 | +from causal_testing.testing.causal_effect import NoEffect, SomeEffect |
| 21 | +from causal_testing.testing.causal_test_case import CausalTestCase |
14 | 22 |
|
15 | 23 | Node = Union[str, int] # Node type hint: A node is a string or an int |
16 | 24 |
|
@@ -123,9 +131,10 @@ class CausalDAG(nx.DiGraph): |
123 | 131 | ensures it is acyclic. A CausalDAG must be specified as a dot file. |
124 | 132 | """ |
125 | 133 |
|
126 | | - def __init__(self, file_path: str = None, ignore_cycles: bool = False, **attr): |
| 134 | + def __init__(self, file_path: str = None, ignore_cycles: bool = False, datatypes: pd.Series = None, **attr): |
127 | 135 | super().__init__(**attr) |
128 | 136 | self.ignore_cycles = ignore_cycles |
| 137 | + self.datatypes = datatypes |
129 | 138 | if file_path: |
130 | 139 | if file_path.endswith(".dot"): |
131 | 140 | graph = nx.DiGraph(nx.nx_pydot.read_dot(file_path)) |
@@ -533,3 +542,155 @@ def to_dot_string(self) -> str: |
533 | 542 |
|
534 | 543 | def __str__(self): |
535 | 544 | return f"Nodes: {self.nodes}\nEdges: {self.edges}" |
| 545 | + |
| 546 | + def _estimator(self, base_test_case: BaseTestCase, nodes_to_ignore: set) -> Estimator: |
| 547 | + treatment_variable = base_test_case.treatment_variable.name |
| 548 | + outcome_variable = base_test_case.outcome_variable.name |
| 549 | + |
| 550 | + if self.datatypes is None or outcome_variable not in self.datatypes: |
| 551 | + raise ValueError(f"No datatype specified for {outcome_variable}.") |
| 552 | + |
| 553 | + adj_sets = self.direct_effect_adjustment_sets( |
| 554 | + [treatment_variable], [outcome_variable], nodes_to_ignore=nodes_to_ignore |
| 555 | + ) |
| 556 | + if not adj_sets: |
| 557 | + return None, None |
| 558 | + |
| 559 | + min_adj_set = sorted(list(map(lambda s: sorted(list(s)), adj_sets)))[0] |
| 560 | + |
| 561 | + if pd.api.types.is_bool_dtype(self.datatypes[outcome_variable]): |
| 562 | + return ( |
| 563 | + LogisticRegressionEstimator(base_test_case=base_test_case, adjustment_set=min_adj_set), |
| 564 | + "unit_odds_ratio", |
| 565 | + ) |
| 566 | + if pd.api.types.is_categorical_dtype(self.datatypes[outcome_variable]) or pd.api.types.is_object_dtype( |
| 567 | + self.datatypes[outcome_variable] |
| 568 | + ): |
| 569 | + return ( |
| 570 | + MultinomialRegressionEstimator(base_test_case=base_test_case, adjustment_set=min_adj_set), |
| 571 | + "unit_odds_ratio", |
| 572 | + ) |
| 573 | + if pd.api.types.is_numeric_dtype(self.datatypes[outcome_variable]): |
| 574 | + return LinearRegressionEstimator(base_test_case=base_test_case, adjustment_set=min_adj_set), "coefficient" |
| 575 | + raise ValueError(f"Invalid datatype for {outcome_variable}: {self.datatypes[outcome_variable]}") |
| 576 | + |
| 577 | + def generate_causal_test( # pylint: disable=R0912 |
| 578 | + self, u: str, v: str, nodes_to_ignore: set = None, **kwargs |
| 579 | + ) -> list[CausalTestCase]: |
| 580 | + """ |
| 581 | + Construct a metamorphic relation for a given node pair implied by the Causal DAG, or None if no such relation |
| 582 | + can be constructed (e.g. because every valid adjustment set contains a node to ignore). |
| 583 | +
|
| 584 | + :param u: The treatment node. |
| 585 | + :param v: The outcome node. |
| 586 | + :param nodes_to_ignore: Set of nodes which will be excluded from causal tests. |
| 587 | + :param kwargs: Keyword arguments to be passed through to test case generation. |
| 588 | +
|
| 589 | + :return: A list containing ShouldCause and ShouldNotCause metamorphic relations. |
| 590 | + """ |
| 591 | + |
| 592 | + if nodes_to_ignore is None: |
| 593 | + nodes_to_ignore = set() |
| 594 | + |
| 595 | + causal_tests = [] |
| 596 | + |
| 597 | + # Create a ShouldNotCause relation for each pair of nodes that are not directly connected |
| 598 | + if ((u, v) not in self.edges) and ((v, u) not in self.edges): |
| 599 | + u_in_ancestors = u in nx.ancestors(self, v) |
| 600 | + v_in_ancestors = v in nx.ancestors(self, u) |
| 601 | + |
| 602 | + # Case 1: U --> ... --> V or U _||_ V |
| 603 | + if u_in_ancestors or (not u_in_ancestors and not v_in_ancestors): |
| 604 | + base_test_case = BaseTestCase(Input(u, None), Output(v, None)) |
| 605 | + estimator, estimate_type = self._estimator(base_test_case, nodes_to_ignore) |
| 606 | + if estimator and estimate_type: |
| 607 | + causal_tests.append( |
| 608 | + CausalTestCase( |
| 609 | + base_test_case=base_test_case, |
| 610 | + expected_causal_effect=NoEffect(), |
| 611 | + estimator=estimator, |
| 612 | + estimate_type=estimate_type, |
| 613 | + **kwargs, |
| 614 | + ), |
| 615 | + ) |
| 616 | + |
| 617 | + # Case 2: V --> ... --> U or U _||_ V |
| 618 | + if v in nx.ancestors(self, u) or (not u_in_ancestors and not v_in_ancestors): |
| 619 | + base_test_case = BaseTestCase(Input(v, None), Output(u, None)) |
| 620 | + estimator, estimate_type = self._estimator(base_test_case, nodes_to_ignore) |
| 621 | + if estimator and estimate_type: |
| 622 | + causal_tests.append( |
| 623 | + CausalTestCase( |
| 624 | + base_test_case=base_test_case, |
| 625 | + expected_causal_effect=NoEffect(), |
| 626 | + estimator=estimator, |
| 627 | + estimate_type=estimate_type, |
| 628 | + **kwargs, |
| 629 | + ), |
| 630 | + ) |
| 631 | + |
| 632 | + # Create a ShouldCause relation for each edge (u, v) or (v, u) |
| 633 | + elif (u, v) in self.edges: |
| 634 | + base_test_case = BaseTestCase(Input(u, None), Output(v, None)) |
| 635 | + estimator, estimate_type = self._estimator(base_test_case, nodes_to_ignore) |
| 636 | + if estimator and estimate_type: |
| 637 | + causal_tests.append( |
| 638 | + CausalTestCase( |
| 639 | + base_test_case=base_test_case, |
| 640 | + expected_causal_effect=SomeEffect(), |
| 641 | + estimator=estimator, |
| 642 | + estimate_type=estimate_type, |
| 643 | + **kwargs, |
| 644 | + ), |
| 645 | + ) |
| 646 | + else: |
| 647 | + base_test_case = BaseTestCase(Input(v, None), Output(u, None)) |
| 648 | + estimator, estimate_type = self._estimator(base_test_case, nodes_to_ignore) |
| 649 | + if estimator and estimate_type: |
| 650 | + causal_tests.append( |
| 651 | + CausalTestCase( |
| 652 | + base_test_case=base_test_case, |
| 653 | + expected_causal_effect=SomeEffect(), |
| 654 | + estimator=estimator, |
| 655 | + estimate_type=estimate_type, |
| 656 | + **kwargs, |
| 657 | + ), |
| 658 | + ) |
| 659 | + return causal_tests |
| 660 | + |
| 661 | + def generate_causal_tests( |
| 662 | + self, nodes_to_ignore: set = None, threads: int = 0, nodes_to_test: set = None, **kwargs: dict |
| 663 | + ) -> list[CausalTestCase]: |
| 664 | + """ |
| 665 | + Construct a list of metamorphic relations implied by the Causal DAG. |
| 666 | + This list of metamorphic relations contains a ShouldCause relation for every edge, and a ShouldNotCause |
| 667 | + relation for every (minimal) conditional independence relation implied by the structure of the DAG. |
| 668 | +
|
| 669 | + :param nodes_to_ignore: Set of nodes which will be excluded from causal tests. |
| 670 | + :param threads: Number of threads to use (if generating in parallel). |
| 671 | + :param nodes_to_test: Set of nodes to test the relationships between (defaults to all nodes). |
| 672 | + :param kwargs: Keyword arguments to be passed through to test case generation. |
| 673 | +
|
| 674 | + :return: A list containing ShouldCause and ShouldNotCause metamorphic relations. |
| 675 | + """ |
| 676 | + |
| 677 | + if nodes_to_ignore is None: |
| 678 | + nodes_to_ignore = set() |
| 679 | + nodes_to_ignore = nodes_to_ignore.union(set(self.cycle_nodes())) |
| 680 | + |
| 681 | + if nodes_to_test is None: |
| 682 | + nodes_to_test = self.nodes |
| 683 | + |
| 684 | + if threads < 2: |
| 685 | + causal_tests = [ |
| 686 | + self.generate_causal_test(u, v, nodes_to_ignore, **kwargs) |
| 687 | + for u, v in combinations(filter(lambda node: node not in nodes_to_ignore, nodes_to_test), 2) |
| 688 | + ] |
| 689 | + else: |
| 690 | + with Pool(threads) as pool: |
| 691 | + causal_tests = pool.starmap( |
| 692 | + partial(self.generate_causal_test, nodes_to_ignore=nodes_to_ignore, **kwargs), |
| 693 | + combinations(filter(lambda node: node not in nodes_to_ignore, nodes_to_test), 2), |
| 694 | + ) |
| 695 | + |
| 696 | + return [item for items in causal_tests for item in items] |
0 commit comments