brown/force.py

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

from coefficients import c

force_function = UFuncWrapper(0, c)
potential_function = UFuncWrapper(2, c)

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()