Merge pull request #168 from geometric-intelligence/pirnn_refactor

Pirnn refactor

.gitignore

@@ -8,6 +8,8 @@ neurometry/datasets/rnn_grid_cells/Dual agent path integration high res/*
 neurometry/datasets/rnn_grid_cells/Single agent path integration high res/*
 neurometry/curvature/grid-cells-curvature/models/xu_rnn/results/*
 neurometry/curvature/grid-cells-curvature/multi-agent/*
+neurometry/neuroai/piRNNs/models/results/*
+neurometry/neuroai/piRNNs/multi-agent/*
 notebooks/
 
 *viewer*
@@ -36,6 +38,9 @@ neurometry/datasets/rnn_grid_cells/Single agent path integration/*
 
 neurometry/curvature/grid-cells-curvature/models/xu_rnn/logs/*
 neurometry/curvature/grid-cells-curvature/models/xu_rnn/wandb/*
+neurometry/neuroai/piRNNs/models/logs/*
+neurometry/neuroai/piRNNs/models/wandb/*
+neurometry/neuroai/piRNNs/models/pretrained/*
 
 
 # Result files

neurometry/curvature/datasets/gridcells.py

@@ -3,11 +3,11 @@
 import numpy as np
 import pandas as pd
 
+import neurometry.curvature.datasets.structures as structures
+
 os.environ["GEOMSTATS_BACKEND"] = "pytorch"
 import geomstats.backend as gs
 
-import neurometry.curvature.datasets.structures as structures
-
 
 # TODO
 def load_grid_cells_synthetic(

neurometry/curvature/datasets/synthetic.py

@@ -3,20 +3,21 @@
 import logging
 import os
 
-os.environ["GEOMSTATS_BACKEND"] = "pytorch"
-import geomstats.backend as gs
 import numpy as np
 import pandas as pd
 import skimage
 import torch
-from geomstats.geometry.special_orthogonal import SpecialOrthogonal
 from torch.distributions.multivariate_normal import MultivariateNormal
 
 from neurometry.topology.persistent_homology import (
     cohomological_circular_coordinates,
     cohomological_toroidal_coordinates,
 )
 
+os.environ["GEOMSTATS_BACKEND"] = "pytorch"
+import geomstats.backend as gs
+from geomstats.geometry.special_orthogonal import SpecialOrthogonal
+
 
 def load_projected_images(n_scalars=5, n_angles=1000, img_size=128):
     """Load a dataset of 2D images projected into 1D projections.

neurometry/curvature/evaluate.py

@@ -4,19 +4,19 @@
 import numpy as np
 import torch
 
+from neurometry.curvature.datasets.synthetic import (
+    get_s1_synthetic_immersion,
+    get_s2_synthetic_immersion,
+    get_t2_synthetic_immersion,
+)
+
 os.environ["GEOMSTATS_BACKEND"] = "pytorch"
 import geomstats.backend as gs
 from geomstats.geometry.base import ImmersedSet
 from geomstats.geometry.euclidean import Euclidean
 from geomstats.geometry.pullback_metric import PullbackMetric
 from geomstats.geometry.special_orthogonal import SpecialOrthogonal
 
-from neurometry.curvature.datasets.synthetic import (
-    get_s1_synthetic_immersion,
-    get_s2_synthetic_immersion,
-    get_t2_synthetic_immersion,
-)
-
 
 class NeuralManifoldIntrinsic(ImmersedSet):
     def __init__(self, dim, neural_embedding_dim, neural_immersion, equip=True):