run_3_step_python.py

import json
import logging
import os
import re
import time
import argparse
from datetime import datetime
import numpy as np
import pandas as pd

from scipy.optimize import minimize

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

import fenics
import ufl

from ase import units as ase_units

from rdkit import Chem
from rdkit.Chem import Descriptors
from tqdm import tqdm

# FEniCS/logging configuration
logging.getLogger('FFC').setLevel(logging.WARNING)
logging.getLogger('UFL').setLevel(logging.WARNING)
fenics.set_log_level(fenics.LogLevel.ERROR)

# =============================================================================
# USER PARAMETERS — safe to modify for a new monomer/experimental setup
# =============================================================================

# Catalyst parameters (GC2 dissociation transition state, experimental)
GC2_deltaH_ts = 27.0   # kcal/mol    — enthalpy of GC2 dissociation TS
GC2_deltaS_ts = 13.0   # cal/(mol·K) — entropy of GC2 dissociation TS

# Additive parameters (TBP binding equilibrium, M06-2X/def2-TZVP toluene)
tbp_deltaH = 31.0      # kcal/mol    — enthalpy of TBP binding
tbp_deltaS = 56.7      # cal/(mol·K) — entropy of TBP binding

# Material properties of the monomer resin (defaults are values commonly used to model DCPD)
kappa = 0.15           # W/(m·K)     — thermal conductivity
rho   = 980.0          # kg/m³       — density
Cp    = 1600.0         # J/(kg·K)    — specific heat capacity

# Simulation domain
length         = 0.001  # m   — 1D domain length
element_size_x = 1e-6   # m   — spatial resolution (1 µm)
dt             = 1e-5   # s   — time step

# Initial conditions
T0_celsius = 25.0       # °C  — initial/ambient temperature
alpha0     = 0.0        # —   — initial conversion (0 = unreacted)

# =============================================================================
# FIXED PARAMETERS — derived constants, do not modify
# =============================================================================

KCAL_TO_J    = 4184.0
CAL_TO_J     = KCAL_TO_J / 1000
CELSIUS_TO_K = 273.15

GC2_deltaH_ts *= KCAL_TO_J
GC2_deltaS_ts *= CAL_TO_J
tbp_deltaH    *= KCAL_TO_J
tbp_deltaS    *= CAL_TO_J

kb_ = ase_units.kB * ase_units._e
h_  = ase_units._hplanck
R_  = ase_units.kB * ase_units._e * ase_units._Nav
_EYRING_PREFACTOR = kb_ / h_

num_elements_x = int(length / element_size_x)
T0 = CELSIUS_TO_K + T0_celsius

ALPHA_COMPLETION_THRESHOLD = 0.85
MIN_FRONT_NODES            = 5

