Merge pull request #141 from geometric-intelligence/ninamiolane-lint

Lint majority of repo

Author: ninamiolane

neurometry/curvature/datasets/gridcells.py

@@ -6,9 +6,9 @@
 import pandas as pd
 
 os.environ["GEOMSTATS_BACKEND"] = "pytorch"
-import geomstats.backend as gs
+import geomstats.backend as gs  # noqa: E402
 
-import neurometry.curvature.datasets.structures as structures
+import neurometry.curvature.datasets.structures as structures  # noqa: E402
 
 
 # TODO
@@ -96,7 +96,7 @@ def generate_all_grids(
     grids : numpy.ndarray, shape=(num_cells, num_fields_per_cell = (ceil(dims[0]/lx)+1)*(ceil(dims[1]/ly)+1),2)
         All the grid cell lattices.
     """
-
+    lx = ly = 10   # TODO: FIX, these values are only placeholders.
     # ref_lattice = create_reference_lattice(lx, ly, arena_dims, lattice_type = lattice_type)
     ref_lattice = structures.get_lattice(
         scale=grid_scale, lattice_type=lattice_type, dimensions=arena_dims

neurometry/curvature/datasets/synthetic.py

@@ -4,12 +4,13 @@
 import os
 
 os.environ["GEOMSTATS_BACKEND"] = "pytorch"
-import geomstats.backend as gs
-import numpy as np
-import pandas as pd
-import torch
-from geomstats.geometry.special_orthogonal import SpecialOrthogonal  # NOQA
-from torch.distributions.multivariate_normal import MultivariateNormal
+import geomstats.backend as gs  # noqa: E402
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import skimage  # noqa: E402
+import torch  # noqa: E402
+from geomstats.geometry.special_orthogonal import SpecialOrthogonal  # noqa: E402
+from torch.distributions.multivariate_normal import MultivariateNormal  # noqa: E402
 
 
 def load_projected_images(n_scalars=5, n_angles=1000, img_size=128):

neurometry/curvature/datasets/utils.py

@@ -6,6 +6,8 @@
 import torch
 from scipy.signal import savgol_filter
 
+import neurometry.curvature.datasets as datasets
+
 
 def load(config):
     """Load dataset according to configuration in config.
@@ -161,29 +163,29 @@ def load(config):
         "s1_synthetic",
     ) or config.dataset_name in ("three_place_cells_synthetic"):
         train = []
-        for d, l in zip(train_dataset, train_labels["angles"], strict=False):
-            train.append([d, float(l)])
+        for data, label in zip(train_dataset, train_labels["angles"], strict=False):
+            train.append([data, float(label)])
         test = []
-        for d, l in zip(test_dataset, test_labels["angles"], strict=False):
-            test.append([d, float(l)])
+        for data, label in zip(test_dataset, test_labels["angles"], strict=False):
+            test.append([data, float(label)])
     elif config.dataset_name in ("s2_synthetic", "t2_synthetic"):
         train = []
