some improvements to acb_lambertw; make arb_nonnegative_part public

acb/lambertw.c

@@ -374,7 +374,7 @@ acb_lambertw_cleared_cut(acb_t res, const acb_t z, const fmpz_t k, int flags, sl
     acb_t ez1;
     acb_init(ez1);
 
-    /* compute z*e + 1 */
+    /* compute e*z + 1 */
     arb_const_e(acb_realref(ez1), prec);
     acb_mul(ez1, ez1, z, prec);
     acb_add_ui(ez1, ez1, 1, prec);
@@ -411,6 +411,81 @@ acb_lambertw_cleared_cut(acb_t res, const acb_t z, const fmpz_t k, int flags, sl
     acb_clear(ez1);
 }
 
+/* Extremely close to the branch point at -1/e, use the series expansion directly. */
+int
+acb_lambertw_try_near_branch_point(acb_t res, const acb_t z,
+    const acb_t ez1, const fmpz_t k, int flags, slong prec)
+{
+    if (fmpz_is_zero(k) || (fmpz_is_one(k) && arb_is_negative(acb_imagref(z)))
+                        || (fmpz_equal_si(k, -1) && arb_is_nonnegative(acb_imagref(z))))
+    {
+        if (acb_contains_zero(ez1) ||
+            (arf_cmpabs_2exp_si(arb_midref(acb_realref(ez1)), -prec / 4.5) < 0 &&
+             arf_cmpabs_2exp_si(arb_midref(acb_imagref(ez1)), -prec / 4.5) < 0))
+        {
+            acb_t t;
+            acb_init(t);
+            acb_mul_2exp_si(t, ez1, 1);
+            acb_sqrt(t, t, prec);
+            if (!fmpz_is_zero(k))
+                acb_neg(t, t);
+            acb_lambertw_branchpoint_series(res, t, 1, prec);
+            acb_clear(t);
+            return 1;
+        }
+    }
+
+    return 0;
+}
+
+void
+acb_lambertw_cleared_cut_fix_small(acb_t res, const acb_t z,
+    const acb_t ez1, const fmpz_t k, int flags, slong prec)
+{
+    acb_t zz, zmid, zmide1;
+    arf_t eps;
+
+    acb_init(zz);
+    acb_init(zmid);
+    acb_init(zmide1);
+    arf_init(eps);
+
+    arf_mul_2exp_si(eps, arb_midref(acb_realref(z)), -prec);
+    acb_set(zz, z);
+
+    if (arf_sgn(arb_midref(acb_realref(zz))) < 0 &&
+        (!fmpz_is_zero(k) || arf_sgn(arb_midref(acb_realref(ez1))) < 0) &&
+        arf_cmpabs(arb_midref(acb_imagref(zz)), eps) < 0)
+    {
+        /* now the value must be in [0,2eps] */
+        arf_get_mag(arb_radref(acb_imagref(zz)), eps);
+        arf_set_mag(arb_midref(acb_imagref(zz)), arb_radref(acb_imagref(zz)));
+
+        if (arf_sgn(arb_midref(acb_imagref(z))) >= 0)
+        {
+            acb_lambertw_cleared_cut(res, zz, k, flags, prec);
+        }
+        else
+        {
+            fmpz_t kk;
+            fmpz_init(kk);
+            fmpz_neg(kk, k);
+            acb_lambertw_cleared_cut(res, zz, kk, flags, prec);
+            acb_conj(res, res);
+            fmpz_clear(kk);
+        }
+    }
+    else
+    {
+        acb_lambertw_cleared_cut(res, zz, k, flags, prec);
+    }
+
+    acb_clear(zz);
+    acb_clear(zmid);
+    acb_clear(zmide1);
+    arf_clear(eps);
+}
+
 void
 _acb_lambertw(acb_t res, const acb_t z, const acb_t ez1, const fmpz_t k, int flags, slong prec)
 {
@@ -483,77 +558,56 @@ _acb_lambertw(acb_t res, const acb_t z, const acb_t ez1, const fmpz_t k, int fla
     }
 
     /* Extremely close to the branch point at -1/e, use the series expansion directly. */
-    if (fmpz_is_zero(k) || (fmpz_is_one(k) && arb_is_negative(acb_imagref(z)))
-                        || (fmpz_equal_si(k, -1) && arb_is_nonnegative(acb_imagref(z))))
+    if (acb_lambertw_try_near_branch_point(res, z, ez1, k, flags, goal))
     {
-        if (acb_contains_zero(ez1) ||
-            (arf_cmpabs_2exp_si(arb_midref(acb_realref(ez1)), -goal / 4.5) < 0 &&
-             arf_cmpabs_2exp_si(arb_midref(acb_imagref(ez1)), -goal / 4.5) < 0))
-        {
-            acb_t t;
-            acb_init(t);
-            acb_mul_2exp_si(t, ez1, 1);
-            acb_sqrt(t, t, goal);
-            if (!fmpz_is_zero(k))
-                acb_neg(t, t);
-            acb_lambertw_branchpoint_series(res, t, 1, goal);
-            acb_set_round(res, res, prec);
-            acb_clear(t);
-            return;
-        }
+        acb_set_round(res, res, prec);
+        return;
     }
 
-    /* todo: compute union of two results */
+    /* compute union of both sides */
     if (acb_lambertw_branch_crossing(z, ez1, k))
     {
-        acb_indeterminate(res);
+        acb_t za, zb, eza1, ezb1;
+        fmpz_t kk;
+
+        acb_init(za);
+        acb_init(zb);
+        acb_init(eza1);
+        acb_init(ezb1);
+        fmpz_init(kk);
+
+        fmpz_neg(kk, k);
+
+        acb_set(za, z);
+        acb_conj(zb, z);
+        arb_nonnegative_part(acb_imagref(za), acb_imagref(za));
+        arb_nonnegative_part(acb_imagref(zb), acb_imagref(zb));
+
+        acb_set(eza1, ez1);
+        acb_conj(ezb1, ez1);
+        arb_nonnegative_part(acb_imagref(eza1), acb_imagref(eza1));
+        arb_nonnegative_part(acb_imagref(ezb1), acb_imagref(ezb1));
+
+        /* Check series expansion again, because now there is no crossing. */
+        if (!acb_lambertw_try_near_branch_point(res, za, eza1, k, flags, goal))
+            acb_lambertw_cleared_cut_fix_small(za, za, eza1, k, flags, goal);
+
+        if (!acb_lambertw_try_near_branch_point(res, zb, ezb1, kk, flags, goal))
+            acb_lambertw_cleared_cut_fix_small(zb, zb, ezb1, kk, flags, goal);
+
+        acb_conj(zb, zb);
+        acb_union(res, za, zb, prec);
+
+        acb_clear(za);
+        acb_clear(zb);
+        acb_clear(eza1);
+        acb_clear(ezb1);
+        fmpz_clear(kk);
     }
     else
     {
-        acb_t zz, zmid, zmide1;
-        arf_t eps;
-
-        acb_init(zz);
-        acb_init(zmid);
-        acb_init(zmide1);
-        arf_init(eps);
-
-        arf_mul_2exp_si(eps, arb_midref(acb_realref(z)), -goal);
-        acb_set(zz, z);
-
-        if (arf_sgn(arb_midref(acb_realref(zz))) < 0 &&
-            (!fmpz_is_zero(k) || arf_sgn(arb_midref(acb_realref(ez1))) < 0) &&
-            arf_cmpabs(arb_midref(acb_imagref(zz)), eps) < 0)
-        {
-            /* now the value must be in [0,2eps] */
-            arf_get_mag(arb_radref(acb_imagref(zz)), eps);
-            arf_set_mag(arb_midref(acb_imagref(zz)), arb_radref(acb_imagref(zz)));
-
-            if (arf_sgn(arb_midref(acb_imagref(z))) >= 0)
-            {
-                acb_lambertw_cleared_cut(res, zz, k, flags, goal);
-            }
-            else
-            {
-                fmpz_t kk;
-                fmpz_init(kk);
-                fmpz_neg(kk, k);
-                acb_lambertw_cleared_cut(res, zz, kk, flags, goal);
-                acb_conj(res, res);
-                fmpz_clear(kk);
-            }
-        }
-        else
-        {
-            acb_lambertw_cleared_cut(res, zz, k, flags, goal);
-        }
-
+        acb_lambertw_cleared_cut_fix_small(res, z, ez1, k, flags, goal);
         acb_set_round(res, res, prec);
-
-        acb_clear(zz);
-        acb_clear(zmid);
-        acb_clear(zmide1);
-        arf_clear(eps);
     }
 }
 

arb.h

@@ -327,6 +327,8 @@ arb_get_lbound_arf(arf_t u, const arb_t x, long prec)
     arf_sub(u, arb_midref(x), t, prec, ARF_RND_FLOOR);
 }
 
