88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
from prsctrl.utility.prsdata import PrsData
|
|
import os
|
|
import re
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
def process_results(
|
|
data_dir,
|
|
dir_regex=r".*",
|
|
out_dir=None,
|
|
get_varied_param=lambda data: data.metadata["lock-in_settings"]["frequency_Hz"],
|
|
fig_suptitle="Frequency scan: <qty>",
|
|
xlabel="$f$ [Hz]",
|
|
what=["theta", "stheta", "dR_R", "sdR_R"],
|
|
):
|
|
"""
|
|
Process the results of several spectra recorded with one varying measurement or lock-in parameter.
|
|
|
|
:param data_dir: Directory containing all data directories of the individual files
|
|
:param dir_regex: Regex for filtering data directories
|
|
:param out_dir: Directory into which the plots are saved
|
|
:param get_varied_param: lambda function getting the varied parameter value from a loaded PrsData object
|
|
:param fig_suptitle: Figure suptitle, <qty> will be replaced by the name of the quantity
|
|
:param xlabel: xlabel for the varied parameter
|
|
:param what: For which quantities to create plots
|
|
:return:
|
|
"""
|
|
data_dir = os.path.expanduser(data_dir)
|
|
if out_dir is None:
|
|
out_dir = data_dir
|
|
paths = os.listdir(data_dir)
|
|
data_dirs = []
|
|
for p in paths:
|
|
full_path = os.path.join(data_dir, p)
|
|
if not os.path.isdir(full_path): continue
|
|
m = re.fullmatch(dir_regex, p)
|
|
if m:
|
|
data_dirs.append(full_path)
|
|
else:
|
|
print(f"Unmatched directory {p}")
|
|
assert len(data_dirs) > 0
|
|
data_dirs.sort()
|
|
|
|
varied_params = []
|
|
data = {}
|
|
shape = None
|
|
wls = None
|
|
for d in data_dirs:
|
|
print(f"Getting data from {d}")
|
|
pd = PrsData(load_data_path=d)
|
|
# f = pd.metadata["lock-in_settings"]["frequency_Hz"]
|
|
varied_param = get_varied_param(pd)
|
|
# print(d, f)
|
|
sdata = pd.get_spectrum_data()
|
|
print(pd.wavelengths)
|
|
print(pd.data.keys())
|
|
if wls is None: wls = sdata[:,0]
|
|
if shape is None: shape = sdata.shape
|
|
else:
|
|
if shape != sdata.shape:
|
|
print(f"ERROR Shape mismatch for param={varied_param}: {shape} != {sdata.shape}")
|
|
continue
|
|
# raise ValueError(f"Shape mismatch for {d}: {shape} != {sdata.shape}")
|
|
varied_params.append(varied_param)
|
|
data[varied_param] = sdata
|
|
data_per_wl_and_vp = np.empty((shape[0], len(varied_params), shape[1]))
|
|
varied_params.sort()
|
|
for i in range(shape[0]):
|
|
for j, f in enumerate(varied_params):
|
|
data_per_wl_and_vp[i, j, :] = data[f][i,:]
|
|
print(f"Found wavelengths: {wls}")
|
|
n_cols = 2
|
|
n_rows = wls.shape[0] // n_cols + wls.shape[0] % n_cols
|
|
for qty in what:
|
|
fig, axs = plt.subplots(n_rows, n_cols, figsize=(8, 8))
|
|
axs = axs.flatten()
|
|
qty_idx = PrsData.default_spectrum_columns.index(qty)
|
|
fig.suptitle(fig_suptitle.replace("<qty>", PrsData.key_names[qty]))
|
|
for i, wl in enumerate(wls):
|
|
ax = axs[i]
|
|
ax.set_xlabel(xlabel)
|
|
ax.set_ylabel(PrsData.labels[qty])
|
|
ax.plot(varied_params, data_per_wl_and_vp[i, :, qty_idx])
|
|
ax.set_title(f"$\\lambda = {wl}$ nm")
|
|
fig.tight_layout()
|
|
fig.savefig(os.path.join(out_dir, f"result_{qty}.pdf"))
|
|
print(varied_params)
|