pyzag.ode

Module for solving ordinary differential equations using numerical integration

class pyzag.ode.BackwardEulerODE(*args, **kwargs)

Applies a backward Euler time integration scheme to an ODE

The input is a torch Module which defines the rate form of the ODE. The forward function must return both the time rate of change of the state and the derivative of the rate of change with respect to the state as a function of time and the current state.

The input and output for the underlying ODE must meet our global batch convention, where the first dimension of the input is used to vectorize time/step.

Parameters:

ode (torch.nn.Module) – module defining the system of ODEs

forward(x, t)

Provide the blocked residual and Jacobian for the time integration scheme

Parameters:

• x (torch.tensor) – (nchunk+self.n,…,nstate) tensor. As always, the first batch dimension is used to vectorize over steps (for an ODE, time steps). Then the operator must accept arbitrary batch dimensions. The final dimension is the problem state size.

• t (torch.tensor) – (nchunk+self.n,…,) tensor of times as driving forces

Returns:

tuple giving R the (nchunk,…,nstate ) tensor giving the nonlinear residual and J the (self.n, nchunk,…,nstate,nstate) tensor giving the Jacobians

Return type:

(R,J) tuple(torch.tensor, torch.tensor)

class pyzag.ode.ForwardEulerODE(*args, **kwargs)

Applies a forward Euler time integration scheme to an ODE

The input is a torch Module which defines the rate form of the ODE. The forward function must return both the time rate of change of the state and the derivative of the rate of change with respect to the state as a function of time and the current state.

The input and output for the underlying ODE must meet our global batch convention, where the first dimension of the input is used to vectorize time/step.

Parameters:

ode (torch.nn.Module) – module defining the system of ODEs

forward(x, t)

Provide the blocked residual and Jacobian for the time integration scheme

Parameters:

• x (torch.tensor) – (nchunk+self.n,…,nstate) tensor. As always, the first batch dimension is used to vectorize over steps (for an ODE, time steps). Then the operator must accept arbitrary batch dimensions. The final dimension is the problem state size.

• t (torch.tensor) – (nchunk+self.n,…,) tensor of times as driving forces

Returns:

tuple giving R the (nchunk,…,nstate ) tensor giving the nonlinear residual and J the (self.n, nchunk,…,nstate,nstate) tensor giving the Jacobians

Return type:

(R,J) tuple(torch.tensor, torch.tensor)

class pyzag.ode.IntegrateODE(ode, *args, **kwargs)

Maps an ODE to a nonlinear function with some numerical integration scheme

The input is a torch Module which defines the rate form of the ODE. The forward function must return both the time rate of change of the state and the derivative of the rate of change with respect to the state as a function of time and the current state.

The input and output for the underlying ODE must meet our global batch convention, where the first dimension of the input is used to vectorize time/step.

Parameters:

ode (torch.nn.Module) – module defining the system of ODEs