Returns self in fit() methods. (#15)

ahwillia · Jun 29, 2021 · 2dc1f76 · 2dc1f76
1 parent bd7a24b
commit 2dc1f76
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/affinewarp/multiwarp.py b/affinewarp/multiwarp.py
@@ -57,6 +57,8 @@ def fit(self, data, iterations=20, verbose=True):
             self.loss_hist.append(
                 np.mean((self.predict() - data)**2))
 
+        return self
+
     def predict(self, skip=[]):
         K = len(self.models[0].shifts)
         T, N = self.models[0].template.shape

diff --git a/affinewarp/piecewisewarp.py b/affinewarp/piecewisewarp.py
@@ -174,6 +174,8 @@ def fit(self, data, iterations=50, warp_iterations=200, verbose=True,
                 pbar.set_description(
                     "Loss improvement: {0:.2f}%".format(rel_imp))
 
+        return self
+
     def _fit_warps(self, data, warp_iterations, neuron_idx):
         """Fit warping functions by local random search.
 

diff --git a/affinewarp/shiftwarp.py b/affinewarp/shiftwarp.py
@@ -154,6 +154,8 @@ def fit(self, data, iterations=20, verbose=True, warp_iterations=None,
         self.fractional_shifts = self.shifts / T
         self._losses = None
 
+        return self
+
     def _fit_warps(self, data):
         """Updates shift parameters."""