how to do curve fitting using google jax?

Question

Extending the examples from http://implicit-layers-tutorial.org/neural_odes/ I am tying to mimic the curve fitting function in scipy , scipy.optimize.curve_fit ,using google jax. The function to be fitted is a first order ODE.

#Generate toy data for first order ode.

import jax.numpy as jnp
import jax
import numpy as np


#input  data 
u = np.zeros(100)  
u[10:50] = 1
t = np.arange(len(u))
u = jnp.array(u)

#first order ODE
def f(y,t,k,tau,u):
 
  return (k*u[t]-y)/tau
  
#Euler integration
def odeint_euler(f, y0, t, *args):
  def step(state, t):
    y_prev, t_prev = state
    dt = t - t_prev
    y = y_prev + dt * f(y_prev, t_prev, *args)
    return (y, t), y
  _, ys = jax.lax.scan(step, (y0, t[0]), t[1:])
  return ys

pred = odeint_euler(f, jnp.array([0.0]),t,2.,5.,u) 
pred_noise = pred.reshape(-1) +  0.05* np.random.randn(len(pred)) # this is the  data to be fitted

# define loss function 
def loss_function(params,u,targets):
  k,tau = params
  
  pred = odeint_euler(f, jnp.array([0.0]),t,k,tau,u)
  return jnp.sum((pred-targets)**2)      


def update(params, u, targets):
  grads = jax.grad(loss_function)(params,u, targets)
  return [w - 0.0001 * dw for w,dw  in zip(params, grads)] 


updated_params = jnp.array([1.0,2.0]) #initial parameters
for i in range(100):
  updated_params = update(updated_params, u, pred_noise)
print(updated_params)

The code works fine. However , this runs pretty slow when compared to scipy curve fit. The accuracy of the solution is not good even after 500, 1000 iterations. What is wrong with the above code ? Any idea how to make the code run faster and to get more accurate solution? Is there any better way of doing the curve fitting with jax?

user_2937831 · Answer

I see two overall issues with your approach:

The reason your code is running slowly is because you are doing your looping in Python, which incurs JAX's dispatch overhead every loop. I'd recommend using JAX's built-in tools for minimization of loss functions; for example:

from jax.scipy.optimize import minimize
result = minimize(
    loss_function, x0=jnp.array([1.0,2.0]),
    method='BFGS', args=(u, pred_noise))

The reason your accuracy does not approach that of scipy is likely because JAX defaults to 32-bit computations (See Double (64 bit) Precision). To run your code in 64-bit, you can run this block before any other imports:

from jax import config
config.update('jax_enable_x64', True)

how to do curve fitting using google jax?

1 Answers