correct_buds cleans buddings and bud_volume

correct_buds.py

+import matplotlib.pylab as plt
+import numpy as np
+from scipy.interpolate import interp1d
+
+min_cell_cycle_in_tps = 3
+
+
+def drop_label(df, label="mother_label"):
+    """Drop label from index of data frame."""
+    df = df.reset_index(label).drop(columns=[label])
+    return df
+
+
+def skip_budding(buddings, bud_volume):
+    """
+    Skip a budding event if it is too close to another.
+
+    Two adjacent buddings events are implausible.
+    We ignore the one corresponding to the larger bud volume.
+    """
+    new_buddings = buddings.copy()
+    new_bud_volume = bud_volume.copy()
+    for ind in buddings.index:
+        ibuddings = np.concatenate(
+            (np.where(buddings.loc[ind] > 0)[0], [buddings.size])
+        )
+        for ib, ib_end in zip(ibuddings, ibuddings[1:]):
+            if ib_end < ib + min_cell_cycle_in_tps:
+                print(" Skipping budding event for", ind)
+                try:
+                    vol = bud_volume.loc[ind].values.astype("float")
+                except KeyError:
+                    continue
+                # ignore budding event with larger volume
+                if vol[ib] > vol[ib_end]:
+                    # ignore earlier budding
+                    new_buddings.loc[ind].iloc[ib] = 0
+                    new_bud_volume.loc[ind].iloc[ib] = np.nan
+                else:
+                    # ignore later budding
+                    new_buddings.loc[ind].iloc[ib_end] = 0
+    return new_buddings, new_bud_volume
+
+
+def interp_bud_volume(ind, buddings, bud_volume, figs=True):
+    """Use linear interpolation to fill NaN in time series of bud volumes."""
+    ibuddings = np.where(buddings.loc[ind] > 0)[0]
+    ibuddings_end = np.concatenate((ibuddings, [buddings.size]))
+    success = False
+    if np.any(ibuddings):
+        try:
+            bv = bud_volume.loc[ind]
+        except KeyError:
+            print(f" No bud volume for {str(ind)}")
+            return success, np.nan
+        t = bv.index.values.astype("float")
+        vol = bv.values.astype("float")
+        # interpolate to fix
+        new_bv = bv.copy()
+        for ib, ib_end in zip(ibuddings_end, ibuddings_end[1:]):
+            if ib_end > ib + min_cell_cycle_in_tps:
+                # interpolate to fix any NaN
+                bud_vol = vol[ib:ib_end]
+                bud_t = t[ib:ib_end]
+                # find last non-NaN data point
+                ib_real = np.max(np.where(~np.isnan(bud_vol))[0])
+                t_bv = bud_t[: ib_real + 1][~np.isnan(bud_vol[: ib_real + 1])]
+                v_bv = bud_vol[: ib_real + 1][~np.isnan(bud_vol[: ib_real + 1])]
+                int_bud = interp1d(t_bv, v_bv)
+                # interpolate missing NaN within bud time series
+                new_bv[ib : ib + ib_real + 1] = int_bud(t[ib : ib + ib_real + 1])
+        # plot
+        if figs:
+            plot_buds(ind, buddings, bud_volume, new_bv)
+        # return interpolated values
+        success = True
+        return success, new_bv
+    else:
+        # no buddings
+        return success, np.nan
+
+
+def plot_buds(ind, buddings, bud_volume, new_bv=False):
+    """
+    Plot bud volume showing budding events for a given index.
+
+    Optional new_bv gives interpolated bud volumes for index ind.
+    """
+    ibuddings = np.where(buddings.loc[ind] > 0)[0]
+    bv = bud_volume.loc[ind]
+    plt.figure()
+    for ib in ibuddings:
+        # budding events
+        plt.plot(bv.index[ib], bv.values[ib], "ro")
+    bs = np.concatenate(([0], ibuddings, [buddings.size]))
+    for j, (start, end) in enumerate(zip(bs, bs[1:])):
+        if np.any(new_bv):
+            # plot interpolated data differently
+            plt.plot(new_bv[start:end], "b--", alpha=0.5)
+        plt.plot(bv[start:end], "b.-")
+    plt.xlabel("time")
+    plt.ylabel("bud volume")
+    plt.grid()
+    plt.title(str(ind))
+    plt.show(block=False)
+
+
+def correct_buds(buddings, bud_volume, figs=True):
+    # remove mother label
+    bud_volume = drop_label(bud_volume)
+    # check bud volumes and buddings
+    new_buddings, new_bud_volume = skip_budding(buddings, bud_volume)
+    for ind in new_buddings.index:
+        success, new_bv = interp_bud_volume(
+            ind,
+            new_buddings,
+            new_bud_volume,
+            figs,
+        )
+        if success:
+            new_bud_volume.loc[ind] = new_bv
+    return new_buddings, new_bud_volume