STAT540: Computing in Statistics

Python Lecture 3: Defining Functions, Plotting with Matplotlib, and Conditional Programming"

---

Outline

1. Making plots with Matplotlib
2. Defining new functions
3. Conditional programming (if / elif / else)
4. Putting it together: piecewise functions
5. Practice

Making Plots: Getting Started

To make plots, we import matplotlib --- the most widely used Python plotting library (matplotlib.org).

x=np.linspace(0,4, 5)

x

array([0., 1., 2., 3., 4.])

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 2*np.pi, 361)
sinx = np.sin(x)
plt.plot(x, sinx)
plt.show()

• plt.plot(x, y) draws a line plot
• plt.show() displays the figure

sinx = np.sin(x)
cosx = np.cos(x)
plt.plot(x, sinx)
plt.plot(x, cosx)
plt.show()

Adding Features and a Legend

u = np.random.random(30) * 2*np.pi
v = 0.8*np.sin(u) + 0.2*np.cos(u) + np.random.normal(0, 1/3, u.size)

x  = np.linspace(0, 2*np.pi, 361)
fx = 0.8*np.sin(x) + 0.2*np.cos(x)

plt.figure()  # start a new, empty figure
plt.plot(u, v, 'bo', label='observations')
plt.plot(x, fx, label='true function')
plt.xlabel("x"); plt.ylabel("response")
plt.title("Demo plot")
plt.legend()
plt.show()

A `label=` on each plotted feature builds the legend automatically.

The Figure/Axes Approach

The more flexible --- and more common in tutorials --- way to build a plot: set up a figure and a set of axes first, then add to the axes.

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(u, v, 'bo', label='Observations')
ax.plot(x, fx, label='True function', linestyle="--")
ax.set_xlabel("x")
ax.set_ylabel("Response")
ax.set_title("Demo plot")
ax.legend()
plt.show()

Multi-Panel Figures

Setting up the figure first lets us build figures with several panels (each panel is an axes):

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

ax1.plot(u, v, 'bo', label='observations')
ax1.set_title("Demo plot"); ax1.legend()

ax2.plot(x, fx, 'g', label='true function')
ax2.legend()
plt.show()

A "mosaic" arrangement (plt.subplot_mosaic) allows even more flexible panel layouts, with panels referenced by name instead of position.

Outline

1. Making plots with Matplotlib
2. Defining new functions
3. Conditional programming (if / elif / else)
4. Putting it together: piecewise functions
5. Practice

Defining a Simple Function

A one-liner function:

def f(x): return x * np.sin(x)

f(1)   

0.8414709848078965

Key rule: A Python function must return a value explicitly. If you do not write a return statement, Python automatically returns a special value called None.

def print_lyrics():
    print("I'm a lumberjack, and I'm okay.")
    print("I sleep all night and I work all day.")

print_lyrics()

I'm a lumberjack, and I'm okay.
I sleep all night and I work all day.

Defining a Python Function

With more than one line, we indent the body --- no curly braces as in R. The function ends when the indentation stops.

def hyp(a, b):
    csq = a**2 + b**2
    c = np.sqrt(csq)
    return c

Returning Multiple Values

List every object to return, separated by commas:

def eucl(r, th):
    x = r * np.cos(th)
    y = r * np.sin(th)
    return x, y

x, y = eucl(2, np.pi/3)

print(x)
print(y)

1.0000000000000002
1.7320508075688772

Python automatically packages multiple return values into a tuple, which we can unpack on the left-hand side --- just as we did with multiple assignment earlier.

Default Argument Values

Default values work just as they do in R:

def logistic(x, a=0, b=1):
    l   = a + x*b
    val = np.exp(l) / (1 + np.exp(l))
    return val

x   = np.linspace(-4, 4, 200)
fx1 = logistic(x)               # uses defaults a=0, b=1
fx2 = logistic(x, a=-2, b=3)    # override the defaults

Arguments can be supplied positionally or by name (a=-2); naming them makes calls more readable.

Outline

1. Making plots with Matplotlib
2. Defining new functions
3. Conditional programming (if / elif / else)
4. Putting it together: piecewise functions
5. Practice

if / else

The logic of conditional programming is the same in every language --- only the syntax changes.

#| echo: true
#| eval: true
def rolldice():
    r1 = np.random.choice(np.arange(1, 7), 1)
    r2 = np.random.choice(np.arange(1, 7), 1)

    if r1 == r2:
        print("Doubles!")
    else:
        print("No match.")

    print(r1, r2)