+void arb_nonnegative_part(arb_t res, const arb_t x);
+
 slong arb_rel_error_bits(const arb_t x);
 
 ARB_INLINE slong

arb/nonnegative_part.c

@@ -0,0 +1,49 @@
+/*
+    Copyright (C) 2017 Fredrik Johansson
+
+    This file is part of Arb.
+
+    Arb is free software: you can redistribute it and/or modify it under
+    the terms of the GNU Lesser General Public License (LGPL) as published
+    by the Free Software Foundation; either version 2.1 of the License, or
+    (at your option) any later version.  See <http://www.gnu.org/licenses/>.
+*/
+
+#include "arb.h"
+
+void
+arb_nonnegative_part(arb_t res, const arb_t x)
+{
+    if (!arb_contains_negative(x))
+    {
+        arb_set(res, x);
+    }
+    else if (!arb_is_finite(x))
+    {
+        arb_indeterminate(res);
+    }
+    else
+    {
+        arf_t t;
+        arf_init(t);
+
+        arf_set_mag(t, arb_radref(x));
+        arf_add(arb_midref(res), arb_midref(x), t, MAG_BITS, ARF_RND_CEIL);
+
+        if (arf_sgn(arb_midref(res)) <= 0)
+        {
+            arf_zero(arb_midref(res));
+            mag_zero(arb_radref(res));
+        }
+        else
+        {
+            arf_mul_2exp_si(arb_midref(res), arb_midref(res), -1);
+            arf_get_mag(arb_radref(res), arb_midref(res));
+            /* needed since arf_get_mag is inexact */
+            arf_set_mag(arb_midref(res), arb_radref(res));
+        }
+
+        arf_clear(t);
+    }
+}
+

