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,54 +226,70 @@ 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:
+            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
+
+                # Skips bad profiles
+                elif is_good == 0:
+                    continue
+
+                # Takes the first isgood profile as the post storm profile
+                else:
+                    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])
-
-                    if profile_type == "prestorm":
-                        z = z_prestorm
-                    else:
-                        z = z_poststorm
-
-                    # Keep a record of the where the center of the profile is located, and the locations of the land
-                    # and sea
-
-                    # TODO: This code isn't very transferrable. What if we don't have lat/lons at 200 m? Relook at this
-                    if x[0] == 200:
-                        x_200_lat = lat[0]
-                        x_200_lon = lon[0]
-                    elif x[0] == 0:
-                        orientation["land_easting"] = easting[0]
-                        orientation["land_northing"] = northing[0]
-                    elif x[0] == 400:
-                        orientation["sea_easting"] = easting[0]
-                        orientation["sea_northing"] = northing[0]
-
-                    profile_rows.append(
-                        {
-                            "site_id": site_id,
-                            "lon": lon[0],
-                            "lat": lat[0],
-                            "profile_type": profile_type,
-                            "x": x[0],
-                            "z": z[0],
-                            "land_lim": land_lim,
-                            "survey_datetime": survey_datetime,
-                        }
-                    )
+
+                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
+
+                # TODO: This code isn't very transferrable. What if we don't have lat/lons at 200 m? Relook at this
+                if x[0] == 200:
+                    x_200_lat = lat[0]
+                    x_200_lon = lon[0]
+                elif x[0] == 0:
+                    orientation["land_easting"] = easting[0]
+                    orientation["land_northing"] = northing[0]
+                elif x[0] == 400:
+                    orientation["sea_easting"] = easting[0]
+                    orientation["sea_northing"] = northing[0]
+
+                profile_rows.append(
+                    {
+                        "site_id": site_id,
+                        "lon": lon[0],
+                        "lat": lat[0],
+                        "profile_type": profile_type,
+                        "x": x[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(
@@ -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