def solve_PDE(conc, params, pde_steps, t_trig):
    """Run the 3-step kinetic + heat-diffusion PDE for a given monomer.

    Solves a 1D reaction-diffusion system on [0, length] using Euler time-stepping
    for the ODE kinetics and FEniCS CG1 elements for the heat equation.

    Kinetic scheme (Eyring rates):
        1. GC2 dissociation (irreversible): C_pcy3 → A + pcy3
        2. TBP equilibrium (analytical):    C_tbp  ⇌ A + tbp
        3. Initiation:                      A + M  → N
        4. Propagation:                     N + M  → N  (tracked as α)
        5. Retro/decomposition:             M      → Md

    A Dirichlet BC (T = T_trig) is applied at the left boundary for t < t_trig
    to trigger the front; after that the domain evolves freely.

    Early stop: loop exits once α > ALPHA_COMPLETION_THRESHOLD across the entire first half of the domain.

    Args:
        conc: [M0, C_pcy3_0, tbp_0] concentrations (mol/L, scaled).
        params: dict of kinetic/thermodynamic parameters (SI units).
        pde_steps: max number of time steps.
        t_trig: time (s) at which the trigger BC is released.

    Returns:
        T_list (°C), alpha_list, Md_list — each shape (actual_steps, nodes).
    """
    mesh = fenics.IntervalMesh(num_elements_x, 0.0, length)
    left = fenics.CompiledSubDomain("on_boundary && near(x[0], side, tol)",
                                    side=0.0, tol=1e-7)
    facets = fenics.MeshFunction("size_t", mesh, mesh.topology().dim() - 1, 0)
    facets.set_all(0)
    left.mark(facets, 1)
    V = fenics.FunctionSpace(mesh, "CG", 1)
    grad, dot, dx = ufl.grad, ufl.dot, ufl.dx

    M0 = conc[0]
    bc_left = fenics.DirichletBC(V, fenics.Constant(float(params["T_trig"])), facets, 1)
    t = 0.0
    vector_shape = V.dim()
    mid = vector_shape // 2
    T_prev = np.full(vector_shape, float(T0), dtype=np.float64)
    alpha_prev = np.full(vector_shape, float(alpha0), dtype=np.float64)

    initial_conc = {
        "M": conc[0],        # Monomer
        "C_pcy3": conc[1],   # Inactive catalyst (GC2)
        "C_tbp": 0.0,        # Inactive catalyst bound to TBP
        "A": 0.0,            # Active catalyst
        "pcy3": 0.0,         # Released PCy3 inhibitor
        "tbp": conc[2],      # TBP inhibitor
        "N": 0.0,            # Initiated complex
        "Md": 0.0,           # Decomposed monomer
        "alpha": alpha0      # Conversion
    }

    state = {k: np.full(vector_shape, v) for k, v in initial_conc.items()}

    # Only track the three outputs needed downstream
    T_tracked     = np.zeros((pde_steps, vector_shape))
    alpha_tracked = np.zeros((pde_steps, vector_shape))
    Md_tracked    = np.zeros((pde_steps, vector_shape))

    # Pre-allocate FEniCS objects — reused every step
    T_prev_f     = fenics.Function(V)
    alpha_f      = fenics.Function(V)
    alpha_prev_f = fenics.Function(V)
    T_new_f      = fenics.Function(V)
    Hr_f         = fenics.Constant(float(params["Hr"]))

    u = fenics.TrialFunction(V)
    v = fenics.TestFunction(V)

    # Bilinear form depends only on constants — assemble once
    a_form = (kappa * dot(grad(v), grad(u)) + rho * Cp / dt * v * u) * dx
    L_form = (rho * Cp / dt * v * T_prev_f
              + rho * Hr_f / dt * v * (alpha_f - alpha_prev_f)) * dx

    A_free = fenics.assemble(a_form)
    A_bc   = A_free.copy()
    bc_left.apply(A_bc)  # bake BC into matrix once

    pbar = tqdm(total=pde_steps)
    i = 0

    while i < pde_steps:
        # Eyring rates: k = (kb/h) * T * exp(-dH/RT) * exp(dS/R)
        K_tbp  = np.exp(-params["tbp_deltaH"] / (R_ * T_prev)) * np.exp(params["tbp_deltaS"]   / R_)  # TBP equilibrium constant
        kd     = _EYRING_PREFACTOR * T_prev * np.exp(-params["GC2_deltaH_ts"] / (R_ * T_prev)) * np.exp(params["GC2_deltaS_ts"] / R_)  # GC2 dissociation
        ki     = _EYRING_PREFACTOR * T_prev * np.exp(-params["Hi_eff"]        / (R_ * T_prev)) * np.exp(params["Si_eff"]        / R_)  # initiation
        kp     = _EYRING_PREFACTOR * T_prev * np.exp(-params["Hp_eff"]        / (R_ * T_prev)) * np.exp(params["Sp_eff"]        / R_)  # propagation
        kretro = _EYRING_PREFACTOR * T_prev * np.exp(-params["H_retro"]       / (R_ * T_prev)) * np.exp(params["S_retro"]       / R_)  # retro/decomposition

        # Reaction 1: GC2 dissociation C_pcy3 → A + pcy3 (irreversible, clamped to available C_pcy3)
        rate1f = np.minimum(kd * state["C_pcy3"] * dt, state["C_pcy3"])

        # TBP equilibrium shift (analytical): C_tbp ⇌ A + tbp
        tbp_shift = -0.5 * (state["A"] + state["tbp"] + K_tbp) + \
                     0.5 * ((state["A"] + state["tbp"] + K_tbp)**2 - 4.0 * ((state["A"] * state["tbp"]) - (K_tbp * state["C_tbp"])))**0.5
        tbp_shift = np.clip(tbp_shift, -np.minimum(state["A"], state["tbp"]), state["C_tbp"])

        # Reaction 2: Initiation A + M → N (clamped to available M and A)
        rate2 = np.minimum(ki * state["A"] * state["M"] * dt, np.minimum(state["M"], state["A"] + tbp_shift))

        # Reaction 3: Propagation N + M → N (clamped to available M and N)
        rate3 = np.minimum(kp * state["N"] * state["M"] * dt, np.minimum(state["M"], state["N"]))

        # Reaction 4: Retro/decomposition M → Md (clamped to available M)
        rate4 = np.minimum(kretro * state["M"] * dt, state["M"])

        # Update concentrations
        state["C_pcy3"] -= rate1f
        state["C_tbp"]  -= tbp_shift
        state["A"]      += rate1f + tbp_shift - rate2
        state["tbp"]    += tbp_shift
        state["pcy3"]   += rate1f
        state["M"]      -= rate2 + rate3 + rate4
        state["N"]      += rate2
        state["Md"]     += rate4
        state["alpha"]  += (rate2 + rate3) / M0

        # Store tracked outputs (reversed so left boundary = index 0)
        T_tracked[i, :]     = T_prev[::-1]
        alpha_tracked[i, :] = state["alpha"][::-1]
        Md_tracked[i, :]    = state["Md"][::-1]

        # Assemble RHS and solve — stiffness matrix A already assembled
        T_prev_f.vector().set_local(T_prev)
        alpha_f.vector().set_local(state["alpha"])
        alpha_prev_f.vector().set_local(alpha_prev)
        b = fenics.assemble(L_form)
        if t < t_trig:
            bc_left.apply(b)
            fenics.solve(A_bc, T_new_f.vector(), b, "cg", "hypre_amg")
        else:
            fenics.solve(A_free, T_new_f.vector(), b, "cg", "hypre_amg")
        pbar.update(1)

        # Early stop: front has swept the first half of the domain
        if np.all(state["alpha"][-mid:] > ALPHA_COMPLETION_THRESHOLD):
            break

        alpha_prev = state["alpha"].copy()
        T_prev = T_new_f.vector().get_local()
        t += dt
        i += 1

    pbar.close()

    return T_tracked[:i] - CELSIUS_TO_K, alpha_tracked[:i], Md_tracked[:i]
    
