<jupyter_start><jupyter_text>Table of Contents 1&nbsp;&nbsp;Solving an equation, numerically or with the Lambert W function1.1&nbsp;&nbsp;Loading packages and configuring plot sizes1.2&nbsp;&nbsp;Plotting the function first1.3&nbsp;&nbsp;Solving numerically?1.4&nbsp;&nbsp;How many solutions for a given a ?1.5&nbsp;&nbsp;Number of solutions as function of a1.6&nbsp;&nbsp;Plot of solution(s) as function of a1.7&nbsp;&nbsp;Solving formally with the Lambert W function1.8&nbsp;&nbsp;Asymptotic behaviors and approximations1.9&nbsp;&nbsp;Conclusion  Solving an equation, numerically or with the Lambert W functionI 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<jupyter_code>%load_ext watermark
%watermark -a "<NAME> (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)<jupyter_output><empty_output><jupyter_text>Plotting the function first<jupyter_code>def objective(x, a):
    return np.exp(- a * x**2) - x<jupyter_output><empty_output><jupyter_text>First, let's have a look to its plot for some values of $a$:<jupyter_code>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()<jupyter_output><empty_output><jupyter_text>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:<jupyter_code>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()<jupyter_output><empty_output><jupyter_text>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:<jupyter_code>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()<jupyter_output><empty_output><jupyter_text>Solving numerically?We can start to try to use [`scipy.optimize.root`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.htmlscipy.optimize.root) to numerically solve this equation.<jupyter_code>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}).")<jupyter_output><empty_output><jupyter_text>Let's check that there is no solution for $a < 0$ too small.<jupyter_code>one_solution(-1, verb=True)<jupyter_output>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])<jupyter_text>It can find a solution, but only one (depending on the starting point $x_0$) and not both:<jupyter_code>one_solution(-0.1, x0=0, verb=True)
one_solution(-0.1, x0=10, verb=True)<jupyter_output>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])<jupyter_text>For $a > 0$, the equation seems to have a unique solution:<jupyter_code>one_solution(1, x0=0)
one_solution(1, x0=-100)
one_solution(1, x0=100)<jupyter_output><empty_output><jupyter_text>We can just hack and try different values for $x_0$, expecting to find all the roots.<jupyter_code>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)<jupyter_output><empty_output><jupyter_text>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:<jupyter_code>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)<jupyter_output>/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp<jupyter_text>I think having $3$ (or more) solutions is a numerical error.<jupyter_code>amin = -100
amax = 100
gap_points = thresholds(amin=amin, amax=amax, delta=0.1)
gap_points<jupyter_output>/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp<jupyter_text>As we will see below, even having two solutions is nothing but a numerical error.  Number of solutions as function of aWe can plot the (estimated) number of solution as a function of $a$, to start wit, thanks to the [`matplotlib.pyplot.hlines`](https://matplotlib.org/api/pyplot_api.htmlmatplotlib.pyplot.hlines) function:<jupyter_code>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)<jupyter_output>{0, 1, 2}
{0: [-0.20000000000567297], 2: [-28.20000000000408, -0.10000000000567866]}
{2: [-0.10000000000567866], 0: [-28.100000000004087], 1: [-5.6843418860808015e-12]}<jupyter_text>Plot of solution(s) as function of aNow we can try to use this to plot the solution(s) as function of $a$.<jupyter_code>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()<jupyter_output>No handles with labels found to put in legend.<jupyter_text>This shows the numerical solution to the equation, and we will check below that the formal solution coincides.  Solving formally with the Lambert W functionLuckily, 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](https://en.wikipedia.org/wiki/Lambert_W_function), or [this article](https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf).As for (almost) all the special function, we don't need to write it ourself: it is in scipy! [`scipy.special.lambertw`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.lambertw.htmlscipy.special.lambertw)<jupyter_code>from scipy.special import lambertw<jupyter_output><empty_output><jupyter_text>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 \Leftrightarrow2 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) \Leftrightarrowu = W(2a) \Leftrightarrowy = \frac{W(2a)}{2a} \Leftrightarrowx(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):<jupyter_code>def formal_solution(a):
    return np.sqrt(lambertw(2 * a) / (2 * a))<jupyter_output><empty_output><jupyter_text>We can check some values:<jupyter_code>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}")<jupyter_output>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<jupyter_text>Asymptotic behaviors and approximationsWe 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](https://en.wikipedia.org/wiki/Lambert_W_functionAsymptotic_expansions):  $$ \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)}}} $$<jupyter_code>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))<jupyter_output><empty_output><jupyter_text>We can plot all this.<jupyter_code>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()<jupyter_output>/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)
{'fraction_non_alphanumeric': 0.17029499921123206, 'fraction_numerical': 0.04251459220697271, 'mean_word_length': 4.083801122694466, 'pattern_counts': {'":': 0, '<': 74, '<?xml version=': 0, '>': 76, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 5, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '19', 'dir': 'jupyter-structured-clean-dedup', 'id': '458219', 'license': 'MIT', 'n_tokens_mistral': 5353, 'n_tokens_neox': 4942, 'n_words': 1285, 'path': 'Solving_an_equation_and_the_Lambert_W_function.ipynb', 'repo_name': 'doc22940/notebooks-2', 'size': '12685'}