brown/force.py

from brown.interaction import UFuncWrapper
import numpy as np
import matplotlib.pyplot as plt

from coefficients import c

# This is the quite unreadable way to create 
# UFuncs with given parameters. FIXME: add this to another module.
force_function = UFuncWrapper(0, c)
potential_function = UFuncWrapper(2, c)

# Plot the force and potential.
r = np.arange(0, 100, 0.02, dtype=np.float16)
f, = plt.plot(r, force_function(r), label="force")
p, = plt.plot(r, potential_function(r), label="potential")
plt.legend(handles=[f, p])
plt.show()