Simulation that explores Fermi Dirac statistics

fermi.py

from pylab import *
from scipy import diff
from collections import Counter
from sympy.combinatorics.partitions import IntegerPartition
from scipy.optimize import curve_fit

rc('text', usetex=True)
rc('font', family='serif', serif='cm')

def fermi_dirac(E, μ, T):
    return 1/(exp((E - μ)/T) + 1)

def simulate(q):
    # The individual particle energies are all possible
    # integer partitions of q
    p = IntegerPartition(q, [q])
    sols = []

    while True:
        # energy of kth particle is a[k]
        a = array(p.partition)
        a = concatenate([a, zeros(q - len(a))])
        # adjust by initial coordinates to get final coordinates
        sols.append(a + arange(0, -q, -1))
        p = p.next_lex()
        if p.partition == [q]:
            break

    # Count frequencies of resulting coordinates
    c = Counter(flatten(sols))
    x, n = array(list(c.items())).T
    p = n / len(sols)

    return x, p, sols

fig = figure()
ax1 = fig.add_subplot(2,1,1)

# Plot and fit for q = 20
x, p, samples = simulate(20)
ax1.set_title('Distribution $(q = 20)$')
ax1.plot(x, p,"ro")
ylabel(r"$\langle n\rangle$")

popt, pcov = curve_fit(fermi_dirac, x, p)

μ = popt[0]
T = popt[1]

xs = linspace(-20, 20, 1000)
ax1.plot(xs, fermi_dirac(xs, μ, T), 'b-')

ax1.text(5, 0.8, '$\\mu = {:.5f}$'.format(μ))
ax1.text(5, 0.65, '$T = {:.5f}$'.format(T))

# Calculation and plot of entropy
qs = arange(1, 25)
xs = linspace(0, 25, 100)
Sb = []

for q in qs:
    x, p, samples = simulate(q)
    Sb.append(log(len(samples)))

ax2 = fig.add_subplot(2,1,2)
ax2.plot(qs, Sb, 'bo')
ax2.set_title('Entropy')
xlabel("$q$")
ylabel("$S$")

m = diff(Sb)[19]
ax2.plot(20, Sb[19], 'go')
ax2.plot(xs, xs * m + Sb[19] - 20 * m, 'r-')
ax2.text(15, 3.5, r'$\frac{\partial S}{\partial E} = 1/T$')
ax2.text(15, 2.5, r'$T = {:.5f}$'.format(1/m))
xlim(0, 25)

tight_layout()
savefig("fermi-dirac.pdf")
clf()

# Display the energy levels
for i in [4, 5, 6]:
    title('$q = {}$'.format(i))
    x, p, samples = simulate(i)
    for k, s in enumerate(samples):
        plot([k + 1] * len(s), s, 'ro', markersize=15)
    grid()
    ax = gca()
    ax.xaxis.set_major_locator(MultipleLocator(1.0))
    ax.yaxis.set_major_locator(MultipleLocator(1.0))
    xlim(0, len(samples)+1)
    ylim(-i, i+1)
    tight_layout()
    savefig('samples_{}.pdf'.format(i))
    clf()

This seems interesting.