Python: Visualising data with Matplotlib (Kronig-Penney model)

Context

NOTE: This post assumes a bit of mathematical familiarity.

NOTE2: MathML is required for this page to render correctly.

This blog post is a follow up on my entries on learning python. In the previous two posts (1, 2) I wrote a small Python script with a C extension to compute vector cross products, and then explored how to distribute the resulting shared object.

The objective here is to write a script that can leverage the matplotlib extension of python to visualize the results outputted by a python script.

The chosen case study here is the well-known Kronig-Penney model. We’ll write a simple program to do the calculation of the model’s band structures, and plot the result using matplotlib.

Quick recap on the Kronig-Penney model

The Kronig–Penney model is a one dimensional lattice of square wells/barriers, which makes it a periodic medium. Bloch's idea is that a wave doesn't need to repeat exactly every cell, it only needs to come back up to a phase. Floquet's theory is more abstract: it concerns the structure of the space of solutions, and the properties of solutions, of a linear system of differential equations with periodic coefficients.

The Kronig–Penney model neatly illustrates where these meet: a concrete periodic Schrödinger problem that exhibits Bloch/Floquet behaviour in a setting simple enough to compute and plot.

To begin with, we must state the Schrodinger Equation as a form of a Hill's Differential Equation, which is an ordinary second-order differential equation with a periodic function, instead of a constant coefficient.

\begin{matrix} - \frac{ℏ^{2}}{2 m} ψ^{''} (x) + [V (x) - E] ψ (x) = 0 & V (x) = V (x + b) \end{matrix}

\begin{matrix} ψ^{''} (x) - \frac{2 m}{ℏ^{2}} [V (x) - E] ψ (x) = 0 & \to ψ^{''} (x) + \frac{2 m}{ℏ^{2}} [E - V (x)] ψ (x) = 0 \end{matrix}

If we let:

f (x) = \frac{2 m}{ℏ^{2}} [E - V (x)]

Then that also implies that f(x) = f(x + b) which nets us with Hill's Differential Equation:

ψ^{''} (x) + f (x) ψ (x) = 0 f (x) = f (x + b)

Floquet theory allows us to represent Hill's Differential Equation as a matrix equation. In a nutshell, instead of tackling the differential equation as is, we're getting a matrix that represents the evolution of the solution over one period of the coefficients. The eigenvalues of this matrix are called the characteristic multiplier.

The magnitude of the characteristic multiplier determines the stability of the system's periodic solutions. If we take the logarithm of the characteristic multiplier, we can observe the growth or decay of the solutions.

The system formulated as a matrix equation:

