Merge pull request #19 from FrederikWarburg/master

Extend nnj functionalities

Author: SkafteNicki

README.md

@@ -125,7 +125,7 @@ If you want to cite the framework feel free to use this (but only if you loved i
 ```bibtex
 @article{software:stochman,
   title={StochMan},
-  author={Nicki S. Detlefsen and Alison Pouplin and Cilie W. Feldager and Cong Geng and Dimitris Kalatzis and Helene Hauschultz and Miguel González Duque and Søren Hauberg},
+  author={Nicki S. Detlefsen and Alison Pouplin and Cilie W. Feldager and Cong Geng and Dimitris Kalatzis and Helene Hauschultz and Miguel González Duque and Frederik Warburg and Marco Miani and Søren Hauberg},
   journal={GitHub. Note: https://github.com/MachineLearningLifeScience/stochman/},
   year={2021}
 }

stochman/curves.py

@@ -39,10 +39,7 @@ def __init__(
             _begin = begin.detach()  # BxD
             _end = end.detach()  # BxD
         else:
-            raise ValueError(
-                "BasicCurve.__init__ requires begin and end points to have "
-                "the same shape"
-            )
+            raise ValueError("BasicCurve.__init__ requires begin and end points to have " "the same shape")
 
         # register begin and end as buffers
         self.register_buffer("begin", _begin)  # BxD
@@ -82,6 +79,7 @@ def plot(
         """
         with torch.no_grad():
             import matplotlib.pyplot as plt
+
             t = torch.linspace(t0, t1, N, dtype=self.begin.dtype, device=self.device)
             points = self(t)  # NxD or BxNxD
 
@@ -126,7 +124,7 @@ def euclidean_length(self, t0: float = 0.0, t1: float = 1.0, N: int = 100) -> to
         if not is_batched:
             points = points.unsqueeze(0)  # 1xNxD
         delta = points[:, 1:] - points[:, :-1]  # Bx(N-1)xD
-        energies = (delta ** 2).sum(dim=2)  # Bx(N-1)
+        energies = (delta**2).sum(dim=2)  # Bx(N-1)
         lengths = energies.sqrt().sum(dim=1)  # B
         return lengths
 

stochman/manifold.py

@@ -444,7 +444,7 @@ def curve_energy(self, curve: BasicCurve, reduction: Optional[str] = "sum", dt=N
         emb_curve = self.embed(curve)  # BxNxD
         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
+        energy = (delta**2).sum((1, 2)) * dt  # B
         return tensor_reduction(energy, reduction)
 
     def curve_length(self, curve: BasicCurve, dt=None):
@@ -541,7 +541,7 @@ def __init__(self, data, sigma, rho, device=None):
         """
         super().__init__()
         self.data = data
-        self.sigma2 = sigma ** 2
+        self.sigma2 = sigma**2
         self.rho = rho
         self.device = device
 
@@ -586,7 +586,7 @@ def metric(self, c, return_deriv=False):
             if return_deriv:
                 weighted_delta = (w_p / sigma2).reshape(-1, 1).expand(-1, D) * delta  # NxD
                 dSdc = 2.0 * torch.diag(w_p.mm(delta).flatten()) - weighted_delta.t().mm(delta2)  # DxD
-                dM = dSdc.t() * (m ** 2).reshape(-1, 1).expand(-1, D)  # DxD
+                dM = dSdc.t() * (m**2).reshape(-1, 1).expand(-1, D)  # DxD
                 dMdc.append(dM.reshape(1, D, D))
 
         if return_deriv: