Merge pull request #89 from e3nn/so3

Minor fixes for SO3Signal

Author: ameya98

e3nn_jax/_src/so3grid.py

@@ -127,15 +127,18 @@ def __truediv__(self, other: float) -> "SO3Signal":
     def integrate_over_angles(self) -> SphericalSignal:
         """Integrate the signal over the angles in the axis-angle parametrization."""
         # Account for angle-dependency in Haar measure.
-        grid_values = self.s2_signals.grid_values * (1 - jnp.cos(self.grid_theta))[..., None, None]
+        grid_values = (
+            self.s2_signals.grid_values
+            * (1 - jnp.cos(self.grid_theta))[..., None, None]
+        )
 
         # Trapezoidal rule for integration.
         delta_theta = self.grid_theta[1] - self.grid_theta[0]
         return self.s2_signals.replace_values(
             grid_values=jnp.sum(grid_values, axis=-3) * delta_theta
         )
 
-    def integrate(self) -> SphericalSignal:
+    def integrate(self) -> float:
         """Numerically integrate the signal over SO(3)."""
         # Integrate over angles.
         s2_signal_integrated = self.integrate_over_angles()
@@ -153,7 +156,7 @@ def integrate(self) -> SphericalSignal:
         integral = integral / (8 * jnp.pi**2)
         return integral
 
-    def sample(self, rng: jax.random.PRNGKey):
+    def sample(self, rng: jax.random.PRNGKey) -> jnp.ndarray:
         """Sample a random rotation from SO(3) using the given probability distribution."""
         # Integrate over angles.
         s2_signal_integrated = self.integrate_over_angles()

tests/_src/so3grid_test.py

@@ -96,4 +96,3 @@ def test_division_scalar():
     sig2 = sig1 / 2.7
     integral2 = sig2.integrate()
     assert jnp.isclose(integral2, integral1 / 2.7)
-