ndivisors.py

import numpy as np


def isqrt(n):
    if n <= 1:
        return n

    def f(x):
        return (x + (n // x)) // 2
    x = n // 2
    while True:
        x1 = f(x)

        if abs(x1 - x) <= 1 or x <= x1:
            x = x1
            break
        x = x1

    x -= 1 if x * x > n else 0
    return x


def get_primes(B):
    is_prime = np.ones(B + 1, dtype=np.bool)
    is_prime[:2] = False
    for p in range(isqrt(B) + 1):
        if not is_prime[p]:
            continue
        is_prime[2 * p::p] = False

    for p in range(B + 1):
        if is_prime[p]:
            yield p


def datashader_plot(fn, seqx, seqy, w=2**11):
    import datashader as ds
    import pandas as pd
    import datashader.utils as utils
    from colorcet import fire
    from datashader import transfer_functions as tf
    df = pd.DataFrame({'x': pd.Series(seqx), 'y': pd.Series(seqy)})
    cvs = ds.Canvas(plot_width=w, plot_height=w)
    agg = cvs.points(df, 'x', 'y')
    img = tf.set_background(tf.shade(agg, cmap=fire), 'black')
    utils.export_image(img=img, filename=fn)


def main():
    B = 2**24
    counts = np.ones(B + 1, dtype=np.int64)
    q = 1
    while q <= B:
        q *= 2
        q2 = q * 2
        k1 = q2 - 1
        counts[::q2] += 1
        counts[k1::q2] += 1

    p_counts = np.empty(B + 1, dtype=np.int64)
    for p in get_primes(B):
        if p == 2:
            continue

        print(p, B)

        for k in [0, -1]:
            p_counts[(k % p)::p] = 1
            q = 1
            while q <= B:
                q *= p
                k1 = k % q
                p_counts[k1::q] += 1

            counts[(k % p)::p] *= p_counts[(k % p)::p]

    xs = np.arange(len(counts))

    datashader_plot('ndivisors', xs[1:], xs[1:].astype(np.float64) / counts[1:], w=2**11)


if __name__ == '__main__':
    main()