optimizer.py

import math
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import r2_score
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split

class Protocol:
    def __init__(self):
        self.steps = []
class Step:
    def __init__(self, temp, time):
        self.temperature = temp
        self.duration = time
class A3aTool:
    def __init__(self, b=1.37755, c=0.0276528, d=0.0526316, e=0.1188, f=0.0174, g=0.1063, h=0.0356):
        self.dt = 0.1
        # Parameters for the kinetic model
        self.b = b
        self.c = c
        self.d = d
        self.e = e
        self.f = f
        self.g = g
        self.h = h
        self.target_conversion = 99.8
        self.objective = 5.0
        self.a3a_amount = 1.0 # not implemented yet
        self.dna_mass = 200
        self.max_step_time = 60
        self.temp_grid = [35, 37, 40, 42, 45, 50, 55, 60]
        self.time_grid = [2, 5, 10, 15, 30, 60, 180]
    def k_rate(self, T):
        return self.b
    def k_deact(self, T):
        return self.c*math.exp(self.d*T)
    def k_chh(self, T):
        return self.f*math.exp(self.e*T)
    def k_chg(self, T):
        return self.h*math.exp(self.g*T)
    # eventually add f_a3a() to add variability in a3a concentration
    def simulate(self, protocol, model):
        E = 1.0
        S = 1.0
        chh = 0
        chg = 0
        sim_chh = 0
        sim_chg = 0
        d_chh = 0
        d_chg = 0
        corr_chh = 0
        corr_chg = 0
        for step in protocol.steps:
            d_chh, d_chg = model.predict(step.temperature, step.duration)
            corr_chh += d_chh
            corr_chg += d_chg
            update = step.duration / self.dt
            for i in range(int(update)):
                rate = self.k_rate(step.temperature) * E * S
                sim_chh += self.k_chh(step.temperature) * rate * self.dt
                sim_chg += self.k_chg(step.temperature) * rate * self.dt
                E *= math.exp(-self.k_deact(step.temperature) * self.dt) # dE/dt = -k_deact(T) * E
                S -= rate * self.dt
                S = max(0, S)  # dS/dt = -k_rate(T) * E * S
        conversion = 1 - S
        print(sim_chh, corr_chh, sim_chg, corr_chg)
        chh = max(0, sim_chh + corr_chh)
        chg = max(0, sim_chg + corr_chg)
        return conversion * 100, chh, chg
    def optimize(self, time_or_off, model):
        best_protocol = None
        best_cpg = 0.0
        best_chg = float('inf')
        best_chh = float('inf')
        best_time = float('inf')
        best_objective = float('inf')
        for T1 in self.temp_grid:
            for T2 in self.temp_grid:
                if T2 < T1:
                    continue
                for T3 in self.temp_grid:
                    if T3 < T2:
                        continue
                    for t1 in self.time_grid:
                        for t2 in self.time_grid:
                            for t3 in self.time_grid:
                                protocol = Protocol()
                                if t1 > 0:
                                    protocol.steps.append(Step(T1, t1))
                                if t2 > 0:
                                    protocol.steps.append(Step(T2, t2))
                                if t3 > 0:
                                    protocol.steps.append(Step(T3, t3))
                                cpg, chh, chg = self.simulate(protocol, model)
                                total_time = t1 + t2 + t3
                                objective = chh + chg
                                if cpg >= self.target_conversion:
                                    if objective < best_objective:
                                        best_protocol = protocol
                                        best_cpg = cpg
                                        best_chh = chh
                                        best_chg = chg
                                        best_objective = objective
                                        best_time = total_time
                                    elif (abs(objective - best_objective) < 1e-6 and total_time < best_time):
                                        best_protocol = protocol
                                        best_cpg = cpg
                                        best_chh = chh
                                        best_chg = chg
                                        best_time = total_time
        return best_protocol, best_cpg, best_chh, best_chg, best_time
class ResidualLayer:
    def __init__(self):
        self.chh_model = None
        self.chg_model = None
        self.X = None
        self.y_chh = None
        self.y_chg = None
    def train(self, df):
        self.X = df[["Temp", "Time"]].to_numpy()
        self.y_chh = df["chh_res"]
        self.y_chg = df["chg_res"]
        self.chh_model = RandomForestRegressor(n_estimators=100, max_depth=3, random_state=42)
        self.chg_model = RandomForestRegressor(n_estimators=100, max_depth=3, random_state=42)
        self.chh_model.fit(self.X, self.y_chh)
        self.chg_model.fit(self.X, self.y_chg)
    def predict(self, temp, time):
        X = np.array([[temp, time]])
        d_chh = self.chh_model.predict(X)[0]
        d_chg = self.chg_model.predict(X)[0]
        return d_chh, d_chg
    def test(self):
        pred_chh = self.chh_model.predict(self.X)
        print("CHH R2: ", r2_score(self.y_chh, pred_chh))
        pred_chg = self.chg_model.predict(self.X)
        print(" CHG R2: ", r2_score(self.y_chg, pred_chg))
    def hybrid_fit_quality(self, df):
        test_rows = self.X.index
        test_df = df.loc[test_rows]
        hybrid_chh = []
        hybrid_chg = []
        for _, row in test_df.iterrows():
            delta_chh, delta_chg = self.predict(row["Temp"], row["Time"])
            hybrid_chh.append(row["chh_pred"] + delta_chh)
            hybrid_chg.append(row["chg_pred"] + delta_chg)
        print("hyb chh r2: ", r2_score(test_df["CHH"], hybrid_chh))
        print("hyb chg r2: ", r2_score(test_df["CHG"], hybrid_chg))
        print("ode chh r2: ", r2_score(test_df["CHH"], test_df["chh_pred"]))
        print("ode chg r2: ", r2_score(test_df["CHG"], test_df["chg_pred"]))
