66 lines
2.3 KiB
Python
66 lines
2.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"202.-..-.._f-scan_f=(\d+)_Hz", out_dir=None):
|
|
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()
|
|
|
|
frequencies = []
|
|
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"]
|
|
# 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 f={f}: {shape} != {sdata.shape}")
|
|
continue
|
|
# raise ValueError(f"Shape mismatch for {d}: {shape} != {sdata.shape}")
|
|
frequencies.append(f)
|
|
data[f] = sdata
|
|
data_per_wl_and_f = np.empty((shape[0], len(frequencies), shape[1]))
|
|
frequencies.sort()
|
|
for i in range(shape[0]):
|
|
for j, f in enumerate(frequencies):
|
|
data_per_wl_and_f[i, j, :] = data[f][i,:]
|
|
print(f"Found wavelengths: {wls}")
|
|
n_cols = 2
|
|
for qty in ["theta", "stheta", "dR_R", "sdR_R"]:
|
|
fig, axs = plt.subplots(wls.shape[0]//n_cols, n_cols, sharex=True, figsize=(8, 8))
|
|
axs = axs.flatten()
|
|
qty_idx = PrsData.default_spectrum_columns.index(qty)
|
|
fig.suptitle(f"Frequency scan: {PrsData.key_names[qty]}")
|
|
axs[-1].set_xlabel("Modulation Frequency $f$ [Hz]")
|
|
for i, wl in enumerate(wls):
|
|
ax = axs[i]
|
|
ax.set_ylabel(PrsData.labels[qty])
|
|
ax.plot(frequencies, data_per_wl_and_f[i, :, qty_idx])
|
|
ax.set_title(f"$\\lambda = {wl}$ nm")
|
|
fig.tight_layout()
|
|
fig.savefig(out_dir + f"result_{qty}.pdf")
|
|
print(frequencies)
|