XAS without dimension stacking

src/toolbox_scs/routines/XAS.py

@@ -153,15 +153,18 @@ def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3peaks', nrjkey='nrj',
                 muIo: the mean of the Io
                 counts: the number of events in each bin
     """
-
     if 'pulseId' in nrun:
-        stacked = nrun.stack(dummy_=['trainId', 'pulseId'])
+        pId = 'pulseId'
     else:
-        stacked = nrun.stack(dummy_=['trainId', 'sa3_pId'])
+        pId = 'sa3_pId'
+
+    Io = nrun[Iokey].values.flatten() + Iooffset
+    It = nrun[Itkey].values.flatten()
 
-    Io = stacked[Iokey].values + Iooffset
-    It = stacked[Itkey].values
-    nrj = stacked[nrjkey].values
+    if pId in nrun[nrjkey].dims:
+        nrj = nrun[nrjkey].values.flatten()
+    else:
+        nrj = np.repeat(nrun[nrjkey].values.flatten(), nrun.dims[pId])
 
     names_list = ['nrj', 'Io', 'It']
     rundata = np.vstack((nrj, Io, It))