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