diff --git a/probabilistic-analysis/probabilistic_assessment.py b/probabilistic-analysis/probabilistic_assessment.py
index 9315829..cc2efd0 100644
--- a/probabilistic-analysis/probabilistic_assessment.py
+++ b/probabilistic-analysis/probabilistic_assessment.py
@@ -486,18 +486,17 @@ def process(beach_name, beach_scenario, n_runs, start_year, end_year,
                     df_out.loc[idx, 'segment_gaps'] = True
 
                 # Specify which profiles to plot
-                df_out = df_out.assign(plot_stats=False)
-                for b, p in zip(plot_stats['block'], plot_stats['profile']):
+                df_out = df_out.assign(diagnostics=False)
+                for b, p in zip(diagnostics['block'], diagnostics['profile']):
                     idx = (df_out['block'] == str(b)) & (df_out['profile']
                                                          == p)
-                    df_out.loc[idx, 'plot_stats'] = True
+                    df_out.loc[idx, 'diagnostics'] = True
 
                 # Specify which profiles to omit from shapefiles
-                df_out = df_out.assign(omit_from_shp=False)
-                for b, p in zip(omit_from_shp['block'],
-                                omit_from_shp['profile']):
+                df_out = df_out.assign(omit=False)
+                for b, p in zip(omit['block'], omit['profile']):
                     idx = (df_out['block'] == b) & (df_out['profile'] == p)
-                    df_out.loc[idx, 'omit_from_shp'] = True
+                    df_out.loc[idx, 'omit'] = True
 
                 # Specify additional points to be included in shapefiles
                 df_out = df_out.assign(insert_points='').astype('object')