Merge pull request #155 from geometric-intelligence/half_plane_s_x

Half plane s x

Author: franciscoeacosta

neurometry/curvature/grid-cells-curvature/models/xu_rnn/configs/rnn_isometry.py

@@ -10,11 +10,13 @@ def d(**kwargs):
 def get_config():
     """Get the hyperparameters for the model"""
     config = ml_collections.ConfigDict()
-    config.gpu = 4
+    config.gpu = 5
 
     # training config
     config.train = d(
-        num_steps_train=30000,  # 100000
+        load_pretrain=True,
+        pretrain_dir="logs/rnn_isometry/20240418-180712/ckpt/model/checkpoint-step25000.pth",
+        num_steps_train=40000,  # 100000
         lr=0.006,
         lr_decay_from=10000,
         steps_per_logging=20,
@@ -45,20 +47,22 @@ def get_config():
         block_size=12,
         sigma=0.07,
         w_kernel=1.05,
-        w_trans=0.1,
+        w_trans=10,#0.1
         w_isometry=0.005,
         w_reg_u=0.2,
         reg_decay_until=15000,
         adaptive_dr=True,
-        s_0 = 0.5,
-        x_star = torch.tensor([0.2, 0.5]),
-        sigma_star = 0.3,
-        reward_step = 11000,
+        s_0 = 5,
+        x_star = torch.tensor([0.5, 0.5]),
+        sigma_star_x = 0.2,
+        sigma_star_y = 0.2,
+        reward_step = 3000,
+        saliency_type = "left_half",
     )
 
     # path integration
     config.integration = d(
-        n_inte_step=30,
+        n_inte_step=30, # 50
         n_traj=100,
         n_inte_step_vis=30,
         n_traj_vis=5,

neurometry/curvature/grid-cells-curvature/models/xu_rnn/experiment.py

@@ -58,6 +58,18 @@ def __init__(self, rng, config: ml_collections.ConfigDict, device):
                 self.model.parameters(), lr=config.train.lr, momentum=0.9
             )
 
+        if config.train.load_pretrain:
+            logging.info("==== load pretrain model ====")
+            ckpt_model_path = config.train.pretrain_dir
+            logging.info(f"Loading pretrain model from {ckpt_model_path}")
+            ckpt = torch.load(ckpt_model_path, map_location=device)
+            self.model.load_state_dict(ckpt["state_dict"])
+            logging.info("==== load pretrained optimizer ====")
+            self.optimizer.load_state_dict(ckpt["optimizer"])
+            self.starting_step = ckpt["step"]
+        else:
+            self.starting_step = 1
+
     def train_and_evaluate(self, workdir):
         logging.info("==== Experiment.train_and_evaluate() ===")
 
@@ -81,7 +93,7 @@ def train_and_evaluate(self, workdir):
 
         logging.info("==== Start of training ====")
         with metric_writers.ensure_flushes(writer):
-            for step in range(1, config.num_steps_train + 1):
+            for step in range(self.starting_step, config.num_steps_train + self.starting_step):
                 batch_data = utils.dict_to_device(next(self.train_iter), self.device)
 
                 if 120000 > step > 10000:
@@ -151,7 +163,7 @@ def train_and_evaluate(self, workdir):
                     )
                     train_metrics = []
 
-                if step % config.steps_per_large_logging == 0:
+                if step == self.starting_step or step % config.steps_per_large_logging == 0:
                     ckpt_dir = os.path.join(workdir, "ckpt")
                     if not tf.io.gfile.exists(ckpt_dir):
                         tf.io.gfile.makedirs(ckpt_dir)
@@ -166,8 +178,8 @@ def visualize(activations):
                             )[:10, :10]
                             return utils.draw_heatmap(activations)
 
-                        images_v = visualize(self.model.encoder.v, "v")
-                        images_u = visualize(self.model.decoder.u, "u")
+                        images_v = visualize(self.model.encoder.v)
+                        images_u = visualize(self.model.decoder.u)
 
                         writer.write_images(step, {"v": images_v})
                         wandb.log({"v": wandb.Image(images_v)}, step=step)
@@ -239,7 +251,7 @@ def visualize(activations):
                             step=step,
                         )
 
