STAT 540 — Programming Basics in Python

Building a Grading Machine

Topics: if / elif / else · vectorized conditionals · for loops · functions · map / comprehensions / apply

This notebook is the Python companion to the R lecture programming-basics_9_09_2026.R. Every R idea has a Python counterpart; where they differ, a R vs. Python note points it out.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats

np.set_printoptions(precision=4, suppress=True)

---

1. if / else

The reciprocal example

# R
a <- 0
if (a != 0) {
  print(1/a)
} else {
  print("No reciprocal for 0.")
}

Python uses indentation instead of braces, and a colon after the condition.

a = 0

if a != 0:
    print(1 / a)
else:
    print("No reciprocal for 0.")

No reciprocal for 0.

R vs. Python | R | Python | |---|---| | if (cond) { ... } else { ... } | if cond: … else: | | braces { } delimit blocks | indentation (4 spaces) delimits blocks | | !=, ==, >= | same | | &&, ||, ! | and, or, not |

The murders data

dslabs::murders is an R package dataset. In Python we read the same data with pandas.

# Reads the local copy if it is sitting next to this notebook,
# otherwise pulls it straight from the dslabs GitHub repo.
try:
    murders = pd.read_csv("murders.csv")
except FileNotFoundError:
    murders = pd.read_csv(
        "https://raw.githubusercontent.com/rafalab/dslabs/master/inst/extdata/murders.csv"
    )

murders.head()

	state	abb	region	population	total
0	Alabama	AL	South	4779736	135
1	Alaska	AK	West	710231	19
2	Arizona	AZ	West	6392017	232
3	Arkansas	AR	South	2915918	93
4	California	CA	West	37253956	1257

murders.iloc[:5]

	state	abb	region	population	total
0	Alabama	AL	South	4779736	135
1	Alaska	AK	West	710231	19
2	Arizona	AZ	West	6392017	232
3	Arkansas	AR	South	2915918	93
4	California	CA	West	37253956	1257

murders.iloc[:5, :3]

	state	abb	region
0	Alabama	AL	South
1	Alaska	AK	West
2	Arizona	AZ	West
3	Arkansas	AR	South
4	California	CA	West

# R: murder_rate <- murders$total / murders$population * 100000
murder_rate = murders["total"] / murders["population"] * 100000
murder_rate.head()

0    2.824424
1    2.675186
2    3.629527
3    3.189390
4    3.374138
dtype: float64

which.min() → argmin()

R vs. Python R indexes from 1; Python indexes from 0. which.min(x) returns a 1-based position, np.argmin(x) / Series.argmin() returns a 0-based one.

ind = murder_rate.argmin()          # 0-based position of the smallest rate
print("position:", ind)
print("state:", murders["state"][ind])
print("rate:", murder_rate[ind])

position: 45
state: Vermont
rate: 0.31962105727449536

if murder_rate[ind] < 0.5:
    print(murders["state"][ind])
else:
    print("No state has murder rate that low")

Vermont

if murder_rate[ind] < 0.25:
    print(murders["state"][ind])
else:
    print("No state has a murder rate that low.")

No state has a murder rate that low.

---

2. Exercise — Build a Grading Machine (Pass / Fail)

scores = [55, 82, 91, 40, 67, 78, 100, 62, 35, 88]
scores

[55, 82, 91, 40, 67, 78, 100, 62, 35, 88]

Task. Write an if/else statement that reports PASS (score ≥ 60) or FAIL for the second score.

Careful! The R version used scores[2]. In Python the second element is scores[1].

if scores[1] >= 60:        # scores[1] is 82 -- the 2nd student
    print("PASS")
else:
    print("FAIL")

PASS

---

3. Exercise 2 — Letter grades with elif

score >= 90  ->  A
score >= 80  ->  B
score >= 70  ->  C
otherwise    ->  F

R vs. Python R's else if becomes a single keyword in Python: elif.

if scores[8] >= 90:        # scores[8] is 35 -- the 9th student
    print("A")
elif scores[8] >= 80:
    print("B")
elif scores[8] >= 70:
    print("C")
else:
    print("F")

F

---

4. Vectorized conditionals: ifelse() → np.where()

R's ifelse() does two different jobs. Python splits them:

R	Python (one value)	Python (whole vector)
ifelse(a > 0, 1/a, NA)	1/a if a > 0 else np.nan	np.where(a > 0, 1/a, np.nan)

