[Python] Why is the JAX and JAXOPT based code so slow?

I am writing a jax and jaxopt based optimization. The same code takes ~3-4 seconds in R. I am not sure why this code takes ~90 seconds in jax. I am new to python, jax and jaxopt, any suggestions to include the code quality or speed will be much appreciated.

Thank you to the community!

The input file (.csv) is pasted below first.

ID,AMT,TIME,DV,WT,MDV,EVID
1,4.02,0,0.74,79.6,1,1
1,.,0.25,2.84,.,0,0
1,.,0.57,6.57,.,0,0
1,.,1.12,10.5,.,0,0
1,.,2.02,9.66,.,0,0
1,.,3.82,8.58,.,0,0
1,.,5.1,8.36,.,0,0
1,.,7.03,7.47,.,0,0
1,.,9.05,6.89,.,0,0
1,.,12.12,5.94,.,0,0
1,.,24.37,3.28,.,0,0
2,4.4,0,0,72.4,1,1
2,.,0.27,1.72,.,0,0
2,.,0.52,7.91,.,0,0
2,.,1,8.31,.,0,0
2,.,1.92,8.33,.,0,0
2,.,3.5,6.85,.,0,0
2,.,5.02,6.08,.,0,0
2,.,7.03,5.4,.,0,0
2,.,9,4.55,.,0,0
2,.,12,3.01,.,0,0
2,.,24.3,0.9,.,0,0
3,4.53,0,0,70.5,1,1
3,.,0.27,4.4,.,0,0
3,.,0.58,6.9,.,0,0
3,.,1.02,8.2,.,0,0
3,.,2.02,7.8,.,0,0
3,.,3.62,7.5,.,0,0
3,.,5.08,6.2,.,0,0
3,.,7.07,5.3,.,0,0
3,.,9,4.9,.,0,0
3,.,12.15,3.7,.,0,0
3,.,24.17,1.05,.,0,0
4,4.4,0,0,72.7,1,1
4,.,0.35,1.89,.,0,0
4,.,0.6,4.6,.,0,0
4,.,1.07,8.6,.,0,0
4,.,2.13,8.38,.,0,0
4,.,3.5,7.54,.,0,0
4,.,5.02,6.88,.,0,0
4,.,7.02,5.78,.,0,0
4,.,9.02,5.33,.,0,0
4,.,11.98,4.19,.,0,0
4,.,24.65,1.15,.,0,0
5,5.86,0,0,54.6,1,1
5,.,0.3,2.02,.,0,0
5,.,0.52,5.63,.,0,0
5,.,1,11.4,.,0,0
5,.,2.02,9.33,.,0,0
5,.,3.5,8.74,.,0,0
5,.,5.02,7.56,.,0,0
5,.,7.02,7.09,.,0,0
5,.,9.1,5.9,.,0,0
5,.,12,4.37,.,0,0
5,.,24.35,1.57,.,0,0
6,4.,0,0,80.,1,1
6,.,0.27,1.29,.,0,0
6,.,0.58,3.08,.,0,0
6,.,1.15,6.44,.,0,0
6,.,2.03,6.32,.,0,0
6,.,3.57,5.53,.,0,0
6,.,5,4.94,.,0,0
6,.,7,4.02,.,0,0
6,.,9.22,3.46,.,0,0
6,.,12.1,2.78,.,0,0
6,.,23.85,0.92,.,0,0
7,4.95,0,0.15,64.6,1,1
7,.,0.25,0.85,.,0,0
7,.,0.5,2.35,.,0,0
7,.,1.02,5.02,.,0,0
7,.,2.02,6.58,.,0,0
7,.,3.48,7.09,.,0,0
7,.,5,6.66,.,0,0
7,.,6.98,5.25,.,0,0
7,.,9,4.39,.,0,0
7,.,12.05,3.53,.,0,0
7,.,24.22,1.15,.,0,0
8,4.53,0,0,70.5,1,1
8,.,0.25,3.05,.,0,0
8,.,0.52,3.05,.,0,0
8,.,0.98,7.31,.,0,0
8,.,2.02,7.56,.,0,0
8,.,3.53,6.59,.,0,0
8,.,5.05,5.88,.,0,0
8,.,7.15,4.73,.,0,0
8,.,9.07,4.57,.,0,0
8,.,12.1,3,.,0,0
8,.,24.12,1.25,.,0,0
9,3.1,0,0,86.4,1,1
9,.,0.3,7.37,.,0,0
9,.,0.63,9.03,.,0,0
9,.,1.05,7.14,.,0,0
9,.,2.02,6.33,.,0,0
9,.,3.53,5.66,.,0,0
9,.,5.02,5.67,.,0,0
9,.,7.17,4.24,.,0,0
9,.,8.8,4.11,.,0,0
9,.,11.6,3.16,.,0,0
9,.,24.43,1.12,.,0,0
10,5.5,0,0.24,58.2,1,1
10,.,0.37,2.89,.,0,0
10,.,0.77,5.22,.,0,0
10,.,1.02,6.41,.,0,0
10,.,2.05,7.83,.,0,0
10,.,3.55,10.21,.,0,0
10,.,5.05,9.18,.,0,0
10,.,7.08,8.02,.,0,0
10,.,9.38,7.14,.,0,0
10,.,12.1,5.68,.,0,0
10,.,23.7,2.42,.,0,0
11,4.92,0,0,65.,1,1
11,.,0.25,4.86,.,0,0
11,.,0.5,7.24,.,0,0
11,.,0.98,8,.,0,0
11,.,1.98,6.81,.,0,0
11,.,3.6,5.87,.,0,0
11,.,5.02,5.22,.,0,0
11,.,7.03,4.45,.,0,0
11,.,9.03,3.62,.,0,0
11,.,12.12,2.69,.,0,0
11,.,24.08,0.86,.,0,0
12,5.3,0,0,60.5,1,1
12,.,0.25,1.25,.,0,0
12,.,0.5,3.96,.,0,0
12,.,1,7.82,.,0,0
12,.,2,9.72,.,0,0
12,.,3.52,9.75,.,0,0
12,.,5.07,8.57,.,0,0
12,.,7.07,6.59,.,0,0
12,.,9.03,6.11,.,0,0
12,.,12.05,4.57,.,0,0
12,.,24.15,1.17,.,0,0


