Why is the number of timesteps smaller than the number of forward calls?

[sambaPython24]

Hey, thank you for the work that you put into this package!

I have a general question to understand how the odeint solver works. Suppose, we have the following example of an oscillator:

from torchdiffeq import odeint
import matplotlib.pyplot as plt

device = torch.device('cuda'  if torch.cuda.is_available() else 'cpu')

t       = torch.linspace(0, 10, 101).to(device)

true_y0 = torch.tensor([[np.pi - 0.1, 0.0]]).to(device)

true_A  = torch.tensor([[-0.1, 2.0],[-2.0, -0.1]]).to(device)

class Lambda(nn.Module):
    def __init__(self):
        super(Lambda, self).__init__() # type: ignore        
        self.index = 0    
    def forward(self, t : torch.Tensor, y : torch.Tensor):
        self.index += 1
        b   = 0.25
        c   = 5.0
        theta, omega = y[0,0],y[0,1]
        dydt = torch.tensor([omega, -b*omega - c*torch.sin(theta)],device = y.device)
        dydt = dydt.unsqueeze(0)
        return dydt

lambda_object = Lambda()
with torch.no_grad():
    true_y = odeint(func = lambda_object, y0 = true_y0, t = t, method='dopri5')
print(f"The forward method was called {lambda_object.index} times whule the length of the t is {len(t)}.")
plt.title("ODE Integration (torchdiffeq)")
plt.plot(t.detach().cpu(),true_y[:,0,0].detach().cpu(), 'b', label='theta(t)')
plt.plot(t.detach().cpu(),true_y[:,0,1].detach().cpu(), 'g', label='omega(t)')
plt.legend(loc='best')
plt.xlabel('t')
plt.grid()
plt.show()

When I count the number of times the forward function is called, I would expect, that it equals the number of time steps in the $t$ array, but the two numbers differ.

Why is this the case?

Ultimatively, I would like to add an external force that is called by an self.index in the forward method.

[psv4]

Dopri5 Tableau

Depending on your method, the forward function can be called multiple times for each integration. Starting with the simplest case of explicit Euler, your forward function is called one time and uses that gradient to integrate (one forward call per time step). For explicit midpoint, your forward function is called two times (initial gradient to get the y value halfway, then the gradient at that halfway point). So, for the explicit midpoint method, it is now two forward calls per time step.

dopri5 is an adaptive step method. It also computes a second estimate for error and can recompute the integration if the error is too large. With these methods, the number of forward calls can change depending on your error tolerance as well.