# --- single value: the conditional expression ---
a = 0
result = 1 / a if a > 0 else np.nan
result

nan

# --- whole array: np.where ---
a = np.array([0, 1, 2, -4, 5], dtype=float)

with np.errstate(divide="ignore"):      # silence the 1/0 warning
    result = np.where(a > 0, 1 / a, np.nan)

result

array([nan, 1. , 0.5, nan, 0.2])

b = np.array([0, 1, 1, 0])
ans = np.where(b == 0, "MALE", "FEMALE")
ans

array(['MALE', 'FEMALE', 'FEMALE', 'MALE'], dtype='<U6')

Replacing missing values

R: no_nas <- ifelse(is.na(na_example), 0, na_example)

In Python, missing numeric data is np.nan, tested with np.isnan() (or pd.isna()).

# a small stand-in for dslabs::na_example
na_example = np.array([2.0, np.nan, 3.0, 1.0, np.nan, 3.0, 2.0, np.nan, 5.0, 2.0])

print("missing before:", np.isnan(na_example).sum())

no_nas = np.where(np.isnan(na_example), 0, na_example)

print("missing after: ", np.isnan(no_nas).sum())
print(no_nas)

missing before: 3
missing after:  0
[2. 0. 3. 1. 0. 3. 2. 0. 5. 2.]

# pandas does the same thing with .fillna()
s = pd.Series(na_example)
s.fillna(0).tolist()

any() and all()

z = [True, True, False]

print(any(z))   # R: any(z)
print(all(z))   # R: all(z)

True
False

---

5. for loops

# R
for (i in 1:5) {
  print(i)
}

R vs. Python 1:5 in R includes 5. Python's range(1, 6) stops before 6.

for i in range(1, 6):
    print(i)

1
2
3
4
5

Building up a vector inside a loop

R:

m <- 25
s_n <- vector(length = m)
for (n in 1:m) {
  s_n[n] <- sum(1:n)
}

m = 25
s_n = np.zeros(m)              # an empty vector of the right size

for n in range(1, m + 1):
    s_n[n - 1] = sum(range(1, n + 1))   # careful: position n-1 holds the n-th answer

s_n

array([  1.,   3.,   6.,  10.,  15.,  21.,  28.,  36.,  45.,  55.,  66.,
        78.,  91., 105., 120., 136., 153., 171., 190., 210., 231., 253.,
       276., 300., 325.])

# Avoid the n-1 bookkeeping with enumerate()
s_n = np.zeros(m)
for i, n in enumerate(range(1, m + 1)):
    s_n[i] = np.arange(1, n + 1).sum()

s_n[:5]

m=25
for i, n in enumerate(range(1, m + 1)):
    print(i, n)

0 1
1 2
2 3
3 4
4 5
5 6
6 7
7 8
8 9
9 10
10 11
11 12
12 13
13 14
14 15
15 16
16 17
17 18
18 19
19 20
20 21
21 22
22 23
23 24
24 25

---

6. Exercise — a for loop over all the scores

Return the letter grade for every student.

scores = [55, 82, 91, 40, 67, 78, 100, 62, 35, 88]

for i in scores:
    if i >= 90:
        print("A")
    elif i >= 80:
        print("B")
    elif i >= 70:
        print("C")
    else:
        print("F")

F
B
A
F
F
C
A
F
F
B

---

7. Functions

# R
newfunction <- function(x, y){
  ans <- x^2 + y^2
  return(ans)
}

R vs. Python function(x, y) → def newfunction(x, y): · ^ → ** · return(ans) → return ans

def newfunction(x, y):
    ans = x**2 + y**2
    return ans


newfunction(3, 4)

25

def square_me(x):
    return x**2


x = np.arange(-3, 3.01, 0.01)

plt.figure(figsize=(5, 3.2))
plt.plot(x, square_me(x))
plt.xlabel("x"); plt.ylabel("square_me(x)")
plt.tight_layout()
plt.show()

Writing our own average

R returns the last expression automatically; Python needs an explicit return.

def avg(x):
    s = sum(x)
    n = len(x)
    return s / n          # R would return s/n implicitly -- Python will not


x = list(range(1, 101))
print(avg(x), np.mean(x))
print(np.isclose(np.mean(x), avg(x)))     # R: identical(mean(x), avg(x))

50.5 50.5
True

Variables inside a function are local

s = 3
print(avg(range(1, 11)))   # the s inside avg() is a different s
print(s)                   # ... the global s is untouched

