zebra/pyzebra
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Performance optimization

Browse Source 
Use lists as intermediate data structure to avoid lots of numpy array
allocations
This commit is contained in:
usov_i 2022-04-20 13:15:48 +02:00
parent 49ff319230
commit 098314e30d
1 changed files with 21 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
pyzebra/ccl_process.py
View File

@ -169,23 +169,26 @@ def merge_h5_scans(scan_into, scan_from):
scan_motor = scan_into["scan_motor"] # the same as scan_from["scan_motor"] scan_motor = scan_into["scan_motor"] # the same as scan_from["scan_motor"]
pos_all = scan_into["init_scan"][scan_motor] pos_all = [scan_into["init_scan"][scan_motor]]
val_all = scan_into["init_scan"]["counts"] val_all = [scan_into["init_scan"]["counts"]]
err_all = scan_into["init_scan"]["counts_err"] ** 2 err_all = [scan_into["init_scan"]["counts_err"] ** 2]
for scan in scan_into["merged_scans"]: for scan in scan_into["merged_scans"]:
pos_all = np.append(pos_all, scan[scan_motor]) pos_all.append(scan[scan_motor])
val_all = np.concatenate((val_all, scan["counts"])) val_all.append(scan["counts"])
err_all = np.concatenate((err_all, scan["counts_err"] ** 2)) err_all.append(scan["counts_err"] ** 2)
pos_all = np.concatenate(pos_all)
val_all = np.concatenate(val_all)
err_all = np.concatenate(err_all)
sort_index = np.argsort(pos_all) sort_index = np.argsort(pos_all)
pos_all = pos_all[sort_index] pos_all = pos_all[sort_index]
val_all = val_all[sort_index] val_all = val_all[sort_index]
err_all = err_all[sort_index] err_all = err_all[sort_index]
pos_tmp = pos_all[:1] pos_tmp = [pos_all[0]]
val_tmp = val_all[:1] val_tmp = [val_all[:1]]
err_tmp = err_all[:1] err_tmp = [err_all[:1]]
num_tmp = np.array([1]) num_tmp = [1]
for pos, val, err in zip(pos_all[1:], val_all[1:], err_all[1:]): for pos, val, err in zip(pos_all[1:], val_all[1:], err_all[1:]):
if pos - pos_tmp[-1] < MOTOR_POS_PRECISION: if pos - pos_tmp[-1] < MOTOR_POS_PRECISION:
# the repeated motor position # the repeated motor position
@ -194,10 +197,14 @@ def merge_h5_scans(scan_into, scan_from):
num_tmp[-1] += 1 num_tmp[-1] += 1
else: else:
# a new motor position # a new motor position
pos_tmp = np.append(pos_tmp, pos) pos_tmp.append(pos)
val_tmp = np.concatenate((val_tmp, val[None, :])) val_tmp.append(val[None, :])
err_tmp = np.concatenate((err_tmp, err[None, :])) err_tmp.append(err[None, :])
num_tmp = np.append(num_tmp, 1) num_tmp.append(1)
pos_tmp = np.array(pos_tmp)
val_tmp = np.concatenate(val_tmp)
err_tmp = np.concatenate(err_tmp)
num_tmp = np.array(num_tmp)
scan_into[scan_motor] = pos_tmp scan_into[scan_motor] = pos_tmp
scan_into["counts"] = val_tmp / num_tmp[:, None, None] scan_into["counts"] = val_tmp / num_tmp[:, None, None]