Code for optimization

## using LBFGSB from jaxopt for optimization

import jax
import jax.numpy as jnp
import pandas as pd
import jaxopt as jaxopt
import time
# import os
# os.environ["JAX_ENABLE_X64"] = "True"
jax.config.update("jax_enable_x64", True) # Enable 64-bit precision


# Calculate the sqrt of the inverse of a matrix
@jax.jit
def mat_sqrt_inv(mat):
eigvals, eigvecs = jnp.linalg.eigh(mat)
d2 = 1.0 / jnp.sqrt(jnp.abs(eigvals))
return eigvecs @ jnp.diag(d2) @ eigvecs.T

# Calculate the PRED value
@jax.jit
def PRED(THETA, ETA, DATA):
DOSE = 320.0
TIME = DATA[:, 1] # TIME values

KA = THETA[0] * jnp.exp(ETA[0])
V = THETA[1] * jnp.exp(ETA[1])
K = THETA[2] * jnp.exp(ETA[2])

F = DOSE / V * KA / (KA - K) * (jnp.exp(-K * TIME) - jnp.exp(-KA * TIME))

G1 = ((DOSE / V / (KA - K) - DOSE / V * KA / (KA - K) ** 2) * (jnp.exp(-K * TIME) - jnp.exp(-KA * TIME)) + DOSE / V * KA / (KA - K) * (jnp.exp(-KA * TIME) * TIME)) * KA
G2 = -(DOSE / V ** 2 * KA / (KA - K) * (jnp.exp(-K * TIME) - jnp.exp(-KA * TIME))) * V
G3 = (DOSE / V * KA / (KA - K) ** 2 * (jnp.exp(-K * TIME) - jnp.exp(-KA * TIME)) - DOSE / V * KA / (KA - K) * (jnp.exp(-K * TIME) * TIME)) * K
H1 = F
H2 = jnp.ones_like(F)