5.5
3

Default arguments

R:

avg <- function(x, arithmetic = TRUE){
  n <- length(x)
  ifelse(arithmetic, sum(x)/n, prod(x)^(1/n))
}

def avg(x, arithmetic=True):
    x = np.asarray(x, dtype=float)
    n = len(x)
    return x.sum() / n if arithmetic else np.prod(x) ** (1 / n)


print("arithmetic mean:", avg([1, 2, 3, 4]))
print("geometric mean: ", avg([1, 2, 3, 4], arithmetic=False))

arithmetic mean: 2.5
geometric mean:  2.213363839400643

def compute_s_n(n):
    return sum(range(1, n + 1))


compute_s_n(100)

5050

n = np.arange(1, m + 1)

plt.figure(figsize=(5, 3.2))
plt.scatter(n, s_n, s=18)
plt.xlabel("n"); plt.ylabel("s_n")
plt.tight_layout()
plt.show()

Vectorization

Just like R, NumPy applies arithmetic element by element — no loop required.

x = np.arange(1, 11)
y = np.arange(1, 11)

print(np.sqrt(x))
print(x * y)

[1.     1.4142 1.7321 2.     2.2361 2.4495 2.6458 2.8284 3.     3.1623]
[  1   4   9  16  25  36  49  64  81 100]

Watch out: a plain Python list is not a vector. [1,2,3] * 2 repeats the list; np.array([1,2,3]) * 2 doubles each element.

print([1, 2, 3] * 2)
print(np.array([1, 2, 3]) * 2)

[1, 2, 3, 1, 2, 3]
[2 4 6]

---

8. Exercise — the grade_it function

def grade_it(grade):
    if grade >= 90:
        return "A"
    elif grade >= 80:
        return "B"
    elif grade >= 70:
        return "C"
    else:
        return "F"


grade_it(30)

'F'

---

9. Exercise — a for loop that uses grade_it

R:

letter_grade <- rep("NA", length(scores))
for (i in 1:length(scores)) {
  letter_grade[i] <- grade_it(scores[i])
}

print(scores)

[55, 82, 91, 40, 67, 78, 100, 62, 35, 88]

letter_grade = ["NA"] * len(scores)          # R: rep("NA", length(scores))

for i in range(len(scores)):
    letter_grade[i] = grade_it(scores[i])

letter_grade

['F', 'B', 'A', 'F', 'F', 'C', 'A', 'F', 'F', 'B']

# More Pythonic: append as you go, no placeholder needed
letter_grade = []
for score in scores:
    letter_grade.append(grade_it(score))

letter_grade

['F', 'B', 'A', 'F', 'F', 'C', 'A', 'F', 'F', 'B']

---

10. sapply() → map, comprehensions, and .apply()

R:

x <- 1:10
sapply(x, sqrt)

Python offers four natural replacements. Pick whichever reads best.

x = list(range(1, 11))

print([np.sqrt(i) for i in x])        # 1. list comprehension  (most common)
print(np.sqrt(x))                     # 3. vectorized ufunc    (fastest)
print(pd.Series(x).apply(np.sqrt).tolist())   # 4. pandas .apply()

[np.float64(1.0), np.float64(1.4142135623730951), np.float64(1.7320508075688772), np.float64(2.0), np.float64(2.23606797749979), np.float64(2.449489742783178), np.float64(2.6457513110645907), np.float64(2.8284271247461903), np.float64(3.0), np.float64(3.1622776601683795)]
[np.float64(1.0), np.float64(1.4142135623730951), np.float64(1.7320508075688772), np.float64(2.0), np.float64(2.23606797749979), np.float64(2.449489742783178), np.float64(2.6457513110645907), np.float64(2.8284271247461903), np.float64(3.0), np.float64(3.1622776601683795)]
[1.     1.4142 1.7321 2.     2.2361 2.4495 2.6458 2.8284 3.     3.1623]
[1.0, 1.4142135623730951, 1.7320508075688772, 2.0, 2.23606797749979, 2.449489742783178, 2.6457513110645907, 2.8284271247461903, 3.0, 3.1622776601683795]

print(list(map(np.sqrt, x)))          # 2. map()

[np.float64(1.0), np.float64(1.4142135623730951), np.float64(1.7320508075688772), np.float64(2.0), np.float64(2.23606797749979), np.float64(2.449489742783178), np.float64(2.6457513110645907), np.float64(2.8284271247461903), np.float64(3.0), np.float64(3.1622776601683795)]

