Understanding and computing the Riemann zeta function

riemann_zeta_function.py

"""
Script discussed in this video:

https://youtu.be/Fl3XgPpvSNI

https://en.wikipedia.org/wiki/Riemann_zeta_function
http://www.claymath.org/sites/default/files/ezeta.pdf
http://www.claymath.org/sites/default/files/zeta.pdf

Note:
$ \phi_n = \log(n) $

$ n^{i\,t}={\mathrm e}^{\log(n^{i\,t})}={\mathrm e}^{i\,t\,\phi_n}=\cos(t\,\phi_n)+i\sin(t\,\phi_n) $

Define functions:
$ \zeta(s) := \sum_{n=1}^\infty n^{-s} $
$ \eta(s) := \sum_{n=1}^\infty n^{-s}\cdot(-1)^{n-1} $

$s\equiv x+iy, x,y\in{\mathbb R}$
Theorem:
$x > 0 \implies \zeta(s) = \dfrac{1}{1-2^{1-s}}\cdot \eta(s)$

So in terms of $x,y$,
$ \zeta(x+i\,y) = \dfrac{2^{i\,y}}{2^{i\,y}-2^{1-x}}\cdot\sum_{n=1}^\infty \dfrac{(-1)^{n-1}}{n^x}\left[\cos(y\,\phi_n)-i\sin(y\,\phi_n)\right]$

E.g. $x=1/2, y=0$, so that the sum is over alternating $1/\sqrt{n}$,
(1/(1-Sqrt[2]))Sum[(-1)^(n-1)/Sqrt[n],{n,1,Infinity}] == Pi^2 / 6
"""

import matplotlib.pyplot as plt
import numpy as np


def _zeta_over_eta_factor(x, y):
    X = -2**(1 - x)
    Y = 2**(1j * y)
    res = Y / (X + Y)
    """
    P = y * np.log(2)
    # Alternatively, isolating the imaginary part via 2^(jy) = cos(P) + j sin(P)
    Re = X * np.cos(P) + 1
    Im = X * np.sin(P)
    res_ = (Re + 1j * Im) / (X**2 + 2 * Re - 1) # real denominator
    print(res - res_)
    """
    return res


def zeta(x, y):
    assert x > 0, f"The currently used formula derivation assumed x > 0, but x = {x}"

    TERMS = 5 * 10**3

    Re = 0
    Im = 0
    for n in range(1, TERMS): # can be optimized
        A = (-1)**(n-1) / n**x
        Py = y * np.log(n)
        Re += A * np.cos(Py)
        Im += A * np.sin(Py)

    eta = Re - 1j * Im

    res = _zeta_over_eta_factor(x, y) * eta

    return res.real, res.imag


def test():
    TOL = 1e-3

    z1 = zeta(1.2, 3.7)
    z2 = zeta(0.8, -2.5) # wolframalpha Zeta[0.8-2.5I]
    z3 = zeta(2, 0)

    test1 = abs(z1[0] - 0.69205300) < TOL and abs(z1[1] - 0.000544329) < TOL
    test2 = abs(z2[0] - 0.56365400) < TOL and abs(z2[1] - 0.204752000) < TOL
    test3 = abs(z3[0] - np.pi**2/6) < TOL and abs(z3[1] - 0.000000000) < TOL

    msg = "test1: {}\ntest2: {}\ntest2: {}"
    print(msg.format(test1, test2, test3))



def plot(x):
    POINTS = 100
    PREC = 3

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_aspect(1)

    reals = []
    imags = []
    for t in range(POINTS):
        y = t / PREC
        z = zeta(x, y)
        reals.append(z[0])
        imags.append(z[1])

        if t % PREC == 0:
            label = "y={}".format(round(y, 1))
            ax.text(*z, label, fontsize=8)

    plt.plot(reals, imags, 'r')
    plt.xlabel('Re')
    plt.ylabel('Im')
    plt.grid(True)
    plt.show()


if __name__=='__main__':
    test()
    plot(4 / 5)

Coming from your nice video I have some advises for your code.

In general your are using numpy wrong. Using numpy functions like np.log is not realy fast on scalars like ints or floats. These functions are designed to be used on an numpy array, distributed on each element.

As example I've timed your zeta function on my pc and I got this results:

19.9 ms ± 330 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

That is not that bad, but given that this calculation is executed several thousand times, pretty slow.

Here is my numpy version of your zeta function:

def zeta(x, y, max_iter=5*10**3):
    
    ns = np.arange(1, max_iter)
    As = (-1)**(ns-1) / ns**x
    ps = y * np.log(ns)
    res = As * np.cos(ps)
    ims = As * np.sin(ps)
    
    re = res.sum()
    im = ims.sum()
    eta = re - 1j * im

    result = zeta_over_eta(x, y) * eta

    return result.real, result.imag

Essentially the logic didn't changed, the important thing was to put all n's into an numpy array and execute the functions on all elements. In the end res and ims are numpy arrays with the different values, to get to the meat I just summed them up.

The new time is this:
455 µs ± 4.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Magnitudes faster.

Now plotting 10000 points with a precision of 200 took roughly as long as your plotting in the video. And the spiral is pretty smooth.

Numpy is designed to do this pretty fast.

Also as recommendation, you should look into jupyter notebooks, then you can better live code and you have inline media such like markdown, latex, videos and plots of course. Runs in your browser, uses your local python instance. Pretty awesome.

Keep on with your work, the zeta video was interesting.

@KleinerNull: Cool, thx.
I've used Jupiter before, I mostly didn't want to get myself into some installation hole with those packets, but you're right that it could help with presentation in videos.

@Nikolaj-K

The installation should be just python3 -m pip install notebook.

python3 -m to pick the pip for the right version of python. Some operation systems have python2 and python3 installed and the regular pip is sometimes mapped to an unexpected version.

I have Anaconda and turns out I already have notebook.

Cool, if you have Anaconda, I also recommend a nicer plotting lib than matplotlib, HoloViews is pretty neat. Easy to use and pretty.

Here a sample plot function with a interactive mapping to change the zeta resolution:

import holoviews as hv
hv.extension('bokeh')

def plotv(x, POINTS=100, PREC=3):
    
    points = np.arange(POINTS)
    ys = points / PREC
    
    curve_map = {}
    
    for i in range(2, 5):
        applied_zeta = lambda y: zeta(x, y, 10**i)  # prefill zeta with x and max_iter values
        vzeta = np.vectorize(applied_zeta)  # stuff to make it usable on a numpy array
        reals, imags = vzeta(ys)
        curve_map[10**i] = hv.Curve((reals, imags)).opts(width=600, height=600)
        
    hmap = hv.HoloMap(curve_map, kdims='resolution')
    
    return hmap

(bokeh and holoviews need to be installed)

Looks like that your eta approach needs a ton of zeta values to get a more realistic spiral. Is there something that approaches faster? Is the prime product still valid for x=0.5?

it seems that zeta(0) should equal -0.5 while in your function, it equals -1.
Am I doing anything wrong?

@don-mao Line 21 says x > 0. I don't recall this years old script by head, but I imagine the series made use of doesn't behave well beyond this.