Improving runtime of Nested optimization with diffrax and optax.

[DylanEsguerra]

Hello

I am working on a nested optimization script with autodiff and wanted to try out JAX as an alternative to Julia. The reason a Nested optimization is important here is because that although fitting all 3 parameters together is faster, kf easily becomes stuck in a local minimum when the initial guess is poor and I want to be robust to this when fitting non simulated data. I am new to jax and want to see if anyone can identify any glaring reasons why this code would run slowly. It is significantly slower than expected and I haven't been able to run it long enough yet to see if the fit is correct. I have experimented with the optimistix package as well but have had less success so far.

Thank you for the help!

import jax
import jax.numpy as jnp
import diffrax as dfx
import optax
import matplotlib.pyplot as plt

@jax.jit
def sys(t, u, args):
p, p_fix, kf, plateau, kd, start_time, end_time = args
forcing = forcing_function(t, plateau, kd, start_time, end_time)

du_1 = kf[0] * (1.0 - forcing) - p[0] * u[0]
du_2 = p[0] * u[0] - p_fix[0] * u[1]
du_3 = p_fix[0] * u[1] - p[1] * u[2]

return jnp.stack([du_1, du_2, du_3])

@jax.jit
def forcing_function(t, plateau, kd, start_time, end_time):
return jnp.where(
t < start_time,
plateau * t,
jnp.where(
t < end_time,
plateau,
plateau * jnp.exp(-kd * (t - end_time))
)
)

#@jax.jit
def Steady_State(args):
p, p_fix, kf, plateau, kd, start_time, end_time = args
new_args = (p, p_fix, kf, 0.0, kd, start_time, end_time) # Set plateau to 0 for SS

term = dfx.ODETerm(sys)
solver = dfx.Tsit5()
cond_fn = dfx.steady_state_event()
controller = dfx.PIDController(rtol=1e-6, atol=1e-9)
event = dfx.Event(cond_fn)
adjoint = dfx.ImplicitAdjoint()
t0, t1 = 0, 1000
dt0 = 0.1
max_steps = 10000
u0 = jnp.array([1000.0, 750.0, 20.0])  # Initial guess

sol = dfx.diffeqsolve(term, solver=solver, t0=t0, t1=t1, dt0=dt0, y0=u0, args=new_args,
                      max_steps=max_steps, stepsize_controller=controller, event=event, adjoint=adjoint)
return sol.ys[-1, :]

def solve_ode(args, u0, ts):
term = dfx.ODETerm(sys)
solver = dfx.Tsit5()
saveat = dfx.SaveAt(ts=ts)
t0, t1 = ts[0], ts[-1]
dt0 = 0.1
sol = dfx.diffeqsolve(term, solver, t0, t1, dt0, y0=u0, args=args, saveat=saveat, adjoint=dfx.DirectAdjoint())
return sol

#@jax.jit
def set_kf(kf_guess, args):
p, p_fix, plateau, kd, u02, start_time, end_time = args
args_ss = (p, p_fix, kf_guess, 0.0, kd, start_time, end_time)
sol = Steady_State(args_ss)
return jnp.sum((sol[2] - u02)**2)

@jax.jit
def update_kf_step(kf, opt_state, args):
loss_value, grads = jax.value_and_grad(set_kf)(kf, args)
updates, opt_state = optimizer_kf.update(grads, opt_state)
kf = optax.apply_updates(kf, updates)
return kf, opt_state, loss_value

def optimize_kf(initial_kf, args, num_iterations=500):
kf = initial_kf
opt_state = optimizer_kf.init(kf)

for i in range(num_iterations):
    kf, opt_state, loss_value = update_kf_step(kf, opt_state, args)
    if i % 100 == 0:
        print(f"KF Iteration {i}, Loss: {loss_value}")

return kf

#@jax.jit
def loss_fn(p, fixed_args):
p_fix, kf_guess, plateau, kd, u02, start_time, end_time, Tmeas, Ymeas = fixed_args

# Optimize kf
args_kf = (p, p_fix, plateau, kd, u02, start_time, end_time)
kf_optimized = optimize_kf(kf_guess, args_kf)

