# STAT3530 -- choosing a nonlinear specification

library(splines)

site <- 'https://tdruiz-stat3530.share.connect.posit.cloud/_data/'

# 88 runs of a single-cylinder test engine. NOx is nitrogen oxide emissions in
# micrograms per joule. E is the equivalence ratio: the fuel-to-air ratio,
# scaled so that E = 1 is exactly enough air to burn all the fuel.
load(url(paste0(site, 'ethanol.RData')))

# plot the data
plot(NOx ~ E, data = ethanol)

# fit a linear model
fit.lin <- lm(NOx ~ E, data = ethanol)

# slightly different visualization: predict for a grid of x, then draw a line
# through the points
grid <- data.frame(E = seq(min(ethanol$E), max(ethanol$E), length = 200))
preds.lin <- predict(fit.lin, grid)
lines(x = grid$E, y = preds.lin, col = 'blue', lwd = 2)

# ===========================================================================
# 1. HOW FAR UP THE POLYNOMIAL LADDER?
# ===========================================================================

## example: quadratic model

# quadratic fit
fit.quad <- lm(NOx ~ poly(E, 2), data = ethanol)
summary(fit.quad)

# make a fresh plot and visualize
plot(NOx ~ E, data = ethanol)
preds.quad <- ...
lines(x = ..., y = ..., col = 'blue', lwd = 2)

# test quadratic compared with linear
anova(fit.lin, fit.quad)

## your turn: now fit cubic and quartic models. which degree is best?

# cubic fit
fit.cubic <- ...

# visualize
plot(NOx ~ E, data = ethanol)

# quartic fit
fit.quart <- ...

# visualize
plot(NOx ~ E, data = ethanol)

# test: note you can compare multiple (nested) models in one go
anova(fit.lin, fit.quad, fit.cubic, fit.quart)

# ===========================================================================
# 2. BROKEN STICK: WHERE DOES THE KNOT GO?
# ===========================================================================
# A broken stick commits to exactly one bend, at a location you choose. Theory
# says NOx peaks at E = 1, where the mixture is exactly balanced. The
# scatterplot looks like it turns a little earlier than that.

## example: a knot at E = 1

# broken stick fit: pmax(E - 1, 0) is flat below the knot, so it only changes
# the slope above it
stick.00 <- lm(NOx ~ E + pmax(E - 1.00, 0), data = ethanol)
summary(stick.00)

# visualize
plot(NOx ~ E, data = ethanol)
preds.00 <- predict(stick.00, grid)
lines(x = grid$E, y = preds.00, col = 'blue', lwd = 2)
abline(v = 1.00, lty = 3)

# slope below the knot
coef(stick.00)[2]

# slope above the knot: the second coefficient is the CHANGE in slope
coef(stick.00)[2] + coef(stick.00)[3]

## your turn: try a knot at E = 0.9 instead

# broken stick fit
stick.90 <- ...

# visualize
plot(NOx ~ E, data = ethanol)

# slopes below and above the knot


## your turn: which knot do the data prefer?

# the two fits are not nested, so no F test compares them head to head. but a
# model with BOTH knots contains each of them
stick.both <- lm(NOx ~ E + pmax(E - 0.90, 0) + pmax(E - 1.00, 0),
                 data = ethanol)

# is the knot at 1.00 enough on its own?
anova(stick.00, stick.both)

# is the knot at 0.90 enough on its own?
anova(..., ...)

# write the null hypothesis for each of those tests in terms of coefficients


# an engine does not change behaviour abruptly at one value of E, so the
# corner is a fiction. what is the stick buying you anyway?


# ===========================================================================
# 3. CUBIC SPLINE: HOW MANY KNOTS, THEN WHERE?
# ===========================================================================
# A spline basis buys bends without corners. Two decisions: how many knots,
# and then where to put them.

## example: how many knots?

# bs(x, df = d) places d - 3 interior knots at quantiles of x, so df sets how
# much flexibility you are asking for
sp4 <- lm(NOx ~ bs(E, df = 4), data = ethanol)   # 1 interior knot
sp6 <- lm(NOx ~ bs(E, df = 6), data = ethanol)   # 3 interior knots
sp8 <- lm(NOx ~ bs(E, df = 8), data = ethanol)   # 5 interior knots

# these are nested, so compare with an F test
anova(sp4, sp6, sp8)

# choose one and visualize
plot(NOx ~ E, data = ethanol)
preds.sp <- ...
lines(x = grid$E, y = preds.sp, col = 'blue', lwd = 2)

## your turn: now fine-tune where those three knots go

# check knots for preferred fit
knots.auto <- quantile(ethanol$E, probs = c(...))
knots.auto

# giving those knots directly refits the same model
sp.auto <- lm(NOx ~ bs(E, knots = knots.auto, degree = 3), data = ethanol)
sp.auto

# pick three knots of your own, bunched where the curve bends hardest
knots.manual <- c(..., ..., ...)
sp.manual <- ...

# overlay both fits
plot(NOx ~ E, data = ethanol)
lines(x = grid$E, y = predict(sp.auto, grid), col = 'blue', lwd = 2)
lines(x = ..., y = ..., col = 'red', lwd = 2)
abline(v = knots.manual, lty = 3)

# ===========================================================================
# YOUR TURN: HOW MUCH FERTILIZER IS WORTH BUYING?
# ===========================================================================
# 44 plots of corn at seven nitrogen application rates, 0 to 300 lb/acre.
# Yield climbs and then flattens out, so here the knot is not a modelling
# convenience: it estimates the rate past which fertilizer stops paying.
load(url(paste0(site, 'cornnit.RData')))

# plot the data
plot(yield ~ nitrogen, data = cornnit)

# grid for visualizing fits
grid.n <- data.frame(nitrogen = seq(0, 300, length = 200))

## your turn: specify and fit an appropriate model