# Source code for diofant.functions.special.zeta_functions

""" Riemann zeta and related function. """

from ...core import (Add, Dummy, E, Function, I, Integer, Rational, expand_mul,
oo, pi, sympify, zoo)
from ...core.function import ArgumentIndexError
from ..combinatorial.numbers import bernoulli, factorial, harmonic
from ..elementary.exponential import exp, exp_polar, log

###############################################################################
# #################### LERCH TRANSCENDENT ################################### #
###############################################################################

[docs]class lerchphi(Function):
r"""
Lerch transcendent (Lerch phi function).

For \operatorname{Re}(a) > 0, |z| < 1 and s \in \mathbb{C}, the
Lerch transcendent is defined as

.. math :: \Phi(z, s, a) = \sum_{n=0}^\infty \frac{z^n}{(n + a)^s},

where the standard branch of the argument is used for n + a,
and by analytic continuation for other values of the parameters.

A commonly used related function is the Lerch zeta function, defined by

.. math:: L(q, s, a) = \Phi(e^{2\pi i q}, s, a).

**Analytic Continuation and Branching Behavior**

It can be shown that

.. math:: \Phi(z, s, a) = z\Phi(z, s, a+1) + a^{-s}.

This provides the analytic continuation to \operatorname{Re}(a) \le 0.

Assume now \operatorname{Re}(a) > 0. The integral representation

.. math:: \Phi_0(z, s, a) = \int_0^\infty \frac{t^{s-1} e^{-at}}{1 - ze^{-t}}
\frac{\mathrm{d}t}{\Gamma(s)}

provides an analytic continuation to \mathbb{C} - [1, \infty).
Finally, for x \in (1, \infty) we find

.. math:: \lim_{\epsilon \to 0^+} \Phi_0(x + i\epsilon, s, a)
-\lim_{\epsilon \to 0^+} \Phi_0(x - i\epsilon, s, a)
= \frac{2\pi i \log^{s-1}{x}}{x^a \Gamma(s)},

using the standard branch for both \log{x} and
\log{\log{x}} (a branch of \log{\log{x}} is needed to
evaluate \log{x}^{s-1}).
This concludes the analytic continuation. The Lerch transcendent is thus
branched at z \in \{0, 1, \infty\} and
a \in \mathbb{Z}_{\le 0}. For fixed z, a outside these
branch points, it is an entire function of s.

========

polylog, zeta

References
==========

* Bateman, H.; Erdélyi, A. (1953), Higher Transcendental Functions,
Vol. I, New York: McGraw-Hill. Section 1.11.
* https://dlmf.nist.gov/25.14
* https://en.wikipedia.org/wiki/Lerch_transcendent

Examples
========

The Lerch transcendent is a fairly general function, for this reason it does
not automatically evaluate to simpler functions. Use expand_func() to
achieve this.

If z=1, the Lerch transcendent reduces to the Hurwitz zeta function:

>>> from diofant.abc import s
>>> expand_func(lerchphi(1, s, a))
zeta(s, a)

More generally, if z is a root of unity, the Lerch transcendent
reduces to a sum of Hurwitz zeta functions:

>>> expand_func(lerchphi(-1, s, a))
2**(-s)*zeta(s, a/2) - 2**(-s)*zeta(s, a/2 + 1/2)

If a=1, the Lerch transcendent reduces to the polylogarithm:

>>> expand_func(lerchphi(z, s, 1))
polylog(s, z)/z

More generally, if a is rational, the Lerch transcendent reduces
to a sum of polylogarithms:

>>> expand_func(lerchphi(z, s, Rational(1, 2)))
2**(s - 1)*(polylog(s, sqrt(z))/sqrt(z) -
polylog(s, sqrt(z)*exp_polar(I*pi))/sqrt(z))
>>> expand_func(lerchphi(z, s, Rational(3, 2)))
-2**s/z + 2**(s - 1)*(polylog(s, sqrt(z))/sqrt(z) -
polylog(s, sqrt(z)*exp_polar(I*pi))/sqrt(z))/z

The derivatives with respect to z and a can be computed in
closed form:

>>> lerchphi(z, s, a).diff(z)
(-a*lerchphi(z, s, a) + lerchphi(z, s - 1, a))/z
>>> lerchphi(z, s, a).diff(a)
-s*lerchphi(z, s + 1, a)

"""

