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])
# Lennard Jones force
F = lambda r: 48 * (pow(r, -14) - 0.5*pow(r, -8))
# truncated Lennard-Jones force
F_c = lambda r, r_c: numpy.piecewise(r, [r < r_c, r >= r_c], [F, 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))
# first derivative of smoothing function
Dg_xi = lambda xi: 4 * pow(xi, 3) / pow(1 + pow(xi, 4), 2)
# first derivative of smoothing function with scale parameter
Dg = lambda r, r_c, h: Dg_xi((r - r_c) / (h * r_c)) / (h * r_c)

# smoothed Lennard-Jones force
F_s = lambda r, r_c, h: F_c(r, r_c) * g(r, r_c, h) - V_c(r, r_c) * Dg(r, r_c, h)

# 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 forces
ax = plot.axes()
ax.axhline(0, 0, 1, color="k", lw=0.5, ls="--")
ax.plot(r, F_c(r, r_c), label=r"$h \rightarrow 0$")
for h in h:
    ax.plot(r, F_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.25, 2))

plot.xlabel(r"$r / \sigma$")
plot.ylabel(r"$|\vec{F}(\vec{r})| / \epsilon \sigma^{-1}$")
plot.show()