ssl-based mlpf first iteration (#158)

* add ssl notebook * add ssl notebook * Squashed commit of the following: commit a4d5077 Author: Farouk <farouk.mokhtar@gmail.com> Date: Thu Nov 3 15:26:54 2022 +0100 up commit 9012e78 Author: Farouk <farouk.mokhtar@gmail.com> Date: Thu Nov 3 15:04:17 2022 +0100 up commit 2915130 Author: Farouk <farouk.mokhtar@gmail.com> Date: Thu Nov 3 14:29:01 2022 +0100 add ssl-VICreg notebook * up * VICReg apples to apples * major cleanup and reorganization * ude full dataset for VICReg * update readme pyg * add first iteration of ssl * before merge * cleanup and ssl organization * add flake8 * revert changes * up Former-commit-id: ee293a2
jpata · Jan 12, 2023 · dcdd941 · dcdd941
1 parent 34d8695
commit dcdd941
Show file tree

Hide file tree

Showing 38 changed files with 6,963 additions and 2,318 deletions.
diff --git a/.gitignore b/.gitignore
@@ -23,5 +23,8 @@ nohup.out
 
 *.pkl
 *.pkl.bz2
+*.parquet
 
 slurm-*.out
+
+.vscode
diff --git a/mlpf/data_clic/postprocessing.py b/mlpf/data_clic/postprocessing.py
@@ -189,6 +189,13 @@ def flatten_event(df_tr, df_cl, df_gen, df_pfs, pairs):
 
 
 def prepare_data_clic(fn):
+    """
+    Processing function that takes as input a raw parquet file and processes it.
+
+    Returns
+        a list of events, each containing three arrays [Xs, ygen, ycand].
+
+    """
 
     data = awkward.from_parquet(fn)
 
@@ -308,6 +315,7 @@ def prepare_data_clic(fn):
             # print(df_gen.loc[gp])
 
         Xs, ys_gen, ys_cand = flatten_event(df_tr, df_cl, df_gen, df_pfs, pairs)
+
         ret.append([Xs, ys_gen, ys_cand])
 
     return ret
diff --git a/mlpf/data_cms/postprocessing2.py b/mlpf/data_cms/postprocessing2.py
@@ -810,12 +810,7 @@ def parse_args():
     import argparse
 
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--input",
-        type=str,
-        help="Input file from PFAnalysis",
-        required=True,
-    )
+    parser.add_argument("--input", type=str, help="Input file from PFAnalysis", required=True)
     parser.add_argument("--outpath", type=str, default="raw", help="output path")
     parser.add_argument(
         "--save-full-graph",

diff --git a/mlpf/lrp/__init__.py b/mlpf/lrp/__init__.py
diff --git a/mlpf/lrp/lrp_mlpf.py b/mlpf/lrp/lrp_mlpf.py