def _eval_expand_func(self, **hints):
from ...polys import Poly
from .. import floor, unpolarify
z, s, a = self.args
if z == 1:
return zeta(s, a)
if s.is_Integer and s <= 0:
t = Dummy('t')
p = Poly((t + a)**(-s), t)
start = 1/(1 - t)
res = Integer(0)
for c in reversed(p.all_coeffs()):
res += c*start
start = t*start.diff(t)
return res.subs({t: z})

if a.is_Rational:
# See section 18 of
#   Kelly B. Roach.  Hypergeometric Function Representations.
#   In: Proceedings of the 1997 International Symposium on Symbolic and
#   Algebraic Computation, pages 205-211, New York, 1997. ACM.
# TODO should something be polarified here?
mul = Integer(1)
# First reduce a to the interaval (0, 1]
if a > 1:
n = floor(a)
if n == a:
n -= 1
a -= n
mul = z**(-n)
elif a <= 0:
n = floor(-a) + 1
a += n
mul = z**n
add = Add(*[z**(n - 1 - k)/(a - k - 1)**s for k in range(n)])

m, n = Integer(a.numerator), Integer(a.denominator)
zet = exp_polar(2*pi*I/n)
root = z**(1/n)
*[polylog(s, zet**k*root)._eval_expand_func(**hints)
/ (unpolarify(zet)**k*root)**m for k in range(n)])

# TODO use minpoly instead of ad-hoc methods when issue sympy/sympy#5888 is fixed
if z.is_Pow and z.base is E and (z.exp/(pi*I)).is_Rational or z in [-1, I, -I]:
# TODO reference?
if z == -1:
p, q = Integer(1), Integer(2)
elif z == I:
p, q = Integer(1), Integer(4)
elif z == -I:
p, q = Integer(-1), Integer(4)
else:
arg = z.exp/(2*pi*I)
p, q = Integer(arg.numerator), Integer(arg.denominator)
for k in range(q)])

return lerchphi(z, s, a)

def fdiff(self, argindex=1):
z, s, a = self.args
if argindex == 3:
return -s*lerchphi(z, s + 1, a)
elif argindex == 1:
return (lerchphi(z, s - 1, a) - a*lerchphi(z, s, a))/z
else:
raise ArgumentIndexError

def _eval_rewrite_helper(self, z, s, a, target):
res = self._eval_expand_func()
if res.has(target):
return res

def _eval_rewrite_as_zeta(self, z, s, a):
return self._eval_rewrite_helper(z, s, a, zeta)

def _eval_rewrite_as_polylog(self, z, s, a):
return self._eval_rewrite_helper(z, s, a, polylog)

###############################################################################
# #################### POLYLOGARITHM ######################################## #
###############################################################################

[docs]class polylog(Function):
r"""
Polylogarithm function.

For |z| < 1 and s \in \mathbb{C}, the polylogarithm is
defined by

.. math:: \operatorname{Li}_s(z) = \sum_{n=1}^\infty \frac{z^n}{n^s},

where the standard branch of the argument is used for n. It admits
an analytic continuation which is branched at z=1 (notably not on the
sheet of initial definition), z=0 and z=\infty.

The name polylogarithm comes from the fact that for s=1, the
polylogarithm is related to the ordinary logarithm (see examples), and that

.. math:: \operatorname{Li}_{s+1}(z) =
\int_0^z \frac{\operatorname{Li}_s(t)}{t} \mathrm{d}t.

The polylogarithm is a special case of the Lerch transcendent:

.. math:: \operatorname{Li}_{s}(z) = z \Phi(z, s, 1)

========

zeta, lerchphi

Examples
========

For z \in \{0, 1, -1\}, the polylogarithm is automatically expressed
using other functions:

>>> from diofant.abc import s
>>> polylog(s, 0)
0
>>> polylog(s, 1)
zeta(s)
>>> polylog(s, -1)
-dirichlet_eta(s)

If s is a negative integer, 0 or 1, the
polylogarithm can be expressed using elementary functions. This can be
done using expand_func():

>>> expand_func(polylog(1, z))
-log(-z + 1)
>>> expand_func(polylog(0, z))
z/(-z + 1)

The derivative with respect to z can be computed in closed form:

>>> polylog(s, z).diff(z)
polylog(s - 1, z)/z

The polylogarithm can be expressed in terms of the lerch transcendent:

>>> polylog(s, z).rewrite(lerchphi)
z*lerchphi(z, s, 1)

References
==========

* https://en.wikipedia.org/wiki/Polylogarithm
* http://mathworld.wolfram.com/Polylogarithm.html

"""

