""" Extended mechanism data generation for ECFlow. Generates training data for 7 additional electrochemical mechanisms beyond the original 6 (Nernst, BV, MHC, Ads, EC, LH) in generate_dataset_diffec.py. Extended mechanisms (from EC-CVAE): 6: EE - Two sequential electron transfers: A -e-> B -e-> C 7: EC_prime - Catalytic EC: A -e-> B, B -> A with rate kc 8: CE - Chemical-Electrochemical: Y -> A (kf, Keq), A -e-> B 9: ECE - A -e-> B, B -> C (kc), C -e-> D Mixed mechanisms (from EC-CVAE): 10: EC_LH - Langmuir coverage-modulated BV + chemical decay 11: MHC_EC - MHC rate theory + chemical follow-up 12: MHC_LH - MHC rate theory + Langmuir coverage Output format matches ECFlow's multi-scan convention: data_extended/{train,val,test}/sample_NNNNNN.npz Usage: python generate_extended_mechanisms.py --output_dir data_extended --n_samples 10000 python generate_extended_mechanisms.py --mechanisms EE CE --n_samples 5000 python generate_extended_mechanisms.py --test """ import os import json import argparse import numpy as np import scipy.linalg from tqdm import tqdm from multiprocessing import Pool, cpu_count import warnings warnings.filterwarnings("ignore", message=".*[Ii]ll.conditioned.*") # ============================================================================= # Simulation constants (must match ECFlow's generate_dataset_diffec.py) # ============================================================================= DELTA_X = 2e-6 DELTA_THETA = 5e-2 EXPANDING_GRID_FACTOR = 1.05 SIMULATION_SPACE_MULTIPLE = 6.0 N_SPATIAL_OUT = 64 # ============================================================================= # Mechanism registry # ============================================================================= MECHANISM_LIST = ['EE', 'EC_prime', 'CE', 'ECE', 'EC_LH', 'MHC_EC', 'MHC_LH'] MECHANISM_TO_ID = { 'EE': 6, 'EC_prime': 7, 'CE': 8, 'ECE': 9, 'EC_LH': 10, 'MHC_EC': 11, 'MHC_LH': 12, } # ============================================================================= # Grid and helper functions (consistent with ECFlow) # ============================================================================= def gen_grid(Xi, deltaX, maxX, expanding_grid_factor): n = 1 current_X = Xi dX = deltaX while current_X < maxX: current_X += dX dX *= expanding_grid_factor n += 1 X_grid = np.zeros(n) X_grid[0] = Xi dX = deltaX for i in range(1, n): X_grid[i] = X_grid[i - 1] + dX dX *= expanding_grid_factor return X_grid, n def resample_concentration(X_grid, conc_history, n_spatial_out): n_sim = len(X_grid) indices = np.linspace(0, n_sim - 1, n_spatial_out).astype(int) x_grid_out = X_grid[indices] if conc_history.ndim == 1: return x_grid_out, conc_history[indices] return x_grid_out, conc_history[:, indices] def ini_conc(n, C_A_bulk, C_B_bulk): conc = np.zeros(2 * n) conc[:n] = C_A_bulk conc[n:] = C_B_bulk return conc, conc.copy() def ini_coeff(n): A_matrix = np.zeros((2 * n, 2 * n)) aA, bA, cA = np.zeros(n), np.zeros(n), np.zeros(n) aB, bB, cB = np.zeros(n), np.zeros(n), np.zeros(n) return A_matrix, aA, bA, cA, aB, bB, cB def calc_abc_linear(n, X_grid, deltaT, a, b, c, D): for i in range(1, n - 1): dX_m = X_grid[i] - X_grid[i - 1] dX_p = X_grid[i + 1] - X_grid[i] a[i] = D * (-2.0 * deltaT) / (dX_m * (dX_m + dX_p)) c[i] = D * (-2.0 * deltaT) / (dX_p * (dX_m + dX_p)) b[i] = 1.0 - a[i] - c[i] return a, b, c def calc_flux(conc, n, dA, X_grid): return -dA * (conc[n - 2] - conc[n - 1]) / (X_grid[1] - X_grid[0]) def _make_potential_waveform(theta_i, theta_v, cycles=1): nTimeSteps = int(2 * abs(theta_v - theta_i) / DELTA_THETA) + 1 Esteps = np.arange(nTimeSteps) E = np.where( Esteps < nTimeSteps / 2.0, theta_i - DELTA_THETA * Esteps, theta_v + DELTA_THETA * (Esteps - nTimeSteps / 2.0), ) return np.tile(E, cycles) def _clamp_bv_rates(K_red, K_ox, max_rate=1e6): return min(K_red, max_rate), min(K_ox, max_rate) # ============================================================================= # MHC rate constants (Gauss-Hermite quadrature) # ============================================================================= def _mhc_integrand_red(Theta_eff, reorg_e, degree=50): pts, wts = np.polynomial.hermite.hermgauss(degree) arg = -(reorg_e * (pts * 2.0 / np.sqrt(reorg_e) - 1.0) - Theta_eff) arg = np.clip(arg, -500, 500) y = 2 * np.sqrt(reorg_e) / (1.0 + np.exp(arg)) return np.sum(wts * y) def _mhc_integrand_ox(Theta_eff, reorg_e, degree=50): pts, wts = np.polynomial.hermite.hermgauss(degree) arg = -reorg_e * (pts * 2.0 / np.sqrt(reorg_e) - 1.0) - Theta_eff arg = np.clip(arg, -500, 500) y = -2.0 * np.sqrt(reorg_e) / (1.0 + np.exp(arg)) return np.sum(wts * y) def calc_mhc_rates(Theta, K0, reorg_e): I_red = _mhc_integrand_red(Theta, reorg_e) I_red0 = _mhc_integrand_red(0.0, reorg_e) I_ox = _mhc_integrand_ox(Theta, reorg_e) I_ox0 = _mhc_integrand_ox(0.0, reorg_e) K_red = K0 * I_red / I_red0 if abs(I_red0) > 1e-30 else 0.0 K_ox = K0 * I_ox / I_ox0 if abs(I_ox0) > 1e-30 else 0.0 return K_red, K_ox def _build_matrix_with_rates(A_matrix, X_grid, n, aA, bA, cA, dA, aB, bB, cB, dB, K_red, K_ox): """Build coefficient matrix given pre-computed rate constants.""" A_matrix[:] = 0.0 rows_A = np.arange(n - 2, 0, -1) A_matrix[rows_A, rows_A - 1] = cA[1:n - 1] A_matrix[rows_A, rows_A] = bA[1:n - 1] A_matrix[rows_A, rows_A + 1] = aA[1:n - 1] rows_B = np.arange(n + 1, 2 * n - 1) A_matrix[rows_B, rows_B - 1] = aB[1:n - 1] A_matrix[rows_B, rows_B] = bB[1:n - 1] A_matrix[rows_B, rows_B + 1] = cB[1:n - 1] X0 = X_grid[1] - X_grid[0] A_matrix[n - 1, n - 2] = -1.0 A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red A_matrix[n - 1, n] = -X0 / dA * K_ox A_matrix[n, n - 1] = -X0 / dB * K_red A_matrix[n, n] = 1.0 + X0 / dB * K_ox A_matrix[n, n + 1] = -1.0 A_matrix[0, 0] = 1.0 A_matrix[2 * n - 1, 2 * n - 1] = 1.0 return A_matrix # ============================================================================= # Simulator: EE (two sequential electron transfers) # ============================================================================= def run_ee_simulation( sigma, K0_1, alpha_1, K0_2, alpha_2, E0_2_offset, dA=1.0, dB=1.0, dC=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, C_C_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): beta_1 = 1.0 - alpha_1 beta_2 = 1.0 - alpha_2 deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) aA, bA, cA = np.zeros(n), np.zeros(n), np.zeros(n) aB, bB, cB = np.zeros(n), np.zeros(n), np.zeros(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) A_matrix = np.zeros((2 * n, 2 * n)) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] for idx in range(total_steps): Theta1 = E[idx] Theta2 = E[idx] - E0_2_offset K_red_1 = K0_1 * np.exp(np.clip(-alpha_1 * Theta1, -500, 500)) K_ox_1 = K0_1 * np.exp(np.clip(beta_1 * Theta1, -500, 500)) K_red_1, K_ox_1 = _clamp_bv_rates(K_red_1, K_ox_1) K_red_2 = K0_2 * np.exp(np.clip(-alpha_2 * Theta2, -500, 500)) K_ox_2 = K0_2 * np.exp(np.clip(beta_2 * Theta2, -500, 500)) K_red_2, K_ox_2 = _clamp_bv_rates(K_red_2, K_ox_2) A_matrix[:] = 0.0 X0 = X_grid[1] - X_grid[0] rows_A = np.arange(n - 2, 0, -1) A_matrix[rows_A, rows_A - 1] = cA[1:n - 1] A_matrix[rows_A, rows_A] = bA[1:n - 1] A_matrix[rows_A, rows_A + 1] = aA[1:n - 1] rows_B = np.arange(n + 1, 2 * n - 1) A_matrix[rows_B, rows_B - 1] = aB[1:n - 1] A_matrix[rows_B, rows_B] = bB[1:n - 1] A_matrix[rows_B, rows_B + 1] = cB[1:n - 1] A_matrix[n - 1, n - 2] = -1.0 A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red_1 A_matrix[n - 1, n] = -X0 / dA * K_ox_1 A_matrix[n, n - 1] = -X0 / dB * K_red_1 A_matrix[n, n] = 1.0 + X0 / dB * K_ox_1 + X0 / dB * K_red_2 A_matrix[n, n + 1] = -1.0 A_matrix[0, 0] = 1.0 A_matrix[2 * n - 1, 2 * n - 1] = 1.0 conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk conc = scipy.linalg.solve(A_matrix, conc_d) flux_1 = calc_flux(conc, n, dA, X_grid) flux_2 = K_red_2 * conc[n] - K_ox_2 * 0.0 fluxes[idx] = flux_1 + flux_2 cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) time_arr = np.arange(total_steps) * deltaT return { 'potential': E, 'flux': fluxes, 'time': time_arr, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: EC' (catalytic EC with B -> A regeneration) # ============================================================================= def run_ec_prime_simulation( sigma, K0, alpha, kc, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): beta = 1.0 - alpha deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) aA, bA, cA = np.zeros(n), np.zeros(n), np.zeros(n) aB, bB, cB = np.zeros(n), np.zeros(n), np.zeros(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) A_matrix = np.zeros((2 * n, 2 * n)) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] decay_factor = np.exp(-kc * deltaT) for idx in range(total_steps): Theta = E[idx] K_red = K0 * np.exp(np.clip(-alpha * Theta, -500, 500)) K_ox = K0 * np.exp(np.clip(beta * Theta, -500, 500)) K_red, K_ox = _clamp_bv_rates(K_red, K_ox) A_matrix[:] = 0.0 X0 = X_grid[1] - X_grid[0] rows_A = np.arange(n - 2, 0, -1) A_matrix[rows_A, rows_A - 1] = cA[1:n - 1] A_matrix[rows_A, rows_A] = bA[1:n - 1] A_matrix[rows_A, rows_A + 1] = aA[1:n - 1] rows_B = np.arange(n + 1, 2 * n - 1) A_matrix[rows_B, rows_B - 1] = aB[1:n - 1] A_matrix[rows_B, rows_B] = bB[1:n - 1] A_matrix[rows_B, rows_B + 1] = cB[1:n - 1] A_matrix[n - 1, n - 2] = -1.0 A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red A_matrix[n - 1, n] = -X0 / dA * K_ox A_matrix[n, n - 1] = -X0 / dB * K_red A_matrix[n, n] = 1.0 + X0 / dB * K_ox A_matrix[n, n + 1] = -1.0 A_matrix[0, 0] = 1.0 A_matrix[2 * n - 1, 2 * n - 1] = 1.0 conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk conc = scipy.linalg.solve(A_matrix, conc_d) fluxes[idx] = calc_flux(conc, n, dA, X_grid) # Operator-split catalytic step: B -> A in interior b_interior = conc[n + 2:2 * n - 1].copy() amount = b_interior * (1.0 - decay_factor) conc[n + 2:2 * n - 1] -= amount conc[1:n - 2] += amount[::-1] cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) time_arr = np.arange(total_steps) * deltaT return { 'potential': E, 'flux': fluxes, 'time': time_arr, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: CE (preceding chemical reaction Y -> A, then A -e-> B) # ============================================================================= def run_ce_simulation( sigma, K0, alpha, kf, Keq, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_Y_bulk=1.0, C_A_bulk=None, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): if C_A_bulk is None: C_A_bulk = Keq * C_Y_bulk / (1.0 + Keq) beta = 1.0 - alpha deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) aA, bA, cA = np.zeros(n), np.zeros(n), np.zeros(n) aB, bB, cB = np.zeros(n), np.zeros(n), np.zeros(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) A_matrix = np.zeros((2 * n, 2 * n)) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] C_A_eq = Keq * C_Y_bulk / (1.0 + Keq) for idx in range(total_steps): Theta = E[idx] K_red = K0 * np.exp(np.clip(-alpha * Theta, -500, 500)) K_ox = K0 * np.exp(np.clip(beta * Theta, -500, 500)) K_red, K_ox = _clamp_bv_rates(K_red, K_ox) A_matrix[:] = 0.0 X0 = X_grid[1] - X_grid[0] rows_A = np.arange(n - 2, 0, -1) A_matrix[rows_A, rows_A - 1] = cA[1:n - 1] A_matrix[rows_A, rows_A] = bA[1:n - 1] A_matrix[rows_A, rows_A + 1] = aA[1:n - 1] rows_B = np.arange(n + 1, 2 * n - 1) A_matrix[rows_B, rows_B - 1] = aB[1:n - 1] A_matrix[rows_B, rows_B] = bB[1:n - 1] A_matrix[rows_B, rows_B + 1] = cB[1:n - 1] A_matrix[n - 1, n - 2] = -1.0 A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red A_matrix[n - 1, n] = -X0 / dA * K_ox A_matrix[n, n - 1] = -X0 / dB * K_red A_matrix[n, n] = 1.0 + X0 / dB * K_ox A_matrix[n, n + 1] = -1.0 A_matrix[0, 0] = 1.0 A_matrix[2 * n - 1, 2 * n - 1] = 1.0 conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk conc = scipy.linalg.solve(A_matrix, conc_d) fluxes[idx] = calc_flux(conc, n, dA, X_grid) # Operator-split preceding chemical step: Y -> A for i in range(2, n - 1): c_a = conc[n - 1 - i] source = kf * (C_A_eq - c_a) * deltaT conc[n - 1 - i] = max(c_a + source, 0.0) cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) time_arr = np.arange(total_steps) * deltaT return { 'potential': E, 'flux': fluxes, 'time': time_arr, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: ECE (A -e-> B, B -> C with kc, C -e-> D) # ============================================================================= def run_ece_simulation( sigma, K0_1, alpha_1, K0_2, alpha_2, E0_2_offset, kc, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): beta_1 = 1.0 - alpha_1 deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) aA, bA, cA = np.zeros(n), np.zeros(n), np.zeros(n) aB, bB, cB = np.zeros(n), np.zeros(n), np.zeros(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) A_matrix = np.zeros((2 * n, 2 * n)) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] decay_factor = np.exp(-kc * deltaT) for idx in range(total_steps): Theta1 = E[idx] K_red_1 = K0_1 * np.exp(np.clip(-alpha_1 * Theta1, -500, 500)) K_ox_1 = K0_1 * np.exp(np.clip(beta_1 * Theta1, -500, 500)) K_red_1, K_ox_1 = _clamp_bv_rates(K_red_1, K_ox_1) # Operator-split chemical decay B -> C before diffusion step conc[n:2 * n] *= decay_factor A_matrix[:] = 0.0 X0 = X_grid[1] - X_grid[0] rows_A = np.arange(n - 2, 0, -1) A_matrix[rows_A, rows_A - 1] = cA[1:n - 1] A_matrix[rows_A, rows_A] = bA[1:n - 1] A_matrix[rows_A, rows_A + 1] = aA[1:n - 1] rows_B = np.arange(n + 1, 2 * n - 1) A_matrix[rows_B, rows_B - 1] = aB[1:n - 1] A_matrix[rows_B, rows_B] = bB[1:n - 1] A_matrix[rows_B, rows_B + 1] = cB[1:n - 1] A_matrix[n - 1, n - 2] = -1.0 A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red_1 A_matrix[n - 1, n] = -X0 / dA * K_ox_1 A_matrix[n, n - 1] = -X0 / dB * K_red_1 A_matrix[n, n] = 1.0 + X0 / dB * K_ox_1 A_matrix[n, n + 1] = -1.0 A_matrix[0, 0] = 1.0 A_matrix[2 * n - 1, 2 * n - 1] = 1.0 conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk conc = scipy.linalg.solve(A_matrix, conc_d) fluxes[idx] = calc_flux(conc, n, dA, X_grid) cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) time_arr = np.arange(total_steps) * deltaT return { 'potential': E, 'flux': fluxes, 'time': time_arr, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: EC_LH (Langmuir coverage-modulated BV + chemical decay) # ============================================================================= def run_ec_lh_simulation( sigma, K0, alpha, KA_eq, KB_eq, kc, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): beta = 1.0 - alpha deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) A_matrix, aA, bA, cA, aB, bB, cB = ini_coeff(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] decay_factor = np.exp(-kc * deltaT) for idx in range(total_steps): Theta = E[idx] conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk cA_surf = max(conc[n - 1], 0.0) cB_surf = max(conc[n], 0.0) denom = 1.0 + KA_eq * cA_surf + KB_eq * cB_surf K_red = K0 * np.exp(-alpha * Theta) K_ox = K0 * np.exp(beta * Theta) K_red_eff = K_red * KA_eq / denom K_ox_eff = K_ox * KB_eq / denom A_mat = _build_matrix_with_rates(A_matrix, X_grid, n, aA, bA, cA, dA, aB, bB, cB, dB, K_red_eff, K_ox_eff) conc = scipy.linalg.solve(A_mat, conc_d) conc[n:2 * n - 1] *= decay_factor conc[2 * n - 1] = C_B_bulk fluxes[idx] = calc_flux(conc, n, dA, X_grid) cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) return { 'potential': E, 'flux': fluxes, 'time': np.arange(total_steps) * deltaT, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: MHC_EC (MHC kinetics + chemical follow-up) # ============================================================================= def run_mhc_ec_simulation( sigma, K0, reorg_e, kc, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) A_matrix, aA, bA, cA, aB, bB, cB = ini_coeff(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] decay_factor = np.exp(-kc * deltaT) for idx in range(total_steps): Theta = E[idx] conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk K_red, K_ox = calc_mhc_rates(Theta, K0, reorg_e) A_mat = _build_matrix_with_rates(A_matrix, X_grid, n, aA, bA, cA, dA, aB, bB, cB, dB, K_red, K_ox) conc = scipy.linalg.solve(A_mat, conc_d) conc[n:2 * n - 1] *= decay_factor conc[2 * n - 1] = C_B_bulk fluxes[idx] = calc_flux(conc, n, dA, X_grid) cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) return { 'potential': E, 'flux': fluxes, 'time': np.arange(total_steps) * deltaT, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator: MHC_LH (MHC kinetics + Langmuir coverage) # ============================================================================= def run_mhc_lh_simulation( sigma, K0, reorg_e, KA_eq, KB_eq, dA=1.0, dB=1.0, theta_i=20.0, theta_v=-20.0, C_A_bulk=1.0, C_B_bulk=0.0, cycles=1, n_spatial_out=N_SPATIAL_OUT, ): deltaT = DELTA_THETA / sigma maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma E = _make_potential_waveform(theta_i, theta_v, cycles) total_steps = len(E) X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT), EXPANDING_GRID_FACTOR) conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk) A_matrix, aA, bA, cA, aB, bB, cB = ini_coeff(n) aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA) aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB) fluxes = np.zeros(total_steps) cA_hist, cB_hist = [], [] for idx in range(total_steps): Theta = E[idx] conc_d[:] = conc[:] conc_d[n - 1] = 0.0 conc_d[n] = 0.0 conc_d[0] = C_A_bulk conc_d[2 * n - 1] = C_B_bulk cA_surf = max(conc[n - 1], 0.0) cB_surf = max(conc[n], 0.0) denom = 1.0 + KA_eq * cA_surf + KB_eq * cB_surf K_red, K_ox = calc_mhc_rates(Theta, K0, reorg_e) K_red_eff = K_red * KA_eq / denom K_ox_eff = K_ox * KB_eq / denom A_mat = _build_matrix_with_rates(A_matrix, X_grid, n, aA, bA, cA, dA, aB, bB, cB, dB, K_red_eff, K_ox_eff) conc = scipy.linalg.solve(A_mat, conc_d) fluxes[idx] = calc_flux(conc, n, dA, X_grid) cA_hist.append(conc[:n].copy()) cB_hist.append(conc[n:].copy()) cA_hist = np.stack(cA_hist) cB_hist = np.stack(cB_hist) x_out, c_ox = resample_concentration(X_grid, cA_hist, n_spatial_out) _, c_red = resample_concentration(X_grid, cB_hist, n_spatial_out) return { 'potential': E, 'flux': fluxes, 'time': np.arange(total_steps) * deltaT, 'c_ox': c_ox, 'c_red': c_red, 'x_grid_out': x_out, } # ============================================================================= # Simulator dispatch table # ============================================================================= SIMULATORS = { 'EE': run_ee_simulation, 'EC_prime': run_ec_prime_simulation, 'CE': run_ce_simulation, 'ECE': run_ece_simulation, 'EC_LH': run_ec_lh_simulation, 'MHC_EC': run_mhc_ec_simulation, 'MHC_LH': run_mhc_lh_simulation, } # ============================================================================= # Parameter sampling # ============================================================================= def _sample_common_params(rng): log_sigma = rng.uniform(-1, 2) sigma = 10 ** log_sigma d_ratio = 10 ** rng.uniform(-0.3, 0.3) theta_center = rng.uniform(-5, 5) theta_range = rng.uniform(15, 25) return { 'sigma': sigma, 'C_A_bulk': 1.0, 'C_B_bulk': 0.0, 'dA': 1.0, 'dB': d_ratio, 'theta_i': theta_center + theta_range / 2, 'theta_v': theta_center - theta_range / 2, } def sample_ee_params(rng): p = _sample_common_params(rng) p['K0_1'] = 10 ** rng.uniform(-2, 2) p['alpha_1'] = rng.uniform(0.3, 0.7) p['K0_2'] = 10 ** rng.uniform(-2, 2) p['alpha_2'] = rng.uniform(0.3, 0.7) p['E0_2_offset'] = rng.uniform(-5, 5) p['dC'] = 10 ** rng.uniform(-0.3, 0.3) return p def sample_ec_prime_params(rng): p = _sample_common_params(rng) p['K0'] = 10 ** rng.uniform(-2, 2) p['alpha'] = rng.uniform(0.3, 0.7) p['kc'] = 10 ** rng.uniform(-2, 2) return p def sample_ce_params(rng): p = _sample_common_params(rng) p['K0'] = 10 ** rng.uniform(-2, 2) p['alpha'] = rng.uniform(0.3, 0.7) p['kf'] = 10 ** rng.uniform(-2, 2) p['Keq'] = 10 ** rng.uniform(-1, 2) return p def sample_ece_params(rng): p = _sample_common_params(rng) p['K0_1'] = 10 ** rng.uniform(-2, 2) p['alpha_1'] = rng.uniform(0.3, 0.7) p['K0_2'] = 10 ** rng.uniform(-2, 2) p['alpha_2'] = rng.uniform(0.3, 0.7) p['E0_2_offset'] = rng.uniform(-5, 5) p['kc'] = 10 ** rng.uniform(-2, 2) return p def sample_ec_lh_params(rng): p = _sample_common_params(rng) p['K0'] = 10 ** rng.uniform(-2, 2) p['alpha'] = rng.uniform(0.3, 0.7) p['KA_eq'] = 10 ** rng.uniform(-1, 2) p['KB_eq'] = 10 ** rng.uniform(-1, 2) p['kc'] = 10 ** rng.uniform(-2, 2) return p def sample_mhc_ec_params(rng): p = _sample_common_params(rng) p['K0'] = 10 ** rng.uniform(-2, 2) p['reorg_e'] = 10 ** rng.uniform(0.5, 2.0) p['kc'] = 10 ** rng.uniform(-2, 2) return p def sample_mhc_lh_params(rng): p = _sample_common_params(rng) p['K0'] = 10 ** rng.uniform(-2, 2) p['reorg_e'] = 10 ** rng.uniform(0.5, 2.0) p['KA_eq'] = 10 ** rng.uniform(-1, 2) p['KB_eq'] = 10 ** rng.uniform(-1, 2) return p PARAM_SAMPLERS = { 'EE': sample_ee_params, 'EC_prime': sample_ec_prime_params, 'CE': sample_ce_params, 'ECE': sample_ece_params, 'EC_LH': sample_ec_lh_params, 'MHC_EC': sample_mhc_ec_params, 'MHC_LH': sample_mhc_lh_params, } # Which K0-like keys need sqrt(sigma) rescaling for each mechanism _K0_KEYS = { 'EE': ['K0_1', 'K0_2'], 'EC_prime': ['K0'], 'CE': ['K0'], 'ECE': ['K0_1', 'K0_2'], 'EC_LH': ['K0'], 'MHC_EC': ['K0'], 'MHC_LH': ['K0'], } # Which chemical rate keys need linear sigma rescaling _KC_KEYS = { 'EE': [], 'EC_prime': ['kc'], 'CE': ['kf'], 'ECE': ['kc'], 'EC_LH': ['kc'], 'MHC_EC': ['kc'], 'MHC_LH': [], } # ============================================================================= # Noise augmentation # ============================================================================= def _add_noise(flux, rng, noise_range=(0.001, 0.02)): sigma_noise = rng.uniform(*noise_range) peak = np.max(np.abs(flux)) + 1e-20 noise = sigma_noise * peak * rng.standard_normal(flux.shape) return flux + noise.astype(flux.dtype), float(sigma_noise) # ============================================================================= # Multi-scan sample generation # ============================================================================= def _build_sim_params(mechanism, base_params, mech_params, sigma_val, K0_at_1_map, kc_at_1_map): """Build simulator kwargs for a given scan rate, rescaling K0 and kc.""" sim_params = {} # Common params (excluding sigma, which we override) for k in ('dA', 'dB', 'theta_i', 'theta_v', 'C_A_bulk', 'C_B_bulk'): if k in base_params: sim_params[k] = base_params[k] sim_params['sigma'] = float(sigma_val) # Mechanism-specific params (copy all, then override rescaled ones) for k, v in mech_params.items(): sim_params[k] = v sqrt_sigma = np.sqrt(sigma_val) for k0_key in _K0_KEYS[mechanism]: if k0_key in K0_at_1_map: sim_params[k0_key] = K0_at_1_map[k0_key] / sqrt_sigma for kc_key in _KC_KEYS[mechanism]: if kc_key in kc_at_1_map: sim_params[kc_key] = kc_at_1_map[kc_key] / sigma_val return sim_params def generate_single_sample(sample_idx, mechanism, seed, n_scan_rates=3, add_noise=True, n_spatial_out=N_SPATIAL_OUT): """Generate one multi-scan-rate sample for a given mechanism.""" rng = np.random.default_rng(seed) sampler = PARAM_SAMPLERS[mechanism] simulator = SIMULATORS[mechanism] all_params = sampler(rng) base_sigma = all_params['sigma'] # Separate base (common) params from mechanism-specific params base_keys = {'sigma', 'C_A_bulk', 'C_B_bulk', 'dA', 'dB', 'theta_i', 'theta_v'} base_params = {k: v for k, v in all_params.items() if k in base_keys} mech_params = {k: v for k, v in all_params.items() if k not in base_keys} # Convert K0 and kc to canonical sigma=1 K0_at_1_map = {} for k0_key in _K0_KEYS[mechanism]: if k0_key in mech_params: K0_at_1_map[k0_key] = mech_params[k0_key] * np.sqrt(base_sigma) kc_at_1_map = {} for kc_key in _KC_KEYS[mechanism]: if kc_key in mech_params: kc_at_1_map[kc_key] = mech_params[kc_key] * base_sigma # Sample scan rates log_sigmas = np.sort(rng.uniform(-1, 2, size=n_scan_rates)) sigmas = 10 ** log_sigmas all_potential, all_flux, all_time = [], [], [] all_c_ox, all_c_red = [], [] lengths = [] x_grid_out = None for sigma_val in sigmas: sim_params = _build_sim_params(mechanism, base_params, mech_params, sigma_val, K0_at_1_map, kc_at_1_map) try: result = simulator(**sim_params, n_spatial_out=n_spatial_out) except Exception: return None pot = result['potential'] flux = result['flux'] if np.any(np.isnan(flux)) or np.any(np.isinf(flux)) or np.any(np.abs(flux) > 1e10): return None flux = flux.astype(np.float32) if add_noise: flux, _ = _add_noise(flux, rng) all_potential.append(pot.astype(np.float32)) all_flux.append(flux) all_time.append(result['time'].astype(np.float32)) all_c_ox.append(result['c_ox'].astype(np.float32)) all_c_red.append(result['c_red'].astype(np.float32)) lengths.append(len(pot)) x_grid_out = result['x_grid_out'].astype(np.float32) # Pad to max length max_len = max(lengths) n_s = len(sigmas) pot_arr = np.zeros((n_s, max_len), dtype=np.float32) flux_arr = np.zeros((n_s, max_len), dtype=np.float32) time_arr = np.zeros((n_s, max_len), dtype=np.float32) c_ox_arr = np.zeros((n_s, max_len, n_spatial_out), dtype=np.float32) c_red_arr = np.zeros((n_s, max_len, n_spatial_out), dtype=np.float32) for i in range(n_s): L = lengths[i] pot_arr[i, :L] = all_potential[i] flux_arr[i, :L] = all_flux[i] time_arr[i, :L] = all_time[i] c_ox_arr[i, :L] = all_c_ox[i] c_red_arr[i, :L] = all_c_red[i] # Store params at canonical sigma=1 stored_params = dict(all_params) stored_params['kinetics'] = mechanism for k0_key, val in K0_at_1_map.items(): stored_params[k0_key] = float(val) for kc_key, val in kc_at_1_map.items(): stored_params[kc_key] = float(val) return { 'potential': pot_arr, 'flux': flux_arr, 'time': time_arr, 'c_ox': c_ox_arr, 'c_red': c_red_arr, 'x_grid': x_grid_out, 'sigmas': sigmas.astype(np.float32), 'lengths': np.array(lengths, dtype=np.int32), 'params': stored_params, 'mechanism_id': np.int32(MECHANISM_TO_ID[mechanism]), 'n_scan_rates': np.int32(n_scan_rates), } # ============================================================================= # Worker functions for multiprocessing # ============================================================================= def _worker_generate(args): """Generate one sample and save directly to disk to avoid OOM.""" (sample_idx, mechanism, seed, n_scan_rates, add_noise_flag, n_spatial_out, tmp_dir) = args try: result = generate_single_sample( sample_idx, mechanism, seed, n_scan_rates=n_scan_rates, add_noise=add_noise_flag, n_spatial_out=n_spatial_out, ) if result is None: return False fpath = os.path.join(tmp_dir, f"{mechanism}_{sample_idx:06d}.npz") np.savez_compressed( fpath, potential=result['potential'], flux=result['flux'], time=result['time'], c_ox=result['c_ox'], c_red=result['c_red'], x_grid=result['x_grid'], sigmas=result['sigmas'], lengths=result['lengths'], params=result['params'], mechanism_id=result['mechanism_id'], n_scan_rates=result['n_scan_rates'], ) return True except Exception: return False # ============================================================================= # Dataset generation # ============================================================================= def generate_dataset( output_dir='data_extended', mechanisms=None, n_samples=10000, n_scan_rates=3, n_workers=None, seed=42, split_ratio=(0.8, 0.1, 0.1), add_noise=True, n_spatial_out=N_SPATIAL_OUT, ): if mechanisms is None: mechanisms = MECHANISM_LIST if n_workers is None: n_workers = max(1, cpu_count() - 1) rng_master = np.random.default_rng(seed) tmp_dir = os.path.join(output_dir, '_tmp') os.makedirs(tmp_dir, exist_ok=True) total_samples = n_samples * len(mechanisms) work_list = [] for mech in mechanisms: seeds = rng_master.integers(0, 2**31, size=n_samples) for i in range(n_samples): work_list.append((i, mech, int(seeds[i]), n_scan_rates, add_noise, n_spatial_out, tmp_dir)) print(f"Generating {n_samples} samples x {len(mechanisms)} mechanisms " f"= {total_samples} total") print(f"Mechanisms: {mechanisms}") print(f"Scan rates per sample: {n_scan_rates}") print(f"Workers: {n_workers}") print(f"Output: {output_dir}") print(flush=True) n_success = 0 n_fail = 0 n_workers = min(n_workers, total_samples) if n_workers <= 1: for args in tqdm(work_list, desc="Generating"): ok = _worker_generate(args) if ok: n_success += 1 else: n_fail += 1 else: with Pool(processes=n_workers) as pool: for ok in tqdm( pool.imap_unordered(_worker_generate, work_list, chunksize=8), total=total_samples, desc="Generating", ): if ok: n_success += 1 else: n_fail += 1 print(f"\nSuccessful: {n_success}/{total_samples} (failed: {n_fail})") if n_success == 0: print("No successful samples generated. Exiting.") return from glob import glob as _glob all_files = sorted(_glob(os.path.join(tmp_dir, '*.npz'))) rng_split = np.random.default_rng(seed + 2) rng_split.shuffle(all_files) n_train = int(len(all_files) * split_ratio[0]) n_val = int(len(all_files) * split_ratio[1]) splits = { 'train': all_files[:n_train], 'val': all_files[n_train:n_train + n_val], 'test': all_files[n_train + n_val:], } import shutil mech_counts = {s: {} for s in splits} for split_name, file_list in splits.items(): split_dir = os.path.join(output_dir, split_name) os.makedirs(split_dir, exist_ok=True) for file_idx, src_path in enumerate( tqdm(file_list, desc=f"Moving {split_name}") ): dst_path = os.path.join(split_dir, f"sample_{file_idx:06d}.npz") shutil.move(src_path, dst_path) basename = os.path.basename(src_path) mech = basename.rsplit('_', 1)[0] mech_counts[split_name][mech] = \ mech_counts[split_name].get(mech, 0) + 1 shutil.rmtree(tmp_dir, ignore_errors=True) summary = { 'n_samples_per_mechanism': n_samples, 'mechanisms': mechanisms, 'mechanism_ids': {m: MECHANISM_TO_ID[m] for m in mechanisms}, 'n_scan_rates': n_scan_rates, 'add_noise': add_noise, 'seed': seed, 'split_ratio': list(split_ratio), 'n_total': n_success, 'n_train': len(splits['train']), 'n_val': len(splits['val']), 'n_test': len(splits['test']), 'mechanism_counts': mech_counts, } with open(os.path.join(output_dir, 'metadata.json'), 'w') as f: json.dump(summary, f, indent=2) print(f"\nDataset saved to {output_dir}/") print(f" train: {len(splits['train'])} samples") print(f" val: {len(splits['val'])} samples") print(f" test: {len(splits['test'])} samples") for split_name in ('train', 'val', 'test'): print(f" {split_name} mechanism counts: {mech_counts[split_name]}") # ============================================================================= # Test mode # ============================================================================= def run_test(): """Run a quick test simulation for each mechanism and print diagnostics.""" rng = np.random.default_rng(42) for mech in MECHANISM_LIST: print(f"\n{'='*60}") print(f"Testing {mech} (mechanism_id={MECHANISM_TO_ID[mech]})") print(f"{'='*60}") sampler = PARAM_SAMPLERS[mech] simulator = SIMULATORS[mech] params = sampler(rng) sim_kwargs = {k: v for k, v in params.items() if k not in ('C_A_bulk', 'C_B_bulk')} sim_kwargs['C_A_bulk'] = params.get('C_A_bulk', 1.0) sim_kwargs['C_B_bulk'] = params.get('C_B_bulk', 0.0) try: result = simulator(**sim_kwargs) flux = result['flux'] print(f" Time steps: {len(result['time'])}") print(f" Potential range: [{result['potential'].min():.2f}, " f"{result['potential'].max():.2f}]") print(f" Flux range: [{flux.min():.6f}, {flux.max():.6f}]") print(f" NaN in flux: {np.any(np.isnan(flux))}") print(f" Inf in flux: {np.any(np.isinf(flux))}") print(f" c_ox shape: {result['c_ox'].shape}") print(f" Status: OK") except Exception as e: print(f" Status: FAILED - {e}") # Test multi-scan generation print(f"\n{'='*60}") print("Testing multi-scan sample generation (EE, 3 scan rates)") print(f"{'='*60}") sample = generate_single_sample(0, 'EE', seed=42, n_scan_rates=3) if sample is not None: print(f" potential shape: {sample['potential'].shape}") print(f" flux shape: {sample['flux'].shape}") print(f" sigmas: {sample['sigmas']}") print(f" lengths: {sample['lengths']}") print(f" mechanism_id: {sample['mechanism_id']}") print(f" params keys: {list(sample['params'].keys())}") print(f" Status: OK") else: print(f" Status: FAILED (returned None)") # ============================================================================= # Main # ============================================================================= def main(): parser = argparse.ArgumentParser( description="Generate extended mechanism training data for ECFlow" ) parser.add_argument( '--output_dir', type=str, default='data_extended', help='Output directory (default: data_extended)' ) parser.add_argument( '--mechanisms', nargs='+', default=None, choices=MECHANISM_LIST, help='Mechanisms to generate (default: all 7)' ) parser.add_argument( '--n_samples', type=int, default=10000, help='Number of samples per mechanism (default: 10000)' ) parser.add_argument( '--n_scan_rates', type=int, default=3, help='Number of scan rates per sample (default: 3)' ) parser.add_argument( '--n_workers', type=int, default=None, help='Number of parallel workers (default: num_cpus - 1)' ) parser.add_argument( '--seed', type=int, default=42, help='Random seed (default: 42)' ) parser.add_argument( '--split_ratio', nargs=3, type=float, default=[0.8, 0.1, 0.1], help='Train/val/test split ratio (default: 0.8 0.1 0.1)' ) parser.add_argument( '--no_noise', action='store_true', help='Disable noise augmentation' ) parser.add_argument( '--test', action='store_true', help='Run test simulations for all mechanisms' ) args = parser.parse_args() if args.test: run_test() else: generate_dataset( output_dir=args.output_dir, mechanisms=args.mechanisms, n_samples=args.n_samples, n_scan_rates=args.n_scan_rates, n_workers=args.n_workers, seed=args.seed, split_ratio=tuple(args.split_ratio), add_noise=not args.no_noise, ) if __name__ == '__main__': main()