# Calculate steady state
args_ss = (p, p_fix, kf_optimized, plateau, kd, start_time, end_time)
u0_ss = Steady_State(args_ss)

# Solve ODE with optimized kf and steady state initial conditions
args_ode = (p, p_fix, kf_optimized, plateau, kd, start_time, end_time)
sol = solve_ode(args_ode, u0_ss, Tmeas)

# Calculate predicted values
Y_pred = (u02 - sol.ys[:, 2]) / u02

# Calculate loss
return jnp.sum((Ymeas - Y_pred) ** 2)

@jax.jit
def update_p_step(p, opt_state, fixed_args):
loss_value, grads = jax.value_and_grad(loss_fn)(p, fixed_args)
updates, opt_state = optimizer_p.update(grads, opt_state)
p = optax.apply_updates(p, updates)
return p, opt_state, loss_value

def optimize_p(initial_p, fixed_args, num_iterations=1000):
p = initial_p
opt_state = optimizer_p.init(p)

for i in range(num_iterations):
    p, opt_state, loss_value = update_p_step(p, opt_state, fixed_args)
    if i % 100 == 0:
        print(f"P Iteration {i}, Loss: {loss_value}")

return p

def main():
# Set fixed parameters
kf0 = jnp.array([10.0]) # True Value
p_fix = jnp.array([0.25]) # Known parameter
p_var = jnp.array([1.0, 0.5]) # Unknown parameters
u0_guess = jnp.array([1000.0, 750.0, 20.0])
u02 = 20.0
plateau, kd = 0.5, 1.0
start_time, end_time = 1.0, 10.0

# Set up ODE Problem
t_end = 36.0
ts = jnp.linspace(0.0, t_end, 300)

# Simulate data
args_sim = (p_var, p_fix, kf0, plateau, kd, start_time, end_time)
u0_ss = Steady_State(args_sim)
sol_sim = solve_ode(args_sim, u0_ss, ts)

X_sim = sol_sim.ys
T_sim = sol_sim.ts
Y_sim = (u02 - X_sim[:, 2]) / u02

Tmeas = jnp.arange(37)
Ymeas = jnp.interp(Tmeas, ts, Y_sim)

# Set up optimization
p_guess = jnp.array([10.0, 3.00])
kf_guess = jnp.array([100.0])

learning_rate_p = 0.003#optax.exponential_decay(0.1, transition_steps=100, decay_rate=0.9)
learning_rate_kf = 0.003#optax.exponential_decay(0.1, transition_steps=100, decay_rate=0.9)
global optimizer_p, optimizer_kf
optimizer_p = optax.adam(learning_rate_p)
optimizer_kf = optax.adam(learning_rate_kf)

fixed_args = (p_fix, kf_guess, plateau, kd, u02, start_time, end_time, Tmeas, Ymeas)

# Perform optimization
p_optimized = optimize_p(p_guess, fixed_args, num_iterations=500)

print("Final p value:", p_optimized)

# Find final kf and initial value for optimal p
args_kf_final = (p_optimized, p_fix, plateau, kd, u02, start_time, end_time)
kf_optimized = optimize_kf(kf_guess, args_kf_final)
print("Final kf value:", kf_optimized)

args_ss_final = (p_optimized, p_fix, kf_optimized, plateau, kd, start_time, end_time)
u0_ss_optimized = Steady_State(args_ss_final)
print("Final u0 value:", u0_ss_optimized)

# Solve ODE with optimized parameters
args_optimized = (p_optimized, p_fix, kf_optimized, plateau, kd, start_time, end_time)
sol = solve_ode(args_optimized, u0_ss_optimized, ts)

X = sol.ys
T = sol.ts
Y = (u02 - X[:, 2]) / u02

# Plot results
plt.figure(figsize=(10, 6))
plt.plot(T, Y, label="Fitted")
plt.plot(T_sim, Y_sim, label="True", linestyle='--')
plt.scatter(Tmeas, Ymeas, label="Measured labeling")
plt.xlabel("Time (h)")
plt.ylabel("Fraction of Labeled Protein")
plt.title("Fit to ODE System")
plt.legend()
plt.show()

if name == "main":
main()