@classmethod
def eval(cls, s, z):
from .. import unpolarify
if z == 1:
return zeta(s)
elif z == -1:
return -dirichlet_eta(s)
elif z == 0:
return Integer(0)

# branch handling
if (1 - abs(z)).is_nonnegative:
newz = unpolarify(z)
if newz != z:
return cls(s, newz)

def fdiff(self, argindex=1):
s, z = self.args
if argindex == 2:
return polylog(s - 1, z)/z
raise ArgumentIndexError

def _eval_rewrite_as_lerchphi(self, s, z):
return z*lerchphi(z, s, 1)

def _eval_expand_func(self, **hints):
s, z = self.args
if s == 1:
return -log(1 - z)
if s.is_Integer and s <= 0:
u = Dummy('u')
start = u/(1 - u)
for _ in range(-s):
start = u*start.diff(u)
return expand_mul(start).subs({u: z})
return polylog(s, z)

###############################################################################
# #################### HURWITZ GENERALIZED ZETA FUNCTION #################### #
###############################################################################

[docs]class zeta(Function):
r"""
Hurwitz zeta function (or Riemann zeta function).

For \operatorname{Re}(a) > 0 and \operatorname{Re}(s) > 1, this function is defined as

.. math:: \zeta(s, a) = \sum_{n=0}^\infty \frac{1}{(n + a)^s},

where the standard choice of argument for n + a is used. For fixed
a with \operatorname{Re}(a) > 0 the Hurwitz zeta function admits a
meromorphic continuation to all of \mathbb{C}, it is an unbranched
function with a simple pole at s = 1.

Analytic continuation to other a is possible under some circumstances,
but this is not typically done.

The Hurwitz zeta function is a special case of the Lerch transcendent:

.. math:: \zeta(s, a) = \Phi(1, s, a).

This formula defines an analytic continuation for all possible values of
s and a (also \operatorname{Re}(a) < 0), see the documentation of
:class:lerchphi for a description of the branching behavior.

If no value is passed for a, by this function assumes a default value
of a = 1, yielding the Riemann zeta function.

========

dirichlet_eta, lerchphi, polylog

References
==========

* https://dlmf.nist.gov/25.11
* https://en.wikipedia.org/wiki/Hurwitz_zeta_function

Examples
========

For a = 1 the Hurwitz zeta function reduces to the famous Riemann
zeta function:

.. math:: \zeta(s, 1) = \zeta(s) = \sum_{n=1}^\infty \frac{1}{n^s}.

>>> from diofant.abc import s
>>> zeta(s, 1)
zeta(s)
>>> zeta(s)
zeta(s)

The Riemann zeta function can also be expressed using the Dirichlet eta
function:

>>> zeta(s).rewrite(dirichlet_eta)
dirichlet_eta(s)/(-2**(-s + 1) + 1)

The Riemann zeta function at positive even integer and negative odd integer
values is related to the Bernoulli numbers:

>>> zeta(2)
pi**2/6
>>> zeta(4)
pi**4/90
>>> zeta(-1)
-1/12

The specific formulae are:

.. math:: \zeta(2n) = (-1)^{n+1} \frac{B_{2n} (2\pi)^{2n}}{2(2n)!}
.. math:: \zeta(-n) = -\frac{B_{n+1}}{n+1}

At negative even integers the Riemann zeta function is zero:

>>> zeta(-4)
0

No closed-form expressions are known at positive odd integers, but
numerical evaluation is possible:

>>> zeta(3).evalf()
1.20205690315959

The derivative of \zeta(s, a) with respect to a is easily
computed:

>>> zeta(s, a).diff(a)
-s*zeta(s + 1, a)

However the derivative with respect to s has no useful closed form
expression:

>>> zeta(s, a).diff(s)
Derivative(zeta(s, a), s)

The Hurwitz zeta function can be expressed in terms of the Lerch transcendent,
:class:diofant.functions.special.zeta_functions.lerchphi:

>>> zeta(s, a).rewrite(lerchphi)
lerchphi(1, s, a)

"""