arb/sqrtpos.c

@@ -11,38 +11,6 @@
 
 #include "arb.h"
 
-static __inline__ void
-arb_nonnegative_part(arb_t z, const arb_t x, slong prec)
-{
-    if (arb_contains_negative(x))
-    {
-        arf_t t;
-        arf_init(t);
-
-        arf_set_mag(t, arb_radref(x));
-        arf_add(arb_midref(z), arb_midref(x), t, MAG_BITS, ARF_RND_CEIL);
-
-        if (arf_sgn(arb_midref(z)) <= 0)
-        {
-            mag_zero(arb_radref(z));
-        }
-        else
-        {
-            arf_mul_2exp_si(arb_midref(z), arb_midref(z), -1);
-            arf_get_mag(arb_radref(z), arb_midref(z));
-
-            /* XXX: needed since arf_get_mag is inexact */
-            arf_set_mag(arb_midref(z), arb_radref(z));
-        }
-
-        arf_clear(t);
-    }
-    else
-    {
-        arb_set(z, x);
-    }
-}
-
 void
 arb_sqrtpos(arb_t z, const arb_t x, slong prec)
 {
@@ -81,6 +49,6 @@ arb_sqrtpos(arb_t z, const arb_t x, slong prec)
         arb_sqrt(z, x, prec);
     }
 
-    arb_nonnegative_part(z, z, prec);
+    arb_nonnegative_part(z, z);
 }
 

doc/source/arb.rst

@@ -342,6 +342,14 @@ Radius and interval operations
     Otherwise zero is returned and the value of *z* is undefined.
     If *x* or *y* contains NaN, the result is NaN.
 
+.. function:: void arb_nonnegative_part(arb_t res, const arb_t x)
+
+    Sets *res* to the intersection of *x* with `[0,\infty]`. If *x* is
+    nonnegative, an exact copy is made. If *x* is finite and contains negative
+    numbers, an interval of the form `[r/2 \pm r/2]` is produced, which
+    certainly contains no negative points.
+    In the special case when *x* is strictly negative, *res* is set to zero.
+
 .. function:: void arb_get_abs_ubound_arf(arf_t u, const arb_t x, slong prec)
 
     Sets *u* to the upper bound for the absolute value of *x*,
@@ -1114,8 +1122,8 @@ Lambert W function
     The principal branch is real-valued for `x \ge -1/e`
     (taking values in `[-1,+\infty)`) and the `W_{-1}` branch is real-valued
     for `-1/e \le x < 0` and takes values in `(-\infty,-1]`.
-    Elsewhere, the Lambert W function is complex and
-    cannot be evaluated by this implementation.
+    Elsewhere, the Lambert W function is complex and :func:`acb_lambertw`
+    should be used.
 
     The implementation first computes a floating-point approximation
     heuristically and then computes a rigorously certified enclosure around