Using batches in a typical nonlinear optimization

[richinex]

I have provided a minimal working example of a nonlinear optimization I try to do in Jax. Basically I fit sets of complex spectra using a rational polynomial in order to obtain the optimal coefficients. In this example I have just 3 spectra in the list but in reality I could have 100s. My question is whether it is possible to use a batch size and how that would look like and secondly whether further optimization might help the code run faster?
Here is my MWE:

import numpy as np
import scipy.sparse as sps
import jax
from jax import value_and_grad
from jax import numpy as jnp
from jax.example_libraries import optimizers as jax_opt
jax.config.update("jax_enable_x64", True)
from time import time

# Load data

F = [jnp.array([4.000000e+00, 1.250000e+01, 2.100000e+01, 2.950000e+01,3.800000e+01, 4.650000e+01, 6.350000e+01, 8.050000e+01,
        1.230000e+02, 1.570000e+02, 1.995000e+02, 2.505000e+02,3.185000e+02, 3.940000e+02, 4.965000e+02, 6.260000e+02,7.925000e+02, 9.975000e+02, 1.256000e+03, 1.581000e+03,
        1.990500e+03, 2.506000e+03, 3.155000e+03, 3.971500e+03,5.000000e+03, 6.294500e+03, 7.924500e+03, 9.976500e+03,1.255950e+04, 1.581150e+04, 1.990550e+04, 2.505950e+04,
        3.154800e+04, 3.971650e+04, 5.000000e+04]),
    jnp.array([4.000000e+00, 1.250000e+01, 2.100000e+01, 2.950000e+01,3.800000e+01, 4.650000e+01, 6.350000e+01, 8.050000e+01,1.230000e+02, 1.570000e+02, 1.995000e+02, 2.505000e+02,
        3.185000e+02, 3.940000e+02, 4.965000e+02, 6.260000e+02,7.925000e+02, 9.975000e+02, 1.256000e+03, 1.581000e+03,1.990500e+03, 2.506000e+03, 3.155000e+03, 3.971500e+03,
        5.000000e+03, 6.294500e+03, 7.924500e+03, 9.976500e+03,1.255950e+04, 1.581150e+04, 1.990550e+04, 2.505950e+04,3.154800e+04, 3.971650e+04, 5.000000e+04]),
    jnp.array([4.000000e+00, 1.250000e+01, 2.100000e+01, 2.950000e+01,3.800000e+01, 4.650000e+01, 6.350000e+01, 8.050000e+01,1.230000e+02, 1.570000e+02, 1.995000e+02, 2.505000e+02,
        3.185000e+02, 3.940000e+02, 4.965000e+02, 6.260000e+02,7.925000e+02, 9.975000e+02, 1.256000e+03, 1.581000e+03,1.990500e+03, 2.506000e+03, 3.155000e+03, 3.971500e+03,
        5.000000e+03, 6.294500e+03, 7.924500e+03, 9.976500e+03,1.255950e+04, 1.581150e+04, 1.990550e+04, 2.505950e+04,3.154800e+04, 3.971650e+04, 5.000000e+04]),
]


