ODE Class

class ODE(method='dopri5', adjoint=False, requires_grad=True, default_type=torch.float64)

Bases: Module

Base class to define, solve, and visualize a system of ODEs.

Parameters:

• method (str) – Name of ODE solver to use. Default is dopri5.

• adjoint (bool) – Whether or not to use the adjoint method for backpropagating through ODE solutions. Default is False.

• requires_grad (bool) – Whether or not gradients should be computed for the tensors defining the ODE system.

• default_type (type) – Set the default torch.Tensor type.

Variables:

odeint (torchdiffeq.odeint or torchdiffeq.odeint_adjoint) – Numerical integrator for a system of ODEs given an initial value. The adjoint method will not be used if requires_grad is False.

get_batch(batch_time, batch_size)

Sample a batch of t, y0, and y.

Parameters:

• batch_time (int) – Number of time points per batch.

• batch_size (int) – Number of initial conditions per batch.

Returns:

• t (torch.tensor of shape (batch_time)) – Tensor of batch_time time points.

• y0 (torch.tensor of shape (batch_size,...,D)) – Initial states for batch_size initial conditions.

• y (torch.tensor of shape (batch_time,batch_size,...,D)) – Solution evaluated at batch_time time points for batch_size initial conditions.

plot_frame(ax, y, ntype=None, vmin=None, vmax=None, alpha=0.9, extent=None)

Plot a single frame of an ODE solution.

Parameters:

• ax (matplotlib.axes) – Axis object on which to display the solution.

• y (torch.tensor) – Solution of the system to plot. The input should be either reshaped to the dimensions of the simulation box or have a final dimension of 2 (only valid for ntype = none).

• ntype (str) –

  Type of normalization to apply when displaying the image. The options are:

  • none – No normalization; point cloud data

  • sym – Symmetric linear normalization scale

  • unit – Linear normalization between 0 and 1

  • mod – Linear normalization modulo \(2 \pi\)

  • log – Logarithmic normalization scale

  • symlog – Symmetric logarithmic normalization scale

  • None – Linear normalization between the min. and max. values of y

  The default is None.

• vmin (float, optional) – Minimum normalization value.

• vmax (float, optional) – Maximum normalization value.

• alpha (float between 0 and 1, optional) – Opacity of the image.

• extent (list or tuple of floats (left, right, bottom, top), optional) – Coordinates of the bounding box of the image.

Returns:

sm – Object mapping scalar data to RGBA color values.

Return type:

matplotlib.cm.ScalarMappable

plot_series(y, ntype=None, vmin=None, vmax=None, clabel=None, colors=None)

Plot a time series of frames of an ODE solution.

Parameters:

• y (torch.tensor) – Solution of the system to plot. The input should be either reshaped to the dimensions of the simulation box or have a final dimension of 2 (only valid for ntype = none). Alternately, y can be a list of 2 solution tensors, which will then be overlaid (only valid for ntype = none).

• ntype (str) –

  Type of normalization to apply when displaying the image. The options are:

  • none – No normalization; point cloud data

  • sym – Symmetric linear normalization scale

  • unit – Linear normalization between 0 and 1

  • mod – Linear normalization modulo \(2 \pi\)

  • log – Logarithmic normalization scale

  • symlog – Symmetric logarithmic normalization scale

  • None – Linear normalization between the min. and max. values of y

  The default is None.

• vmin (float, optional) – Minimum normalization value.

• vmax (float, optional) – Maximum normalization value.

• clabel (str, optional) – Text used to label the colorbar.

solve(t, y0=None, device='cpu', rtol=1e-07, atol=1e-09)

Numerically integrate the ODE system and return the solution at times t.

Parameters:

• t (torch.tensor) – 1-dimensional tensor of evaluation times.

• y0 (torch.tensor of shape (M,...,D)) – Initial state of the system for M initial conditions. D denotes the flattened system size.

• device (str) – The name of the device on which the computation will be performed (e.g. cpu or cuda).

• rtol (float) – Upper bound on relative error. Default is \(10^{-7}\).

• atol (float) – Upper bound on absolute error. Default is \(10^{-9}\).

Returns:

y – Solution evaluated at T time points for M initial conditions. D denotes the flattened system size.

Return type:

torch.tensor of shape (T,M,...,D)

trim(t0=0)

Trim the solution at early time points to exclude initial transients.

Parameters:

t0 (int) – The first index into t at which to return the solution in order to exclude initial transients. The initial state will be set to the solution at this point and initial time set to 0 at this point.

class Kuramoto(args, method='dopri5', default_type=torch.float64)

Bases: ODE

Class to define, initialize, solve, and visualize the Kuramoto model of coupled oscillators:

\begin{eqnarray*} \frac{d\theta_i}{dt} = \omega + K\sum_{j\in N(i)}\sin(\theta_j-\theta_i) \end{eqnarray*}

Parameters:

• args (dict) –

  Dictionary of parameters defining the ODE system:

  • N (int) – Dimension of the simulation box (N x N)

  • L (float) – Length of the real-space simulation box (L x L)

  • v (float) – Intrinsic frequency of the oscillators

  • K (float) – Coupling strength

  • s (float) – Length scale of striped pattern

