diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py
index 189a47b..673b917 100644
--- a/src/mdevaluate/extra/free_energy_landscape.py
+++ b/src/mdevaluate/extra/free_energy_landscape.py
@@ -103,7 +103,7 @@ def find_maxima(
     return maxima_df
 
 
-def calc_energies(maxima_indices, maxima_df, bins):
+def calc_energies(maxima_indices, maxima_df, bins, box):
     points = np.array(maxima_df[["x", "y", "z"]])
     tree = KDTree(points, boxsize=box)
     maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
@@ -125,7 +125,7 @@ def calc_energies(maxima_indices, maxima_df, bins):
         energy_hist = np.histogram(current_distances, bins=bins, weights=energy)[0]
         occupied_bins_hist = np.histogram(current_distances, bins=bins)[0]
         result = energy_hist / occupied_bins_hist
-        return r, result
+        return result
     for i, maxima_occupation in enumerate(maxima_occupations):
         current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
         current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
@@ -141,6 +141,22 @@ def calc_energies(maxima_indices, maxima_df, bins):
     return result
 
 
+def add_distances(occupation_df, pore_geometry, origin):
+    new_df = occupation_df.copy()
+    if pore_geometry is "cylindrical":
+        new_df["distance"] = (
+            (new_df["x"] - origin[0]) ** 2 + (new_df["y"] - origin[1]) ** 2
+        ) ** (1 / 2)
+    elif pore_geometry is "slit":
+        new_df["distance"] = np.abs(new_df["z"] - origin[2])
+    else:
+        raise ValueError(
+            f"pore_geometry is {pore_geometry}, should either be "
+            f"'cylindrical' or 'slit'"
+        )
+    return new_df
+
+
 def get_fel(
     traj,
     path,