Table of Contents¶
Solving an equation, numerically or with the Lambert W function¶
I want to solve the equation $\exp(-ax^2)=x$ and find its solution(s) as a function of $a\in\mathbb{R}$.
Loading packages and configuring plot sizes¶
%load_ext watermark
%watermark -a "Lilian Besson (Naereen)" -i -v -p numpy,matplotlib,scipy,seaborn
import numpy as np
from scipy import optimize as opt
import matplotlib as mpl
mpl.rcParams['figure.figsize'] = (15, 8)
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(context="notebook", style="darkgrid", palette="hls", font="sans-serif", font_scale=1.8)
Plotting the function first¶
def objective(x, a):
return np.exp(- a * x**2) - x
First, let's have a look to its plot for some values of $a$:
X = np.linspace(-2, 2, 2000)
for a in [0, -0.1, 0.1, -1, 1]:
plt.plot(X, objective(X, a), 'o-', label=f"$a={a:.3g}$", markevery=50)
plt.legend()
plt.xlabel("$x$"); plt.ylabel("$y$")
plt.title(r"Function $\exp(- a x^2) - x$ for different $a$")
plt.show()
We can see that a solution to $\exp(-a x^2) = x$ has to be positive, as $\exp(-a x^2) > 0$ for any $x,a$. We also check that if $a < 0$, $\exp(-a x^2) - x$ seems to always be positive, but if $a \geq 0$, it seems to have a unique root. Let's zoom a little bit:
X = np.linspace(0, 1.5, 2000)
for a in [0, -0.1, 0.1, -1, 1]:
plt.plot(X, objective(X, a), 'o-', label=f"$a={a:.3g}$", markevery=50)
plt.legend()
plt.xlabel("$x$"); plt.ylabel("$y$")
plt.title(r"Function $\exp(- a x^2) - x$ for different $a$")
plt.show()
The curve for $a=-0.1$ seems to stay negative, but that's not possible as for $a<0$ and $x\to\infty$, $\exp(-a x^2)$ dominates over $-x$. We can check that it will have a second root:
X = np.linspace(0, 5, 2000)
for a in [0, -0.1, 0.1, 1]:
plt.plot(X, objective(X, a), 'o-', label=f"$a={a:.3g}$", markevery=50)
plt.legend()
plt.xlabel("$x$"); plt.ylabel("$y$")
plt.title(r"Function $\exp(- a x^2) - x$ for different $a$")
plt.show()
Solving numerically?¶
We can start to try to use scipy.optimize.root to numerically solve this equation.
def one_solution(a, x0=0, verb=False):
sol = opt.root(objective, x0, args=(a,))
if verb: print(sol)
if sol.success:
return sol.x
else:
raise ValueError(f"No solution was found for a = {a:.3g} (and starting at x0 = {x0:.3g}).")
Let's check that there is no solution for $a < 0$ too small.
one_solution(-1, verb=True)
It can find a solution, but only one (depending on the starting point $x_0$) and not both:
one_solution(-0.1, x0=0, verb=True)
one_solution(-0.1, x0=10, verb=True)
For $a > 0$, the equation seems to have a unique solution:
one_solution(1, x0=0)
one_solution(1, x0=-100)
one_solution(1, x0=100)
We can just hack and try different values for $x_0$, expecting to find all the roots.
def solutions(a, x0s=None, tol=1e-10, verb=False):
nbdigits = int(np.log10(1. / tol))
sols = set()
if x0s is None:
x0s = [-10, -5, -2, -1, 0, 1, 2, 5, 10]
for x0 in x0s:
sol = opt.root(objective, x0, args=(a,))
if sol.success:
approx = np.round(float(sol.x), nbdigits)
sols.add(approx)
if verb and len(sols) == 0:
print(f"No solution was found for a = {a:.3g} (and starting at x0 = {x0:.3g}).")
return sols
solutions(-10)
solutions(-0.1)
solutions(0)
solutions(1)
solutions(2)
How many solutions for a given a ?¶
We can use this to try to find the threshold value for $a$ from $0$ to $2$ and from $2$ to $1$ solution:
def thresholds(amin=-10, amax=10, delta=0.01):
gap_points = dict()
prev_a = amin
prev_nb_sol = len(solutions(prev_a))
for a in np.arange(amin, amax, delta):
nb_sol = len(solutions(a))
if nb_sol != prev_nb_sol:
gap_points[(prev_nb_sol, nb_sol)] = (prev_a, a)
prev_nb_sol = nb_sol
prev_a = a
return gap_points
thresholds(amin=-10, amax=10, delta=0.01)
thresholds(amin=-8, amax=1, delta=0.01)
I think having $3$ (or more) solutions is a numerical error.
amin = -100
amax = 100
gap_points = thresholds(amin=amin, amax=amax, delta=0.1)
gap_points
As we will see below, even having two solutions is nothing but a numerical error.
Number of solutions as function of a¶
We can plot the (estimated) number of solution as a function of $a$, to start wit, thanks to the matplotlib.pyplot.hlines function:
def plot_gap_points(gap_points, amin, amax):
ys = set()
for ym, yM in gap_points.keys():
ys.add(ym)
ys.add(yM)
print(ys)
xleft = dict()
xright = dict()
for (ym, yM), (xm, xM) in gap_points.items():
xleft[ym] = xleft.get(ym, []) + [xm]
xright[yM] = xright.get(yM, []) + [xM]
for ym, yM in gap_points.keys():
xleft[ym].sort()
xright[yM].sort()
print(xleft)
print(xright)
min_xleft = min(sum(list(xleft.values()), []))
max_xright = min(sum(list(xright.values()), []))
plt.figure()
for y in ys:
if y not in xleft and y in xright:
for x in xright[y]:
plt.hlines(y, x, amax)
if y in xleft and min_xleft in xleft[y]:
plt.hlines(y, amin, min_xleft)
del xleft[y][0]
#if y in xright and max_xright in xright[y]:
# plt.hlines(y, max_xright, amax)
# del xright[y][-1]
if y in xleft and y in xright:
for xmin, xmax in zip(xleft[y], xright[y]):
plt.hlines(y, xmin, xmax)
plt.xlabel("Value of $a$")
plt.ylabel("Number of solution")
plt.title(r"Number of solutions to $\exp(- a x^2) = x$, as function of $a$")
return ys, xleft, min_xleft, xright, max_xright
ys, xleft, min_xleft, xright, max_xright = plot_gap_points(gap_points, amin, amax)
Plot of solution(s) as function of a¶
Now we can try to use this to plot the solution(s) as function of $a$.
def plot_multivalued_function(X, f, maxnboutput=1, **kwargs):
Y = np.zeros((maxnboutput, len(X)))
Y.fill(np.nan)
for i, x in enumerate(X):
ys = sorted(list(f(x)))
for j, y in enumerate(ys):
Y[j, i] = y
for j in range(maxnboutput):
plt.plot(X, Y[j], 'o-', **kwargs)
A = np.linspace(-100, 100, 1000)
plot_multivalued_function(A, solutions, maxnboutput=2, markevery=10)
plt.legend()
plt.xlabel("Parameter $a$"); plt.ylabel("Solution(s)")
plt.title(r"Solution(s) to $\exp(- a x^2) = x$, as function of $a$")
plt.show()
A = np.linspace(0, 20, 2000)
plot_multivalued_function(A, solutions, maxnboutput=2, markevery=20)
plt.legend()
plt.xlabel("Parameter $a$"); plt.ylabel("Solution(s)")
plt.title(r"Solution(s) to $\exp(- a x^2) = x$, as function of $a$")
plt.show()
This shows the numerical solution to the equation, and we will check below that the formal solution coincides.
Solving formally with the Lambert W function¶
Luckily, we can transform this equation to solve it with the Lambert $W$ function, defined as $W(x) = z \Leftrightarrow x = z \mathrm{e}^{z}$. For more details, please see this page, or this article.
As for (almost) all the special function, we don't need to write it ourself: it is in scipy! scipy.special.lambertw
from scipy.special import lambertw
As the only possible solution are $x>0$ $$ \exp(-a x^2) = x \Leftrightarrow \left(\exp(-a x^2)\right)^2= \exp(-2 a x^2) = x^2 \Leftrightarrow 2 a y \exp(2 a x) = 2 a \;\;(\text{with}\;\; y := x^2) \Leftrightarrow \\ u \exp(u) = 2 a \;\;(\text{with}\;\; u := 2 a y) \Leftrightarrow u = W(2a) \Leftrightarrow y = \frac{W(2a)}{2a} \Leftrightarrow x(a) := \sqrt{\frac{W(2a)}{2a}}. $$
And so it is quite easy to compute, for $a > 0$ (the behavior at $0$ is undefined without a more careful study):
def formal_solution(a):
return np.sqrt(lambertw(2 * a) / (2 * a))
We can check some values:
for a in [0.5, 1, 2, 3, 4]:
xa = formal_solution(a)
assert np.isclose(exp(-a * xa**2), xa)
print(f"a = {a:.3g} gives x(a) = {float(xa):.3g}")
Asymptotic behaviors and approximations¶
We can try to approximate the solution for small $a$ or large $a$:
For small $a$, $W(2a) \simeq 2a - 4a^2$ so $x(a) \simeq 1 - a$.
For large $a$, we have this bound: $$ \forall x \geq \mathrm{e},{\displaystyle \ln(x)-\ln {\bigl (}\ln(x){\bigr )}+{\frac {\ln {\bigl (}\ln(x){\bigr )}}{2\ln(x)}}\leq W(x)\leq \ln(x)-\ln {\bigl (}\ln(x){\bigr )}+{\frac {e}{e-1}}{\frac {\ln {\bigl (}\ln(x){\bigr )}}{\ln(x)}}} $$
e = np.exp(1)
def upper_bound(a):
up_b = np.log(2*a) - np.log(np.log(2*a)) + (e / (e - 1)) * (np.log(np.log(2*a)) / np.log(2*a))
return np.sqrt(up_b / (2*a))
def lower_bound(a):
lo_b = np.log(2*a) - np.log(np.log(2*a)) + np.log(np.log(2*a)) / (2 * np.log(2*a))
return np.sqrt(lo_b / (2*a))
We can plot all this.
A = np.linspace(0, 20, 4000)
A1 = A[A <= 0.5]
A2 = A[A >= 1]
Ae = A[A >= e]
plt.plot(A, formal_solution(A), label="Solution", markevery=20)
plt.plot(A1, 1 - A1, 'b--', label=r"Tangent at $0$: $x(a) \simeq 1 - a$", markevery=20)
#plt.plot(A2, np.sqrt((np.log(2*A2) - np.log(np.log(2*A2)))/(2*A2)), 'g--', label=r"Asymptote at $+\infty$", markevery=20)
plt.plot(Ae, lower_bound(Ae), 'g--', label=r"Lower-bound for $a \geq e$", markevery=20)
plt.plot(Ae, upper_bound(Ae), 'c--', label=r"Upper-bound for $a \geq e$", markevery=20)
plt.legend()
plt.xlabel("Parameter $a$"); plt.ylabel(r"Solution $x(a) = \sqrt{\frac{W(2a)}{2a}}$")
plt.title(r"Solution to $\exp(- a x^2) = x$, as function of $a$")
plt.show()
And voilà.