Code cleanup example

Original code, based on code from the free E-book Computational Physics: Problem Solving with Python, 3rd edition, by Landau, Páez, and Bordeianu. Modified primarily to use matplotlib instead of VPython.

from numpy import *
import matplotlib.pyplot as plt

Xmax = 40.
Xmin = 0.25
step = 0.1                                 # Global class variables
order = 10; start = 50       # Plot j_order
graph1, ax1 = plt.subplots(figsize = (5, 5))
ax1.set_title('Sperical Bessel, \
               L = 1 (red), 10')
ax1.set_xlabel("x"); ax1.set_ylabel('j(x)')
ax1.set_xlim(left=Xmin , \
             right=Xmax)
ax1.set_ylim(bottom=-0.2 , top=0.5)


def down(x, n, m):                   # Method down, recurs downward
    j = zeros( (start   +   2), float)
    j[m   +   1] = j[m] = 1.                  # Start with anything
    for k in range(m, 0,  -1):
        j[k - 1] = ( (2.*k + 1.)/x)*j[k] - j[k + 1]
    scale = (sin(x)/x)/j[0]          # Scale solution to known j[0]
    return j[n] * scale


for x in arange(Xmin, Xmax, step):
    ax1.plot(x, down(x, order, start), "r.")

for x in arange(Xmin, Xmax, step):
    ax1.plot(x, down(x,1,start), "g.")

plt.show()

Try to clean it up. Two possible solutions: one closer to the original (but fixing some bugs, like the plot title color being wrong!):

Show code cell content Hide code cell content

import numpy as np
import matplotlib.pyplot as plt

x_max = 40.0
x_min = 0.25
step = 0.1
start = 50  # Plot j_order


def down(x, order, start):
    """
    Method down, recurse downward.
    """
    j = np.zeros((start + 2), dtype=float)

    # Start with anything
    j[start + 1] = j[start] = 1.0

    for k in range(start, 0, -1):
        j[k - 1] = ((2.0 * k + 1.0) / x) * j[k] - j[k + 1]

    # Scale solution to known j[0]
    scale = (np.sin(x) / x) / j[0]
    return j[order] * scale


fig, ax = plt.subplots(figsize=(5, 5))
ax.set_title("Spherical Bessel, L = 1, 10 (red)")
ax.set_xlabel("x")
ax.set_ylabel("j(x)")
ax.set_xlim(x_min, x_max)
ax.set_ylim(-0.2, 0.5)

x = np.arange(x_min, x_max, step)

# Warning! red/green are bad colors to use, default colors are better
ax.plot(x, [down(x_i, 10, start) for x_i in x], "r.")
ax.plot(x, np.vectorize(down)(x, 1, start), "g.")

plt.show()

And one going a bit further and vectorizing the function properly, and fixing the color display:

Show code cell content Hide code cell content

import numpy as np
import matplotlib.pyplot as plt

x_min = 0.25
x_max = 40.0
step = 0.1


def down(x, order, start):
    """
    Method down, recurse downward.
    """
    j = np.zeros((start + 2, len(x)), dtype=float)

    # Start with anything
    j[start + 1] = j[start] = 1.0

    for k in range(start, 0, -1):
        j[k - 1] = ((2.0 * k + 1.0) / x) * j[k] - j[k + 1]

    # Scale solution to known j[0]
    scale = np.sin(x) / (x * j[0])
    return j[order] * scale


fig, ax = plt.subplots(figsize=(5, 5))
ax.set_title("Spherical Bessel")
ax.set_xlabel("x")
ax.set_ylabel("j(x)")
ax.set_xlim(x_min, x_max)
ax.set_ylim(-0.2, 0.5)

x = np.arange(x_min, x_max, step)

ax.plot(x, down(x, 10, 50), "r.", label="L=10")
ax.plot(x, down(x, 1, 50), "g.", label="L=1")
ax.legend()

plt.show()