Merge pull request #143 from geometric-intelligence/backup-before-det…

…aching-submodules

Backup before detaching submodules

neurometry/datasets/load_rnn_grid_cells.py

@@ -33,6 +33,7 @@ def load_activations(epochs, version="single", verbose=True):
     activations = []
     rate_maps = []
     state_points = []
+    positions = []
 
     if version == "single":
         activations_dir = (
@@ -51,53 +52,95 @@ def load_activations(epochs, version="single", verbose=True):
         rate_map_epoch_path = (
             activations_dir + f"rate_map_{version}_agent_epoch_{epoch}.npy"
         )
+        positions_epoch_path = (
+            activations_dir + f"positions_{version}_agent_epoch_{epoch}.npy"
+        )
+
         if os.path.exists(activations_epoch_path) and os.path.exists(
             rate_map_epoch_path
-        ):
+        ) and os.path.exists(positions_epoch_path):
             activations.append(np.load(activations_epoch_path))
             rate_maps.append(np.load(rate_map_epoch_path))
+            positions.append(np.load(positions_epoch_path))
             if verbose:
-                print(f"Epoch {epoch} found!!! :D")
+                print(f"Epoch {epoch} found!")
         else:
             print(f"Epoch {epoch} not found. Loading ...")
             parser = config.parser
             options, _ = parser.parse_known_args()
             options.run_ID = utils.generate_run_ID(options)
-            if type == "single":
+            if version == "single":
                 (
                     activations_single_agent,
                     rate_map_single_agent,
+                    positions_single_agent,
                 ) = single_agent_activity.main(options, epoch=epoch)
                 activations.append(activations_single_agent)
                 rate_maps.append(rate_map_single_agent)
