Brains_time_constant.py

#!/usr/bin/env python3
"""EMA Time-Constant Fitting for DCD Species Learning Analysis.

This script:
- loads subject-level binned PropCorrect data from ./data/IndividualBinnedData_DCD.csv
- builds per-subject PropCorrect time series for selected species (fit_species)
- scans EMA coefficients (alpha) from 0.05 to 0.95
- computes mean squared one-step-ahead prediction error for each alpha
- reports the alpha with minimal loss and its equivalent time constant tau (in bins)
- plots predictive loss versus alpha so the chosen value can be inspected visually.

The fitted alpha is intended to be used as EXP_ALPHA in the main analysis script
when TTC_MODE == "ema".
"""

import numpy as np
import csv
import matplotlib.pyplot as plt

from pathlib import Path

# Base directory = parent folder of this script (the main Brains directory)
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"

# Data files now live in BASE_DIR / "data"
raw_learning_path = DATA_DIR / "IndividualBinnedData_DCD.csv"


# ---------- loader (same pattern as in Brains_predict_test_from_training.py) ----------

def load_structured_csv(filepath, dtypes):
    converters = []
    for name, dt in dtypes:
        if dt in (float, int):
            converters.append(dt)
        else:
            converters.append(str)
    rows = []
    with open(filepath, newline="") as f:
        reader = csv.reader(f)
        _header = next(reader)
        for line in reader:
            try:
                parsed = [conv(line[i]) for i, conv in enumerate(converters)]
                rows.append(tuple(parsed))
            except (ValueError, IndexError):
                continue
    return np.array(rows, dtype=dtypes)

# ---------- paths (adjust if needed) ----------
# Data file is expected to live in the same folder as this script.
# Using a path relative to the script makes the code portable across machines.


raw_learning_dtypes = [
    ("SubjectID", "U20"),
    ("Species", "U20"),
    ("Bin", int),
    ("CorrectNum", int),
    ("TotalTrials", int),
    ("PropCorrect", float),
]

raw_data = load_structured_csv(raw_learning_path, raw_learning_dtypes)
print("Raw learning data rows:", raw_data.shape[0])

# plural -> singular mapping from your main script
species_name_map = {
    "bumblebees":   "bumblebee",
    "honeybees":    "honeybee",
    "salamanders":  "salamander",
    "chickens":     "chicken",
    "hummingbirds": "hummingbird",
    "bluejays":     "bluejay",
    "tortoises":    "tortoise",
    "capuchins":    "capuchin",
}

# which species to use to fit alpha (you can tweak this list)
# Species to include when fitting alpha. This controls which subject time series
# contribute to the predictive loss; you can tweak this list if needed.
fit_species = [
    "bumblebee",
    "honeybee",
    "salamander",
    "chicken",
    "hummingbird",
    "bluejay",
    "tortoise",
    "capuchin",
]


def get_plural_name(sing_name):
    for plural, sing in species_name_map.items():
        if sing == sing_name:
            return plural
    raise ValueError(f"No plural mapping for {sing_name}")


# ---------- collect per-subject PropCorrect sequences ----------

def get_subject_series():
    """
    Collect subject-level PropCorrect time series for each species.

    Returns
    -------
    series_by_species : dict
        Keys are species (singular names, from fit_species), values are lists
        of 1D numpy arrays y, where each y contains PropCorrect across bins
        for a single subject, sorted by Bin. Only subjects with at least 3
        bins are retained.
    """
    series_by_species = {sp: [] for sp in fit_species}

    for sp in fit_species:
        plural = get_plural_name(sp)
        d_sp = raw_data[raw_data["Species"] == plural]
        if d_sp.size == 0:
            continue

        subj_ids = np.unique(d_sp["SubjectID"])
        for sid in subj_ids:
            d_subj = d_sp[d_sp["SubjectID"] == sid]
            # sort by Bin
            order = np.argsort(d_subj["Bin"])
            bins = d_subj["Bin"][order]
            y = d_subj["PropCorrect"][order].astype(float)

            # # require at least 3 bins to contribute something meaningful
            # if y.size >= 3:
            series_by_species[sp].append(y)

    return series_by_species

series_by_species = get_subject_series()

print("\nSubject series per species (for fitting alpha):")
for sp in fit_species:
    n = len(series_by_species.get(sp, []))
    print(f"{sp:12s}: {n:3d} subjects with >=3 bins")

# ---------- loss for a given alpha ----------

def predictive_loss(alpha, series_by_species):
    """
    For a given alpha, compute mean squared one-step-ahead prediction error.

    For each subject time series:
    - Run a causal EMA on PropCorrect (same recurrence as exp_running_mean
      in Brains_predict_test_from_training.py).
    - After seeing bin t, use the EMA state to predict PropCorrect at bin t+1.
    The loss is the average of (y[t+1] - EMA_t)^2 over all subjects and bins.
    """
    total_sq_err = 0.0
    total_count = 0

    for sp, series_list in series_by_species.items():
        for y in series_list:
            # y: 1D array of PropCorrect over bins
            # initialise state with first observation
            s = float(y[0])
            # predict y[t+1] from state after seeing y[t]
            for t in range(0, len(y) - 1):
                # update state with current bin
                s = (1.0 - alpha) * s + alpha * float(y[t])
                pred = s
                err = float(y[t+1]) - pred
                total_sq_err += err * err
                total_count += 1

    if total_count == 0:
        return np.inf
    return total_sq_err / total_count  # mean squared error

# ---------- scan alpha grid and find best ----------
# We evaluate predictive_loss(alpha) on a simple grid from 0.05 to 0.95.
# The alpha with the smallest loss is our fitted EMA coefficient.
# After running this script, copy best_alpha into the main analysis as EXP_ALPHA.
#
# Grid of EMA coefficients to evaluate: 0.05, 0.10, ..., 0.95.
alphas = np.linspace(0.05, 0.95, 19)  # 0.05, 0.10, ..., 0.95
losses = []

print("\nScanning alpha values:")
for a in alphas:
    L = predictive_loss(a, series_by_species)
    losses.append(L)
    print(f"alpha = {a:5.2f},  MSE = {L:.6f}")

losses = np.array(losses)
best_idx = int(np.argmin(losses))
best_alpha = float(alphas[best_idx])
best_loss = float(losses[best_idx])

# Map alpha to a continuous-time time constant tau (in bins) using the exact
# relation for an exponential filter: alpha = 1 - exp(-1/tau).
tau = -1.0 / np.log(1.0 - best_alpha)

print("\n=== BEST EMA PARAMETER FROM PREDICTIVE FIT ===")
print(f"best alpha = {best_alpha:.4f}")
print(f"min MSE    = {best_loss:.6f}")
print(f"tau (bins) = {tau:.3f}")

# Note: best_alpha is the recommended EXP_ALPHA value for the main analysis script
# when TTC_MODE == "ema".

print("\nNeighbourhood around best alpha (grid values):")
start = max(0, best_idx - 2)
end = min(len(alphas), best_idx + 3)
for i in range(start, end):
    mark = "*" if i == best_idx else " "
    print(f"{mark} alpha = {alphas[i]:.4f},  MSE = {losses[i]:.6f}")

plt.figure()
plt.plot(alphas, losses, marker="o", linestyle="-")
plt.axvline(best_alpha, linestyle="--")
plt.xlabel("alpha")
plt.ylabel("MSE (one-step-ahead)")
plt.title("EMA predictive loss vs alpha")
plt.tight_layout()
plt.show()