diff --git a/README.md b/README.md index a79628d..b76747a 100644 --- a/README.md +++ b/README.md @@ -90,13 +90,30 @@ and **Bonferroni** correction for planned comparisons. NeuroAI-Project13/ ├── README.md ├── environment.yml # pinned conda environment +├── requirements.txt # pip dependencies +├── pyproject.toml ├── src/ -│ ├── models/ # CTRNN definition + Euler integration +│ ├── models/ # CTRNN definition and architecture │ ├── training/ # training loop, loss functions, BPTT -│ └── analysis/ # PID / ΦID / MI / Fisher pipelines -├── notebooks/ # exploratory analysis + throwaway/example files -├── results/ # saved model weights/activations, metric outputs +│ ├── tasks/ # NeuroGym task wrappers and stimulus preprocessing +│ └── analysis/ # PID computation, stats, plotting +├── notebooks/ # exploratory scratch analysis + throwaway/example files +├── results/ +│ ├── model_weights/ # saved .pt checkpoint files per CTRNN seed +│ │ ├── perceptual/ +│ │ └── context/ +│ ├── model_activations/ # saved hidden state tensors per CTRNN seed and trial +│ │ ├── perceptual/ +│ │ └── context/ +│ ├── stimulus_coherences/ # saved coherence arrays per CTRNN seed and trial +│ │ ├── perceptual/ +│ │ └── context/ +│ └── pid_outputs/ # saved PID atom arrays per CTRNN seed +│ ├── perceptual/ +│ └── context/ ├── figures/ # final figures +│ ├── learning_curves/ # CTRNN train/test loss and accuracy curves +│ └── pid/ # PID atom plots └── docs/ # technical note, references ``` diff --git a/figures/.gitkeep b/figures/learning_curves/.gitkeep similarity index 100% rename from figures/.gitkeep rename to figures/learning_curves/.gitkeep diff --git a/figures/pid/.gitkeep b/figures/pid/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt index e3d05d2..ded76ed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,11 +3,11 @@ torch==2.3.1; python_version >= '3.11' --extra-index-url https://download.pytorch.org/whl/cu121 # Numerics & stats -numpy==1.26.4; python_version >= '3.11' +numpy==2.2.6; python_version >= '3.11' scipy==1.13.1; python_version >= '3.11' # Visualization -matplotlib==3.9.2; python_version >= '3.11' +matplotlib==3.10.9; python_version >= '3.11' seaborn==0.13.2; python_version >= '3.11' # NeuroGym tasks diff --git a/results/stimulus_coherences/context/.gitkeep b/results/stimulus_coherences/context/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/results/stimulus_coherences/perceptual/.gitkeep b/results/stimulus_coherences/perceptual/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/src/analysis/PID_figures.py b/src/analysis/PID_figures.py index 3cd4671..7ead7d5 100644 --- a/src/analysis/PID_figures.py +++ b/src/analysis/PID_figures.py @@ -38,7 +38,43 @@ plt.gca().spines['top'].set_visible(False) plt.gca().spines['right'].set_visible(False) plt.tight_layout() -plt.savefig(f'{dir}pid_bars_{SYSTEM}.png', dpi=300) +# plt.savefig(f'{dir}pid_bars_{SYSTEM}.png', dpi=300) plt.show() + + +# Measure Gaussianity of hidden unit activations across all timesteps, for each unit. +import torch +from scipy import stats + +h = torch.load(f'{dir}hidden_Sinusoid.pt', weights_only=True).numpy() # (300, 10) +n_units = h.shape[1] + +fig, axes = plt.subplots(2, 5, figsize=(14, 5)) +axes = axes.ravel() + +for i in range(n_units): + ax = axes[i] + unit_acts = h[:, i] # 300 values across timesteps + + # histogram of activations + ax.hist(unit_acts, bins=20, color='#4C78A8', alpha=0.75, density=True, zorder=2) + + # overlay best-fit Gaussian for visual comparison + mu, sigma = unit_acts.mean(), unit_acts.std() + x = np.linspace(unit_acts.min(), unit_acts.max(), 200) + ax.plot(x, stats.norm.pdf(x, mu, sigma), color='#E45756', lw=1.8, zorder=3) + + # Shapiro-Wilk p-value (tests departure from normality; p<0.05 = non-Gaussian) + _, p = stats.shapiro(unit_acts) + ax.set_title(f'Unit {i+1} p={p:.2f}', fontsize=9) + ax.tick_params(labelsize=7) + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + +fig.suptitle('Hidden unit activation distributions (blue=data, red=fitted Gaussian)\n' + 'Shapiro-Wilk p-value: p>0.05 consistent with Gaussian', fontsize=11) +fig.tight_layout() +# fig.savefig('unit_activations_gaussianity.png', dpi=150, bbox_inches='tight') +plt.show() \ No newline at end of file diff --git a/src/analysis/RNN.py b/src/analysis/RNN.py deleted file mode 100644 index 07b3b18..0000000 --- a/src/analysis/RNN.py +++ /dev/null @@ -1,58 +0,0 @@ -import torch -import torch.nn as nn - - -class ElmanRNN(nn.Module): - def __init__(self, dim, system='VanDerPol', hidden_dim=None, num_layers=1): - super(ElmanRNN, self).__init__() - ############## - ############## - - # store basic info - self.dim = dim - self.system = system - - # choose hidden size (you can tune these) - if hidden_dim is None: - hidden_dim = 10 - - self.hidden_dim = hidden_dim - self.num_layers = num_layers - - # classic RNN cell with tanh nonlinearity and a linear readout - self.rnn = nn.RNN( - input_size=dim, - hidden_size=hidden_dim, - num_layers=num_layers, - nonlinearity="tanh", - batch_first=False, - ) - self.readout = nn.Linear(hidden_dim, dim) - - ############## - ############## - - def forward(self, x, h0=None): - ############## - ############## - - # x is expected to have shape (n_timesteps, dim) - # or (dim,) for a single state in testing mode - squeezed = False - if x.dim() == 1: - x = x.unsqueeze(0) - squeezed = True - - # add batch dimension for the RNN: (seq_len, batch, dim) - x = x.unsqueeze(1) - out, hn = self.rnn(x, h0) - y = self.readout(out).squeeze(1) - - if squeezed: - y = y.squeeze(0) - - ############## - ############## - return y, hn - - \ No newline at end of file diff --git a/src/analysis/compute_pid_sinusoid.py b/src/analysis/compute_pid_sinusoid.py deleted file mode 100644 index 11414c2..0000000 --- a/src/analysis/compute_pid_sinusoid.py +++ /dev/null @@ -1,124 +0,0 @@ -""" -compute_pid_sinusoid.py -======================= - -Test the Gaussian analytic PID implementation on the sinusoid-trained RNN. - -Goal: compute PID at the "last timestep" of the RNN's hidden-state trajectory, -using the hidden activations as the two source groups and the input sinusoid as -the target. - -Important modelling note ------------------------- -PID is estimated from a "covariance over an ensemble of observations". A single -timestep of a single trajectory is one observation, which cannot define a -covariance. The sinusoid RNN was trained on one trajectory, so to obtain a valid -"PID at the last timestep" we must build an ensemble. We do this exactly the way -the our NeuroGym project will: we run the trained RNN over many "trials" and use -the trial dimension as the sample/ensemble axis. Each trial is the same sinusoid -with a random phase and a small amount of input noise (mirroring the noisy, -parametrically-varied stimuli of the decision-making tasks). The PID at the last -timestep is then computed across trials: - - sources = hidden activations h(T) at the last timestep (n_trials x n_units) - target = (clean) sinusoid value at the last timestep (n_trials,) - -Run `main.py` first so that `RNN_Sinusoid.pt` exists. -""" - -# Track run time -import time -start_time = time.time() - -import torch -import numpy as np -import matplotlib.pyplot as plt -from RNN import ElmanRNN -from gaussian_pid import gaussian_pid_rnn -# Define input data directory -dir = "src/analysis/results/" - -# --------------------------------------------------------------------------- # -# Settings (kept explicit so every choice is visible) -# --------------------------------------------------------------------------- # -SYSTEM = "Sinusoid" # which trained model file to load (RNN_.pt) -DIM = 1 # the sinusoid is a single (1-D) channel -N_TIMESTEPS = 300 # trajectory length, must match how the RNN was trained -PERIOD = 50.0 # sinusoid period in timesteps, must match training -N_TRIALS = 500 # ensemble size: number of phase-randomised trials -INPUT_NOISE_STD = 0.05 # std of additive input noise per timestep (keeps MI finite) -N_BIPARTITIONS = 200 # number of random 5/5 unit splits to average PID over -SEED = 0 # reproducibility for phases, noise and bipartitions -LOG_BASE = 2 # report information in bits - -if __name__ == "__main__": - - # --- reproducibility ---------------------------------------------------- # - rng = np.random.default_rng(SEED) # numpy RNG for phases/noise - torch.manual_seed(SEED) # torch RNG (defensive; eval is deterministic) - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # CPU is fine - - # --- load the trained sinusoid RNN -------------------------------------- # - # The constructor's default hidden size is 10, matching the trained model. - model = ElmanRNN(dim=DIM, system=SYSTEM).to(device) # build the architecture - model_path = f"{dir}RNN_{SYSTEM}.pt" # path to the saved weights - state = torch.load(model_path, map_location=device, weights_only=True) # load weights - model.load_state_dict(state) # restore the parameters - model.eval() # inference mode - n_units = model.hidden_dim # number of hidden units (10) - - # --- build the trial ensemble (phase-randomised, slightly noisy sinusoids) # - t = np.arange(N_TIMESTEPS) # timestep index 0..T-1 - phases = rng.uniform(0.0, PERIOD, size=N_TRIALS) # one random phase per trial - # clean target signal per trial and timestep: sin(2*pi*(t + phase)/period) - clean = np.sin(2 * np.pi * (t[None, :] + phases[:, None]) / PERIOD) # (n_trials, T) - # the RNN input is the clean signal plus small observation noise - noise = INPUT_NOISE_STD * rng.standard_normal((N_TRIALS, N_TIMESTEPS)) # (n_trials, T) - inputs_np = clean + noise # noisy input stream (n_trials, T) - - # --- run the RNN over the whole ensemble in one batched pass ------------ # - # nn.RNN (batch_first=False) expects input shaped (seq_len, batch, input_size). - x = torch.tensor(inputs_np.T[:, :, None], dtype=torch.float32, device=device) # (T, n_trials, 1) - with torch.no_grad(): # no gradients needed at eval - h_seq, _ = model.rnn(x) # hidden states: (T, n_trials, n_units) - # reorder to the wrapper's convention (n_samples, n_timesteps, n_units) - activations = h_seq.permute(1, 0, 2).cpu().numpy() # (n_trials, T, n_units) - - # --- target: the clean sinusoid value (the underlying signal, like coherence) # - target = clean # (n_trials, T); univariate per step - - # --- compute PID at the LAST timestep, averaged over random 5/5 splits --- # - out = gaussian_pid_rnn( - activations, # sources: (n_trials, T, n_units) - target, # target : (n_trials, T) - timestep=-1, # -1 selects the last timestep - bipartitions="random", # average over random balanced unit splits - n_bipartitions=N_BIPARTITIONS, - balanced=True, # 5 units vs 5 units for a 10-unit RNN - seed=SEED, - log_base=LOG_BASE, # bits - ) - - # --- report ------------------------------------------------------------- # - unit = "bits" if LOG_BASE == 2 else "nats" - print(f"Gaussian analytic PID at the last timestep (t = {N_TIMESTEPS - 1})") - print(f" ensemble: {N_TRIALS} phase-randomised trials, " - f"{n_units} units split 5/5 over {N_BIPARTITIONS} bipartitions") - print(f" (values in {unit}; '+/-' is the spread across bipartitions)\n") - print(f" Redundancy : {out['redundancy']:.4f} +/- {out['redundancy_std']:.4f}") - print(f" Unique 1 : {out['unique1']:.4f} +/- {out['unique1_std']:.4f}") - print(f" Unique 2 : {out['unique2']:.4f} +/- {out['unique2_std']:.4f}") - print(f" Synergy : {out['synergy']:.4f} +/- {out['synergy_std']:.4f}") - print(f" -----") - print(f" I(X1;Y) : {out['mi_1']:.4f}") - print(f" I(X2;Y) : {out['mi_2']:.4f}") - print(f" I(X1,X2;Y) : {out['mi_joint']:.4f} (= sum of the four atoms)") - - # Report run time - end_time = time.time() - elapsed = end_time - start_time - print(f"\nTotal run time: {elapsed:.2f} seconds") - - # Save PID results - np.savez(f"{dir}pid_{SYSTEM}.npz", **out) \ No newline at end of file diff --git a/src/analysis/main.py b/src/analysis/main.py deleted file mode 100644 index 46439a8..0000000 --- a/src/analysis/main.py +++ /dev/null @@ -1,219 +0,0 @@ -import torch -import torch.nn as nn -import torch.optim as optim -import matplotlib.pyplot as plt -from RNN import ElmanRNN - -# Define output directory -dir = "src/analysis/results/" - -# hyperparameters -BATCH_SIZE = 1 # not really used here, since we train on a single trajectory -EPOCHS = 2500 -LEARNING_RATE = 5e-3 - -# Choose dynamical system -system = 'Sinusoid' - - -def r2_score(y_pred, y_true): - """ - Coefficient of determination (R²) between predictions and targets. - - R² measures the fraction of variance in y_true that is explained by - y_pred, and is the standard 'accuracy' analog for regression tasks: - R² = 1 → perfect prediction - R² = 0 → model is no better than predicting the mean of y_true - R² < 0 → model is worse than predicting the mean - - We use R² rather than normalised MSE because it is scale-free and has a - natural [−∞, 1] range with a meaningful zero point, making curves from - different runs directly comparable. - - Parameters - ---------- - y_pred : torch.Tensor (any shape) - y_true : torch.Tensor (same shape as y_pred) - - Returns - ------- - float - """ - # total variance of the targets (denominator) - ss_tot = ((y_true - y_true.mean()) ** 2).sum() - # residual variance left unexplained by the predictions (numerator) - ss_res = ((y_true - y_pred) ** 2).sum() - # R² = 1 − (unexplained / total); clamp denominator to avoid div-by-zero - return (1.0 - ss_res / ss_tot.clamp(min=1e-8)).item() - - -if __name__ == "__main__": - - # Generate a simple sinusoid trajectory to test the PID metrics. Size (300, 1) - n_timesteps = 300 - period = 50.0 - t = torch.linspace(0, n_timesteps - 1, n_timesteps) - y_true = torch.sin(2 * torch.pi * t / period).unsqueeze(1) - # print(y_true.shape) - - n_timesteps, dim = y_true.shape - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - y_true = y_true.to(device) - - # Create RNN - model = ElmanRNN(dim=dim, system=system).to(device) - - # Optimizer and loss - optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE) # create the optimizer - loss_function = nn.MSELoss() # loss function (MSE) - - # For tracking over epochs - training_losses = [] - training_r2s = [] # R² at each epoch (teacher-forced predictions vs targets) - - # ===== TRAINING LOOP ===== - for epoch in range(EPOCHS): - model.train() - if epoch == 1000: - for param_group in optimizer.param_groups: - param_group["lr"] = 1e-3 - if epoch == 2000: - for param_group in optimizer.param_groups: - param_group["lr"] = 5e-4 - optimizer.zero_grad() - - # Teacher forcing over the full sequence: predict x_{t+1} from x_t - x_in = y_true[:-1] - target = y_true[1:] - pred, _ = model(x_in) - loss = loss_function(pred, target) - - # Backpropagate once after summing over all timesteps - loss.backward() - optimizer.step() - - # Track loss and R² for this epoch. - # R² is computed on the same teacher-forced predictions used for the - # loss, so it reflects how well the model fits the one-step-ahead - # prediction task it is actually being trained on. - training_losses.append(loss.item()) - training_r2s.append(r2_score(pred.detach(), target)) - - if (epoch + 1) % 50 == 0 or epoch == 0: - print( - f"Epoch {epoch+1}/{EPOCHS}, " - f"training loss = {loss.item():.6e}, " - f"training R² = {training_r2s[-1]:.4f}" - ) - - # Save trained RNN - torch.save(model.state_dict(), f"{dir}RNN_{system}.pt") - print(f"Trained RNN saved to {dir}RNN_{system}.pt") - - - # ===== EVALUATION: AUTOREGRESSIVE ROLLOUT ===== - model.eval() - with torch.no_grad(): - y_pred = torch.zeros_like(y_true) - # initial condition: first value in dataset - y_pred[0] = y_true[0] - - # always feed prediction at n-1 as input to predict state at n - h = None - for t in range(n_timesteps - 1): - step_pred, h = model(y_pred[t], h0=h) - y_pred[t+1] = step_pred - - # Evaluation loss and R² on the autoregressive rollout. - # This is harder than teacher-forced training: the model must sustain - # the oscillation from its own outputs, so errors accumulate over time. - mse_eval = loss_function(y_pred, y_true).item() - r2_eval = r2_score(y_pred, y_true) - print(f"Evaluation MSE over full trajectory: {mse_eval:.6e}") - print(f"Evaluation R² over full trajectory: {r2_eval:.4f}") - - - # ===== COLLECT HIDDEN STATES (for testing the PID metrics) ===== - model.eval() - with torch.no_grad(): - # full hidden-state trajectory under teacher forcing over the sinusoid - x_seq = y_true.unsqueeze(1) # (seq_len, batch=1, dim) - h_seq, _ = model.rnn(x_seq) # (seq_len, batch=1, hidden_dim) - h_seq = h_seq.squeeze(1) # (seq_len, hidden_dim) - torch.save(h_seq.cpu(), f"{dir}hidden_{system}.pt") - print(f"Hidden states {tuple(h_seq.shape)} saved to {dir}hidden_{system}.pt") - - - # ===== FIGURE 1: TRAINING CURVES (loss + R²) ===== - # Both metrics are plotted against epoch on a shared x-axis. - # Loss (left y-axis) uses a log scale so the steep early drop and the - # slow late convergence are both visible. R² (right y-axis) is linear. - epochs_axis = range(1, EPOCHS + 1) - - fig, ax1 = plt.subplots(figsize=(10, 5)) - - # Left axis: MSE loss (log scale) - color_loss = 'C0' - ax1.set_xlabel('Epoch') - ax1.set_ylabel('MSE loss (log scale)', color=color_loss) - ax1.semilogy(epochs_axis, training_losses, color=color_loss, - linewidth=1.2, label='Training loss') - ax1.tick_params(axis='y', labelcolor=color_loss) - - # Right axis: R² - ax2 = ax1.twinx() - color_r2 = 'C1' - ax2.set_ylabel('R²', color=color_r2) - ax2.plot(epochs_axis, training_r2s, color=color_r2, - linewidth=1.2, label='Training R²') - ax2.tick_params(axis='y', labelcolor=color_r2) - ax2.set_ylim(-0.05, 1.05) # R² lives in (−∞, 1]; clip low outliers - - # Combined legend for both axes - lines1, labels1 = ax1.get_legend_handles_labels() - lines2, labels2 = ax2.get_legend_handles_labels() - ax1.legend(lines1 + lines2, labels1 + labels2, loc='center right') - - plt.title(f'Training curves — {system} RNN') - plt.tight_layout() - plt.savefig(f"{dir}sinusoid_training_curves_{system}.png", dpi=300) - plt.show() - - - # ===== FIGURE 2: TRUE VS PREDICTED TRAJECTORY ===== - y_true_np = y_true.detach().cpu().numpy() - y_pred_np = y_pred.detach().cpu().numpy() - time = range(n_timesteps) - - # Plot for Sinusoid system - if system == 'Sinusoid': - fig, ax = plt.subplots(figsize=(10, 6)) - ax.plot(time, y_true_np[:, 0], label='y_true', c='C0') - ax.plot(time, y_pred_np[:, 0], '--', label='y_predicted', c='C1') - - # Evaluation metrics shown as a text box in the upper-right corner. - # Using a semi-transparent box so the label is readable over the curve. - metrics_text = ( - f"Eval MSE = {mse_eval:.4e}\n" - f"Eval R² = {r2_eval:.4f}" - ) - ax.text( - 0.98, 0.97, # upper-right in axes-fraction coordinates - metrics_text, - transform=ax.transAxes, # interpret x,y as fractions of the axes - fontsize=10, - verticalalignment='top', - horizontalalignment='right', - bbox=dict(boxstyle='round,pad=0.4', facecolor='white', - alpha=0.75, edgecolor='0.7'), - ) - - ax.set_ylabel('y') - ax.set_title(f'Sinusoid trajectory: true vs predicted ({system} RNN)') - ax.set_xlabel('Time (s)') - ax.legend(loc='upper left') - - plt.tight_layout() - plt.savefig(f"{dir}sinusoid_trajectory_{system}.png", dpi=300) - plt.show() diff --git a/src/analysis/results/RNN100_Sinusoid.pt b/src/analysis/results/RNN100_Sinusoid.pt deleted file mode 100644 index b6d48e0..0000000 Binary files a/src/analysis/results/RNN100_Sinusoid.pt and /dev/null differ diff --git a/src/analysis/results/RNN_Sinusoid.pt b/src/analysis/results/RNN_Sinusoid.pt deleted file mode 100644 index 31c507c..0000000 Binary files a/src/analysis/results/RNN_Sinusoid.pt and /dev/null differ diff --git a/src/analysis/results/hidden100_Sinusoid.pt b/src/analysis/results/hidden100_Sinusoid.pt deleted file mode 100644 index 5e3d791..0000000 Binary files a/src/analysis/results/hidden100_Sinusoid.pt and /dev/null differ diff --git a/src/analysis/results/hidden_Sinusoid.pt b/src/analysis/results/hidden_Sinusoid.pt deleted file mode 100644 index 8d6bcab..0000000 Binary files a/src/analysis/results/hidden_Sinusoid.pt and /dev/null differ diff --git a/src/analysis/results/pid100_Sinusoid.npz b/src/analysis/results/pid100_Sinusoid.npz deleted file mode 100644 index d3ada93..0000000 Binary files a/src/analysis/results/pid100_Sinusoid.npz and /dev/null differ diff --git a/src/analysis/results/pid100_bars_Sinusoid.png b/src/analysis/results/pid100_bars_Sinusoid.png deleted file mode 100644 index 9274db3..0000000 Binary files a/src/analysis/results/pid100_bars_Sinusoid.png and /dev/null differ diff --git a/src/analysis/results/pid_Sinusoid.npz b/src/analysis/results/pid_Sinusoid.npz deleted file mode 100644 index ae90ed7..0000000 Binary files a/src/analysis/results/pid_Sinusoid.npz and /dev/null differ diff --git a/src/analysis/results/pid_bars_Sinusoid.png b/src/analysis/results/pid_bars_Sinusoid.png deleted file mode 100644 index a9cc18e..0000000 Binary files a/src/analysis/results/pid_bars_Sinusoid.png and /dev/null differ diff --git a/src/analysis/results/sinusoid100_training_curves_Sinusoid.png b/src/analysis/results/sinusoid100_training_curves_Sinusoid.png deleted file mode 100644 index 70e9505..0000000 Binary files a/src/analysis/results/sinusoid100_training_curves_Sinusoid.png and /dev/null differ diff --git a/src/analysis/results/sinusoid100_trajectory_Sinusoid.png b/src/analysis/results/sinusoid100_trajectory_Sinusoid.png deleted file mode 100644 index 380a64b..0000000 Binary files a/src/analysis/results/sinusoid100_trajectory_Sinusoid.png and /dev/null differ diff --git a/src/analysis/results/sinusoid_training_curves_Sinusoid.png b/src/analysis/results/sinusoid_training_curves_Sinusoid.png deleted file mode 100644 index e54f4f9..0000000 Binary files a/src/analysis/results/sinusoid_training_curves_Sinusoid.png and /dev/null differ diff --git a/src/analysis/results/sinusoid_trajectory_Sinusoid.png b/src/analysis/results/sinusoid_trajectory_Sinusoid.png deleted file mode 100644 index 120f580..0000000 Binary files a/src/analysis/results/sinusoid_trajectory_Sinusoid.png and /dev/null differ diff --git a/src/analysis/test_gaussianity.py b/src/analysis/test_gaussianity.py new file mode 100644 index 0000000..b9210f3 --- /dev/null +++ b/src/analysis/test_gaussianity.py @@ -0,0 +1,212 @@ +""" +test_gaussianity.py +=================== + +Tests whether the hidden-unit activations of the sinusoid-trained RNN are +approximately Gaussian distributed across trials at the final timestep. + +Why this is the right check +--------------------------- +Gaussian PID assumes that the joint distribution of (X1, X2, Y) is multivariate +Gaussian across the *trial* ensemble. The relevant random variable for each unit +is therefore h_i(t*) evaluated across many independent trials (each with a +different phase and noise draw), NOT the activation of unit i across timesteps +within a single trajectory. The latter is autocorrelated and mixes temporal +dynamics with the distributional question; the former is the actual sample the +covariance matrix is built from. + +What this script does +--------------------- +1. Rebuilds the same 500-trial ensemble used in compute_pid_sinusoid.py. +2. Extracts the across-trial snapshot at the final timestep: shape (500, 10). +3. For each unit: + - plots a histogram of its 500 across-trial activations + - overlays the best-fit Gaussian + - runs a Shapiro-Wilk normality test (exact for n<=5000) + - runs a D'Agostino-Pearson K² test (sensitive to skew + kurtosis) + - prints both p-values; p > 0.05 is consistent with Gaussianity +4. Produces a Q-Q plot for each unit (points on the diagonal = Gaussian). +5. Saves both figures. + +Run `main.py` first so that `RNN_.pt` exists. +""" + +import time +start_time = time.time() + +import torch +import numpy as np +import matplotlib.pyplot as plt +from scipy import stats +from RNN import ElmanRNN + +# Define directories +dir = "src/analysis/results/" + +# --------------------------------------------------------------------------- # +# Settings — must match compute_pid_sinusoid.py exactly so the ensemble is +# identical and the gaussianity check applies to the same data PID is run on +# --------------------------------------------------------------------------- # +SYSTEM = "Sinusoid" +DIM = 1 +N_TIMESTEPS = 300 +PERIOD = 50.0 +N_TRIALS = 500 +INPUT_NOISE_STD = 0.05 +SEED = 0 + +if __name__ == "__main__": + + # --- reproducibility ---------------------------------------------------- # + rng = np.random.default_rng(SEED) + torch.manual_seed(SEED) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # --- load the trained RNN ----------------------------------------------- # + model = ElmanRNN(dim=DIM, system=SYSTEM).to(device) + model_path = f"{dir}RNN_{SYSTEM}.pt" + state = torch.load(model_path, map_location=device, weights_only=True) + model.load_state_dict(state) + model.eval() + n_units = model.hidden_dim # 10 + + # --- build the identical trial ensemble --------------------------------- # + # Must use the same SEED as compute_pid_sinusoid.py so this is checking the + # exact distribution PID was estimated from. + t_idx = np.arange(N_TIMESTEPS) + phases = rng.uniform(0.0, PERIOD, size=N_TRIALS) + clean = np.sin(2 * np.pi * (t_idx[None, :] + phases[:, None]) / PERIOD) + noise = INPUT_NOISE_STD * rng.standard_normal((N_TRIALS, N_TIMESTEPS)) + inputs_np = clean + noise # (n_trials, T) + + # --- run the RNN over all trials in one batched pass ------------------- # + x = torch.tensor(inputs_np.T[:, :, None], dtype=torch.float32, device=device) + with torch.no_grad(): + h_seq, _ = model.rnn(x) # (T, n_trials, n_units) + activations = h_seq.permute(1, 0, 2).cpu().numpy() # (n_trials, T, n_units) + + # --- extract the snapshot at the FINAL timestep ------------------------- # + # This is the (500, 10) matrix whose covariance Gaussian PID is built from. + # Each column is one unit's activation across 500 independent trials at t*. + snapshot = activations[:, -1, :] # (n_trials, n_units) + + # --- normality tests per unit ------------------------------------------- # + print(f"Normality tests on across-trial activations at t = {N_TIMESTEPS - 1}") + print(f"(n = {N_TRIALS} trials per unit)\n") + print(f"{'Unit':<6} {'SW p-val':>10} {'SW pass':>8} {'DP p-val':>10} {'DP pass':>8}") + print("-" * 52) + + sw_pvals = [] + dp_pvals = [] + for i in range(n_units): + unit_acts = snapshot[:, i] # 500 across-trial values + + # Shapiro-Wilk: general normality test, exact for n<=5000 + _, sw_p = stats.shapiro(unit_acts) + + # D'Agostino-Pearson K²: tests skewness and excess kurtosis jointly; + # complementary to Shapiro-Wilk which is more sensitive to tail behaviour + _, dp_p = stats.normaltest(unit_acts) + + sw_pvals.append(sw_p) + dp_pvals.append(dp_p) + + sw_pass = "YES" if sw_p > 0.05 else "NO " + dp_pass = "YES" if dp_p > 0.05 else "NO " + print(f" {i+1:<4} {sw_p:>10.4f} {sw_pass:>8} {dp_p:>10.4f} {dp_pass:>8}") + + n_pass_sw = sum(p > 0.05 for p in sw_pvals) + n_pass_dp = sum(p > 0.05 for p in dp_pvals) + print(f"\n {n_pass_sw}/{n_units} units pass Shapiro-Wilk at α=0.05") + print(f" {n_pass_dp}/{n_units} units pass D'Agostino-Pearson at α=0.05") + print(f"\n Interpretation: p > 0.05 means we cannot reject Gaussianity.") + print(f" Units that fail both tests are the most problematic for Gaussian PID.") + + # ======================================================================== # + # FIGURE 1: histograms + fitted Gaussian overlay, one panel per unit + # ======================================================================== # + fig1, axes = plt.subplots(2, 5, figsize=(14, 5.5)) + axes = axes.ravel() + + for i in range(n_units): + ax = axes[i] + vals = snapshot[:, i] + mu, sigma = vals.mean(), vals.std() + + # histogram of across-trial activations + ax.hist(vals, bins=25, color='#4C78A8', alpha=0.75, + density=True, zorder=2) + + # best-fit Gaussian overlay + x_fit = np.linspace(vals.min(), vals.max(), 300) + ax.plot(x_fit, stats.norm.pdf(x_fit, mu, sigma), + color='#E45756', lw=1.8, zorder=3) + + # Shapiro-Wilk result in title; red title = non-Gaussian + sw_p = sw_pvals[i] + color = '#333333' if sw_p > 0.05 else '#E45756' + ax.set_title(f'Unit {i+1} SW p={sw_p:.3f}', + fontsize=9, color=color) + ax.tick_params(labelsize=7) + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.set_xlabel('activation', fontsize=7) + + fig1.suptitle( + f'Across-trial activation distributions at t* = {N_TIMESTEPS-1} ' + f'(N={N_TRIALS} trials)\n' + 'Blue = data histogram Red = fitted Gaussian ' + 'Red title = Shapiro-Wilk rejects Gaussianity (p < 0.05)', + fontsize=10 + ) + fig1.tight_layout() + # fig1.savefig(f'{dir}gaussianity_histograms_{SYSTEM}.png', + # dpi=300, bbox_inches='tight') + + # ======================================================================== # + # FIGURE 2: Q-Q plots, one panel per unit + # A Q-Q plot maps the empirical quantiles against the theoretical Gaussian + # quantiles. Points lying on the diagonal line indicate Gaussianity. + # Deviations at the tails indicate heavier/lighter tails than Gaussian. + # ======================================================================== # + fig2, axes2 = plt.subplots(2, 5, figsize=(14, 5.5)) + axes2 = axes2.ravel() + + for i in range(n_units): + ax = axes2[i] + vals = snapshot[:, i] + + # stats.probplot returns (quantiles, (slope, intercept, r)) and plots + # onto the given axis; r is the correlation with the Gaussian line + (osm, osr), (slope, intercept, r) = stats.probplot(vals, dist='norm') + ax.scatter(osm, osr, s=4, alpha=0.5, color='#4C78A8', zorder=2) + # diagonal reference line + x_line = np.array([osm[0], osm[-1]]) + ax.plot(x_line, slope * x_line + intercept, + color='#E45756', lw=1.5, zorder=3) + + sw_p = sw_pvals[i] + color = '#333333' if sw_p > 0.05 else '#E45756' + ax.set_title(f'Unit {i+1} r={r:.3f}', fontsize=9, color=color) + ax.tick_params(labelsize=7) + ax.set_xlabel('theoretical quantiles', fontsize=7) + ax.set_ylabel('sample quantiles', fontsize=7) + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + + fig2.suptitle( + f'Q-Q plots at t* = {N_TIMESTEPS-1} (N={N_TRIALS} trials)\n' + 'Points on the diagonal = Gaussian. ' + 'Red title = Shapiro-Wilk rejects Gaussianity (p < 0.05)', + fontsize=10 + ) + fig2.tight_layout() + # fig2.savefig(f'{dir}gaussianity_qq_{SYSTEM}.png', + # dpi=300, bbox_inches='tight') + + plt.show() + + end_time = time.time() + print(f"\nTotal run time: {end_time - start_time:.2f} seconds") + print(f"Figures saved to {dir}")