Merge pull request #12 from MachineLearningLifeScience/plotting-in-curves

Adds plotting in individual axes

Author: SkafteNicki

stochman/curves.py

@@ -3,6 +3,7 @@
 from typing import Optional, Tuple
 
 import torch
+from matplotlib.axis import Axis
 from torch import nn
 
 
@@ -47,7 +48,9 @@ def __len__(self):
         """Returns the batch dimension e.g. the number of curves"""
         return self.begin.shape[0]
 
-    def plot(self, t0: float = 0.0, t1: float = 1.0, N: int = 100, *plot_args, **plot_kwargs):
+    def plot(
+        self, t0: float = 0.0, t1: float = 1.0, N: int = 100, ax: Axis = None, *plot_args, **plot_kwargs
+    ):
         """Plot the curve.
 
         Args:
@@ -70,15 +73,20 @@ def plot(self, t0: float = 0.0, t1: float = 1.0, N: int = 100, *plot_args, **plo
             if len(points.shape) == 2:
                 points.unsqueeze_(0)  # 1xNxD
 
+            plot_in = ax or plt
+            if ax is not None:
+                t = t.detach().numpy()
+                points = points.detach().numpy()
+
             figs = []
             if points.shape[-1] == 1:
                 for b in range(points.shape[0]):
-                    fig = plt.plot(t, points[b], *plot_args, **plot_kwargs)
+                    fig = plot_in.plot(t, points[b], *plot_args, **plot_kwargs)
                     figs.append(fig)
                 return figs
             if points.shape[-1] == 2:
                 for b in range(points.shape[0]):
-                    fig = plt.plot(points[b, :, 0], points[b, :, 1], *plot_args, **plot_kwargs)
+                    fig = plot_in.plot(points[b, :, 0], points[b, :, 1], *plot_args, **plot_kwargs)
                     figs.append(fig)
                 return figs
 

tests/test_curves.py

@@ -61,13 +61,13 @@ def test_plot_func(self, curve_class, dim):
 
     @pytest.mark.parametrize("batch_size", [1, 5])
     def test_fit_func(self, curve_class, batch_size):
-        """ test fit function """
+        """test fit function"""
         c = curve_class(torch.randn(batch_size, 2), torch.randn(batch_size, 2), 20)
         loss = c.fit(torch.linspace(0, 1, 10), torch.randn(5, 10, 2))
         assert isinstance(loss, torch.Tensor)
 
     def test_getindex_func(self, curve_class):
-        """ test __getidx__ function """
+        """test __getidx__ function"""
         batched_c = curve_class(torch.randn(5, 2), torch.randn(5, 2))
         for i in range(len(batched_c)):
             c = batched_c[i]
@@ -77,14 +77,14 @@ def test_getindex_func(self, curve_class):
             assert c.device == batched_c.device
 
     def test_setindex_func(self, curve_class):
-        """ test __setidx__ function """
+        """test __setidx__ function"""
         batched_c = curve_class(torch.randn(5, 2), torch.randn(5, 2))
         for i in range(len(batched_c)):
             batched_c[i] = curve_class(torch.randn(1, 2), torch.randn(1, 2))
             assert batched_c[i]
 
     def test_to_other(self, curve_class):
-        """ test .tospline and .todiscrete """
+        """test .tospline and .todiscrete"""
         c = curve_class(torch.randn(1, 2), torch.randn(1, 2), 20)
         if curve_class == curves.DiscreteCurve:
             new_c = c.tospline()
@@ -112,3 +112,19 @@ def test_constant_speed(self, curve_class):
         assert isinstance(Ct, torch.Tensor)
         assert new_t.shape == (batch_size, timesteps)
         assert Ct.shape == (batch_size, timesteps, dim)
+
+    def test_plotting_in_axis(self, curve_class):
+        batch_size = 5
+        dim = 2
+        begin = torch.randn(batch_size, dim)
+        end = torch.randn(batch_size, dim)
+        c = curve_class(begin, end, 20)
+        try:
+            import torchplot as plt
+
+            fig, ax = plt.subplots(1, 1)
+            c.plot(ax=ax)
+            plt.close(fig)
+            assert True
+        except Exception as e:
+            assert False, e