Adds instrumentation and some other stuff

Author: PTNobel

Diff for: randalo/adelie_integration.py

@@ -67,13 +67,14 @@ def _matmul_impl(self, v):
                 penalty=np.zeros(S),
                 lmda_path=[0], progress_bar=False, n_threads=32, intercept=False)
         import pickle
-        pickle.dump({
-            'fit_active': state.benchmark_fit_active, 
-            'fit_screen': state.benchmark_fit_screen, 
-            'invariance': state.benchmark_invariance, 
-            'kkt': state.benchmark_kkt, 
-            'screen': state.benchmark_screen,
-        }, f'/scratch/groups/candes/parth/benchmark{_i}.pkl')
+        with open(f'/scratch/groups/candes/parth/benchmark{_i}.pkl', 'wb') as fd:
+            pickle.dump({
+                'fit_active': state.benchmark_fit_active, 
+                'fit_screen': state.benchmark_fit_screen, 
+                'invariance': state.benchmark_invariance, 
+                'kkt': state.benchmark_kkt, 
+                'screen': state.benchmark_screen,
+            }, fd)
         _i += 1
         B = np.array(
             self.X_S @ state.betas.toarray()[0] #.reshape((S, -1), order='C')

Diff for: randalo/reductions.py

@@ -95,11 +95,23 @@ def _matmul_impl(self, rhs):
         # TODO: Split cholesky/minres code paths into seperate ones
         if self.inverse_method == 'minres':
             if constraints is None and hessians is None:
+                import time
+                print("Line 1", time.monotonic(), flush=True)
                 sqrt_d2loss_dy_hat2 = torch.sqrt(d2loss_dy_hat2)[:, None]
+                print("Line 2", time.monotonic(), flush=True)
                 tilde_X = sqrt_d2loss_dy_hat2 * X_mask
-                return ((
-                    X @ minres(X.T @ X, (X.T @ (rhs_scaled / sqrt_d2loss_dy_hat2)))
-                ) / sqrt_d2loss_dy_hat2).to(rhs.dtype)
+                print("Line 3", time.monotonic(), flush=True)
+                p1 = (rhs_scaled / sqrt_d2loss_dy_hat2)
+                print("Line 4", time.monotonic(), flush=True)
+                p2 = X.T @ p1
+                print("Line 5", time.monotonic(), flush=True)
+                p3 = X.T @ X
+                print("Line 6", time.monotonic(), flush=True)
+                p4 = minres(p3, p2)
+                print("Line 7", time.monotonic(), flush=True)
+                p5 = X @ p4
+                print("Line 8", time.monotonic(), flush=True)
+                return (p5 / sqrt_d2loss_dy_hat2).to(rhs.dtype)
             else:
                 raise NotImplementedError()
 

Diff for: utils/run_data_collection.sh

@@ -4,7 +4,7 @@
 #SBATCH --ntasks=1
 #SBATCH --cpus-per-task=32
 #SBATCH --mem=512GB
-#SBATCH --partition=candes
+#SBATCH --partition=candes,pilanci
 
 BASE_DIR=$HOME/adelie_alo/benchmarking/lasso_sweep
 RESULTS_DIR=$BASE_DIR/results

Diff for: utils/run_parameter_sweep.sh

@@ -26,8 +26,8 @@
 df = df.drop('ethnicity', axis=1)
 covars_dense = np.array(
     df[['age', 'age_squared', 'sex'] + [f'PC{i}' for i in range(1, 11)]].to_numpy(),
-    dtype=np.float32)
-y = np.array(df['height'].to_numpy(), dtype=np.float32)
+    dtype=np.float64)
+y = np.array(df['height'].to_numpy(), dtype=np.float64)
 
 chromosomes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 
@@ -37,11 +37,12 @@
         [
             ad.matrix.snp_unphased(
                 ad.io.snp_unphased(
-                    os.path.join(cache_dir, f"EUR_subset_chr{chr}.snpdat"),
-                ), n_threads=32, dtype=np.float32
+                    os.path.join(cache_dir, f"EUR_subset_chr{chr}.snpdat"), "mmap"
+                ), n_threads=32, dtype=np.float64
             )
             for chr in chromosomes],
         axis=1,
+        n_threads=32
 )
 print(f'{X.shape=}')
 
@@ -57,7 +58,7 @@
 print(f'{X_train.shape=}')
 print(f'{X_test.shape=}')
 
-model_cache = f'/scratch/groups/candes/parth/fit_model_{task_id}.pkl'
+model_cache = f'/scratch/groups/candes/parth/fit_model_{task_id}_v3.pkl'
 
 if os.path.exists(model_cache):
     class fake_state:
@@ -78,11 +79,11 @@ def __init__(self):
     ti_solve = time.monotonic()
     state = ad.grpnet(
         X=X_train,
-        glm=ad.glm.gaussian(y_train, dtype=np.float32),
+        glm=ad.glm.gaussian(y_train, dtype=np.float64),
         early_exit=False,
-        min_ratio=1e-6,
+        min_ratio=1e-9,
         n_threads=32,
-        lmda_path_size=241,
+        lmda_path_size=101,
     )
     tf_solve = time.monotonic()
 
@@ -100,7 +101,7 @@ def __init__(self):
     ins[i] = loss(torch.from_numpy(y_hat_train[i]), torch.from_numpy(y_train))
 
 ti_alo = time.monotonic()
-ld, alo, ts, r2 = ai.get_alo_for_sweep_v2(y_train, state, loss, 80)
+ld, alo, ts, r2 = ai.get_alo_for_sweep(y_train, state, loss, 20)
 tf_alo = time.monotonic()
 
-#np.savez(sys.argv[1], alo_lamda=ld, full_lamda=state.lmda_path, alo=alo, oos=oos, in_sample=ins, ts=ts, r2=r2, solve_time=tf_solve - ti_solve, alo_time=tf_alo - ti_alo)
+np.savez(sys.argv[1], alo_lamda=ld, full_lamda=state.lmda_path, alo=alo, oos=oos, in_sample=ins, ts=ts, r2=r2, solve_time=tf_solve - ti_solve, alo_time=tf_alo - ti_alo)