Take post-storm profile based on isgood parameter in .mat file

src/data/parse_mat.py

@@ -198,6 +198,17 @@ def parse_profiles_and_sites(profiles_mat):
     site_rows = []
     site_counter = 0
 
+    # Our z values can come from these columns, depending on the isgood flag.
+    # Let's reoganise them into a list of list
+    z_names = ["Zpre", 'Zpost', 'Zrec1', 'Zrec2', 'Zrec3', 'Zrec4']
+    z_cols = [mat_data[col] for col in z_names]
+    z_sites = []
+    for cols in zip(*z_cols):
+        z_vals = []
+        for z_vector in zip(*cols):
+            z_vals.append([z[0] for z in z_vector])
+        z_sites.append(z_vals)
+
     for i, site in enumerate(mat_data["site"]):
         logger.debug("Processing site {} of {}".format(i + 1, len(mat_data["site"])))
 
@@ -215,27 +226,38 @@ def parse_profiles_and_sites(profiles_mat):
         # Want to calculation the orientation
         orientation = {}
 
-        for x, lat, lon, z_prestorm, z_poststorm, easting, northing in zip(
+
+
+        for x, lat, lon, z_site, easting, northing in zip(
             mat_data["x"][i],
             mat_data["lats"][i],
             mat_data["lons"][i],
-            mat_data["Zpre"][i],
-            mat_data["Zpost"][i],
+            z_sites[i],
             mat_data["eastings"][i],
             mat_data["northings"][i],
         ):
 
-            # Only extract pre and post storm profile
-            for j, profile_type in enumerate(["prestorm", "poststorm"]):
 
-                if mat_data["isgood"][i][j] == 1:
-                    land_lim = mat_data["landlims"][i][j]
-                    survey_datetime = matlab_datenum_to_datetime(mat_data["surveydates"][i][j])
+            profile_type = None
+            for j, is_good in enumerate([1] + mat_data["isgood"][i]):
+
+                # Assumes the first profile is always good and is the prestorm profike
+                if j == 0:
+                    profile_type = 'prestorm'
+                    z = z_site[j]
+                    land_lim = np.nan
 
-                    if profile_type == "prestorm":
-                        z = z_prestorm
+                # Skips bad profiles
+                elif is_good == 0:
+                    continue
+
+                # Takes the first isgood profile as the post storm profile
                 else:
-                        z = z_poststorm
+                    profile_type = 'poststorm'
+                    z = z_site[j]
+                    land_lim = mat_data["landlims"][i][j]
+
+                survey_datetime = matlab_datenum_to_datetime(mat_data["surveydates"][i][j])
 
                 # Keep a record of the where the center of the profile is located, and the locations of the land
                 # and sea
@@ -258,12 +280,17 @@ def parse_profiles_and_sites(profiles_mat):
                         "lat": lat[0],
                         "profile_type": profile_type,
                         "x": x[0],
-                            "z": z[0],
+                        "z": z,
                         "land_lim": land_lim,
                         "survey_datetime": survey_datetime,
                     }
                 )
 
+                # Stop looking at profiles if we've got our post-storm profile
+                if profile_type == 'poststorm':
+                    break
+
+
         orientation = math.degrees(
             math.atan2(
                 orientation["land_northing"] - orientation["sea_northing"],
@@ -306,8 +333,9 @@ def remove_zeros(df_profiles):
             df_profiles.index.get_level_values("profile_type") == key[1]
         )
         df_profile = df_profiles[idx_site]
-        x_last_ele = df_profile[df_profile.z != 0].index.get_level_values("x")[-1]
-        df_profiles.loc[idx_site & (df_profiles.index.get_level_values("x") > x_last_ele), "z"] = np.nan
+        x_last_ele = df_profile[df_profile.z == 0].index.get_level_values("x")[0]
+        df_profiles.loc[idx_site & (df_profiles.index.get_level_values("x") > x_last_ele),
+                        "z"] = np.nan
     logger.info("Removed zeros from end of profiles")
 
     return df_profiles