The Lambert W function compiles, and W_0 left of -2/e is on its branch

BREAKING-CHANGES.md

@@ -684,7 +684,8 @@ and below `-1/e` there is none ([#2182](https://github.com/asc-community/Angouri
 al. (1996), as Mathematica's `ProductLog` and mpmath's `lambertw` have them: `W_0` is real on `[-1/e, +oo)` and `W_(-1)`
 on `[-1/e, 0)`. It is exact at `0` and where `z` is `c e^c` for a real `c` the branch takes, so `lambertw(e)` is `1`
 and `lambertw(-1/e)` is `-1`, and it is evaluated numerically elsewhere. Its derivative is `e^(-W(z))/(1 + W(z))`
-([#2182](https://github.com/asc-community/AngouriMath/issues/2182)).
+([#2182](https://github.com/asc-community/AngouriMath/issues/2182)). It compiles, in double precision on the same
+branches, in both compilers ([#2195](https://github.com/asc-community/AngouriMath/issues/2195)).
 
 | Input | Was (2.6.0) | Now |
 |---|---|---|

Sources/.editorconfig

@@ -363,6 +363,12 @@ file_header_template=\nCopyright (c) 2019-2026 Angouri.\nAngouriMath is licensed
 [AngouriMath/Numerics/CombinatorialFunctions.cs]
 file_header_template=\nCopyright (c) 2019-2026 Angouri.\nAngouriMath is licensed under MIT.\nDetails: https://github.com/asc-community/AngouriMath/blob/master/LICENSE.md.\nWebsite: https://am.angouri.org.\n
 
+[AngouriMath/Numerics/LambertWFunction.cs]
+file_header_template=\nCopyright (c) 2019-2026 Angouri.\nAngouriMath is licensed under MIT.\nDetails: https://github.com/asc-community/AngouriMath/blob/master/LICENSE.md.\nWebsite: https://am.angouri.org.\n
+
+[Tests/UnitTests/Common/LambertWCompilationTest.cs]
+file_header_template=\nCopyright (c) 2019-2026 Angouri.\nAngouriMath is licensed under MIT.\nDetails: https://github.com/asc-community/AngouriMath/blob/master/LICENSE.md.\nWebsite: https://am.angouri.org.\n
+
 [Tests/UnitTests/Common/CompileMatrixTest.cs]
 file_header_template=\nCopyright (c) 2019-2026 Angouri.\nAngouriMath is licensed under MIT.\nDetails: https://github.com/asc-community/AngouriMath/blob/master/LICENSE.md.\nWebsite: https://am.angouri.org.\n
 

Sources/AngouriMath/Core/Entity/Continuous/Number/LambertW.cs

@@ -29,10 +29,11 @@ partial record Number
             /// on <c>[-1/e, 0)</c>, and a value on a branch cut is the one approached from above.
             /// </para>
             /// <para>
-            /// The branch is found in double precision, where it is cheap: from the series in
-            /// <c>p = sqrt(2 (e z + 1))</c> about the branch point <c>-1/e</c>, from <c>ln(1 + z)</c> elsewhere on
-            /// the principal branch short of <c>|z| = 3</c>, and from <c>ln z + 2 pi i k - ln(ln z + 2 pi i k)</c> otherwise,
-            /// each then taken to a root by Halley's iteration. That root is refined by the same iteration
+            /// The branch is found in double precision, where it is cheap, by the kernel the compilers use,
+            /// <see cref="AngouriMath.Numerics.LambertWFunction"/>: from the series in <c>p = sqrt(2 (e z + 1))</c>
+            /// about the branch point <c>-1/e</c>, from <c>ln(1 + z)</c> elsewhere on the principal branch short of
+            /// <c>|z| = 3</c>, and from <c>ln z + 2 pi i k - ln(ln z + 2 pi i k)</c> otherwise, each then taken to a
+            /// root by Halley's iteration. That root is refined by the same iteration
             /// at the working precision, which triples the digits with each step and does not leave the
             /// branch. Near the branch point, where double precision cannot resolve <c>z + 1/e</c>, the
             /// series is evaluated at the working precision instead, and refined from there.
@@ -78,7 +79,7 @@ partial record Number
                 }
                 else
                 {
-                    var start = InDoublePrecision(k, new NumericsComplex(re.ToDouble(), im.ToDouble()));
+                    var start = AngouriMath.Numerics.LambertWFunction.LambertW(new NumericsComplex(k, 0), new NumericsComplex(re.ToDouble(), im.ToDouble()));
                     if (double.IsNaN(start.Real) || double.IsNaN(start.Imaginary) || double.IsInfinity(start.Real) || double.IsInfinity(start.Imaginary))
                         return null;
                     w = (EDecimal.FromDouble(start.Real), EDecimal.FromDouble(start.Imaginary));
@@ -107,40 +108,6 @@ partial record Number
                 return Complex.Create(w.Re.RoundToPrecision(context), w.Im.RoundToPrecision(context));
             }
 
-            /// <summary>
-            /// W_k(z) in double precision, from a start on the branch and Halley's iteration: the
-            /// branch is chosen here, and the working precision only refines the root.
-            /// </summary>
-            internal static NumericsComplex InDoublePrecision(int k, NumericsComplex z)
-            {
-                NumericsComplex w;
-                var p = NumericsComplex.Sqrt(2 * (Math.E * z + 1));
-                // The series about the branch point, within 1/e of it, on the principal branch; ln(1 + z)
-                // short of the asymptotic reach, which W_0 is close to there and the asymptotic start is
-                // not: from that one, 0.5i went to another branch. Between -1 and -1/e on the axis
-                // ln(1 + z) is real where W_0 is not, and only the series reaches the root.
-                if (k == 0 && NumericsComplex.Abs(Math.E * z + 1) < 1)
-                    w = -1 + p - p * p / 3 + 11 * p * p * p / 72;
-                else if (k == 0 && NumericsComplex.Abs(z) < 3)
-                    w = NumericsComplex.Log(1 + z);
-                else
-                {
-                    var l1 = NumericsComplex.Log(z) + new NumericsComplex(0, 2 * Math.PI * k);
-                    var l2 = NumericsComplex.Log(l1);
-                    w = l1 - l2 + l2 / l1;
-                }
-                for (var step = 0; step < 64; step++)
-                {
-                    var ew = NumericsComplex.Exp(w);
-                    var f = w * ew - z;
-                    var change = f / (ew * (w + 1) - (w + 2) * f / (2 * w + 2));
-                    w -= change;
-                    if (NumericsComplex.Abs(change) <= 1e-15 * Math.Max(NumericsComplex.Abs(w), 1))
-                        break;
-                }
-                return w;
-            }
-
             private static (EDecimal Re, EDecimal Im) Exp((EDecimal Re, EDecimal Im) w, EContext work)
             {
                 var modulus = w.Re.Exponential(work);

Sources/AngouriMath/Functions/Compilation/IntoFE/Compiler.cs

@@ -604,6 +604,19 @@ private protected override void CompileNode(Compiler compiler)
             }
         }
 
+        // The Lambert W function on its branch. https://github.com/asc-community/AngouriMath/issues/2195
+        public partial record LambertWf
+        {
+            private protected override bool CompilesItself => true;
+
+            private protected override void CompileNode(Compiler compiler)
+            {
+                Argument.InnerCompile(compiler);
+                Branch.InnerCompile(compiler);
+                compiler.Instructions.Add(new(InstructionType.CALL_LAMBERTW));
+            }
+        }
+
         public partial record Maxf
         {
             private protected override bool CompilesItself => true;

Sources/AngouriMath/Functions/Compilation/IntoFE/FastExpression.cs

@@ -81,6 +81,7 @@ internal enum InstructionType
             CALL_BINOMIAL,
             CALL_FALLING_FACTORIAL,
             CALL_RISING_FACTORIAL,
+            CALL_LAMBERTW,
 
             // 3-arg: a value where its predicate holds, and the one below it elsewhere.
             CALL_SELECT = 150,
@@ -355,6 +356,9 @@ public System.Numerics.Complex Substitute(params System.Numerics.Complex[] value
                     case InstructionType.CALL_RISING_FACTORIAL:
                         stack.Push(AngouriMath.Numerics.CombinatorialFunctions.RisingFactorial(stack.Pop(), stack.Pop()));
                         break;
+                    case InstructionType.CALL_LAMBERTW:
+                        stack.Push(AngouriMath.Numerics.LambertWFunction.LambertW(stack.Pop(), stack.Pop()));
+                        break;
                     case InstructionType.CALL_DOUBLE_FACTORIAL:
                         stack.Push(AngouriMath.Numerics.CombinatorialFunctions.DoubleFactorial(stack.Pop()));
                         break;

Sources/AngouriMath/Functions/Compilation/IntoLinq/CompilationProtocol.cs

@@ -242,6 +242,8 @@ FallingFactorialf when ShouldBeAtLeastDouble(left) is var newLeft && ShouldBeAtL
                     => Expression.Call(GetDef("FallingFactorial", 2, newLeft.Type), newLeft, newRight),
                 RisingFactorialf when ShouldBeAtLeastDouble(left) is var newLeft && ShouldBeAtLeastDouble(right) is var newRight
                     => Expression.Call(GetDef("RisingFactorial", 2, newLeft.Type), newLeft, newRight),
+                LambertWf when ShouldBeAtLeastDouble(left) is var newLeft && ShouldBeAtLeastDouble(right) is var newRight
+                    => Expression.Call(GetDef("LambertW", 2, newLeft.Type), newLeft, newRight),
                 Maxf => Expression.Call(GetDef("Max", 2, left.Type), left, right),
                 Minf => Expression.Call(GetDef("Min", 2, left.Type), left, right),
 

Sources/AngouriMath/Functions/Compilation/IntoLinq/MathAllMethods.cs

@@ -342,6 +342,11 @@ public static System.Numerics.Complex Round(System.Numerics.Complex a)
         public static System.Numerics.Complex RisingFactorial(System.Numerics.Complex a, System.Numerics.Complex b) => AngouriMath.Numerics.CombinatorialFunctions.RisingFactorial(a, b);
         public static double RisingFactorial(double a, double b) => AngouriMath.Numerics.CombinatorialFunctions.RisingFactorial(a, b);
 
+        // The Lambert W function on the branch the first argument names.
+        // https://github.com/asc-community/AngouriMath/issues/2195
+        public static System.Numerics.Complex LambertW(System.Numerics.Complex a, System.Numerics.Complex b) => AngouriMath.Numerics.LambertWFunction.LambertW(a, b);
+        public static double LambertW(double a, double b) => AngouriMath.Numerics.LambertWFunction.LambertW(a, b);
+
         /// <summary>
         /// True where <c>System.Numerics.Complex.IsNaN</c> is, spelled out because that overload is
         /// .NET 7 and later and this assembly also targets netstandard2.0. A complex number is NaN
@@ -534,6 +539,8 @@ public static bool IsNaN(System.Numerics.Complex a)
             { ("FallingFactorial", 2, typeof(double)), Def<double>((a, b) => FallingFactorial(a, b)) },
             { ("RisingFactorial", 2, typeof(System.Numerics.Complex)), Def<System.Numerics.Complex>((a, b) => RisingFactorial(a, b)) },
             { ("RisingFactorial", 2, typeof(double)), Def<double>((a, b) => RisingFactorial(a, b)) },
+            { ("LambertW", 2, typeof(System.Numerics.Complex)), Def<System.Numerics.Complex>((a, b) => LambertW(a, b)) },
+            { ("LambertW", 2, typeof(double)), Def<double>((a, b) => LambertW(a, b)) },
         };
 
         // An expression tree is how C# spells "the MethodInfo of this method" without a string:

Sources/AngouriMath/Functions/Compilation/IntoLinq/MathAllMethods.tt

@@ -232,6 +232,11 @@ namespace AngouriMath.Core.Compilation.IntoLinq
         public static double <#= name #>(double a, double b) => AngouriMath.Numerics.CombinatorialFunctions.<#= name #>(a, b);
 <# } #>
 
+        // The Lambert W function on the branch the first argument names.
+        // https://github.com/asc-community/AngouriMath/issues/2195
+        public static System.Numerics.Complex LambertW(System.Numerics.Complex a, System.Numerics.Complex b) => AngouriMath.Numerics.LambertWFunction.LambertW(a, b);
+        public static double LambertW(double a, double b) => AngouriMath.Numerics.LambertWFunction.LambertW(a, b);
+
         /// <summary>
         /// True where <c>System.Numerics.Complex.IsNaN</c> is, spelled out because that overload is
         /// .NET 7 and later and this assembly also targets netstandard2.0. A complex number is NaN
@@ -285,6 +290,9 @@ namespace AngouriMath.Core.Compilation.IntoLinq
 <# foreach (var type in new[]{ "System.Numerics.Complex", "double" }) { #>
             { ("<#= name #>", 2, typeof(<#= type #>)), Def<<#= type #>>((a, b) => <#= name #>(a, b)) },
 <# } #>
+<# } #>
+<# foreach (var type in new[]{ "System.Numerics.Complex", "double" }) { #>
+            { ("LambertW", 2, typeof(<#= type #>)), Def<<#= type #>>((a, b) => LambertW(a, b)) },
 <# } #>
         };
 

Sources/AngouriMath/Numerics/LambertWFunction.cs

@@ -0,0 +1,94 @@
+//
+// Copyright (c) 2019-2026 Angouri.
+// AngouriMath is licensed under MIT.
+// Details: https://github.com/asc-community/AngouriMath/blob/master/LICENSE.md.
+// Website: https://am.angouri.org.
+//
+
+using System;
+using NumericsComplex = System.Numerics.Complex;
+
+namespace AngouriMath.Numerics
+{
+    /// <summary>
+    /// The Lambert W function in double precision, for the compilers, on the branches the interpreter
+    /// takes: those of Corless et al. (1996), <c>W_0</c> real on <c>[-1/e, +oo)</c> and <c>W_(-1)</c> on
+    /// <c>[-1/e, 0)</c>, <c>NaN</c> on a branch that is not whole, and <c>-oo</c> at 0 off the principal one.
+    /// https://github.com/asc-community/AngouriMath/issues/2195
+    /// </summary>
+    internal static class LambertWFunction
+    {
+        [AngouriMath.Core.ConstantField] private static readonly NumericsComplex NaN = new(double.NaN, 0);
+
+        public static NumericsComplex LambertW(NumericsComplex branch, NumericsComplex z)
+        {
+            if (branch.Imaginary != 0 || double.IsNaN(branch.Real) || double.IsInfinity(branch.Real)
+                || Math.Floor(branch.Real) != branch.Real || Math.Abs(branch.Real) > int.MaxValue)
+                return NaN;
+            if (double.IsNaN(z.Real) || double.IsNaN(z.Imaginary) || double.IsInfinity(z.Real) || double.IsInfinity(z.Imaginary))
+                return NaN;
+            var k = (int)branch.Real;
+            if (z == 0)
+                return k == 0 ? 0 : double.NegativeInfinity;
+            // Without a signed zero, which a computation can leave behind: z on the axis is approached from above.
+            if (z.Imaginary == 0)
+                z = new NumericsComplex(z.Real, 0);
+
+            var nearTheBranchPoint = Math.E * z + 1;
+            // W_0 from either side, W_(-1) from above and on the axis, W_1 from below meet at -1/e.
+            var meetsTheBranchPoint = k == 0 || k == -1 && z.Imaginary >= 0 || k == 1 && z.Imaginary < 0;
+            // -1/e is not a double, and within its rounding the value is the branch point's.
+            if (meetsTheBranchPoint && NumericsComplex.Abs(nearTheBranchPoint) < 1e-15)
+                return -1;
+            var isReal = z.Imaginary == 0 && nearTheBranchPoint.Real >= 0 && (k == 0 || k == -1 && z.Real < 0);
+
+            // The start, by the regions SciPy's lambertw takes on the principal branch: the series about the
+            // branch point within 0.3 of it (and within a quarter of it on the other branches that meet
+            // there), ln(1 + z) in a box right of -1 about the axis, and the asymptotic start elsewhere. ln(1 + z)
+            // is no start near -1, where it is ln 0 and its iteration reaches another branch, nor on the axis
+            // from -1 to -1/e, where it is real and W_0 is not.
+            NumericsComplex w;
+            var branchPointSeries = k == 0 ? NumericsComplex.Abs(z + 1 / Math.E) < 0.3 : meetsTheBranchPoint && NumericsComplex.Abs(nearTheBranchPoint) < 0.25;
+            if (branchPointSeries)
+            {
+                // p = sqrt(2 (e z + 1)), and w = -1 + p - p^2/3 + 11 p^3/72 on W_0, -p on the other branch.
+                var p = NumericsComplex.Sqrt(2 * nearTheBranchPoint);
+                if (k != 0)
+                    p = -p;
+                w = -1 + p - p * p / 3 + 11 * p * p * p / 72;
+            }
+            else if (k == 0 && z.Real > -1 && z.Real < 1.5 && Math.Abs(z.Imaginary) < 1 && z.Real > -2.5 * Math.Abs(z.Imaginary) - 0.2)
+                w = NumericsComplex.Log(1 + z);
+            else
+            {
+                var l1 = NumericsComplex.Log(z) + new NumericsComplex(0, 2 * Math.PI * k);
+                var l2 = NumericsComplex.Log(l1);
+                w = l1 - l2 + l2 / l1;
+            }
+            if (isReal)
+                w = new NumericsComplex(w.Real, 0);
+            for (var step = 0; step < 64; step++)
+            {
+                // Halley: w - f / (e^w (w + 1) - (w + 2) f / (2 w + 2)), with f = w e^w - z.
+                var ew = NumericsComplex.Exp(w);
+                var f = w * ew - z;
+                if (w + 1 == 0)
+                    break;
+                var change = f / (ew * (w + 1) - (w + 2) * f / (2 * w + 2));
+                w -= change;
+                if (isReal)
+                    w = new NumericsComplex(w.Real, 0);
+                if (NumericsComplex.Abs(change) <= 1e-15 * Math.Max(NumericsComplex.Abs(w), 1))
+                    break;
+            }
+            return w;
+        }
+
+        // A real value where the branch is real there, and NaN where it is not.
+        public static double LambertW(double branch, double z)
+        {
+            var w = LambertW(new NumericsComplex(branch, 0), new NumericsComplex(z, 0));
+            return w.Imaginary == 0 ? w.Real : double.NaN;
+        }
+    }
+}

Sources/Tests/AotSmokeTest/Program.cs

@@ -309,6 +309,8 @@ private static void Compilation()
             var linqBinomial = "binomial(x, 2) + fallingfactorial(x, 3) + risingfactorial(x, 2) + x!!".Compile<Complex, Complex>(x);
             Check("compiled binomial(x, 2) + x(x-1)(x-2) + x(x+1) + x!! at x = 5", fastBinomial.Call(5) == 10 + 60 + 30 + 15 && linqBinomial(5) == 10 + 60 + 30 + 15,
                 $"{fastBinomial.Call(5)}, {linqBinomial(5)}");
+            var compiledW = "lambertw(-1, x)".Compile<double, double>(x);
+            CheckClose("compiled lambertw(-1, x) at x = -0.3", compiledW(-0.3), -1.7813370234216276);
             // FastExpression: a statement is its truth value, and a piecewise the case that holds.
             var fastPiecewise = "piecewise(x provided x > 0, -x)".Compile(x);
             var fastStatement = "x > 1 or x < 0".Compile(x);