return jnp.stack([F, G1, G2, G3, H1, H2], axis=1)

# Read data
# Read data
df = pd.read_csv("THEOPH.csv")
df = df[df['EVID'] == 0]
ID = df['ID'].unique()
NETA = 3
NEPS = 2

selected_columns = ['ID', 'TIME', 'DV']
DATASET = jnp.array(df[selected_columns].values)

_subject_data_list = []
for subject_id_val in sorted(ID): # Iterate through unique IDs in sorted order
_subject_data_list.append(DATASET[DATASET[:, 0] == subject_id_val])

# Final NONMEM estimates
THETA = jnp.array([2.8864797758451806, 33.693922520723312, 8.7260954693777815E-002])
OMEGA = jnp.array([
[0.88639937186299267, 0, 0],
[0, 2.0903952408947064E-002, 0],
[0, 0, 6.8919408612682198E-002]
])
SIGMA = jnp.array([
[9.9259732296116312E-003, 0],
[0, 3.2931599977662498E-002]
])


# @jax.jit
def obj_fn(params, DATASET):
THETA = params[:3]
OMEGA_diag = params[3:6]
SIGMA_diag = params[6:8]


OBJ = 0.0
for i in ID:
# Select the data for the current subject
DATA = _subject_data_list[i-1]

# DATA = select_id(DATASET, i)
ETA = jnp.zeros(NETA)
FGH = PRED(THETA, ETA, DATA)
Y = DATA[:, 2] # DV values
Fpred = FGH[:, 0]
G = FGH[:, 1:4]
H = FGH[:, 4:6]

OM = jnp.diag(OMEGA_diag)
SG = jnp.diag(SIGMA_diag)
RES = Y - Fpred
COV = G @ OM @ G.T + jnp.diag(jnp.diag(H @ SG @ H.T))
# Add small value to diagonal for numerical stability
COV += jnp.eye(COV.shape[0]) * 1e-8
# WRES = mat_sqrt_inv(COV) @ RES
# OBJI = jnp.log(jnp.linalg.det(COV)) + WRES.T @ WRES
OBJI = jnp.log(jnp.linalg.det(COV)) + RES @ jnp.linalg.solve(COV, RES)
OBJ += OBJI

# print(f"Objective function: {OBJ}")
return jnp.array(OBJ, dtype=jnp.float32)

# JIT the objective function for performance
obj_fn_jit = jax.jit(obj_fn)

# Objective function from final NONMEM estimates
# print("Objective function from final NONMEM estimates:")
# print(obj_fn(jnp.concatenate([THETA, OMEGA.diagonal(), SIGMA.diagonal()]), DATASET))
print(obj_fn_jit(jnp.concatenate([THETA, OMEGA.diagonal(), SIGMA.diagonal()]), DATASET))
# Initial parameters
initpar = jnp.array([
3.0, 35.0, 0.09, # THETA
0.9, 0.02, 0.06, # OMEGA diagonal
0.01, 0.03 # SIGMA diagonal
])

# print("Objective function from initial parameters:")
# print(obj_fn_jit(initpar, DATASET))

def obj_fn_grad(params, DATASET):
# Compute both value and gradient of the objective function
value, grad = jax.value_and_grad(obj_fn, argnums=0)(params, DATASET)
return(value, grad)

obj_fn_grad_jit = jax.jit(obj_fn_grad)
# print("Objective function and gradient from initial parameters:")
# print(obj_fn_grad_jit(initpar, DATASET))


