diff --git a/Trajectories/cost_function.py b/Trajectories/cost_function.py
index 44b404af3e445d7fc3483e3d637d3274178e0415..52bc7f3f94b6939abd2adc0715c6c73fab7c5e07 100755
--- a/Trajectories/cost_function.py
+++ b/Trajectories/cost_function.py
@@ -91,12 +91,12 @@ args = parser.parse_args()
 with open("e_overestim.txt") as f: e_overestim = int(f.read())
 
 # Set some values needed for the cost function:
-delta_cent_mean = 3.8481 # [in km]
+delta_cent_mean = 3.8481 # in km
 delta_cent_std = 8.0388
 delta_ro_mean = -0.0025965
 delta_ro_std = 5.2168
-delta_r_mean = -0.0094709 * 1000 # [in m]
-delta_r_std = 8.6953 * 1000
+delta_r_mean = -9.4709 # in m
+delta_r_std = 8.6953e3
 
 # Load the graph_tool file:
 
@@ -189,17 +189,14 @@ for edge in g.edges():
 
     # because of the wrapping issue (360° wrapping incorrectly to 0°),
     # we check for that here
-    lon_diff = abs(g.vp.pos_last[source_node][0] \
+    lon_diff = abs(g.vp.pos_last[source_node][0]
                    - g.vp.pos_first[target_node][0])
-    if (lon_diff > 300):
-        lon_diff = 360 - lon_diff
+    if lon_diff > 300: lon_diff = 360 - lon_diff
 
     Delta_Cent = math.sqrt((lon_diff * 111.32 * math.cos(lat_for_conv))**2
                            + ((g.vp.pos_last[source_node][1]
                                - g.vp.pos_first[target_node][1]) * 110.574)**2)
-
-    # calculate the first term
-    first_term = ((Delta_Cent - delta_cent_mean)/delta_cent_std) ** 2
+    first_term = ((Delta_Cent - delta_cent_mean)/delta_cent_std)**2
 
     # Rossbies:
     if (g.vp.first_av_ros[target_node] and g.vp.last_av_ros[source_node]):
@@ -210,16 +207,13 @@ for edge in g.edges():
         # Delta_Ro = delta_ro_mean
         Delta_Ro = 0
 
-    # Calculate the second term
-    second_term = ((Delta_Ro - delta_ro_mean)/delta_ro_std ) ** 2
+    second_term = ((Delta_Ro - delta_ro_mean)/delta_ro_std)**2
 
     # R_Vmax 1 and 2 already exist, just get the delta
-
     Delta_R_Vmax = g.vp.last_av_rad[source_node] \
         - g.vp.first_av_rad[target_node]
 
-    # Calculate the third term
-    third_term = ((Delta_R_Vmax - delta_r_mean)/delta_r_std) ** 2
+    third_term = ((Delta_R_Vmax - delta_r_mean)/delta_r_std)**2
 
     # Calculate the cost function and assign as weight to the edge:
     g.ep.cost_function[edge] = math.sqrt(first_term + second_term + third_term)
diff --git a/Trajectories/draw_segments.py b/Trajectories/draw_segments.py
index 90f7057759e8c490490dcb31541ce0fb8a8a1903..4f7f6c970129a517c354363f46964794e36684ef 100755
--- a/Trajectories/draw_segments.py
+++ b/Trajectories/draw_segments.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+"""Read a graph of segments with cost functions."""
+
 import graph_tool
 import sys
 import pygraphviz as pgv