from .functions import eV_to_nm
from .prsdata import *
from .plot import plot_rta, plot_spectrum
from .calc import calc_absorption_data
from devctrl.utility.file_io import listdir_abs
from .check import sanity_check_data_gen_report
import os
import shutil
import re
[docs]
def get_data_name(s, date=False, time=False, sample_type=True, sample_batch=True, info=True, mode=True, replace_underscore=True):
r = r"(\d{4}-\d{2}-\d{2}_)?(\d{2}-\d{2}_)?([^_]+_)([^_]+_)?(.*?)(?:(-...))?$"
match = re.match(s, r)
if match:
name = ""
if date: name += match.groups()[0]
if time: name += match.groups()[1]
if sample_type: name += match.groups()[2]
if sample_batch: name += match.groups()[3]
if info: name += match.groups()[4]
if mode: name += match.groups()[5]
else:
log.debug(f"Failed to match data name: '{s}'")
name = s
if replace_underscore:
name = name.replace("_", " ")
return name
[docs]
def get_data_name_template(s, template="<sample_type> (<sample_batch>)", replace_underscore=True):
r = r"(\d{4}-\d{2}-\d{2}_)?(\d{2}-\d{2}_)?([^_]+_)([^_]+_?)?(.*?)(?:-(...))?$"
match = re.match(r, s)
if match:
repls = [ "<date>", "<time>", "<sample_type>", "<sample_batch>", "<info>", "<mode>"]
groups = [(g or "").removesuffix("_") for g in match.groups()]
name = template
for i, repl in enumerate(repls):
name = name.replace(repl, groups[i] or "")
else:
log.warning(f"Failed to match data name: '{s}'")
name = s
if replace_underscore:
name = name.replace("_", " ")
return name
#
# LOAD/FIND DATA
#
[docs]
def load_rta(parent_data_dir, norm_data_dir=None, force_recalculate_absorption=False, **abs_prsdata_kw):
"""
Load reflection, transmission and absorption data from the path to the parent directory.
The parent directoy must contain two directories with the same as itself, with suffixes -ref and -tra.
The absorption data may exist (suffix -abs) or not, in the latter case it will be calculated from ref and tra.
:param force_recalculate_absoprtion: Calculate absorption data from ref and tra even if it exists
:param norm_data_dir: Data directory with reference data for normalizing the absorption spectrum
:return: Reflection, Transmission, Absorptíon data
"""
dirname = os.path.basename(parent_data_dir.removesuffix("/"))
ref_data = PrsData(os.path.join(parent_data_dir, dirname + "-" + REFLECTION))
tra_data = PrsData(os.path.join(parent_data_dir, dirname + "-" + TRANSMISSION))
abs_dir = os.path.join(parent_data_dir, dirname + "-" + ABSORPTION)
if os.path.isdir(abs_dir) and not force_recalculate_absorption:
abs_data = PrsData(abs_dir)
else:
if norm_data_dir:
norm_data = PrsData(norm_data_dir)
else:
norm_data = None
abs_data = calc_absorption_data(ref_data=ref_data, tra_data=tra_data, norm_data=norm_data, **abs_prsdata_kw)
return ref_data, tra_data, abs_data
[docs]
def find_all_data_dirs(base_dir) -> list[str]:
"""
Recursively find all PrsData directories in base_dir
"""
data_dirs = []
is_data_dir = False
for file in listdir_abs(base_dir):
if os.path.isfile(file) and (file.endswith(FULL_DATA_SUFFIX) or file.endswith(FULL_DATA_SUFFIX+".gz")):
is_data_dir = True
elif os.path.isfile(file) and file.startswith(PARTIAL_PREFIX):
is_data_dir = True
elif os.path.isdir(file):
data_dirs += find_all_data_dirs(file)
if is_data_dir: data_dirs.append(base_dir)
return data_dirs
#
# ABSOPRTION
#
[docs]
def find_all_rta_data_dirs(base_dir: str) -> list[tuple[str,str,str,str|None]]:
"""
Find data directories where reflection and transmission data exists.
The directory paths must habe the same stem and only differ by the -ref -tra -abs suffixes.
:return: list[(shared_part_of_dirname, ref_data_dir, tra_data_dir, abs_data_dir/None)]
"""
all_data_dirs: list[str] = find_all_data_dirs(base_dir)
all_data_dirs.sort()
rta_data_shared_names: list[tuple[str,str,str,str|None]] = []
ref_suffix = f"-{REFLECTION}"
abs_suffix = f"-{ABSORPTION}"
tra_suffix = f"-{TRANSMISSION}"
for i, ddir in enumerate(all_data_dirs):
if not ddir.endswith(ref_suffix): continue
abs_dir = None
tra_dir = None
shared_name = ddir.removesuffix(ref_suffix)
# rely on the fact that data is sorted, abs will come before ref and tra after ref
if i > 0:
if all_data_dirs[i-1].startswith(shared_name) and all_data_dirs[i-1].endswith(abs_suffix):
abs_dir = all_data_dirs[i-1]
if i+1 < len(all_data_dirs):
if all_data_dirs[i+1].startswith(shared_name) and all_data_dirs[i+1].endswith(tra_suffix):
tra_dir = all_data_dirs[i+1]
if tra_dir: # have at least ref and tra
rta_data_shared_names.append((shared_name, ddir, tra_dir, abs_dir))
return rta_data_shared_names
[docs]
def add_all_absorption_datas(base_dir: str, norm_data_dir:str|None=None, **add_kw):
"""
- Search recursively through base_dir
- For all directories that contains reflection and transmission data, do:
- Add an absorption directory
- Calculate the absorption data using reference data from norm_data_dir
- If plot: plot the absorption, and R/T/A combined
"""
if norm_data_dir:
norm_data = PrsData(norm_data_dir)
else:
norm_data = None
rta_dirs = find_all_rta_data_dirs(base_dir)
for shared_dirpath, ref_dir, tra_dir, abs_dir in rta_dirs:
abs_dir = shared_dirpath + f"-{ABSORPTION}"
add_absorption_data(ref_dir, tra_dir, abs_dir, norm_data, **add_kw)
[docs]
def add_absorption_data(ref_data_or_dir: str|PrsData|None, tra_data_or_dir: str|PrsData, abs_dir: str, norm_data_or_dir: str|PrsData|None, replace_existing=False, replot_existing=True, plot=True, csv=True, subplots_kw=None, **plot_spectrum_kw) -> PrsData|None:
ref_data: PrsData = PrsData(ref_data_or_dir) if type(ref_data_or_dir) == str else ref_data_or_dir
tra_data: PrsData = PrsData(tra_data_or_dir) if type(tra_data_or_dir) == str else tra_data_or_dir
if len(ref_data.wavelengths) == 0:
log.warning(f"Not adding absorption data, reflection data has no wavelength data: '{ref_data.dirpath}'")
return None
if len(tra_data.wavelengths) == 0:
log.warning(f"Not adding absorption data, transmission data has no wavelength data: '{ref_data.dirpath}'")
return None
norm_data: PrsData|None = PrsData(norm_data_or_dir) if type(norm_data_or_dir) == str else norm_data_or_dir
if replace_existing or abs_dir is None:
abs_data = calc_absorption_data(ref_data=ref_data, tra_data=tra_data, norm_data=norm_data, file_mode="w", data_path=abs_dir)
abs_data.write_full_file()
abs_data.save_csv(columns=["wl", "E", "dI-X_I", "dI-X", "I", "sdI-X_I", "sdI-X", "sI"])
elif plot and replot_existing:
abs_data = PrsData(abs_dir)
else:
return None
if csv:
write_rta_csv(ref_data, tra_data, abs_data)
if plot:
# Plot of abs
cols = ["wl", "E", "dI-X_I", "sdI-X_I"]
sdata, cols = abs_data.get_spectrum_data(columns=cols)
fig1, _ = plot_spectrum(sdata, columns=cols, title=abs_data.name, what=["dI-X_I"], mode=abs_data.mode, **plot_spectrum_kw)
fig1_path = os.path.join(abs_data.dirpath, abs_data.dirname + ".pdf")
fig1.savefig(fig1_path)
# Plot of ref, tra and abs
plot_rta(ref_data, tra_data, abs_data, fig=None, axs=None, title=True, single_ax=False, subplots_kw=subplots_kw or {}, label=None, save_fig=True, **plot_spectrum_kw)
return abs_data
[docs]
def write_rta_csv(ref_data: PrsData, tra_data: PrsData, abs_data: PrsData|None, save_file=True, columns=["dI-X_I"]) -> str:
datas: list[PrsData] = [ref_data, tra_data]
if abs_data:
datas.append(abs_data)
# get common wavelengths
wls = np.array(datas[0].wavelengths)
for i in range(1, len(datas)):
wls = np.intersect1d(wls, datas[i].wavelengths)
fixed_cols = ["wl", "E"]
# get spectrum data
get_columns = fixed_cols + columns
sdatas = [data.get_spectrum_data(wavelengths=wls, columns=get_columns)[0] for data in datas]
# header line
csv = ",".join([PrsData.get_column_label_with_unit(col) for col in fixed_cols])
for i, data in enumerate(datas):
for col in columns:
csv += "," + data.get_column_label_with_unit(col)
for j, wl in enumerate(wls):
csv += "\n"
# Energy, wavelength
csv += str(sdatas[0][j,0]) + "," + str(sdatas[0][j,1])
for _, sdata in enumerate(sdatas):
for k, col in enumerate(columns):
csv += "," + str(sdata[j,k+len(fixed_cols)])
if save_file == True:
save_file = os.path.join(os.path.dirname(ref_data.dirpath), ref_data.dirname.removesuffix("-" + REFLECTION) + ".csv")
if save_file:
log.info(f"Saving RTA csv as '{save_file}'")
with open(save_file, "w") as f:
f.write(csv)
return csv
#
# SAVING DATA
#
[docs]
def archive_data(data_dir_path: str, archive_dirs: str|list[str], n_wavelengths_threshold=0):
"""
Save a zip of the directory in the archive directory, if it contains enough partial files.
This is restriction is to prevent the archive being flooded with small test measurements.
:param data_dir_path:
:param archive_dirs: archive directory or list of archive directories
:param n_wavelengths_threshold: Minimum number of wavelengths the data must have to be archived
:return:
"""
if n_wavelengths_threshold <= 0:
n_part_files = 0
for f in os.listdir(data_dir_path):
if f.startswith(PARTIAL_PREFIX): n_part_files += 1
if n_part_files < n_wavelengths_threshold:
log.info(f"Not archiving data since it only contains {n_part_files} partial files, threshold is set to {n_wavelengths_threshold}")
return
if not type(archive_dirs) is list:
archive_dirs = [archive_dirs]
if len(archive_dirs) == 0:
return
archive_dir = os.path.expanduser(archive_dirs[0])
if not os.path.isdir(archive_dir):
os.makedirs(archive_dir)
data_dirname = os.path.basename(data_dir_path)
zipname = os.path.join(archive_dir, data_dirname)
log.info(f"Archiving data at '{zipname}.zip'. This might take a while...")
# this is so slow
# import zipfile
# zipfile.main(['-c', zipname, data_dir_path + "/"])
# this is so slow too
shutil.make_archive(zipname, "zip", data_dir_path)
# if more than one archive dirs, copy to the others since zipping is so slow
for i in range(1, len(archive_dirs)):
archive_dir_2 = os.path.expanduser(archive_dirs[i])
log.info(f"Copying archive from '{zipname+'.zip'}' to '{archive_dir_2}'")
shutil.copy(zipname + ".zip", archive_dir_2)
[docs]
def save_data(data_dir_path, copy_data_dirs: str|list[str], dont_save_partial_files=True):
"""
Copy the data from data_dir_path to the data_dir from the configuration file. Partial files can be skipped.
:param data_dir_path:
:param dont_save_partial_files:
:param copy_data_dirs: Directory(ies) to copy the data to
:return:
"""
if not type(copy_data_dirs) is list:
copy_data_dirs = [copy_data_dirs]
for d in copy_data_dirs:
d = os.path.expanduser(d)
new_data_dir = os.path.join(d, os.path.basename(data_dir_path))
if os.path.isdir(new_data_dir):
log.error(f"Data directory '{new_data_dir}' already exists, skipping.")
return
os.makedirs(new_data_dir)
import shutil
log.info(f"Copying data from '{d}' to '{new_data_dir}'")
for f in os.listdir(data_dir_path):
if dont_save_partial_files and f.startswith(PARTIAL_PREFIX): continue
log.debug(f"Copying '{f}' to '{new_data_dir}'")
shutil.copy(os.path.join(data_dir_path, f), new_data_dir)
[docs]
def save_data_rta(ref_data: PrsData, tra_data: PrsData, copy_data_dirs: str|list[str], norm_data_dir: str|None, dont_save_partial_files=True, **plot_spectrum_kw):
"""
Copy the data ref and tra data into a subdir in the data_dir_final(s) from the configuration file.
Partial files can be skipped.
Reflection and transmission data will be placed in a subdirectory and the absorption data will be calculated.
:param data_dir_path_no_suffix: Path to the reference and transmission data, without "-ref"/"-tra")
:param copy_data_dirs: Directory(ies) to copy the data to
:param dont_save_partial_files:
:return*:
"""
# get shared name
data_dir_path_no_suffix_ref = ref_data.dirpath.removesuffix(f"-{REFLECTION}")
data_dir_path_no_suffix_tra = tra_data.dirpath.removesuffix(f"-{TRANSMISSION}")
if data_dir_path_no_suffix_ref != data_dir_path_no_suffix_tra:
raise RuntimeError(f"ref and tra data must share common name, but names are '{data_dir_path_no_suffix_ref}' and '{data_dir_path_no_suffix_tra}")
data_dir_path_no_suffix: str = data_dir_path_no_suffix_ref
dirname = os.path.basename(data_dir_path_no_suffix)
# get target directories
if not type(copy_data_dirs) is list:
copy_data_dirs = [copy_data_dirs]
copy_data_dirs = [os.path.expanduser(d) for d in copy_data_dirs]
# copy and calc in the first dir
new_data_dir_base = os.path.join(copy_data_dirs[0], dirname)
if os.path.isdir(new_data_dir_base):
log.error(f"Data directory '{new_data_dir_base}' already exists, skipping.")
return
os.makedirs(new_data_dir_base)
for mode in [REFLECTION, TRANSMISSION]:
data_dir = data_dir_path_no_suffix + f"-{mode}"
new_data_dir = os.path.join(new_data_dir_base, dirname + f"-{mode}")
log.info(f"Copying data from '{data_dir}' to '{new_data_dir}'")
os.makedirs(new_data_dir)
for f in os.listdir(data_dir):
if dont_save_partial_files and f.startswith(PARTIAL_PREFIX): continue
log.debug(f"Copying '{f}' to '{new_data_dir}'")
shutil.copy(os.path.join(data_dir, f), new_data_dir)
# calculate absorption
if norm_data_dir:
norm_data = PrsData(norm_data_dir)
else:
norm_data = None
abs_dir = os.path.join(new_data_dir_base, dirname + f"-{ABSORPTION}")
abs_data = calc_absorption_data(ref_data=ref_data, tra_data=tra_data, norm_data=norm_data, file_mode="w", data_path=abs_dir)
# save absorption
abs_data.write_full_file()
abs_data.save_csv(columns=["wl", "E", "dI-X_I", "dI-X", "I", "sdI-X_I", "sdI-X", "sI"])
plot_file = os.path.join(new_data_dir_base, dirname + ".pdf")
plot_rta(ref_data, tra_data, abs_data, fig=None, axs=None, title=True, single_ax=False, label=None, save_fig=plot_file, **plot_spectrum_kw)
csv_file = os.path.join(new_data_dir_base, dirname + ".csv")
# combined CSV
try:
write_rta_csv(ref_data, tra_data, abs_data, save_file=csv_file, columns=["dI-X_I"])
except Exception as e:
log.error(f"Failed to write combined csv: {type(e)}: {e}")
# do a combined error report
try:
report = sanity_check_data_gen_report(ref_data)
report += "\n\n" + sanity_check_data_gen_report(tra_data)
report += "\n\n" + sanity_check_data_gen_report(abs_data)
report = report.strip("\n")
if report:
with open(os.path.join(new_data_dir_base, dirname + "-report.md"), "w") as file:
file.write(report)
except Exception as e:
log.error(f"Failed to assemble or write report: {type(e)}: {e}")
# copy all three directories to all other final data directories
for i in range(1, len(copy_data_dirs)):
new_data_dir_base_2 = os.path.join(copy_data_dirs[i], dirname)
log.info(f"Copying from '{new_data_dir_base}' to '{new_data_dir_base_2}'")
shutil.copytree(new_data_dir_base, new_data_dir_base_2)