# 2. Optimizer Initialization
optimizer = jaxopt.LBFGSB(
# fun=obj_fn_jit, # Objective function to minimize
fun=obj_fn_grad_jit, # Objective function to minimize
value_and_grad=True, # Return both value and gradient
tol=1e-6, # Tolerance for convergence
# maxiter=1000, # Maximum number of iterations
# verbose=True, # Print optimization progress
jit=True
)

l_bound = jnp.array([
0.0, # Lower bound for THETA[0]
0.0, # Lower bound for THETA[1]
0.0, # Lower bound for THETA[2]
-10.0, # Lower bound for OMEGA[0,0]
-10.0, # Lower bound for OMEGA[1,1]
-10.0, # Lower bound for OMEGA[2,2]
-10.0, # Lower bound for SIGMA[0,0]
-10.0 # Lower bound for SIGMA[1,1]
])
u_bound = jnp.array([
100.0, # Upper bound for THETA[0]
100.0, # Upper bound for THETA[1]
100.0, # Upper bound for THETA[2]
10.0, # Upper bound for OMEGA[0,0]
10.0, # Upper bound for OMEGA[1,1]
10.0, # Upper bound for OMEGA[2,2]
10.0, # Upper bound for SIGMA[0,0]
10.0 # Upper bound for SIGMA[1,1]
])

# 3. Step function for optimization
# opt_state = optimizer.init(initpar)
start_time = time.time()
solution = optimizer.run(init_params=initpar, bounds=(l_bound, u_bound), DATASET=DATASET)
execution_time = time.time() - start_time
print(f"Optimization completed in {execution_time:.4f} seconds.")
print("Final parameter values:",solution.params)
print("Final objective function value:", solution.state.value)


## To print the optimization progress step by step ---------------------------
current_params = initpar

# Initialize the optimizer's state.
# DATASET and bounds are passed as they are needed by the objective function
# or the LBFGSB algorithm during initialization and updates.
# The fun (obj_fn_grad_jit) will be called here to get the initial value and grad.
opt_state = optimizer.init_state(current_params, bounds=(l_bound, u_bound), DATASET=DATASET)

max_total_updates = 1000 # Define a maximum number of total updates for our loop
print_every_n_updates = 5
total_updates_done = 0

print(f"Initial parameters: {current_params}")
print(f"Initial objective value: {opt_state.value:.6f}")
if hasattr(opt_state, 'grad') and opt_state.grad is not None:
print(f"Initial gradient norm: {jnp.linalg.norm(opt_state.grad):.6e}")
print(f"Initial optimality error: {opt_state.error:.6e}")

start_time = time.time()

# Outer loop: controls how many times we print results
for reporting_step in range(max_total_updates // print_every_n_updates + 1):
if total_updates_done >= max_total_updates:
break

# Inner loop: perform `print_every_n_updates`
for _ in range(print_every_n_updates):
if total_updates_done >= max_total_updates:
break

current_params, opt_state = optimizer.update(
params=current_params,
state=opt_state,
bounds=(l_bound, u_bound),
DATASET=DATASET
)
total_updates_done += 1

if opt_state.error < optimizer.tol: # Check for convergence after each step
break

print(f"\n--- After {total_updates_done} updates (Reporting step {reporting_step + 1}) ---")
print(f" Current parameters: {current_params}")
print(f" Objective value: {opt_state.value:.6f}")
print(f" Optimality error: {opt_state.error:.6e}")
print(f" Solver iterations (total LBFGSB updates): {opt_state.iter_num}")

if opt_state.error < optimizer.tol:
print(f"\nConvergence reached after {total_updates_done} updates.")
break

execution_time = time.time() - start_time
print(f"Optimization completed in {execution_time:.4f} seconds.")
print(f"Final parameter values: {solution.params}")
print(f"Final objective function value: {solution.state.value:.6f}")

Continue reading...