""" data.py - workbook loading and the prediction-target registry. The workbook (AI_Data_Training.xlsx, one row per recording) is read into a tidy table and every predictable quantity in it is turned into a named target. The registry is data-driven: any numeric column that is not an identifier becomes a presence (>0) target and a log-level regression target, and any two-class text column becomes a binary target. Adding a new column to the workbook (for example a Varroa mite count) therefore adds targets automatically, without code changes. """ import os import re import numpy as np import pandas as pd MONTHS = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"] MONTH_NUM = {m: i + 1 for i, m in enumerate(MONTHS)} _MONTH_ALIASES = { "jan": "January", "feb": "February", "mar": "March", "apr": "April", "may": "May", "jun": "June", "june": "June", "jul": "July", "july": "July", "aug": "August", "sep": "September", "sept": "September", "oct": "October", "nov": "November", "dec": "December", } COUNTRY_CODE = {"IA": "USA", "NZ": "NZ"} ID_COLUMNS = {"colony", "month", "country", "sound_file", "year", "month_num", "t_index", "has_audio", "recording_batch"} KNOWN_BINARY_TEXT = {"y03": "Terminal_or_Survivor_0_to_3_months", "y46": "Terminal_or_Survivor_4_to_6_months"} FILENAME_FIXES = { # workbook file name (lower-case) -> file on disk, if ever needed } def normalize_month(raw): if raw is None or (isinstance(raw, float) and np.isnan(raw)): return None s = re.sub(r"20\d{2}", "", str(raw)).strip().lower() s = re.sub(r"[^a-z]", "", s) if s in MONTH_NUM or s.capitalize() in MONTH_NUM: return s.capitalize() return _MONTH_ALIASES.get(s[:4] if s[:4] in _MONTH_ALIASES else s[:3]) def parse_filename(fname): """'12.IA.Sept.2022.wav' -> dict(colony='12', country='USA', month='September', year=2022).""" base = os.path.basename(str(fname)) stem = re.sub(r"\.wav$", "", base, flags=re.IGNORECASE) parts = stem.split(".") out = {"colony": None, "country": None, "month": None, "year": None} if len(parts) >= 1: out["colony"] = parts[0].strip() if len(parts) >= 2: out["country"] = COUNTRY_CODE.get(parts[1].strip().upper(), parts[1].strip()) if len(parts) >= 3: out["month"] = normalize_month(parts[2]) m = re.search(r"20\d{2}", stem) out["year"] = int(m.group()) if m else None return out def _clean_columns(df): df = df.copy() df.columns = [str(c).strip() for c in df.columns] ren = {} for c in df.columns: lc = c.lower() if lc in ("sound file", "sound_file", "file", "filename"): ren[c] = "sound_file" elif lc == "colony": ren[c] = "colony" elif lc == "month": ren[c] = "month" elif lc == "country": ren[c] = "country" elif c == KNOWN_BINARY_TEXT["y03"]: ren[c] = "y03" elif c == KNOWN_BINARY_TEXT["y46"]: ren[c] = "y46" return df.rename(columns=ren) def load_workbook(xlsx_path, audio_dir=None, sheet=0): """Read the workbook into a tidy DataFrame (one row per recording).""" df = pd.read_excel(xlsx_path, sheet_name=sheet) df = _clean_columns(df) if "sound_file" not in df.columns: raise ValueError("workbook needs a 'sound file' column") df["sound_file"] = df["sound_file"].astype(str).str.strip() df["sound_file"] = df["sound_file"].map(lambda s: FILENAME_FIXES.get(s.lower(), s)) parsed = pd.DataFrame([parse_filename(f) for f in df["sound_file"]]) if "colony" not in df.columns: df["colony"] = parsed["colony"] df["colony"] = df["colony"].astype(str).str.strip() if "country" not in df.columns: df["country"] = parsed["country"] df["country"] = df["country"].astype(str).str.strip().str.upper().replace({"IA": "USA", "US": "USA"}) if "month" not in df.columns: df["month"] = parsed["month"] df["month"] = df["month"].map(normalize_month) df["month_num"] = df["month"].map(MONTH_NUM).astype("Int64") df["year"] = parsed["year"].astype("Int64") df["t_index"] = ((df["year"] - 2022) * 12 + df["month_num"]).astype("Int64") # Aug 2022 = 8 df["recording_batch"] = df["country"] + "_" + df["year"].astype(str) + "_" + df["month_num"].astype(str).str.zfill(2) for k in ("y03", "y46"): if k in df.columns: df[k] = df[k].astype("string").str.strip().str.upper() # numeric pathogen / physiology columns: coerce, keep NaN (imputation is a modelling choice) for c in df.columns: if c in ID_COLUMNS or c in ("y03", "y46"): continue if df[c].dtype == object: conv = pd.to_numeric(df[c], errors="coerce") if conv.notna().sum() >= 0.9 * df[c].notna().sum(): df[c] = conv if audio_dir is not None: present = {f.lower() for f in os.listdir(audio_dir)} df["has_audio"] = df["sound_file"].str.lower().isin(present) else: df["has_audio"] = True return df # --------------------------------------------------------------------------- targets def build_targets(df, high_quantile=0.75): """Return a list of target dicts: name, kind ('binary'|'regression'), y (float array with NaN where undefined), source, description.""" targets = [] n = len(df) def add(name, kind, y, source, desc): y = np.asarray(y, dtype=float) targets.append(dict(name=name, kind=kind, y=y, source=source, description=desc)) for key, col in KNOWN_BINARY_TEXT.items(): if key in df.columns: v = df[key] y = np.where(v == "T", 1.0, np.where(v == "S", 0.0, np.nan)) horizon = "0-3 months" if key == "y03" else "0-6 months (nested: every 0-3 month Terminal row is also Terminal here)" add(f"terminal_{key[1:]}", "binary", y, col, f"Colony is Terminal (dies) within {horizon} of the recording (1) vs Survivor (0)") if "y03" in df.columns and "y46" in df.columns: y = np.where((df["y46"] == "T") & (df["y03"] == "S"), 1.0, np.where((df["y46"] == "S") & (df["y03"] == "S"), 0.0, np.nan)) add("terminal_4_6_only", "binary", y, "Terminal_or_Survivor_4_to_6_months minus 0_to_3", "Colony dies 4-6 months after the recording (1) vs survives 6 months (0); rows already Terminal within 3 months excluded") numeric_cols = [c for c in df.columns if c not in ID_COLUMNS and c not in ("y03", "y46") and pd.api.types.is_numeric_dtype(df[c])] for c in numeric_cols: v = pd.to_numeric(df[c], errors="coerce").astype(float).values key = re.sub(r"[^a-z0-9]+", "_", c.lower()).strip("_") present = np.where(np.isnan(v), np.nan, (v > 0).astype(float)) add(f"{key}_present", "binary", present, c, f"{c} detected (> 0) in the recording month") add(f"{key}_log_level", "regression", np.log1p(np.clip(v, 0, None)), c, f"log(1 + {c}) level, all recordings") pos = v > 0 if pos.sum() >= 30: thr = float(np.nanquantile(v[pos], high_quantile)) hi = np.where(pos, (v >= thr).astype(float), np.nan) # among positives only add(f"{key}_high", "binary", hi, c, f"Among recordings with {c} > 0: high load (at or above the {int(high_quantile*100)}th percentile of positive values, {thr:.4g}) vs low positive") among = np.where(pos, np.log1p(v), np.nan) add(f"{key}_log_level_among_positive", "regression", among, c, f"log(1 + {c}) among recordings where {c} > 0 (n={int(pos.sum())})") if len(numeric_cols) >= 2: M = df[numeric_cols].apply(pd.to_numeric, errors="coerce").values.astype(float) anyv = np.where(np.isnan(M).all(axis=1), np.nan, (np.nan_to_num(M) > 0).any(axis=1).astype(float)) add("any_pathogen_present", "binary", anyv, ",".join(numeric_cols), "Any of the measured pathogens detected (> 0); largely a site proxy in this data set") # other two-class text columns become binary targets automatically for c in df.columns: if c in ID_COLUMNS or c in ("y03", "y46") or c in numeric_cols: continue if df[c].dtype == object or str(df[c].dtype) == "string": vals = df[c].dropna().astype(str).str.strip().unique() if len(vals) == 2: pos_label = sorted(vals)[-1] y = np.where(df[c].isna(), np.nan, (df[c].astype(str).str.strip() == pos_label).astype(float)) key = re.sub(r"[^a-z0-9]+", "_", c.lower()).strip("_") add(key, "binary", y, c, f"{c} == '{pos_label}' (1) vs '{sorted(vals)[0]}' (0)") # confound-detectability probes (not colony-health targets): can audio identify site / season? add("probe_country_is_nz", "binary", (df["country"] == "NZ").astype(float).values, "Country", "Probe: recording made in NZ (1) vs USA (0); measures site/recording-condition detectability") if "recording_batch" in df.columns: add("probe_batch_oct2023", "binary", (df["recording_batch"] == "USA_2023_10").astype(float).values, "recording_batch", "Probe: recording belongs to the single-visit October 2023 USA batch (1) vs all other recordings (0)") if df["month_num"].notna().any(): mn = df["month_num"].astype(float).values add("probe_month_sin", "regression", np.sin(2 * np.pi * mn / 12), "Month", "Probe: seasonal phase (sin) of the recording month") add("probe_month_cos", "regression", np.cos(2 * np.pi * mn / 12), "Month", "Probe: seasonal phase (cos) of the recording month") return targets def metadata_matrix(df): """Site, season (two harmonics), year and the October 2023 batch as the metadata-only baseline.""" mn = df["month_num"].astype(float).fillna(0).values ang = 2 * np.pi * mn / 12.0 yr = df["year"].astype(float).fillna(2022).values X = np.column_stack([ (df["country"] == "NZ").astype(float).values, np.sin(ang), np.cos(ang), np.sin(2 * ang), np.cos(2 * ang), (yr - 2022.0), (df["recording_batch"] == "USA_2023_10").astype(float).values, ]) return X.astype(np.float32), ["is_nz", "month_sin", "month_cos", "month_sin2", "month_cos2", "year_offset", "is_oct2023_batch"]