Adaptive timestepping and hijacking the "step" method

[tbridel]

Hi again @Nicholaswogan !

I was wondering if you ever thought about the adaptive timestepping capability that exists in Scipy LSODA ? Is anything like this possible with your implementation ?

Thanks !
Best,
T. Bridel-Bertomeu

[Nicholaswogan]

Can you explain what you mean by adaptive time stepping?

[tbridel]

Yes of course.

I saw that in all the examples you expose, you specify manullay an array of instants in time whereat the solver has to solve the ODE (the t_eval argument I believe) - basically you are setting dt.
I would like to just be able to specify t0, tmax and let the solver do its job to infer from each state n the proper timestep to reach state n+1 until it reaches tmax.

[Nicholaswogan]

Ive implemented this with the DOP853 solver but not lsoda. Check out this notebook here

https://github.com/Nicholaswogan/numbalsoda/blob/main/solve_ivp.ipynb

It is sort of a clone of solve_ivp from scipy, but much faster. It only works with DOP853 method right now. With this, you don’t have to specify the times to compute the solution.

[tbridel]

Ah, great, thank you, I'll have a look.

[tbridel]

I changed the title because I have a second question actually.

For now the way I use scipy functionality is that I declare my solver like solver = scipy.integrate.RK45(...) and then I write the loop myself because I need to do some extra stuff after one step of the solver:

def solve(...):
  while iteration_number < max_iteration_number + 1:
    solver.step()
    do_action_1()
    do_action_2()

   if solver.status is finished or failed:
     break
  
   iteration_number += 1

 do_action_3_before_leaving_solve()

From what I understood from your implementation, the speed of execution also comes from the fact that you have a low-level implementation of the entire loop, so I guess exposing just the stepping part would not really make sense for the logic of your work - right ?
What would you suggest to go in this direction ?

[Nicholaswogan]

It’s possible to expose each step in the way you describe and also have high performance. I want to do this some day. But there is an issue with numba that prevents it:

numba/numba#8470 (comment)

Basically, you have some memory safety issues if it were implemented with current versions of numba

[tbridel]

OK thank you, I see you wanted to head that way fairly recently too ! So is it a feature you see being added to numbalsoda in the future with high certainty ?

Other possibility maybe, tell me what you think: in the same way you pass the rhs function all the way down to the F90 code of DOP853, could we imagine passing down a "postStep" function, originally written in the python code of the user and with a given signature ?

[Nicholaswogan]

I will eventually implement this new way. But not sure how long until I get around to it, and if numba will end up fixing the issue I linked.

what would the signature of this post step method be?

[tbridel]

Would it make sense to expect a method that takes as arguments the current state vector y (which would then be the state right after the step), the current time t (again, right after the step) and the current iteration number ?

I believe it is fair that an ODE solver could pass this to a third-party function, but also it seems reasonable in the sense that we cannot expect any more variables from an ODE solver in a general context.

So I would maybe suggest to go for:

poststep_sig = types.void(types.CPointer(double), types.double, types.int64)

What do you think ?

[Nicholaswogan]

So would the function be used to just print stuff? Or would the purpose be to save the state of the calculation after the step? I think to save the state of the calculation, we would need a different signature than the one you put.

Another possibility is to implement everything in an alternative way which exposes the core of the integrator, which has the possibility of being memory unsafe. This is probably bad form in Python, because people are not use to memory issues.

[tbridel]

No the function could be used for more complex actions. Imagine for instance this use case, where the whole solving logic is embedded in a class and relies on instances of other classes - just to put this in context, my application focuses around flight dynamics here.

class Solver():
  def __init__(self, t0, tf, itmax, **vehicle_characs, **atmos_characs, **planet_characs):
    self.t0 = t0
    self.tf = tf
    self.itmax = itmax
    self.itnum = 0
    self.vehicle = Vehicle(**vehicle_characs)
    self.atmosphere = Atmosphere(**atmos_characs)
    self.planet = Planet(**planet_characs)
    self.solver = RK45(
      self.rhs, 
      self.t0,
      self.vehicle.state,
      self.tf
    )

  def rhs(self, t, y):
    x, y, z = get_position(y)
    vx, vy, vz = get_velocity(y)
    
    gx, gy, gz = self.planet.gravity(x, y, z)
    
    self.atmosphere.update(
      x, y, z, t - self.t0
    )

    self.vehicle.compute_body_deformation_with_fem(...)

    ...
   
    def _assemble_flux():
      return np.array([vx, vy, vz, gx, gy, gz, <stuff for rotation dof>, <stuff for rotation velocity>])

    return _assemble_flux()

  def solve(self):
    while self.itnum < self.itmax + 1:
      self.solver.step()

      self.vehicle.store_in_history(self.solver.t, self.solver.y, self.itnum)
      self.vehicle.compute_auxiliary_characteristics(self.solver.y)
      ... # other auxiliary computations for the vehicle, atmosphere, etc... instances

So basically the post step here serves as

• storage of the computation state (yes indeed),
• update auxiliary characteristics of all the objects involved in the computation.

[tbridel]

Writing this answer actually made me realize that if I re-arrange the code the slightest bit, everything could be in the RHS 🤔