Exercise — Build a grading machine with "sapply"

scores = [55, 82, 91, 40, 67, 78, 100, 62, 35, 88]

letter_grade = [grade_it(s) for s in scores]     # R: sapply(scores, grade_it)
letter_grade

['F', 'B', 'A', 'F', 'F', 'C', 'A', 'F', 'F', 'B']

# the map() version
list(map(grade_it, scores))

['F', 'B', 'A', 'F', 'F', 'C', 'A', 'F', 'F', 'B']

# the pandas version -- and a tidy gradebook to go with it
gradebook = pd.DataFrame({"score": scores})
gradebook["letter"] = gradebook["score"].apply(grade_it)
gradebook["pass"] = np.where(gradebook["score"] >= 60, "PASS", "FAIL")
gradebook

	score	letter	pass
0	55	F	FAIL
1	82	B	PASS
2	91	A	PASS
3	40	F	FAIL
4	67	F	PASS
5	78	C	PASS
6	100	A	PASS
7	62	F	PASS
8	35	F	FAIL
9	88	B	PASS

gradebook["letter"].value_counts()

letter
F    5
B    2
A    2
C    1
Name: count, dtype: int64

gradebook["letter"].value_counts().sort_index() ### sort_index: sort by index alphabetically,

letter
A    2
B    2
C    1
F    5
Name: count, dtype: int64

---

11. apply() over a matrix → axis=

R:

mat <- matrix(1:20, ncol = 2)
apply(mat, 2, mean)   # column means
apply(mat, 1, mean)   # row means

R vs. Python R fills a matrix column by column; NumPy's reshape fills row by row. Use order="F" to reproduce R's layout.

mat = np.arange(1, 21).reshape(10, 2, order="F")   # same as R's matrix(1:20, ncol=2)
mat

array([[ 1, 11],
       [ 2, 12],
       [ 3, 13],
       [ 4, 14],
       [ 5, 15],
       [ 6, 16],
       [ 7, 17],
       [ 8, 18],
       [ 9, 19],
       [10, 20]])

cMean = mat.mean(axis=0)     # R: apply(mat, 2, mean) -- collapse the rows -> column means
rMean = mat.mean(axis=1)     # R: apply(mat, 1, mean) -- collapse the cols -> row means

print("cMean:", cMean)
print("rMean:", rMean)

cMean: [ 5.5 15.5]
rMean: [ 6.  7.  8.  9. 10. 11. 12. 13. 14. 15.]

Remember the axis rule: axis=0 works down the rows (giving one number per column), axis=1 works across the columns (one number per row). That is the mirror image of R's apply(mat, 1, ...) / apply(mat, 2, ...) numbering.

# np.apply_along_axis() is the literal translation of apply(), and works
# for any function -- not just the ones NumPy has built in.
print(np.apply_along_axis(np.mean, 0, mat))
print(np.apply_along_axis(np.mean, 1, mat))

[ 5.5 15.5]
[ 6.  7.  8.  9. 10. 11. 12. 13. 14. 15.]

# pandas keeps R's 1/2 numbering readable with axis names
df = pd.DataFrame(mat, columns=["V1", "V2"])
print(df.apply(np.mean, axis="index").tolist())    # column means
print(df.apply(np.mean, axis="columns").tolist())  # row means

[5.5, 15.5]
[6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0]

---

12. Putting it together — a simulation

Task

1. Generate 30 random numbers from a normal distribution.
2. Randomly assign 15 to control and 15 to treatment.
3. Run a t-test and report the p-value.

rng = np.random.default_rng(123)                 # R: set.seed(123)

obs = rng.normal(loc=0, scale=1, size=30)        # R: rnorm(30, mean=0, sd=1)
obs[:5]

array([-0.9891, -0.3678,  1.2879,  0.194 ,  0.9202])

# R: sample(c("A","B","C"), 1)
rng.choice(["A", "B", "C"], size=1)

array(['B'], dtype='<U1')

⚠️ Not good!

control <- sample(obs, 15)
trt     <- sample(obs, 15)

Drawing twice independently lets the same observation land in both groups — and leaves some observations in neither. The groups are not a partition of the data.

control_bad = rng.choice(obs, 15, replace=False)
trt_bad = rng.choice(obs, 15, replace=False)

overlap = np.intersect1d(control_bad, trt_bad)
print("observations in BOTH groups:", len(overlap))     # should be 0, usually isn't