• method (str) – Name of ODE solver to use. Default is dopri5.

• default_type (type) – Set the default torch.Tensor type.

Variables:

conv (torch.nn.Conv2d) – Convolution operator coupling neighboring oscillators.

forward(t, y)

Evaluate the ODE system at a specified time t and state y.

Parameters:

• t (torch.tensor) – 1-dimensional tensor of the evaluation time points.

• y (torch.tensor of shape (M,1,D)) – State of the system for M initial conditions. D denotes the flattened system size.

Returns:

dy/dt – Derivative of the system.

Return type:

torch.tensor of shape (M,1,D)

init_state(M=1, seed=12)

Randomly generate the initial state(s) of the ODE system.

Parameters:

• M (int) – Number of initial conditions to generate. Default is 1.

• seed (int, optional) – Default seed used to set the state of a random number generator. Default is 12.

Variables:

y0 (torch.tensor of shape (M,1,D)) – Initial state of the system for M initial conditions. D denotes the flattened system size.

class LotkaVolterra(args, method='dopri5', default_type=torch.float64)

Bases: ODE

Class to define, initialize, solve, and visualize a point cloud evolving according to the Lotka-Volterra model:

\begin{eqnarray*} \frac{dx}{dt} & = & \alpha x - \beta xy \\ \frac{dy}{dt} & = & \delta xy - \gamma y \end{eqnarray*}

Parameters:

• args (dict) –

  Dictionary of parameters defining the ODE system:

  • N (int) – Dimension of the simulation box (N x N)

  • L (float) – Length of the real-space simulation box (L x L)

  • R (float) – Radius of particles

  • alpha, beta, gamma, delta (float) – Interaction parameters

• method (str) – Name of ODE solver to use. Default is dopri5.

• default_type (type) – Set the default torch.Tensor type.

forward(t, y)

Evaluate the ODE system at a specified time t and state y.

Parameters:

• t (torch.tensor) – 1-dimensional tensor of the evaluation time points.

• y (torch.tensor of shape (M,N,2)) – State of the system for M initial conditions.

Returns:

dy/dt – Derivative of the system.

Return type:

torch.tensor of shape (M,N,2)

init_state(M=1, seed=12)

Randomly generate the initial state(s) of the ODE system.

Parameters:

• M (int) – Number of initial conditions to generate. Default is 1.

• seed (int, optional) – Default seed used to set the state of a random number generator. Default is 12.

Variables:

y0 (torch.tensor of shape (M,N,2)) – Initial state of the system for M initial conditions.

class ODEGraph(method='dopri5', adjoint=False, requires_grad=True, default_type=torch.float64)

Bases: MessagePassing

Base class to define, solve, and visualize a system of ODEs on a graph.

Parameters:

• method (str) – Name of ODE solver to use. Default is dopri5.

• adjoint (bool) – Whether or not to use the adjoint method for backpropagating through ODE solutions. Default is False.

• requires_grad (bool) – Whether or not gradients should be computed for the tensors defining the ODE system.

• default_type (type) – Set the default torch.Tensor type.

Variables:

odeint (torchdiffeq.odeint or torchdiffeq.odeint_adjoint) – Numerical integrator for a system of ODEs given an initial value. The adjoint method will not be used if requires_grad is False.

forward(x, edge_index)

Update node embeddings x as specified by the message. edge_index gives the indices of source and target nodes connected by edges.

Parameters:

• x (torch.tensor) – Current node embeddings.

• edge_index (torch.tensor of shape (N,2)) – Edge indices from N sources (first column) to N targets (second column).

Returns:

x’ – Updated node embeddings.

Return type:

torch.tensor

get_batch(batch_time, batch_size)

Sample a batch of t, y0, and y.

Parameters:

• batch_time (int) – Number of time points per batch.

• batch_size (int) – Number of initial conditions per batch.

Returns:

• t (torch.tensor of shape (batch_time)) – Tensor of batch_time time points.

• y0 (torch.tensor of shape (batch_size,...,D)) – Initial states for batch_size initial conditions.

• y (torch.tensor of shape (batch_time,batch_size,...,D)) – Solution evaluated at batch_time time points for batch_size initial conditions.

get_eval(T, n=6, d=10)

Get an array of logarithmically spaced evaluation time points.

Parameters:

• T (int) – Maximum time.

• n (int) – Number of time points. Default is 6.

• d (int) – All time points will be rounded to the nearest d. Default is 10.

Returns:

t – Numpy array of n time points.

Return type:

numpy.array of shape (n)

plot_frame(ax, y, ntype=None, vmin=None, vmax=None, alpha=0.9, extent=None)

Plot a single frame of an ODE solution.

Parameters:

• ax (matplotlib.axes) – Axis object on which to display the solution.

• y (torch.tensor) – Solution of the system to plot. The input should be either reshaped to the dimensions of the simulation box or have a final dimension of 3.

