Lambert W Function Source

import "dart:math";

//
// Approximation for both branches of the real valued Lambert W function.
//
// See: https://en.wikipedia.org/wiki/Lambert_W_function
//
// These approximations were found in "Having Fun with Lambert W(x) Function".
//
// See: https://arxiv.org/abs/1003.1628
//
// Lambert's W function is needed to solve certain exponential equations.  It
// is defined so that W(z * e^z) = z. So if you can put an equation into the
// form inside W on one side you can apply W to both sides and solve for z.
//
// But it's never that simple, is it, the W function is not really a function,
// because there are many possible values of z that satisfy the equation
// z' = z * e^z.  But if we restrict z to being a real number, then there are
// only ever zero, one, or two solutions.  The functions W_negative_one and
// W_zero return the two possible solutions.
//

//
// Evaluate a polynomial in [x] using the coefficients given in [a].
//
double _polynomial(Iterable<double> a, double x)
{
    double result = 0.0;

    for (int i = a.length - 1; i >= 0; --i)
        result = a.elementAt(i) + x * result;

    return result;
}

//
// Evaluate two polynomials and return their ratio.
//
double _rational(Iterable<double> a, Iterable<double> b, double x)
{
    return _polynomial(a, x) / _polynomial(b, x);
}

//
// One step of Fritsch's iteration
//
// This refines an initial estimate of W by enough that only one iteration
// is required to go from the initial approximations below to a final value
// of full double precision.
//
double _fritsch_iteration(double x, double W)
{
    final double z = log(x / W) - W;
    final double q = 2.0 * (1.0 + W) * (1.0 + W + 2.0 * z / 3.0);

    if (z == 0.0 && q == 0.0)
        return W;

    final double e = (z / (1 + W)) * ((q - z) / (q - 2.0 * z));

    return W * (1 + e);
}

//
// A simple recursion that comes from the definition of W.
//
double _W_negative_one_recursive(double x, int n)
{
    //
    // I'm not sure how to correctly terminate the recursion.  Just ignoring
    // the denominator (treating it as 1) works.
    //
    if (n <= 0)
        return log(-x);
    else
        return log(-x / -_W_negative_one_recursive(x, n - 1));
}

//
// Polynomial expansion of the function around the point [-1/e, -1], this
// can be used to evaluate either W_negative_one or W_zero, based on the sign
// of p, since the expansion is symetric.  The value p is generated by
// transforming x (the input value to W(x)) so that p is rescaled and centered
// on the branch point.
//
double _W_branch_point(double p)
{
    return _polynomial(const <double>[-1.0,
                                      1.0,
                                      -1.0/3,
                                      11.0/72,
                                      -43.0/540,
                                      769.0/17280,
                                      -221.0/8505,
                                      680863.0/43545600,
                                      -1963.0/204120,
                                      226287557.0/37623398400],
                       p);
}

double _W_asymptotic(double a, double b)
{
    final double A0 =            a;
    final double A1 = A0 * 2.0 * a;
    final double A2 = A1 * 3.0 * a;
    final double A3 = A2 * 2.0 * a;
    final double A4 = A3 * 5.0 * a;

    const List<double> B1 = const <double>[ -2.0,    1.0];
    const List<double> B2 = const <double>[  6.0,   -9.0,   2.0];
    const List<double> B3 = const <double>[-12.0,   36.0, -22.0,    3.0];
    const List<double> B4 = const <double>[ 60.0, -300.0, 350.0, -125.0, 12.0];

    return (a -
            b +
            b / a +
            (b * _polynomial(B1, b) / A1) +
            (b * _polynomial(B2, b) / A2) +
            (b * _polynomial(B3, b) / A3) +
            (b * _polynomial(B4, b) / A4));
}

//
// Apply the various approximations of W_negative_one in the ranges that they
// are most accurate.
//
double _W_negative_one_approximate(double x)
{
    if (x < -1.0/E)
        return double.NAN;

    if (x < -0.30298541769)
        return _W_branch_point(-sqrt(2.0 * (1.0 + E * x)));

    if (x < -0.051012917658221676)
        return _rational(const <double>[-7.81417672390744,
                                        253.88810188892484,
                                        657.9493176902304],
                         const <double>[1.0,
                                        -60.43958713690808,
                                        99.9856708310761,
                                        682.6073999909428,
                                        962.1784396969866,
                                        1477.9341280760887],
                         x);

    if (x < 0)
        return _W_negative_one_recursive(x, 9);

    return double.NAN;
}

//
// Apply the various approximations of W_zero in the ranges that they are most
// accurate.
//
double _W_zero_approximate(double x)
{
    if (x < -1.0/E)
        return double.NAN;

    if (x < -0.32358170806015724)
        return _W_branch_point(sqrt(2.0 * (1.0 + E * x)));

    if (x < 0.14546954290661823)
        return x * _rational(const <double>[1.0,
                                            5.931375839364438,
                                            11.39220550532913,
                                            7.33888339911111,
                                            0.653449016991959],
                             const <double>[1.0,
                                            6.931373689597704,
                                            16.82349461388016,
                                            16.43072324143226,
                                            5.115235195211697],
                             x);

    if (x < 8.706658967856612)
        return x * _rational(const <double>[1.0,
                                            2.445053070726557,
                                            1.343664225958226,
                                            0.148440055397592,
                                            0.0008047501729130],
                             const <double>[1.0,
                                            3.444708986486002,
                                            3.292489857371952,
                                            0.916460018803122,
                                            0.0530686404483322],
                             x);

    return _W_asymptotic(log(x), log(log(x)));
}

double W_negative_one(double x)
{
    return _fritsch_iteration(x, _W_negative_one_approximate(x));
}

double W_zero(double x)
{
    return _fritsch_iteration(x, _W_zero_approximate(x));
}