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¶
In [1]:
%load_ext watermark
%watermark -a "Lilian Besson (Naereen)" -i -v -p numpy,matplotlib,scipy,seaborn
Lilian Besson (Naereen) 2018-01-31T14:56:14+01:00 CPython 3.6.3 IPython 6.2.1 numpy 1.14.0 matplotlib 2.1.2 scipy 1.0.0 seaborn 0.8.1
In [2]:
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¶
In [3]:
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$:
In [4]:
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()
Out[4]:
[<matplotlib.lines.Line2D at 0x7f1fd8042240>]
Out[4]:
[<matplotlib.lines.Line2D at 0x7f1fa146bbe0>]
Out[4]:
[<matplotlib.lines.Line2D at 0x7f1fa146bcf8>]
Out[4]:
[<matplotlib.lines.Line2D at 0x7f1fa14051d0>]
Out[4]:
[<matplotlib.lines.Line2D at 0x7f1fa14059b0>]
Out[4]:
<matplotlib.legend.Legend at 0x7f1fa1405e48>
Out[4]:
Text(0.5,0,'$x$')
Out[4]:
Text(0,0.5,'$y$')
Out[4]:
Text(0.5,1,'Function $\\exp(- a x^2) - x$ for different $a$')
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:
In [5]:
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()
Out[5]:
[<matplotlib.lines.Line2D at 0x7f1fa10bc278>]
Out[5]:
[<matplotlib.lines.Line2D at 0x7f1fa10bc828>]
Out[5]:
[<matplotlib.lines.Line2D at 0x7f1fa10bc9e8>]
Out[5]:
[<matplotlib.lines.Line2D at 0x7f1fa10bce80>]
Out[5]:
[<matplotlib.lines.Line2D at 0x7f1fa109b668>]
Out[5]:
<matplotlib.legend.Legend at 0x7f1fa1156f28>
Out[5]:
Text(0.5,0,'$x$')
Out[5]:
Text(0,0.5,'$y$')
Out[5]:
Text(0.5,1,'Function $\\exp(- a x^2) - x$ for different $a$')
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:
In [6]:
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()
Out[6]:
[<matplotlib.lines.Line2D at 0x7f1fa0ba26d8>]
Out[6]:
[<matplotlib.lines.Line2D at 0x7f1fa0ba2c88>]
Out[6]:
[<matplotlib.lines.Line2D at 0x7f1fa0ba2e48>]
Out[6]:
[<matplotlib.lines.Line2D at 0x7f1fa0b65668>]
Out[6]:
<matplotlib.legend.Legend at 0x7f1fa1053cf8>
Out[6]:
Text(0.5,0,'$x$')
Out[6]:
Text(0,0.5,'$y$')
Out[6]:
Text(0.5,1,'Function $\\exp(- a x^2) - x$ for different $a$')
Solving numerically?¶
We can start to try to use scipy.optimize.root to numerically solve this equation.
In [7]:
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.
In [8]:
one_solution(-1, verb=True)
fjac: array([[-1.]])
fun: array([0.77291476])
message: 'The iteration is not making good progress, as measured by the \n improvement from the last ten iterations.'
nfev: 28
qtf: array([-0.77291498])
r: array([-0.00495227])
status: 5
success: False
x: array([0.41894532])
--------------------------------------------------------------------------- ValueError Traceback (most recent call last) <ipython-input-8-4cc3b5bba819> in <module>() ----> 1 one_solution(-1, verb=True) <ipython-input-7-c8cc0197ed60> in one_solution(a, x0, verb) 5 return sol.x 6 else: ----> 7 raise ValueError(f"No solution was found for a = {a:.3g} (and starting at x0 = {x0:.3g}).") ValueError: No solution was found for a = -1 (and starting at x0 = 0).
It can find a solution, but only one (depending on the starting point $x_0$) and not both:
In [9]:
one_solution(-0.1, x0=0, verb=True)
fjac: array([[-1.]])
fun: array([2.22044605e-16])
message: 'The solution converged.'
nfev: 8
qtf: array([-1.57868385e-10])
r: array([0.7408292])
status: 1
success: True
x: array([1.13835649])
Out[9]:
array([1.13835649])
In [10]:
one_solution(-0.1, x0=10, verb=True)
fjac: array([[-1.]])
fun: array([0.])
message: 'The solution converged.'
nfev: 23
qtf: array([-2.18643326e-10])
r: array([-1.54264248])
status: 1
success: True
x: array([3.56555841])
Out[10]:
array([3.56555841])
For $a > 0$, the equation seems to have a unique solution:
In [11]:
one_solution(1, x0=0)
one_solution(1, x0=-100)
one_solution(1, x0=100)
Out[11]:
array([0.65291864])
Out[11]:
array([0.65291864])
Out[11]:
array([0.65291864])
We can just hack and try different values for $x_0$, expecting to find all the roots.
In [12]:
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
In [13]:
solutions(-10)
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp
Out[13]:
set()
In [14]:
solutions(-0.1)
Out[14]:
{1.1383564947, 3.5655584119}
In [15]:
solutions(0)
Out[15]:
{1.0}
In [16]:
solutions(1)
Out[16]:
{0.6529186404}
In [17]:
solutions(2)
Out[17]:
{0.5482170814}
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:
In [18]:
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
In [21]:
thresholds(amin=-10, amax=10, delta=0.01)
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp
Out[21]:
{(0, 1): (-0.3700000000002053, -0.3600000000002055),
(0, 2): (-0.1900000000002091, -0.18000000000020933),
(1, 0): (-0.3400000000002059, -0.33000000000020613),
(2, 0): (-7.030000000000063, -7.0200000000000635),
(2, 1): (-0.020000000000212736, -0.01000000000021295),
(2, 3): (-0.0500000000002121, -0.04000000000021231),
(3, 2): (-0.04000000000021231, -0.030000000000212523)}
In [22]:
thresholds(amin=-8, amax=1, delta=0.01)
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp
Out[22]:
{(0, 1): (-0.37000000000016264, -0.36000000000016286),
(0, 2): (-0.19000000000016648, -0.1800000000001667),
(1, 0): (-0.3400000000001633, -0.3300000000001635),
(2, 0): (-7.030000000000021, -7.020000000000021),
(2, 1): (-0.020000000000170104, -0.010000000000170317),
(2, 3): (-0.050000000000169464, -0.04000000000016968),
(3, 2): (-0.04000000000016968, -0.03000000000016989)}
I think having $3$ (or more) solutions is a numerical error.
In [23]:
amin = -100
amax = 100
gap_points = thresholds(amin=amin, amax=amax, delta=0.1)
gap_points
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp
Out[23]:
{(0, 2): (-0.20000000000567297, -0.10000000000567866),
(2, 0): (-28.20000000000408, -28.100000000004087),
(2, 1): (-0.10000000000567866, -5.6843418860808015e-12)}
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:
In [24]:
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)
{0, 1, 2}
{0: [-0.20000000000567297], 2: [-28.20000000000408, -0.10000000000567866]}
{2: [-0.10000000000567866], 0: [-28.100000000004087], 1: [-5.6843418860808015e-12]}
Plot of solution(s) as function of a¶
Now we can try to use this to plot the solution(s) as function of $a$.
In [25]:
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)
In [26]:
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()
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp No handles with labels found to put in legend.
Out[26]:
<matplotlib.legend.Legend at 0x7f1fa06891d0>
Out[26]:
Text(0.5,0,'Parameter $a$')
Out[26]:
Text(0,0.5,'Solution(s)')
Out[26]:
Text(0.5,1,'Solution(s) to $\\exp(- a x^2) = x$, as function of $a$')
In [27]:
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()
No handles with labels found to put in legend.
Out[27]:
<matplotlib.legend.Legend at 0x7f1fa13a1e10>
Out[27]:
Text(0.5,0,'Parameter $a$')
Out[27]:
Text(0,0.5,'Solution(s)')
Out[27]:
Text(0.5,1,'Solution(s) to $\\exp(- a x^2) = x$, as function of $a$')
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
In [28]:
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):
In [30]:
def formal_solution(a):
return np.sqrt(lambertw(2 * a) / (2 * a))
We can check some values:
In [35]:
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}")
a = 0.5 gives x(a) = 0.753 a = 1 gives x(a) = 0.653 a = 2 gives x(a) = 0.548 a = 3 gives x(a) = 0.489 a = 4 gives x(a) = 0.448
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:3: ComplexWarning: Casting complex values to real discards the imaginary part This is separate from the ipykernel package so we can avoid doing imports until /usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:4: ComplexWarning: Casting complex values to real discards the imaginary part after removing the cwd from sys.path.
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)}}} $$
In [36]:
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.
In [39]:
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()
/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: invalid value encountered in true_divide /usr/local/lib/python3.6/dist-packages/numpy/core/numeric.py:492: ComplexWarning: Casting complex values to real discards the imaginary part return array(a, dtype, copy=False, order=order)
Out[39]:
[<matplotlib.lines.Line2D at 0x7f1fa1053ac8>]
Out[39]:
[<matplotlib.lines.Line2D at 0x7f1fa0bb9e80>]
Out[39]:
[<matplotlib.lines.Line2D at 0x7f1fa10a5c50>]
Out[39]:
[<matplotlib.lines.Line2D at 0x7f1fa0695128>]
Out[39]:
<matplotlib.legend.Legend at 0x7f1fa0695748>
Out[39]:
Text(0.5,0,'Parameter $a$')
Out[39]:
Text(0,0.5,'Solution $x(a) = \\sqrt{\\frac{W(2a)}{2a}}$')
Out[39]:
Text(0.5,1,'Solution to $\\exp(- a x^2) = x$, as function of $a$')
And voilà.