diff --git a/src/analysis/runup_models.py b/src/analysis/runup_models.py
index e088aeb..9176adb 100644
--- a/src/analysis/runup_models.py
+++ b/src/analysis/runup_models.py
@@ -2,52 +2,51 @@ import numpy as np
 import pandas as pd
 
 
-def sto06_individual(Hs0, Tp, beta):
+def sto06(Hs0, Tp, beta):
+    """
+    :param Hs0: List or float of offshore significant wave height values
+    :param Tp: List or float of peak wave period
+    :param beta: List of float of beach slope
+    :return: Float or list of R2, setup, S_total, S_inc and S_ig values
+    """
 
-    Lp = 9.8 * Tp ** 2 / 2 / np.pi
+    df = pd.DataFrame({"Hs0": Hs0, "Tp": Tp, "beta": beta}, index=[x for x in range(0, np.size(Hs0))])
 
-    S_ig = 0.06 * np.sqrt(Hs0 * Lp)
-    S_inc = 0.75 * beta * np.sqrt(Hs0 * Lp)
+    df["Lp"] = 9.8 * df['Tp'] ** 2 / 2 / np.pi
+
+    # General equation
+    df["S_ig"] = pd.to_numeric(0.06 * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
+    df["S_inc"] = pd.to_numeric(0.75 * df["beta"] * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
+    df["setup"] = pd.to_numeric(0.35 * df["beta"] * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
+    df["S_total"] = np.sqrt(df["S_inc"] ** 2 + df["S_ig"] ** 2)
+    df["R2"] = 1.1 * (df["setup"] + df["S_total"] / 2)
 
     # Dissipative conditions
-    if beta / (Hs0 / Lp) ** (0.5) <= 0.3:
-        setup = 0.016 * (Hs0 * Lp) ** 0.5
-        S_total = 0.046 * (Hs0 * Lp) ** 0.5
-        R2 = 0.043 * (Hs0 * Lp) ** 0.5
-    else:
-        setup = 0.35 * beta * (Hs0 * Lp) ** 0.5
-        S_total = np.sqrt(S_inc ** 2 + S_ig ** 2)
-        R2 = 1.1 * (setup + S_total / 2)
+    dissipative = df["beta"] / (df["Hs0"] / df["Lp"]) ** (0.5) <= 0.3
+
+    df.loc[dissipative, "setup"] = 0.016 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 16
+    df.loc[dissipative, "S_total"] = 0.046 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 17
+    df.loc[dissipative, "R2"] = 0.043 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 18
 
-    return R2, setup, S_total, S_inc, S_ig
+    return (
+        float_or_list(df["R2"].tolist()),
+        float_or_list(df["setup"].tolist()),
+        float_or_list(df["S_total"].tolist()),
+        float_or_list(df["S_inc"].tolist()),
+        float_or_list(df["S_ig"].tolist()),
+    )
 
 
-def sto06(df, Hs0_col, Tp_col, beta_col):
+def float_or_list(a):
     """
-    Vectorized version of Stockdon06 which can be used with dataframes
-    :param df:
-    :param Hs0_col:
-    :param Tp_col:
-    :param beta_col:
+    If only one value in the array, return the float, else return a list
+    :param a:
     :return:
     """
-
-    Lp = 9.8 * df[Tp_col] ** 2 / 2 / np.pi
-
-    # General equation
-    S_ig = pd.to_numeric(0.06 * np.sqrt(df[Hs0_col] * Lp), errors="coerce")
-    S_inc = pd.to_numeric(0.75 * df[beta_col] * np.sqrt(df[Hs0_col] * Lp), errors="coerce")
-    setup = pd.to_numeric(0.35 * df[beta_col] * np.sqrt(df[Hs0_col] * Lp), errors="coerce")
-    S_total = np.sqrt(S_inc ** 2 + S_ig ** 2)
-    R2 = 1.1 * (setup + S_total / 2)
-
-    # Dissipative conditions
-    dissipative = df[beta_col] / (df[Hs0_col] / Lp) ** (0.5) <= 0.3
-    setup.loc[dissipative, :] = 0.016 * (df[Hs0_col] * Lp) ** (0.5)  # eqn 16
-    S_total.loc[dissipative, :] = 0.046 * (df[Hs0_col] * Lp) ** (0.5)  # eqn 17
-    R2.loc[dissipative, :] = 0.043 * (df[Hs0_col] * Lp) ** (0.5)  # eqn 18
-
-    return R2, setup, S_total, S_inc, S_ig
+    if len(a) == 1:
+        return a[0]
+    else:
+        return list(a)
 
 
 if __name__ == "__main__":