Code:
// =============================================
// Elastic Net Multiple Linear Regression
// for the HP Prime Graphing Calculator (PPL)
// =============================================
// Minimizes: (1/(2n))*||y - X*b||^2
// + lambda1*||b||_1
// + lambda2*||b||_2^2
//
// Uses coordinate descent with soft
// thresholding. Lambda1 and lambda2 are
// estimated via BIC over a grid search.
// =============================================
// -----------------------------------------
// Soft-thresholding operator
// S(z, lam) = sign(z)*max(|z|-lam, 0)
// -----------------------------------------
softThresh(z, lam)
BEGIN
IF z > lam THEN
RETURN z - lam;
END;
IF z < -lam THEN
RETURN z + lam;
END;
RETURN 0;
END;
// -----------------------------------------
// Coordinate descent on standardized data.
// xn: n x m standardized matrix
// yn: n x 1 centered response
// lam1, lam2: penalty parameters
// maxIt: max iterations
// Returns m x 1 coefficient vector
// -----------------------------------------
cdSolve(xn, yn, lam1, lam2, maxIt)
BEGIN
LOCAL sz, n, m, i, j, k;
LOCAL beta, res, zj, bNew, dif;
LOCAL maxCh, done;
sz := SIZE(xn);
n := sz(1);
m := sz(2);
// Initialize beta = 0, residual = yn
beta := MAKEMAT(0, m, 1);
res := MAKEMAT(0, n, 1);
FOR i FROM 1 TO n DO
res(i, 1) := yn(i, 1);
END;
done := 0;
FOR k FROM 1 TO maxIt DO
IF done == 0 THEN
maxCh := 0;
FOR j FROM 1 TO m DO
// Compute z_j = beta_j + (X_j' * res)/n
LOCAL dot;
dot := 0;
FOR i FROM 1 TO n DO
dot := dot + xn(i, j) * res(i, 1);
END;
zj := beta(j, 1) + dot / n;
// Apply soft-thresholding
bNew := softThresh(zj, lam1)
/ (1 + 2 * lam2);
// Update residual and beta
dif := bNew - beta(j, 1);
IF ABS(dif) > maxCh THEN
maxCh := ABS(dif);
END;
IF ABS(dif) > 1E-15 THEN
FOR i FROM 1 TO n DO
res(i, 1) := res(i, 1)
- xn(i, j) * dif;
END;
END;
beta(j, 1) := bNew;
END; // j
IF maxCh < 1E-9 THEN
done := 1;
END;
END; // done check
END; // k
RETURN beta;
END;
// -----------------------------------------
// elasticNetMlr(x, y)
//
// x : n x m predictor matrix
// y : n x 1 response vector
//
// Returns a list:
// { intercept,
// coefficients (m x 1 matrix),
// adjusted R-squared,
// RMSE,
// best lambda1,
// best lambda2 }
// -----------------------------------------
EXPORT elasticNetMlr(x, y)
BEGIN
LOCAL sz, n, m, i, j;
LOCAL xm, xsd, ym;
LOCAL xn, yn;
LOCAL lamMax, dot;
LOCAL lam1v, lam2v, nL1, nL2;
LOCAL iL1, iL2, lam1, lam2;
LOCAL beta, res;
LOCAL rss, df, bic;
LOCAL bestBIC, bestL1, bestL2;
LOCAL bCoef, b0;
LOCAL yh, sse, sst, r2, adjR2, rmse;
sz := SIZE(x);
n := sz(1);
m := sz(2);
// --- Compute column means of x ---
xm := MAKEMAT(0, 1, m);
FOR j FROM 1 TO m DO
LOCAL s;
s := 0;
FOR i FROM 1 TO n DO
s := s + x(i, j);
END;
xm(1, j) := s / n;
END;
// --- Compute population std dev of x ---
// (so that X_j'X_j = n after scaling)
xsd := MAKEMAT(0, 1, m);
FOR j FROM 1 TO m DO
LOCAL ss;
ss := 0;
FOR i FROM 1 TO n DO
ss := ss + (x(i, j) - xm(1, j))^2;
END;
xsd(1, j) := ss / n;
IF xsd(1, j) < 1E-15 THEN
xsd(1, j) := 1;
END;
END;
// --- Compute mean of y ---
ym := 0;
FOR i FROM 1 TO n DO
ym := ym + y(i, 1);
END;
ym := ym / n;
// --- Standardize x, center y ---
xn := MAKEMAT(0, n, m);
yn := MAKEMAT(0, n, 1);
FOR i FROM 1 TO n DO
FOR j FROM 1 TO m DO
xn(i, j) := (x(i, j) - xm(1, j))
/ xsd(1, j);
END;
yn(i, 1) := y(i, 1) - ym;
END;
// --- Compute lambda_max ---
// lambda_max = max_j |X_j' y_c| / n
// Above this value all coefficients = 0
lamMax := 0;
FOR j FROM 1 TO m DO
dot := 0;
FOR i FROM 1 TO n DO
dot := dot + xn(i, j) * yn(i, 1);
END;
IF ABS(dot) / n > lamMax THEN
lamMax := ABS(dot) / n;
END;
END;
// Prevent lamMax = 0 (perfect zero corr.)
IF lamMax < 1E-12 THEN
lamMax := 1;
END;
// --- Grid search over (lam1, lam2) ---
// lam1 as fractions of lamMax
lam1v := {0.001, 0.005, 0.01, 0.05,
0.1, 0.2, 0.4};
// lam2 absolute values
lam2v := {0.0001, 0.001, 0.01, 0.05,
0.1, 0.5};
nL1 := 7;
nL2 := 6;
bestBIC := 1E30;
bestL1 := 0.01 * lamMax;
bestL2 := 0.01;
FOR iL1 FROM 1 TO nL1 DO
FOR iL2 FROM 1 TO nL2 DO
lam1 := lam1v(iL1) * lamMax;
lam2 := lam2v(iL2);
// Run coordinate descent (coarse)
beta := cdSolve(xn, yn, lam1, lam2,
80);
// Compute RSS on standardized data
rss := 0;
FOR i FROM 1 TO n DO
LOCAL ri;
ri := yn(i, 1);
FOR j FROM 1 TO m DO
ri := ri - xn(i, j) * beta(j, 1);
END;
rss := rss + ri^2;
END;
// Count non-zero coefficients
df := 0;
FOR j FROM 1 TO m DO
IF ABS(beta(j, 1)) > 1E-12 THEN
df := df + 1;
END;
END;
// BIC = n*ln(RSS/n) + (df+1)*ln(n)
bic := n * LN(rss / n + 1E-30)
+ (df + 1) * LN(n);
IF bic < bestBIC THEN
bestBIC := bic;
bestL1 := lam1;
bestL2 := lam2;
END;
END; // iL2
END; // iL1
// --- Final fit with best lambdas ---
beta := cdSolve(xn, yn, bestL1, bestL2,
300);
// --- Convert to original scale ---
// b_orig_j = beta_std_j / xsd_j
// b0 = ym - sum(b_orig_j * xm_j)
bCoef := MAKEMAT(0, m, 1);
b0 := ym;
FOR j FROM 1 TO m DO
bCoef(j, 1) := beta(j, 1) / xsd(1, j);
b0 := b0 - bCoef(j, 1) * xm(1, j);
END;
// --- Compute fit statistics ---
sse := 0;
sst := 0;
FOR i FROM 1 TO n DO
yh := b0;
FOR j FROM 1 TO m DO
yh := yh + bCoef(j, 1) * x(i, j);
END;
sse := sse + (y(i, 1) - yh)^2;
sst := sst + (y(i, 1) - ym)^2;
END;
// R-squared
IF sst > 1E-30 THEN
r2 := 1 - sse / sst;
ELSE
r2 := 0;
END;
// Effective df = non-zero coefficients
df := 0;
FOR j FROM 1 TO m DO
IF ABS(bCoef(j, 1)) > 1E-12 THEN
df := df + 1;
END;
END;
// Adjusted R-squared
IF n - df - 1 > 0 THEN
adjR2 := 1 - (1 - r2) * (n - 1)
/ (n - df - 1);
ELSE
adjR2 := r2;
END;
// RMSE
rmse := sse / n;
RETURN {b0, bCoef, adjR2, rmse,
bestL1, bestL2};
END;
// =============================================
// Test Program
// =============================================
EXPORT testElasticNet()
BEGIN
LOCAL x, y, res;
LOCAL b0, bCoef, adjR2, rmse;
LOCAL l1, l2;
LOCAL n, m, j;
LOCAL txt;
// --- Test 1: Clean linear data ---
// y = 2 + 3*x1 + 5*x2
// 8 observations, 2 predictors
PRINT();
PRINT("=== Test 1: y = 2 + 3x1 + 5x2 ===");
x := [[1, 2],
[2, 1],
[3, 3],
[4, 2],
[5, 5],
[6, 4],
[7, 6],
[8, 5]];
y := [[2 + 3*1 + 5*2],
[2 + 3*2 + 5*1],
[2 + 3*3 + 5*3],
[2 + 3*4 + 5*2],
[2 + 3*5 + 5*5],
[2 + 3*6 + 5*4],
[2 + 3*7 + 5*6],
[2 + 3*8 + 5*5]];
res := elasticNetMlr(x, y);
b0 := res(1);
bCoef := res(2);
adjR2 := res(3);
rmse := res(4);
l1 := res(5);
l2 := res(6);
PRINT("Intercept : " + STRING(b0));
n := SIZE(bCoef);
FOR j FROM 1 TO n(1) DO
PRINT(" b(" + STRING(j) + ") = "
+ STRING(bCoef(j, 1)));
END;
PRINT("Adj R² : " + STRING(adjR2));
PRINT("RMSE : " + STRING(rmse));
PRINT("Lambda1 : " + STRING(l1));
PRINT("Lambda2 : " + STRING(l2));
PRINT("Expected : b0˜2, b1˜3, b2˜5");
PRINT("");
// --- Test 2: Noisy data, 3 predictors ---
// y ˜ 10 + 2*x1 - 3*x2 + 0*x3
// x3 is irrelevant (elastic net may
// shrink its coefficient toward zero)
PRINT("=== Test 2: 3 predictors ===");
PRINT(" (x3 irrelevant)");
x := [[1, 5, 10],
[2, 4, 20],
[3, 3, 30],
[4, 2, 15],
[5, 1, 25],
[6, 6, 12],
[7, 7, 18],
[8, 8, 22],
[9, 9, 28],
[10, 10, 35]];
// y = 10 + 2*x1 - 3*x2 + noise
y := [[10 + 2*1 - 3*5 + 0.5],
[10 + 2*2 - 3*4 - 0.3],
[10 + 2*3 - 3*3 + 0.8],
[10 + 2*4 - 3*2 + 0.1],
[10 + 2*5 - 3*1 - 0.6],
[10 + 2*6 - 3*6 + 0.4],
[10 + 2*7 - 3*7 - 0.2],
[10 + 2*8 - 3*8 + 0.7],
[10 + 2*9 - 3*9 + 0.3],
[10 + 2*10 - 3*10 - 0.5]];
res := elasticNetMlr(x, y);
b0 := res(1);
bCoef := res(2);
adjR2 := res(3);
rmse := res(4);
l1 := res(5);
l2 := res(6);
PRINT("Intercept : " + STRING(b0));
n := SIZE(bCoef);
FOR j FROM 1 TO n(1) DO
PRINT(" b(" + STRING(j) + ") = "
+ STRING(bCoef(j, 1)));
END;
PRINT("Adj R² : " + STRING(adjR2));
PRINT("RMSE : " + STRING(rmse));
PRINT("Lambda1 : " + STRING(l1));
PRINT("Lambda2 : " + STRING(l2));
PRINT("Expected : b0˜10, b1˜2,");
PRINT(" b2˜-3, b3˜0");
PRINT("");
// --- Test 3: Correlated predictors ---
// x2 ˜ 2*x1, testing grouping effect
PRINT("=== Test 3: Correlated x ===");
x := [[1, 2.1],
[2, 3.9],
[3, 6.2],
[4, 7.8],
[5, 10.1],
[6, 12.3],
[7, 13.8],
[8, 16.2],
[9, 17.9],
[10, 20.0],
[11, 22.1],
[12, 24.3]];
// y = 5 + 4*(x1 + x2/2) + noise
y := [[5 + 4*1 + 4*2.1/2 + 0.3],
[5 + 4*2 + 4*3.9/2 - 0.5],
[5 + 4*3 + 4*6.2/2 + 0.2],
[5 + 4*4 + 4*7.8/2 + 0.7],
[5 + 4*5 + 4*10.1/2 - 0.1],
[5 + 4*6 + 4*12.3/2 + 0.4],
[5 + 4*7 + 4*13.8/2 - 0.8],
[5 + 4*8 + 4*16.2/2 + 0.6],
[5 + 4*9 + 4*17.9/2 + 0.1],
[5 + 4*10 + 4*20.0/2 - 0.3],
[5 + 4*11 + 4*22.1/2 + 0.5],
[5 + 4*12 + 4*24.3/2 - 0.4]];
res := elasticNetMlr(x, y);
b0 := res(1);
bCoef := res(2);
adjR2 := res(3);
rmse := res(4);
l1 := res(5);
l2 := res(6);
PRINT("Intercept : " + STRING(b0));
n := SIZE(bCoef);
FOR j FROM 1 TO n(1) DO
PRINT(" b(" + STRING(j) + ") = "
+ STRING(bCoef(j, 1)));
END;
PRINT("Adj R² : " + STRING(adjR2));
PRINT("RMSE : " + STRING(rmse));
PRINT("Lambda1 : " + STRING(l1));
PRINT("Lambda2 : " + STRING(l2));
PRINT("Note: b1+b2 should total ˜6");
PRINT("(elastic net groups correlated vars)");
PRINT("");
PRINT("Done.");
END;