Neil G - 4 years ago 167
Python Question

# How do I use Theano to solve an ordinary differential equation?

Here is my Python code:

import numpy as np

from scipy.integrate import ode
if rng is None:
rng = np.random

size = n * dimensions
y0 = np.zeros((2 * size))
y0[:size] = rng.uniform(0, 0.1, size=size)
t0 = 0.0

mean = np.zeros(dimensions)
variance = np.eye(dimensions) * 0

def decode(y):
positions = np.remainder(y[:size].reshape((n, dimensions)), 1)
velocities = y[size:].reshape((n, dimensions))
return positions, velocities

def get_forces(positions, velocities):
delta_positions = positions[:, np.newaxis, :] - positions[np.newaxis, :, :]
delta_positions = np.remainder(delta_positions + 0.5, 1) - 0.5 # wrapping
distances = np.linalg.norm(delta_positions, axis=2, ord=2)
distances += 1e-5
pairwise_forces = delta_positions * (distances ** -3)[:, :, np.newaxis]
magnetic_forces = np.sum(pairwise_forces, axis=1) / (n ** 2)

velocity_mags = np.linalg.norm(velocities, axis=1, ord=2)
drag_forces = -velocities * velocity_mags[:, np.newaxis]

forces = magnetic_forces * 0.01 + drag_forces * 10.0
return forces

def f(t, y):
positions, velocities = decode(y)
forces = get_forces(positions, velocities)
retval = np.zeros((2 * size))
retval[:size] = velocities.reshape(size)
retval[size:] = forces.reshape(size)
return retval

r = ode(f).set_integrator('vode', method='bdf')
r.set_initial_value(y0, t0)
t_max = 20
dt = 1
while r.successful() and r.t < t_max:
positions, velocities = decode(r.y)
forces = get_forces(positions, velocities)
total_speed = np.sum(np.linalg.norm(velocities, axis=1, ord=2))
total_force = np.sum(np.linalg.norm(forces, axis=1, ord=2))
#print("\nt", r.t, "\np", positions, "\nv", velocities,
# "\nf", forces
print(total_speed, total_force)
if total_force < n * 1e-4:
print("breaking")
#break
r.integrate(r.t + dt)
print("converged after", r.t, total_speed, total_force)
return positions

Is it possible to use Theano to solve the ODE?

Here is a VERY simple ode solver in theano:

import numpy
import theano

# the right-hand side
def f(x, t):
return x*(1-x)

x = theano.tensor.matrix() # why not a matrix
dt = theano.tensor.scalar()
t = theano.tensor.scalar()

x_next = x + f(x, t)*dt # implement your favourite RK method here!

# matrix of random initial values
# store it on the device
x_shared = theano.shared(numpy.random.rand(10, 10))

step = theano.function([t, dt], [],
givens=[(x, x_shared)],
on_unused_input='warn')

t = 0.0
dt = 0.01

while t < 10:
step(t, dt)
t += dt
# test halt condition here

print(x_shared.get_value()) # read back the result

I hope it helps. Basically you have to implement your Runge-Kutta method.

And mind that the strength of theano is in vectorization, so I won't bother implementing the time loop in theano. That's why I used a simple python while loop, although I could use theano scan. Anyway, depending on your goal, the optimization can be tricky. I'm not 100% convinced that theano is a good choice for an ODE solver. Numpy does the vectorizations on your force and position matrices anyway, at least on CPU. With a theano implementation you can utilize GPU, but that's not a guarantee for a speedup.

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download