Includes input energy parameter in the model and adds non-linearities

pes_to_spec/test/offline_analysis.py

@@ -143,7 +143,8 @@ def main():
     parser.add_argument('-X', '--xgm', type=str, metavar='NAME', default="SA3_XTD10_XGM/XGM/DOOCS:output", help='XGM name')
     parser.add_argument('-o', '--offset', type=int, metavar='INT', default=0, help='Train ID offset')
     parser.add_argument('-c', '--xgm_cut', type=float, metavar='INTENSITY', default=500, help='XGM intensity threshold in uJ.')
-    parser.add_argument('-e', '--poly', action="store_true", default=False, help='Wheteher to expand PES data in higher order polynomials.')
+    parser.add_argument('-e', '--bnn', action="store_true", default=False, help='Use BNN?')
+    parser.add_argument('-w', '--weight', action="store_true", default=False, help='Whether to reweight data as a function of the pulse energy to make it invariant to that.')
 
     args = parser.parse_args()
 
@@ -233,7 +234,7 @@ def main():
     t = list()
     t_names = list()
 
-    model = Model(poly=args.poly)
+    model = Model(bnn=args.bnn)
 
     train_idx = np.isin(tids, train_tids) & (xgm_flux[:,0] > args.xgm_cut)
     # we just need this for training and we need to avoid copying it, which blows up the memoray usage
@@ -245,6 +246,8 @@ def main():
         print("Fitting")
         start = time_ns()
         w = model.uniformize(xgm_flux[train_idx])
+        if not args.weight:
+            w[...] = 1.0
         print("w", np.amin(w), np.amax(w), np.median(w))
         model.fit(pes_raw,
                   #{k: v[train_idx]