diff --git a/src/analysis/PID_figures.py b/src/analysis/PID_figures.py new file mode 100644 index 0000000..3cd4671 --- /dev/null +++ b/src/analysis/PID_figures.py @@ -0,0 +1,44 @@ +import matplotlib.pyplot as plt +import numpy as np + +# Define input data directory +dir = "src/analysis/results/" + +# Load PID results +SYSTEM = "Sinusoid" +pid_path = f"{dir}pid_{SYSTEM}.npz" +out = np.load(pid_path) + +# Plot PID bars with std error bars, and total MI as a title annotation +unit = 'bits' + +labels = ['Redundancy', 'Unique1', 'Unique2', 'Synergy'] +values = [out['redundancy'], out['unique1'], out['unique2'], out['synergy']] +stds = [out['redundancy_std'], out['unique1_std'], out['unique2_std'], out['synergy_std']] +total = out['mi_joint'] + +plt.figure(figsize=(8, 6)) +bars = plt.bar(labels, values, yerr=stds, capsize=5, width=0.6) +# Add value labels above bars +for bar, val, std in zip(bars, values, stds): + plt.text(bar.get_x() + bar.get_width() / 2, val + std + 0.02, f'{val:.3f}', ha='center', va='bottom', fontsize=12) +# Give different colors to each bar +colors = ['#4C78A8', '#F58518', '#E45756', '#54A24B'] +for bar, color in zip(bars, colors): + bar.set_color(color) + +plt.title(f'PID at last timestep - {SYSTEM} (10 unit) RNN\nTotal I(X1,X2;Y) = {total:.3f} {unit}', fontsize=15, fontweight='bold') +# Fontsize and limits +plt.ylabel(f'Information ({unit})', fontsize=15) +plt.yticks(fontsize=12) +plt.ylim(0, max(v + s for v, s in zip(values, stds)) * 1.1) +plt.xticks(fontsize=15) + +# Remove top and right edges +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.show() + + diff --git a/src/analysis/RNN.py b/src/analysis/RNN.py index a7d5361..07b3b18 100644 --- a/src/analysis/RNN.py +++ b/src/analysis/RNN.py @@ -2,9 +2,9 @@ import torch.nn as nn -class TrueRNN(nn.Module): +class ElmanRNN(nn.Module): def __init__(self, dim, system='VanDerPol', hidden_dim=None, num_layers=1): - super(TrueRNN, self).__init__() + super(ElmanRNN, self).__init__() ############## ############## @@ -55,77 +55,4 @@ def forward(self, x, h0=None): ############## return y, hn - def calculate_jacobian(self, x): - """" Calculates Jacobians along trajectory x. - For every timestep the Jacobian is calculated as the derivative of the predictions wrt. inputs - and has a shape (n_dim, n_dim) at every timestep. - The final Jacobian (return variable) contains the concatenated Jacobians for all timesteps - and has a shape (n_timesteps, n_dim, n_dim). - - Parameters: - ----------- - x : torch.Tensor - trajectory values, with shape (n_timesteps, n_dim) - - - Return: - ----------- - jacobian : torch.Tensor - computed jacobians for the given trajectory - shape = (n_timesteps, n_dim, n_dim) - - """ - jacobian = None - ############## - ############## - - assert x.dim() == 2, "x must have shape (n_timesteps, n_dim)" - n_timesteps, n_dim = x.shape - device = x.device - dtype = x.dtype - - # storage for all Jacobians - jacobian = torch.zeros(n_timesteps, n_dim, n_dim, device=device, dtype=dtype) - - # initial hidden state - h = torch.zeros(self.num_layers, 1, self.hidden_dim, device=device, dtype=dtype) - - # start from the first point of the trajectory as initial state - # then always use the model's own prediction as the next state - state = x[0].clone().detach().requires_grad_(True) - - for t in range(n_timesteps): - - # prediction of the next state from the current state - with torch.enable_grad(): - h = h.detach() - pred, h = self.forward(state, h0=h) # pred shape (n_dim,) - - # compute Jacobian d(pred) / d(state) - for i in range(n_dim): - grad_output = torch.zeros_like(pred) # Partial derivative of - # predicted dim i wrt to all previous state dims - grad_output[i] = 1.0 - # If dim=2, grad_output is [1,0] for i=0 and [0,1] for i=1 - # This allows to select which output dimension to differentiate, - # i.e., which predicted state (y1(t) or y2(t)) to differentiate wrt. - # the two input states (y1(t-1), y2(t-1)) hence, which row of - # the Jacobian to compute. - - grads = torch.autograd.grad( - outputs=pred, # vector, shape (n_dim,) - inputs=state, # vector, shape (n_dim,) - grad_outputs=grad_output, - retain_graph=(i < n_dim - 1), - create_graph=False, - allow_unused=False, - )[0] # same shape as state: (n_dim,) - - jacobian[t, i, :] = grads # row i of the Jacobian at time t - - # roll the system forward: next state is the prediction - state = pred.detach().requires_grad_(True) - - ############## - ############## - return jacobian + \ No newline at end of file diff --git a/src/analysis/compute_pid_sinusoid.py b/src/analysis/compute_pid_sinusoid.py index 8e1e0d4..11414c2 100644 --- a/src/analysis/compute_pid_sinusoid.py +++ b/src/analysis/compute_pid_sinusoid.py @@ -4,32 +4,39 @@ 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, +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 +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 real NeuroGym project will: we run the trained RNN over many *trials* and use +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 = the (clean) sinusoid value at the last timestep (n_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 -from trueRNN import TrueRNN +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) @@ -54,8 +61,8 @@ # --- load the trained sinusoid RNN -------------------------------------- # # The constructor's default hidden size is 10, matching the trained model. - model = TrueRNN(dim=DIM, system=SYSTEM).to(device) # build the architecture - model_path = f"RNN_{SYSTEM}.pt" # path to the saved weights + 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 @@ -107,3 +114,11 @@ 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/gaussian_pid.py b/src/analysis/gaussian_pid.py index 3122ebd..b6668a2 100644 --- a/src/analysis/gaussian_pid.py +++ b/src/analysis/gaussian_pid.py @@ -5,7 +5,7 @@ Gaussian analytic Partial Information Decomposition (PID), using the Minimum-Mutual-Information (MMI) redundancy of Barrett (2015), i.e. the same Gaussian machinery that the `phyid` library (Mediano/Rosas/Luppi) uses for its -Integrated Information Decomposition (PhiID), but specialised to *static* PID: +Integrated Information Decomposition (PhiID), but specialised to "static" PID: two source groups -> one target, instead of past -> future. Why this is the right reduction @@ -42,14 +42,14 @@ Known limitations (worth stating in the report) ------------------------------------------------ -* MMI forces one of the two unique terms to be exactly zero and forces +1) MMI forces one of the two unique terms to be exactly zero and forces redundancy to be the smaller of the two source-target MIs, regardless of how the sources actually relate. This is a rigidity of the MMI definition. -* Barrett's exactness result is for a *univariate* target. The code accepts a +2) Barrett's exactness result is for a "univariate" target. The code accepts a multivariate target mechanically, but the MMI interpretation is cleanest for a univariate target (which is what this project uses: signed coherence / stimulus value). -* The Gaussian assumption is only an approximation for bounded tanh activations. +3) The Gaussian assumption is only an approximation for bounded tanh activations. """ import numpy as np @@ -123,8 +123,8 @@ def gaussian_pid(sources_1, sources_2, target, """ Gaussian analytic MMI-PID for two source groups and one target. - This is the atomic primitive. It treats every *row* as one observation - (e.g. one trial) and every *column* as one variable (e.g. one RNN unit, or + This is the atomic primitive. It treats every "row" as one observation + (e.g. one trial) and every "column" as one variable (e.g. one RNN unit, or the scalar target). It estimates a single joint covariance over [sources_1 | sources_2 | target] and returns the four PID atoms in closed form. @@ -267,26 +267,26 @@ def gaussian_pid_rnn(activations, target, ---------- activations : array_like Hidden activations. Either - * (n_samples, n_timesteps, n_units) -> per-timestep PID is available, or - * (n_samples, n_units) -> already a single snapshot. + a) (n_samples, n_timesteps, n_units) -> per-timestep PID is available, or + b) (n_samples, n_units) -> already a single snapshot. `n_samples` is the ensemble dimension over which covariance is estimated (in the real project: trials; for a single-trajectory test RNN: time, if you pass the trajectory in the samples axis). target : array_like The target / stimulus. Either - * (n_samples, n_timesteps) (a value per trial and timestep), - * (n_samples,) (one value per trial, shared across time), or - * (n_samples, n_timesteps, dt) / (n_samples, dt) for a multivariate target. + a) (n_samples, n_timesteps) (a value per trial and timestep), + b) (n_samples,) (one value per trial, shared across time), or + c) (n_samples, n_timesteps, dt) / (n_samples, dt) for a multivariate target. timestep : int or None, optional - * int -> compute PID only at this timestep (supports negative indexing, + int -> compute PID only at this timestep (supports negative indexing, e.g. -1 = last timestep). Returns scalar atoms. - * None -> compute PID at every timestep. Returns arrays of shape + None -> compute PID at every timestep. Returns arrays of shape (n_timesteps,). Requires 3-D `activations`. bipartitions : {"random", "half"} or list of (idx1, idx2), optional How to split the units into two source groups: - * "random" (default): draw `n_bipartitions` random splits. - * "half": a single split into first-half / second-half units. - * explicit list: each element is a pair (idx1, idx2) of index arrays. + a) "random" (default): draw `n_bipartitions` random splits. + b) "half": a single split into first-half / second-half units. + c) explicit list: each element is a pair (idx1, idx2) of index arrays. n_bipartitions : int, optional Number of random splits to average over when bipartitions="random". balanced : bool, optional diff --git a/src/analysis/main.py b/src/analysis/main.py index e7bae1c..46439a8 100644 --- a/src/analysis/main.py +++ b/src/analysis/main.py @@ -2,18 +2,50 @@ import torch.nn as nn import torch.optim as optim import matplotlib.pyplot as plt -from trueRNN import TrueRNN +from RNN import ElmanRNN -# Define task directory -dir = "" +# Define output directory +dir = "src/analysis/results/" # hyperparameters BATCH_SIZE = 1 # not really used here, since we train on a single trajectory -EPOCHS = 1000 -LEARNING_RATE = 1e-3 +EPOCHS = 2500 +LEARNING_RATE = 5e-3 # Choose dynamical system -system = 'Sinusoid' # 'VanDerPol', 'Lorenz' or 'Sinusoid' +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__": @@ -30,20 +62,25 @@ y_true = y_true.to(device) # Create RNN - model = TrueRNN(dim=dim, system=system).to(device) + 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 + # For tracking over epochs training_losses = [] - jacobian_norm_max_list = [] - + 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 @@ -56,27 +93,18 @@ loss.backward() optimizer.step() - # ---- Track Jacobian after this epoch ---- - model.eval() - # do NOT use torch.no_grad() here, autograd is needed inside calculate_jacobian - with torch.enable_grad(): - # detach so we don't reuse the training graph - x_for_jacobian = y_true.detach() - jacobian = model.calculate_jacobian(x_for_jacobian) # shape: (n_timesteps, dim, dim) - # Frobenius norm over last two dims -> shape (n_timesteps,) - jacobian_norms = torch.linalg.norm(jacobian, dim=(1, 2)) - # maximum norm over all timesteps - jacobian_norm_max = jacobian_norms.max().item() - - # Print loss and max Jacobian norm every 50 epochs + # 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()) - jacobian_norm_max_list.append(jacobian_norm_max) + 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"max Jacobian norm = {jacobian_norm_max:.6e}" + f"training R² = {training_r2s[-1]:.4f}" ) # Save trained RNN @@ -97,8 +125,13 @@ 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) ===== @@ -112,55 +145,75 @@ print(f"Hidden states {tuple(h_seq.shape)} saved to {dir}hidden_{system}.pt") - # ===== PLOTTING TRUE VS PREDICTED TRAJECTORY ===== - y_true_np = y_true.detach().cpu().numpy() - y_pred_np = y_pred.detach().cpu().numpy() - time = range(n_timesteps) + # ===== 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) - # Plot for Van der Pol system - if system == 'VanDerPol': - plt.figure(figsize=(10, 6)) - plt.plot(time, y_true_np[:, 0], label='y1_true', c='C0') - plt.plot(time, y_pred_np[:, 0], '--', label='y1_predicted', c='C1') - plt.plot(time, y_true_np[:, 1], label='y2_true', c='C0') - plt.plot(time, y_pred_np[:, 1], '--', label='y2_predicted', c='C1') - - plt.ylabel(f'y1, y2') - plt.title('Van der Pol trajectory: true vs predicted') - plt.xlabel('Time (s)') - plt.legend(loc='best') - - plt.tight_layout() - plt.show() + fig, ax1 = plt.subplots(figsize=(10, 5)) - # Plot for Lorenz system - elif system == 'Lorenz': - from mpl_toolkits.mplot3d import Axes3D # needed for 3D plotting + # 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) - fig = plt.figure(figsize=(10, 6)) - ax = fig.add_subplot(111, projection='3d') - ax.plot(y_true_np[:, 0], y_true_np[:, 1], y_true_np[:, 2], label='True trajectory', c='C0') - ax.plot(y_pred_np[:, 0], y_pred_np[:, 1], y_pred_np[:, 2], '--', label='Predicted trajectory', c='C1') + # 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 - ax.set_xlabel('X') - ax.set_ylabel('Y') - ax.set_zlabel('Z') - ax.set_title('Lorenz trajectory: true vs predicted') - ax.legend(loc='best') + # 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.tight_layout() - plt.show() + plt.title(f'Training curves — {system} RNN') + plt.tight_layout() + plt.savefig(f"{dir}sinusoid_training_curves_{system}.png", dpi=300) + plt.show() - # Plot for Sinusoid system - elif system == 'Sinusoid': - plt.figure(figsize=(10, 6)) - plt.plot(time, y_true_np[:, 0], label='y_true', c='C0') - plt.plot(time, y_pred_np[:, 0], '--', label='y_predicted', c='C1') - plt.ylabel(f'y') - plt.title('Sinusoid trajectory: true vs predicted') - plt.xlabel('Time (s)') - plt.legend(loc='best') + # ===== 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 new file mode 100644 index 0000000..b6d48e0 Binary files /dev/null and b/src/analysis/results/RNN100_Sinusoid.pt differ diff --git a/src/analysis/results/RNN_Sinusoid.pt b/src/analysis/results/RNN_Sinusoid.pt new file mode 100644 index 0000000..31c507c Binary files /dev/null and b/src/analysis/results/RNN_Sinusoid.pt differ diff --git a/src/analysis/results/hidden100_Sinusoid.pt b/src/analysis/results/hidden100_Sinusoid.pt new file mode 100644 index 0000000..5e3d791 Binary files /dev/null and b/src/analysis/results/hidden100_Sinusoid.pt differ diff --git a/src/analysis/results/hidden_Sinusoid.pt b/src/analysis/results/hidden_Sinusoid.pt new file mode 100644 index 0000000..8d6bcab Binary files /dev/null and b/src/analysis/results/hidden_Sinusoid.pt differ diff --git a/src/analysis/results/pid100_Sinusoid.npz b/src/analysis/results/pid100_Sinusoid.npz new file mode 100644 index 0000000..d3ada93 Binary files /dev/null and b/src/analysis/results/pid100_Sinusoid.npz differ diff --git a/src/analysis/results/pid100_bars_Sinusoid.png b/src/analysis/results/pid100_bars_Sinusoid.png new file mode 100644 index 0000000..9274db3 Binary files /dev/null and b/src/analysis/results/pid100_bars_Sinusoid.png differ diff --git a/src/analysis/results/pid_Sinusoid.npz b/src/analysis/results/pid_Sinusoid.npz new file mode 100644 index 0000000..ae90ed7 Binary files /dev/null and b/src/analysis/results/pid_Sinusoid.npz differ diff --git a/src/analysis/results/pid_bars_Sinusoid.png b/src/analysis/results/pid_bars_Sinusoid.png new file mode 100644 index 0000000..a9cc18e Binary files /dev/null and b/src/analysis/results/pid_bars_Sinusoid.png differ diff --git a/src/analysis/results/sinusoid100_training_curves_Sinusoid.png b/src/analysis/results/sinusoid100_training_curves_Sinusoid.png new file mode 100644 index 0000000..70e9505 Binary files /dev/null and b/src/analysis/results/sinusoid100_training_curves_Sinusoid.png differ diff --git a/src/analysis/results/sinusoid100_trajectory_Sinusoid.png b/src/analysis/results/sinusoid100_trajectory_Sinusoid.png new file mode 100644 index 0000000..380a64b Binary files /dev/null and b/src/analysis/results/sinusoid100_trajectory_Sinusoid.png differ diff --git a/src/analysis/results/sinusoid_training_curves_Sinusoid.png b/src/analysis/results/sinusoid_training_curves_Sinusoid.png new file mode 100644 index 0000000..e54f4f9 Binary files /dev/null and b/src/analysis/results/sinusoid_training_curves_Sinusoid.png differ diff --git a/src/analysis/results/sinusoid_trajectory_Sinusoid.png b/src/analysis/results/sinusoid_trajectory_Sinusoid.png new file mode 100644 index 0000000..120f580 Binary files /dev/null and b/src/analysis/results/sinusoid_trajectory_Sinusoid.png differ diff --git a/src/analysis/utils.py b/src/analysis/utils.py deleted file mode 100644 index 0791e71..0000000 --- a/src/analysis/utils.py +++ /dev/null @@ -1,82 +0,0 @@ -import numpy as np - -SMOOTHING_SIGMA = 2 -FREQUENCY_CUTOFF = 500 - - -def ensure_length_is_even(x): - n = len(x) - if n % 2 != 0: - x = x[:-1] - n = len(x) - x = np.reshape(x, (n,)) - return x - - -def gauss(x, sigma=1): - return 1 / np.sqrt(2 * np.pi * sigma ** 2) * np.exp(-1 / 2 * (x / sigma) ** 2) - - -def get_kernel(sigma): - size = sigma * 10 + 1 - kernel = list(range(size)) - kernel = [float(k) - int(size / 2) for k in kernel] - kernel = [gauss(k, sigma) for k in kernel] - kernel = [k / np.sum(kernel) for k in kernel] - return kernel - - -def kernel_smoothen(data, kernel_sigma=1): - kernel = get_kernel(kernel_sigma) - data_final = data.copy() - data_conv = np.convolve(data[:], kernel) - pad = int(len(kernel) / 2) - data_final[:] = data_conv[pad:-pad] - data = data_final - return data - - -def fft_smoothed(x): - x = ensure_length_is_even(x) - fft_real = np.fft.rfft(x, norm='ortho') - fft_magnitude = np.abs(fft_real) ** 2 * 2 / len(x) - fft_smoothed = kernel_smoothen(fft_magnitude, kernel_sigma=SMOOTHING_SIGMA) - - return fft_smoothed - - -def get_average_spectrum(trajectories): - spectrum = [] - for trajectory in trajectories: - trajectory = (trajectory - trajectory.mean()) / trajectory.std() - fft = fft_smoothed(trajectory) - spectrum.append(fft) - spectrum = np.nanmean(np.array(spectrum), axis=0) - - return spectrum - - -def power_spectrum_error_per_dim(x_gen, x_true): - x_gen = np.array([x_gen]) - x_true = np.array([x_true]) - - assert x_true.shape[1] == x_gen.shape[1] - assert x_true.shape[2] == x_gen.shape[2] - dim_x = x_gen.shape[2] - pse_corrs_per_dim = [] - for dim in range(dim_x): - spectrum_true = get_average_spectrum(x_true[:, :, dim]) - spectrum_gen = get_average_spectrum(x_gen[:, :, dim]) - spectrum_true = spectrum_true[:FREQUENCY_CUTOFF] - spectrum_gen = spectrum_gen[:FREQUENCY_CUTOFF] - BC = np.trapz(np.sqrt(spectrum_true * spectrum_gen)) - hellinger_dist = np.sqrt(1 - BC) - - pse_corrs_per_dim.append(hellinger_dist) - - return pse_corrs_per_dim - - -def power_spectrum_error(x_gen, x_true): - pse_errors_per_dim = power_spectrum_error_per_dim(x_gen, x_true) - return np.array(pse_errors_per_dim).mean(axis=0)