def loss(tool, dataframe):
    error_cpg = 0
    error_chh = 0
    error_chg = 0
    for _, row in dataframe.iterrows():
        step = Step(temp=row["Temp"], time=row["Time"])
        protocol = Protocol()
        protocol.steps.append(step)
        pred_cpg, pred_chh, pred_chg = tool.simulate(protocol, model=None)
        error_cpg += (pred_cpg - row["Conversion"])**2
        error_chh += (pred_chh - row["CHH"])**2
        error_chg += (pred_chg - row["CHG"])**2
    return error_cpg, error_chh, error_chg
def a_b_param_est(tool, dataframe):
    best_loss = float("inf")
    best_b = tool.b
    print("starting k_rate() est")
    for b in np.linspace(0.5, 1.5, 50):
        tool.b = b
        current_loss, _, _ = loss(tool, dataframe)
        if current_loss < best_loss:
            best_loss = current_loss
            best_b = b
    tool.b = best_b
    return best_b, best_loss
def c_d_param_est(tool, dataframe):
    best_loss = float("inf")
    best_c = tool.c
    best_d = tool.d
    print("starting k_deact() est")
    for c in np.logspace(np.log10(0.003), np.log10(0.05), 20):
        for d in np.linspace(0.04, 0.08, 20):
            tool.c = c
            tool.d = d
            current_loss, _, _ = loss(tool, dataframe)
            if current_loss < best_loss:
                best_loss = current_loss
                best_c = c
                best_d = d
    tool.c = best_c
    tool.d = best_d
    return best_c, best_d, best_loss
def e_f_param_est(tool, dataframe):
    best_loss = float("inf")
    best_e = tool.e
    best_f = tool.f
    print("starting k_chh() est")
    for e in np.linspace(0.1, 0.13, 20):
        for f in np.linspace(0.005, 0.025, 20):
            tool.e = e
            tool.f = f
            _, current_loss, _ = loss(tool, dataframe)
            if current_loss < best_loss:
                best_loss = current_loss
                best_e = e
                best_f = f
    tool.e = best_e
    tool.f = best_f
    return best_e, best_f, best_loss
def g_h_param_est(tool, dataframe):
    best_loss = float("inf")
    best_g = tool.g
    best_h = tool.h
    print("starting k_chg() est")
    for g in np.linspace(0.09, 0.12, 20):
        for h in np.linspace(0.01, 0.045, 20):
            tool.g = g
            tool.h = h
            _, _, current_loss = loss(tool, dataframe)
            if current_loss < best_loss:
                best_loss = current_loss
                best_g = g
                best_h = h
    tool.g = best_g
    tool.h = best_h
    return best_g, best_h, best_loss
def param_est():
    tool = A3aTool(0.7, 0.007, 0.064)
    df_global = pd.read_csv("/Users/nxanthopoulos/Desktop/A3A_offtarget.csv")
    rel_change = 0
    prev_loss = float("inf")
    for iteration in range(20):
        e, f,_ = e_f_param_est(tool, df_global)
        g, h,_ = g_h_param_est(tool, df_global)
        _, chh_loss, chg_loss = loss(tool, df_global)
        global_loss = chh_loss + chg_loss
        print(f"iter={iteration}", f"loss={global_loss:.2f}", f"e={e:.4f}", f"f={f:.6f}", f"g={g:.6f}", f"h={h:.6f}")
        if prev_loss != float("inf"):
            rel_change = (abs(prev_loss - global_loss) / prev_loss)
            if rel_change < 0.001:
                break
        prev_loss = global_loss
    return tool       
def experiment():
    tool = A3aTool()
    model = ResidualLayer()
    print("==============================")
    print("A3A Protocol Tool")
    print("==============================")
    print(f"")
    sim_or_opt = input("Protocol Simulator or Optimizer (S or O): ")
    return sim_or_opt, tool, model