-        for d, t, p in zip(
+        for data, theta, phi in zip(
             train_dataset, train_labels["thetas"], train_labels["phis"], strict=False
         ):
-            train.append([d, torch.tensor([float(t), float(p)])])
+            train.append([data, torch.tensor([float(theta), float(phi)])])
         test = []
-        for d, t, p in zip(
+        for data, theta, phi in zip(
             test_dataset, test_labels["thetas"], test_labels["phis"], strict=False
         ):
-            test.append([d, torch.tensor([float(t), float(p)])])
+            test.append([data, torch.tensor([float(theta), float(phi)])])
     elif config.dataset_name == "grid_cells":
         train = []
-        for d, l in zip(train_dataset, train_labels["no_labels"], strict=False):
-            train.append([d, float(l)])
+        for data, label in zip(train_dataset, train_labels["no_labels"], strict=False):
+            train.append([data, float(label)])
         test = []
-        for d, l in zip(test_dataset, test_labels["no_labels"], strict=False):
-            test.append([d, float(l)])
+        for data, label in zip(test_dataset, test_labels["no_labels"], strict=False):
+            test.append([data, float(label)])
 
     train_loader = torch.utils.data.DataLoader(train, batch_size=config.batch_size)
     test_loader = torch.utils.data.DataLoader(test, batch_size=config.batch_size)

neurometry/curvature/evaluate.py

@@ -1,18 +1,20 @@
 import os
+import time
+
+import numpy as np
+import torch
 
 os.environ["GEOMSTATS_BACKEND"] = "pytorch"
-import geomstats.backend as gs
+import geomstats.backend as gs  # noqa: E402
 
 # import gph
-import numpy as np
-import torch
-from datasets.synthetic import (
+from datasets.synthetic import (  # noqa: E402
     get_s1_synthetic_immersion,
     get_s2_synthetic_immersion,
     get_t2_synthetic_immersion,
 )
-from geomstats.geometry.pullback_metric import PullbackMetric
-from geomstats.geometry.special_orthogonal import SpecialOrthogonal  # NOQA
+from geomstats.geometry.pullback_metric import PullbackMetric  # noqa: E402
+from geomstats.geometry.special_orthogonal import SpecialOrthogonal  # noqa: E402
 
 
 def get_learned_immersion(model, config):

neurometry/curvature/grid-cells-curvature/notebooks/12_path_int.ipynb

@@ -26,7 +26,9 @@ def __init__(self, dim, validate_args=None, device=None):
         self._dim = dim
         self.device = device
 
-    def sample(self, shape=torch.Size()):
+    def sample(self, shape=None):
+        if shape is None:
+            shape = torch.Size()
         output = (
             torch.distributions.Normal(0, 1)
             .sample(

neurometry/curvature/hyperspherical/distributions/hyperspherical_uniform.py

@@ -44,11 +44,15 @@ def __init__(self, loc, scale, validate_args=None, k=1):
 
         super().__init__(self.loc.size(), validate_args=validate_args)
 
-    def sample(self, shape=torch.Size()):
+    def sample(self, shape=None):
+        if shape is None:
+            shape = torch.Size()
         with torch.no_grad():
             return self.rsample(shape)
 
-    def rsample(self, shape=torch.Size()):
+    def rsample(self, shape=None):
+        if shape is None:
+            shape = torch.Size()
         shape = shape if isinstance(shape, torch.Size) else torch.Size([shape])
 
         w = (

neurometry/curvature/hyperspherical/distributions/von_mises_fisher.py

@@ -2,8 +2,10 @@
 
 import torch
 
-from neurometry.curvature.hyperspherical.distributions import hyperspherical_uniform
-from neurometry.curvature.hyperspherical.distributions import von_mises_fisher
+from neurometry.curvature.hyperspherical.distributions import (
+    hyperspherical_uniform,
+    von_mises_fisher,
+)
 
 
 def elbo(x, x_mu, posterior_params, z, labels, config):
@@ -25,9 +27,11 @@ def elbo(x, x_mu, posterior_params, z, labels, config):
     x : array-like, shape=[batch_size, data_dim]
         Input data.
     gen_likelihood_params : tuple
-        Learned distributional parameters of generative model. (e.g., (x_mu,x_logvar) for Gaussian).
+        Learned distributional parameters of generative model.
+        (e.g., (x_mu,x_logvar) for Gaussian).
     posterior_params : tuple
-        Learned distributional parameters of approximate posterior. (e.g., (z_mu,z_logvar) for Gaussian).
+        Learned distributional parameters of approximate posterior.
+        (e.g., (z_mu,z_logvar) for Gaussian).
     config : module
         Module specifying various model hyperparameters
 
@@ -52,9 +56,10 @@ def elbo(x, x_mu, posterior_params, z, labels, config):
 
     if config.posterior_type == "toroidal":
         z_theta_mu, z_theta_kappa, z_phi_mu, z_phi_kappa = posterior_params
-        q_z_theta = VonMisesFisher(z_theta_mu, z_theta_kappa)
-        q_z_phi = VonMisesFisher(z_phi_mu, z_phi_kappa)
-        p_z = HypersphericalUniform(config.latent_dim - 1, device=config.device)
+        q_z_theta = von_mises_fisher.VonMisesFisher(z_theta_mu, z_theta_kappa)
+        q_z_phi = von_mises_fisher.VonMisesFisher(z_phi_mu, z_phi_kappa)
+        p_z = hyperspherical_uniform.HypersphericalUniform(
+            config.latent_dim - 1, device=config.device)
         kld_theta = torch.distributions.kl.kl_divergence(q_z_theta, p_z).mean()
         kld_phi = torch.distributions.kl.kl_divergence(q_z_phi, p_z).mean()
         kld = kld_theta + kld_phi
@@ -149,7 +154,8 @@ def moving_forward_loss(z, config):
     """
     if config.dataset_name != "experimental":
         # print(
-        #     "WARNING: Moving forward loss only implemented for experimental data --> Skipped."
+        #     "WARNING: Moving forward loss only implemented for experimental data
+        #     --> Skipped."
         # )
         return torch.zeros(1).to(config.device)
     if len(z) == 1:
@@ -176,7 +182,8 @@ def dynamic_loss(labels, z, config):
     """
     if config.dataset_name != "experimental":
         # print(
-        #     "WARNING: Dynamic loss only implemented for experimental data --> Skipped."
+        #     "WARNING: Dynamic loss only implemented for experimental data
+        #     --> Skipped."
         # )
         return torch.zeros(1).to(config.device)
     latent_angles = (torch.atan2(z[:, 1], z[:, 0]) + 2 * torch.pi) % (2 * torch.pi)

neurometry/curvature/losses.py

@@ -4,20 +4,11 @@
 import json
 import logging
 import os
-import time
-
-os.environ["GEOMSTATS_BACKEND"] = "pytorch"
 import random
+import time
 
-import datasets.utils
-import default_config
-import evaluate
-import geomstats.backend as gs
 import matplotlib
 import matplotlib.pyplot as plt
-import models.klein_bottle_vae
-import models.neural_vae
-import models.toroidal_vae
 import numpy as np
 import pandas as pd
 
@@ -30,6 +21,15 @@
 from ray.tune.schedulers import AsyncHyperBandScheduler
 from ray.tune.search.hyperopt import HyperOptSearch
 
+os.environ["GEOMSTATS_BACKEND"] = "pytorch"
+import datasets.utils  # noqa: E402
+import default_config  # noqa: E402
+import evaluate  # noqa: E402
+import geomstats.backend as gs  # noqa: E402
+import models.klein_bottle_vae  # noqa: E402
+import models.neural_vae  # noqa: E402
+import models.toroidal_vae  # noqa: E402
+
 # Required to make matplotlib figures in threads:
 matplotlib.use("Agg")
 
@@ -92,7 +92,8 @@ def main():
                         f"Manifold cannot be embedded in {embedding_dim} dimensions"
                     )
                     continue
-                sweep_name = f"{dataset_name}_noise_var_{noise_var}_embedding_dim_{embedding_dim}"
+                sweep_name = f"{dataset_name}_noise_var_{noise_var}"
+                sweep_name += f"_embedding_dim_{embedding_dim}"
                 logging.info(f"\n---> START training for ray sweep: {sweep_name}.")
                 main_sweep(
                     sweep_name=sweep_name,
@@ -120,7 +121,8 @@ def main():
                 default_config.field_width,
                 default_config.resolution,
             ):
-                sweep_name = f"{dataset_name}_orientation_std_{grid_orientation_std}_ncells_{n_cells}"
+                sweep_name = f"{dataset_name}_orientation_std_{grid_orientation_std}"
+                sweep_name += f"_ncells_{n_cells}"
                 logging.info(f"\n---> START training for ray sweep: {sweep_name}.")
                 main_sweep(
                     sweep_name=sweep_name,
@@ -219,7 +221,8 @@ def main_sweep(
         "grid_orientation_std": grid_orientation_std,
         "field_width": field_width,
         "resolution": resolution,
-        # Parameters fixed across runs and sweeps (unique value depending on dataset_name):
+        # Parameters fixed across runs and sweeps
+        # (unique value depending on dataset_name):
         "manifold_dim": default_config.manifold_dim[dataset_name],
         "latent_dim": default_config.latent_dim[dataset_name],
         "posterior_type": default_config.posterior_type[dataset_name],

neurometry/curvature/main.py

@@ -629,10 +629,8 @@ def plot_persistence_diagrams(diagrams_df, density=False):
 
     handles, _ = ax.get_legend_handles_labels()
 
-    legend_labels = []
-    legend_labels.append("Infinity")
-    for dimension in plot_df["Dimension"].unique():
-        legend_labels.append(f"Dimension {dimension}")
+    legend_labels = ["Infinity"] + [
+        f"Dimension {dimension}" for dimension in plot_df["Dimension"].unique()]
 
     ax.legend(handles, legend_labels, loc="lower right")