Made generation options configurable

Author: jmafoster1

causal_testing/__main__.py

@@ -21,7 +21,15 @@ def main() -> None:
 
     if args.command == Command.GENERATE:
         logging.info("Generating causal tests")
-        generate_causal_tests(args.dag_path, args.output, args.ignore_cycles, args.threads)
+        generate_causal_tests(
+            args.dag_path,
+            args.output,
+            args.ignore_cycles,
+            args.threads,
+            effect_type=args.effect_type,
+            estimate_type=args.estimate_type,
+            estimator=args.estimator,
+        )
         logging.info("Causal test generation completed successfully")
         return
 

causal_testing/main.py

@@ -490,16 +490,36 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         help="The action you want to run - call `causal_testing {action} -h` for further details", dest="command"
     )
 
+    # Generation
     parser_generate = subparsers.add_parser(Command.GENERATE.value, help="Generate causal tests from a DAG")
     parser_generate.add_argument("-D", "--dag_path", help="Path to the DAG file (.dot)", required=True)
     parser_generate.add_argument("-o", "--output", help="Path for output file (.json)", required=True)
+    parser_generate.add_argument(
+        "-e",
+        "--estimator",
+        help="The name of the estimator class to use when evaluating tests (defaults to LinearRegressionEstimator)",
+        default="LinearRegressionEstimator",
+    )
+    parser_generate.add_argument(
+        "-T",
+        "--effect_type",
+        help="The effect type to estimate {direct, total}",
+        default="direct",
+    )
+    parser_generate.add_argument(
+        "-E",
+        "--estimate_type",
+        help="The estimate type to use when evaluating tests (defaults to coefficient)",
+        default="coefficient",
+    )
     parser_generate.add_argument(
         "-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False
     )
     parser_generate.add_argument(
         "--threads", "-t", type=int, help="The number of parallel threads to use.", required=False, default=0
     )
 
+    # Testing
     parser_test = subparsers.add_parser(Command.TEST.value, help="Run causal tests")
     parser_test.add_argument("-D", "--dag_path", help="Path to the DAG file (.dot)", required=True)
     parser_test.add_argument("-o", "--output", help="Path for output file (.json)", required=True)

causal_testing/testing/metamorphic_relation.py

@@ -37,13 +37,25 @@ def __eq__(self, other):
 class ShouldCause(MetamorphicRelation):
     """Class representing a should cause metamorphic relation."""
 