def print_residuals():
    temps = [35, 36.7, 39.6, 44.5, 50.1, 54.6, 57.9, 60]
    times = [5, 10, 30, 60, 180]
    obs_df = pd.read_csv("/Users/nxanthopoulos/Desktop/A3A_param_est.csv")
    train_df = pd.read_csv("/Users/nxanthopoulos/Desktop/A3A_res_calc.csv")
    obs_df["chh_pred"] = np.nan
    obs_df["chg_pred"] = np.nan
    obs_df['hybrid_chh'] = np.nan
    obs_df['hybrid_chg'] = np.nan
    obs_df['sim_chh_res'] = np.nan
    obs_df['sim_chg_res'] = np.nan
    obs_df['model_chh_res'] = np.nan
    obs_df['model_chg_res'] = np.nan
    obs_df['cpg_pred'] = np.nan
    tool = A3aTool()
    model = ResidualLayer()
    model.train(train_df)
    row = 0
    for temp in temps:
        for time in times:
            protocol = Protocol()
            protocol.steps.append(Step(temp, time))
            cpg, chh, chg = tool.simulate(protocol)
            delta_chh, delta_chg = model.predict(temp, time)
            obs_df.loc[row, "cpg_pred"] = cpg
            obs_df.loc[row, "chh_pred"] = chh
            obs_df.loc[row, "chg_pred"] = chg
            obs_df.loc[row, "hybrid_chh"] = chh + delta_chh
            obs_df.loc[row, "hybrid_chg"] = chg + delta_chg
            obs_df.loc[row, 'sim_chh_res'] = obs_df.loc[row, 'CHH'] - chh
            obs_df.loc[row, 'sim_chg_res'] = obs_df.loc[row, 'CHG'] - chg
            obs_df.loc[row, 'model_chh_res'] = obs_df.loc[row, 'CHH'] - (chh + delta_chh)
            obs_df.loc[row, 'model_chg_res'] = obs_df.loc[row, 'CHG'] - (chg + delta_chg)
            row += 1
    obs_df.to_csv("/Users/nxanthopoulos/Desktop/A3A_model_accuracy.csv", index=False)
    print("chh hybrid r2: ", r2_score(obs_df["CHH"], obs_df["hybrid_chh"]))
    print("chh sim r2: ", r2_score(obs_df["CHH"], obs_df["chh_pred"]))
    print("chg hybrid r2: ", r2_score(obs_df["CHG"], obs_df["hybrid_chg"]))
    print("chg sim r2: ", r2_score(obs_df["CHG"], obs_df["chg_pred"]))
def main():
    sim_or_opt, tool, model = experiment()
    df = pd.read_csv("/Users/nxanthopoulos/Desktop/A3A_res_calc.csv")
    model.train(df)
    if sim_or_opt == "S":
        print(f"Propose a three-step protocol:")
        temps = []
        times = []
        temps.append(int(input("Step 1 temperature (°C): ")))
        times.append(int(input("Step 1 duration (minutes): ")))
        temps.append(int(input("Step 2 temperature (°C): ")))
        times.append(int(input("Step 2 duration (minutes): ")))
        temps.append(int(input("Step 3 temperature (°C): ")))
        times.append(int(input("Step 3 duration (minutes): ")))
        step1 = Step(temps[0], times[0])
        step2 = Step(temps[1], times[1])
        step3 = Step(temps[2], times[2])
        protocol = Protocol()
        protocol.steps.append(step1)
        protocol.steps.append(step2)
        protocol.steps.append(step3)
        print(f"Current protocol: ")
        print(f"Step 1: {protocol.steps[0].temperature}°C for {protocol.steps[0].duration} minutes")
        print(f"Step 2: {protocol.steps[1].temperature}°C for {protocol.steps[1].duration} minutes")
        print(f"Step 3: {protocol.steps[2].temperature}°C for {protocol.steps[2].duration} minutes")
        print(f"")
        print(f"Simulating current protocol...")
        cpg, chh, chg = tool.simulate(protocol, model)
        print(f"CPG: {cpg:.2f}")
        print(f"CHG: {chg:.2f}")
        print(f"CHH: {chh:.2f}")
        print(f"Total time: {protocol.steps[0].duration + protocol.steps[1].duration + protocol.steps[2].duration} minutes")
    elif sim_or_opt == "O":
        time_or_off = input("Minimize protocol time or off-target conversion (TIME or OFF): ")
        print(f"Optimizing protocol...")
        print(f"")
        if time_or_off == "OFF":
            best_protocol, best_cpg, best_chg, best_chh, best_time = tool.optimize(time_or_off, model)
            if best_protocol is None:
                print("No protocol found that meets the target conversion and off-target constraints.")
                return
            print(f"Best protocol: ")
            for i, step in enumerate(best_protocol.steps):
                print(f"Step {i+1}: {step.temperature}°C for {step.duration} minutes")
            print(f"Best CPG: {best_cpg:.2f}")
            print(f"Best CHG: {best_chg:.2f}")
            print(f"Best CHH: {best_chh:.2f}")
            print(f"Best time: {best_time} minutes")
if __name__ == "__main__":
    main()