diff --git a/Trajectories/Tests/perf_tests.json b/Trajectories/Tests/perf_tests.json
index 0cdc81de3fa7fb24ab284ee898341a5b5f6fd85c..1896728b4924c23ec17e3eac207a8c8a70d7d21f 100644
--- a/Trajectories/Tests/perf_tests.json
+++ b/Trajectories/Tests/perf_tests.json
@@ -15,14 +15,12 @@
 	"command":
 	[
 	    "$src_dir/Trajectories/cost_function.py",
-	    "$PWD/Global_1993_2020_SHPC", "Anticyclones"
+	    "$PWD/Global_1993_2020_SHPC", "Anticyclones",
+	    "$tests_old_dir/Perf_segments/segments.gt",
+	    "segments_cost_functions.gt"
 	],
 	"env": {"PYTHONPATH": "$src_dir/Overlap:$src_dir/Common"},
-	"symlink":
-	[
-	    "$PWD/Global_1993_2020_graph_anti/e_overestim.txt",
-	    "$tests_old_dir/Perf_segments/segments.gt"
-	]
+	"symlink": ["$PWD/Global_1993_2020_graph_anti/e_overestim.txt"]
     },
     {
 	"title": "Perf_trajectories",
diff --git a/Trajectories/Tests/tests.json b/Trajectories/Tests/tests.json
index a6d147e91cae9776b1bed4aeab6a742029f13f95..9041f1274958efaf49bee5e378cf4e1f96057b0d 100644
--- a/Trajectories/Tests/tests.json
+++ b/Trajectories/Tests/tests.json
@@ -48,16 +48,13 @@
 	"title": "Cost_functions",
 	"command":
 	[
-	    "$src_dir/Trajectories/cost_function.py", "--graphml",
+	    "$src_dir/Trajectories/cost_function.py",
 	    "$tests_old_dir/Extraction_eddies_region_2_noise/SHPC",
-	    "Anticyclones"
+	    "Anticyclones", "$tests_old_dir/Segments/segments.graphml",
+	    "segments_cost_functions.graphml"
 	],
 	"env": {"PYTHONPATH": "$src_dir/Overlap:$src_dir/Common"},
-	"symlink":
-	[
-	    "$tests_old_dir/EGr2nd15md4/e_overestim.txt",
-	    "$tests_old_dir/Segments/segments.graphml"
-	],
+	"symlink": ["$tests_old_dir/EGr2nd15md4/e_overestim.txt"],
 	"exclude_cmp": ["timings.txt"]
     },
     {
@@ -76,15 +73,13 @@
 	"This is a case where the speed is not defined on a whole segment.",
 	"command":
 	[
-	    "$src_dir/Trajectories/cost_function.py", "--graphml",
-	    "$tests_old_dir/Greece/SHPC", "Anticyclones"
+	    "$src_dir/Trajectories/cost_function.py",
+	    "$tests_old_dir/Greece/SHPC", "Anticyclones",
+	    "$tests_old_dir/Greece_segments/segments.graphml",
+	    "segments_cost_functions.graphml"
 	],
 	"env": {"PYTHONPATH": "$src_dir/Overlap:$src_dir/Common"},
-	"symlink":
-	[
-	    "$tests_old_dir/Greece_graph/e_overestim.txt",
-	    "$tests_old_dir/Greece_segments/segments.graphml"
-	]
+	"symlink": ["$tests_old_dir/Greece_graph/e_overestim.txt"]
     },
     {
 	"title": "Trajectories",
diff --git a/Trajectories/cost_function.py b/Trajectories/cost_function.py
index b7517db6b35bc1733963be32afe5387904337699..a176bc9a932ef10e122a6ead5b6d70467a200125 100755
--- a/Trajectories/cost_function.py
+++ b/Trajectories/cost_function.py
@@ -86,8 +86,10 @@ timings = open("timings.txt", "w")
 parser = argparse.ArgumentParser()
 parser.add_argument("SHPC_dir")
 parser.add_argument("orientation", choices = ["Anticyclones", "Cyclones"])
-parser.add_argument("--graphml", action = "store_true",
-                    help = "save to graphml format")
+parser.add_argument("input_segments", help = "input graph of segments without "
+                    "cost functions, suffix .gt (graph-tool) or .graphml")
+parser.add_argument("output_segments", help = "output graph of segments with "
+                    "cost functions, suffix .gt (graph-tool) or .graphml")
 args = parser.parse_args()
 with open("e_overestim.txt") as f: e_overestim = int(f.read())
 
@@ -102,12 +104,7 @@ delta_r_std = 8.6953e3
 # Load the graph_tool file:
 
 print('Loading graph...')
-
-try:
-    g = graph_tool.load_graph('segments.gt')
-except FileNotFoundError:
-    g = graph_tool.load_graph('segments.graphml')
-
+g = graph_tool.load_graph(args.input_segments)
 print('Loading done...')
 print("Input graph:")
 print("Number of vertices:", g.num_vertices())
@@ -223,12 +220,7 @@ t1 = time.perf_counter()
 timings.write(f"iterating on edges: {t1 - t0:.0f} s\n")
 t0 = t1
 print("Saving...")
-
-if args.graphml:
-    g.save('segments_cost_functions.graphml')
-else:
-    g.save('segments_cost_functions.gt')
-
+g.save(args.output_segments)
 print('All done')
 t1 = time.perf_counter()
 timings.write(f"saving: {t1 - t0:.0f} s\n")