-    def to_json_stub(self, skip=True) -> dict:
-        """Convert to a JSON frontend stub string for user customisation"""
+    def to_json_stub(
+        self,
+        skip: bool = True,
+        estimate_type: str = "coefficient",
+        effect_type: str = "direct",
+        estimator: str = "LinearRegressionEstimator",
+    ) -> dict:
+        """
+        Convert to a JSON frontend stub string for user customisation.
+        :param skip: Whether to skip the test
+        :param effect_type: The type of causal effect to consider (total or direct)
+        :param estimate_type: The estimate type to use when evaluating tests
+        :param estimator: The name of the estimator class to use when evaluating the test
+        """
         return {
             "name": str(self),
-            "estimator": "LinearRegressionEstimator",
-            "estimate_type": "coefficient",
-            "effect": "direct",
+            "estimator": estimator,
+            "estimate_type": estimate_type,
+            "effect": effect_type,
             "treatment_variable": self.base_test_case.treatment_variable,
             "expected_effect": {self.base_test_case.outcome_variable: "SomeEffect"},
             "formula": (
@@ -63,13 +75,25 @@ def __str__(self):
 class ShouldNotCause(MetamorphicRelation):
     """Class representing a should cause metamorphic relation."""
 
-    def to_json_stub(self, skip=True) -> dict:
-        """Convert to a JSON frontend stub string for user customisation"""
+    def to_json_stub(
+        self,
+        skip: bool = True,
+        estimate_type: str = "coefficient",
+        effect_type: str = "direct",
+        estimator: str = "LinearRegressionEstimator",
+    ) -> dict:
+        """
+        Convert to a JSON frontend stub string for user customisation.
+        :param skip: Whether to skip the test
+        :param effect_type: The type of causal effect to consider (total or direct)
+        :param estimate_type: The estimate type to use when evaluating tests
+        :param estimator: The name of the estimator class to use when evaluating the test
+        """
         return {
             "name": str(self),
-            "estimator": "LinearRegressionEstimator",
-            "estimate_type": "coefficient",
-            "effect": "direct",
+            "estimator": estimator,
+            "estimate_type": estimate_type,
+            "effect": effect_type,
             "treatment_variable": self.base_test_case.treatment_variable,
             "expected_effect": {self.base_test_case.outcome_variable: "NoEffect"},
             "formula": (
@@ -179,7 +203,15 @@ def generate_metamorphic_relations(
     return [item for items in metamorphic_relations for item in items]
 
 
-def generate_causal_tests(dag_path: str, output_path: str, ignore_cycles: bool = False, threads: int = 0):
+def generate_causal_tests(
+    dag_path: str,
+    output_path: str,
+    ignore_cycles: bool = False,
+    threads: int = 0,
+    estimate_type: str = "coefficient",
+    effect_type: str = "direct",
+    estimator: str = "LinearRegressionEstimator",
+):
     """
     Generate and output causal tests for a given DAG.
 
@@ -190,6 +222,9 @@ def generate_causal_tests(dag_path: str, output_path: str, ignore_cycles: bool =
                           be omitted from the test set.
     :param threads: The number of threads to use to generate tests in parallel. If unspecified, tests are generated in
                     serial. This is tylically fine unless the number of tests to be generated is >10000.
+    :param effect_type: The type of causal effect to consider (total or direct)
+    :param estimate_type: The estimate type to use when evaluating tests
+    :param estimator: The name of the estimator class to use when evaluating the test
     """
     causal_dag = CausalDAG(dag_path, ignore_cycles=ignore_cycles)
 
@@ -212,7 +247,7 @@ def generate_causal_tests(dag_path: str, output_path: str, ignore_cycles: bool =
         relations = generate_metamorphic_relations(causal_dag, nodes_to_test=dag_nodes_to_test, threads=threads)
 
     tests = [
-        relation.to_json_stub(skip=False)
+        relation.to_json_stub(skip=False, estimate_type=estimate_type, effect_type=effect_type, estimator=estimator)
         for relation in relations
         if len(list(causal_dag.graph.predecessors(relation.base_test_case.outcome_variable))) > 0
     ]

tests/main_tests/test_main.py

@@ -356,6 +356,28 @@ def test_parse_args_generation(self):
                 main()
                 self.assertTrue(os.path.exists(os.path.join(tmp, "tests.json")))
 
+    def test_parse_args_generation_non_default(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            with unittest.mock.patch(
+                "sys.argv",
+                [
+                    "causal_testing",
+                    "generate",
+                    "--dag_path",
+                    str(self.dag_path),
+                    "--output",
+                    os.path.join(tmp, "tests_non_default.json"),
+                    "--estimator",
+                    "LogisticRegressionEstimator",
+                    "--estimate_type",
+                    "unit_odds_ratio",
+                    "--effect_type",
+                    "total",
+                ],
+            ):
+                main()
+                self.assertTrue(os.path.exists(os.path.join(tmp, "tests_non_default.json")))
+
     def tearDown(self):
         if self.output_path.parent.exists():
             shutil.rmtree(self.output_path.parent)

tests/testing_tests/test_metamorphic_relations.py

@@ -50,9 +50,9 @@ def test_should_not_cause_json_stub(self):
         causal_dag = CausalDAG(self.dag_dot_path)
         causal_dag.graph.remove_nodes_from(["X2", "X3"])
         adj_set = list(causal_dag.direct_effect_adjustment_sets(["X1"], ["Z"])[0])
-        should_not_cause_MR = ShouldNotCause(BaseTestCase("X1", "Z"), adj_set)
+        should_not_cause_mr = ShouldNotCause(BaseTestCase("X1", "Z"), adj_set)
         self.assertEqual(
-            should_not_cause_MR.to_json_stub(),
+            should_not_cause_mr.to_json_stub(),
             {
                 "effect": "direct",
                 "estimate_type": "coefficient",
@@ -66,15 +66,39 @@ def test_should_not_cause_json_stub(self):
             },
         )
 
+    def test_should_not_cause_logistic_json_stub(self):
+        """Test if the ShouldCause MR passes all metamorphic tests where the DAG perfectly represents the program
+        and there is only a single input."""
+        causal_dag = CausalDAG(self.dag_dot_path)
+        causal_dag.graph.remove_nodes_from(["X2", "X3"])
+        adj_set = list(causal_dag.direct_effect_adjustment_sets(["X1"], ["Z"])[0])
+        should_not_cause_mr = ShouldNotCause(BaseTestCase("X1", "Z"), adj_set)
+        self.assertEqual(
+            should_not_cause_mr.to_json_stub(
+                effect_type="total", estimate_type="unit_odds_ratio", estimator="LogisticRegressionEstimator"
+            ),
+            {
+                "effect": "total",
+                "estimate_type": "unit_odds_ratio",
+                "estimator": "LogisticRegressionEstimator",
+                "expected_effect": {"Z": "NoEffect"},
+                "treatment_variable": "X1",
+                "name": "X1 _||_ Z",
+                "formula": "Z ~ X1",
+                "alpha": 0.05,
+                "skip": True,
+            },
+        )
+
     def test_should_cause_json_stub(self):
         """Test if the ShouldCause MR passes all metamorphic tests where the DAG perfectly represents the program
         and there is only a single input."""
         causal_dag = CausalDAG(self.dag_dot_path)
         causal_dag.graph.remove_nodes_from(["X2", "X3"])
         adj_set = list(causal_dag.direct_effect_adjustment_sets(["X1"], ["Z"])[0])
-        should_cause_MR = ShouldCause(BaseTestCase("X1", "Z"), adj_set)
+        should_cause_mr = ShouldCause(BaseTestCase("X1", "Z"), adj_set)
         self.assertEqual(
-            should_cause_MR.to_json_stub(),
+            should_cause_mr.to_json_stub(),
             {
                 "effect": "direct",
                 "estimate_type": "coefficient",
@@ -87,6 +111,29 @@ def test_should_cause_json_stub(self):
             },
         )
 
+    def test_should_cause_logistic_json_stub(self):
+        """Test if the ShouldCause MR passes all metamorphic tests where the DAG perfectly represents the program
+        and there is only a single input."""
+        causal_dag = CausalDAG(self.dag_dot_path)
+        causal_dag.graph.remove_nodes_from(["X2", "X3"])
+        adj_set = list(causal_dag.direct_effect_adjustment_sets(["X1"], ["Z"])[0])
+        should_cause_mr = ShouldCause(BaseTestCase("X1", "Z"), adj_set)
+        self.assertEqual(
+            should_cause_mr.to_json_stub(
+                effect_type="total", estimate_type="unit_odds_ratio", estimator="LogisticRegressionEstimator"
+            ),
+            {
+                "effect": "total",
+                "estimate_type": "unit_odds_ratio",
+                "estimator": "LogisticRegressionEstimator",
+                "expected_effect": {"Z": "SomeEffect"},
+                "formula": "Z ~ X1",
+                "treatment_variable": "X1",
+                "name": "X1 --> Z",
+                "skip": True,
+            },
+        )
+
     def test_all_metamorphic_relations_implied_by_dag(self):
         dag = CausalDAG(self.dag_dot_path)
         dag.add_edge("Z", "Y")  # Add a direct path from Z to Y so M becomes a mediator