Rates_2step_ODE

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from tqdm import tqdm
import os
from scipy.integrate import solve_ivp

# Parameters
SM = 0.0
INT = 5.0
PROD = -2.0

TS1_values = np.linspace(5, 30, 40)
TS2_values = np.linspace(5, 40, 40)

R = 8.314e-3
T = 298.0
RT = R * T
A = 1.0

# ODE system
def reaction_odes(t, y, k1f, k1r, k2f, k2r):
    SM_conc, INT_conc, PROD_conc = y
    dSMdt = -k1f * SM_conc + k1r * INT_conc
    dINTdt = k1f * SM_conc - k1r * INT_conc - k2f * INT_conc + k2r * PROD_conc
    dPRODdt = k2f * INT_conc - k2r * PROD_conc
    return [dSMdt, dINTdt, dPRODdt]

y0 = [1.0, 0.0, 0.0]
t_span = (0, 1000.0)
t_eval = [1000.0]

TS1_grid, TS2_grid = np.meshgrid(TS1_values, TS2_values)
rate_grid = np.zeros_like(TS1_grid)

# Calculate rates
for i in tqdm(range(len(TS2_values)), desc="Rows"):
    for j in range(len(TS1_values)):
        TS1_val = TS1_grid[i, j]
        TS2_val = TS2_grid[i, j]

        if (TS1_val >= 5) and (TS2_val >= 5) and ((TS2_val - INT) > (TS1_val - SM)):
            k1f = A * np.exp(-(TS1_val - SM) / RT)
            k1r = A * np.exp(-(TS1_val - INT) / RT)
            k2f = A * np.exp(-(TS2_val - INT) / RT)
            k2r = A * np.exp(-(TS2_val - PROD) / RT)

            sol = solve_ivp(reaction_odes, t_span, y0, t_eval=t_eval, args=(k1f, k1r, k2f, k2r))
            if sol.success:
                SM_final, INT_final, PROD_final = sol.y[:, -1]
                rate = (k2f * INT_final) - (k2r * PROD_final)
            else:
                rate = np.nan
        else:
            rate = np.nan

        rate_grid[i, j] = rate

# Stats
valid_rates = rate_grid[~np.isnan(rate_grid)]
mean_rate = np.mean(valid_rates) if valid_rates.size > 0 else np.nan
std_rate = np.std(valid_rates) if valid_rates.size > 0 else np.nan
max_rate = np.nanmax(rate_grid)
min_rate = np.nanmin(rate_grid)

print("Summary Statistics:")
print(f"Mean rate: {mean_rate:.4e}")
print(f"Std dev of rate: {std_rate:.4e}")
print(f"Max rate: {max_rate:.4e}")
print(f"Min rate: {min_rate:.4e}")

# Plot
fig = plt.figure(figsize=(10, 7))
ax = fig.add_subplot(111, projection='3d')

masked_rate = np.ma.masked_invalid(rate_grid)
surf = ax.plot_surface(TS1_grid, TS2_grid, masked_rate, cmap=cm.viridis, edgecolor='none')
fig.colorbar(surf, shrink=0.5, aspect=5)

ax.set_xlabel('TS1 Free Energy')
ax.set_ylabel('TS2 Free Energy')
ax.set_zlabel('Net Rate')
ax.set_title('Rate Dependence (Assuming A=1)')

ax.invert_xaxis()
ax.invert_yaxis()
ax.view_init(elev=30, azim=210)

plt.tight_layout()

output_path = "3d_rate_plot_odes_reversible.png"
plt.savefig(output_path, dpi=300)
print(f"Plot saved as {os.path.abspath(output_path)}")

plt.show()