\begin{matrix} {\hat{Ψ}}^{'} (x) & = [\begin{matrix} 0 & 1 \\ - f (x) & 0 \end{matrix}] [\begin{matrix} ψ_{1} (x) & ψ_{2} (x) \\ ψ_{1}^{'} (x) & ψ_{2}^{'} (x) \end{matrix}] = \hat{A} (x) \hat{Ψ} (x) \end{matrix}

Will have the following property:

\begin{matrix} \vec{ψ} (x + b) & = \tilde{λ} \vec{ψ} (x) & x \in R & \tilde{λ} \in C \end{matrix}

Where lambda is the characteristic multiplier. The solutions are two linearly independent solutions ψ_1,2(x) with the following properties:

[\begin{matrix} ψ_{1} (x + b) \\ ψ_{2} (x + b) \end{matrix}] = [\begin{matrix} {\tilde{α}}_{1} & {\tilde{α}}_{2} \\ {\tilde{β}}_{1} & {\tilde{β}}_{2} \end{matrix}] [\begin{matrix} ψ_{1} (x) \\ ψ_{2} (x) \end{matrix}]

[\begin{matrix} ψ_{1}^{'} (x + b) \\ ψ_{2}^{'} (x + b) \end{matrix}] = [\begin{matrix} {\tilde{α}}_{1} & {\tilde{α}}_{2} \\ {\tilde{β}}_{1} & {\tilde{β}}_{2} \end{matrix}] [\begin{matrix} ψ_{1}^{'} (x) \\ ψ_{2}^{'} (x) \end{matrix}]

It is considered that ψ₁(x) and ψ₂(x) cannot be proportional to each other, as it would rule out the possibility that they can be used to build a general solution:

[\hat{Ψ} (x) - \tilde{λ} \hat{I}] [\begin{matrix} \tilde{A} \\ \tilde{B} \end{matrix}] = [\begin{matrix} ({\tilde{α}}_{1} - \tilde{λ}) & {\tilde{β}}_{1} \\ {\tilde{α}}_{2} & ({\tilde{β}}_{2} - \tilde{λ}) \end{matrix}] [\begin{matrix} \tilde{A} \\ \tilde{B} \end{matrix}] = [\begin{matrix} 0 \\ 0 \end{matrix}]

Solving for linear independent solutions arrives at:

\begin{matrix} {\tilde{λ}}^{2} - ({\tilde{α}}_{1} + {\tilde{β}}_{2}) \tilde{λ} + 1 = 0 & {\tilde{λ}}_{1} {\tilde{λ}}_{2} = 1 \end{matrix}

The result can be further simplified to give the equation for a discriminant Δ(λ):

Δ (\tilde{λ}) = {\tilde{α}}_{1} + {\tilde{β}}_{2} = {\tilde{λ}}_{1} + {\tilde{λ}}_{2}

Where by Floquet Theory, we know that λ has the form e^±ikb.
There are three possible outcomes concerning the discriminant Δ(λ):

-> |Δ(λ)| > 2: unstable solution, meaning that the characteristic exponent itself is a complex number, or in Bloch terms, a forbidden band
-> |Δ(λ)| < 2: stable solution, meaning that the characteristic exponent is a real number, or in Bloch terms, the bloch functions are bounded oscillating functions, an allowed band
-> |Δ(λ)| = ±2: inconclusive, meaning that this is a boundary between regions of stability and instability

This yields a quadratic equation with two solutions for k(E) equal in magnitude and opposite in direction. It is something we can easily implement as code.

Implenting the calculation

The main calculation, kp.c:

#include <stdio.h>
#include <math.h>
#include <complex.h>
#include "kp.h"

void solve(const double E, double complex k[2], double b, double V_0, double Delta) {
    double complex q = csqrt(2.0 * E);
    double complex kappa = csqrt(2.0 * (E - V_0));

    double complex psi_11, psi_12, psi_21, psi_22;

    psi_11 = cexp(I * q * (b - Delta)) / (4.0 * q * kappa)
        * ( cexp(I * kappa * Delta)   * (q + kappa) * (q + kappa)
          - cexp(-I * kappa * Delta)  * (q - kappa) * (q - kappa) );

    psi_22 = conj(psi_11);

    psi_12 = -I * cexp(I * q * (b - Delta)) / (2.0 * q * kappa)
        * (q * q - kappa * kappa) * csin(kappa * Delta);

    psi_21 = conj(psi_12);

    // quadratic determinantal eqn
    double complex u = -(psi_11 + psi_22);
    double complex v = (psi_11 * psi_22 - psi_12 * psi_21);
    double complex lam[2];

    lam[0] = (-u + csqrt(u * u - 4.0 * v)) / 2.0;
    lam[1] = (-u - csqrt(u * u - 4.0 * v)) / 2.0;

    for (int x = 0; x < 2; x++) {
        k[x] = clog(lam[x]) / (I * b);
    }
}

void gen_data(double b, double V_0, double Delta) {
    const double pi = 4.0 * atan(1.0);
    const double dE = 0.01;
    double E = dE;
    const int N = 3000;

    FILE *fp = fopen("data.csv", "w");
    if (!fp) {
        fp = fopen("/dev/null", "w");
    } 

    for (int y = 0; y < N; y++) {
        double complex q[2];
        solve(E, q, b, V_0, Delta);

        double rq = creal(q[0]);

        if (rq > 0 && rq < pi / b) {
            fprintf(fp, "%.15g,%.15g\n", rq, E);
            fprintf(fp, "%.15g,%.15g\n", -rq, E);
        }

        rq = creal(q[1]);

        if (rq > 0 && rq < pi / b) {
            fprintf(fp, "%.15g,%.15g\n", rq, E);
            fprintf(fp, "%.15g,%.15g\n", -rq, E);
        }

        E += dE;
    }

    fclose(fp);
}

The header file kp.h is:

#include <complex.h>

void gen_data(double b, double V_0, double Delta);

The python binding needed to integrate this into our future python script:

#include <complex.h>
#include "kp.h"

static PyObject* py_gen_data(PyObject* self, PyObject* args, PyObject* kwargs) {
    static char* kwlist[] = {"b", "V_0", "Delta", NULL};
    double b = 1.0, V_0 = -5.0, Delta = 0.2;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|ddd", kwlist, &b, &V_0, &Delta)) {
        return NULL;
    }

    gen_data(b, V_0, Delta);

    Py_RETURN_NONE;
}