Y = [jnp.array([0.00495074+0.00290374j, 0.00724701+0.00289439j,0.00821288+0.00279885j, 0.00877054+0.00276919j,
        0.00921332+0.0027551j , 0.00953043+0.00274739j,0.01002155+0.00274946j, 0.01038829+0.00279736j,
        0.01103745+0.00293741j, 0.01143682+0.00304808j,0.01185019+0.00321095j, 0.01222892+0.00340771j,
        0.01264666+0.00365856j, 0.01312294+0.00390083j,0.01356835+0.00423682j, 0.01414305+0.00459166j,
        0.01475416+0.00502188j, 0.01544523+0.0054795j ,0.01620464+0.00597393j, 0.01707565+0.00650766j,
        0.01800564+0.00707323j, 0.01907494+0.00766403j,0.0202539 +0.00824607j, 0.02156295+0.00882627j,
        0.02293967+0.0093636j , 0.02446602+0.00988404j,0.02606663+0.01034258j, 0.02778773+0.01073912j,
        0.0295645 +0.01105176j, 0.03142458+0.01130524j,0.03332406+0.01142638j, 0.03529196+0.01141756j,
        0.03725344+0.01128458j, 0.03917468+0.01100424j,0.04104471+0.0105539j  ]),
        jnp.array([0.00591629+0.00332144j, 0.00851045+0.00319077j,
        0.0095824 +0.00303868j, 0.01017932+0.00296819j,0.01064447+0.00291851j, 0.0109805 +0.00289757j,
        0.01149782+0.00285981j, 0.01187611+0.00288041j,0.01255219+0.00297316j, 0.01296647+0.00305608j,
        0.01335927+0.00317665j, 0.01375042+0.00334631j,0.01416365+0.00354991j, 0.01462764+0.00375336j,
        0.01508156+0.00403475j, 0.01560866+0.00434549j,0.01617377+0.00471426j, 0.01681709+0.00512406j,
        0.01752344+0.00555375j, 0.01832037+0.00605304j,0.01916483+0.00655218j, 0.02014014+0.00711006j,
        0.0211973 +0.0076624j , 0.02239568+0.00822725j,0.0236615 +0.00875432j, 0.025081  +0.00927988j,
        0.02656711+0.00977482j, 0.02817824+0.010206j  ,0.02987373+0.01057232j, 0.03166605+0.01084335j,
        0.03351033+0.01101296j, 0.03539926+0.01105692j,0.0373132 +0.01098819j, 0.03920738+0.01075104j,
        0.0410484 +0.01034544j]),
        jnp.array([0.00698985+0.00375178j, 0.00987582+0.00348102j,0.01101949+0.00328484j, 0.01167011+0.00315813j,
        0.0121601 +0.00307674j, 0.01251025+0.00303722j,0.01305459+0.00295732j, 0.01344384+0.00295235j,
        0.01412694+0.00299687j, 0.01455009+0.00306043j,0.01493542+0.00313616j, 0.01533892+0.00327267j,
        0.01575224+0.0034428j , 0.01616527+0.00359421j,0.01664215+0.0038313j , 0.01713806+0.00410226j,
        0.01766022+0.00442292j, 0.01825988+0.00478992j,0.01892454+0.00516821j, 0.01963565+0.00558726j,
        0.02042078+0.00605703j, 0.02129398+0.00655357j,0.02225692+0.00707153j, 0.02332775+0.00760602j,
        0.02449876+0.00812741j, 0.02577426+0.00864679j,0.0271547 +0.00916218j, 0.02865093+0.00963409j,
        0.03025404+0.01002012j, 0.03193923+0.0103418j,0.03371385+0.01055021j, 0.03552156+0.01065117j,
        0.03738471+0.01063324j, 0.03924574+0.01046488j,0.04104824+0.01014013j])]


sigma = [jnp.array([2.43219802e-06, 3.84912892e-06, 4.65468565e-06, 5.23176095e-06,5.68508176e-06, 6.05872401e-06, 6.64642994e-06, 7.11385064e-06,
        7.95151846e-06, 8.43719499e-06, 8.92535354e-06, 9.37367440e-06,9.85436691e-06, 1.02790955e-05, 1.07571723e-05, 1.12416874e-05,
        1.17638756e-05, 1.22902720e-05, 1.28422944e-05, 1.34043157e-05,1.39690355e-05, 1.45196518e-05, 1.50516798e-05, 1.55538437e-05,
        1.60391119e-05, 1.65177389e-05, 1.69736650e-05, 1.74361339e-05,1.78881437e-05, 1.83461307e-05, 1.87868263e-05, 1.92436037e-05,
        1.96903675e-05, 2.01275998e-05, 2.05546330e-05]),
        jnp.array([2.43658201e-06, 3.86081138e-06, 4.67159089e-06, 5.25138466e-06,5.70892826e-06, 6.08476012e-06, 6.67752010e-06, 7.14463113e-06,
        7.99083227e-06, 8.47998126e-06, 8.96919300e-06, 9.42881525e-06,9.90748958e-06, 1.03379216e-05, 1.08229442e-05, 1.13118876e-05,
        1.18282487e-05, 1.23585069e-05, 1.29251703e-05, 1.34936708e-05,1.40533757e-05, 1.45997838e-05, 1.51278800e-05, 1.56280821e-05,
        1.61002881e-05, 1.65664605e-05, 1.70139119e-05, 1.74691468e-05,1.79091330e-05, 1.83603006e-05, 1.88002414e-05, 1.92456519e-05,
        1.96931851e-05, 2.01247949e-05, 2.05527485e-05]),
        jnp.array([2.44242051e-06, 3.87566070e-06, 4.69175257e-06, 5.27434395e-06,5.73617535e-06, 6.11477162e-06, 6.71176349e-06, 7.18220417e-06,
        8.03600778e-06, 8.53506663e-06, 9.01912881e-06, 9.48978050e-06,9.97460302e-06, 1.04078345e-05, 1.08886052e-05, 1.13899423e-05,
        1.19043279e-05, 1.24448907e-05, 1.30168282e-05, 1.35899490e-05,1.41519176e-05, 1.46950069e-05, 1.52207831e-05, 1.57159302e-05,
        1.61797980e-05, 1.66293976e-05, 1.70720359e-05, 1.75116747e-05,1.79413837e-05, 1.83749271e-05, 1.88138765e-05, 1.92483785e-05,
        1.96921701e-05, 2.01242128e-05, 2.05496508e-05])]


