Jmafoster1/388 fix large scenarios

causal_testing/__main__.py

@@ -1,18 +1,163 @@
 """This module contains the main entrypoint functionality to the Causal Testing Framework."""
 
-import json
+import argparse
 import logging
-import os
-import tempfile
+from enum import Enum
 from importlib.metadata import entry_points
-from pathlib import Path
+from typing import Optional, Sequence
 
 import networkx as nx
 import pandas as pd
 
+from causal_testing.causal_testing_framework import CausalTestingFramework, read_dataframe
 from causal_testing.testing.metamorphic_relation import generate_causal_tests
 
-from .main import CausalTestingFramework, CausalTestingPaths, Command, parse_args, setup_logging
+logger = logging.getLogger(__name__)
+
+
+class Command(Enum):
+    """
+    Enum for supported CTF commands.
+    """
+
+    TEST = "test"
+    GENERATE = "generate"
+    DISCOVER = "discover"
+
+
+def setup_logging(level: str) -> None:
+    """Set up logging configuration."""
+    logging.basicConfig(
+        level=level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+    )
+
+
+def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    """Parse command line arguments."""
+    main_parser = argparse.ArgumentParser(
+        add_help=True,
+        description="Causal Testing Framework - "
+        "A causal inference-driven framework for functional black-box testing of complex software.",
+    )
+
+    main_parser.add_argument(
+        "-l",
+        "--log_level",
+        default="WARNING",
+        type=str.upper,
+        choices=["NONE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Set the logging level (default: WARNING).",
+    )
+
+    subparsers = main_parser.add_subparsers(
+        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)
+    parser_test.add_argument("-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False)
+    parser_test.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True)
+    parser_test.add_argument("-t", "--test-config", help="Path to test configuration file (.json)", required=True)
+    parser_test.add_argument("-q", "--query", help="Query string to filter data (e.g. 'age > 18')", type=str)
+    parser_test.add_argument(
+        "-a", "--adequacy", help="Calculate causal test adequacy for each test case", action="store_true", default=False
+    )
+    parser_test.add_argument(
+        "-b",
+        "--adequacy-bootstrap-size",
+        dest="bootstrap_size",
+        help="Number of bootstrap samples for causal test adequacy. Defaults to 100",
+        type=int,
+    )
+    parser_test.add_argument(
+        "-s",
+        "--silent",
+        action="store_true",
+        help="Do not crash on error. If set to true, errors are recorded as test results.",
+        default=False,
+    )
+
+    # Discovery
+    parser_discover = subparsers.add_parser(Command.DISCOVER.value, help="Discover causal structures from data")
+    parser_discover.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True)
+    parser_discover.add_argument(
+        "-a",
+        "--alpha",
+        help=(
+            "The significance level of the confidence intervals used to determine causality. "
+            "This should be a value between 0 and 1. Defaults to 0.05 for 95%% confidence intervals."
+        ),
+        default=0.05,
+    )
+    parser_discover.add_argument(
+        "-t",
+        "--technique",
+        help="The name of the technique to use. Currently supported are 'HillClimberDiscovery' and 'NSGADiscovery'",
+        required=True,
+    )
+    parser_discover.add_argument(
+        "-V",
+        "--variables",
+        help="The subset of variables from the data to consider. Defaults to all.",
+        nargs="*",
+        default=[],
+    )
+    parser_discover.add_argument("-o", "--output", help="Path for output DAG file (.dot)", required=True)
+    parser_discover.add_argument(
+        "-i", "--include-edges", help="Path to file containing edges to include", required=False
+    )
+    parser_discover.add_argument(
+        "-e", "--exclude-edges", help="Path to file containing edges to exclude", required=False
+    )
+    parser_discover.add_argument(
+        "--technique-kwargs",
+        help="Keywords for the discovery technique. These should be specified as `arg1=value1 arg2=value2...`.",
+        nargs="*",
+        default=[],
+    )
+
+    args = main_parser.parse_args(args)
+
+    # Assume the user wants test adequacy if they're setting bootstrap_size
+    if getattr(args, "bootstrap_size", None) is not None:
+        args.adequacy = True
+    if getattr(args, "adequacy", False) and getattr(args, "bootstrap_size", None) is None:
+        # Need this here rather than a default value because otherwise the above always sets adequacy to True
+        args.bootstrap_size = 100
+
+    args.command = Command(args.command)
+    return args
 
 
 def main() -> None:
@@ -60,9 +205,14 @@ def main() -> None:
             logging.info("Discovering causal structure")
             # Need to reset index to allow for multiple files having the same index (i.e. starting at zero).
             # Otherwise you end up with duplicate indices, which causes problems further down the line
-            df = pd.concat([pd.read_csv(path) for path in args.data_paths]).reset_index()
+            df = pd.concat([read_dataframe(path) for path in args.data_paths]).reset_index()
             if args.variables:
                 df = df[args.variables]
+            # Drop unnamed columns
+            unnamed_columns = [c for c in df.columns if c.startswith("Unnamed: ")]
+            if unnamed_columns:
+                logger.warning(f"Dropping unnamed columns: {unnamed_columns}")
+            df = df.drop(unnamed_columns)
 
             discover_class = discover_map[args.technique].load()
             discover = discover_class(
@@ -80,55 +230,20 @@ def main() -> None:
             discover.write_dot(evolved_dag, args.output)
             logging.info("Causal structure discovery completed successfully.")
         case Command.TEST:
-            # Create paths object
-            paths = CausalTestingPaths(
+            # Create and setup framework
+            framework = CausalTestingFramework()
+
+            framework.setup(
                 dag_path=args.dag_path,
                 data_paths=args.data_paths,
-                test_config_path=args.test_config,
-                output_path=args.output,
+                test_cases_path=args.test_config,
+                query=args.query,
+                ignore_cycles=args.ignore_cycles,
             )
 
-            # Create and setup framework
-            framework = CausalTestingFramework(paths, ignore_cycles=args.ignore_cycles, query=args.query)
-            framework.setup()
-
-            # Load and run tests
-            framework.load_tests()
-
-            if args.batch_size > 0:
-                logging.info(f"Running tests in batches of size {args.batch_size}")
-                with tempfile.TemporaryDirectory() as tmpdir:
-                    output_files = []
-                    for i, results in enumerate(
-                        framework.run_tests_in_batches(
-                            batch_size=args.batch_size,
-                            silent=args.silent,
-                            adequacy=args.adequacy,
-                            bootstrap_size=args.bootstrap_size,
-                        )
-                    ):
-                        temp_file_path = os.path.join(tmpdir, f"output_{i}.json")
-                        framework.save_results(results, temp_file_path)
-                        output_files.append(temp_file_path)
-                        del results
-
-                    # Now stitch the results together from the temporary files
-                    all_results = []
-                    for file_path in output_files:
-                        with open(file_path, "r", encoding="utf-8") as f:
-                            all_results.extend(json.load(f))
-
-                    output_path = Path(args.output)
-                    output_path.parent.mkdir(parents=True, exist_ok=True)
-
-                    with open(args.output, "w", encoding="utf-8") as f:
-                        json.dump(all_results, f, indent=4)
-            else:
-                logging.info("Running tests in regular mode")
-                results = framework.run_tests(
-                    silent=args.silent, adequacy=args.adequacy, bootstrap_size=args.bootstrap_size
-                )
-                framework.save_results(results)
+            logging.info("Running tests")
+            framework.run_tests(silent=args.silent, adequacy=args.adequacy, bootstrap_size=args.bootstrap_size)
+            framework.save_results(args.output)
 
             logging.info("Causal testing completed successfully.")