static PyMethodDef KPMethods[] = {
    {"gen_data", (PyCFunction)py_gen_data, METH_VARARGS | METH_KEYWORDS,
     "gen_data(b=1.0, V_0=-5.0, Delta=0.2) -> None\n"
     "Generates ./data.csv using hardcoded path; silent on I/O failure."},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef kp_module = {
    PyModuleDef_HEAD_INIT,
    "kp",
    "Kronig–Penney data generator",
    -1,
    KPMethods
};

PyMODINIT_FUNC PyInit_kp(void) {
    return PyModule_Create(&kp_module);
}

The full project files needed to build this yourself are available on my git repo.

On my machine, creating the shared object file for python boils down to:

gcc -fPIC -shared $(python3-config --includes) kp.c kp_py.c -o kp$(python3-config --extension-suffix)

Assembling the python script

In order to create a plot from the data generated by the simple kronig-penney calculator above, we're going to need to install two python extensions from the package manager:

sudo apt install python3-numpy python3-matplotlib

NumPy

Numpy is a mathematical library used for working with arrays. This essentially allows us to turn the following:

import csv, math
k, E = [], []
with open("data.csv") as f:
    for row in csv.reader(f):
        if not row: 
            continue
        kk, ee = map(float, row)
        k.append(kk); E.append(ee)

into a one-liner like:

import numpy as np
k, E = np.loadtxt("data.csv", delimiter=",", unpack=True)

The result is an array, rather than a python list, which is what matplotlib works with under the hood anyway.

Matplotlib

This is a visualisation library for creating various types of 2D and 3D plots from sets of data. The goal here is to turn the data file generated by the kronig-penney calculator into a 2D plot of the band structure.

Running the following code

import kp
kp.gen_data(b=1.0, V_0=-5.0, Delta=0.2);
import numpy as np, matplotlib.pyplot as plt
k, E = np.loadtxt("data.csv", delimiter=",", unpack=True)
plt.style.use('dark_background')
plt.scatter(k, E, s=2, color='red')
plt.xlabel("k"); plt.ylabel("E"); plt.xlim(-np.pi, np.pi); plt.tight_layout();
plt.show()

should result in a nice scatter plot.

The matplotlib options are fairly self-explanatory:

-> style.use() is a colour scheme setting, my blog has a dark background, so it makes sense to align a plot to be the same way.
-> scatter() This is where the magic happens, it creates a scatter plot of y against x, where s is a parameter for the size of each marker of the scatter.
-> xlabel() and ylabel() will print the given text onto the relevant axis.
-> xlim() allows to manually set how wide the scale is allowed to get on the x axis.
-> show() creates the plot into a new window.

The plot looks the following way:

plot of the Kronig-Penney model band structure

Conclusion

It is quite nice to see how little it really takes to visualise the Kronig–Penney band structure reproducibly. The C code computes the per-cell multipliers and maps them to Bloch wavenumbers k, which gets output as a CSV for Python to load and render as a figure. Tweak b, delta (well width), or V_0 and re-run the script to see how the band structure shifts.