Just as with functions and loops, the body of an if/else block is defined by indentation, not braces.

elif: "Else If"

Python spells "else if" as a single word, elif: \small

def rolldice():
    r1 = np.random.choice(np.arange(1, 7), 1)
    r2 = np.random.choice(np.arange(1, 7), 1)

    if r1 == r2:
        if r1 == 1:
            print("Snake-eyes!")
        elif r1 == 6:
            print("Can't beat it!")
        else:
            print("Doubles!")
    elif (r1 + r2) <= 7:
        print("Sad rollin'.")
    else:
        print("Mighty fine!")

Conditions can be nested just like in R.

Booleans as Numbers

The function bool() coerces its argument to True/False:

bool(1)   # True
bool(0)   # False

Just as in R, a logical value coerces to $0$ or $1$ in arithmetic. This lets us write piecewise-defined functions compactly, without an explicit if statement, using indicator-style expressions.

Outline

1. Making plots with Matplotlib
2. Defining new functions
3. Conditional programming (if / elif / else)
4. Putting it together: piecewise functions
5. Practice

Example: Soft-Thresholding

$$ S(x, \lambda) = \begin{cases} x + \lambda, & x < -\lambda \\ 0, & -\lambda \le x \le \lambda \\ x - \lambda, & \lambda < x \end{cases} \quad (\lambda > 0) $$

def softthresh(x, lam):
    return (x + lam)*(x < -lam) + (x - lam)*(x > lam)

x  = np.linspace(-2, 2, 200)
fx = softthresh(x, lam=1)

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(x, fx)
ax.axhline(0, linestyle="--", linewidth=0.5)
plt.show()

The boolean expressions (x < -lam) etc. act as $0/1$ indicators.

Result

#| echo: false
#| eval: true
def softthresh(x, lam):
    return (x + lam)*(x < -lam) + (x - lam)*(x > lam)

x  = np.linspace(-2, 2, 200)
fx = softthresh(x, lam=1)

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(x, fx)
ax.axhline(0, linestyle="--", linewidth=0.5)
plt.show()

Example: A Bivariate Density Surface

def bivn(z1, z2, rho=0):
    r = 1 - rho**2
    a = 1 / (2*np.pi*np.sqrt(r))
    b = (z1**2 - 2*rho*z1*z2 + z2**2) / r
    return a * np.exp(-b/2)

gs = 50
z       = np.linspace(-3, 3, gs)
X, Y    = np.meshgrid(z, z)       # build a 2-d grid from 1-d grids
Z       = bivn(X, Y, rho=-0.5)   # evaluate the function on the grid

fig = plt.figure(figsize=(8, 8))
ax  = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, Z, cmap=plt.cm.YlGnBu_r)
plt.show()

Result

def bivn(z1, z2, rho=0):
    r = 1 - rho**2
    a = 1 / (2*np.pi*np.sqrt(r))
    b = (z1**2 - 2*rho*z1*z2 + z2**2) / r
    return a * np.exp(-b/2)

gs = 50
z       = np.linspace(-3, 3, gs)
X, Y    = np.meshgrid(z, z)       # build a 2-d grid from 1-d grids
Z       = bivn(X, Y, rho=-0.5)   # evaluate the function on the grid

fig = plt.figure(figsize=(8, 8))
ax  = fig.add_subplot(projection='3d')
ax.plot_surface(X, Y, Z, cmap=plt.cm.YlGnBu_r)
plt.show()

Outline

1. Making plots with Matplotlib
2. Defining new functions
3. Conditional programming (if / elif / else)
4. Putting it together: piecewise functions
5. Practice

Practice: Write a function for Bootstrap Median

\scriptsize

def bootstrap_ci(x, n_boot=2000, seed=42, alpha=0.05):
    rng = np.random.default_rng(seed)     
    x = np.asarray(x, dtype=float)
    n = x.size
    obs_median = float(np.median(x))
    boot_meds = np.empty(n_boot)          # pre-allocate storage

    for b in range(n_boot):
        idx = rng.integers(0, n, size=n)  # 0-based, high exclusive
        x_star = x[idx]
        boot_meds[b] = np.median(x_star)

    lo, hi = np.quantile(boot_meds, [alpha/2, 1 - alpha/2])
    return {"estimate": obs_median, "ci_lower": float(lo),
            "ci_upper": float(hi), "boot_dist": boot_meds}
