Adds timing and r2

examples/adelie_example.py

@@ -36,7 +36,7 @@
 import randalo.adelie_integration as ai
 import torch
 
-ld, alo = ai.get_alo_for_sweep(y, state, torch.nn.MSELoss())
+ld, alo, ts, r2 = ai.get_alo_for_sweep(y, state, torch.nn.MSELoss())
 dg = ad.diagnostic.diagnostic(state)
 dg.plot_devs()
 plt.plot(-np.log(ld), alo)

randalo/adelie_integration.py

@@ -76,8 +76,7 @@ def adelie_state_to_jacobian(y, state, adelie_state):
 
     return loss, J
 
-def adelie_state_to_randalo(y, state, adelie_state, loss, J, index, rng=None):
-    y_hat = (state.X @ state.betas[index].T).squeeze()
+def adelie_state_to_randalo(y, y_hat, state, adelie_state, loss, J, index, rng=None):
     adelie_state.set_index(index)
     randalo = ra.RandALO(
             loss,
@@ -92,15 +91,18 @@ def get_alo_for_sweep(y, state, risk_fun):
     L, _ = state.betas.shape
     adelie_state = AdelieState(state)
     loss, J = adelie_state_to_jacobian(y, state, adelie_state)
+    y_hat = ad.diagnostic.predict(state.X, state.betas, state.intercepts)
 
     output = np.empty(L)
     times = np.empty(L)
+    r2 = np.empty(L)
 
     for i in range(L):
         t0 = time.monotonic()
-        randalo = adelie_state_to_randalo(y, state, adelie_state, loss, J, i)
+        randalo = adelie_state_to_randalo(y, y_hat[i], state, adelie_state, loss, J, i)
         output[i] = randalo.evaluate(risk_fun)
         times[i] = time.monotonic() - t0
+        r2[i] = 1 - np.square(y - y_hat[i]).sum() / np.square(y - np.mean(y)).sum()
 
-    return state.lmda_path[:L], output, times
+    return state.lmda_path[:L], output, times, r2
 

utils/sherlock_script.py

@@ -67,6 +67,6 @@
     oos[i] = loss(torch.from_numpy(y_hat_test[:, i]), torch.from_numpy(y_test))
     ins[i] = loss(torch.from_numpy(y_hat_train[:, i]), torch.from_numpy(y_train))
 
-ld, alo, ts = ai.get_alo_for_sweep(y_train, state, loss)
+ld, alo, ts, r2 = ai.get_alo_for_sweep(y_train, state, loss)
 
-np.savez(sys.argv[-1], lamda=ld, alo=alo, oos=oos, in_sample=ins, ts=ts)
+np.savez(sys.argv[-1], lamda=ld, alo=alo, oos=oos, in_sample=ins, ts=ts, r2=r2)