Neural ODEs with non-uniform time intervals

[jalving]

Greetings! I have been using Neuromancer to perform nonlinear system identification with considerable success. I am now interested in learning and using neural controlled ODEs with non-uniform control points which has led to a few questions about how the integrators work with respect to time rollouts.

• Does Neuromancer support time rollouts over arbitrary intervals? I would like to do something like in the below figure where I apply a control input over two separate time intervals with different ranges for the integration limits $t_{N}$ and $t_{N_2}$. Perhaps put another way, is it possible to prescribe input times for controls and output times to observe states?

• If the above works, is it possible to also train neural ODEs with data that is non-uniformly sampled in time? This is not necessary for my application; I am just curious.

• How is the integration step handled internally? I notice the examples pick the step size to be the sample time of the data, but is this strictly required? It seems to train fine with different time steps, so I'm curious how the integration step translates to the rollout step.

Thanks for any clarification. This has been a great tool for my work.

[drgona]

Hi @jalving,

Thanks for your interest in the library! As of now, our integrators support fixed timesteps for training NODEs.

There are three ways you could implement NODE with variable timesteps in Neuromancer.

Option 1: Modify existing integrators to take timestep as extra input instead of attribute h.
This would require modifying the forward pass of the selected integrator. For instance, Euler would need to be modified as:

class EulerVT(Integrator):
    def __init__(self, block):
        """
        Euler integrator with timestep as control input
        """
        super().__init__(block=block, interp_u=None, h=1.0)

def integrate(self, x, h, *args):
        k1 = self.block(*[x, *args])       
        return x + h*k1

https://github.com/pnnl/neuromancer/blob/master/src/neuromancer/dynamics/integrators.py#L86

You would need to introduce a new input variable (lets label it by 'h') in your dataset that will provide the corresponding timestep values, then you can treat the integrator with this extra positional input 'h' for timestep together with control input 'u'. Something like:

integrator = integrators.EulerVT(ode)
Node(integrator, ['xn', 'h', 'u'], ['xn'])

Option 2: Modify torchdiffeq integrator to support the architecture you need.
Neuromancer provides a wrapper for torchdiffeq integrator from the authors of the original NODE paper.
torchdiffeq is nice because it allows you to specify a vector of timepoints at which to solve the ODE.
https://github.com/pnnl/neuromancer/blob/master/src/neuromancer/dynamics/integrators.py#L77

It can take irregularly sampled intervals, this is straightforward for systems without control inputs. However, for systems with controls you need to handle controls via interpolations. There is a related issue at torchdiffeq github repo that might be relevant:
rtqichen/torchdiffeq#128
The proposed solution is to include additional closure function to interpolate the controls at corresponding timesteps.

We used to have this interpolation supported via interp_u method in our integrators, but it was hard to maintain for all integrators and we deprecated this functionality.
https://github.com/pnnl/neuromancer/blob/master/src/neuromancer/dynamics/integrators.py#L27C12-L27C21
https://github.com/pnnl/neuromancer/blob/master/src/neuromancer/dynamics/interpolation.py

Option 3:
Leverage the Node abstraction to create a symbolic computational graph where each Node integrates the NODE using integrators with different timesteps. Simplest approach, but a bit clumsy with more code.
Here instead of doing rollout using the System class over nsteps (our symbolic RNN constructor), we can build the unrolled computational graph manually.

For instance, lets take this example:
https://colab.research.google.com/github/pnnl/neuromancer/blob/master/examples/ODEs/Part_2_param_estim_ODE.ipynb#scrollTo=f62e2ab0&line=4&uniqifier=1
I can create multiple integrators with different timesteps and plug them together in the computational graph using System class with nsteps=1. You can now inject controls at different timesteps

int1 = integrators.RK4(ode, h=0.1)
int2 = integrators.RK4(ode, h=0.5)
dynamics_model = System([Node(int1, ['xn', 'u1'], ['xn']), Node(int2, ['xn', 'u2'], ['xn']) ], nsteps=1)

[jalving]

Thanks @drgona! This more than answers all of my questions. I will experiment with these options; getting the torchdiffeq integrators to work seems particularly promising. I think the approach will be quite useful for performing direct multiple shooting with trained neuromancer neural ODEs.