introduce reduction arg and move utils around

Author: SkafteNicki

stochman/manifold.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 from abc import ABC, abstractmethod
 from typing import Optional, Tuple, Union
 
@@ -8,7 +7,7 @@
 
 from stochman.curves import BasicCurve, CubicSpline
 from stochman.geodesic import geodesic_minimizing_energy, shooting_geodesic
-from stochman.utilities import squared_manifold_distance
+from stochman.utils import squared_manifold_distance, tensor_reduction
 
 
 class Manifold(ABC):
@@ -17,13 +16,14 @@ class Manifold(ABC):
     from this abstract base class abstraction.
     """
 
-    def curve_energy(self, curve: BasicCurve) -> torch.Tensor:
+    def curve_energy(self, curve: BasicCurve, reduction: Optional[str] = "sum") -> torch.Tensor:
         """
         Compute the discrete energy of a given curve.
 
-        Input:
-            curve:      a Nx(d) torch Tensor representing a curve or
-                        a BxNx(d) torch Tensor representing B curves.
+        Args:
+            curve: a Nx(d) torch Tensor representing a curve or a BxNx(d) torch Tensor representing B curves.
+            reduction: how to reduce the curve energy over the batch dimension. Choose between
+                `'sum'`, `'mean'`, `'none'` or `None` (where the last two will return the individual scores)
 
         Output:
             energy:     a scalar corresponding to the energy of
@@ -44,7 +44,7 @@ def curve_energy(self, curve: BasicCurve) -> torch.Tensor:
         delta = curve[:, 1:] - curve[:, :-1]  # Bx(N-1)x(d)
         flat_delta = delta.view(-1, d)  # (B*(N-1))x(d)
         energy = self.inner(curve[:, :-1].reshape(-1, d), flat_delta, flat_delta)  # B*(N-1)
-        return energy.sum()  # scalar
+        return tensor_reduction(energy, reduction)
 
     def curve_length(self, curve: BasicCurve) -> torch.Tensor:
         """
@@ -417,13 +417,14 @@ class EmbeddedManifold(Manifold, ABC):
     should inherit from this abstract base class abstraction.
     """
 
-    def curve_energy(self, curve: BasicCurve, dt=None):
+    def curve_energy(self, curve: BasicCurve, reduction: Optional[str] = "sum", dt=None):
         """
         Compute the discrete energy of a given curve.
 
-        Input:
-            curve:      a Nx(d) torch Tensor representing a curve or
-                        a BxNx(d) torch Tensor representing B curves.
+        Args:
+            curve: a Nx(d) torch Tensor representing a curve or a BxNx(d) torch Tensor representing B curves.
+            reduction: how to reduce the curve energy over the batch dimension. Choose between
+                `'sum'`, `'mean'`, `'none'` or `None` (where the last two will return the individual scores)
 
         Output:
             energy:     a scalar corresponding to the energy of
@@ -444,8 +445,7 @@ def curve_energy(self, curve: BasicCurve, dt=None):
         B, N, D = emb_curve.shape
         delta = emb_curve[:, 1:, :] - emb_curve[:, :-1, :]  # Bx(N-1)xD
         energy = (delta ** 2).sum((1, 2)) * dt  # B
-
-        return energy
+        return tensor_reduction(energy, reduction)
 
     def curve_length(self, curve: BasicCurve, dt=None):
         """

stochman/utilities/__init__.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+from multiprocessing.sharedctypes import Value
 import torch
 
 
@@ -51,3 +51,16 @@ def squared_manifold_distance(manifold, p0: torch.Tensor, p1: torch.Tensor):
     """
     distance_op = __Dist2__()
     return distance_op.apply(manifold, p0, p1)
+
+
+def tensor_reduction(x: torch.Tensor, reduction: str):
+    if reduction == 'sum':
+        return x.sum()
+    elif reduction == 'mean':
+        return x.mean()
+    elif reduction is None or reduction == "none":
+        return x
+    else:
+        raise ValueError(
+            f"Expected `reduction` to either be `'mean'`, `'sum'`, `'none'` or `None` but got {reduction}"
+        )