PID preliminary analysis on sinusoid-trained RNN

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()
+
+

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
+

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)

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