@classmethod
def eval(cls, z, a_=None):
if a_ is None:
z, a = list(map(sympify, (z, 1)))
else:
z, a = list(map(sympify, (z, a_)))

if a.is_Number:
if a == 1 and a_ is not None:
return cls(z)
# TODO Should a == 0 return nan as well?

if z.is_Number:
if z is oo:
return Integer(1)
elif z == 0:
return Rational(1, 2) - a
elif z == 1:
return zoo
elif z.is_Integer:
if a.is_Integer:
if z.is_negative:
zeta = (-1)**z * bernoulli(-z + 1)/(-z + 1)
elif z.is_even:
B, F = bernoulli(z), factorial(z)
zeta = 2**(z - 1) * abs(B) * pi**z / F
else:
return

if a.is_negative:
return zeta + harmonic(abs(a), z)
else:
return zeta - harmonic(a - 1, z)

def _eval_rewrite_as_dirichlet_eta(self, s, a=1):
if a != 1:
return self
s = self.args[0]
return dirichlet_eta(s)/(1 - 2**(1 - s))

def _eval_rewrite_as_lerchphi(self, s, a=1):
return lerchphi(1, s, a)

def fdiff(self, argindex=1):
if len(self.args) == 2:
s, a = self.args
else:
s, a = self.args + (1,)
if argindex == 2:
return -s*zeta(s + 1, a)
else:
raise ArgumentIndexError

def _eval_rewrite_as_tractable(self, s, a=1):
if len(self.args) == 1:
return _zetas(exp(s))

class _zetas(Function):
def _eval_rewrite_as_intractable(self, s):
return zeta(log(s))

def _eval_aseries(self, n, args0, x, logx):
from ...series import Order
point = args0[0]

# Expansion at oo
if point is oo:
if n < 1:
return Order(1, x)
z = self.args[0]
l = [(1/z)**log(p) for p in range(1, n)]
o = Order(1/z**log(n), x)
return (Add(*l))._eval_nseries(x, n, logx) + o

# All other points are not handled
return super()._eval_aseries(n, args0, x, logx)

[docs]class dirichlet_eta(Function):
r"""
Dirichlet eta function.

For \operatorname{Re}(s) > 0, this function is defined as

.. math:: \eta(s) = \sum_{n=1}^\infty \frac{(-1)^{n-1}}{n^s}.

It admits a unique analytic continuation to all of \mathbb{C}.
It is an entire, unbranched function.

========

zeta

References
==========

* https://en.wikipedia.org/wiki/Dirichlet_eta_function
* http://mathworld.wolfram.com/DirichletEtaFunction.html

Examples
========

The Dirichlet eta function is closely related to the Riemann zeta function:

>>> from diofant.abc import s
>>> dirichlet_eta(s).rewrite(zeta)
(-2**(-s + 1) + 1)*zeta(s)

"""

@classmethod
def eval(cls, s):
if s == 1:
return log(2)
z = zeta(s)
if not z.has(zeta):
return (1 - 2**(1 - s))*z

def _eval_rewrite_as_zeta(self, s):
return (1 - 2**(1 - s)) * zeta(s)