diff --git a/causal_testing/__main__.py b/causal_testing/__main__.py
index 3ecd443a..bd49bc6a 100644
--- a/causal_testing/__main__.py
+++ b/causal_testing/__main__.py
@@ -58,18 +58,31 @@ def main() -> None:
                 kwargs[split[0]] = split[1]
 
             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()
-            if args.variables:
-                df = df[args.variables]
+            exclude_edges = list(nx.nx_pydot.read_dot(args.exclude_edges).edges()) if args.exclude_edges is not None else []
+            
+            if args.context:
+                dfs = []
+                for i, path in enumerate(args.data_paths):
+                    temp_df = pd.read_csv(path)
+                    temp_df['file_index'] = i
+                    dfs.append(temp_df)
+
+                df = pd.concat(dfs, ignore_index=True)
+                
+                if args.variables:
+                    df = df[list(set(args.variables + ['file_index']))]
+
+                exclude_edges.append((".*", "file_index"))
+            else:
+                df = pd.concat((pd.read_csv(path) for path in args.data_paths), ignore_index=True)
+                
+                if args.variables:
+                    df = df[args.variables]
 
             discover_class = discover_map[args.technique].load()
             discover = discover_class(
                 df=df,
-                exclude_edges=(
-                    list(nx.nx_pydot.read_dot(args.exclude_edges).edges()) if args.exclude_edges is not None else []
-                ),
+                exclude_edges=exclude_edges,
                 include_edges=(
                     list(nx.nx_pydot.read_dot(args.include_edges).edges()) if args.include_edges is not None else []
                 ),
diff --git a/causal_testing/discovery/abstract_discovery.py b/causal_testing/discovery/abstract_discovery.py
index cb0c049e..f9835315 100644
--- a/causal_testing/discovery/abstract_discovery.py
+++ b/causal_testing/discovery/abstract_discovery.py
@@ -166,6 +166,8 @@ def write_dot(self, individual: CausalDAG, output_file: str):
                     else:
                         raise ValueError(f"Invalid test outcome {test['result']}")
 
+        if "file_index" in individual.nodes:
+            individual.remove_node("file_index")
         nx.drawing.nx_pydot.write_dot(individual, output_file)
 
     def _json_stub_params(self, outcome: str) -> str:
diff --git a/causal_testing/main.py b/causal_testing/main.py
index 9d5ed77e..f4c96bf3 100644
--- a/causal_testing/main.py
+++ b/causal_testing/main.py
@@ -590,6 +590,14 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
     # 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(
+        "-c",
+        "--context",
+        help="Combine data from multiple files into a single DataFrame using a context node rather than concatenation",
+        action="store_true",
+        default=False,
+        required=False,
+    )
     parser_discover.add_argument(
         "-a",
         "--alpha",