-            elif type == "dual":
-                activations_dual_agent, rate_map_dual_agent = dual_agent_activity.main(
+                positions.append(positions_single_agent)
+            elif version == "dual":
+                activations_dual_agent, rate_map_dual_agent, positions_dual_agent = dual_agent_activity.main(
                     options, epoch=epoch
                 )
                 activations.append(activations_dual_agent)
                 rate_maps.append(rate_map_dual_agent)
+                positions.append(positions_dual_agent)
+        print(len(activations))
         state_points_epoch = activations[-1].reshape(activations[-1].shape[0], -1)
         state_points.append(state_points_epoch)
 
     if verbose:
         print(f"Loaded epochs {epochs} of {version} agent model.")
         print(
-            f"There are {activations[0].shape[0]} grid cells with {activations[0].shape[1]} x {activations[0].shape[2]} environment resolution, averaged over {activations[0].shape[3]} trajectories."
+            f"activations has shape {activations[0].shape}. There are {activations[0].shape[0]} grid cells with {activations[0].shape[1]} x {activations[0].shape[2]} environment resolution, averaged over {activations[0].shape[3]} trajectories."
         )
         print(
-            f"There are {state_points[0].shape[1]} data points in the {state_points[0].shape[0]}-dimensional state space."
+            f"state_points has shape {state_points[0].shape}. There are {state_points[0].shape[1]} data points in the {state_points[0].shape[0]}-dimensional state space."
         )
         print(
-            f"There are {rate_maps[0].shape[1]} data points averaged over {activations[0].shape[3]} trajectories in the {rate_maps[0].shape[0]}-dimensional state space."
+            f"rate_maps has shape {rate_maps[0].shape}. There are {rate_maps[0].shape[1]} data points averaged over {activations[0].shape[3]} trajectories in the {rate_maps[0].shape[0]}-dimensional state space."
         )
+        print(f"positions has shape {positions[0].shape}.")
+
+    return activations, rate_maps, state_points, positions
+
+
+# def plot_rate_map(indices, num_plots, activations, title):
+#     rng = np.random.default_rng(seed=0)
+#     if indices is None:
+#         idxs = rng.integers(0, 4095, num_plots)
+#     else:
+#         idxs = indices
+#         num_plots = len(indices)
+
+#     rows = 4
+#     cols = num_plots // rows + (num_plots % rows > 0)
 
-    return activations, rate_maps, state_points
+#     plt.rcParams["text.usetex"] = False
 
+#     fig, axes = plt.subplots(rows, cols, figsize=(20, 8))
 
-def plot_rate_map(indices, num_plots, activations):
+#     for i in range(rows):
+#         for j in range(cols):
+#             if i * cols + j < num_plots:
+#                 gc = np.mean(activations[idxs[i * cols + j]], axis=2)
+#                 axes[i, j].imshow(gc)
+#                 axes[i, j].set_title(f"grid cell id: {idxs[i * cols + j]}")
+#                 axes[i, j].axis("off")
+#             else:
+#                 axes[i, j].axis("off")
+
+#     plt.suptitle(title)
+#     plt.tight_layout()
+#     plt.show()
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+def plot_rate_map(indices, num_plots, activations, title):
     rng = np.random.default_rng(seed=0)
     if indices is None:
-        idxs = rng.integers(0, 4095, num_plots)
+        idxs = rng.integers(0, activations.shape[0]-1, num_plots)
     else:
         idxs = indices
         num_plots = len(indices)
@@ -112,12 +155,24 @@ def plot_rate_map(indices, num_plots, activations):
     for i in range(rows):
         for j in range(cols):
             if i * cols + j < num_plots:
-                gc = np.mean(activations[idxs[i * cols + j]], axis=2)
-                axes[i, j].imshow(gc)
-                axes[i, j].set_title(f"grid cell id: {idxs[i * cols + j]}")
-                axes[i, j].axis("off")
+                if len(activations.shape) == 4:
+                    gc = np.mean(activations[idxs[i * cols + j]], axis=2)
+                else:
+                    gc = activations[idxs[i * cols + j]]
+                if axes.ndim > 1:  # Check if axes is a 2D array
+                    ax = axes[i, j]
+                else:  # If axes is flattened (e.g., only one row of subplots)
+                    ax = axes[i * cols + j]
+                ax.imshow(gc)
+                ax.set_title(f"grid cell id: {idxs[i * cols + j]}", fontsize=10)
+                ax.axis("off")
             else:
-                axes[i, j].axis("off")
+                if axes.ndim > 1:
+                    axes[i, j].axis("off")
+                else:
+                    axes[i * cols + j].axis("off")
 
+    fig.suptitle(title, fontsize=30) 
     plt.tight_layout()
     plt.show()
+

neurometry/datasets/rnn_grid_cells/config.py

@@ -21,9 +21,10 @@ class Config:
     periodic = False  # trajectories with periodic boundary conditions
     box_width = 2.2  # width of training environment
     box_height = 2.2  # height of training environment
-    device = (
-        "cuda" if torch.cuda.is_available() else "cpu"
-    )  # device to use for training
+    # device = (
+    #     "cuda" if torch.cuda.is_available() else "cpu"
+    # )  # device to use for training
+    device = torch.device('cuda:8')
     n_avg = 50  # number of trajectories to average over for rate maps
 
 

neurometry/datasets/rnn_grid_cells/dual_agent_activity.py

@@ -3,14 +3,14 @@
 
 import numpy as np
 import torch
-from config import parser
-from model_dual_path_integration import RNN
-from place_cells_dual_path_integration import PlaceCells
-from scores import GridScorer
+from .config import parser
+from .model_dual_path_integration import RNN
+from .place_cells_dual_path_integration import PlaceCells
+from .scores import GridScorer
 from tqdm import tqdm
-from trajectory_generator_dual_path_integration import TrajectoryGenerator
-from utils import generate_run_ID
-from visualize import compute_ratemaps
+from .trajectory_generator_dual_path_integration import TrajectoryGenerator
+from .utils import generate_run_ID
+from .visualize import compute_ratemaps
 
 parent_dir = os.getcwd() + "/"
 model_folder = "Dual agent path integration disjoint PCs/Seed 1 weight decay 1e-06/"
@@ -44,7 +44,7 @@ def main(options, epoch="final", res=20):
 
     Ng = options.Ng
     n_avg = options.n_avg
-    activations_dual_agent, rate_map_dual_agent, _, _ = compute_ratemaps(
+    activations_dual_agent, rate_map_dual_agent, _, positions_dual_agent = compute_ratemaps(
         model_dual_agent,
         trajectory_generator,
         options,
@@ -63,11 +63,15 @@ def main(options, epoch="final", res=20):
     np.save(
         activations_dir + f"rate_map_dual_agent_epoch_{epoch}.npy", rate_map_dual_agent
     )
+    np.save(
+        activations_dir + f"positions_dual_agent_epoch_{epoch}.npy",
+        positions_dual_agent,
+    )
 
     # #   activations is in the shape [number of grid cells (Ng) x res x res x n_avg]
     # #   ratemap is in the shape [Ng x res^2]
 
-    return activations_dual_agent, rate_map_dual_agent
+    return activations_dual_agent, rate_map_dual_agent, positions_dual_agent
 
 
 def compute_grid_scores(res, rate_map_dual_agent, scorer):

neurometry/datasets/rnn_grid_cells/scores.py

@@ -161,7 +161,7 @@ def get_grid_scores_for_mask(self, sac, rotated_sacs, mask):
         masked_sac_centered = (masked_sac - masked_sac_mean) * mask
         variance = np.sum(masked_sac_centered**2) / ring_area + 1e-5
         corrs = dict()
-        for angle, rotated_sac in zip(self._corr_angles, rotated_sacs, strict=False):
+        for angle, rotated_sac in zip(self._corr_angles, rotated_sacs,strict=False):
             masked_rotated_sac = (rotated_sac - masked_sac_mean) * mask
             cross_prod = np.sum(masked_sac_centered * masked_rotated_sac) / ring_area
             corrs[angle] = cross_prod / variance

neurometry/datasets/rnn_grid_cells/single_agent_activity.py

@@ -5,16 +5,17 @@
 
 import numpy as np
 import torch
-from config import parser
-from model import RNN
-from place_cells import PlaceCells
-from scores import GridScorer
+from .config import parser
+from .model import RNN
+from .place_cells import PlaceCells
+from .scores import GridScorer
 from tqdm import tqdm
-from trajectory_generator import TrajectoryGenerator
-from utils import generate_run_ID
-from visualize import compute_ratemaps
+from .trajectory_generator import TrajectoryGenerator
+from .utils import generate_run_ID
+from .visualize import compute_ratemaps
 
-parent_dir = os.getcwd() + "/"
+# parent_dir = os.getcwd() + "/"
+parent_dir = "/scratch/facosta/rnn_grid_cells/"
 
 
 model_folder = "Single agent path integration/Seed 1 weight decay 1e-06/"
@@ -38,6 +39,7 @@ def main(options, epoch="final", res=20):
 
     model_single_agent = model.to(options.device)
 
+
     model_name = "final_model.pth" if epoch == "final" else f"epoch_{epoch}.pth"
     saved_model_single_agent = torch.load(
         parent_dir + model_folder + model_parameters + model_name
@@ -49,7 +51,7 @@ def main(options, epoch="final", res=20):
     Ng = options.Ng
     n_avg = options.n_avg
 
-    activations_single_agent, rate_map_single_agent, _, _ = compute_ratemaps(
+    activations_single_agent, rate_map_single_agent, _, positions_single_agent = compute_ratemaps(
         model_single_agent,
         trajectory_generator,
         options,
@@ -69,10 +71,12 @@ def main(options, epoch="final", res=20):
         activations_dir + f"rate_map_single_agent_epoch_{epoch}.npy",
         rate_map_single_agent,
     )
+
+    np.save(activations_dir + f"positions_single_agent_epoch_{epoch}.npy", positions_single_agent)
     # #   activations is in the shape [number of grid cells (Ng) x res x res x n_avg]
     # #   ratemap is in the shape [Ng x res^2]
 
-    return activations_single_agent, rate_map_single_agent
+    return activations_single_agent, rate_map_single_agent, positions_single_agent
 
 
 def compute_grid_scores(res, rate_map_single_agent, scorer):

neurometry/datasets/synthetic.py

@@ -5,7 +5,7 @@
 import torch  # noqa: E402
 from geomstats.geometry.euclidean import Euclidean  # noqa: E402
 from geomstats.geometry.hypersphere import Hypersphere  # noqa: E402
-from geomstats.geometry.klein_bottle import KleinBottle  # noqa: E402
+#from geomstats.geometry.klein_bottle import KleinBottle  # noqa: E402
 from geomstats.geometry.product_manifold import ProductManifold  # noqa: E402