autograd and odeint not working properly

[aoguedao]

Hi! First of all, thanks to the developers for this package, I had no idea about it and I have been exploring it since I want to use for a workshop with some undergraduate student.

For some context, I am writing a very simplex gradient descent algorithm to estimate parameters from a SIR model. For the loss function, I get a numerical approximation using scipy.integrate.odeint, which as far as I know is supported (here).

import autograd
import autograd.numpy as np
import matplotlib.pyplot as plt

from scipy.integrate import odeint
from autograd import grad, jacobian


def sir(y, t, beta, gamma):
  S, I, R = y
  dS_dt = - beta * S * I
  dI_dt = beta * S  * I - gamma * I
  dR_dt = gamma * I
  return np.array([dS_dt, dI_dt, dR_dt])

def loss(params, Y0, t, y_obs):
  beta, gamma = params
  sol = odeint(sir, y0=Y0, t=t, args=(beta, gamma))
  err = np.linalg.norm(y_obs - sol, 2)
  return err

# Generate data
np.random.seed(42)
Y0 = np.array([0.95, 0.05, 0.0])
t = np.linspace(0, 30, 101)
beta, gamma = 0.5, 1/14
sol = odeint(sir, y0=Y0, t=t, args=(beta, gamma))
y_obs = sol + np.random.normal(0, 0.05, size=sol.shape)
# plt.plot(t, y_obs)

Then, what I would like to do is something like this

# --- THIS DOES NOT WORK ---
params = np.array([0.4, 0.1])  # Initial guess of beta and gamma 
loss_grad = grad(loss, argnum=0)  # params is the first argument of loss
# Perform gradient descent
for i in range(n_iterations):
  grads = loss_grad(params, Y0, t, y_obs)  # Compute gradients
  params -= learning_rate * grads  # Update parameters

However, evaluating the gradient gives me an error

params = np.array([0.4, 0.1])
loss_grad = grad(loss, argnum=0)
grads = loss_grad(params, Y0, t, y_obs)

I get the following error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
TypeError: float() argument must be a string or a real number, not 'ArrayBox'

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In [5], [line 8](vscode-notebook-cell:?execution_count=5&line=8)
      [5](vscode-notebook-cell:?execution_count=5&line=5) gamma_init = 1.0
      [6](vscode-notebook-cell:?execution_count=5&line=6) params = np.array([beta_init, gamma_init])
----> [8](vscode-notebook-cell:?execution_count=5&line=8) grads = loss_grad(params, Y0, t, y_obs)
      [9](vscode-notebook-cell:?execution_count=5&line=9) # loss_jac(params, Y0, t, y_obs)

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     [18](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:18) else:
     [19](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:19)     x = tuple(args[i] for i in argnum)
---> [20](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:20) return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:28, in grad(fun, x)
     [21](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:21) @unary_to_nary
     [22](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:22) def grad(fun, x):
     [23](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:23)     """
     [24](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:24)     Returns a function which computes the gradient of `fun` with respect to
     [25](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:25)     positional argument number `argnum`. The returned function takes the same
     [26](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:26)     arguments as `fun`, but returns the gradient instead. The function `fun`
     [27](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:27)     should be scalar-valued. The gradient has the same type as the argument."""
---> [28](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:28)     vjp, ans = _make_vjp(fun, x)
     [29](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:29)     if not vspace(ans).size == 1:
     [30](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:30)         raise TypeError("Grad only applies to real scalar-output functions. "
     [31](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/differential_operators.py:31)                         "Try jacobian, elementwise_grad or holomorphic_grad.")

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:10, in make_vjp(fun, x)
      [8](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:8) def make_vjp(fun, x):
      [9](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:9)     start_node = VJPNode.new_root()
---> [10](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:10)     end_value, end_node =  trace(start_node, fun, x)
     [11](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:11)     if end_node is None:
     [12](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/core.py:12)         def vjp(g): return vspace(x).zeros()

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:10, in trace(start_node, fun, x)
      [8](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:8) with trace_stack.new_trace() as t:
      [9](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:9)     start_box = new_box(x, t, start_node)
---> [10](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:10)     end_box = fun(start_box)
     [11](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:11)     if isbox(end_box) and end_box._trace == start_box._trace:
     [12](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/tracer.py:12)         return end_box._value, end_box._node

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:15, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f(x)
     [13](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:13) else:
     [14](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:14)     subargs = subvals(args, zip(argnum, x))
---> [15](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/autograd/wrap_util.py:15) return fun(*subargs, **kwargs)

Cell In [2], [line 11](vscode-notebook-cell:?execution_count=2&line=11)
      [9](vscode-notebook-cell:?execution_count=2&line=9) beta, gamma = params
     [10](vscode-notebook-cell:?execution_count=2&line=10) # Solve the ODE system using odeint
---> [11](vscode-notebook-cell:?execution_count=2&line=11) sol = odeint(sir, y0=Y0, t=t, args=(beta, gamma))
     [13](vscode-notebook-cell:?execution_count=2&line=13) # Compute the L2 norm error between the observed and predicted values
     [14](vscode-notebook-cell:?execution_count=2&line=14) err = np.linalg.norm(y_obs - sol, 2)

File ~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:254, in odeint(func, y0, t, args, Dfun, col_deriv, full_output, ml, mu, rtol, atol, tcrit, h0, hmax, hmin, ixpr, mxstep, mxhnil, mxordn, mxords, printmessg, tfirst)
    [251](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:251) y0 = copy(y0)
    [253](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:253) with ODE_LOCK:
--> [254](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:254)     output = _odepack.odeint(func, y0, t, args, Dfun, col_deriv, ml, mu,
    [255](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:255)                             full_output, rtol, atol, tcrit, h0, hmax, hmin,
    [256](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:256)                             ixpr, mxstep, mxhnil, mxordn, mxords,
    [257](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:257)                             int(bool(tfirst)))
    [258](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:258) if output[-1] < 0:
    [259](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:259)     warning_msg = (f"{_msgs[output[-1]]} Run with full_output = 1 to "
    [260](https://vscode-remote+wsl-002bubuntu.vscode-resource.vscode-cdn.net/home/alonsolml/git/quantitative-biology/chapters/chapter-03/~/mambaforge/envs/nc-book/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:260)                    f"get quantitative information.")

ValueError: setting an array element with a sequence.

I am working on a .ipynb file on VSCode using WSL, Python 3.10.9 (mamba venv) and

Package                       Version
----------------------------- -----------
autograd                      1.7.0
numpy                         1.26.4
scipy                         1.15.2

PS: I already asked on stackoverflow (here). The proposed solution didn't work either since (afaik) assignation is not supported in autograd.

Thanks in advance!

[aoguedao]

I noticed it I was using the scipy version and I just tried using the wrapped one with from autograd.scipy.integrate import odeint and it didn't work either :(

[fjosw]

Hey @aoguedao, looks like you found a solution for your problem in https://stackoverflow.com/a/79489985 ? I'll close the issue for now, feel free to reopen it if you still run into problems.