-                if step % config.steps_per_integration == 0 or step == 1:
+                if step % config.steps_per_integration == 0 or step == self.starting_step:
                     # perform path integration
                     with torch.no_grad():
                         eval_data = utils.dict_to_device(

neurometry/curvature/grid-cells-curvature/models/xu_rnn/model.py

@@ -23,8 +23,10 @@ class GridCellConfig:
     adaptive_dr: bool
     s_0: float
     x_star: torch.Tensor
-    sigma_star: float
+    sigma_star_x: float
+    sigma_star_y: float
     reward_step: int
+    saliency_type: str
 
 
 class GridCell(nn.Module):
@@ -222,7 +224,7 @@ def _loss_trans_rnn(self, traj, step):
             heatmap_reshape = heatmap.reshape((heatmap.shape[0], -1))  # (N, 1600)
             y_hat = heatmap_reshape # actual "place cell" activity over the grid (linear readout of grid cells)
 
-            saliency_kernel = self._saliency_kernel(x_grid).unsqueeze(0).to(traj.device)  # (1, 1600)
+            saliency_kernel = self._saliency_kernel(x_grid, config.saliency_type).unsqueeze(0).to(traj.device)  # (1, 1600)
             if step < config.reward_step:
                 L_error = (y - y_hat) ** 2
             else:
@@ -233,12 +235,43 @@ def _loss_trans_rnn(self, traj, step):
 
         return loss_trans * config.w_trans
 
-    def _saliency_kernel(self, x_grid):
+    def _saliency_kernel(self, x_grid, saliency_type):
+        if saliency_type == "gaussian":
+            return self._saliency_kernel_gaussian(x_grid)
+        if saliency_type == "left_half":
+            return self._saliency_kernel_left_half(x_grid)
+        raise NotImplementedError
+
+    # def _saliency_kernel_gaussian(self, x_grid, sigma_star_x, sigma_star_y):
+    #     config = self.config
+    #     s_0 = config.s_0
+    #     x_star = config.x_star
+    #     sigma_star = config.sigma_star
+    #     s_x = s_0*torch.exp(-torch.sum((x_grid - x_star)**2, dim=1)/(2*sigma_star**2))/np.sqrt(2*np.pi*sigma_star**2)
+    #     return 1 + s_x
+
+    def _saliency_kernel_gaussian(self, x_grid):
         config = self.config
         s_0 = config.s_0
         x_star = config.x_star
-        sigma_star = config.sigma_star
-        s_x = s_0*torch.exp(-torch.sum((x_grid - x_star)**2, dim=1)/(2*sigma_star**2))/np.sqrt(2*np.pi*sigma_star**2)
+        sigma_star_x = config.sigma_star_x
+        sigma_star_y = config.sigma_star_y
+
+        # Calculate the squared differences, scaled by respective sigma values
+        diff = x_grid - x_star
+        scaled_diff_sq = (diff[:, 0]**2 / sigma_star_x**2) + (diff[:, 1]**2 / sigma_star_y**2)
+
+        # Compute the Gaussian function
+        normalization_factor = 2 * np.pi * sigma_star_x * sigma_star_y
+        s_x = s_0 * torch.exp(-0.5 * scaled_diff_sq) / normalization_factor
+
+        return 1 + s_x
+
+
+    def _saliency_kernel_left_half(self, x_grid):
+        config = self.config
+        s_0 = config.s_0
+        s_x = s_0 * (x_grid[:, 0] < 0.5).float()
         return 1 + s_x
 
     def _loss_trans_lstm(self, traj):

neurometry/curvature/grid-cells-curvature/models/xu_rnn/utils.py

@@ -118,32 +118,32 @@ def block_diagonal(matrices):
     return blocked
 
 
-def draw_heatmap(data, save_path, xlabels=None, ylabels=None):
-    # data = np.clip(data, -0.05, 0.05)
-    cmap = cm.get_cmap("rainbow", 1000)
-    figure = plt.figure(facecolor="w")
-    ax = figure.add_subplot(1, 1, 1, position=[0.1, 0.15, 0.8, 0.8])
-    if xlabels is not None:
-        ax.set_xticks(range(len(xlabels)))
-        ax.set_xticklabels(xlabels)
-    if ylabels is not None:
-        ax.set_yticks(range(len(ylabels)))
-        ax.set_yticklabels(ylabels)
-
-    vmax = data[0][0]
-    vmin = data[0][0]
-    for i in data:
-        for j in i:
-            if j > vmax:
-                vmax = j
-            if j < vmin:
-                vmin = j
-    map = ax.imshow(
-        data, interpolation="nearest", cmap=cmap, aspect="auto", vmin=vmin, vmax=vmax
-    )
-    plt.colorbar(mappable=map, cax=None, ax=None, shrink=0.5)
-    plt.savefig(save_path)
-    plt.close()
+# def draw_heatmap(data, save_path, xlabels=None, ylabels=None):
+#     # data = np.clip(data, -0.05, 0.05)
+#     cmap = cm.get_cmap("rainbow", 1000)
+#     figure = plt.figure(facecolor="w")
+#     ax = figure.add_subplot(1, 1, 1, position=[0.1, 0.15, 0.8, 0.8])
+#     if xlabels is not None:
+#         ax.set_xticks(range(len(xlabels)))
+#         ax.set_xticklabels(xlabels)
+#     if ylabels is not None:
+#         ax.set_yticks(range(len(ylabels)))
+#         ax.set_yticklabels(ylabels)
+
+#     vmax = data[0][0]
+#     vmin = data[0][0]
+#     for i in data:
+#         for j in i:
+#             if j > vmax:
+#                 vmax = j
+#             if j < vmin:
+#                 vmin = j
+#     map = ax.imshow(
+#         data, interpolation="nearest", cmap=cmap, aspect="auto", vmin=vmin, vmax=vmax
+#     )
+#     plt.colorbar(mappable=map, cax=None, ax=None, shrink=0.5)
+#     plt.savefig(save_path)
+#     plt.close()
 
 
 def shape_mask(size, shape):
@@ -452,39 +452,39 @@ def _draw_real_pred_pairs(real, pred, area_size: int):
     ax.set_aspect(1)
 
 
-# def draw_heatmap(weights):
-#     # weights should a 4-D tensor: [M, N, H, W]
-#     nrow, ncol = weights.shape[0], weights.shape[1]
-#     fig = plt.figure(figsize=(ncol, nrow))
-
-#     for i in range(nrow):
-#         for j in range(ncol):
-#             plt.subplot(nrow, ncol, i * ncol + j + 1)
-#             weight = weights[i, j]
-#             vmin, vmax = weight.min() - 0.01, weight.max()
-
-#             cmap = cm.get_cmap("rainbow", 1000)
-#             cmap.set_under("w")
-
-#             plt.imshow(
-#                 weight,
-#                 interpolation="nearest",
-#                 cmap=cmap,
-#                 aspect="auto",
-#                 vmin=vmin,
-#                 vmax=vmax,
-#             )
-#             plt.axis("off")
-
-#     fig.canvas.draw()
-#     image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-#     image_from_plot = image_from_plot.reshape(
-#         fig.canvas.get_width_height()[::-1] + (3,)
-#     )
-#     # plt.show()
-#     plt.close(fig)
+def draw_heatmap(weights):
+    # weights should a 4-D tensor: [M, N, H, W]
+    nrow, ncol = weights.shape[0], weights.shape[1]
+    fig = plt.figure(figsize=(ncol, nrow))
 
-#     return np.expand_dims(image_from_plot, axis=0)
+    for i in range(nrow):
+        for j in range(ncol):
+            plt.subplot(nrow, ncol, i * ncol + j + 1)
+            weight = weights[i, j]
+            vmin, vmax = weight.min() - 0.01, weight.max()
+
+            cmap = cm.get_cmap("rainbow", 1000)
+            cmap.set_under("w")
+
+            plt.imshow(
+                weight,
+                interpolation="nearest",
+                cmap=cmap,
+                aspect="auto",
+                vmin=vmin,
+                vmax=vmax,
+            )
+            plt.axis("off")
+
+    fig.canvas.draw()
+    image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    image_from_plot = image_from_plot.reshape(
+        fig.canvas.get_width_height()[::-1] + (3,)
+    )
+    # plt.show()
+    plt.close(fig)
+
+    return np.expand_dims(image_from_plot, axis=0)
 
 
 def average_appended_metrics(metrics):

neurometry/curvature/grid-cells-curvature/models/xu_rnn/viz.py

@@ -0,0 +1,97 @@
+import os
+import pickle
+
+import matplotlib.pyplot as plt
+import numpy as np
+from moviepy.editor import ImageSequenceClip
+
+from neurometry.datasets.load_rnn_grid_cells import plot_rate_map
+
+logs_dir = "logs/rnn_isometry"
+
+activations_dir = "ckpt/activations"
+
+def load_matrices(file_path):
+    with open(file_path, "rb") as file:
+        data = pickle.load(file)
+    return data["u"], data["v"]
+
+
+# def save_matrix_as_image(matrix, image_path):
+#     plt.imshow(matrix, cmap='viridis')
+#     plt.colorbar()
+#     plt.savefig(image_path)
+#     plt.close()
+
+
+def create_video(image_folder, output_file, fps=10):
+    image_files = [os.path.join(image_folder, img) for img in sorted(os.listdir(image_folder))]
+    clip = ImageSequenceClip(image_files, fps=fps)
+    clip.write_videofile(output_file, codec="libx264")
+
+
+# def generate_videos(run_id, start_epoch=25000, end_epoch=65000, step=500):
+
+#     data_dir = os.path.join(logs_dir, run_id, activations_dir)
+
+#     print(data_dir)
+
+#     output_dir = data_dir
+#     os.makedirs(output_dir, exist_ok=True)
+#     u_images_dir = os.path.join(output_dir, 'u_images')
+#     v_images_dir = os.path.join(output_dir, 'v_images')
+#     os.makedirs(u_images_dir, exist_ok=True)
+#     os.makedirs(v_images_dir, exist_ok=True)
+
+#     # Generate images
+#     for epoch in range(start_epoch, end_epoch, step):
+#         file_name = f'activations-step{epoch}.pkl'
+#         file_path = os.path.join(data_dir, file_name)
+#         u, v = load_matrices(file_path)
+#         save_matrix_as_image(u, os.path.join(u_images_dir, f'u_{epoch}.png'))
+#         save_matrix_as_image(v, os.path.join(v_images_dir, f'v_{epoch}.png'))
+
+#     # Create videos
+#     create_video(u_images_dir, os.path.join(output_dir, 'u_video.mp4'))
+#     create_video(v_images_dir, os.path.join(output_dir, 'v_video.mp4'))
+
+
+def save_rate_maps_as_image(indices, num_plots, activations, image_path, title, seed=None):
+    plot_rate_map(indices, num_plots, activations, title, seed=seed)
+    plt.savefig(image_path)
+    plt.close()
+
+
+def generate_videos(run_id, start_epoch=25000, end_epoch=65000, step=500, num_cells_per_image=20, seed=None):
+
+    data_dir = os.path.join(logs_dir, run_id, activations_dir)
+    #os.makedirs(output_dir, exist_ok=True)
+    output_dir = data_dir
+    u_images_dir = os.path.join(output_dir, "u_images")
+    v_images_dir = os.path.join(output_dir, "v_images")
+    os.makedirs(u_images_dir, exist_ok=True)
+    os.makedirs(v_images_dir, exist_ok=True)
+
+    rng = np.random.default_rng(seed=seed)
+    idxs = rng.integers(0, 1799, num_cells_per_image)
+
+    # Iterate through epochs
+    for epoch in range(start_epoch, end_epoch, step):
+        file_name = f"activations-step{epoch}.pkl"
+        file_path = os.path.join(data_dir, file_name)
+        u, v = load_matrices(file_path)
+
+        # Save images using your custom plotting function
+        u_image_path = os.path.join(u_images_dir, f"u_{epoch}.png")
+        v_image_path = os.path.join(v_images_dir, f"v_{epoch}.png")
+        save_rate_maps_as_image(idxs, num_cells_per_image, u, u_image_path, f"U matrices at epoch {epoch}", seed=seed)
+        save_rate_maps_as_image(idxs, num_cells_per_image, v, v_image_path, f"V matrices at epoch {epoch}", seed=seed)
+
+    # Create videos
+    create_video(u_images_dir, os.path.join(output_dir, "u_video.mp4"))
+    create_video(v_images_dir, os.path.join(output_dir, "v_video.mp4"))
+
+
+
+
+