Working plot_fits method (#100)

gpfit/fit.py

@@ -76,17 +76,21 @@ def residual(self, params):
 
     def plot_fit(self):
         """Plots fit"""
-        cstrt, _ = fit(np.log(udata), np.log(wdata), K, fitclass)
-        uu = np.linspace(min(udata), max(udata), 1000)
-        WW = super().plot_fit_data(cstrt, uu)
-        stringlist = super().print_fit()
         f, ax = plt.subplots()
+        udata = np.exp(self.xdata)
+        wdata = np.exp(self.ydata)
         ax.plot(udata, wdata, "+r")
-        for ww in WW:
-            ax.plot(uu, ww)
+
+        xx = np.linspace(min(self.xdata), max(self.xdata), 10)
+        yy, _ = self.evaluate(xx, self.params)
+        uu = np.exp(xx)
+        ww = np.exp(yy)
+        ax.plot(uu, ww)
+
+        stringlist = self.print_fit()
         ax.set_xlabel("u")
+        ax.set_ylabel("w")
         ax.legend(["Data"] + stringlist, loc="best")
-
         return f, ax
 
 
@@ -162,16 +166,6 @@ def print_fit(self):
             print(print_string)
         return string_list
 
-    #def plot_fit_data(self):
-    #    """Return data for plotting"""
-    #    (uvarkey,) = cstrt[0].left.varkeys
-    #    A = [c.left.exps[0][uvarkey] for c in cstrt]
-    #    B = np.log([c.left.cs[0] for c in cstrt])
-    #    WW = []
-    #    for k in range(K):
-    #        WW += [np.exp(B[k])*uu**A[k]]
-    #    return WW
-
 
 class SoftmaxAffine(Fit):
     """Softmax Affine fit class"""
@@ -255,20 +249,6 @@ def print_fit(self):
         print(print_string)
         return ["".join(string_list)]
 
-    #def plot_fit_data(self):
-    #    """Return data for plotting"""
-    #    (wexps,) = cstrt[0].left.exps
-    #    alpha = list(wexps.values())[0]
-    #    (uvarkey,) = cstrt[0].right.varkeys
-    #    A = [d[uvarkey]/alpha for d in cstrt[0].right.exps]
-    #    B = np.log(cstrt[0].right.cs)/alpha
-
-    #    ww = 0
-    #    for k in range(K):
-    #        ww += np.exp(alpha*B[k])*uu**(alpha*A[k])
-    #    WW = [ww**(1/alpha)]
-    #    return WW
-
 
 class ImplicitSoftmaxAffine(Fit):
     """Implicit Softmax Affine fit class"""
@@ -351,7 +331,3 @@ def print_fit(self):
             string_list[k] = print_string
         print(print_string)
         return ["".join(string_list)]
-
-    def plot_fit_data(self):
-        """Return data for plotting"""
-        raise NotImplementedError

gpfit/tests/run_tests.py

@@ -18,8 +18,8 @@
 from gpfit.tests import t_print_fit
 TESTS += t_print_fit.TESTS
 
-#from gpfit.tests import t_plot_fit
-#TESTS += t_plot_fit.TESTS
+from gpfit.tests import t_plot_fit
+TESTS += t_plot_fit.TESTS
 
 from gpfit.tests import t_examples
 TESTS += t_examples.TESTS

gpfit/tests/t_plot_fit.py

@@ -1,17 +1,38 @@
-"unit tests for gpfit.print_fit module"
+"unit tests for plot fit methods"
 import unittest
 import numpy as np
-from gpfit.plot_fit import plot_fit_1d
+from gpfit.fit import MaxAffine, SoftmaxAffine, ImplicitSoftmaxAffine
+
+SEED = 33404
 
 
 class TestPlotFit(unittest.TestCase):
-    "Unit tests for plot_fit_1d"
+    "Unit tests for plot_fit methods"
+
+    np.random.seed(SEED)
     N = 51
     u = np.logspace(0, np.log10(3), N)
     w = (u**2 + 3)/(u + 1)**2
-
-    def test_plot_fit_1d(self):
-        plot_fit_1d(self.u, self.w, K=2, fitclass='SMA', plotspace="linear")
+    x = np.log(u)
+    y = np.log(w)
+    K = 2
+
+    def test_max_affine(self):
+        f = MaxAffine(self.x, self.y, self.K)
+        fig, _ = f.plot_fit()
+        fig.savefig("plots/ma_test.png")
+
+    def test_softmax_affine(self):
+        f = SoftmaxAffine(self.x, self.y, self.K)
+        f.plot_fit()
+        fig, _ = f.plot_fit()
+        fig.savefig("plots/sma_test.png")
+
+    def test_implicit_softmax_affine(self):
+        f = ImplicitSoftmaxAffine(self.x, self.y, self.K)
+        f.plot_fit()
+        fig, _ = f.plot_fit()
+        fig.savefig("plots/isma_test.png")
 
 
 TESTS = [TestPlotFit]