Fix tests

Author: ninamiolane

neurometry/curvature/hyperspherical/__init__.py

@@ -1,2 +0,0 @@
-import hyperspherical.distributions
-import hyperspherical.ops

neurometry/curvature/hyperspherical/distributions/__init__.py

@@ -1,2 +0,0 @@
-from hyperspherical.distributions.hyperspherical_uniform import HypersphericalUniform
-from hyperspherical.distributions.von_mises_fisher import VonMisesFisher

neurometry/curvature/hyperspherical/distributions/von_mises_fisher.py

@@ -1,10 +1,13 @@
 import math
 
 import torch
-from hyperspherical.distributions.hyperspherical_uniform import HypersphericalUniform
-from hyperspherical.ops.ive import ive
 from torch.distributions.kl import register_kl
 
+from neurometry.curvature.hyperspherical.distributions.hyperspherical_uniform import (
+    HypersphericalUniform,
+)
+from neurometry.curvature.hyperspherical.ops.ive import ive
+
 
 class VonMisesFisher(torch.distributions.Distribution):
     arg_constraints = {

neurometry/curvature/hyperspherical/ops/__init__.py

@@ -1 +0,0 @@
-from hyperspherical.ops.ive import ive

neurometry/curvature/losses.py

@@ -1,7 +1,11 @@
 """Losses."""
 
 import torch
-from hyperspherical.distributions import HypersphericalUniform, VonMisesFisher
+
+from neurometry.curvature.hyperspherical.distributions import (
+    HypersphericalUniform,
+    VonMisesFisher,
+)
 
 
 def elbo(x, x_mu, posterior_params, z, labels, config):

notebooks/test_plot.py renamed to notebooks/plot_klein_bottle.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 
-def klein_sphere_points(num_points, radius = 1):
+def klein_sphere_points(num_points, radius=1):
     """
     Generate points on a Klein sphere.
 
@@ -13,8 +13,8 @@ def klein_sphere_points(num_points, radius = 1):
     - points: Array of points on the Klein sphere.
     """
     # Generate random points on a 2D plane
-    theta = np.random.uniform(0, 2*np.pi, num_points)
-    phi = np.random.uniform(0, 2*np.pi, num_points)
+    theta = np.random.uniform(0, 2 * np.pi, num_points)
+    phi = np.random.uniform(0, 2 * np.pi, num_points)
 
     # Parametric equations for a Klein sphere
     x = radius * (np.cos(theta) * (1 + np.sin(phi)))
@@ -34,7 +34,7 @@ def plot_klein_sphere(points):
     fig = plt.figure()
     ax = fig.add_subplot(111, projection="3d")
 
-    ax.scatter(points[:,0], points[:,1], points[:,2], c="r", marker="o")
+    ax.scatter(points[:, 0], points[:, 1], points[:, 2], c="r", marker="o")
 
     ax.set_xlabel("X")
     ax.set_ylabel("Y")
@@ -44,38 +44,42 @@ def plot_klein_sphere(points):
 
     plt.show()
 
+
 # Generate points on the Klein sphere
 num_points = 1000
-scales = [0.5,1,0.5]
+scales = [0.5, 1, 0.5]
 points = klein_sphere_points(num_points, scales)
 
 # Plot the generated points
 plot_klein_sphere(points)
 
-#https://mathworld.wolfram.com/KleinBottle.html
-#Klein Bagel
+# https://mathworld.wolfram.com/KleinBottle.html
+# Klein Bagel
+
 
 def klein_bottle_points(num_points, scale=1):
-    u = np.linspace(0, 2*np.pi, num_points)
-    v = np.linspace(0, 2*np.pi, num_points)
+    u = np.linspace(0, 2 * np.pi, num_points)
+    v = np.linspace(0, 2 * np.pi, num_points)
     U, V = np.meshgrid(u, v)
 
-    X = (scale + np.cos(U/2) *np.sin(V) - np.sin(U/2)*np.sin(2*V))*np.cos(U)
-    Y = (scale + np.cos(U/2) *np.sin(V) - np.sin(U/2)*np.sin(2*V))*np.sin(U)
-    Z = np.sin(U/2)*np.sin(V) + np.cos(U/2)*np.sin(2*V)
+    X = (scale + np.cos(U / 2) * np.sin(V) - np.sin(U / 2) * np.sin(2 * V)) * np.cos(U)
+    Y = (scale + np.cos(U / 2) * np.sin(V) - np.sin(U / 2) * np.sin(2 * V)) * np.sin(U)
+    Z = np.sin(U / 2) * np.sin(V) + np.cos(U / 2) * np.sin(2 * V)
     return X, Y, Z
 
+
 def plot_klein_bottle(X, Y, Z):
     fig = plt.figure()
     ax = fig.add_subplot(111, projection="3d")
-    ax.plot_surface(X, Y, Z, cmap="viridis", alpha = 0.7)
+    ax.plot_surface(X, Y, Z, cmap="viridis", alpha=0.7)
 
     ax.set_xlabel("X")
     ax.set_ylabel("Y")
     ax.set_zlabel("Z")
     ax.set_title("Klein Bottle")
     plt.show()
 
+
 # Example usage
 num_points = 100
 scale = 3

pyproject.toml

@@ -58,7 +58,8 @@ doc = [
     "pydata-sphinx-theme"
 ]
 lint = [
-    "pre-commit"
+    "pre-commit",
+    "ruff"
 ]
 test = [
     "pytest",

tests/test_losses.py renamed to tests/test_curvature.py

@@ -1,7 +1,7 @@
 import numpy as np
 import torch
 
-from neurometry.losses import latent_regularization_loss
+from neurometry.curvature.losses import latent_regularization_loss
 
 
 class AttrDict(dict):