bmi = np.round(np.random.default_rng(1).gamma(30.0, 1.0, size=5000), 1)
res = bootstrap_ci(bmi, n_boot=2000, seed=42)
res

{'estimate': 29.6,
 'ci_lower': 29.4,
 'ci_upper': 29.8,
 'boot_dist': array([29.5, 29.9, 29.7, ..., 29.8, 29.5, 29.5])}

Practice: Read Code

Predict the output of the following code before running it. \small

#| echo: true
#| eval: false
def star(p, s):
    th = np.linspace(np.pi/2, np.pi/2 + 2*np.pi*s, s*p + 1)
    x  = np.cos(th[::s])
    y  = np.sin(th[::s])
    return x, y

x, y = star(5, 2)

fig, ax = plt.subplots(figsize=(4, 4))
ax.set_aspect('equal')
ax.axis('off')
ax.plot(x, y)
plt.show()

p=5
s=2
th = np.linspace(np.pi/2, np.pi/2 + 2*np.pi*s, s*p + 1)
print(th)
th.size
th[::s]

[ 1.57079633  2.82743339  4.08407045  5.34070751  6.59734457  7.85398163
  9.1106187  10.36725576 11.62389282 12.88052988 14.13716694]

array([ 1.57079633,  4.08407045,  6.59734457,  9.1106187 , 11.62389282,
       14.13716694])

Result

def star(p, s):
    th = np.linspace(np.pi/2, np.pi/2 + 2*np.pi*s, s*p + 1)
    x  = np.cos(th[::s])
    y  = np.sin(th[::s])
    return x, y

x, y = star(5, 2)

fig, ax = plt.subplots(figsize=(4, 4))
ax.set_aspect('equal')
ax.axis('off')
ax.plot(x, y)
plt.show()

%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch

GARNET, BLUE = "#73000A", "#1B5E9E"
INK, MUTED, FAINT = "#1A1A1A", "#9AA0A6", "#DDE1E5"

plt.rcParams.update({
    "font.family": "DejaVu Sans",
    "font.size": 9,
    "axes.titlesize": 10,
    "figure.facecolor": "white",
    "axes.facecolor": "white",
})


def angles(p, s):
    """The angle array the original star() builds, before and after th[::s]."""
    th = np.linspace(np.pi / 2, np.pi / 2 + 2 * np.pi * s, s * p + 1)
    return th, th[::s]


def frame(ax, lim=1.42):
    ax.set_aspect("equal")
    ax.set_xlim(-lim, lim)
    ax.set_ylim(-lim, lim)
    ax.axis("off")


def guide_circle(ax, r=1.0):
    t = np.linspace(0, 2 * np.pi, 400)
    ax.plot(r * np.cos(t), r * np.sin(t), color=FAINT, lw=1, zorder=0)


def caption(ax, text, y=-1.33):
    ax.text(0, y, text, ha="center", va="top", fontsize=8.5, color=MUTED)

th_all, th_sub = angles(5, 2)

deg = np.degrees(th_all) - 90          # measured from the starting angle
print("all 11 angles (deg from start):", np.round(deg, 1))
print("spacing:", round(deg[1] - deg[0], 1), "deg")
print()

deg_sub = np.degrees(th_sub) - 90
print("after th[::2]:", np.round(deg_sub, 1))
print("spacing:", round(deg_sub[1] - deg_sub[0], 1), "deg")
print()

print("vertex index visited at each step:",
      [int(round(d / 72)) % 5 for d in deg_sub])

all 11 angles (deg from start): [  0.  72. 144. 216. 288. 360. 432. 504. 576. 648. 720.]
spacing: 72.0 deg

after th[::2]: [  0. 144. 288. 432. 576. 720.]
spacing: 144.0 deg

vertex index visited at each step: [0, 2, 4, 1, 3, 0]

fig = plt.figure(figsize=(14, 4.6))
gs = fig.add_gridspec(1, 4, wspace=0.12, top=0.86, bottom=0.12)

th_all, th_sub = angles(5, 2)
r_ramp = np.linspace(0.70, 1.12, len(th_all))
xs, ys = r_ramp * np.cos(th_all), r_ramp * np.sin(th_all)

