quick data load refactor

This commit is contained in:
Cal W 2024-10-17 22:58:30 +10:00
parent a04be0889a
commit ba1136b0ea

29
main.py
View File

@ -24,6 +24,7 @@ DATA_PATH = "./data"
DATA_INFO = "_info.yaml" DATA_INFO = "_info.yaml"
TUNNEL_INFO_FILE = "./tunnel-info.yaml" TUNNEL_INFO_FILE = "./tunnel-info.yaml"
SAMPLES_TO_AVG = 500 SAMPLES_TO_AVG = 500
CANNY_TIME_OFFSET = 50 #us
with open(TUNNEL_INFO_FILE, 'r') as file: with open(TUNNEL_INFO_FILE, 'r') as file:
TUNNEL_INFO = yaml.safe_load(file) TUNNEL_INFO = yaml.safe_load(file)
@ -36,15 +37,12 @@ data_to_load = [
] ]
# ==== Data Loading & Processing ==== # ==== Data Loading & Processing ====
def load_data(data_to_load: list[str]) -> dict: def load_data(data_path: str, data={}) -> dict:
data = {}
for dp in data_to_load:
data_path = f"{DATA_PATH}/{dp}/"
data_info_path = data_path + DATA_INFO data_info_path = data_path + DATA_INFO
if not os.path.exists(data_info_path): if not os.path.exists(data_info_path):
print(f"[ERR] Could not find data info file: '{data_info_path}'") print(f"[ERR] Could not find data info file: '{data_info_path}'")
print(f"[WARN] Not Loading Data '{dp}'") print(f"[WARN] Not Loading Data '{data_path}'")
continue return None
# Load Shot Data Info YAML File (Cal) # Load Shot Data Info YAML File (Cal)
with open(data_info_path, 'r') as file: with open(data_info_path, 'r') as file:
@ -192,12 +190,11 @@ def load_data(data_to_load: list[str]) -> dict:
probeCh2 = data[x2_shot]["data"]["scope"][2] probeCh2 = data[x2_shot]["data"]["scope"][2]
#[HACK] For detection #[HACK] For detection
TIME_OFFSET = 50 #us
if i > 0: if i > 0:
privPoint = dataInfo["probe-info"]["locations"][i-1] privPoint = dataInfo["probe-info"]["locations"][i-1]
time_offset = data[x2_shot]["shock-point"][f"{privPoint}-g1"][1] + TIME_OFFSET time_offset = data[x2_shot]["shock-point"][f"{privPoint}-g1"][1] + CANNY_TIME_OFFSET
post_pres = 0.03 post_pres = 0.2
sigma = 7 sigma = 1
doPlot = True doPlot = True
else: # These work for the first probe int he chain else: # These work for the first probe int he chain
time_offset = None time_offset = None
@ -216,7 +213,7 @@ def load_data(data_to_load: list[str]) -> dict:
#[HACK] For detection #[HACK] For detection
if i > 0: if i > 0:
time_offset = data[x2_shot]["shock-point"][f"{privPoint}-g2"][1] + TIME_OFFSET time_offset = data[x2_shot]["shock-point"][f"{privPoint}-g2"][1] + CANNY_TIME_OFFSET
shock_point = np.where(probeCh2 >= 0.3)[0] # + offset shock_point = np.where(probeCh2 >= 0.3)[0] # + offset
@ -256,9 +253,15 @@ def load_data(data_to_load: list[str]) -> dict:
print() print()
# Return the data & the successfully loaded data keys # Return the data & the successfully loaded data keys
return data, tuple(data.keys()) return data #, tuple(data.keys())
data = {}
for dp in data_to_load:
pdp = f"{DATA_PATH}/{dp}/"
load_data(pdp, data)
loaded_data = tuple(data.keys())
data, loaded_data = load_data(data_to_load)
print("Loaded Data") print("Loaded Data")