observations in BOTH groups: 8

✅ This is good!

Sample the indices once, then take the chosen ones for control and the rest for treatment.

R vs. Python R's negative index obs[-icontrol] means "everything except". In Python a negative index counts from the end, so use a boolean mask or np.setdiff1d() instead.

icontrol = rng.choice(30, size=15, replace=False)    # R: sample(1:30, 15)

mask = np.zeros(30, dtype=bool)
mask[icontrol] = True

control = obs[mask]        # R: obs[icontrol]
trt = obs[~mask]           # R: obs[-icontrol]   (~ means "not")

res = stats.ttest_ind(control, trt)
print(res)                                  # R: str(res)
print("p-value:", res.pvalue)               # R: res$p.value

TtestResult(statistic=np.float64(-0.8220123018732333), pvalue=np.float64(0.41801502605208213), df=np.float64(28.0))
p-value: 0.41801502605208213

Do it 1000 times — with a for loop

R:

sim <- 1000; n <- 15
pvalue <- numeric(sim)
for (i in 1:sim) { ... pvalue[i] <- res$p.value }

sim = 1000
n = 15
rng = np.random.default_rng(1)

pvalue = np.zeros(sim)                       # R: numeric(sim)

for i in range(sim):
    obs = rng.normal(size=2 * n)
    icontrol = rng.choice(2 * n, size=n, replace=False)
    mask = np.zeros(2 * n, dtype=bool)
    mask[icontrol] = True
    res = stats.ttest_ind(obs[mask], obs[~mask])
    pvalue[i] = res.pvalue

pvalue[:5]

array([0.2491, 0.0894, 0.9008, 0.2565, 0.0359])

The same thing as a function

Note that iter is never used inside the function — it is just a placeholder so that map/sapply has something to iterate over. This is a common pattern in simulations.

Naming note: iter is a built-in Python function, so we call the argument i instead.

def getpvalue(i, n, rng):
    obs = rng.normal(size=2 * n)
    icontrol = rng.choice(2 * n, size=n, replace=False)
    mask = np.zeros(2 * n, dtype=bool)
    mask[icontrol] = True
    res = stats.ttest_ind(obs[mask], obs[~mask])
    return res.pvalue

# R: pvalue_sapply <- sapply(1:1000, getpvalue, n = 15)
rng = np.random.default_rng(1)
pvalue_listcomp = np.array([getpvalue(i, n=15, rng=rng) for i in range(1000)])

# R's replicate() equivalent
rng = np.random.default_rng(1)
pvalue_map = np.array(list(map(lambda i: getpvalue(i, 15, rng), range(1000))))

# R: identical(pvalue_sapply, pvalue_replicate)
print(np.array_equal(pvalue_listcomp, pvalue_map))

True

print("mean p-value:      ", pvalue_listcomp.mean())
print("proportion < 0.05: ", (pvalue_listcomp < 0.05).mean(), " (should be about 0.05)")

plt.figure(figsize=(5, 3.2))
plt.hist(pvalue_listcomp, bins=20, edgecolor="white")
plt.xlabel("p-value"); plt.ylabel("frequency")
plt.title("1000 null simulations: p-values are uniform")
plt.tight_layout()
plt.show()

mean p-value:       0.5038851300705502
proportion < 0.05:  0.055  (should be about 0.05)

---

Cheat sheet

R	Python
if (x) { } else { }	if x: / else: (indentation, no braces)
else if	elif
ifelse(cond, a, b)	np.where(cond, a, b) · a if cond else b
is.na(x)	np.isnan(x) · pd.isna(x)
any(z) / all(z)	any(z) / all(z)
for (i in 1:5)	for i in range(1, 6):
length(x)	len(x)
x^2	x**2
function(x) { }	def f(x):
implicit last value returned	explicit return required
sapply(x, f)	[f(i) for i in x] · list(map(f, x)) · Series.apply(f)
apply(m, 2, mean)	m.mean(axis=0)
apply(m, 1, mean)	m.mean(axis=1)
matrix(1:20, ncol=2)	np.arange(1,21).reshape(10,2, order="F")
rnorm(n)	rng.normal(size=n)
sample(1:n, k)	rng.choice(n, size=k, replace=False)
t.test(a, b)$p.value	stats.ttest_ind(a, b).pvalue
set.seed(1)	np.random.default_rng(1)
x[-i] ("drop i")	boolean mask x[~mask]
1-based indexing	0-based indexing