diff --git a/src/toolbox_scs/detectors/xgm.py b/src/toolbox_scs/detectors/xgm.py index 9d0d794b7e91e62358aa48554e5818f7fbdd1cc0..9c604d5525c43a48fce60e156a13bea84fa93b8d 100644 --- a/src/toolbox_scs/detectors/xgm.py +++ b/src/toolbox_scs/detectors/xgm.py @@ -13,8 +13,7 @@ import xarray as xr import matplotlib.pyplot as plt from ..misc.bunch_pattern_external import is_sase_1, is_sase_3 -from ..mnemonics_machinery import (mnemonics_to_process, - mnemonics_for_run) +from ..mnemonics_machinery import mnemonics_for_run from toolbox_scs.load import (get_array, load_bpt, get_sase_pId) __all__ = [ @@ -25,6 +24,7 @@ __all__ = [ log = logging.getLogger(__name__) + def get_xgm(run, mnemonics=None, merge_with=None, indices=slice(0, None)): """ @@ -75,7 +75,7 @@ def get_xgm(run, mnemonics=None, merge_with=None, if len(m2) == 0: log.info('no XGM mnemonics to process. Skipping.') return merge_with - mnemonics = list(set(m2)) + mnemonics = list(set(m2)) # Prepare the dataset of non-XGM data to merge with if bool(merge_with): ds_mw = merge_with.drop(mnemonics, errors='ignore') @@ -106,8 +106,8 @@ def get_xgm(run, mnemonics=None, merge_with=None, ds_mw = ds_mw.merge(bpt, join='inner') else: xgm_val = da_xgm.values - xgm_val[xgm_val==1] = np.nan - xgm_val[xgm_val==0] = np.nan + xgm_val[xgm_val == 1] = np.nan + xgm_val[xgm_val == 0] = np.nan da_xgm.values = xgm_val da_xgm = da_xgm.dropna(dim='XGMbunchId', how='all') ds_xgm = da_xgm.fillna(0).sel(XGMbunchId=indices).to_dataset() @@ -143,15 +143,15 @@ def load_xgm_array(run, xgm, mnemonic, sase1, sase3): the dataset containing the aligned XGM variable(s). """ xgm_val = xgm.values - xgm_val[xgm_val==1] = np.nan - xgm_val[xgm_val==0] = np.nan + xgm_val[xgm_val == 1] = np.nan + xgm_val[xgm_val == 0] = np.nan xgm.values = xgm_val xgm = xgm.dropna(dim='XGMbunchId', how='all') xgm = xgm.fillna(0) if 'XGM' in mnemonic: sase1_3 = np.sort(np.concatenate([sase1, sase3])) - sase1_idx = [np.argwhere(sase1_3==i)[0][0] for i in sase1] - sase3_idx = [np.argwhere(sase1_3==i)[0][0] for i in sase3] + sase1_idx = [np.argwhere(sase1_3 == i)[0][0] for i in sase1] + sase3_idx = [np.argwhere(sase1_3 == i)[0][0] for i in sase3] xgm_sa1 = xgm.isel(XGMbunchId=sase1_idx).rename(XGMbunchId='sa1_pId') xgm_sa1 = xgm_sa1.assign_coords(sa1_pId=sase1) xgm_sa1 = xgm_sa1.rename(mnemonic.replace('XGM', 'SA1'))