# --- 1. the raw angles ---------------------------------------------------
ax = fig.add_subplot(gs[0]); frame(ax); guide_circle(ax)
ax.plot(xs, ys, color=BLUE, lw=1.6, alpha=0.65, zorder=1)
ax.scatter(xs, ys, s=46, color=BLUE, zorder=3, edgecolor="white", linewidth=1.2)
for i in range(len(th_all)):
    rr = r_ramp[i] + 0.13
    ax.text(rr*np.cos(th_all[i]), rr*np.sin(th_all[i]), str(i),
            ha="center", va="center", fontsize=8, color=BLUE)
ax.set_title("1.  11 angles, 72° apart", color=INK, pad=12)
caption(ax, "Two full laps (2π·s with s=2).")

# --- 2. the subsample ----------------------------------------------------
ax = fig.add_subplot(gs[1]); frame(ax); guide_circle(ax)
keep = np.zeros(len(th_all), dtype=bool); keep[::2] = True
ax.plot(xs, ys, color=FAINT, lw=1.4, zorder=1)
ax.scatter(xs[~keep], ys[~keep], s=42, facecolor="white",
           edgecolor=MUTED, linewidth=1.4, zorder=3)
ax.scatter(xs[keep], ys[keep], s=52, color=GARNET, zorder=4,
           edgecolor="white", linewidth=1.2)
for i in range(len(th_all)):
    rr = r_ramp[i] + 0.13
    ax.text(rr*np.cos(th_all[i]), rr*np.sin(th_all[i]), str(i),
            ha="center", va="center", fontsize=8,
            color=GARNET if keep[i] else MUTED,
            fontweight="bold" if keep[i] else "normal")
ax.set_title("2.  th[::2] keeps every other one", color=INK, pad=12)
caption(ax, "Six survive, now 2 × 72° = 144° apart.")

# --- 3. the visit order --------------------------------------------------
ax = fig.add_subplot(gs[2]); frame(ax); guide_circle(ax)
vx, vy = np.cos(th_sub), np.sin(th_sub)
for k in range(5):
    ax.add_patch(FancyArrowPatch((vx[k], vy[k]), (vx[k+1], vy[k+1]),
                 arrowstyle="-|>", mutation_scale=13, lw=1.8, color=GARNET,
                 shrinkA=9, shrinkB=11, zorder=2, alpha=0.45 + 0.11*k))
ax.scatter(vx[:5], vy[:5], s=58, color=INK, zorder=4,
           edgecolor="white", linewidth=1.2)
for step, vert in enumerate([0, 2, 4, 1, 3]):
    a = np.pi/2 + vert * 2*np.pi/5
    ax.text(1.24*np.cos(a), 1.24*np.sin(a), f"v{vert}", ha="center",
            va="center", fontsize=9, color=INK, fontweight="bold")
    ax.text(0.44*np.cos(a), 0.44*np.sin(a), f"step {step}", ha="center",
            va="center", fontsize=7, color=GARNET, zorder=5,
            bbox=dict(boxstyle="round,pad=0.22", facecolor="white",
                      edgecolor=FAINT, linewidth=0.8))
ax.set_title("3.  Each 144° step skips a vertex", color=INK, pad=12)
caption(ax, "v0 → v2 → v4 → v1 → v3 → v0.")

# --- 4. the result -------------------------------------------------------
ax = fig.add_subplot(gs[3]); frame(ax); guide_circle(ax)
ax.plot(vx, vy, color=GARNET, lw=2.2, zorder=2)
ax.scatter(vx[:5], vy[:5], s=46, color=GARNET, zorder=3,
           edgecolor="white", linewidth=1.2)
ax.set_title("4.  ax.plot joins them — a pentagram", color=INK, pad=12)
caption(ax, "5 × 144° = 720°, so the path closes.")

plt.show()

x, y = star(7, 3)

fig, ax = plt.subplots(figsize=(4, 4))
ax.set_aspect('equal')
ax.axis('off')
ax.plot(x, y)
plt.show()

import matplotlib.pyplot as plt
import numpy as np

# 1. Generate data
x = np.linspace(0, 10, 100)
y = np.sin(x)

# 2. Create the figure and axes objects
fig, ax = plt.subplots()

# 3. Plot data on the specific axes
ax.plot(x, y, label="Sine Wave")

# 4. Customize the specific axes
ax.set_title("Simple Line Plot")
ax.set_xlabel("X Axis")
ax.set_ylabel("Y Axis")
ax.legend()

# 5. Display the window
plt.show()