Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,11 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Auto detect text files and perform LF normalization
+* text=auto
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true
+docs/LSTM_hexagons.png filter=lfs diff=lfs merge=lfs -text
+docs/RNNgrids.png filter=lfs diff=lfs merge=lfs -text
+docs/poisson_spiking.gif filter=lfs diff=lfs merge=lfs -text
+models/example_pc_centers.npy filter=lfs diff=lfs merge=lfs -text
+models/example_trained_weights.npy filter=lfs diff=lfs merge=lfs -text
+models/steps_20_batch_200_RNN_4096_relu_rf_012_DoG_True_periodic_False_lr_00001_weight_decay_00001/1.png filter=lfs diff=lfs merge=lfs -text
+models/steps_20_batch_200_RNN_4096_relu_rf_012_DoG_True_periodic_False_lr_00001_weight_decay_00001/most_recent_model.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,123 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+images/
+movies/
+data/
+fixed_point_finder_data/
+ben_decimate_1000_step/
+ben_true_100_step/
+ben_100*/
+data/
+nohup.out
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+
+# by user
+saved
+*.ipynb
+*.zip
+# pixi environments
+.pixi/*
+!.pixi/config.toml

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,69 @@
----
-license: mit
----
+[![DOI](https://zenodo.org/badge/217773694.svg)](https://zenodo.org/badge/latestdoi/217773694)
+
+# Grid cells in RNNs trained to path integrate
+
+Code to reproduce the trained RNN in [**a unified theory for the origin of grid cells through the lens of pattern formation (NeurIPS '19)**](https://papers.nips.cc/paper/9191-a-unified-theory-for-the-origin-of-grid-cells-through-the-lens-of-pattern-formation) and additional analysis described in this [**preprint**](https://www.biorxiv.org/content/10.1101/2020.12.29.424583v1). 
+
+
+Quick start:
+
+<img src="./docs/poisson_spiking.gif" width="300" align="right">
+
+* [**inspect_model.ipynb**](inspect_model.ipynb):
+  Train a model and visualize its hidden unit ratemaps. 
+ 
+* [**main.py**](main.py):
+  or, train a model from the command line.
+  
+* [**pattern_formation.ipynb**](pattern_formation.ipynb):
+  Numerical simulations of pattern-forming dynamics.
+  
+  
+Includes:
+
+* [**trajectory_generator.py**](trajectory_generator.py):
+  Generate simulated rat trajectories in a rectangular environment.
+
+* [**place_cells.py**](place_cells.py):
+  Tile a set of simulated place cells across the training environment. 
+  
+* [**model.py**](model.py):
+  Contains the vanilla RNN model architecture, as well as an LSTM.
+  
+* [**trainer.py**](model.py):
+  Contains model training loop.
+  
+* [**models/example_trained_weights.npy**](models/example_trained_weights.npy)
+  Contains a set of pre-trained weights.
+
+## Running
+
+We recommend creating a virtual environment:
+
+```shell
+$ virtualenv env
+$ source env/bin/activate
+$ pip install --upgrade pip
+```
+
+Then, install the dependencies automatically with `pip install -r requirements.txt`
+or manually with:
+
+```shell
+$ pip install --upgrade numpy==1.17.2
+$ pip install --upgrade tensorflow==2.0.0rc2
+$ pip install --upgrade scipy==1.4.1
+$ pip install --upgrade matplotlib==3.0.3
+$ pip install --upgrade imageio==2.5.0
+$ pip install --upgrade opencv-python==4.1.1.26
+$ pip install --upgrade tqdm==4.36.0
+$ pip install --upgrade opencv-python==4.1.1.26
+$ pip install --upgrade torch==1.10.0
+```
+
+If you want to train your own models, make sure to properly set the default
+save directory in `main.py`! 
+
+## Result
+
+![grid visualization](./docs/RNNgrids.png)

main.py

@@ -0,0 +1,103 @@
+import numpy as np
+import tensorflow as tf
+import torch.cuda
+import argparse
+
+
+from utils import generate_run_ID
+from place_cells import PlaceCells
+from trajectory_generator import TrajectoryGenerator
+from model import RNN
+from trainer import Trainer
+
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--save_dir",
+    # default='/mnt/fs2/bsorsch/grid_cells/models/',
+    default="models/",
+    help="directory to save trained models",
+)
+parser.add_argument(
+    "--n_epochs", default=100, help="number of training epochs", type=int
+)
+parser.add_argument("--n_steps", default=1000, help="batches per epoch", type=int)
+parser.add_argument(
+    "--batch_size", default=200, help="number of trajectories per batch", type=int
+)
+parser.add_argument(
+    "--sequence_length", default=20, help="number of steps in trajectory", type=int
+)
+parser.add_argument(
+    "--learning_rate", default=1e-4, help="gradient descent learning rate", type=float
+)
+parser.add_argument("--Np", default=512, help="number of place cells", type=int)
+parser.add_argument("--Ng", default=4096, help="number of grid cells", type=int)
+parser.add_argument(
+    "--place_cell_rf",
+    default=0.12,
+    help="width of place cell center tuning curve (m)",
+    type=float,
+)
+parser.add_argument(
+    "--surround_scale",
+    default=2,
+    help="if DoG, ratio of sigma2^2 to sigma1^2",
+    type=int,
+)
+parser.add_argument("--RNN_type", default="RNN", help="RNN or LSTM")
+parser.add_argument("--activation", default="relu", help="recurrent nonlinearity")
+parser.add_argument(
+    "--weight_decay",
+    default=1e-4,
+    help="strength of weight decay on recurrent weights",
+    type=float,
+)
+parser.add_argument(
+    "--DoG", default=True, help="use difference of gaussians tuning curves"
+)
+parser.add_argument(
+    "--periodic", default=False, help="trajectories with periodic boundary conditions"
+)
+parser.add_argument(
+    "--box_width", default=2.2, help="width of training environment", type=float
+)
+parser.add_argument(
+    "--box_height", default=2.2, help="height of training environment", type=float
+)
+parser.add_argument(
+    "--device",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    help="device to use for training",
+)
+parser.add_argument(
+    "--seed", default=None, help="seed number for all numpy random number generator"
+)
+
+options = parser.parse_args()
+options.run_ID = generate_run_ID(options)
+
+print(f"Using device: {options.device}")
+
+if options.seed:
+    np.random.seed(int(options.seed))
+
+place_cells = PlaceCells(options)
+if options.RNN_type == "RNN":
+    model = RNN(options, place_cells)
+elif options.RNN_type == "LSTM":
+    # model = LSTM(options, place_cells)
+    raise NotImplementedError
+
+# Put model on GPU if using GPU
+if options.device == "cuda":
+    print("Using CUDA")
+model = model.to(options.device)
+
+trajectory_generator = TrajectoryGenerator(options, place_cells)
+
+trainer = Trainer(options, model, trajectory_generator)
+
+# Train
+trainer.train(n_epochs=options.n_epochs, n_steps=options.n_steps)

model.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+import torch
+
+
+class RNN(torch.nn.Module):
+    def __init__(self, options, place_cells):
+        super(RNN, self).__init__()
+        self.Ng = options.Ng
+        self.Np = options.Np
+        self.sequence_length = options.sequence_length
+        self.weight_decay = options.weight_decay
+        self.place_cells = place_cells
+
+        # Input weights
+        self.encoder = torch.nn.Linear(self.Np, self.Ng, bias=False)
+        self.RNN = torch.nn.RNN(
+            input_size=2,
+            hidden_size=self.Ng,
+            nonlinearity=options.activation,
+            bias=False,
+        )
+        # Linear read-out weights
+        self.decoder = torch.nn.Linear(self.Ng, self.Np, bias=False)
+
+        self.softmax = torch.nn.Softmax(dim=-1)
+
+    def g(self, inputs):
+        """
+        Compute grid cell activations.
+        Args:
+            inputs: Batch of 2d velocity inputs with shape [batch_size, sequence_length, 2].
+
+        Returns:
+            g: Batch of grid cell activations with shape [batch_size, sequence_length, Ng].
+        """
+        v, p0 = inputs
+        init_state = self.encoder(p0)[None]
+        g, _ = self.RNN(v, init_state)
+        return g
+
+    def predict(self, inputs):
+        """
+        Predict place cell code.
+        Args:
+            inputs: Batch of 2d velocity inputs with shape [batch_size, sequence_length, 2].
+
+        Returns:
+            place_preds: Predicted place cell activations with shape
+                [batch_size, sequence_length, Np].
+        """
+        place_preds = self.decoder(self.g(inputs))
+
+        return place_preds
+
+    def set_weights(self, weights):
+        """
+        Load weights from a numpy array (e.g. from the provided example weights).
+        Assumes weights are in the order: [encoder, rnn_ih, rnn_hh, decoder]
+        and transposed (TF/Keras format).
+        """
+        with torch.no_grad():
+            # Encoder: (Np, Ng) -> (Ng, Np)
+            self.encoder.weight.copy_(torch.from_numpy(weights[0].T).float())
+            
+            # RNN input: (2, Ng) -> (Ng, 2)
+            self.RNN.weight_ih_l0.copy_(torch.from_numpy(weights[1].T).float())
+            
+            # RNN hidden: (Ng, Ng) -> (Ng, Ng)
+            self.RNN.weight_hh_l0.copy_(torch.from_numpy(weights[2].T).float())
+            
+            # Decoder: (Ng, Np) -> (Np, Ng)
+            self.decoder.weight.copy_(torch.from_numpy(weights[3].T).float())
+
+    def compute_loss(self, inputs, pc_outputs, pos):
+        """
+        Compute avg. loss and decoding error.
+        Args:
+            inputs: Batch of 2d velocity inputs with shape [batch_size, sequence_length, 2].
+            pc_outputs: Ground truth place cell activations with shape
+                [batch_size, sequence_length, Np].
+            pos: Ground truth 2d position with shape [batch_size, sequence_length, 2].
+
+        Returns:
+            loss: Avg. loss for this training batch.
+            err: Avg. decoded position error in cm.
+        """
+        y: torch.Tensor = pc_outputs
+        preds: torch.Tensor = self.predict(inputs)
+        loss = torch.nn.functional.cross_entropy(preds.flatten(0, 1), y.flatten(0, 1))
+
+        # Weight regularization
+        loss += self.weight_decay * (self.RNN.weight_hh_l0**2).sum()
+
+        # Compute decoding error
+        pred_pos = self.place_cells.get_nearest_cell_pos(preds)
+        err = torch.sqrt(((pos - pred_pos) ** 2).sum(-1)).mean()
+
+        return loss, err

models/example_pc_centers.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cafa64a4894d5f1c4aaf4befdf036fb43e9c9ed2f2825b79feb02ad1f7199997
+size 4224

models/example_trained_weights.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5277a433615f00894ffe13d55b46d5becda9b362e099c7c73a6191e61bf51ff2
+size 83919352

models/steps_20_batch_200_RNN_4096_relu_rf_012_DoG_True_periodic_False_lr_00001_weight_decay_00001/most_recent_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f7050cd3c660329c5c49e1bc935e66cb9eaebeb42caff69930d6ee793c8f6c1
+size 83954107

place_cells.py

@@ -0,0 +1,124 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import torch
+import scipy
+
+
+class PlaceCells(object):
+    def __init__(self, options, us=None):
+        self.Np = options.Np
+        self.sigma = options.place_cell_rf
+        self.surround_scale = options.surround_scale
+        self.box_width = options.box_width
+        self.box_height = options.box_height
+        self.is_periodic = options.periodic
+        self.DoG = options.DoG
+        self.device = options.device
+        self.softmax = torch.nn.Softmax(dim=-1)
+
+        # Randomly tile place cell centers across environment
+        usx = np.random.uniform(-self.box_width / 2, self.box_width / 2, (self.Np,))
+        usy = np.random.uniform(-self.box_width / 2, self.box_width / 2, (self.Np,))
+        self.us = torch.tensor(np.vstack([usx, usy]).T)
+        # If using a GPU, put on GPU
+        self.us = self.us.to(self.device)
+        # self.us = torch.tensor(np.load('models/example_pc_centers.npy')).cuda()
+
+    def get_activation(self, pos):
+        """
+        Get place cell activations for a given position.
+
+        Args:
+            pos: 2d position of shape [batch_size, sequence_length, 2].
+
+        Returns:
+            outputs: Place cell activations with shape [batch_size, sequence_length, Np].
+        """
+        d = torch.abs(pos[:, :, None, :] - self.us[None, None, ...]).float()
+
+        if self.is_periodic:
+            dx = d[:, :, :, 0]
+            dy = d[:, :, :, 1]
+            dx = torch.minimum(dx, self.box_width - dx)
+            dy = torch.minimum(dy, self.box_height - dy)
+            d = torch.stack([dx, dy], axis=-1)
+
+        norm2 = (d**2).sum(-1)
+
+        # Normalize place cell outputs with prefactor alpha=1/2/np.pi/self.sigma**2,
+        # or, simply normalize with softmax, which yields same normalization on
+        # average and seems to speed up training.
+        outputs = self.softmax(-norm2 / (2 * self.sigma**2))
+
+        if self.DoG:
+            # Again, normalize with prefactor
+            # beta=1/2/np.pi/self.sigma**2/self.surround_scale, or use softmax.
+            outputs -= self.softmax(-norm2 / (2 * self.surround_scale * self.sigma**2))
+
+            # Shift and scale outputs so that they lie in [0,1].
+            min_output, _ = outputs.min(-1, keepdims=True)
+            outputs += torch.abs(min_output)
+            outputs /= outputs.sum(-1, keepdims=True)
+        return outputs
+
+    def get_nearest_cell_pos(self, activation, k=3):
+        """
+        Decode position using centers of k maximally active place cells.
+
+        Args:
+            activation: Place cell activations of shape [batch_size, sequence_length, Np].
+            k: Number of maximally active place cells with which to decode position.
+
+        Returns:
+            pred_pos: Predicted 2d position with shape [batch_size, sequence_length, 2].
+        """
+        _, idxs = torch.topk(activation, k=k)
+        pred_pos = self.us[idxs].mean(-2)
+        return pred_pos
+
+    def grid_pc(self, pc_outputs, res=32):
+        """Interpolate place cell outputs onto a grid"""
+        coordsx = np.linspace(-self.box_width / 2, self.box_width / 2, res)
+        coordsy = np.linspace(-self.box_height / 2, self.box_height / 2, res)
+        grid_x, grid_y = np.meshgrid(coordsx, coordsy)
+        grid = np.stack([grid_x.ravel(), grid_y.ravel()]).T
+
+        # Convert to numpy
+        pc_outputs = pc_outputs.reshape(-1, self.Np)
+
+        T = pc_outputs.shape[0]  # T vs transpose? What is T? (dim's?)
+        pc = np.zeros([T, res, res])
+        for i in range(len(pc_outputs)):
+            gridval = scipy.interpolate.griddata(self.us.cpu(), pc_outputs[i], grid)
+            pc[i] = gridval.reshape([res, res])
+
+        return pc
+
+    def compute_covariance(self, res=30):
+        """Compute spatial covariance matrix of place cell outputs"""
+        pos = np.array(
+            np.meshgrid(
+                np.linspace(-self.box_width / 2, self.box_width / 2, res),
+                np.linspace(-self.box_height / 2, self.box_height / 2, res),
+            )
+        ).T
+
+        pos = torch.tensor(pos)
+
+        # Put on GPU if available
+        pos = pos.to(self.device)
+
+        # Maybe specify dimensions here again?
+        pc_outputs = self.get_activation(pos).reshape(-1, self.Np).cpu()
+
+        C = pc_outputs @ pc_outputs.T
+        Csquare = C.reshape(res, res, res, res)
+
+        Cmean = np.zeros([res, res])
+        for i in range(res):
+            for j in range(res):
+                Cmean += np.roll(np.roll(Csquare[i, j], -i, axis=0), -j, axis=1)
+
+        Cmean = np.roll(np.roll(Cmean, res // 2, axis=0), res // 2, axis=1)
+
+        return Cmean

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+authors = [{name = "Mio", email = "violet172309@gmail.com"}]
+dependencies = ["numpy", "tensorflow", "scipy", "matplotlib", "imageio", "opencv-python", "tqdm", "torch"]
+name = "grid-pattern-formation"
+requires-python = ">= 3.11"
+version = "0.1.0"
+
+[build-system]
+build-backend = "hatchling.build"
+requires = ["hatchling"]
+
+[tool.pixi.workspace]
+channels = ["conda-forge"]
+platforms = ["linux-64"]
+
+[tool.pixi.system-requirements]
+cuda = "12.6"
+
+[tool.pixi.pypi-dependencies]
+grid_pattern_formation = { path = ".", editable = true }
+
+[tool.pixi.tasks]

requirements.txt

@@ -0,0 +1,8 @@
+numpy
+tensorflow
+scipy
+matplotlib
+imageio
+opencv-python
+tqdm
+torch

scores.py

@@ -0,0 +1,272 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Grid score calculations.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.signal
+import scipy.ndimage as ndimage
+
+
+def circle_mask(size, radius, in_val=1.0, out_val=0.0):
+  """Calculating the grid scores with different radius."""
+  sz = [math.floor(size[0] / 2), math.floor(size[1] / 2)]
+  x = np.linspace(-sz[0], sz[1], size[1])
+  x = np.expand_dims(x, 0)
+  x = x.repeat(size[0], 0)
+  y = np.linspace(-sz[0], sz[1], size[1])
+  y = np.expand_dims(y, 1)
+  y = y.repeat(size[1], 1)
+  z = np.sqrt(x**2 + y**2)
+  z = np.less_equal(z, radius)
+  vfunc = np.vectorize(lambda b: b and in_val or out_val)
+  return vfunc(z)
+
+
+class GridScorer(object):
+  """Class for scoring ratemaps given trajectories."""
+
+  def __init__(self, nbins, coords_range, mask_parameters, min_max=False):
+    """Scoring ratemaps given trajectories.
+    Args:
+      nbins: Number of bins per dimension in the ratemap.
+      coords_range: Environment coordinates range.
+      mask_parameters: parameters for the masks that analyze the angular
+        autocorrelation of the 2D autocorrelation.
+      min_max: Correction.
+    """
+    self._nbins = nbins
+    self._min_max = min_max
+    self._coords_range = coords_range
+    self._corr_angles = [30, 45, 60, 90, 120, 135, 150]
+    # Create all masks
+    self._masks = [(self._get_ring_mask(mask_min, mask_max), (mask_min,
+                                                              mask_max))
+                   for mask_min, mask_max in mask_parameters]
+    # Mask for hiding the parts of the SAC that are never used
+    self._plotting_sac_mask = circle_mask(
+        [self._nbins * 2 - 1, self._nbins * 2 - 1],
+        self._nbins,
+        in_val=1.0,
+        out_val=np.nan)
+
+  def calculate_ratemap(self, xs, ys, activations, statistic='mean'):
+    return scipy.stats.binned_statistic_2d(
+        xs,
+        ys,
+        activations,
+        bins=self._nbins,
+        statistic=statistic,
+        range=self._coords_range)[0]
+
+  def _get_ring_mask(self, mask_min, mask_max):
+    n_points = [self._nbins * 2 - 1, self._nbins * 2 - 1]
+    return (circle_mask(n_points, mask_max * self._nbins) *
+            (1 - circle_mask(n_points, mask_min * self._nbins)))
+
+  def grid_score_60(self, corr):
+    if self._min_max:
+      return np.minimum(corr[60], corr[120]) - np.maximum(
+          corr[30], np.maximum(corr[90], corr[150]))
+    else:
+      return (corr[60] + corr[120]) / 2 - (corr[30] + corr[90] + corr[150]) / 3
+
+  def grid_score_90(self, corr):
+    return corr[90] - (corr[45] + corr[135]) / 2
+
+  def calculate_sac(self, seq1):
+    """Calculating spatial autocorrelogram."""
+    seq2 = seq1
+
+    def filter2(b, x):
+      stencil = np.rot90(b, 2)
+      return scipy.signal.convolve2d(x, stencil, mode='full')
+
+    seq1 = np.nan_to_num(seq1)
+    seq2 = np.nan_to_num(seq2)
+
+    ones_seq1 = np.ones(seq1.shape)
+    ones_seq1[np.isnan(seq1)] = 0
+    ones_seq2 = np.ones(seq2.shape)
+    ones_seq2[np.isnan(seq2)] = 0
+
+    seq1[np.isnan(seq1)] = 0
+    seq2[np.isnan(seq2)] = 0
+
+    seq1_sq = np.square(seq1)
+    seq2_sq = np.square(seq2)
+
+    seq1_x_seq2 = filter2(seq1, seq2)
+    sum_seq1 = filter2(seq1, ones_seq2)
+    sum_seq2 = filter2(ones_seq1, seq2)
+    sum_seq1_sq = filter2(seq1_sq, ones_seq2)
+    sum_seq2_sq = filter2(ones_seq1, seq2_sq)
+    n_bins = filter2(ones_seq1, ones_seq2)
+    n_bins_sq = np.square(n_bins)
+
+    std_seq1 = np.power(
+        np.subtract(
+            np.divide(sum_seq1_sq, n_bins),
+            (np.divide(np.square(sum_seq1), n_bins_sq))), 0.5)
+    std_seq2 = np.power(
+        np.subtract(
+            np.divide(sum_seq2_sq, n_bins),
+            (np.divide(np.square(sum_seq2), n_bins_sq))), 0.5)
+    covar = np.subtract(
+        np.divide(seq1_x_seq2, n_bins),
+        np.divide(np.multiply(sum_seq1, sum_seq2), n_bins_sq))
+    x_coef = np.divide(covar, np.multiply(std_seq1, std_seq2))
+    x_coef = np.real(x_coef)
+    x_coef = np.nan_to_num(x_coef)
+    return x_coef
+
+  def rotated_sacs(self, sac, angles):
+    return [
+        scipy.ndimage.interpolation.rotate(sac, angle, reshape=False)
+        for angle in angles
+    ]
+
+  def get_grid_scores_for_mask(self, sac, rotated_sacs, mask):
+    """Calculate Pearson correlations of area inside mask at corr_angles."""
+    masked_sac = sac * mask
+    ring_area = np.sum(mask)
+    # Calculate dc on the ring area
+    masked_sac_mean = np.sum(masked_sac) / ring_area
+    # Center the sac values inside the ring
+    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):
+      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
+    return self.grid_score_60(corrs), self.grid_score_90(corrs), variance
+
+  def get_scores(self, rate_map):
+    """Get summary of scrores for grid cells."""
+    sac = self.calculate_sac(rate_map)
+    rotated_sacs = self.rotated_sacs(sac, self._corr_angles)
+
+    scores = [
+        self.get_grid_scores_for_mask(sac, rotated_sacs, mask)
+        for mask, mask_params in self._masks  # pylint: disable=unused-variable
+    ]
+    scores_60, scores_90, variances = map(np.asarray, zip(*scores))  # pylint: disable=unused-variable
+    max_60_ind = np.argmax(scores_60)
+    max_90_ind = np.argmax(scores_90)
+
+    return (scores_60[max_60_ind], scores_90[max_90_ind],
+            self._masks[max_60_ind][1], self._masks[max_90_ind][1], sac, max_60_ind)
+
+  def plot_ratemap(self, ratemap, ax=None, title=None, *args, **kwargs):  # pylint: disable=keyword-arg-before-vararg
+    """Plot ratemaps."""
+    if ax is None:
+      ax = plt.gca()
+    # Plot the ratemap
+    ax.imshow(ratemap, interpolation='none', *args, **kwargs)
+    # ax.pcolormesh(ratemap, *args, **kwargs)
+    ax.axis('off')
+    if title is not None:
+      ax.set_title(title)
+
+  def plot_sac(self,
+               sac,
+               mask_params=None,
+               ax=None,
+               title=None,
+               *args,
+               **kwargs):  # pylint: disable=keyword-arg-before-vararg
+    """Plot spatial autocorrelogram."""
+    if ax is None:
+      ax = plt.gca()
+    # Plot the sac
+    useful_sac = sac * self._plotting_sac_mask
+    ax.imshow(useful_sac, interpolation='none', *args, **kwargs)
+    # ax.pcolormesh(useful_sac, *args, **kwargs)
+    # Plot a ring for the adequate mask
+    if mask_params is not None:
+      center = self._nbins - 1
+      ax.add_artist(
+          plt.Circle(
+              (center, center),
+              mask_params[0] * self._nbins,
+              # lw=bump_size,
+              fill=False,
+              edgecolor='k'))
+      ax.add_artist(
+          plt.Circle(
+              (center, center),
+              mask_params[1] * self._nbins,
+              # lw=bump_size,
+              fill=False,
+              edgecolor='k'))
+    ax.axis('off')
+    if title is not None:
+      ax.set_title(title)
+
+
+def border_score(rm, res, box_width):
+	# Find connected firing fields
+    pix_area = 100**2*box_width**2/res**2
+    rm_thresh =  rm>(rm.max()*0.3)
+    rm_comps, ncomps = ndimage.measurements.label(rm_thresh)
+
+    # Keep fields with area > 200cm^2
+    masks = []
+    nfields = 0
+    for i in range(1,ncomps+1):
+        mask = (rm_comps==i).reshape(res,res)
+        if mask.sum()*pix_area > 200:
+            masks.append(mask)
+            nfields += 1
+            
+    # Max coverage of any one field over any one border
+    cm_max = 0
+    for mask in masks:
+        mask = masks[0]
+        n_cov = mask[0].mean()
+        s_cov = mask[-1].mean()
+        e_cov = mask[:,0].mean()
+        w_cov = mask[:,-1].mean()
+        cm = np.max([n_cov,s_cov,e_cov,w_cov])
+        if cm>cm_max:
+            cm_max = cm
+
+    # Distance to nearest wall
+    x,y = np.mgrid[:res,:res] + 1
+    x = x.ravel()
+    y = y.ravel()
+    xmin = np.min(np.vstack([x,res+1-x]),0)
+    ymin = np.min(np.vstack([y,res+1-y]),0)
+    dweight = np.min(np.vstack([xmin,ymin]),0).reshape(res,res)
+    dweight = dweight*box_width/res
+
+    # Mean firing distance
+    dms = []
+    for mask in masks:
+        field = rm[mask]
+        field /= field.sum()   # normalize
+        dm = (field*dweight[mask]).sum()
+        dms.append(dm)
+    dm = np.nanmean(dms) / (box_width/2)
+    border_score = (cm_max-dm)/(cm_max+dm)
+    return border_score, cm_max, dm

trainer.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+import torch
+import numpy as np
+import datetime
+from tqdm import tqdm
+
+from visualize import save_ratemaps
+import os
+
+
+class Trainer(object):
+    def __init__(self, options, model, trajectory_generator, restore=True):
+        self.options = options
+        self.model = model
+        self.trajectory_generator = trajectory_generator
+        lr = self.options.learning_rate
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
+
+        self.loss = []
+        self.err = []
+
+        # Set up checkpoints
+        self.ckpt_dir = os.path.join(options.save_dir, options.run_ID)
+        ckpt_path = os.path.join(self.ckpt_dir, "most_recent_model.pth")
+        pc_path = os.path.join(self.ckpt_dir, "place_cell_centers.npy")
+        if restore and os.path.isdir(self.ckpt_dir) and os.path.isfile(ckpt_path):
+            self.model.load_state_dict(torch.load(ckpt_path))
+            print("Restored trained model from {}".format(ckpt_path))
+            # 加载 place cell 位置
+            if os.path.isfile(pc_path):
+                us = np.load(pc_path)
+                self.model.place_cells.us = torch.tensor(us).to(options.device)
+                print("Restored place cell centers from {}".format(pc_path))
+            else:
+                print("Warning: place_cell_centers.npy not found! Model may not work correctly.")
+        else:
+            if not os.path.isdir(self.ckpt_dir):
+                os.makedirs(self.ckpt_dir, exist_ok=True)
+            print("Initializing new model from scratch.")
+            print("Saving to: {}".format(self.ckpt_dir))
+
+    def train_step(self, inputs, pc_outputs, pos):
+        """
+        Train on one batch of trajectories.
+
+        Args:
+            inputs: Batch of 2d velocity inputs with shape [batch_size, sequence_length, 2].
+            pc_outputs: Ground truth place cell activations with shape
+                [batch_size, sequence_length, Np].
+            pos: Ground truth 2d position with shape [batch_size, sequence_length, 2].
+
+        Returns:
+            loss: Avg. loss for this training batch.
+            err: Avg. decoded position error in cm.
+        """
+        self.model.zero_grad()
+
+        loss, err = self.model.compute_loss(inputs, pc_outputs, pos)
+
+        loss.backward()
+        self.optimizer.step()
+
+        return loss.item(), err.item()
+
+    def train(self, n_epochs: int = 1000, n_steps=10, save=True):
+        """
+        Train model on simulated trajectories.
+
+        Args:
+            n_steps: Number of training steps
+            save: If true, save a checkpoint after each epoch.
+        """
+
+        # Construct generator
+        gen = self.trajectory_generator.get_generator()
+
+        for epoch_idx in range(n_epochs):
+            tbar = tqdm(range(n_steps), leave=False)
+            for step_idx in tbar:
+                inputs, pc_outputs, pos = next(gen)
+                loss, err = self.train_step(inputs, pc_outputs, pos)
+                self.loss.append(loss)
+                self.err.append(err)
+
+                # Log error rate to progress bar
+                tbar.set_description('Error = ' + str(int(100*err)) + 'cm')
+
+            if save and ((epoch_idx + 1) % 10 == 0 or epoch_idx == 0):
+                # Save checkpoint
+                # ckpt_path = os.path.join(
+                #     self.ckpt_dir, "epoch_{}.pth".format(epoch_idx)
+                # )
+                # torch.save(self.model.state_dict(), ckpt_path)
+                torch.save(
+                    self.model.state_dict(),
+                    os.path.join(self.ckpt_dir, "most_recent_model.pth"),
+                )
+                # 保存 place cell 位置
+                np.save(
+                    os.path.join(self.ckpt_dir, "place_cell_centers.npy"),
+                    self.model.place_cells.us.cpu().numpy(),
+                )
+
+                # Save a picture of rate maps
+                save_ratemaps(
+                    self.model,
+                    self.trajectory_generator,
+                    self.options,
+                    step=epoch_idx + 1,
+                )
+
+                print(
+                    "Epoch: {}/{}. Date: {}. Loss: {}. Err: {}cm".format(
+                        epoch_idx + 1,
+                        n_epochs,
+                        str(datetime.datetime.now())[:-7],
+                        np.round(loss, 2),
+                        np.round(100 * err, 2),
+                    )
+                )

trajectory_generator.py

@@ -0,0 +1,243 @@
+# -*- coding: utf-8 -*-
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+class TrajectoryGenerator(object):
+    def __init__(self, options, place_cells):
+        self.options = options
+        self.place_cells = place_cells
+
+    def plot_trajectory(self, traj, box_width, box_height, idx=0, step=2):
+        """
+        Visualize one trajectory from traj dict.
+
+        Args:
+            traj: dictionary containing trajectory info
+            box_width, box_height: dimensions of the environment
+            idx: which trajectory to plot from the batch (default: 0)
+            step: plot an arrow every 'step' frames
+        """
+        # Extract trajectory for one rat
+        x = traj["target_x"][idx]  # shape (samples,)
+        y = traj["target_y"][idx]
+        hd = traj["target_hd"][idx]  # head directions in radians
+
+        # Also add starting point
+        x0 = traj["init_x"][idx, 0]
+        y0 = traj["init_y"][idx, 0]
+        hd0 = traj["init_hd"][idx, 0]
+
+        x = np.concatenate([[x0], x])
+        y = np.concatenate([[y0], y])
+        hd = np.concatenate([[hd0], hd])
+
+        # Plot trajectory
+        plt.figure(figsize=(6, 6))
+        plt.plot(x, y, "-o", markersize=2, label="trajectory")
+
+        # Add arrows for head direction
+        for t in range(0, len(x), step):
+            dx = 0.1 * np.cos(hd[t])
+            dy = 0.1 * np.sin(hd[t])
+            plt.arrow(
+                x[t], y[t], dx, dy, head_width=0.05, head_length=0.08, fc="r", ec="r"
+            )
+
+        # Draw box boundaries
+        plt.axhline(y=-box_height / 2, color="k")
+        plt.axhline(y=box_height / 2, color="k")
+        plt.axvline(x=-box_width / 2, color="k")
+        plt.axvline(x=box_width / 2, color="k")
+
+        plt.xlim([-box_width / 2 - 0.2, box_width / 2 + 0.2])
+        plt.ylim([-box_height / 2 - 0.2, box_height / 2 + 0.2])
+        plt.gca().set_aspect("equal", adjustable="box")
+        plt.xlabel("x position (m)")
+        plt.ylabel("y position (m)")
+        plt.title(f"Trajectory {idx}")
+        plt.legend()
+        plt.show()
+
+    def avoid_wall(self, position, hd, box_width, box_height):
+        """
+        Compute distance and angle to nearest wall
+        """
+        x = position[:, 0]
+        y = position[:, 1]
+        dists = [
+            box_width / 2 - x,
+            box_height / 2 - y,
+            box_width / 2 + x,
+            box_height / 2 + y,
+        ]
+        d_wall = np.min(dists, axis=0)
+        angles = np.arange(4) * np.pi / 2
+        theta = angles[np.argmin(dists, axis=0)]
+        hd = np.mod(hd, 2 * np.pi)
+        a_wall = hd - theta
+        a_wall = np.mod(a_wall + np.pi, 2 * np.pi) - np.pi
+
+        is_near_wall = (d_wall < self.border_region) * (np.abs(a_wall) < np.pi / 2)
+        turn_angle = np.zeros_like(hd)
+        turn_angle[is_near_wall] = np.sign(a_wall[is_near_wall]) * (
+            np.pi / 2 - np.abs(a_wall[is_near_wall])
+        )
+
+        return is_near_wall, turn_angle
+
+    def generate_trajectory(self, box_width, box_height, batch_size):
+        """Generate a random walk in a rectangular box"""
+        samples = self.options.sequence_length
+        dt = 0.02  # time step increment (seconds)
+        sigma = 5.76 * 2  # stdev rotation velocity (rads/sec)
+        b = 0.13 * 2 * np.pi  # forward velocity rayleigh dist scale (m/sec)
+        mu = 0  # turn angle bias
+        self.border_region = 0.03  # meters
+
+        # Initialize variables
+        position = np.zeros([batch_size, samples + 2, 2])
+        head_dir = np.zeros([batch_size, samples + 2])
+        position[:, 0, 0] = np.random.uniform(-box_width / 2, box_width / 2, batch_size)
+        position[:, 0, 1] = np.random.uniform(
+            -box_height / 2, box_height / 2, batch_size
+        )
+        head_dir[:, 0] = np.random.uniform(0, 2 * np.pi, batch_size)
+        velocity = np.zeros([batch_size, samples + 2])
+
+        # Generate sequence of random boosts and turns
+        random_turn = np.random.normal(mu, sigma, [batch_size, samples + 1])
+        random_vel = np.random.rayleigh(b, [batch_size, samples + 1])
+        v = np.abs(np.random.normal(0, b * np.pi / 2, batch_size))
+
+        for t in range(samples + 1):
+            # Update velocity
+            v = random_vel[:, t]
+            turn_angle = np.zeros(batch_size)
+
+            if not self.options.periodic:
+                # If in border region, turn and slow down
+                is_near_wall, turn_angle = self.avoid_wall(
+                    position[:, t], head_dir[:, t], box_width, box_height
+                )
+                v[is_near_wall] *= 0.25
+
+            # Update turn angle
+            turn_angle += dt * random_turn[:, t]
+
+            # Take a step
+            velocity[:, t] = v * dt
+            update = velocity[:, t, None] * np.stack(
+                [np.cos(head_dir[:, t]), np.sin(head_dir[:, t])], axis=-1
+            )
+            position[:, t + 1] = position[:, t] + update
+
+            # Rotate head direction
+            head_dir[:, t + 1] = head_dir[:, t] + turn_angle
+
+        # Periodic boundaries
+        if self.options.periodic:
+            position[:, :, 0] = (
+                np.mod(position[:, :, 0] + box_width / 2, box_width) - box_width / 2
+            )
+            position[:, :, 1] = (
+                np.mod(position[:, :, 1] + box_height / 2, box_height) - box_height / 2
+            )
+
+        head_dir = np.mod(head_dir + np.pi, 2 * np.pi) - np.pi  # Periodic variable
+
+        traj = {}
+        # Input variables
+        traj["init_hd"] = head_dir[:, 0, None]
+        traj["init_x"] = position[:, 1, 0, None]
+        traj["init_y"] = position[:, 1, 1, None]
+
+        traj["ego_v"] = velocity[:, 1:-1]
+        ang_v = np.diff(head_dir, axis=-1)
+        traj["phi_x"], traj["phi_y"] = np.cos(ang_v)[:, :-1], np.sin(ang_v)[:, :-1]
+
+        # Target variables
+        traj["target_hd"] = head_dir[:, 1:-1]
+        traj["target_x"] = position[:, 2:, 0]
+        traj["target_y"] = position[:, 2:, 1]
+
+        # for i in range(5):
+        #     self.plot_trajectory(traj, box_width, box_height, i)
+        # raise Exception("dog")
+
+        return traj
+
+    def get_generator(self, batch_size=None, box_width=None, box_height=None):
+        """
+        Returns a generator that yields batches of trajectories
+        """
+        if not batch_size:
+            batch_size = self.options.batch_size
+        if not box_width:
+            box_width = self.options.box_width
+        if not box_height:
+            box_height = self.options.box_height
+
+        while True:
+            traj = self.generate_trajectory(box_width, box_height, batch_size)
+
+            v = np.stack(
+                [
+                    traj["ego_v"] * np.cos(traj["target_hd"]),
+                    traj["ego_v"] * np.sin(traj["target_hd"]),
+                ],
+                axis=-1,
+            )
+            v = torch.tensor(v, dtype=torch.float32).transpose(0, 1)
+
+            pos = np.stack([traj["target_x"], traj["target_y"]], axis=-1)
+            pos = torch.tensor(pos, dtype=torch.float32).transpose(0, 1)
+            # Put on GPU if GPU is available
+            pos = pos.to(self.options.device)
+            place_outputs = self.place_cells.get_activation(pos)
+
+            init_pos = np.stack([traj["init_x"], traj["init_y"]], axis=-1)
+            init_pos = torch.tensor(init_pos, dtype=torch.float32)
+            init_pos = init_pos.to(self.options.device)
+            init_actv = self.place_cells.get_activation(init_pos).squeeze()
+
+            v = v.to(self.options.device)
+            inputs = (v, init_actv)
+
+            yield (inputs, place_outputs, pos)
+
+    def get_test_batch(self, batch_size=None, box_width=None, box_height=None):
+        """For testing performance, returns a batch of smample trajectories"""
+        if not batch_size:
+            batch_size = self.options.batch_size
+        if not box_width:
+            box_width = self.options.box_width
+        if not box_height:
+            box_height = self.options.box_height
+
+        traj = self.generate_trajectory(box_width, box_height, batch_size)
+
+        v = np.stack(
+            [
+                traj["ego_v"] * np.cos(traj["target_hd"]),
+                traj["ego_v"] * np.sin(traj["target_hd"]),
+            ],
+            axis=-1,
+        )
+        v = torch.tensor(v, dtype=torch.float32).transpose(0, 1)
+
+        pos = np.stack([traj["target_x"], traj["target_y"]], axis=-1)
+        pos = torch.tensor(pos, dtype=torch.float32).transpose(0, 1)
+        pos = pos.to(self.options.device)
+        place_outputs = self.place_cells.get_activation(pos)
+
+        init_pos = np.stack([traj["init_x"], traj["init_y"]], axis=-1)
+        init_pos = torch.tensor(init_pos, dtype=torch.float32)
+        init_pos = init_pos.to(self.options.device)
+        init_actv = self.place_cells.get_activation(init_pos).squeeze()
+
+        v = v.to(self.options.device)
+        inputs = (v, init_actv)
+
+        return (inputs, pos, place_outputs)

utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+
+
+def generate_run_ID(options):
+    ''' 
+    Create a unique run ID from the most relevant
+    parameters. Remaining parameters can be found in 
+    params.npy file. 
+    '''
+    params = [
+        'steps', str(options.sequence_length),
+        'batch', str(options.batch_size),
+        options.RNN_type,
+        str(options.Ng),
+        options.activation,
+        'rf', str(options.place_cell_rf),
+        'DoG', str(options.DoG),
+        'periodic', str(options.periodic),
+        'lr', str(options.learning_rate),
+        'weight_decay', str(options.weight_decay),
+        ]
+    separator = '_'
+    run_ID = separator.join(params)
+    run_ID = run_ID.replace('.', '')
+
+    return run_ID
+
+
+def get_2d_sort(x1,x2):
+    """
+    Reshapes x1 and x2 into square arrays, and then sorts
+    them such that x1 increases downward and x2 increases
+    rightward. Returns the order.
+    """
+    n = int(np.round(np.sqrt(len(x1))))
+    total_order = x1.argsort()
+    total_order = total_order.reshape(n,n)
+    for i in range(n):
+        row_order = x2[total_order.ravel()].reshape(n,n)[i].argsort()
+        total_order[i] = total_order[i,row_order]
+    total_order = total_order.ravel()
+    return total_order
+
+
+def dft(N,real=False,scale='sqrtn'):
+    if not real:
+        return scipy.linalg.dft(N,scale)
+    else:
+        cosines = np.cos(2*np.pi*np.arange(N//2+1)[None,:]/N*np.arange(N)[:,None])
+        sines = np.sin(2*np.pi*np.arange(1,(N-1)//2+1)[None,:]/N*np.arange(N)[:,None])
+        if N%2==0:
+            cosines[:,-1] /= np.sqrt(2)
+        F = np.concatenate((cosines,sines[:,::-1]),1)
+        F[:,0] /= np.sqrt(N)
+        F[:,1:] /= np.sqrt(N/2)
+        return F
+
+
+def skaggs_power(Jsort):
+    F = dft(int(np.sqrt(N)), real=True)
+    F2d = F[:,None,:,None]*F[None,:,None,:]
+
+    F2d_unroll = np.reshape(F2d, (N, N))
+
+    F2d_inv = F2d_unroll.conj().T
+    Jtilde = F2d_inv.dot(Jsort).dot(F2d_unroll)
+
+    return (Jtilde[1,1]**2 + Jtilde[-1,-1]**2) / (Jtilde**2).sum()
+
+
+def skaggs_power_2(Jsort):
+    J_square = np.reshape(Jsort, (n,n,n,n))
+    Jmean = np.zeros([n,n])
+    for i in range(n):
+        for j in range(n):
+            Jmean += np.roll(np.roll(J_square[i,j], -i, axis=0), -j, axis=1)
+
+#     Jmean[0,0] = np.max(Jmean[1:,1:])
+    Jmean = np.roll(np.roll(Jmean, n//2, axis=0), n//2, axis=1)
+    Jtilde = np.real(np.fft.fft2(Jmean))
+    
+    Jtilde[0,0] = 0
+    sk_power = Jtilde[1,1]**2 + Jtilde[0,1]**2 + Jtilde[1,0]**2
+    sk_power += Jtilde[-1,-1]**2 + Jtilde[0,-1]**2 + Jtilde[-1,0]**2
+    sk_power /= (Jtilde**2).sum()
+    
+    return sk_power
+
+
+def calc_err():
+    inputs, _, pos = next(gen)
+    pred = model(inputs)
+    pred_pos = place_cells.get_nearest_cell_pos(pred)
+    return tf.reduce_mean(tf.sqrt(tf.reduce_sum((pos - pred_pos)**2, axis=-1)))
+
+from visualize import compute_ratemaps, plot_ratemaps
+
+
+def compute_variance(res, n_avg):
+    
+    activations, rate_map, g, pos = compute_ratemaps(model, data_manager, options, res=res, n_avg=n_avg)
+
+    counts = np.zeros([res,res])
+    variance = np.zeros([res,res])
+
+    x_all = (pos[:,0] + options['box_width']/2) / options['box_width'] * res
+    y_all = (pos[:,1] + options['box_height']/2) / options['box_height'] * res
+    for i in tqdm(range(len(g))):
+        x = int(x_all[i])
+        y = int(y_all[i])
+        if x >=0 and x < res and y >=0 and y < res:
+            counts[x, y] += 1
+            variance[x, y] += np.linalg.norm(g[i] - activations[:, x, y]) / np.linalg.norm(g[i]) / np.linalg.norm(activations[:,x,y])
+
+    for x in range(res):
+        for y in range(res):
+            if counts[x, y] > 0:
+                variance[x, y] /= counts[x, y]
+                
+    return variance
+
+
+def load_trained_weights(model, trainer, weight_dir):
+    ''' Load weights stored as a .npy file (for github)'''
+
+    # Train for a single step to initialize weights
+    # trainer.train(n_epochs=1, n_steps=1, save=False)
+
+    # Load weights from npy array
+    weights = np.load(weight_dir, allow_pickle=True)
+    model.set_weights(weights)
+    print('Loaded trained weights.')
+

visualize.py

@@ -0,0 +1,178 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from matplotlib import pyplot as plt
+
+import scipy
+import scipy.stats
+from imageio import imsave
+import cv2
+
+
+def concat_images(images, image_width, spacer_size):
+    """ Concat image horizontally with spacer """
+    spacer = np.ones([image_width, spacer_size, 4], dtype=np.uint8) * 255
+    images_with_spacers = []
+
+    image_size = len(images)
+
+    for i in range(image_size):
+        images_with_spacers.append(images[i])
+        if i != image_size - 1:
+            # Add spacer
+            images_with_spacers.append(spacer)
+    ret = np.hstack(images_with_spacers)
+    return ret
+
+
+def concat_images_in_rows(images, row_size, image_width, spacer_size=4):
+    """ Concat images in rows """
+    column_size = len(images) // row_size
+    spacer_h = np.ones([spacer_size, image_width*column_size + (column_size-1)*spacer_size, 4],
+                       dtype=np.uint8) * 255
+
+    row_images_with_spacers = []
+
+    for row in range(row_size):
+        row_images = images[column_size*row:column_size*row+column_size]
+        row_concated_images = concat_images(row_images, image_width, spacer_size)
+        row_images_with_spacers.append(row_concated_images)
+
+        if row != row_size-1:
+            row_images_with_spacers.append(spacer_h)
+
+    ret = np.vstack(row_images_with_spacers)
+    return ret
+
+
+def convert_to_colormap(im, cmap):
+    im = cmap(im)
+    im = np.uint8(im * 255)
+    return im
+
+
+def rgb(im, cmap='jet', smooth=True):
+    cmap = plt.cm.get_cmap(cmap)
+    np.seterr(invalid='ignore')  # ignore divide by zero err
+    im = (im - np.min(im)) / (np.max(im) - np.min(im))
+    if smooth:
+        im = cv2.GaussianBlur(im, (3,3), sigmaX=1, sigmaY=0)
+    im = cmap(im)
+    im = np.uint8(im * 255)
+    return im
+
+
+def plot_ratemaps(activations, n_plots, cmap='jet', smooth=True, width=16):
+    images = [rgb(im, cmap, smooth) for im in activations[:n_plots]]
+    rm_fig = concat_images_in_rows(images, n_plots//width, activations.shape[-1])
+    return rm_fig
+
+
+def compute_ratemaps(model, trajectory_generator, options, res=20, n_avg=None, Ng=512, idxs=None, return_raw=False):
+    '''Compute spatial firing fields
+    
+    Args:
+        model: The RNN model
+        trajectory_generator: Generator for test trajectories
+        options: Training options
+        res: Resolution of the rate map grid
+        n_avg: Number of batches to average over
+        Ng: Number of grid cells to analyze
+        idxs: Indices of specific grid cells to analyze
+        return_raw: If True, also return raw activations (g) and positions (pos).
+                   Warning: This uses significant memory for large batch_size/n_avg.
+                   If False, returns None for g and pos to save memory.
+    
+    Returns:
+        activations: Spatial firing fields [Ng, res, res]
+        rate_map: Flattened rate maps [Ng, res*res]
+        g: Raw activations (None if return_raw=False)
+        pos: Raw positions (None if return_raw=False)
+    '''
+
+    if not n_avg:
+        n_avg = 1000 // options.sequence_length
+
+    if not np.any(idxs):
+        idxs = np.arange(Ng)
+    idxs = idxs[:Ng]
+
+    # Only allocate large arrays if return_raw is True
+    if return_raw:
+        g = np.zeros([n_avg, options.batch_size * options.sequence_length, Ng])
+        pos = np.zeros([n_avg, options.batch_size * options.sequence_length, 2])
+    else:
+        g = None
+        pos = None
+
+    activations = np.zeros([Ng, res, res]) 
+    counts  = np.zeros([res, res])
+
+    for index in range(n_avg):
+        inputs, pos_batch, _ = trajectory_generator.get_test_batch()
+        g_batch = model.g(inputs).detach().cpu().numpy()
+        
+        pos_batch = np.reshape(pos_batch.cpu(), [-1, 2])
+        g_batch = g_batch[:,:,idxs].reshape(-1, Ng)
+        
+        if return_raw:
+            g[index] = g_batch
+            pos[index] = pos_batch
+
+        x_batch = (pos_batch[:,0] + options.box_width/2) / (options.box_width) * res
+        y_batch = (pos_batch[:,1] + options.box_height/2) / (options.box_height) * res
+
+        for i in range(options.batch_size*options.sequence_length):
+            x = x_batch[i]
+            y = y_batch[i]
+            if x >=0 and x < res and y >=0 and y < res:
+                counts[int(x), int(y)] += 1
+                activations[:, int(x), int(y)] += g_batch[i, :]
+
+    for x in range(res):
+        for y in range(res):
+            if counts[x, y] > 0:
+                activations[:, x, y] /= counts[x, y]
+                
+    if return_raw:
+        g = g.reshape([-1, Ng])
+        pos = pos.reshape([-1, 2])
+
+    # # scipy binned_statistic_2d is slightly slower
+    # activations = scipy.stats.binned_statistic_2d(pos[:,0], pos[:,1], g.T, bins=res)[0]
+    rate_map = activations.reshape(Ng, -1)
+
+    return activations, rate_map, g, pos
+
+
+def save_ratemaps(model, trajectory_generator, options, step, res=20, n_avg=None):
+    if not n_avg:
+        n_avg = 1000 // options.sequence_length
+    activations, rate_map, g, pos = compute_ratemaps(model, trajectory_generator,
+                                                     options, res=res, n_avg=n_avg)
+    rm_fig = plot_ratemaps(activations, n_plots=len(activations))
+    imdir = options.save_dir + "/" + options.run_ID
+    imsave(imdir + "/" + str(step) + ".png", rm_fig)
+
+
+def save_autocorr(sess, model, save_name, trajectory_generator, step, flags):
+    starts = [0.2] * 10
+    ends = np.linspace(0.4, 1.0, num=10)
+    coord_range=((-1.1, 1.1), (-1.1, 1.1))
+    masks_parameters = zip(starts, ends.tolist())
+    latest_epoch_scorer = scores.GridScorer(20, coord_range, masks_parameters)
+    
+    res = dict()
+    index_size = 100
+    for _ in range(index_size):
+      feed_dict = trajectory_generator.feed_dict(flags.box_width, flags.box_height)
+      mb_res = sess.run({
+          'pos_xy': model.target_pos,
+          'bottleneck': model.g,
+      }, feed_dict=feed_dict)
+      res = utils.concat_dict(res, mb_res)
+        
+    filename = save_name + '/autocorrs_' + str(step) + '.pdf'
+    imdir = flags.save_dir + '/'
+    out = utils.get_scores_and_plot(
+                latest_epoch_scorer, res['pos_xy'], res['bottleneck'],
+                imdir, filename)