Visualization of discrete logarithms

Visdislog.py

# -*- coding:utf-8 -*-

import matplotlib
matplotlib.use('AGG')

import numpy as np
import sympy as sym
import matplotlib.pyplot as plt
import seaborn as sns

def discrete(p):
    d = dict()
    pr = sym.primitive_root(p)
    for i in range(0, p-1):
        d[(pr ** i) % p] = i
    x, y = [], []
    map(lambda v: (x.append(v[0]), y.append(v[1])), sorted(d.items()))
    return pr, x, y

def plot(b, l):
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')
    sns.despine()
    fig = plt.figure()
    primes = [p for p in sym.primerange(b, l + 1)]
    half = len(primes) / 2
    for i, p in enumerate(primes, 1):
        pr, x, y = discrete(p)
        ax = fig.add_subplot(half, half, i)
        ax.set_title(r'$p=' + str(p) + r',\ g=' + str(pr) + r'$')
        ax.set_ylabel(r'$f=Ind_{' + str(pr) + r'}(n)$')
        ax.set_xlabel(r'$n \equiv ' + str(pr) + r'^f\pmod{' + str(p) + r'}$')
        print "prime:", p, "primitive root:", pr
        ax.plot(x, y)
    fig.tight_layout()

if __name__ == '__main__':
    plot(3, 19)
    plt.savefig('fig.png')