void
acb_lambertw_left(acb_t res, const acb_t z, const fmpz_t k, slong prec)
{
if (acb_contains_zero(z) && !(fmpz_equal_si(k, -1) && acb_is_real(z)))
{
acb_indeterminate(res);
return;
}
if (arb_is_positive(acb_imagref(z)))
{
acb_lambertw(res, z, k, 0, prec);
}
else if (arb_is_nonpositive(acb_imagref(z)))
{
fmpz_t kk;
fmpz_init(kk);
fmpz_add_ui(kk, k, 1);
fmpz_neg(kk, kk);
acb_conj(res, z);
acb_lambertw(res, res, kk, 0, prec);
acb_conj(res, res);
fmpz_clear(kk);
}
else
{
acb_t za, zb;
fmpz_t kk;
acb_init(za);
acb_init(zb);
fmpz_init(kk);
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));
fmpz_add_ui(kk, k, 1);
fmpz_neg(kk, kk);
acb_lambertw(za, za, k, 0, prec);
acb_lambertw(zb, zb, kk, 0, prec);
acb_conj(zb, zb);
acb_union(res, za, zb, prec);
acb_clear(za);
acb_clear(zb);
fmpz_clear(kk);
}
}
void
arb_sqrtpos(arb_t z, const arb_t x, long prec)
{
if (!arb_is_finite(x))
{
if (mag_is_zero(arb_radref(x)) && arf_is_pos_inf(arb_midref(x)))
arb_pos_inf(z);
else
arb_zero_pm_inf(z);
}
else if (arb_contains_nonpositive(x))
{
arf_t t;
arf_init(t);
arf_set_mag(t, arb_radref(x));
arf_add(t, arb_midref(x), t, MAG_BITS, ARF_RND_CEIL);
if (arf_sgn(t) <= 0)
{
arb_zero(z);
}
else
{
arf_sqrt(t, t, MAG_BITS, ARF_RND_CEIL);
arf_mul_2exp_si(t, t, -1);
arf_set(arb_midref(z), t);
arf_get_mag(arb_radref(z), t);
}
arf_clear(t);
}
else
{
arb_sqrt(z, x, prec);
}
arb_nonnegative_part(z, z, prec);
}
void
acb_lambertw_middle(acb_t res, const acb_t z, slong prec)
{
fmpz_t k;
if (acb_contains_zero(z))
{
acb_indeterminate(res);
return;
}
fmpz_init(k);
fmpz_set_si(k, -1);
if (arb_is_positive(acb_imagref(z)))
{
acb_lambertw(res, z, k, 0, prec);
}
else if (arb_is_negative(acb_imagref(z)))
{
acb_conj(res, z);
acb_lambertw(res, res, k, 0, prec);
acb_conj(res, res);
}
else if (arb_is_negative(acb_realref(z)))
{
if (arb_is_nonnegative(acb_imagref(z)))
{
acb_lambertw(res, z, k, 0, prec);
}
else if (arb_is_negative(acb_imagref(z)))
{
acb_conj(res, z);
acb_lambertw(res, res, k, 0, prec);
acb_conj(res, res);
}
else
{
acb_t za, zb;
acb_init(za);
acb_init(zb);
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_lambertw(za, za, k, 0, prec);
acb_lambertw(zb, zb, k, 0, prec);
acb_conj(zb, zb);
acb_union(res, za, zb, prec);
acb_clear(za);
acb_clear(zb);
}
}
else /* re is positive */
{
if (arb_is_positive(acb_imagref(z)))
{
acb_lambertw(res, z, k, 0, prec);
}
else if (arb_is_nonpositive(acb_imagref(z)))
{
acb_conj(res, z);
acb_lambertw(res, res, k, 0, prec);
acb_conj(res, res);
}
else
{
acb_t za, zb;
acb_init(za);
acb_init(zb);
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_lambertw(za, za, k, 0, prec);
acb_lambertw(zb, zb, k, 0, prec);
acb_conj(zb, zb);
acb_union(res, za, zb, prec);
acb_clear(za);
acb_clear(zb);
}
}
fmpz_clear(k);
}
void
_acb_lambertw(acb_t res, const acb_t z, const acb_t ez1, const fmpz_t k, int flags, slong prec)
{
slong goal, ebits, ebits2, ls, lt;
const fmpz * expo;
/* Estimated accuracy goal. */
/* todo: account for exponent bits and bits in k. */
goal = acb_rel_accuracy_bits(z);
goal = FLINT_MAX(goal, 10);
goal = FLINT_MIN(goal, prec);
/* Handle tiny z directly. For k >= 2, |c_k| <= 4^k / 16. */
if (fmpz_is_zero(k)
&& arf_cmpabs_2exp_si(arb_midref(acb_realref(z)), -goal / 2) < 0
&& arf_cmpabs_2exp_si(arb_midref(acb_imagref(z)), -goal / 2) < 0)
{
mag_t err;
mag_init(err);
acb_get_mag(err, z);
mag_mul_2exp_si(err, err, 2);
acb_set(res, z);
acb_submul(res, res, res, prec);
mag_geom_series(err, err, 3);
mag_mul_2exp_si(err, err, -4);
acb_add_error_mag(res, err);
mag_clear(err);
return;
}
if (arf_cmpabs(arb_midref(acb_realref(z)), arb_midref(acb_imagref(z))) >= 0)
expo = ARF_EXPREF(arb_midref(acb_realref(z)));
else
expo = ARF_EXPREF(arb_midref(acb_imagref(z)));
ebits = fmpz_bits(expo);
/* ebits ~= log2(|log(z) + 2 pi i k|) */
/* ebits2 ~= log2(log(log(z))) */
ebits = FLINT_MAX(ebits, fmpz_bits(k));
ebits = FLINT_MAX(ebits, 1) - 1;
ebits2 = FLINT_BIT_COUNT(ebits);
ebits2 = FLINT_MAX(ebits2, 1) - 1;
/* We gain accuracy from the exponent when W ~ log - log log */
if (fmpz_sgn(expo) > 0 || (fmpz_sgn(expo) < 0 && !fmpz_is_zero(k)))
{
goal += ebits - ebits2;
goal = FLINT_MAX(goal, 10);
goal = FLINT_MIN(goal, prec);
/* The asymptotic series with truncation L, M gives us about
t - max(2+lt+L*(2+ls), M*(2+lt)) bits of accuracy where
ls = -ebits, lt = ebits2 - ebits. */
ls = 2 - ebits;
lt = 2 + ebits2 - ebits;
if (ebits - FLINT_MAX(lt + 1*ls, 1*lt) > goal)
{
acb_lambertw_asymp(res, z, k, 1, 1, goal);
acb_set_round(res, res, prec);
return;
}
else if (ebits - FLINT_MAX(lt + 3*ls, 5*lt) > goal)
{
acb_lambertw_asymp(res, z, k, 3, 5, goal);
acb_set_round(res, res, prec);
return;
}
}
/* Extremely close to the branch point at -1/e, use the series expansion directly. */
if (acb_lambertw_try_near_branch_point(res, z, ez1, k, flags, goal))
{
acb_set_round(res, res, prec);
return;
}
/* compute union of both sides */
if (acb_lambertw_branch_crossing(z, ez1, k))
{
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_lambertw_cleared_cut_fix_small(res, z, ez1, k, flags, goal);
acb_set_round(res, res, prec);
}
}