def plot_sol(sol, length, dt, title, kind="T", stride=None, zoom=False, save_path=None):
    """Plot spatial profiles over time with a copper colormap.

    Curves are colored by time (dark=early, light=late). Saves to file if
    save_path is given, otherwise displays interactively.

    Args:
        sol: array of shape (timesteps, nodes).
        length: domain length in metres.
        dt: timestep size in seconds.
        title: plot title string.
        kind: "T" for temperature (y-axis in °C) or "alpha" for conversion.
        stride: plot every nth timestep (default: ~100 curves total).
        zoom: if True, zoom y-axis to [1-1e-13, 1] (useful for near-complete alpha).
        save_path: path to save the figure; if None, show interactively.
    """
    ylabel = "T (°C)" if kind == "T" else "α"
    times = np.arange(len(sol)) * dt
    cmap = plt.cm.copper
    norm = mcolors.Normalize(vmin=times.min(), vmax=times.max())
    x_plot = np.linspace(0, length, sol.shape[1])
    if stride is None:
        stride = max(1, len(sol) // 100)

    fig, ax = plt.subplots(figsize=(6, 4))
    for i in range(0, len(sol), stride):
        ax.plot(x_plot, sol[i], color=cmap(norm(times[i])))
    ax.set_title(title)
    ax.set_xlabel("x (m)")
    ax.set_ylabel(ylabel)
    if zoom:
        ax.set_ylim([1 - 1e-13, 1.0])
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    fig.colorbar(sm, ax=ax, label="time (s)")
    fig.tight_layout()
    if save_path is not None:
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
    else:
        plt.show()
    
def get_reaction_completion_step(alpha_sol):
    """Return the first timestep at which all spatial nodes exceed α=0.85."""
    completed = np.where(np.all(np.array(alpha_sol) > ALPHA_COMPLETION_THRESHOLD, axis=1))[0]
    return completed[0] if len(completed) > 0 else len(alpha_sol)

def process_reaction_front(T_sol, alpha_sol, Md_sol, pde_steps):
    """Extract front propagation metrics from PDE solution arrays."""
    reaction_completion_step = get_reaction_completion_step(alpha_sol)
    last_increasing_index = max(reaction_completion_step-2, 1)
    slopes = np.diff(alpha_sol[last_increasing_index])

    max_neg_slope_index = np.argmin(slopes)
    max_neg_slope = np.min(slopes)

    y = alpha_sol[last_increasing_index][:max(max_neg_slope_index, MIN_FRONT_NODES)]
    point = fit_piecewise_breakpoint(np.arange(len(y)), y)
    point_idx = int(point)
    if point_idx < MIN_FRONT_NODES:  # front didn't form — reset to ambient
        alpha_front_mean = 0.0
        T_front_mean = 25.0  # ambient °C
    else:
        alpha_front_mean = np.mean(alpha_sol[last_increasing_index][:point_idx])
        T_front_mean = np.mean(T_sol[last_increasing_index][:point_idx])
    if max_neg_slope_index < 20:  # front peak near boundary — extend window
        last_increasing_index = min(last_increasing_index + 2, pde_steps-1)

    front_speed = 1000 * element_size_x * max_neg_slope_index / (reaction_completion_step * dt)  # m/s → mm/s
    return {
        'alpha_front_mean': alpha_front_mean,
        'T_front_mean': T_front_mean,
        'intercept_index': point,
        'max_front_timestep': last_increasing_index,
        'reaction_done_timestep': reaction_completion_step,
        'max_neg_slope_index': max_neg_slope_index,
        'max_neg_slope': max_neg_slope,
        "retro_prod": np.mean(np.max(Md_sol, axis=1)),
        "front_speed": front_speed
    }

def process_reaction_pot(T_sol, alpha_sol):
    """Extract pot-cure metrics (max T and max alpha spatially averaged)."""
    return {
        "T_pot": np.mean(np.max(T_sol, axis=0)),
        "alpha_pot": np.mean(np.max(alpha_sol, axis=0)),
    }

def piecewise_linear(x, x0, k1, k2, b1):
    """Two lines joined at x0, with slopes k1 and k2, and intercept b1 on the first segment."""
    return np.where(x < x0, k1 * x + b1, k2 * x + (k1 - k2) * x0 + b1)

def fit_piecewise_breakpoint(x, y):
    """Fits a piecewise linear function and returns the breakpoint x-coordinate."""
    x0_guess = x[len(x) // 2]
    p0 = [x0_guess, 0.0, 0.0, y[0]]  # [x0, k1, k2, b1]

    def objective(params):
        x0, k1, k2, b1 = params
        return np.sum((y - piecewise_linear(x, x0, k1, k2, b1)) ** 2)

    return minimize(objective, p0).x[0]

def process(mat, conc, condition, t_trig, columns, pde_steps, return_data=False, plot_dir=None):
    """Set up and run the PDE for one monomer row, returning result metrics.

    If conc == [0, 0, 0], reads concentrations from mat['exp_conc'] (JSON).
    If plot_dir is given, saves T and alpha spatial plots there.
    """
    mat = dict(zip(columns, mat))

    from_kcal = lambda key: float(mat[key]) * KCAL_TO_J   # kcal/mol → J/mol
    from_cal  = lambda key: float(mat[key]) * CAL_TO_J    # cal/(mol·K) → J/(mol·K)
    if list(conc) == [0.0, 0.0, 0.0]:
        conc = np.array(json.loads(mat["exp_conc"]))

    Hi_eff  = from_kcal("H_init")
    Si_eff  = from_cal("S_init")
    Hp_eff  = from_kcal("H_prop")
    Sp_eff  = from_cal("S_prop")
    H_retro = from_kcal("H_retro")
    S_retro = from_cal("S_retro")
    Hr = mat["Hr"] * 1000  # kJ/mol → J/mol
    m_conc = rho / mat['mol_wt']  # mol/L
    scaled_conc = np.array([1.0,  conc[1] / conc[0],  conc[2] / conc[0]]) * m_conc
    if condition == 'front':
        T_trig = T0 + Hr/Cp * (1-alpha0)
    else:
        T_trig = T0
        t_trig = 0

    params = {
        "tbp_deltaH": tbp_deltaH,  "tbp_deltaS": tbp_deltaS,
        "GC2_deltaH_ts": GC2_deltaH_ts, "GC2_deltaS_ts": GC2_deltaS_ts,
        "Hi_eff": Hi_eff,  "Si_eff": Si_eff,
        "Hp_eff": Hp_eff,  "Sp_eff": Sp_eff,
        "Hr": Hr,
        "H_retro": H_retro, "S_retro": S_retro,
        "T_trig": float(T_trig),
    }

    T_sol, alpha_sol, Md_sol = solve_PDE(scaled_conc, params, pde_steps, t_trig)
    if condition == 'front':
        result = process_reaction_front(T_sol, alpha_sol, Md_sol, pde_steps)
        result['retro_prod'] = result['retro_prod'] / m_conc  # mol/L → fraction of M0
        result['unreacted_monomer'] = 1.0 - result['alpha_front_mean'] - result['retro_prod']
    else:
        result = process_reaction_pot(T_sol, alpha_sol)

    if plot_dir is not None:
        os.makedirs(plot_dir, exist_ok=True)
        monomer_id = re.sub(r'[^\w]', '_', str(mat.get('monomer', 'monomer')))[:60]
        monomer_smiles = mat.get('monomer', '')
        plot_sol(T_sol, length, dt, monomer_smiles,
                 kind="T", save_path=os.path.join(plot_dir, f"{monomer_id}_T.png"))
        plot_sol(alpha_sol, length, dt, monomer_smiles,
                 kind="alpha", save_path=os.path.join(plot_dir, f"{monomer_id}_alpha.png"))

    if return_data:
        result.update({"T_sol": T_sol, "alpha_sol": alpha_sol})

    return result
    
def generate_filename(metadata, extra_label, extension='csv'):
    """Generate a timestamped output filename from run metadata."""
    date_str = datetime.now().strftime('%Y%m%d')
    prefix = 'pot' if metadata['condition'].lower() == 'pot' else 'front'
    if extra_label:
        prefix = prefix + "_" + extra_label
    conc_str = (
        f"{'-'.join([str(c) for c in metadata['conc']])}" 
        if isinstance(metadata['conc'], list) 
        else str(metadata['conc'])
    )

    return f"{prefix}_{np.round(metadata['T0'])}_c{conc_str}_n{metadata['pde_steps']}_{date_str}.{extension}"

def parse_args():
    parser = argparse.ArgumentParser(
        description="Run the 3-step kinetic front model for one or more monomers."
    )
    parser.add_argument("--monomers", required=True,
                        help="Path to monomer CSV (see README for required columns).")
    parser.add_argument("--condition", default="front", choices=["front", "pot"],
                        help="Simulation mode: 'front' (default) or 'pot' (isothermal pot cure).")
    parser.add_argument("--conc", type=float, nargs=3, default=[2500, 1, 1],
                        metavar=("M_0", "GC2_0", "TBP_0"),
                        help="Relative molar concentrations of monomer, GC2 catalyst, and TBP "
                             "additive (default: 2500 1 1). Only ratios matter; absolute scale "
                             "is set by monomer density.")
    parser.add_argument("--pde_steps", type=int, default=300000,
                        help="Maximum number of time steps (default: 300000). Simulation stops "
                             "early once the front has swept the domain.")
    parser.add_argument("--t_trig", type=float, default=0.1,
                        help="Duration (s) for which the left-boundary trigger temperature is "
                             "held before the domain evolves freely (default: 0.1 s).")
    parser.add_argument("--smiles_list", type=str, nargs="*", default=None,
                        help="Optional subset of SMILES to run (must exist in --monomers CSV). "
                             "If omitted, all rows are processed.")
    parser.add_argument("--label", default="run",
                        help="Label appended to output filenames (default: 'run').")
    parser.add_argument("--n_chunks", type=int, default=1,
                        help="Split monomers into N chunks for parallel SLURM array jobs.")
    parser.add_argument("--chunk_idx", type=int, default=0,
                        help="Which chunk this process should handle (0-indexed).")
    parser.add_argument("--plot", action="store_true",
                        help="Save T(x,t) and α(x,t) plots for each monomer.")
    parser.add_argument("--plot_dir", default="plots",
                        help="Directory to write plots into (default: 'plots/'). "
                             "Created automatically if it does not exist.")
    return parser.parse_args()

def main(monomers_csv, pde_steps=300000, t_trig=0.1, conc=None, condition="front",
         n_chunks=1, chunk_idx=0, label="run", smiles_list=None, plot_dir=None):
    """Load monomer CSV, run PDE for each row, and save results to CSV/JSON."""
    if conc is None:
        conc = [2500, 1, 1]   # M : GC2 : TBP molar ratio

    monomers = pd.read_csv(monomers_csv)

    monomers['mol_wt'] = [Descriptors.MolWt(Chem.MolFromSmiles(s)) for s in monomers['monomer']]
    monomers = monomers[monomers['Hr'] > 0]
    if smiles_list:
        monomers = monomers[monomers['monomer'].isin(smiles_list)].reset_index(drop=True)

    new_data = []
    chunks_to_process = np.array_split(monomers.values, n_chunks)
    chunk = chunks_to_process[chunk_idx]
    columns = monomers.columns

    for i, row in enumerate(chunk):
        t0 = time.perf_counter()
        print(f"[chunk {chunk_idx}] starting row {i}", flush=True)
        try:
            result = process(row, conc, condition, t_trig, columns, pde_steps,
                             plot_dir=plot_dir)
        except Exception as e:
            print(f"[chunk {chunk_idx}] row {i} failed: {e}", flush=True)
            result = {}
        dt_row = time.perf_counter() - t0
        print(f"[chunk {chunk_idx}] finished row {i} in {dt_row:.1f} s", flush=True)
        new_data.append({**dict(zip(columns, row)), **result})

    if len(new_data) > 0:
        new_df = pd.DataFrame(new_data)

        metadata = {
            'condition': condition,
            'conc': conc,
            'T0': T0,
            'pde_steps': pde_steps,
            'timestep': dt,
            'length': length,
            'element_size': element_size_x,
            'date': datetime.now().isoformat()
        }
        extra_label = f"{label}_chunk_{chunk_idx}" if n_chunks > 1 else label
        csv_name  = generate_filename(metadata, extra_label=extra_label, extension='csv')
        json_name = generate_filename(metadata, extra_label=extra_label, extension='json')
        with open(json_name, 'w') as f:
            json.dump(metadata, f, indent=2)
        new_df.to_csv(csv_name, index=False)

if __name__ == "__main__":
    args = parse_args()
    main(args.monomers, pde_steps=args.pde_steps, t_trig=args.t_trig, condition=args.condition,
         conc=args.conc, n_chunks=args.n_chunks, chunk_idx=args.chunk_idx,
         smiles_list=args.smiles_list, label=args.label,
         plot_dir=args.plot_dir if args.plot else None)