• ntype (str) –

  Type of normalization to apply when displaying the image. The options are:

  • none – No normalization; point cloud data

  • sym – Symmetric linear normalization scale

  • unit – Linear normalization between 0 and 1

  • mod – Linear normalization modulo \(2 \pi\)

  • log – Logarithmic normalization scale

  • symlog – Symmetric logarithmic normalization scale

  • None – Linear normalization between the min. and max. values of y

  The default is None.

• vmin (float, optional) – Minimum normalization value.

• vmax (float, optional) – Maximum normalization value.

• alpha (float between 0 and 1, optional) – Opacity of the image.

• extent (list or tuple of floats (left, right, bottom, top), optional) – Coordinates of the bounding box of the image.

Returns:

sm – Object mapping scalar data to RGBA color values.

Return type:

matplotlib.cm.ScalarMappable

plot_series(y, ntype=None, vmin=None, vmax=None, clabel=None)

Plot a time series of frames of an ODE solution.

Parameters:

• y (torch.tensor) – Solution of the system to plot. The input should be either reshaped to the dimensions of the simulation box or have a final dimension of 2.

• ntype (str) –

  Type of normalization to apply when displaying the image. The options are:

  • none – No normalization; point cloud data

  • sym – Symmetric linear normalization scale

  • unit – Linear normalization between 0 and 1

  • mod – Linear normalization modulo \(2 \pi\)

  • log – Logarithmic normalization scale

  • symlog – Symmetric logarithmic normalization scale

  • None – Linear normalization between the min. and max. values of y

  The default is None.

• vmin (float, optional) – Minimum normalization value.

• vmax (float, optional) – Maximum normalization value.

• clabel (str, optional) – Text used to label the colorbar.

solve(t, y0=None, device='cpu', rtol=1e-06, atol=1e-06)

Numerically integrate the ODE system and return the solution at times t.

Parameters:

• t (torch.tensor) – 1-dimensional tensor of evaluation times.

• y0 (torch.tensor of shape (M,...,D)) – Initial state of the system for M initial conditions. D denotes the flattened system size.

• device (str) – The name of the device on which the computation will be performed (e.g. cpu or cuda).

• rtol (float) – Upper bound on relative error. Default is \(10^{-6}\).

• atol (float) – Upper bound on absolute error. Default is \(10^{-6}\).

Returns:

y – Solution evaluated at T time points for M initial conditions. D denotes the flattened system size.

Return type:

torch.tensor of shape (T,M,...,D)

trim(t0=0)

Trim the solution at early time points to exclude initial transients.

Parameters:

t0 (int) – The first index into t at which to return the solution in order to exclude initial transients. The initial state will be set to the solution at this point and initial time set to 0 at this point.

class Swarm(args, method='dopri5', default_type=torch.float64)

Bases: ODEGraph

Class to define, solve, and visualize a system of ODEs governing agents with pairwise-interactions.

Parameters:

• args (dict) –

  Dictionary of parameters defining the ODE system:

  • N (int) – Dimension of the simulation box (N x N)

  • L (float) – Length of the real-space simulation box

  • D (int) – Dimensionality of system (2 or 3)

  • R (float) – Particle radius

  • J (float) – Spatial-phase coupling strength

  • K (float) – Phase-phase coupling strength

  • rc (float) – Cutoff radius of interaction

  • Nc (float) – Maximum number of interactions

• method (str) – Name of ODE solver to use. Default is dopri5.

• default_type (type) – Set the default torch.Tensor type.

Variables:

odeint (torchdiffeq.odeint or torchdiffeq.odeint_adjoint) – Numerical integrator for a system of ODEs given an initial value. The adjoint method will not be used if requires_grad is False.

forward(t, y)

Evaluate the ODE system at a specified time t and state y. Constructs the adjacency graphs based on the cutoff radius and current particle positions, then evaluates and propagates messages along constructed edges.

Parameters:

• t (torch.tensor) – 1-dimensional tensor of the evaluation time points.

• y (torch.tensor of shape (M,N,D)) – State of the system for M initial conditions.

Returns:

dy/dt – Derivative of the system.

Return type:

torch.tensor of shape (M,N,D)

init_state(M=1, seed=12)

Randomly generate the initial state(s) of the ODE system.

Parameters:

• M (int) – Number of initial conditions to generate. Default is 1.

• seed (int, optional) – Default seed used to set the state of a random number generator. Default is 12.

Variables:

y0 (torch.tensor of shape (M,1,D)) – Initial state of the system for M initial conditions. D denotes the flattened system size.

message(x_i, x_j)

Construct the message to node i from each neighboring node j.

Parameters:

• x_i (torch.tensor) – Node embedding for node i.

• x_j (torch.tensor) – Node embedding for node j.

Returns:

m_ij – Message between nodes i and j.

Return type:

torch.tensor

potential(x, r=1.0)

Smooth cutoff potential used to specify spatial interaction.

Parameters:

• x (torch.tensor) – Tensor of distances.

• r (float) – Cutoff radius. Default is 1.

Returns:

d – Cutoff potential at each x.

Return type:

torch.tensor