order = 5 # order of the rational function
params_init = jnp.ones(2*order+1) # initial parameters
n_par = len(params_init) # length of initial parameters
n_data = len(F) # number of data to fit
par_mat = jnp.broadcast_to(params_init[:,None], (len(params_init), n_data))

opt_init, opt_update, get_params = jax_opt.adam(1e-3)

opt_state = opt_init(par_mat) # initialize the state based on the input parameters

@jax.jit
def fun(p, x):  
    norder = int((len(p)-1)/2)
    a = jax.nn.softplus(jnp.array(p)[0:norder+1]) # using softplus to constrain the output to be positive
    b = jax.nn.softplus(jnp.concatenate([jnp.array(p)[norder+1:2*norder+1], jnp.array([1])], axis = 0))
    Ypa = jnp.polyval(a,jnp.sqrt(1j*x))/jnp.polyval(b,jnp.sqrt(1j*x))
    return jnp.concatenate([(Ypa).real, (Ypa).imag], axis = 0)


# objective function
@jax.jit
def obj_fun(p, x, y, yerr):
    n_points = len(x)
    dof = (2*n_points-(len(p)))
    y_concat = jnp.concatenate([y.real, y.imag], axis = 0)
    sigma = jnp.concatenate([yerr,yerr], axis = 0)
    y_model = fun(p, x)
    chi_sqr = jnp.linalg.norm(((y_concat - y_model)/sigma))**2
    return (chi_sqr)


# make a multidimensional cost function    
@jax.jit
def cost_fun(P, X, Y, YERR):
    dof = (2*len(X[0])*len(X))-len(P)
    chi = jax.vmap(obj_fun, in_axes=1)(P, jnp.column_stack(X), jnp.column_stack(Y), jnp.column_stack(YERR))
    return (jnp.sum(chi)) / dof



@jax.jit
def train_step(step_i, opt_state, X, Y, YERR):
    net_params = get_params(opt_state)
    loss, grads = value_and_grad(cost_fun, argnums=0)(net_params, X, Y, YERR)
    return loss, opt_update(step_i, grads, opt_state)

from time import time

num_epochs = int(1e6)
loss_history = []
start = time()
for i in range(num_epochs):
    loss, opt_state = train_step(i, opt_state, F, Y, sigma)
    loss_history.append(float(loss))
    if i%int(num_epochs/10)==0:
            print("" + str(i) + ": "
                #+ "loss = " + str(loss)
                + "loss=" + "{:5.3e}".format(loss)
                #+ ", parameters: "+ str(["{:5.3e}".format(p.detach()) for p in param.values() ])
            ) 

end = time()
msg = "training time for {0} iterations = {1:.1f} seconds"
print(msg.format(num_epochs, end-start))

# 0: loss=8.590e+09
# 100000: loss=5.460e+00
# 200000: loss=5.459e+00
# 300000: loss=5.461e+00
# 400000: loss=5.458e+00
# 500000: loss=5.457e+00
# 600000: loss=5.454e+00
# 700000: loss=5.449e+00
# 800000: loss=5.444e+00
# 900000: loss=5.441e+00
# training time for 1000000 iterations = 52.9 seconds

[YouJiacheng]

It seems that you've already used batching. BTW, stack F, Y and sigma outside main optimization loop seems more efficient.