import numpy
import matplotlib
from matplotlib import pyplot as plot

matplotlib.rc('figure', figsize=(6, 5))
matplotlib.rc('text', usetex=True)
matplotlib.rc('font', size=12)

# Lennard Jones potential
V = lambda r: 4 * (pow(r, -12) - pow(r, -6))
# truncated Lennard-Jones potential
V_c = lambda r, r_c: numpy.piecewise(r, [r < r_c, r >= r_c], [lambda r: V(r) - V(r_c), 0])

# smoothing function
g_xi = lambda xi: pow(xi, 4) / (1 + pow(xi, 4))
# smoothing function with scale parameter
g = lambda r, r_c, h: g_xi((r - r_c) / (h * r_c))

# smoothed Lennard-Jones potential
V_s = lambda r, r_c, h: g(r, r_c, h) * V_c(r, r_c)

# particle distance
r = numpy.linspace(1, 1.4, 1000)
# cutoff distance
r_c = pow(2, 1./6)
# smoothing function scale parameters
h = (0.002, 0.005, 0.010)

# plot potentials
ax = plot.axes()
ax.axhline(0, 0, 1, color="k", lw=0.5, ls="--")
ax.plot(r, V_c(r, r_c), label=r"$h \rightarrow 0$")
for h in h:
    ax.plot(r, V_s(r, r_c, h), label=r"$h = {0:.3f}$".format(h))

ax.legend(loc="upper right")

major_formatter = matplotlib.ticker.ScalarFormatter()
ax.yaxis.set_major_formatter(major_formatter)

plot.setp(ax, xlim=(1.08, 1.14 + 1e-15))
plot.setp(ax, ylim=(-0.00499, 0.03 + 1e-15))

plot.xlabel(r"$r / \sigma$")
plot.ylabel(r"$U(r) / \epsilon$")
plot.show()