import numpy
import matplotlib
from matplotlib import pyplot as plot

matplotlib.rc('figure', figsize=(6, 5))
matplotlib.rc('text', usetex=True)
matplotlib.rc('text.latex', preamble=r'\usepackage{amsmath}')

# 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
f_xi = lambda xi: pow(xi, 4) / (1 + pow(xi, 4))
# smoothing function with scale parameter
f = lambda r, r_c, h: f_xi((r - r_c) / (h * r_c))

# smoothed Lennard-Jones potential
V_s = lambda r, r_c, h: f(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.0025, 0.005, 0.0075)

# plot potentials
ax = plot.axes()
ax.axhline(0, 0, 1, color="k", lw=0.5, ls="--")
ax.set_color_cycle(["r", "m", "g", "b"])
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 = %.3g$" % h)

l = ax.legend(loc="upper right")
l.legendPatch.set_alpha(0.7)

major_formatter = matplotlib.ticker.ScalarFormatter()
major_formatter.set_powerlimits((-1, 2))
ax.yaxis.set_major_formatter(major_formatter)

plot.setp(ax, xlim=(1.09, 1.15 + 1e-15))
plot.setp(ax, ylim=(-0.008, 0.04))

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